Added callback functions for DampedSpring and DampedRotarySpring to allow customized force/torque calculations

Author: viblo

CHANGELOG.rst

@@ -5,6 +5,7 @@ Changelog
 .. 
   Fix for 2 static bodies that are changed to dynamic and are attached to constraints
   Updated the fork of Chipmunk2D used by Pymunk, fixing a number of issues, including maxForce on Spring constaints.
+  Added callback functions for DampedSpring and DampedRotarySpring to allow customized force/torque calculations.
   
 Pymunk 6.7.0 (2024-05-01)
 -------------------------

pymunk/_callbacks.py

@@ -292,9 +292,22 @@ def ext_cpConstraintPostSolveFunc(
     constraint._post_solve_func(constraint, constraint.a.space)
 
 
+@ffi.def_extern()
+def ext_cpDampedSpringForceFunc(cp_constraint: ffi.CData, dist: float) -> float:
+    constraint = ffi.from_handle(lib.cpConstraintGetUserData(cp_constraint))
+    return constraint._force_func(constraint, dist)
+
+
+@ffi.def_extern()
+def ext_cpDampedRotarySpringTorqueFunc(
+    cp_constraint: ffi.CData, relative_angle: float
+) -> float:
+    constraint = ffi.from_handle(lib.cpConstraintGetUserData(cp_constraint))
+    return constraint._torque_func(constraint, relative_angle)
+
 
 # Pickle of Arbiters
 @ffi.def_extern()
 def ext_cpArbiterIteratorFunc(_arbiter, data):  # type: ignore
     arbiters = ffi.from_handle(data)
-    arbiters.append(_arbiter)
+    arbiters.append(_arbiter)

pymunk/constraints.py

@@ -82,6 +82,9 @@
 
 _logger = logging.getLogger(__name__)
 
+_TorqueFunc = Callable[["DampedRotarySpring", float], float]
+_ForceFunc = Callable[["DampedSpring", float], float]
+
 
 class Constraint(PickleMixin, TypingAttrMixing, object):
     """Base class of all constraints.
@@ -582,6 +585,8 @@ class DampedSpring(Constraint):
         "damping",
     ]
 
+    _pickle_attrs_skip = Constraint._pickle_attrs_skip + ["_force_func"]
+
     def __init__(
         self,
         a: "Body",
@@ -672,6 +677,34 @@ def _set_damping(self, damping: float) -> None:
         doc="""How soft to make the damping of the spring.""",
     )
 
+    @staticmethod
+    def spring_force(spring: "DampedSpring", dist: float) -> float:
+        """Default damped spring force function."""
+        return lib.defaultSpringForce(spring._constraint, dist)
+
+    def _set_force_func(self, func: _ForceFunc) -> None:
+        if func == DampedSpring.spring_force:
+            lib.cpDampedSpringSetSpringForceFunc(
+                self._constraint,
+                ffi.cast(
+                    "cpDampedSpringForceFunc", ffi.addressof(lib, "defaultSpringForce")
+                ),
+            )
+        else:
+            self._force_func = func
+            lib.cpDampedSpringSetSpringForceFunc(
+                self._constraint, lib.ext_cpDampedSpringForceFunc
+            )
+
+    force_func = property(
+        fset=_set_force_func,
+        doc="""The force callback function. 
+        
+        The force callback function is called each time step and is used to
+        calculate the force of the spring (exclusing any damping).
+        """,
+    )
+
 
 class DampedRotarySpring(Constraint):
     """DampedRotarySpring works like the DammpedSpring but in a angular fashion."""
@@ -682,6 +715,10 @@ class DampedRotarySpring(Constraint):
         "damping",
     ]
 
+    _pickle_attrs_skip = Constraint._pickle_attrs_skip + ["_torque_func"]
+
+    _torque_func: Optional[_TorqueFunc] = None
+
     def __init__(
         self, a: "Body", b: "Body", rest_angle: float, stiffness: float, damping: float
     ) -> None:
@@ -733,6 +770,35 @@ def _set_damping(self, damping: float) -> None:
         doc="""How soft to make the damping of the spring.""",
     )
 
+    @staticmethod
+    def spring_torque(spring: "DampedRotarySpring", relative_angle: float) -> float:
+        """Default damped rotary spring torque function."""
+        return lib.defaultSpringTorque(spring._constraint, relative_angle)
+
+    def _set_torque_func(self, func: _TorqueFunc) -> None:
+        if func == DampedRotarySpring.spring_torque:
+            lib.cpDampedRotarySpringSetSpringTorqueFunc(
+                self._constraint,
+                ffi.cast(
+                    "cpDampedRotarySpringTorqueFunc",
+                    ffi.addressof(lib, "defaultSpringTorque"),
+                ),
+            )
+        else:
+            self._torque_func = func
+            lib.cpDampedRotarySpringSetSpringTorqueFunc(
+                self._constraint, lib.ext_cpDampedRotarySpringTorqueFunc
+            )
+
+    torque_func = property(
+        fset=_set_torque_func,
+        doc="""The torque callback function. 
+        
+        The torque callback function is called each time step and is used to
+        calculate the torque of the spring (exclusing any damping).
+        """,
+    )
+
 
 class RotaryLimitJoint(Constraint):
     """RotaryLimitJoint constrains the relative rotations of two bodies."""

pymunk/tests/test_constraint.py

@@ -338,6 +338,27 @@ def testDamping(self) -> None:
         j.damping = 2
         self.assertEqual(j.damping, 2)
 
+    def testForceFunc(self) -> None:
+        s = p.Space()
+        a, b = p.Body(10, 10), p.Body(20, 20)
+        b.position = 0, 100
+        j = DampedSpring(a, b, (0, 0), (0, 0), 0, 1, 0)
+        s.add(a, b, j)
+
+        def f(spring: p.DampedSpring, dist: float) -> float:
+            self.assertEqual(spring, j)
+            self.assertEqual(dist, 100)
+
+            return 1
+
+        j.force_func = f
+        s.step(1)
+        self.assertEqual(j.impulse, 1)
+
+        j.force_func = DampedSpring.spring_force
+        s.step(1)
+        self.assertAlmostEqual(j.impulse, -100.15)
+
     def testPickle(self) -> None:
         a, b = p.Body(10, 10), p.Body(20, 20)
         j = DampedSpring(a, b, (1, 2), (3, 4), 5, 6, 7)
@@ -376,6 +397,27 @@ def testDamping(self) -> None:
         j.damping = 2
         self.assertEqual(j.damping, 2)
 
+    def testTorqueFunc(self) -> None:
+        s = p.Space()
+        a, b = p.Body(10, 10), p.Body(20, 20)
+        b.angle = 2
+        j = DampedRotarySpring(a, b, 0, 10, 0)
+        s.add(a, b, j)
+
+        def f(spring: p.DampedRotarySpring, relative_angle: float) -> float:
+            self.assertEqual(spring, j)
+            self.assertEqual(relative_angle, -2)
+
+            return 1
+
+        j.torque_func = f
+        s.step(1)
+        self.assertEqual(j.impulse, 1)
+
+        j.torque_func = DampedRotarySpring.spring_torque
+        s.step(1)
+        self.assertAlmostEqual(j.impulse, -21.5)
+
     def testPickle(self) -> None:
         a, b = p.Body(10, 10), p.Body(20, 20)
         j = DampedRotarySpring(a, b, 1, 2, 3)

pymunk_cffi/callbacks_cdef.h

@@ -3,6 +3,8 @@ extern "Python" {
     // cpConstraint.h
     void ext_cpConstraintPreSolveFunc(cpConstraint *constraint, cpSpace *space);
     void ext_cpConstraintPostSolveFunc(cpConstraint *constraint, cpSpace *space);
+    cpFloat ext_cpDampedSpringForceFunc(cpConstraint *constraint, cpFloat dist);
+    cpFloat ext_cpDampedRotarySpringTorqueFunc(cpConstraint *constraint, cpFloat relative_angle);
 
 
     // cpBody.h

pymunk_cffi/extensions.c

@@ -447,4 +447,12 @@ cpArbiter *cpArbiterNew(cpShape *a, cpShape *b)
 cpContact *cpContactArrAlloc(int count)
 {
     return (cpContact *)cpcalloc(count, sizeof(struct cpContact));
+}
+
+cpFloat defaultSpringForce(cpDampedSpring *spring, cpFloat dist){
+	return (spring->restLength - dist)*spring->stiffness;
+}
+
+cpFloat defaultSpringTorque(cpDampedRotarySpring *spring, cpFloat relativeAngle){
+	return (relativeAngle - spring->restAngle)*spring->stiffness;
 }

pymunk_cffi/extensions_cdef.h

@@ -99,3 +99,7 @@ cpArbiter *cpArbiterNew(cpShape *a, cpShape *b);
 typedef struct cpContact cpContact;
 
 cpContact *cpContactArrAlloc(int count);
+
+cpFloat defaultSpringForce(cpDampedSpring *spring, cpFloat dist);
+
+cpFloat defaultSpringTorque(cpDampedRotarySpring *spring, cpFloat relativeAngle);