Source

KODE / src / solvers / SimplePGSSolver.cpp

Full commit
/*
  This file is part of the KODE.

    KODE Physics Library
    Copyright (C) 2013-2014  Daniel Kohler Osmari

    KODE is free software: you can redistribute it and/or modify it
    under the terms of EITHER:

        * the GNU Lesser General Public License as published by the
          Free Software Foundation, either version 3 of the License,
          or (at your option) any later version.

        * the Apache License, Version 2.0.

    This program is distributed in the hope that it will be useful,
    but WITHOUT ANY WARRANTY; without even the implied warranty of
    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
    GNU Lesser General Public License and the Apache License for more
    details.

    You should have received a copy of the GNU Lesser General Public
    License along with this program.  If not, see
    <http://www.gnu.org/licenses/>.

    You may obtain a copy of the Apache License at
       http://www.apache.org/licenses/LICENSE-2.0
*/

#include <iostream>
#include <vector>
#include <stdexcept>
#include <string>

#include <kode/solvers/SimplePGSSolver.hpp>

#include <kode/World.hpp>
#include <kode/Body.hpp>
#include <kode/joints/Joint.hpp>

#ifdef HAVE_CONFIG_H
#include <config.h>
#endif


#ifdef HAVE_EIGEN
#include <Eigen/Dense>
#else
#include "MiniEigen.hpp"
#endif

using std::clog;
using std::endl;

/*
 * This is a "trivial" implementation of a SOR[1] PGS[2] LCP[3] solver.
 *
 * It's written to be readable and as a baseline for comparison.
 *
 * The main missing optimization is the matrix A; it can be a quite
 * large dense matrix, while it's formed from a product of sparse matrices.
 *
 * Although for medium to large islands of bodies it's slow, for smaller islands
 * (with 2 bodies) it might actually beat the sparse solver due to optimized SIMD
 * operations from Eigen.
 *
 * [1] Successive Over-Relaxation
 * [2] Projected Gauss-Seidel; it's like plain Gauss-Seidel, but with
 *     the variables bounded below and above.
 * [3] Linear Complementarity Problem
 */




namespace kode {

#ifdef HAVE_EIGEN
    using EMatrix = Eigen::Matrix<Real, Eigen::Dynamic, Eigen::Dynamic>;
    using EVector = Eigen::Matrix<Real, Eigen::Dynamic, 1>;

    template<class T, class U>
    Real
    dot(const T& a, const U& b) noexcept
    {
        return a.dot(b);
    }
#else
    using EMatrix = detail::Matrix<detail::Dynamic, detail::Dynamic>;
    using EVector = detail::Vector<detail::Dynamic>;
#endif

    namespace {

        void solveLCP_SOR(const EMatrix& A,
                          const EVector& b,
                          const EVector& lo,
                          const EVector& hi,
                          const std::vector<int>& findex,
                          EVector& lambda,
                          unsigned maxIterations,
                          Real relaxation)
        {
            const size_t n = b.size();
            //clog << "b: " << b << endl;

            for (unsigned iter = 0; iter<maxIterations; ++iter) {
                //clog << "Iteration: " << iter << endl;
                for (size_t i=0; i<n; ++i) {
                    lambda(i) += relaxation/A(i,i) * ( b(i) - dot(A.row(i), lambda) );

                    Real loBound, hiBound;

                    if (findex[i] == -1) {
                        loBound = lo(i);
                        hiBound = hi(i);
                    } else {
                        // used for friction constraints
                        hiBound = Math::abs(hi(i) * lambda(findex[i]));
                        loBound = -hiBound;
                    }

                    lambda(i) = Math::clamp(lambda(i), loBound, hiBound);

                    //clog << "lambda(" << i << ") = " << lambda(i) << endl;
                }
            }
        }
    }



