REFAC: PbdConstraintFunctors and Containers

Source/Constraint/PbdConstraints/imstkPbdConstraint.h

@@ -133,6 +133,4 @@ protected:
 
     std::vector<Vec3d> m_dcdx;
 };
-
-using PBDConstraintVector = std::vector<std::shared_ptr<PbdConstraint>>;
-}
+}
\ No newline at end of file

Source/Constraint/PbdConstraints/imstkPbdConstraintContainer.cpp

+/*=========================================================================
+
+   Library: iMSTK
+
+   Copyright (c) Kitware, Inc. & Center for Modeling, Simulation,
+   & Imaging in Medicine, Rensselaer Polytechnic Institute.
+
+   Licensed under the Apache License, Version 2.0 (the "License");
+   you may not use this file except in compliance with the License.
+   You may obtain a copy of the License at
+
+      http://www.apache.org/licenses/LICENSE-2.0.txt
+
+   Unless required by applicable law or agreed to in writing, software
+   distributed under the License is distributed on an "AS IS" BASIS,
+   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+   See the License for the specific language governing permissions and
+   limitations under the License.
+
+=========================================================================*/
+
+#include "imstkPbdConstraintContainer.h"
+#include "imstkGraph.h"
+
+#include <unordered_map>
+
+namespace imstk
+{
+void
+PbdConstraintContainer::addConstraint(std::shared_ptr<PbdConstraint> constraint)
+{
+    m_constraintLock.lock();
+    m_constraints.push_back(constraint);
+    m_constraintLock.unlock();
+}
+
+void
+PbdConstraintContainer::removeConstraint(std::shared_ptr<PbdConstraint> constraint)
+{
+    m_constraintLock.lock();
+    iterator i = std::find(m_constraints.begin(), m_constraints.end(), constraint);
+    if (i != m_constraints.end())
+    {
+        m_constraints.erase(i);
+    }
+    m_constraintLock.unlock();
+}
+
+void
+PbdConstraintContainer::removeConstraintVertices(std::shared_ptr<std::unordered_set<size_t>> vertices)
+{
+    // Remove constraints that contain the given vertices
+    auto removeConstraintFunc = [&](std::shared_ptr<PbdConstraint> constraint)
+                                {
+                                    for (auto i : constraint->getVertexIds())
+                                    {
+                                        if (vertices->find(i) != vertices->end())
+                                        {
+                                            return true;
+                                        }
+                                    }
+                                    return false;
+                                };
+
+    m_constraintLock.lock();
+    m_constraints.erase(std::remove_if(m_constraints.begin(), m_constraints.end(), removeConstraintFunc),
+        m_constraints.end());
+
+    // Also remove partitioned constraints
+    for (auto& pc : m_partitionedConstraints)
+    {
+        pc.erase(std::remove_if(pc.begin(), pc.end(), removeConstraintFunc), pc.end());
+    }
+
+    m_constraintLock.unlock();
+}
+
+void
+PbdConstraintContainer::eraseConstraint(iterator iter)
+{
+    m_constraintLock.lock();
+    m_constraints.erase(iter);
+    m_constraintLock.unlock();
+}
+
+void
+PbdConstraintContainer::partitionConstraints(const int partitionedThreshold)
+{
+    // Form the map { vertex : list_of_constraints_involve_vertex }
+    std::vector<std::shared_ptr<PbdConstraint>>& allConstraints = m_constraints;
+
+    //std::cout << "---------partitionConstraints: " << allConstraints.size() << std::endl;
+
+    std::unordered_map<size_t, std::vector<size_t>> vertexConstraints;
+    for (size_t constrIdx = 0; constrIdx < allConstraints.size(); ++constrIdx)
+    {
+        const auto& constr = allConstraints[constrIdx];
+        for (const auto& vIds : constr->getVertexIds())
+        {
+            vertexConstraints[vIds].push_back(constrIdx);
+        }
+    }
+
+    // Add edges to the constraint graph
+    // Each edge represent a shared vertex between two constraints
+    Graph constraintGraph(allConstraints.size());
+    for (const auto& kv : vertexConstraints)
+    {
+        const auto& constraints = kv.second;     // the list of constraints for a vertex
+        for (size_t i = 0; i < constraints.size(); ++i)
+        {
+            for (size_t j = i + 1; j < constraints.size(); ++j)
+            {
+                constraintGraph.addEdge(constraints[i], constraints[j]);
+            }
+        }
+    }
+    vertexConstraints.clear();
+
+    // do graph coloring for the constraint graph
+    const auto coloring = constraintGraph.doColoring(Graph::ColoringMethod::WelshPowell);
+    // const auto  coloring = constraintGraph.doColoring(Graph::ColoringMethod::Greedy);
+    const auto& partitionIndices = coloring.first;
+    const auto  numPartitions    = coloring.second;
+    assert(partitionIndices.size() == allConstraints.size());
+
+    std::vector<std::vector<std::shared_ptr<PbdConstraint>>>& partitionedConstraints = m_partitionedConstraints;
+    partitionedConstraints.resize(0);
+    partitionedConstraints.resize(static_cast<size_t>(numPartitions));
+
+    for (size_t constrIdx = 0; constrIdx < partitionIndices.size(); ++constrIdx)
+    {
+        const auto partitionIdx = partitionIndices[constrIdx];
+        partitionedConstraints[partitionIdx].push_back(allConstraints[constrIdx]);
+    }
+
+    // If a partition has size smaller than the partition threshold, then move its constraints back
+    // These constraints will be processed sequentially
+    // Because small size partitions yield bad performance upon running in parallel
+    allConstraints.resize(0);
+    for (const auto& constraints : partitionedConstraints)
+    {
+        if (constraints.size() < partitionedThreshold)
+        {
+            for (size_t constrIdx = 0; constrIdx < constraints.size(); ++constrIdx)
+            {
+                allConstraints.push_back(std::move(constraints[constrIdx]));
+            }
+        }
+    }
+
+    // Remove all empty partitions
+    size_t writeIdx = 0;
+    for (size_t readIdx = 0; readIdx < partitionedConstraints.size(); ++readIdx)
+    {
+        if (partitionedConstraints[readIdx].size() >= partitionedThreshold)
+        {
+            if (readIdx != writeIdx)
+            {
+                partitionedConstraints[writeIdx] = std::move(partitionedConstraints[readIdx]);
+            }
+            ++writeIdx;
+        }
+    }
+    partitionedConstraints.resize(writeIdx);
+
+    // Print
+    /*if (print)
+    {
+        size_t numConstraints = 0;
+        int    idx = 0;
+        for (const auto& constraints : partitionedConstraints)
+        {
+            std::cout << "Partition # " << idx++ << " | # nodes: " << constraints.size() << std::endl;
+            numConstraints += constraints.size();
+        }
+        std::cout << "Sequential processing # nodes: " << allConstraints.size() << std::endl;
+        numConstraints += allConstraints.size();
+        std::cout << "Total constraints: " << numConstraints << " | Graph size: "
+            << constraintGraph.size() << std::endl;
+    }*/
+}
+}
\ No newline at end of file

