imstkPbdModel.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 "imstkAbstractDynamicalModel.h"
#include "imstkPbdBody.h"
#include "imstkPbdConstraint.h"

#include <unordered_map>
#include <unordered_set>

namespace imstk
{
class PbdConstraintContainer;
class PbdModelConfig;
class PbdSolver;

class PbdModel : public AbstractDynamicalModel
{
public:
    PbdModel();
    ~PbdModel() override = default;

    void resetToInitialState() override;

    void configure(std::shared_ptr<PbdModelConfig> params);

    std::shared_ptr<PbdBody> addBody();
    void removeBody(std::shared_ptr<PbdBody> body);

    std::shared_ptr<PbdBody> getBody(size_t index) const;

    PbdState& getBodies() { return m_state; }

    PbdParticleId addVirtualParticle(
        const Vec3d& pos, const Quatd& orientation,
        const double mass, const Mat3d inertia,
        const Vec3d& velocity = Vec3d::Zero(), const Vec3d& angularVelocity = Vec3d::Zero(),
        const bool persist    = false);

    PbdParticleId addVirtualParticle(
        const Vec3d& pos, const double mass,
        const Vec3d& velocity = Vec3d::Zero(),
        const bool persist    = false);

    void clearVirtualParticles();

    std::shared_ptr<PbdModelConfig> getConfig() const;

    void addConstraints(std::shared_ptr<std::unordered_set<size_t>> vertices, const int bodyId);

    void setTimeStep(const double timeStep) override;
    double getTimeStep() const override;

    void setVelocityThreshold(double velCap) { m_velocityThreshold = velCap; }
    double getVelocityThreshold() { return m_velocityThreshold; }

    std::shared_ptr<PbdConstraintContainer> getConstraints() { return m_constraints; }

    void integratePosition();
    void integratePosition(PbdBody& body);

    void updateVelocity();
    void updateVelocity(PbdBody& body);

    void solveConstraints();

    bool initialize() override;

    void setConstraintPartitionThreshold(size_t threshold) { m_partitionThreshold = threshold; }

    std::shared_ptr<PbdSolver> getSolver() const { return m_pbdSolver; }

    void setSolver(std::shared_ptr<PbdSolver> solver) { this->m_pbdSolver = solver; }

    std::shared_ptr<TaskNode> getIntegratePositionNode() const { return m_integrationPositionNode; }
    std::shared_ptr<TaskNode> getSolveNode() const { return m_solveConstraintsNode; }
    std::shared_ptr<TaskNode> getUpdateVelocityNode() const { return m_updateVelocityNode; }

protected:
    // Hide this function as PbdModel doesn't require it. It can support multiple bodies
    using AbstractDynamicalModel::setModelGeometry;

    void resizeBodyParticles(PbdBody& body, const int particleCount);

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

    size_t m_partitionThreshold = 16; 

    bool m_modified = true;
    int  m_iterKey  = 0;     

    PbdState m_initialState;
    PbdState m_state;

    std::shared_ptr<PbdSolver>      m_pbdSolver = nullptr;     
    std::shared_ptr<PbdModelConfig> m_config    = nullptr;     
    std::shared_ptr<PbdConstraintContainer> m_constraints;     

    std::shared_ptr<TaskNode> m_integrationPositionNode = nullptr;
    std::shared_ptr<TaskNode> m_solveConstraintsNode    = nullptr;
    std::shared_ptr<TaskNode> m_updateVelocityNode      = nullptr;

    double m_velocityThreshold = 100000.0;
};
} // namespace imstk