Merge pull request #48 from empriselab/justin/particle-group-pos

Particle Group Position - Python Interface & Examples

pyrcareworld/Test/test_cloth_anchor.py

@@ -10,13 +10,32 @@
     # Anchors the cloth's particle group "corner" to the object with id 200.
     cloth.addParticleAnchor("corner", 200)
 
-    # Step 100 times.
-    for _ in range(100):
+    cube = env.create_object(id=300, name="Average Indicator", is_in_scene=True)
+
+    def fetch_and_set():
+        """
+        Fetches the hole particle positions and steps the environment,
+        setting the average cube in the Unity world to the average position.
+        """
+        # Fetch particle positions on the next frame. Need to step before reading.
+        cloth.fetchParticlePositions("hole")
         env.step()
 
+        # Also test cloth particle group position retrieval.
+        data = env.instance_channel.data[100]
+        positions = data["particle_groups"]["hole"]
+        average_position = [sum(x) / len(x) for x in zip(*positions)]
+
+        # Set the cube to this average position in the world.
+        cube.setTransform(position=average_position)
+
+    # Step 200 times.
+    for _ in range(200):
+        fetch_and_set()
+
     # Unanchor. Observe the cloth fall.
     cloth.removeParticleAnchor("corner")
 
     # Continue stepping.
     for _ in range(10000):
-        env.step()
+        fetch_and_set()

pyrcareworld/pyrcareworld/agents/cloth.py

@@ -43,3 +43,14 @@ def initializeParticlePositions(self, mappings: dict):
             particle_indices=mappings.keys(),
             positions=mappings.values(),
         )
+
+    def fetchParticlePositions(self, particle_group_name: str):
+        """
+        Fetches the positions of all particles in the specified particle group
+        in the next frame's collect data.
+        """
+        self.env.instance_channel.set_action(
+            "FetchParticlePositions",
+            id=self.id,
+            particle_group_name=particle_group_name,
+        )

pyrcareworld/pyrcareworld/attributes/cloth_attr.py

@@ -7,7 +7,30 @@
 
 
 def parse_message(msg: IncomingMessage) -> dict:
+    """
+    Fetches the same information as a base_attr, but with the additional
+    `"particle_groups"` key. You can do `dict["particle_groups"][particle_group_name]` to get the list of positions, which are each represented as a 3-length list.
+
+    Returns:
+        dict: The same information as base_attr, but with the additional
+        `"particle_groups"` key.
+    """
     this_object_data = attr.base_attr.parse_message(msg)
+    count = msg.read_int32()
+    this_object_data["particle_groups"] = {}
+    for _ in range(count):
+        # First is the particle group name.
+        name = msg.read_string()
+        # Finally is the xs list, then the ys list, then the zs list.
+        xs = msg.read_float32_list()
+        ys = msg.read_float32_list()
+        zs = msg.read_float32_list()
+
+        # Make a couple of lists
+        positions = [[x, y, z] for x, y, z in zip(xs, ys, zs)]
+
+        this_object_data["particle_groups"][name] = positions
+
     return this_object_data
 
 
@@ -71,3 +94,20 @@ def InitializeParticlePositions(kwargs: dict) -> OutgoingMessage:
     msg.write_float32_list(zs)
 
     return msg
+
+
+def FetchParticlePositions(kwargs: dict) -> OutgoingMessage:
+    """
+    Sends a message containing a request to fetch the current positions of all particles in the cloth actor for the specified particle group.
+    """
+    compulsory_params = ["id", "particle_group_name"]
+    utility.CheckKwargs(kwargs, compulsory_params)
+
+    msg = OutgoingMessage()
+
+    msg.write_int32(kwargs["id"])
+    msg.write_string("FetchParticlePositions")
+
+    msg.write_string(kwargs["particle_group_name"])
+
+    return msg