[FEATURE] get_weld_constraints API #1370
Conversation
LeonLiu4
force-pushed
the
feature/get-weld
branch
from
672a192 to
00f2534
Compare
LeonLiu4
force-pushed
the
feature/get-weld
branch
from
00f2534 to
c176525
Compare
| @@ -5715,6 +5715,68 @@ def _kernel_delete_weld_constraint( | |||
| ] | |||
| self.constraint_solver.ti_n_equalities[i_b] = self.constraint_solver.ti_n_equalities[i_b] - 1 | |||
|
|
|||
| @gs.assert_built | |||
| def get_weld_constraints(self, as_tensor: bool = True, to_torch: bool = True): | |||
There was a problem hiding this comment.
This is still wrong. If you look more closely, there is no such thing as valid_mask in RigidSolver.get_contacts, only in RigidEntity.get_contacts. You should do the same here.
|
|
||
| @ti.kernel | ||
| def _kernel_collect_welds(self, buf: ti.types.ndarray()): | ||
| WELD = gs.EQUALITY_TYPE.WELD |
There was a problem hiding this comment.
Remove all these variables and use their value directly in the implementation.
| "obj_b": obj_b, | ||
| "valid_mask": zeros((n_envs, 0), dtype=bool), | ||
| } | ||
| shape3 = (n_envs, stride, 3) |
There was a problem hiding this comment.
Could you follow more strictly the pattern in RigidSolver.get_contacts? You should not pass a 3D tensor because it would allocate more memory than necessary but most importantly, it will break "C"-contiguity, which would be problematic if forwarded as input argument of another taichi kernel.
| for _ in range(200): | ||
| scene.step() | ||
|
|
||
| vel1 = np.asarray(cube1.get_vel()) |
There was a problem hiding this comment.
No need to call np.asarray here
| vel1 = np.asarray(cube1.get_vel()) | ||
| vel2 = np.asarray(cube2.get_vel()) | ||
|
|
||
| assert_allclose(vel1, vel2, tol=1e-7) |
There was a problem hiding this comment.
Just do assert_allclose(cube1.get_vel(), cube2.get_vel(), tol=1e-7)
| welds_single = cube1.get_weld_constraints(as_tensor=True, to_torch=False) | ||
|
|
||
| assert_allclose( | ||
| np.array([int(welds_single["obj_a"][0]), int(welds_single["obj_b"][0])]), |
|
|
||
| assert_allclose( | ||
| np.array([int(welds_single["obj_a"][0]), int(welds_single["obj_b"][0])]), | ||
| np.array([link_a.item(), link_b.item()], dtype=np.int32), |
|
|
||
| scene.build(n_envs=1) | ||
|
|
||
| rigid = scene.sim.rigid_solver |
There was a problem hiding this comment.
Do not define this alias. This is useless here.
| @pytest.mark.parametrize("backend", [gs.cpu]) | ||
| def test_get_weld_constraints_entity(show_viewer, tol): | ||
| scene = gs.Scene(show_viewer=show_viewer) | ||
|
|
There was a problem hiding this comment.
No need to skip line here.
|
|
||
| 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))) | ||
|
|
Adds RigidSolver.get_weld_constraints() – returns an (N, 3) NumPy array (env_id, link1_idx, link2_idx) of all active weld constraints, enabling quick inspection and test assertions without poking internal Taichi fields.
Resolves #1365