    void
    SimplePGSSolver::solve(const World& world,
                           std::vector<Body*>& bodies,
                           std::vector<Joint*>& joints)
    {
        /*
         * naive Projected Gauss-Seidel implementation:
         * Solve A lambda = b
         *
         * with the full A matrix: A = J W Jt
         * where:
         *  - J  = Jacobian
         *  - W  = inverse of M, the mass matrix
         *  - Jt = transpose of J
         */

        const size_t nbodies = bodies.size();

        // compute inverse mass for each body, in global coordinates
        // that is, invI = b.R b.invI b.R^t


        EMatrix W(6*nbodies, 6*nbodies);
        EVector v(nbodies*6);
        EVector Fext(nbodies*6);

        W.setZero();

        for (const Body* b : bodies) {
            const Matrix3& R = b->getLocalAxes();
            const Mass& invMass = b->getInvMass();
            const Matrix3 invI = R * invMass.inertia * R.transposed();
            const size_t idx = b->idx*6;

            W(idx+0, idx+0) = invMass.total;
            W(idx+1, idx+1) = invMass.total;
            W(idx+2, idx+2) = invMass.total;

            for (unsigned i=0; i<3; ++i)
                for (unsigned j=0; j<3; ++j)
                    W(idx+3+i, idx+3+j) = invI(i, j);


            v(idx + 0) = b->getLinearVel().x;
            v(idx + 1) = b->getLinearVel().y;
            v(idx + 2) = b->getLinearVel().z;

            v(idx + 3) = b->getAngularVel().x;
            v(idx + 4) = b->getAngularVel().y;
            v(idx + 5) = b->getAngularVel().z;

            Fext(idx + 0) = b->getTotalForce().x;
            Fext(idx + 1) = b->getTotalForce().y;
            Fext(idx + 2) = b->getTotalForce().z;

            Fext(idx + 3) = b->getTotalTorque().x;
            Fext(idx + 4) = b->getTotalTorque().y;
            Fext(idx + 5) = b->getTotalTorque().z;
        }


        // now build J

        // first collect joint information
        size_t nconstraints = 0;
        for (const Joint* j : joints)
            nconstraints += j->size();


        EVector c(nconstraints);
        EVector low(nconstraints);
        EVector high(nconstraints);
        std::vector<int> findex(nconstraints, -1);

        EMatrix J(nconstraints, 6*nbodies);
        J.setZero();
        size_t Jrow = 0;
        for (const Joint* j : joints) {
            unsigned conIdx = 0;
            for (const Constraint& con : *j) {

                // fill in body1 part
                if (j->getBody1()) {
                    size_t col = j->getBody1()->idx * 6;
                    J(Jrow, col++) = con.lin1.x;
                    J(Jrow, col++) = con.lin1.y;
                    J(Jrow, col++) = con.lin1.z;

                    J(Jrow, col++) = con.ang1.x;
                    J(Jrow, col++) = con.ang1.y;
                    J(Jrow, col++) = con.ang1.z;
                }

                if (j->getBody2()) {
                    // fill in body2 part
                    size_t col = j->getBody2()->idx * 6;
                    J(Jrow, col++) = con.lin2.x;
                    J(Jrow, col++) = con.lin2.y;
                    J(Jrow, col++) = con.lin2.z;

                    J(Jrow, col++) = con.ang2.x;
                    J(Jrow, col++) = con.ang2.y;
                    J(Jrow, col++) = con.ang2.z;
                }

                // store c, for the right-hand side
                c(Jrow) = con.c * world.getFPS();

                low(Jrow)  = con.low;
                high(Jrow) = con.high;

                if (con.findex && con.findex <= conIdx)
                    findex[Jrow] = Jrow - con.findex;

                ++Jrow;
                ++conIdx;
            }
        }

        // now finally compute A = J W Jt
        EMatrix A = J * W * J.transpose();

        // tweak: add CFM/stepSize to A's diagonal
        Jrow = 0;
        for (const Joint* j : joints) {
            for (const Constraint& con : *j) {
                A(Jrow, Jrow) += con.cfm * world.getFPS();
                ++Jrow;
            }
        }

        // now for the right-hand side: c - J (v/stepSize + W Fext)

        EVector rhs = c - J * (world.getFPS() * v + W * Fext);


        EVector lambda(nconstraints);
        lambda.setZero();

        solveLCP_SOR(A, rhs,
                     low, high,
                     findex,
                     lambda,
                     maxIterations,
                     relaxation);

        // write down the lambdas, useful for warm-starting and to retrieve
        // joint feedbacks
        Jrow = 0;
        for (Joint* j : joints)
            for (Constraint& con : *j)
                con.lambda = lambda(Jrow++);

        // constraint force = Jt lambda
        const EVector cforce = J.transpose()*lambda;

        // accumulate
        for (Body* b : bodies) {
            b->addForce (cforce(b->idx*6+0),
                         cforce(b->idx*6+1),
                         cforce(b->idx*6+2));
            b->addTorque(cforce(b->idx*6+3),
                         cforce(b->idx*6+4),
                         cforce(b->idx*6+5));
        }
    }



    std::string
    SimplePGSSolver::getName() const noexcept
    {
        return "Simple Projected Gauss-Seidel Solver";
    }

}