ENH: Write UsdGeomTetMesh from tetrahedral grids; stress-colored tutorial 18

pyproject.toml

@@ -119,7 +119,7 @@ dependencies = [
 
   # Visualization and USD
   "pyvista[all]>=0.47.0",
-  "usd-core>=23.11",
+  "usd-core>=24.5",
   "trimesh>=4.0.0",
   "pyacvd>=0.4.0",
 

src/monai_physio/convert_vtk_to_usd.py

@@ -89,6 +89,7 @@ def __init__(
         object_names: Optional[Sequence[str]] = None,
         segmenter: Optional[SegmentAnatomyBase] = None,
         log_level: int | str = logging.INFO,
+        output_tet_mesh: bool = False,
     ) -> None:
         """
         Initialize converter.
@@ -130,6 +131,10 @@ def __init__(
                        and materials under ``/World/Looks/{type}/{label}_material``.
                        When None, labels are grouped under a single ``Anatomy`` Xform.
             log_level: Logging level
+            output_tet_mesh: If True, write each tetrahedral unstructured grid
+                      as a UsdGeomTetMesh holding every tetrahedron instead of
+                      its extracted surface. Overrides ``convert_to_surface``
+                      and does not combine with ``mask_ids``.
 
         Raises:
             ValueError: If time_codes or object_names is not None and its length
@@ -162,6 +167,7 @@ def __init__(
         self.mask_ids = mask_ids
         self.compute_normals = compute_normals
         self.convert_to_surface = convert_to_surface
+        self.output_tet_mesh = output_tet_mesh
         self.separate_by = separate_by
         self.solid_color = solid_color
         self.segmenter = segmenter
@@ -250,6 +256,7 @@ def from_files(
         mask_ids: Optional[dict[int, str]] = None,
         segmenter: Optional[SegmentAnatomyBase] = None,
         log_level: int | str = logging.INFO,
+        output_tet_mesh: bool = False,
     ) -> ConvertVTKToUSD:
         """Create a converter by loading VTK files from disk.
 
@@ -273,6 +280,9 @@ def from_files(
             segmenter: Optional SegmentAnatomyBase used to group labeled prims
                        by anatomy type. See the ConvertVTKToUSD constructor.
             log_level: Logging level.
+            output_tet_mesh: If True, keep tetrahedral unstructured grids whole
+                and write them as UsdGeomTetMesh; the surface is then not
+                extracted. See the ConvertVTKToUSD constructor.
 
         Returns:
             ConvertVTKToUSD instance ready to call .convert().
@@ -284,7 +294,11 @@ def from_files(
         meshes: list[pv.DataSet | vtk.vtkDataSet] = []
         for path in file_list:
             mesh = pv.read(str(path))
-            if extract_surface and isinstance(mesh, pv.UnstructuredGrid):
+            if (
+                extract_surface
+                and not output_tet_mesh
+                and isinstance(mesh, pv.UnstructuredGrid)
+            ):
                 mesh = mesh.extract_surface(algorithm="dataset_surface")
             meshes.append(mesh)
 
@@ -306,6 +320,7 @@ def from_files(
             time_codes=resolved_time_codes,
             segmenter=segmenter,
             log_level=log_level,
+            output_tet_mesh=output_tet_mesh,
         )
 
         # Validate topology consistency for multi-frame time series and cache the
@@ -888,8 +903,24 @@ def _vtk_to_mesh_data(
         if isinstance(vtk_mesh, vtk.vtkDataSet):
             vtk_mesh = pv.wrap(vtk_mesh)
 
+        tet_indices: Optional[np.ndarray] = None
+        if self.output_tet_mesh:
+            if not isinstance(vtk_mesh, pv.UnstructuredGrid) or set(
+                np.unique(vtk_mesh.celltypes)
+            ) != {pv.CellType.TETRA}:
+                raise ValueError(
+                    "output_tet_mesh needs an unstructured grid of tetrahedra only"
+                )
+            tet_indices = np.asarray(
+                vtk_mesh.cells_dict[np.uint8(pv.CellType.TETRA)], dtype=np.int32
+            )
+
         # Extract surface if needed
-        if self.convert_to_surface and not isinstance(vtk_mesh, pv.PolyData):
+        if (
+            self.convert_to_surface
+            and tet_indices is None
+            and not isinstance(vtk_mesh, pv.PolyData)
+        ):
             if isinstance(vtk_mesh, pv.UnstructuredGrid) or hasattr(
                 vtk_mesh, "extract_surface"
             ):
@@ -901,7 +932,10 @@ def _vtk_to_mesh_data(
         points = np.array(vtk_mesh.points, dtype=np.float64)
 
         # Get faces
-        if hasattr(vtk_mesh, "faces"):
+        if tet_indices is not None:
+            face_counts = np.array([], dtype=np.int32)
+            face_indices = np.array([], dtype=np.int32)
+        elif hasattr(vtk_mesh, "faces"):
             faces = vtk_mesh.faces
             if len(faces) == 0 and not self.convert_to_surface:
                 raise ValueError("Mesh has no faces - surface extraction may be needed")
@@ -970,6 +1004,7 @@ def _vtk_to_mesh_data(
             normals=normals,
             colors=colors,
             generic_arrays=generic_arrays,
+            tet_vertex_indices=tet_indices,
         )
 
     def _split_by_labels(

src/monai_physio/process_surface_mesh.py

@@ -12,6 +12,7 @@
 from __future__ import annotations
 
 import logging
+import warnings
 from collections.abc import Mapping, Sequence
 from typing import Any, Optional, cast
 
@@ -318,16 +319,19 @@ def remesh_and_smooth_surface(
                 piece = regions.threshold(
                     (region_id, region_id), scalars="RegionId", preference="cell"
                 ).extract_surface(algorithm="geometry")
-                clustering = pyacvd.Clustering(piece)
-                # One cluster per retained point.  A closed surface carries
-                # about twice as many triangles as points, so scaling the
-                # point count by (1 - rate) scales the triangle count by the
-                # same fraction; four is the fewest clusters that can still
-                # close a surface.
-                clustering.cluster(
-                    max(4, round(piece.n_points * (1.0 - surface_reduction_rate)))
-                )
-                remeshed_pieces.append(clustering.create_mesh())
+                # pyacvd 0.4.0 still uses a pyvista-deprecated private API.
+                with warnings.catch_warnings():
+                    warnings.filterwarnings("ignore", message=".*_connectivity_array.*")
+                    clustering = pyacvd.Clustering(piece)
+                    # One cluster per retained point.  A closed surface carries
+                    # about twice as many triangles as points, so scaling the
+                    # point count by (1 - rate) scales the triangle count by the
+                    # same fraction; four is the fewest clusters that can still
+                    # close a surface.
+                    clustering.cluster(
+                        max(4, round(piece.n_points * (1.0 - surface_reduction_rate)))
+                    )
+                    remeshed_pieces.append(clustering.create_mesh())
             conditioned = pv.merge(remeshed_pieces)
             if "SegmentationLabelIds" in original.cell_data:
                 nearest = original.find_closest_cell(conditioned.cell_centers().points)

src/monai_physio/vtk_to_usd/data_structures.py

@@ -108,6 +108,9 @@ class MeshData:
     colors: Optional[NDArray] = None  # Shape: (N, 3) or (N, 4)
     generic_arrays: list[GenericArray] = field(default_factory=list)
     material_id: str = "default_material"
+    # When set, the mesh is written as a UsdGeomTetMesh from these (T, 4) point
+    # indices and the face arrays are ignored.
+    tet_vertex_indices: Optional[NDArray] = None
 
     def __post_init__(self) -> None:
         """Validate mesh data."""

src/monai_physio/vtk_to_usd/usd_mesh_converter.py

@@ -95,8 +95,14 @@ def create_mesh(
         """
         logger.info(f"Creating USD mesh at: {mesh_path}")
 
+        is_tet_mesh = mesh_data.tet_vertex_indices is not None
+
         # Create mesh prim
-        mesh = UsdGeom.Mesh.Define(self.stage, mesh_path)
+        mesh = (
+            UsdGeom.TetMesh.Define(self.stage, mesh_path)
+            if is_tet_mesh
+            else UsdGeom.Mesh.Define(self.stage, mesh_path)
+        )
 
         # Convert points to USD coordinates
         usd_points = lps_points_to_usd(mesh_data.points)
@@ -112,8 +118,15 @@ def create_mesh(
 
         # Set topology as the default value. create_time_varying_mesh() adds
         # time samples on top of this when a series changes topology.
-        mesh.CreateFaceVertexCountsAttr(face_counts_vt)
-        mesh.CreateFaceVertexIndicesAttr(face_indices_vt)
+        if is_tet_mesh:
+            mesh.CreateTetVertexIndicesAttr(
+                Vt.Vec4iArray.FromNumpy(
+                    np.asarray(mesh_data.tet_vertex_indices, dtype=np.int32)
+                )
+            )
+        else:
+            mesh.CreateFaceVertexCountsAttr(face_counts_vt)
+            mesh.CreateFaceVertexIndicesAttr(face_indices_vt)
 
         # Set points (time-varying if time_code provided). Also author a
         # default value for readers that inspect the prim without a time code.
@@ -136,7 +149,21 @@ def create_mesh(
             extent_attr.Set(extent)
 
         # Set mesh attributes
-        mesh.CreateSubdivisionSchemeAttr("none")  # No subdivision
+        if is_tet_mesh:
+            # The viewer draws only surfaceFaceVertexIndices, so author them
+            # from the tetrahedra, and pick the orientation that leaves the
+            # tetrahedra uninverted.
+            tet_mesh = UsdGeom.TetMesh(mesh)
+            default_time = Usd.TimeCode.Default()
+            if len(UsdGeom.TetMesh.FindInvertedElements(tet_mesh, time_code or 0.0)) > (
+                len(np.asarray(mesh_data.tet_vertex_indices)) // 2
+            ):
+                tet_mesh.CreateOrientationAttr(UsdGeom.Tokens.leftHanded)
+            tet_mesh.CreateSurfaceFaceVertexIndicesAttr(
+                UsdGeom.TetMesh.ComputeSurfaceFaces(tet_mesh, default_time)
+            )
+        else:
+            mesh.CreateSubdivisionSchemeAttr("none")  # No subdivision
         mesh.CreateDoubleSidedAttr(True)  # Visible from both sides
 
         # Handle normals
@@ -176,9 +203,13 @@ def create_mesh(
             material = self.material_mgr.material_cache[mesh_data.material_id]
             self.material_mgr.bind_material(mesh, material)
 
+        element_count = (
+            f"{len(mesh_data.tet_vertex_indices)} tetrahedra"
+            if mesh_data.tet_vertex_indices is not None
+            else f"{len(face_counts)} faces"
+        )
         logger.info(
-            f"Created mesh with {len(mesh_data.points)} points, "
-            f"{len(face_counts)} faces"
+            f"Created mesh with {len(mesh_data.points)} points, {element_count}"
         )
 
         return mesh

tutorials/tutorial_18_duke_heart_physics_informed_motion_infer.py

@@ -54,12 +54,15 @@
 
 import contextlib
 import csv
+import itertools
 import logging
 from pathlib import Path
 from typing import Any, Optional, cast
 
+import matplotlib
 import numpy as np
 import pyvista as pv
+from matplotlib.colors import LinearSegmentedColormap
 from parameters_duke_heart_physics_informed import DUKE_HEART_PHYSICS_INFORMED
 
 from monai_physio import (
@@ -84,7 +87,59 @@
 
 # Range the stress colormap is stretched over, in kilopascals.  Fixed rather
 # than per-frame, so a color means the same thing throughout the animation.
-STRESS_RANGE_KPA = (0.0, 50.0)
+STRESS_RANGE_KPA = (1.0, 100.0)
+
+# Colored on log10 of the stress so the low end is resolved and a few hot spots
+# do not dominate the scale.
+LOG_STRESS_ARRAY = "log10_von_mises_stress"
+STRESS_COLORMAP = "heart_heat"
+_stress_cmap = LinearSegmentedColormap.from_list(
+    STRESS_COLORMAP, ["#5a0a10", "#a01418", "#e0481c", "#ffd23c"]
+)
+matplotlib.colormaps.register(_stress_cmap, force=True)
+
+# Fraction of the log range the colormap is shifted up by, which moves more of
+# the surface toward the red end.
+COLORMAP_SHIFT = 0.2
+_log_low = float(np.log10(STRESS_RANGE_KPA[0]))
+_log_high = float(np.log10(STRESS_RANGE_KPA[1]))
+_log_shift = COLORMAP_SHIFT * (_log_high - _log_low)
+LOG_STRESS_RANGE = (_log_low + _log_shift, _log_high + _log_shift)
+
+# Passes of neighbor averaging over the tetrahedral mesh edges applied to the
+# nodal stress tensor, which removes element-to-element noise.
+STRESS_SMOOTHING_ITERATIONS = 20
+
+# Passes of a median filter over each node's mesh neighbors applied to the
+# point positions of the rendered surface, which removes isolated spikes.  It
+# can fold elements, so it is applied to the surface only, never the tet mesh.
+POSITION_MEDIAN_PASSES = 4
+
+
+def _cell_edges(cells: np.ndarray) -> np.ndarray:
+    """Return the ``(n_edges, 2)`` directed node pairs of simplex ``cells``."""
+    edges = np.concatenate(
+        [cells[:, [a, b]] for a, b in itertools.combinations(range(cells.shape[1]), 2)]
+    )
+    return np.unique(np.concatenate([edges, edges[:, ::-1]]), axis=0)
+
+
+def _median_filter_points(points: np.ndarray, cells: np.ndarray) -> np.ndarray:
+    """Return ``points`` with each node set to the per-axis median of itself
+    and its edge neighbors in ``cells``."""
+    points = np.asarray(points, dtype=np.float64)
+    edges = _cell_edges(cells)
+    edges = edges[np.argsort(edges[:, 0], kind="stable")]
+    source = edges[:, 0]
+    degree = np.bincount(source, minlength=len(points))
+    starts = np.cumsum(degree) - degree
+    slot = np.arange(len(source)) - starts[source]
+    for _ in range(POSITION_MEDIAN_PASSES):
+        window = np.full((len(points), degree.max() + 1, 3), np.nan)
+        window[:, 0] = points
+        window[source, slot + 1] = points[edges[:, 1]]
+        points = np.nanmedian(window, axis=1)
+    return points
 
 
 def _stress_at_points(
@@ -126,7 +181,15 @@ def _stress_at_points(
         np.repeat(element_stress, 4, axis=0) * weights[:, None],
     )
     np.add.at(nodal_weight, flat_nodes, weights)
-    return nodal_stress / np.maximum(nodal_weight, 1.0e-12)[:, None]
+    nodal_stress /= np.maximum(nodal_weight, 1.0e-12)[:, None]
+
+    edges = _cell_edges(tets)
+    degree = np.maximum(np.bincount(edges[:, 0], minlength=len(nodal_stress)), 1)
+    for _ in range(STRESS_SMOOTHING_ITERATIONS):
+        neighbor_sum = np.zeros_like(nodal_stress)
+        np.add.at(neighbor_sum, edges[:, 0], nodal_stress[edges[:, 1]])
+        nodal_stress = 0.5 * nodal_stress + 0.5 * neighbor_sum / degree[:, None]
+    return nodal_stress
 
 
 def _render_stress(surface: pv.DataSet, plot_file: Path) -> Path:
@@ -143,10 +206,10 @@ def _render_stress(surface: pv.DataSet, plot_file: Path) -> Path:
         plotter = pv.Plotter(off_screen=True, window_size=[800, 600])
         plotter.add_mesh(
             surface,
-            scalars="von_mises_stress",
-            cmap="rainbow",
-            clim=STRESS_RANGE_KPA,
-            scalar_bar_args={"title": "von Mises (kPa)"},
+            scalars=LOG_STRESS_ARRAY,
+            cmap=_stress_cmap,
+            clim=LOG_STRESS_RANGE,
+            scalar_bar_args={"title": "log10 von Mises (kPa)"},
         )
         plotter.camera_position = "iso"
         plotter.screenshot(str(plot_file))
@@ -339,6 +402,10 @@ def _render_stress(surface: pv.DataSet, plot_file: Path) -> Path:
                 f"the fitted reference carries {reference_model.n_points}; the "
                 "deformation gradient needs them on the same nodes."
             )
+        # compute_von_mises_stress reads cell data before point data, so a
+        # stale per-cell tensor of the same name would shadow the nodal one.
+        if STRESS_ARRAY in predicted_mesh.cell_data:
+            del predicted_mesh.cell_data[STRESS_ARRAY]
         predicted_mesh.point_data[STRESS_ARRAY] = _stress_at_points(
             reference_model, predicted_mesh, tets, law
         )
@@ -364,9 +431,14 @@ def _render_stress(surface: pv.DataSet, plot_file: Path) -> Path:
         stress_meshes,
         frames_per_second=float(len(stress_meshes)),
         log_level=log_level,
+        output_tet_mesh=True,
     )
     converter.compute_von_mises_stress(STRESS_ARRAY)
-    converter.set_colormap("von_mises_stress", "rainbow", STRESS_RANGE_KPA)
+    for mesh in stress_meshes:
+        mesh.point_data[LOG_STRESS_ARRAY] = np.log10(
+            np.maximum(mesh.point_data["von_mises_stress"], STRESS_RANGE_KPA[0])
+        )
+    converter.set_colormap(LOG_STRESS_ARRAY, STRESS_COLORMAP, LOG_STRESS_RANGE)
     converter.convert(str(usd_file))
     tutorial_results["usd_file"] = usd_file
     logger.log_info("USD: %s", usd_file)
@@ -379,6 +451,9 @@ def _render_stress(surface: pv.DataSet, plot_file: Path) -> Path:
         log_level=log_level,
     )
     stress_surface = stress_meshes[-1].extract_surface(algorithm="dataset_surface")
+    stress_surface.points = _median_filter_points(
+        stress_surface.points, stress_surface.faces.reshape(-1, 4)[:, 1:]
+    )
     tutorial_results["screenshots"] = [
         _render_stress(stress_surface, output_dir / "von_mises_stress.png"),
         tt.save_screenshot_mesh(