Merge pull request #109 from Jax922/feature/USD

Support to record the simulation to a usd file and soft body demos

.github/workflows/python-package-conda.yml

@@ -33,6 +33,7 @@ jobs:
             conda install pyparsing
             conda install -c anaconda ipython_genutils 
             conda install -c conda-forge meshplot
+            pip install usd-core
             coverage run  -m unittest python/unit_tests/test_*.py
             coverage report
       - name: Format the code

.gitignore

@@ -137,3 +137,12 @@ dmypy.json
 cpp/external/Eigen3/
 __init__.py
 data/test.mat
+
+# usd*.files
+*.usd*
+
+# Profiling files
+*.prof
+
+# nohup log
+nohup.out

python/rainbow/util/USD.py

@@ -0,0 +1,105 @@
+from pxr import Usd, UsdGeom, Gf, Vt
+import numpy as np
+from numpy.typing import ArrayLike
+
+
+class USD:
+    """
+    Represents a Universal Scene Description (USD) for 3D data.
+
+    Universal Scene Description (USD) is a file format developed by Pixar for representing 3D data. 
+    This class facilitates recording simulation processes in USD format and subsequently writing to a USD file.
+
+    Note:
+        Currently, the class supports only the conversion of position changes in the mesh. Advanced features such as
+        camera, lights, color, etc. are not yet supported.
+
+    Attributes:
+        stage (Usd.Stage): The primary container for composing and interrogating scene description.
+        xform (UsdGeom.Xform): A transform applied to the scene.
+        meshes (dict): A dictionary of meshes added to the scene, with mesh names as keys.
+        file_path (str): Path to save the USD file.
+
+    Example:
+        >>> usd_instance = USD("path/to/save/file.usd")
+        >>> usd_instance.add_mesh("sample_mesh", vertex_positions, triangle_faces)
+        >>> usd_instance.set_mesh_positions("sample_mesh", new_positions, time_stamp)
+        >>> usd_instance.save()
+
+    See Also:
+        [Universal Scene Description (USD) Docs](https://graphics.pixar.com/usd/docs/index.html)
+    """
+
+    def __init__(self, file_path: str):
+        self.stage = Usd.Stage.CreateInMemory(file_path)
+        self.xform = UsdGeom.Xform.Define(self.stage, '/scene/xform')
+        self.meshes = dict()
+        self.file_path = file_path
+
+    def add_mesh(self, name: str, V: ArrayLike, T: ArrayLike) -> None:
+        """ Add a mesh to the scene(or called the stage in USD)
+
+        Args:
+            name (str): The name of the mesh
+            V (ArrayLike): The vertex positions of the mesh
+            T (ArrayLike): The triangle faces of the mesh
+        """
+        mesh = UsdGeom.Mesh.Define(self.stage, f'/scene/xform/{name}')
+        mesh.CreatePointsAttr(V)
+        mesh.CreateFaceVertexCountsAttr([len(face) for face in T])
+        mesh.CreateFaceVertexIndicesAttr(np.concatenate(T))
+        self.meshes[name] = mesh
+
+    def set_mesh_positions(self, name: str, V: ArrayLike, time: float) -> None:
+        """_summary_
+
+        Args:
+            name (str): The name of the mesh
+            V (ArrayLike): The vertex positions of the mesh
+            time (float): The timestamp when the mesh is positioned at V
+
+        Raises:
+            ValueError: If the mesh does not exist in the scene
+        """
+        if name not in self.meshes:
+            raise ValueError(f'Mesh {name} does not exist')
+        vertex_positions = Vt.Vec3fArray(V.tolist())
+        self.meshes[name].GetPointsAttr().Set(vertex_positions, time)
+
+    def get_mesh_positions(self, name: str, time: float) -> ArrayLike:
+        """ Retrieve the positions of a mesh at a given timestamp.
+
+        Args:
+            name (str): The name of the mesh.
+            time (float): The timestamp at which the positions should be retrieved.
+
+        Returns:
+            ArrayLike: An array containing the vertex positions of the mesh.
+
+        Raises:
+            ValueError: If the mesh does not exist in the scene, or if the mesh does not have positions set at the given timestamp.
+        """
+        if name not in self.meshes:
+            raise ValueError(f'Mesh {name} does not exist')
+
+        vertex_positions_attr = self.meshes[name].GetPointsAttr()
+        vertex_positions = vertex_positions_attr.Get(time)
+
+        if vertex_positions:
+            return np.array(vertex_positions, dtype=np.float64)
+        else:
+            raise ValueError(f"No positions set for mesh {name} at time {time}")
+
+    def set_animation_time(self, duration: float) -> None:
+        """ Set the total animation time of the scene
+
+        Args:
+            duration (float): The total animation time of the scene
+        """
+        self.stage.SetStartTimeCode(0)
+        self.stage.SetEndTimeCode(duration)
+
+    def save(self) -> None:
+        """ Save the scene to a USD file
+        """
+        self.stage.GetRootLayer().Export(self.file_path)