Source/Constraint/PbdConstraints/imstkPbdConstraintContainer.h

+/*=========================================================================
+
+   Library: iMSTK
+
+   Copyright (c) Kitware, Inc. & Center for Modeling, Simulation,
+   & Imaging in Medicine, Rensselaer Polytechnic Institute.
+
+   Licensed under the Apache License, Version 2.0 (the "License");
+   you may not use this file except in compliance with the License.
+   You may obtain a copy of the License at
+
+      http://www.apache.org/licenses/LICENSE-2.0.txt
+
+   Unless required by applicable law or agreed to in writing, software
+   distributed under the License is distributed on an "AS IS" BASIS,
+   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+   See the License for the specific language governing permissions and
+   limitations under the License.
+
+=========================================================================*/
+
+#pragma once
+
+#include "imstkPbdConstraint.h"
+
+#include <unordered_set>
+
+namespace imstk
+{
+///
+/// \class PbdConstraintContainer
+///
+/// \brief Container for pbd constraints
+///
+class PbdConstraintContainer
+{
+public:
+    PbdConstraintContainer() = default;
+    virtual ~PbdConstraintContainer() = default;
+
+public:
+    using iterator = std::vector<std::shared_ptr<PbdConstraint>>::iterator;
+
+public:
+    ///
+    /// \brief Adds a constraint to the system, thread safe
+    ///
+    virtual void addConstraint(std::shared_ptr<PbdConstraint> constraint);
+
+    ///
+    /// \brief Linear searches for and removes a constraint from the system, thread safe
+    ///
+    virtual void removeConstraint(std::shared_ptr<PbdConstraint> constraint);
+
+    ///
+    /// \brief Removes all constraints associated with vertex ids
+    ///
+    virtual void removeConstraintVertices(std::shared_ptr<std::unordered_set<size_t>> vertices);
+
+    ///
+    /// \brief Removes a constraint from the system by iterator, thread safe
+    ///
+    virtual void eraseConstraint(iterator iter);
+
+    ///
+    /// \brief Reserve an amount of constraints in the pool, if you know
+    /// ahead of time the number of constraints, or even an estimate, it
+    /// can be faster to first reserve them
+    ///
+    virtual void reserve(const size_t n) { m_constraints.reserve(n); }
+
+    ///
+    /// \brief Returns if there are no constraints
+    ///
+    const bool empty() const { return m_constraints.empty() && m_partitionedConstraints.empty(); }
+
+    ///
+    /// \brief Get the underlying container
+    ///
+    const std::vector<std::shared_ptr<PbdConstraint>>& getConstraints() const { return m_constraints; }
+
+    ///
+    /// \brief Get the partitioned constraints
+    ///
+    const std::vector<std::vector<std::shared_ptr<PbdConstraint>>> getPartitionedConstraints() const { return m_partitionedConstraints; }
+
+    ///
+    /// \brief Partitions pbd constraints into separate vectors via graph coloring
+    /// \param Minimum number of constraints in groups, any under will be dumped back into m_constraints
+    ///
+    void partitionConstraints(const int partitionThreshold);
+
+    ///
+    /// \brief Clear the parition vectors
+    ///
+    void clearPartitions() { m_partitionedConstraints.clear(); }
+
+protected:
+    std::vector<std::shared_ptr<PbdConstraint>> m_constraints;                         ///> Not partitioned constraints
+    std::vector<std::vector<std::shared_ptr<PbdConstraint>>> m_partitionedConstraints; ///> Partitioned pbd constraints
+    ParallelUtils::SpinLock m_constraintLock;                                          ///> Used to deal with concurrent addition/removal of constraints
+};
+}
\ No newline at end of file

