Remove unnecessary gimbal test computation

Author: chschuma-disney

newton/tests/kamino/test_kamino_gimbal.py

@@ -77,6 +77,155 @@ def _build_rotational_d6(
     return builder.finalize(device="cpu"), d6
 
 
+@dataclass(frozen=True)
+class _Fixture:
+    """Model and D6 layout metadata for one conformance fixture."""
+
+    model: newton.Model
+    q_start: int
+    qd_start: int
+    target_q_start: int
+
+
+@dataclass(frozen=True)
+class _Probe:
+    """Raw coordinate, velocity, and effort data from a solver rollout."""
+
+    q: np.ndarray
+    qd: np.ndarray
+    effort: np.ndarray
+    coord_count: int
+    dof_count: int
+
+
+def _build_fixture(
+    fixed_base: bool,
+    axes: tuple[newton.Axis, newton.Axis, newton.Axis],
+    *,
+    stiffness: float = 0.0,
+    drive_damping: float = 0.0,
+    armature: float = 0.0,
+    passive_damping: float = 0.0,
+    lower: float | np.ndarray = -newton.MAXVAL,
+    upper: float | np.ndarray = newton.MAXVAL,
+) -> _Fixture:
+    """Build a collision-free articulated rotational D6."""
+    builder = newton.ModelBuilder(gravity=(0.0, 0.0, 0.0), up_axis=newton.Axis.Z)
+    base = builder.add_link(mass=2.0, inertia=_BASE_INERTIA, label="base")
+    link = builder.add_link(mass=1.0, inertia=_LINK_INERTIA, label="link")
+    root = (
+        builder.add_joint_fixed(parent=-1, child=base, label="root")
+        if fixed_base
+        else builder.add_joint_free(parent=-1, child=base, label="root")
+    )
+    lower_values = np.broadcast_to(lower, 3)
+    upper_values = np.broadcast_to(upper, 3)
+    configs = [
+        newton.ModelBuilder.JointDofConfig(
+            axis=axis,
+            target_pos=0.0,
+            target_vel=0.0,
+            target_ke=stiffness,
+            target_kd=drive_damping,
+            damping=passive_damping,
+            armature=armature,
+            limit_lower=float(lower_values[i]),
+            limit_upper=float(upper_values[i]),
+            limit_ke=1.0e4,
+            limit_kd=100.0,
+        )
+        for i, axis in enumerate(axes)
+    ]
+    d6 = builder.add_joint_d6(base, link, angular_axes=configs, label="d6")
+    builder.add_articulation([root, d6], label="d6")
+    model = builder.finalize(device=_DEVICE)
+    return _Fixture(
+        model,
+        int(model.joint_q_start.numpy()[d6]),
+        int(model.joint_qd_start.numpy()[d6]),
+        int(model.joint_target_q_start.numpy()[d6]),
+    )
+
+
+def _make_solver(backend: str, model: newton.Model) -> SolverKamino | SolverMuJoCo:
+    """Create a configured collision-free conformance solver."""
+    if backend == "kamino":
+        config = SolverKamino.Config(
+            integrator="euler",
+            use_collision_detector=False,
+            use_fk_solver=False,
+            sparse_jacobian=True,
+        )
+        config.constraints.alpha = 0.0
+        config.constraints.beta = 0.1
+        config.padmm.max_iterations = 200
+        config.padmm.primal_tolerance = 1.0e-6
+        config.padmm.dual_tolerance = 1.0e-6
+        config.padmm.compl_tolerance = 1.0e-6
+        return SolverKamino(model, config)
+    if backend == "mjwarp":
+        return SolverMuJoCo(
+            model, disable_contacts=True, integrator="implicitfast", iterations=100, use_mujoco_contacts=False
+        )
+    raise ValueError(f"Unsupported conformance backend: {backend}")
+
+
+def _run(
+    backend: str,
+    fixed_base: bool,
+    axes: tuple[newton.Axis, newton.Axis, newton.Axis],
+    *,
+    q: np.ndarray | None = None,
+    qd: np.ndarray | None = None,
+    effort: np.ndarray | None = None,
+    position_target: np.ndarray | None = None,
+    velocity_target: np.ndarray | None = None,
+    steps: int = 1,
+    **fixture_kwargs,
+) -> _Probe:
+    """Run a D6 rollout and retain the raw D6 coordinate trajectory."""
+    fixture = _build_fixture(fixed_base, axes, **fixture_kwargs)
+    state_in, state_out, control = fixture.model.state(), fixture.model.state(), fixture.model.control()
+    if q is not None:
+        values = state_in.joint_q.numpy()
+        values[fixture.q_start : fixture.q_start + 3] = q
+        state_in.joint_q.assign(values)
+    if qd is not None:
+        values = state_in.joint_qd.numpy()
+        values[fixture.qd_start : fixture.qd_start + 3] = qd
+        state_in.joint_qd.assign(values)
+    if effort is not None:
+        values = np.zeros(fixture.model.joint_dof_count, dtype=np.float32)
+        values[fixture.qd_start : fixture.qd_start + 3] = effort
+        control.joint_f.assign(values)
+    if position_target is not None:
+        values = control.joint_target_q.numpy()
+        values[fixture.target_q_start : fixture.target_q_start + 3] = position_target
+        control.joint_target_q.assign(values)
+    if velocity_target is not None:
+        values = control.joint_target_qd.numpy()
+        values[fixture.qd_start : fixture.qd_start + 3] = velocity_target
+        control.joint_target_qd.assign(values)
+    newton.eval_fk(fixture.model, state_in.joint_q, state_in.joint_qd, state_in)
+    solver = _make_solver(backend, fixture.model)
+    contacts = Contacts(rigid_contact_max=0, soft_contact_max=0, device=_DEVICE) if backend == "mjwarp" else None
+    positions = [state_in.joint_q.numpy()[fixture.q_start : fixture.q_start + 3].copy()]
+    velocities = [state_in.joint_qd.numpy()[fixture.qd_start : fixture.qd_start + 3].copy()]
+    for _ in range(steps):
+        state_in.clear_forces()
+        solver.step(state_in, state_out, control, contacts, _DT)
+        state_in, state_out = state_out, state_in
+        positions.append(state_in.joint_q.numpy()[fixture.q_start : fixture.q_start + 3].copy())
+        velocities.append(state_in.joint_qd.numpy()[fixture.qd_start : fixture.qd_start + 3].copy())
+    return _Probe(
+        np.stack(positions),
+        np.stack(velocities),
+        control.joint_f.numpy()[fixture.qd_start : fixture.qd_start + 3].copy(),
+        fixture.model.joint_coord_count,
+        fixture.model.joint_dof_count,
+    )
+
+
 class TestGimbal(unittest.TestCase):
     """Verify the rotational D6 representation."""
 
@@ -272,154 +421,27 @@ def test_fk_reset_preserves_left_handed_coordinates_and_rates(self):
             state.joint_qd.numpy()[qd_start : qd_start + 3], qd[qd_start : qd_start + 3], atol=1.0e-5
         )
 