python/soft_body_app/hanging_beam_simulation.py

@@ -0,0 +1,33 @@
+from soft_application import SoftBodyApplication
+
+
+def run():
+    """
+    This function provides a demonstration of soft body simulation, specifically,
+    a beam, fixed to a wall on one side, is affected by gravity, causing it to droop or lower on the unfixed side over time.
+
+    Usage:
+        To run the simulation, simply execute this function. Adjust parameters within to modify simulation behaviors,
+        such as the material properties of the beam or the strength of gravity.
+    """
+    app = SoftBodyApplication()
+
+    # Create a soft material named "default_material"
+    app.create_soft_material("default_material", model_name="SNH")
+    app.set_friction_coefficient("default_material", "default_material", 0.5)
+
+    # Add the beam and wall into the scene
+    app.create_soft_body("beam", "default_material", (12, 3, 3, 4.0, 1.0, 1.0))
+    app.create_soft_body("wall", "default_material", (2, 2, 2, 0.1, 8.0, 8.0), is_free=False, center_position=(-2.05, 0, 0))
+
+    # Set dirichlet conditions to fix the beam on the wall
+    app.create_dirichlet_conditions("beam", lambda x: x[0] + 1.9)
+
+    # Do simulation loop
+    app.initialize_viewer()
+    app.initialize_usd('./hanging_beam_simulation.usda')
+    app.simulate(T = 5.0, update_viewer = False, non_update_bodies = ["wall"])
+
+
+if __name__ == "__main__":
+    run()

python/soft_body_app/midpoint_traction_beam_simulation.py

@@ -0,0 +1,39 @@
+import numpy as np
+from soft_application import SoftBodyApplication
+
+
+def run():
+    """
+    A simulation scenario in which a beam has both ends fixed. A traction force is applied 
+    to the middle section of the beam.
+
+    Usage:
+        To run the simulation, simply execute this function. Adjust parameters within to modify simulation behaviors,
+        such as the material properties of the beam or the strength of gravity.
+    """
+    app = SoftBodyApplication()
+
+    # Create a soft material named "default_material"
+    app.create_soft_material("default_material", model_name="SNH")
+    app.set_friction_coefficient("default_material", "default_material", 0.5)
+
+    # Add the beam into the scene
+    app.create_soft_body("beam", "default_material", (20, 2, 2, 4.0, 0.1, 0.1), gravity= 0.0)
+
+    # Set traction conditions to push midpoint of the beam upward
+    def center_beam(x):
+        return -1 if x[0] < 0.05 or x[0] > -0.05 else 1
+    app.create_traction_conditions("beam", center_beam, np.array([0, 2000, 0], dtype=np.float64))
+
+    # Set dirichlet conditions to fix dual ends of the beam
+    app.create_dirichlet_conditions("beam", lambda x: x[0] + 1.9)
+    app.create_dirichlet_conditions("beam", lambda x: 1.9 - x[0])
+
+    # Do simulation loop
+    app.initialize_viewer()
+    app.initialize_usd('./midpoint_traction_beam_simulation.usda')
+    app.simulate(T = 5.0, update_viewer = False)
+
+
+if __name__ == "__main__":
+    run()