typing fixes

Author: LeonLiu4

genesis/engine/entities/rigid_entity/rigid_entity.py

@@ -838,7 +838,7 @@ def get_jacobian(self, link, local_point=None):
         return jacobian
 
     @ti.func
-    def _impl_get_jacobian(self, tgt_link_idx, i_b, p_vec: ti.types.vector(3, ti.f64)):
+    def _impl_get_jacobian(self, tgt_link_idx, i_b, p_vec):
         self._func_get_jacobian(
             tgt_link_idx,
             i_b,
@@ -849,14 +849,14 @@ def _impl_get_jacobian(self, tgt_link_idx, i_b, p_vec: ti.types.vector(3, ti.f64
 
     @ti.kernel
     def _kernel_get_jacobian(self, tgt_link_idx: ti.i32, p_local: ti.types.ndarray()):
-        p_vec = ti.Vector([p_local[0], p_local[1], p_local[2]], dt=ti.f64)
+        p_vec = ti.Vector([p_local[0], p_local[1], p_local[2]], dt=gs.ti_float)
         for i_b in range(self._solver._B):
             self._impl_get_jacobian(tgt_link_idx, i_b, p_vec)
 
     @ti.kernel
     def _kernel_get_jacobian_zero(self, tgt_link_idx: ti.i32):
         for i_b in range(self._solver._B):
-            self._impl_get_jacobian(tgt_link_idx, i_b, ti.Vector.zero(ti.f64, 3))
+            self._impl_get_jacobian(tgt_link_idx, i_b, ti.Vector.zero(gs.ti_float, 3))
 
     @ti.func
     def _func_get_jacobian(self, tgt_link_idx, i_b, p_local, pos_mask, rot_mask):

tests/test_rigid_physics.py

@@ -2169,6 +2169,37 @@ def test_terrain_size(show_viewer, tol):
     assert_allclose((height_ref * 2.0), height_test, tol=tol)
 
 
+@pytest.mark.required
+@pytest.mark.parametrize("backend", [gs.cpu])
+def test_get_weld_constraints_basic(show_viewer, tol):
+    scene = gs.Scene(show_viewer=show_viewer)
+
+    cube1 = scene.add_entity(gs.morphs.Box(size=(0.05,) * 3, pos=(0.0, 0.0, 0.05)))
+    cube2 = scene.add_entity(gs.morphs.Box(size=(0.05,) * 3, pos=(0.2, 0.0, 0.05)))
+
+    scene.build(n_envs=1)
+
+    rigid = scene.sim.rigid_solver
+
+    link_a = torch.tensor([cube1.base_link.idx], dtype=gs.tc_int, device=gs.device)
+    link_b = torch.tensor([cube2.base_link.idx], dtype=gs.tc_int, device=gs.device)
+
+    rigid.add_weld_constraint(link_a, link_b)
+    scene.step()
+
+    welds = rigid.get_weld_constraints(as_tensor=True, to_torch=False)
+
+    env_id = 0 if "env" not in welds else int(welds["env"][0])
+
+    row = np.array(
+        [env_id, int(welds["obj_a"][0]), int(welds["obj_b"][0])],
+        dtype=np.int32,
+    )
+    ref = np.array([0, link_a.item(), link_b.item()], dtype=np.int32)
+
+    assert_allclose(row, ref, tol=tol)
+
+
 @pytest.mark.required
 @pytest.mark.parametrize("backend", [gs.cpu])
 def test_urdf_parsing(show_viewer, tol):
@@ -2302,37 +2333,6 @@ def test_urdf_mimic(show_viewer, tol):
     assert_allclose(gs_qpos[-1], gs_qpos[-2], tol=tol)
 
 
-@pytest.mark.required
-@pytest.mark.merge_fixed_links(False)
-@pytest.mark.parametrize("model_name", ["pendulum"])
-@pytest.mark.parametrize("gs_solver", [gs.constraint_solver.CG])
-@pytest.mark.parametrize("gs_integrator", [gs.integrator.Euler])
-def test_jacobian(gs_sim, tol):
-    pendulum = gs_sim.entities[0]
-    angle = 0.7
-    pendulum.set_qpos(np.array([angle], dtype=np.float64))
-    gs_sim.scene.step()
-
-    link = pendulum.get_link("PendulumArm_0")
-
-    p_local = np.array([0.05, -0.02, 0.12], dtype=np.float64)
-    J_o = pendulum.get_jacobian(link).cpu().numpy()
-    J_p = pendulum.get_jacobian(link, p_local).cpu().numpy()
-
-    c, s = np.cos(angle), np.sin(angle)
-    Rx = np.array([[1, 0, 0], [0, c, -s], [0, s, c]], dtype=np.float64)
-    r_world = Rx @ p_local
-    r_cross = np.array(
-        [[0, -r_world[2], r_world[1]], [r_world[2], 0, -r_world[0]], [-r_world[1], r_world[0], 0]], dtype=np.float64
-    )
-
-    lin_o, ang_o = J_o[:3, 0], J_o[3:, 0]
-    lin_expected = lin_o - r_cross @ ang_o
-
-    assert_allclose(J_p[3:, 0], ang_o, tol=tol)
-    assert_allclose(J_p[:3, 0], lin_expected, tol=tol)
-
-
 @pytest.mark.required
 @pytest.mark.parametrize("backend", [gs.cpu])
 def test_gravity(show_viewer, tol):