Fix TetMesh._infer_frequency for length-1 attribute arrays (#3240)

Co-authored-by: Miguel Angel Zamora Mora <mzamoramora@nvidia.com>

Author: eric-heiden

CHANGELOG.md

@@ -145,6 +145,7 @@
 - Fix `ModelBuilder.add_usd()` silently ignoring authored Newton USD collision schemas when `newton-usd-schemas` is unavailable. USD import now fails clearly instead of continuing without the registered schema defaults required for SDF and hydroelastic configuration.
 - Fix `cable_cross_slide_table` example stability so the cable-driven table reliably tracks its rectangular path and catches drift during regression runs.
 - Fix URDF `package://` mesh fallback resolution without `resolve-robotics-uri-py` so package names only match full path components instead of unrelated directory-name substrings
+- Fix TetMesh USD loading to infer unambiguous length-1 custom attributes as `AttributeFrequency.ONCE`, flatten indexed primvars, and omit uninferable attributes without dropping the soft body. `ModelBuilder.add_usd()` imports only registered TetMesh attributes using their declared frequencies. (#3228)
 - Fix `ModelBuilder.collapse_fixed_joints()` crashing with `IndexError` when a `mujoco:equality_constraint` row omits optional fields (`anchor`, `relpose`) that carry defaults. (#3054)
 - Fix `ViewerGL.set_model()` resetting headless/interactive camera and wind state when switching between models that use the same up-axis. (#2658)
 - Fix bend force calculation error in Style3D solver

newton/_src/geometry/types.py

@@ -1718,11 +1718,15 @@ def _infer_frequency(
         first_dim = arr.shape[0] if arr.ndim >= 1 else 1
         counts = {"vertex_count": vertex_count, "tet_count": tet_count, "tri_count": tri_count}
         matches = [label for label, c in counts.items() if first_dim == c and c > 0]
+        if first_dim == 1:
+            matches.append("ONCE")
         if len(matches) > 1:
             raise ValueError(
                 f"Cannot infer frequency for custom attribute '{name}': array length {first_dim} matches "
                 f"{', '.join(matches)}. Pass an explicit (array, frequency) tuple instead."
             )
+        if "ONCE" in matches:
+            return Model.AttributeFrequency.ONCE
         if first_dim == vertex_count and vertex_count > 0:
             return Model.AttributeFrequency.PARTICLE
         if first_dim == tet_count and tet_count > 0:
@@ -1851,6 +1855,11 @@ def create_from_usd(prim, *, compat_namespaces: Sequence[str] | None = None) ->
         ``k_lambda``) and density on the returned TetMesh. Material properties
         are set to ``None`` if not present.
 
+        Custom primvars use their resolved interpolation to determine attribute
+        frequency. Other custom arrays use length-based inference; arrays whose
+        frequency is ambiguous or cannot be inferred emit a warning and are
+        omitted without preventing the TetMesh from loading.
+
         Material-attribute namespaces (deprecated default): with ``compat_namespaces=None``
         (the default) the legacy vendor namespaces (``omniphysics:`` / ``physxDeformableBody:``)
         are read off any bound material, matching the pre-canonical behavior. That default is

newton/_src/usd/utils.py

@@ -647,6 +647,31 @@ def get_custom_attribute_values(
     return out
 
 
+def _get_tetmesh_custom_attribute_values(
+    prim: Usd.Prim,
+    custom_attributes: Sequence[ModelBuilder.CustomAttribute],
+) -> dict[str, np.ndarray]:
+    """Read builder-declared TetMesh arrays without inferring their frequency."""
+    out: dict[str, np.ndarray] = {}
+    for spec in custom_attributes:
+        usd_name = spec.usd_attribute_name
+        if not usd_name or usd_name == "*":
+            continue
+        usd_attr = prim.GetAttribute(usd_name)
+        if not usd_attr or not usd_attr.HasAuthoredValue():
+            continue
+        primvar = UsdGeom.Primvar(usd_attr)
+        value = primvar.ComputeFlattened() if primvar else usd_attr.Get()
+        if value is None:
+            continue
+        if spec.usd_value_transformer is not None:
+            value = spec.usd_value_transformer(value, {"prim": prim, "attr": spec})
+            if value is None:
+                continue
+        out[spec.key] = np.asarray(value)
+    return out
+
+
 def _newell_normal(P: np.ndarray) -> np.ndarray:
     """Newell's method for polygon normal (not normalized)."""
     x = y = z = 0.0
@@ -1559,7 +1584,12 @@ def _material_authors_unscoped_canonical_attrs(prim: Usd.Prim) -> bool:
     )
 
 
-def get_tetmesh(prim: Usd.Prim, *, compat_namespaces: Sequence[str] | None = None) -> TetMesh:
+def get_tetmesh(
+    prim: Usd.Prim,
+    *,
+    compat_namespaces: Sequence[str] | None = None,
+    _load_custom_attributes: bool = True,
+) -> TetMesh:
     """Load a tetrahedral mesh from a USD prim with the ``UsdGeom.TetMesh`` schema.
 
     Reads vertex positions from the ``points`` attribute and tetrahedral
@@ -1571,6 +1601,11 @@ def get_tetmesh(prim: Usd.Prim, *, compat_namespaces: Sequence[str] | None = Non
     ``k_lambda``) and density on the returned TetMesh. Material properties
     are set to ``None`` if not present.
 
+    Custom primvars use their resolved interpolation to determine attribute
+    frequency. Other custom arrays use length-based inference; arrays whose
+    frequency is ambiguous or cannot be inferred emit a warning and are
+    omitted without preventing the TetMesh from loading.
+
     Material-attribute namespaces (deprecated default): with ``compat_namespaces=None``
     (the default) the legacy vendor namespaces (``omniphysics:`` / ``physxDeformableBody:``)
     are read off any bound material, matching the pre-canonical behavior. That default is
@@ -1716,6 +1751,15 @@ def get_tetmesh(prim: Usd.Prim, *, compat_namespaces: Sequence[str] | None = Non
         # density too; a plain rigid-style physics material is a valid source.
         density = _get_physics_material_density(material_prim)
 
+    if not _load_custom_attributes:
+        return TetMesh(
+            vertices=vertices,
+            tet_indices=tet_indices,
+            k_mu=k_mu,
+            k_lambda=k_lambda,
+            density=density,
+        )
+
     # Read custom primvars and attributes (per-vertex, per-tet, etc.)
     # Primvar interpolation is used to determine the attribute frequency
     # when available, falling back to length-based inference in TetMesh.__init__.
@@ -1738,7 +1782,7 @@ def get_tetmesh(prim: Usd.Prim, *, compat_namespaces: Sequence[str] | None = Non
         name = primvar.GetPrimvarName()
         if name in ("st", "normals"):
             continue  # skip well-known primvars handled elsewhere
-        val = primvar.Get()
+        val = primvar.ComputeFlattened()
         if val is not None:
             arr = np.array(val)
             interp = primvar.GetInterpolation()
@@ -1773,14 +1817,32 @@ def get_tetmesh(prim: Usd.Prim, *, compat_namespaces: Sequence[str] | None = Non
             except (TypeError, ValueError):
                 pass  # skip non-array attributes
 
-    return TetMesh(
+    result = TetMesh(
         vertices=vertices,
         tet_indices=tet_indices,
         k_mu=k_mu,
         k_lambda=k_lambda,
         density=density,
-        custom_attributes=custom_attributes if custom_attributes else None,
     )
+    tri_count = len(result.surface_tri_indices) // 3
+    for name, value in custom_attributes.items():
+        if name in result._RESERVED_ATTR_KEYS:
+            warnings.warn(
+                f"{prim.GetPath()}: custom attribute '{name}' uses a reserved TetMesh name; skipping the attribute.",
+                stacklevel=2,
+            )
+            continue
+        if isinstance(value, tuple):
+            arr, frequency = value
+        else:
+            arr = value
+            try:
+                frequency = result._infer_frequency(arr, result.vertex_count, result.tet_count, tri_count, name)
+            except ValueError as exc:
+                warnings.warn(f"{prim.GetPath()}: {exc}; skipping the attribute.", stacklevel=2)
+                continue
+        result.custom_attributes[name] = (np.asarray(arr), frequency)
+    return result
 
 
 def _find_physics_material_prim(prim: Usd.Prim):

newton/_src/utils/import_usd.py

@@ -1013,7 +1013,11 @@ def _get_tetmesh_cached(prim: Usd.Prim) -> TetMesh:
                     stacklevel=2,
                 )
                 compat_ns = usd.DEFORMABLE_LEGACY_NAMESPACES
-            tetmesh_cache[prim_path] = usd.get_tetmesh(prim, compat_namespaces=compat_ns)
+            tetmesh_cache[prim_path] = usd.get_tetmesh(
+                prim,
+                compat_namespaces=compat_ns,
+                _load_custom_attributes=False,
+            )
         return tetmesh_cache[prim_path]
 
     def _has_visual_material_properties(material_props: dict[str, Any]) -> bool:

newton/_src/utils/import_usd_deformable_volume.py

@@ -16,6 +16,7 @@
 import numpy as np
 import warp as wp
 
+from ..sim.model import Model
 from .import_usd_deformable_utils import (
     _apply_particle_masses,
     _bake_world_points,
@@ -104,16 +105,57 @@ def _deformable_import_volume(ctx: _DeformableImportContext) -> None:
             # whole import; skip the prim like other broken deformable geometry.
             warnings.warn(f"{path}: invalid TetMesh; skipping soft-body import ({exc}).", stacklevel=2)
             continue
+        supported_frequencies = {
+            Model.AttributeFrequency.PARTICLE,
+            Model.AttributeFrequency.TETRAHEDRON,
+            Model.AttributeFrequency.TRIANGLE,
+        }
+        registered_attributes = [
+            attr
+            for attr in builder.custom_attributes.values()
+            if attr.frequency in supported_frequencies and attr.usd_attribute_name != "*"
+        ]
+        registered_values = usd._get_tetmesh_custom_attribute_values(prim, registered_attributes)
+        once_attributes = [
+            attr
+            for attr in builder.custom_attributes.values()
+            if attr.frequency == Model.AttributeFrequency.ONCE and attr.usd_attribute_name != "*"
+        ]
+        once_values = usd._get_tetmesh_custom_attribute_values(prim, once_attributes)
+        for attr in once_attributes:
+            if attr.key in once_values:
+                warnings.warn(
+                    f"{path}: registered TetMesh attribute '{attr.usd_attribute_name}' has ONCE frequency, "
+                    "which cannot be attached per soft body and is not imported.",
+                    stacklevel=2,
+                )
+        frequency_counts = {
+            Model.AttributeFrequency.PARTICLE: (tetmesh.vertex_count, "particle"),
+            Model.AttributeFrequency.TETRAHEDRON: (tetmesh.tet_count, "tetrahedron"),
+            Model.AttributeFrequency.TRIANGLE: (len(tetmesh.surface_tri_indices) // 3, "triangle"),
+        }
+        imported_attributes = {}
+        for attr in registered_attributes:
+            if attr.key not in registered_values:
+                continue
+            arr = np.asarray(registered_values[attr.key])
+            actual_count = arr.shape[0] if arr.ndim >= 1 else 1
+            expected_count, frequency_name = frequency_counts[attr.frequency]
+            if arr.ndim == 0 or actual_count != expected_count:
+                warnings.warn(
+                    f"{path}: registered TetMesh attribute '{attr.usd_attribute_name}' has array length "
+                    f"{actual_count}, which does not match {frequency_name} count {expected_count}; "
+                    "skipping the attribute.",
+                    stacklevel=2,
+                )
+                continue
+            imported_attributes[attr.key] = (arr, attr.frequency)
         tetmesh_for_builder = tetmesh
-        if tetmesh.custom_attributes:
-            filtered_custom_attributes = {
-                k: v for k, v in tetmesh.custom_attributes.items() if k in builder.custom_attributes
-            }
-            if len(filtered_custom_attributes) != len(tetmesh.custom_attributes):
-                # Preserve the cached TetMesh while keeping add_usd's
-                # current behavior of dropping unregistered import attrs.
-                tetmesh_for_builder = copy.copy(tetmesh)
-                tetmesh_for_builder.custom_attributes = filtered_custom_attributes
+        if tetmesh.custom_attributes or imported_attributes:
+            # The builder declaration supplies both the USD name and authoritative frequency.
+            # Preserve the cached TetMesh because another import pass may reuse it.
+            tetmesh_for_builder = copy.copy(tetmesh)
+            tetmesh_for_builder.custom_attributes = imported_attributes
 
         soft_mesh_mat = get_prim_world_mat(prim, None, incoming_world_xform)
         # Bake the full world affine into the tet vertices and pass an identity placement, so a

newton/tests/test_import_usd.py

@@ -11605,13 +11605,14 @@ def test_tetmesh_custom_attributes_constructor(self):
         temperature = np.array([100.0, 200.0, 300.0, 400.0], dtype=np.float32)
         region_id = np.array([7], dtype=np.int32)
 
-        # Single tet: vertex_count == tri_count == 4, so temperature needs explicit frequency
+        # Single tet: vertex_count == tri_count == 4, so temperature needs explicit frequency.
+        # regionId also needs explicit frequency because tet_count == 1 is ambiguous with ONCE.
         tm = newton.TetMesh(
             vertices,
             tet_indices,
             custom_attributes={
                 "temperature": (temperature, newton.Model.AttributeFrequency.PARTICLE),
-                "regionId": region_id,
+                "regionId": (region_id, newton.Model.AttributeFrequency.TETRAHEDRON),
             },
         )
 
@@ -11624,6 +11625,37 @@ def test_tetmesh_custom_attributes_constructor(self):
         assert_np_equal(arr, region_id)
         self.assertEqual(freq, newton.Model.AttributeFrequency.TETRAHEDRON)
 
+    def test_tetmesh_custom_attributes_infer_once(self):
+        """Test that length-1 arrays are inferred as ONCE when unambiguous."""
+        vertices = np.array(
+            [[0, 0, 0], [1, 0, 0], [0, 1, 0], [0, 0, 1], [0.5, 0.5, 0.5]],
+            dtype=np.float32,
+        )
+        tet_indices = np.array([0, 1, 2, 3, 0, 1, 2, 4], dtype=np.int32)
+        constant = np.array([42.0], dtype=np.float32)
+
+        tm = newton.TetMesh(
+            vertices,
+            tet_indices,
+            custom_attributes={"constant": constant},
+        )
+
+        arr, freq = tm.custom_attributes["constant"]
+        assert_np_equal(arr, constant)
+        self.assertEqual(freq, newton.Model.AttributeFrequency.ONCE)
+
+    def test_tetmesh_custom_attributes_ambiguous_once(self):
+        """Test that length-1 arrays raise when tet_count is also 1."""
+        vertices = np.array([[0, 0, 0], [1, 0, 0], [0, 1, 0], [0, 0, 1]], dtype=np.float32)
+        tet_indices = np.array([0, 1, 2, 3], dtype=np.int32)
+
+        with self.assertRaisesRegex(ValueError, "ONCE"):
+            newton.TetMesh(
+                vertices,
+                tet_indices,
+                custom_attributes={"ambig": np.array([1.0], dtype=np.float32)},
+            )
+
     def test_tetmesh_custom_attributes_empty_by_default(self):
         """Test TetMesh has empty custom_attributes when none are provided."""
         vertices = np.array([[0, 0, 0], [1, 0, 0], [0, 1, 0], [0, 0, 1]], dtype=np.float32)
@@ -11669,13 +11701,14 @@ def test_tetmesh_custom_attributes_npz_roundtrip(self):
         temperature = np.array([10.0, 20.0, 30.0, 40.0], dtype=np.float32)
         region_id = np.array([3], dtype=np.int32)
 
-        # Single tet: vertex_count == tri_count == 4, so temperature needs explicit frequency
+        # Single tet: vertex_count == tri_count == 4, so temperature needs explicit frequency.
+        # regionId also needs explicit frequency because tet_count == 1 is ambiguous with ONCE.
         tm = newton.TetMesh(
             vertices,
             tet_indices,
             custom_attributes={
                 "temperature": (temperature, newton.Model.AttributeFrequency.PARTICLE),
-                "regionId": region_id,
+                "regionId": (region_id, newton.Model.AttributeFrequency.TETRAHEDRON),
             },
         )