Source/Constraint/PbdConstraints/imstkPbdVolumeConstraint.cpp

@@ -21,7 +21,7 @@
 
 #include "imstkPbdVolumeConstraint.h"
 
-namespace  imstk
+namespace imstk
 {
 void
 PbdVolumeConstraint::initConstraint(const VecDataArray<double, 3>& initVertexPositions,
@@ -68,7 +68,8 @@ PbdVolumeConstraint::computeValueAndGradient(
     dcdx[2] = onesixth * (x3 - x0).cross(x1 - x0);
     dcdx[3] = onesixth * (x1 - x0).cross(x2 - x0);
 
-    c = dcdx[3].dot(x3 - x0);
+    const double volume = dcdx[3].dot(x3 - x0);
+    c = (volume - m_restVolume);
     return true;
 }
 } // imstk

Source/DynamicalModels/ObjectModels/imstkPbdConstraintFunctor.h

+/*=========================================================================
+
+   Library: iMSTK
+
+   Copyright (c) Kitware, Inc. & Center for Modeling, Simulation,
+   & Imaging in Medicine, Rensselaer Polytechnic Institute.
+
+   Licensed under the Apache License, Version 2.0 (the "License");
+   you may not use this file except in compliance with the License.
+   You may obtain a copy of the License at
+
+      http://www.apache.org/licenses/LICENSE-2.0.txt
+
+   Unless required by applicable law or agreed to in writing, software
+   distributed under the License is distributed on an "AS IS" BASIS,
+   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+   See the License for the specific language governing permissions and
+   limitations under the License.
+
+=========================================================================*/
+
+#pragma once
+
+#include "imstkLineMesh.h"
+#include "imstkLogger.h"
+#include "imstkParallelUtils.h"
+#include "imstkPbdAreaConstraint.h"
+#include "imstkPbdBendConstraint.h"
+#include "imstkPbdConstantDensityConstraint.h"
+#include "imstkPbdConstraint.h"
+#include "imstkPbdDihedralConstraint.h"
+#include "imstkPbdDistanceConstraint.h"
+#include "imstkPbdFEMConstraint.h"
+#include "imstkPbdFETetConstraint.h"
+#include "imstkPbdVolumeConstraint.h"
+#include "imstkPointSet.h"
+#include "imstkSurfaceMesh.h"
+#include "imstkTetrahedralMesh.h"
+#include "imstkPbdConstraintContainer.h"
+
+namespace imstk
+{
+///
+/// \brief PbdConstraintFunctor take input geometry and produce constraints
+///
+struct PbdConstraintFunctor
+{
+    public:
+        PbdConstraintFunctor() = default;
+        virtual ~PbdConstraintFunctor() = default;
+
+        ///
+        /// \brief Appends a set of constraint to the container given a geometry
+        ///
+        virtual void operator()(PbdConstraintContainer& constraints) = 0;
+
+        void setGeometry(std::shared_ptr<PointSet> geom) { m_geom = geom; }
+
+    public:
+        std::shared_ptr<PointSet> m_geom;
+};
+
+struct PbdDistanceConstraintFunctor : public PbdConstraintFunctor
+{
+    public:
+        PbdDistanceConstraintFunctor() = default;
+        ~PbdDistanceConstraintFunctor() override = default;
+
+        virtual void operator()(PbdConstraintContainer& constraints) override
+        {
+            const VecDataArray<double, 3>& vertices = *m_geom->getVertexPositions();
+            auto                           addDistConstraint =
+                [&](std::vector<std::vector<bool>>& E, size_t i1, size_t i2)
+                {
+                    if (i1 > i2) // Make sure i1 is always smaller than i2
+                    {
+                        std::swap(i1, i2);
+                    }
+                    if (E[i1][i2])
+                    {
+                        E[i1][i2] = 0;
+                        auto c = std::make_shared<PbdDistanceConstraint>();
+                        c->initConstraint(vertices, i1, i2, m_stiffness);
+                        constraints.addConstraint(c);
+                    }
+                };
+
+            if (m_geom->getTypeName() == "TetrahedralMesh")
+            {
+                const auto&                    tetMesh  = std::dynamic_pointer_cast<TetrahedralMesh>(m_geom);
+                const VecDataArray<int, 4>&    elements = *tetMesh->getTetrahedraIndices();
+                const auto                     nV       = tetMesh->getNumVertices();
+                std::vector<std::vector<bool>> E(nV, std::vector<bool>(nV, 1));
+
+                for (int k = 0; k < elements.size(); ++k)
+                {
+                    auto& tet = elements[k];
+                    addDistConstraint(E, tet[0], tet[1]);
+                    addDistConstraint(E, tet[0], tet[2]);
+                    addDistConstraint(E, tet[0], tet[3]);
+                    addDistConstraint(E, tet[1], tet[2]);
+                    addDistConstraint(E, tet[1], tet[3]);
+                    addDistConstraint(E, tet[2], tet[3]);
+                }
+            }
+            else if (m_geom->getTypeName() == "SurfaceMesh")
+            {
+                const auto&                    triMesh  = std::dynamic_pointer_cast<SurfaceMesh>(m_geom);
+                const VecDataArray<int, 3>&    elements = *triMesh->getTriangleIndices();
+                const auto                     nV       = triMesh->getNumVertices();
+                std::vector<std::vector<bool>> E(nV, std::vector<bool>(nV, 1));
+
+                for (int k = 0; k < elements.size(); ++k)
+                {
+                    auto& tri = elements[k];
+                    addDistConstraint(E, tri[0], tri[1]);
+                    addDistConstraint(E, tri[0], tri[2]);
+                    addDistConstraint(E, tri[1], tri[2]);
+                }
+            }
+            else if (m_geom->getTypeName() == "LineMesh")
+            {
+                const auto&                    lineMesh = std::dynamic_pointer_cast<LineMesh>(m_geom);
+                const VecDataArray<int, 2>&    elements = *lineMesh->getLinesIndices();
+                const auto&                    nV       = lineMesh->getNumVertices();
+                std::vector<std::vector<bool>> E(nV, std::vector<bool>(nV, 1));
+
+                for (int k = 0; k < elements.size(); k++)
+                {
+                    auto& seg = elements[k];
+                    addDistConstraint(E, seg[0], seg[1]);
+                }
+            }
+        }
+
+        void setStiffness(double stiffness) { m_stiffness = stiffness; }
+
+    protected:
+        double m_stiffness;
+};
+
+struct PbdFemConstraintFunctor : public PbdConstraintFunctor
+{
+    public:
+        PbdFemConstraintFunctor() = default;
+        ~PbdFemConstraintFunctor() override = default;
+
+        virtual void operator()(PbdConstraintContainer& constraints) override
+        {
+            // Check if constraint type matches the mesh type
+            CHECK(m_geom->getTypeName() == "TetrahedralMesh")
+                << "FEM Tetrahedral constraint should come with tetrahedral mesh";
+
+            // Create constraints
+            auto                           tetMesh  = std::dynamic_pointer_cast<TetrahedralMesh>(m_geom);
+            const VecDataArray<double, 3>& vertices = *m_geom->getVertexPositions();
+            const VecDataArray<int, 4>&    elements = *tetMesh->getTetrahedraIndices();
+
+            ParallelUtils::parallelFor(elements.size(),
+                [&](const size_t k)
+            {
+                const Vec4i& tet = elements[k];
+                auto c = std::make_shared<PbdFEMTetConstraint>(m_matType);
+                c->initConstraint(vertices,
+                    tet[0], tet[1], tet[2], tet[3], m_femConfig);
+                constraints.addConstraint(c);
+            }, elements.size() > 100);
+        }
+
+        void setMaterialType(PbdFEMTetConstraint::MaterialType materialType) { m_matType = materialType; }
+        void setFemConfig(std::shared_ptr<PbdFEMConstraintConfig> femConfig) { m_femConfig = femConfig; }
+
+    protected:
+        PbdFEMTetConstraint::MaterialType m_matType;
+        std::shared_ptr<PbdFEMConstraintConfig> m_femConfig;
+};
+
+struct PbdVolumeConstraintFunctor : public PbdConstraintFunctor
+{
+    public:
+        PbdVolumeConstraintFunctor() = default;
+        ~PbdVolumeConstraintFunctor() override = default;
+
+        virtual void operator()(PbdConstraintContainer& constraints) override
+        {
+            // Check if constraint type matches the mesh type
+            CHECK(m_geom->getTypeName() == "TetrahedralMesh")
+                << "Volume constraint should come with volumetric mesh";
+
+            // Create constraints
+            auto                           tetMesh  = std::dynamic_pointer_cast<TetrahedralMesh>(m_geom);
+            const VecDataArray<double, 3>& vertices = *m_geom->getVertexPositions();
+            const VecDataArray<int, 4>&    elements = *tetMesh->getTetrahedraIndices();
+
+            ParallelUtils::parallelFor(elements.size(),
+                [&](const size_t k)
+            {
+                auto& tet = elements[k];
+                auto c    = std::make_shared<PbdVolumeConstraint>();
+                c->initConstraint(vertices,
+                    tet[0], tet[1], tet[2], tet[3], m_stiffness);
+                constraints.addConstraint(c);
+            });
+        }
+
+        void setStiffness(double stiffness) { m_stiffness = stiffness; }
+
+    protected:
+        double m_stiffness;
+};
+
+struct PbdAreaConstraintFunctor : public PbdConstraintFunctor
+{
+    public:
+        PbdAreaConstraintFunctor() = default;
+        ~PbdAreaConstraintFunctor() override = default;
+
+        virtual void operator()(PbdConstraintContainer& constraints) override
+        {
+            // check if constraint type matches the mesh type
+            CHECK(m_geom->getTypeName() == "SurfaceMesh")
+                << "Area constraint should come with a triangular mesh";
+
+            // ok, now create constraints
+            auto                           triMesh  = std::dynamic_pointer_cast<SurfaceMesh>(m_geom);
+            const VecDataArray<double, 3>& vertices = *m_geom->getVertexPositions();
+            const VecDataArray<int, 3>&    elements = *triMesh->getTriangleIndices();
+
+            ParallelUtils::parallelFor(elements.size(),
+                [&](const size_t k)
+            {
+                auto& tri = elements[k];
+                auto c    = std::make_shared<PbdAreaConstraint>();
+                c->initConstraint(vertices, tri[0], tri[1], tri[2], m_stiffness);
+                constraints.addConstraint(c);
+            });
+        }
+
+        void setStiffness(double stiffness) { m_stiffness = stiffness; }
+
+    protected:
+        double m_stiffness;
+};
+
+struct PbdBendConstraintFunctor : public PbdConstraintFunctor
+{
+    public:
+        PbdBendConstraintFunctor() = default;
+        ~PbdBendConstraintFunctor() override = default;
+
+        virtual void operator()(PbdConstraintContainer& constraints) override
+        {
+            CHECK(m_geom->getTypeName() == "LineMesh")
+                << "Bend constraint should come with a line mesh";
+
+            auto                           lineMesh = std::dynamic_pointer_cast<LineMesh>(m_geom);
+            const VecDataArray<double, 3>& vertices = *m_geom->getVertexPositions();
+            const VecDataArray<int, 2>&    elements = *lineMesh->getLinesIndices();
+
+            auto addBendConstraint =
+                [&](const double k, size_t i1, size_t i2, size_t i3)
+                {
+                    // i1 should always come first
+                    if (i2 < i1)
+                    {
+                        std::swap(i1, i2);
+                    }
+                    // i3 should always come last
+                    if (i2 > i3)
+                    {