imstkRigidBodyModel2.h

/*
** This file is part of the Interactive Medical Simulation Toolkit (iMSTK)
** iMSTK is distributed under the Apache License, Version 2.0.
** See accompanying NOTICE for details.
*/

#pragma once

#include "imstkDynamicalModel.h"
#include "imstkRigidBodyState2.h"

#include <list>
#include <unordered_map>

namespace imstk
{
template<typename ScalarType>
class ProjectedGaussSeidelSolver;
class RbdConstraint;
struct RigidBody;

struct RigidBodyModel2Config
{
    double m_dt     = 0.001; 
    Vec3d m_gravity = Vec3d(0.0, -9.8, 0.0);
    unsigned int m_maxNumIterations = 10;
    double m_velocityDamping = 1.0;
    double m_angularVelocityDamping = 1.0;
    double m_epsilon = 1e-4;
    int m_maxNumConstraints = -1;
};

class RigidBodyModel2 : public DynamicalModel<RigidBodyState2>
{
public:
    using StorageIndex = Eigen::SparseMatrix<double>::StorageIndex;

    RigidBodyModel2();
    ~RigidBodyModel2() override = default;

    virtual void setTimeStep(const double timeStep) override { m_config->m_dt = timeStep; }

    virtual double getTimeStep() const override { return m_config->m_dt; }

    std::shared_ptr<RigidBodyModel2Config> getConfig() const { return m_config; }
    const std::list<std::shared_ptr<RbdConstraint>>& getConstraints() const { return m_constraints; }
    std::shared_ptr<ProjectedGaussSeidelSolver<double>> getSolver() const { return m_pgsSolver; }

    std::shared_ptr<RigidBody> addRigidBody();

    void addConstraint(std::shared_ptr<RbdConstraint> constraint) { m_constraints.push_back(constraint); }

    void removeRigidBody(std::shared_ptr<RigidBody> body);

    bool initialize() override;

    void updateMass();

    void configure(std::shared_ptr<RigidBodyModel2Config> config);

    void computeTentativeVelocities();

    void solveConstraints();

    void integrate();

    std::shared_ptr<TaskNode> getComputeTentativeVelocitiesNode() const { return m_computeTentativeVelocities; }
    std::shared_ptr<TaskNode> getSolveNode() const { return m_solveNode; }
    std::shared_ptr<TaskNode> getIntegrateNode() const { return m_integrateNode; }

protected:
    void initGraphEdges(std::shared_ptr<TaskNode> source, std::shared_ptr<TaskNode> sink) override;

    std::shared_ptr<RigidBodyModel2Config> m_config;

    std::shared_ptr<TaskNode> m_computeTentativeVelocities;
    std::shared_ptr<TaskNode> m_solveNode;
    std::shared_ptr<TaskNode> m_integrateNode;

    std::shared_ptr<ProjectedGaussSeidelSolver<double>> m_pgsSolver;
    Eigen::SparseMatrix<double> m_Minv;
    std::list<std::shared_ptr<RbdConstraint>>    m_constraints;
    std::vector<std::shared_ptr<RigidBody>>      m_bodies;
    std::unordered_map<RigidBody*, StorageIndex> m_locations;
    bool   m_modified = true;
    size_t m_maxBodiesParallel = 10; // After 10 bodies, parallel for's are used

    Eigen::VectorXd F;               // Reaction forces
};
} // namespace imstk