newton/_src/viewer/picking.py (1)

103-111: Tighten type annotations for nullable input and Warp array typing

At Line 110 the function accepts None, so the parameter type should reflect that. Also prefer typed Warp array annotations for consistency.

Proposed annotation update

-    def _compute_effective_mass(model: newton.Model) -> wp.array:
+    def _compute_effective_mass(model: newton.Model | None) -> wp.array[float]:

As per coding guidelines: "Use PEP 604 unions (x | None, not Optional[x]) in type annotations." and "Annotate Warp arrays with bracket syntax (wp.array[wp.vec3], wp.array2d[float], wp.array[Any]), not the parenthesized form."

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@newton/_src/viewer/picking.py` around lines 103 - 111, Update the
_compute_effective_mass annotation to reflect that model can be None and use
Warp's bracketed array typing: change the parameter type to "newton.Model |
None" and the return type to "wp.array[float]". Locate the function
_compute_effective_mass in picking.py and replace the current signature
accordingly so callers and type checkers know the input is nullable and the
function returns a typed Warp array.

Fix all unresolved CodeRabbit comments on this PR:

• Push a commit to this branch (recommended)
• Create a new PR with the fixes

newton/_src/viewer/picking.py (1)

107-145: Implementation is correct — optional numpy optimization available.

The articulation mass aggregation logic correctly handles all cases:

• Free bodies retain their own mass
• Articulated bodies get the total mass of their articulation
• The joint_child-only traversal is sufficient since Newton's model ensures every articulation body (including root) appears as a child of exactly one joint.

For large models, the Python loops could be replaced with NumPy vectorized operations for marginal performance gains, though this is initialization-time code that runs once.

♻️ Optional: NumPy-based aggregation

         if model.joint_count > 0:
             joint_child_np = model.joint_child.numpy()
             joint_art_np = model.joint_articulation.numpy()

             # Map each body to its articulation index (-1 if free)
             body_art = np.full(model.body_count, -1, dtype=np.int32)
-            for j in range(model.joint_count):
-                child = joint_child_np[j]
-                if child >= 0:
-                    body_art[child] = joint_art_np[j]
+            valid_children = joint_child_np >= 0
+            body_art[joint_child_np[valid_children]] = joint_art_np[valid_children]

-            # Sum mass per articulation
-            art_mass = {}
-            for b in range(model.body_count):
-                a = body_art[b]
-                if a >= 0:
-                    art_mass[a] = art_mass.get(a, 0.0) + body_mass_np[b]
-
-            # Assign total articulation mass to each body in that articulation
-            for b in range(model.body_count):
-                a = body_art[b]
-                if a >= 0:
-                    effective[b] = art_mass[a]
+            # Sum mass per articulation and assign back
+            in_art = body_art >= 0
+            if np.any(in_art):
+                art_ids = body_art[in_art]
+                art_mass_sum = np.bincount(art_ids, weights=body_mass_np[in_art])
+                effective[in_art] = art_mass_sum[art_ids]

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@newton/_src/viewer/picking.py` around lines 107 - 145, The current
_compute_effective_mass implementation is correct but can be optimized by
replacing the Python loops that build body_art, compute art_mass, and assign
per-body effective masses with NumPy vectorized ops: create body_art from
joint_child_np/joint_art_np by initializing body_art = np.full(...) and using
body_art[joint_child_np >= 0] = joint_art_np[...] (or equivalent indexed
assignment), compute articulation masses with np.bincount on joint_art_np
filtered by child indices and weights from body_mass_np (ensuring the bincount
length covers max articulation index), then assign effective = np.where(body_art
>= 0, art_mass_lookup[body_art], body_mass_np); keep returning
wp.array(effective, dtype=float, device=model.device) and preserve behavior for
model is None and free bodies.