-
-@dataclass(frozen=True)
-class _Fixture:
-    """Model and D6 layout metadata for one conformance fixture."""
-
-    model: newton.Model
-    q_start: int
-    qd_start: int
-    target_q_start: int
-
-
-@dataclass(frozen=True)
-class _Probe:
-    """Raw coordinate, velocity, and effort data from a solver rollout."""
-
-    q: np.ndarray
-    qd: np.ndarray
-    effort: np.ndarray
-    coord_count: int
-    dof_count: int
-
-
-def _build_fixture(
-    fixed_base: bool,
-    axes: tuple[newton.Axis, newton.Axis, newton.Axis],
-    *,
-    stiffness: float = 0.0,
-    drive_damping: float = 0.0,
-    armature: float = 0.0,
-    passive_damping: float = 0.0,
-    lower: float | np.ndarray = -newton.MAXVAL,
-    upper: float | np.ndarray = newton.MAXVAL,
-) -> _Fixture:
-    """Build a collision-free articulated rotational D6."""
-    builder = newton.ModelBuilder(gravity=(0.0, 0.0, 0.0), up_axis=newton.Axis.Z)
-    base = builder.add_link(mass=2.0, inertia=_BASE_INERTIA, label="base")
-    link = builder.add_link(mass=1.0, inertia=_LINK_INERTIA, label="link")
-    root = (
-        builder.add_joint_fixed(parent=-1, child=base, label="root")
-        if fixed_base
-        else builder.add_joint_free(parent=-1, child=base, label="root")
-    )
-    lower_values = np.broadcast_to(lower, 3)
-    upper_values = np.broadcast_to(upper, 3)
-    configs = [
-        newton.ModelBuilder.JointDofConfig(
-            axis=axis,
-            target_pos=0.0,
-            target_vel=0.0,
-            target_ke=stiffness,
-            target_kd=drive_damping,
-            damping=passive_damping,
-            armature=armature,
-            limit_lower=float(lower_values[i]),
-            limit_upper=float(upper_values[i]),
-            limit_ke=1.0e4,
-            limit_kd=100.0,
-        )
-        for i, axis in enumerate(axes)
-    ]
-    d6 = builder.add_joint_d6(base, link, angular_axes=configs, label="d6")
-    builder.add_articulation([root, d6], label="d6")
-    model = builder.finalize(device=_DEVICE)
-    return _Fixture(
-        model,
-        int(model.joint_q_start.numpy()[d6]),
-        int(model.joint_qd_start.numpy()[d6]),
-        int(model.joint_target_q_start.numpy()[d6]),
-    )
-
-
-def _make_solver(backend: str, model: newton.Model) -> SolverKamino | SolverMuJoCo:
-    """Create a configured collision-free conformance solver."""
-    if backend == "kamino":
-        config = SolverKamino.Config(
-            integrator="euler",
-            use_collision_detector=False,
-            use_fk_solver=False,
-            sparse_jacobian=True,
-        )
-        config.constraints.alpha = 0.0
-        config.constraints.beta = 0.1
-        config.padmm.max_iterations = 200
-        config.padmm.primal_tolerance = 1.0e-6
-        config.padmm.dual_tolerance = 1.0e-6
-        config.padmm.compl_tolerance = 1.0e-6
-        return SolverKamino(model, config)
-    if backend == "mjwarp":
-        return SolverMuJoCo(
-            model, disable_contacts=True, integrator="implicitfast", iterations=100, use_mujoco_contacts=False
-        )
-    raise ValueError(f"Unsupported conformance backend: {backend}")
-
-
-def _run(
-    backend: str,
-    fixed_base: bool,
-    axes: tuple[newton.Axis, newton.Axis, newton.Axis],
-    *,
-    q: np.ndarray | None = None,
-    qd: np.ndarray | None = None,
-    effort: np.ndarray | None = None,
-    position_target: np.ndarray | None = None,
-    velocity_target: np.ndarray | None = None,
-    steps: int = 1,
-    **fixture_kwargs,
-) -> _Probe:
-    """Run a D6 rollout and retain the raw D6 coordinate trajectory."""
-    fixture = _build_fixture(fixed_base, axes, **fixture_kwargs)
-    state_in, state_out, control = fixture.model.state(), fixture.model.state(), fixture.model.control()
-    if q is not None:
-        values = state_in.joint_q.numpy()
-        values[fixture.q_start : fixture.q_start + 3] = q
-        state_in.joint_q.assign(values)
-    if qd is not None:
-        values = state_in.joint_qd.numpy()
-        values[fixture.qd_start : fixture.qd_start + 3] = qd
-        state_in.joint_qd.assign(values)
-    if effort is not None:
-        values = np.zeros(fixture.model.joint_dof_count, dtype=np.float32)
-        values[fixture.qd_start : fixture.qd_start + 3] = effort
-        control.joint_f.assign(values)
-    if position_target is not None:
-        values = control.joint_target_q.numpy()
-        values[fixture.target_q_start : fixture.target_q_start + 3] = position_target
-        control.joint_target_q.assign(values)
-    if velocity_target is not None:
-        values = control.joint_target_qd.numpy()
-        values[fixture.qd_start : fixture.qd_start + 3] = velocity_target
-        control.joint_target_qd.assign(values)
-    newton.eval_fk(fixture.model, state_in.joint_q, state_in.joint_qd, state_in)
-    solver = _make_solver(backend, fixture.model)
-    contacts = Contacts(rigid_contact_max=0, soft_contact_max=0, device=_DEVICE) if backend == "mjwarp" else None
-    positions = [state_in.joint_q.numpy()[fixture.q_start : fixture.q_start + 3].copy()]
-    velocities = [state_in.joint_qd.numpy()[fixture.qd_start : fixture.qd_start + 3].copy()]
-    for _ in range(steps):
-        state_in.clear_forces()
-        solver.step(state_in, state_out, control, contacts, _DT)
-        state_in, state_out = state_out, state_in
-        positions.append(state_in.joint_q.numpy()[fixture.q_start : fixture.q_start + 3].copy())
-        velocities.append(state_in.joint_qd.numpy()[fixture.qd_start : fixture.qd_start + 3].copy())
-    return _Probe(
-        np.stack(positions),
-        np.stack(velocities),
-        control.joint_f.numpy()[fixture.qd_start : fixture.qd_start + 3].copy(),
-        fixture.model.joint_coord_count,
-        fixture.model.joint_dof_count,
-    )
+    def test_limits(self):
+        """Drive each D6 coordinate to its own position limit."""
+        lower = np.array([-0.15, -0.25, -0.35], dtype=np.float32)
+        upper = np.array([0.2, 0.3, 0.4], dtype=np.float32)
+        target = np.array([0.6, -0.6, 0.6], dtype=np.float32)
+        expected = np.where(target > 0.0, upper, lower)
+        for axes in (_RH_AXES, _LH_AXES):
+            for fixed_base in (True, False):
+                with self.subTest(axes=axes, fixed_base=fixed_base):
+                    probe = _run(
+                        "kamino",
+                        fixed_base,
+                        axes,
+                        position_target=target,
+                        stiffness=100.0,
+                        drive_damping=15.0,
+                        lower=lower,
+                        upper=upper,
+                        steps=50,
+                    )
+                    np.testing.assert_allclose(probe.q[-1], expected, atol=1.0e-2, rtol=0.0)
 
 
 @unittest.skipUnless(wp.get_cuda_device_count(), "requires CUDA device")
