fix(dynamics): support arbitrary joint counts in paycap

src/roboticstoolbox/robot/Dynamics.py

@@ -1438,93 +1438,58 @@ def paycap(
         frame: int = 1,
         q: ArrayLike | None = None,
     ):
-        """
-        Static payload capacity of a robot
-
-        :param w: The payload wrench
-        :type w: ndarray(6,)
-        :param tauR: Joint torque matrix minimum and maximums
-        :type tauR: ndarray(n,2)
-        :param frame: The frame in which to torques are expressed in when J
-            is not supplied. 'base' means base frame of the robot, 'ee' means
-            end-effector frame
-        :param q: Joint coordinates
-        :type q: ndarray(n,)
-        :returns: The maximum permissible payload wrench
-        :rtype: ndarray(6,)
-
-        ``wmax, joint = paycap(q, w, f, tauR)`` returns the maximum permissible
-        payload wrench ``wmax`` (6) applied at the end-effector, and the index
-        of the joint (zero indexed) which hits its force/torque limit at that
-        wrench. ``q`` (n) is the manipulator pose, ``w`` the payload wrench
-        (6), ``f`` the wrench reference frame and tauR (nx2) is a matrix of
-        joint forces/torques (first col is maximum, second col minimum).
-
-        **Trajectory operation:**
-
-        In the case q is nxm then wmax is Mx6 and J is Mx1 where the rows are
-        the results at the pose given by corresponding row of q.
+        """Compute static payload capacity along a wrench direction.
+
+        :param w: Nonzero payload wrench direction (6,), or (m, 6) for
+            m configurations. Its magnitude is ignored.
+        :param tauR: Joint torque limits (n, 2): maximum in column 0,
+            minimum in column 1.
+        :param frame: Wrench reference frame: 0 for base, 1 for end-effector.
+        :param q: Joint coordinates (n,), or (m, n) for m configurations.
+            Defaults to the robot's current configuration.
+        :returns: Tuple ``(wmax, joint)``. ``wmax`` contains one allowable
+            load scale per joint (n,), and ``joint`` is the zero-based index
+            of the limiting joint. With m configurations, their shapes are
+            (m, n) and (m,), respectively. The maximum load scale in the
+            specified direction is ``wmax.min()`` per configuration.
 
         .. rubric:: Notes
 
-        - Wrench vector and Jacobian must be from the same reference frame
-        - Tool transforms are taken into consideration for frame=1.
-
+        - Wrench and Jacobian must use the same reference frame.
+        - Tool transforms are taken into consideration for ``frame=1``.
+        - Gravity compensation requires a robot model with valid dynamics.
         """
-
-        # TODO rewrite
-        trajn = 1
-
-        if q is None:
-            q = self.q
-        else:
-            q = np.array(q)
-
-        try:
-            q = np.array(getvector(q, self.n, "row"))
-            w = np.array(getvector(w, 6, "row"))
-        except ValueError:
-            trajn = q.shape[1]
-            verifymatrix(q, (trajn, self.n))
-            verifymatrix(w, (trajn, 6))
-
+        q = getmatrix(self.q if q is None else q, (None, self.n))
+        w = getmatrix(w, (None, 6))
+        if q.shape[0] != w.shape[0]:
+            raise ValueError("q and w must have the same number of rows")
         verifymatrix(tauR, (self.n, 2))
 
-        wmax = np.zeros((trajn, 6))
+        w_norm = np.linalg.norm(w, axis=1)
+        if np.any(w_norm == 0):
+            raise ValueError("w must contain nonzero wrench directions")
+
+        trajn = q.shape[0]
+        wmax = np.zeros((trajn, self.n))
         joint = np.zeros(trajn, dtype=int)
 
         for i in range(trajn):
             tauB = self.gravload(q[i, :])
+            tauP = self.pay(w[i, :] / w_norm[i], q=q[i, :], frame=frame)
 
-            # tauP = self.rne(
-            #     np.zeros(self.n), np.zeros(self.n),
-            #     q, grav=[0, 0, 0], fext=w/np.linalg.norm(w))
-
-            tauP = self.pay(w[i, :] / np.linalg.norm(w[i, :]), q=q[i, :], frame=frame)
-
-            M = tauP > 0
-            m = tauP <= 0
-
-            TAUm = np.ones(self.n)
-            TAUM = np.ones(self.n)
-
-            for c in range(self.n):
-                TAUM[c] = tauR[c, 0]
-                TAUm[c] = tauR[c, 1]
-
-            WM = np.zeros(self.n)
-            WM[M] = (TAUM[M] - tauB[M]) / tauP[M]
-            WM[m] = (TAUm[m] - tauB[m]) / tauP[m]
-
+            positive = tauP > 0
+            negative = tauP < 0
+            WM = np.full(self.n, np.inf)
+            WM[positive] = (tauR[positive, 0] - tauB[positive]) / tauP[positive]
+            WM[negative] = (tauR[negative, 1] - tauB[negative]) / tauP[negative]
             WM[WM == -np.inf] = np.inf
 
             wmax[i, :] = WM
             joint[i] = np.argmin(WM)
 
         if trajn == 1:
             return wmax[0, :], joint[0]
-        else:
-            return wmax, joint
+        return wmax, joint
 
     def perturb(self: RobotProto, p=0.1):
         """

