add rod-gravity-surface contact module

Author: skim0119

src/virtual_field/runtime/foraging_elastica/__init__.py

@@ -7,7 +7,7 @@
     idle_policy_like,
     rotate_policy_by_angle,
 )
-from .forcing import SuckerActuation, YSurfaceBallwGravity
+from .forcing import SuckerActuation, YSurfaceBallwGravity, YSurfaceRodwGravity
 
 __all__ = [
     "BaseSphereTether",
@@ -17,6 +17,7 @@
     "SegmentExtensionActuation",
     "SuckerActuation",
     "YSurfaceBallwGravity",
+    "YSurfaceRodwGravity",
     "current_activation",
     "idle_policy_like",
     "rotate_policy_by_angle",

src/virtual_field/runtime/foraging_elastica/forcing.py

@@ -324,3 +324,112 @@ def _apply_sphere_plane_hertz(
             s = -f_t / speed_t
             sphere_external_force[0, 0] += s * tx
             sphere_external_force[2, 0] += s * tz
+
+
+class YSurfaceRodwGravity(NoForces):
+    """Cosserat rod on a horizontal floor with gravity and dry friction.
+
+    Applies gravity in ``-Y`` at rod nodes. Contact uses the existing
+    rod-plane normal response against ``y = plane_origin`` with fixed ``+Y``
+    surface normal. A local Coulomb-style tangential friction term is applied
+    on contacting elements and distributed to adjacent nodes.
+    """
+
+    def __init__(
+        self,
+        k_c: float,
+        nu_c: float,
+        *,
+        mu: float = 5.0,
+        plane_origin: float = 0.0,
+    ) -> None:
+        super().__init__()
+        self.k_c = float(k_c)
+        self.nu_c = float(nu_c)
+        self.mu = float(mu)
+        self.plane_origin = float(plane_origin)
+        self._plane_origin = np.array(
+            [[0.0], [self.plane_origin], [0.0]], dtype=np.float64
+        )
+        self._plane_normal = np.array([0.0, 1.0, 0.0], dtype=np.float64)
+        self._surface_tol = 1e-4
+
+    def apply_forces(
+        self,
+        system: RodType,
+        time: np.float64 = np.float64(0.0),
+    ) -> None:
+        _apply_uniform_gravity_to_nodes(system.external_forces, system.mass, 9.81)
+        _calculate_contact_forces_rod_plane(
+            self._plane_origin,
+            self._plane_normal,
+            self._surface_tol,
+            self.k_c,
+            self.nu_c,
+            system.radius,
+            system.mass,
+            system.position_collection,
+            system.velocity_collection,
+            system.internal_forces,
+            system.external_forces,
+        )
+        _apply_rod_plane_coulomb_friction(
+            system.position_collection,
+            system.velocity_collection,
+            system.radius,
+            system.external_forces,
+            self.plane_origin,
+            self.k_c,
+            self.nu_c,
+            self.mu,
+        )
+
+
+@njit(cache=True)
+def _apply_uniform_gravity_to_nodes(
+    external_forces: np.ndarray,
+    nodal_mass: np.ndarray,
+    gravity_mag: float,
+) -> None:
+    for i in range(nodal_mass.shape[0]):
+        external_forces[1, i] -= gravity_mag * nodal_mass[i]
+
+
+@njit(cache=True, fastmath=True)
+def _apply_rod_plane_coulomb_friction(
+    rod_positions: np.ndarray,
+    rod_velocities: np.ndarray,
+    rod_radii: np.ndarray,
+    rod_external_forces: np.ndarray,
+    plane_origin: float,
+    k_c: float,
+    nu_c: float,
+    mu: float,
+) -> None:
+    element_count = rod_radii.shape[0]
+    for idx in range(element_count):
+        center_y = 0.5 * (rod_positions[1, idx] + rod_positions[1, idx + 1])
+        penetration = rod_radii[idx] - (center_y - plane_origin)
+        if penetration <= 0.0:
+            continue
+
+        vy = 0.5 * (rod_velocities[1, idx] + rod_velocities[1, idx + 1])
+        normal_force = k_c * penetration ** (3.0 / 2.0) - nu_c * vy
+        if normal_force <= 0.0:
+            continue
+
+        vx = 0.5 * (rod_velocities[0, idx] + rod_velocities[0, idx + 1])
+        vz = 0.5 * (rod_velocities[2, idx] + rod_velocities[2, idx + 1])
+        speed_t_sq = vx * vx + vz * vz
+        if speed_t_sq <= 1.0e-12:
+            continue
+
+        speed_t = np.sqrt(speed_t_sq)
+        scale = -(mu * normal_force) / speed_t
+        fx = scale * vx
+        fz = scale * vz
+
+        rod_external_forces[0, idx] += 0.5 * fx
+        rod_external_forces[2, idx] += 0.5 * fz
+        rod_external_forces[0, idx + 1] += 0.5 * fx
+        rod_external_forces[2, idx + 1] += 0.5 * fz

src/virtual_field/runtime/noel_c4_simulation.py

@@ -23,9 +23,7 @@ class NoelObstacleSet:
 
 def load_noel_c4_obstacles() -> NoelObstacleSet:
     obstacle_path = (
-        Path(__file__).resolve().parents[3]
-        / "externals"
-        / "noel-c4-obstacles-aug.npz"
+        Path(__file__).resolve().parents[3] / "externals" / "noel-c4-obstacles-aug.npz"
         # / "noel-c4-obstacles.npz"
     )
     data = np.load(obstacle_path, allow_pickle=True)
@@ -107,7 +105,7 @@ def load_noel_c4_obstacles() -> NoelObstacleSet:
     ends = starts + directions * lengths[:, None]
     center = np.vstack([starts, ends]).mean(axis=0)
     starts[:, 0] -= 0.10
-    starts[:, 1] += 1.0 - center[1] + 0.25
+    starts[:, 1] += 1.0 - center[1] + 0.15
     starts[:, 2] -= 0.75
 
     return NoelObstacleSet(
@@ -123,9 +121,7 @@ def load_noel_c4_obstacles() -> NoelObstacleSet:
 class NoelC4Simulation(DualArmSimulationBase):
     tip_haptic_max_penetration: float = 0.01
     _obstacles: NoelObstacleSet = field(init=False)
-    _tip_penetration_by_arm: dict[str, float] = field(
-        init=False, default_factory=dict
-    )
+    _tip_penetration_by_arm: dict[str, float] = field(init=False, default_factory=dict)
     _haptic_events: list[HapticEvent] = field(init=False, default_factory=list)
 
     def build_simulation(self) -> None:
@@ -138,8 +134,8 @@ def build_simulation(self) -> None:
         # from virtual_field.runtime.spirob_elastica.constraints import (
         #     _SpirobBendConstraint,
         # )
-        from virtual_field.runtime.spirob_elastica.sdf_objects import (
-            SDFObstacleCylinders,
+        from virtual_field.runtime.spirob_elastica.sdf_objects_hash import (
+            SDFObstacleCylindersHash,
         )
 
         self._obstacles = load_noel_c4_obstacles()
@@ -193,7 +189,7 @@ class _Simulator(
         self.simulator.append(self.left_rod)
         self.simulator.append(self.right_rod)
 
-        p_linear = 200.0
+        p_linear = 500.0
         p_angular = 5.0
 
         self.simulator.add_forcing_to(self.left_rod).using(
@@ -226,24 +222,25 @@ class _Simulator(
             constrained_director_idx=(0,),
         )
 
-        self.simulator.detect_contact_between(
-            self.left_rod, self.right_rod
-        ).using(ea.RodRodContact, k=1e4, nu=3)
-        self.simulator.detect_contact_between(
-            self.right_rod, self.right_rod
-        ).using(ea.RodSelfContact, k=1e4, nu=3)
+        self.simulator.detect_contact_between(self.left_rod, self.right_rod).using(
+            ea.RodRodContact, k=1e4, nu=3
+        )
+        self.simulator.detect_contact_between(self.right_rod, self.right_rod).using(
+            ea.RodSelfContact, k=1e4, nu=3
+        )
+        self.simulator.detect_contact_between(self.left_rod, self.left_rod).using(
+            ea.RodSelfContact, k=1e4, nu=3
+        )
         # self.simulator.add_forcing_to(self.left_rod).using(
         #     _SpirobBendConstraint,
         #     kt=0,
         #     allowed_angle_in_deg=30,
         # )
 
         for rod in (self.left_rod, self.right_rod):
-            arm_id = (
-                self.arm_ids[0] if rod is self.left_rod else self.arm_ids[1]
-            )
+            arm_id = self.arm_ids[0] if rod is self.left_rod else self.arm_ids[1]
             self.simulator.add_forcing_to(rod).using(
-                SDFObstacleCylinders,
+                SDFObstacleCylindersHash,
                 starts=self._obstacles.starts,
                 directions=self._obstacles.directions,
                 lengths=self._obstacles.lengths,
@@ -293,7 +290,7 @@ def haptic_events(self) -> list[HapticEvent]:
             arm_id = event.arm_id
             penetration = self._tip_penetration_by_arm.get(arm_id, 0.0)
             intensity = penetration / self.tip_haptic_max_penetration
-            intensity = max(0.0, min(1.0, intensity))
+            intensity = max(0.0, min(0.5, intensity))  # cap at 50%
             event.active = intensity > 0.0
             event.intensity = intensity
         return self._haptic_events