@@ -530,30 +552,6 @@ def test_unwrapped_pd_targets(self):
                         **kwargs,
                     )
 
-    def test_limits(self):
-        """Drive each D6 coordinate to its own position limit."""
-        lower = np.array([-0.15, -0.25, -0.35], dtype=np.float32)
-        upper = np.array([0.2, 0.3, 0.4], dtype=np.float32)
-        target = np.array([0.6, -0.6, 0.6], dtype=np.float32)
-        expected = np.where(target > 0.0, upper, lower)
-        for axes in (_RH_AXES, _LH_AXES):
-            for fixed_base in (True, False):
-                with self.subTest(axes=axes, fixed_base=fixed_base):
-                    for backend in ("mjwarp", "kamino"):
-                        with self.subTest(backend=backend):
-                            probe = _run(
-                                backend,
-                                fixed_base,
-                                axes,
-                                position_target=target,
-                                stiffness=100.0,
-                                drive_damping=15.0,
-                                lower=lower,
-                                upper=upper,
-                                steps=120,
-                            )
-                            np.testing.assert_allclose(probe.q[-1], expected, atol=1.0e-2, rtol=0.0)
-
     def test_passive_damping(self):
         """Match passive-damping D6 behavior."""
         for axes in (_RH_AXES, _LH_AXES):