tests/test_paycap_capacity.py

@@ -0,0 +1,87 @@
+"""Regression tests for payload capacity with arbitrary joint counts."""
+
+import numpy as np
+import numpy.testing as npt
+import pytest
+
+from roboticstoolbox.robot.Dynamics import DynamicsMixin
+
+
+class _DummyRobot:
+    paycap = DynamicsMixin.paycap
+
+    def __init__(self, n):
+        self.n = n
+        self.q = np.zeros(n)
+        self.calls = []
+
+    def gravload(self, q):
+        return np.zeros(self.n)
+
+    def pay(self, w, q, frame):
+        self.calls.append((np.copy(w), np.copy(q), frame))
+        return np.arange(1, self.n + 1) * w[0]
+
+
+def _limits(n):
+    return np.column_stack((np.full(n, 20.0), np.full(n, -10.0)))
+
+
+@pytest.mark.parametrize("n", [3, 6, 7])
+def test_paycap_joint_count(n):
+    robot = _DummyRobot(n)
+    capacity, limiting_joint = robot.paycap(
+        [2.0, 0, 0, 0, 0, 0], _limits(n), frame=0
+    )
+    npt.assert_allclose(capacity, 20.0 / np.arange(1, n + 1))
+    npt.assert_equal(capacity.shape, (n,))
+    npt.assert_equal(limiting_joint, n - 1)
+    npt.assert_allclose(robot.calls[0][0], [1, 0, 0, 0, 0, 0])
+    npt.assert_equal(robot.calls[0][2], 0)
+
+
+def test_paycap_negative_wrench_uses_minimum_torque_limit():
+    robot = _DummyRobot(7)
+    capacity, limiting_joint = robot.paycap(
+        [-10.0, 0, 0, 0, 0, 0], _limits(7)
+    )
+    npt.assert_allclose(capacity, 10.0 / np.arange(1, 8))
+    npt.assert_equal(limiting_joint, 6)
+
+
+def test_paycap_trajectory_uses_each_wrench_direction():
+    robot = _DummyRobot(7)
+    q = np.stack([np.zeros(7), np.ones(7)])
+    wrench = np.array([[2.0, 0, 0, 0, 0, 0], [-10.0, 0, 0, 0, 0, 0]])
+    capacities, joints = robot.paycap(wrench, _limits(7), q=q)
+    npt.assert_equal(capacities.shape, (2, 7))
+    npt.assert_allclose(capacities[0], 20.0 / np.arange(1, 8))
+    npt.assert_allclose(capacities[1], 10.0 / np.arange(1, 8))
+    npt.assert_array_equal(joints, [6, 6])
+    npt.assert_allclose([call[1] for call in robot.calls], q)
+
+
+def test_paycap_trajectory_row_mismatch():
+    robot = _DummyRobot(7)
+    with pytest.raises(ValueError, match="same number of rows"):
+        robot.paycap(np.ones((1, 6)), _limits(7), q=np.zeros((2, 7)))
+
+
+def test_paycap_rejects_zero_wrench_direction():
+    robot = _DummyRobot(7)
+    with pytest.raises(ValueError, match="nonzero wrench"):
+        robot.paycap(np.zeros(6), _limits(7))
+
+
+def test_paycap_joint_unaffected_by_wrench_is_unbounded():
+    robot = _DummyRobot(3)
+    robot.pay = lambda w, q, frame: np.array([1.0, 0.0, 2.0])
+    limits = _limits(3)
+    limits[1, 1] = 0.0
+
+    with np.errstate(divide="raise", invalid="raise"):
+        capacities, joint = robot.paycap([1.0, 0, 0, 0, 0, 0], limits)
+
+    npt.assert_allclose(capacities[[0, 2]], [20.0, 10.0])
+    npt.assert_equal(np.isinf(capacities[1]), True)
+    npt.assert_equal(joint, 2)