USD feature

.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,81 @@
+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 set_animation_time(self, duration: float) -> None:
+        """ Set the total animation time of the scene
+
+        Args:
+            duration (_type_): 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/unit_tests/test_utils_USD.py

@@ -0,0 +1,56 @@
+import unittest
+import os
+import sys
+import numpy as np
+
+sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+
+from rainbow.util.USD import USD
+
+
+class TestUSD(unittest.TestCase):
+
+    def setUp(self):
+        self.file_path = "test.usd"
+        self.usd_instance = USD(self.file_path)
+        self.sample_mesh_name = "sample_mesh"
+        self.sample_vertex_positions = np.array(
+            [[0, 0, 0], [1, 0, 0], [0, 1, 0]])
+        self.sample_triangle_faces = np.array([[0, 1, 2]])
+
+    def test_add_mesh(self):
+        """ Test if a mesh can be added to the USD instance.
+        """
+        self.usd_instance.add_mesh(
+            self.sample_mesh_name, self.sample_vertex_positions, self.sample_triangle_faces)
+        self.assertIn(self.sample_mesh_name, self.usd_instance.meshes,
+                      "Mesh was not added successfully.")
+
+    def test_set_mesh_positions(self):
+        """ Test if the positions of a mesh can be set at a given time stamp.
+        """
+        self.usd_instance.add_mesh(
+            self.sample_mesh_name, self.sample_vertex_positions, self.sample_triangle_faces)
+        new_positions = np.array(
+            [[0.1, 0.1, 0.1], [1.1, 0.1, 0.1], [0.1, 1.1, 0.1]])
+        time_stamp = 1.0
+        self.usd_instance.set_mesh_positions(
+            self.sample_mesh_name, new_positions, time_stamp)
+
+    def test_set_mesh_positions_with_invalid_name(self):
+        """ Test if the mesh does not exist, an error will be raised.
+        """
+        with self.assertRaises(ValueError):
+            self.usd_instance.set_mesh_positions(
+                "invalid_mesh_name", self.sample_vertex_positions, 1.0)
+
+    def test_set_animation_time(self):
+        """ Test if the animation time can be set successfully.
+        """
+        duration = 10.0
+        self.usd_instance.set_animation_time(duration)
+        self.assertEqual(self.usd_instance.stage.GetEndTimeCode(
+        ), duration, "Animation time was not set correctly.")
+
+if __name__ == "__main__":
+    unittest.main()