ENH: Reduce computation for inversion reduction

Author: aylward

src/monai_physio/contour_tools.py

@@ -780,9 +780,15 @@ def repair_inverted_tetrahedra(
             already clears zero volume.
 
         Raises:
-            ValueError: If elements are still inverted or degenerate after
-                *max_iterations* passes.
+            ValueError: If *tetrahedra* has no TETRA cells, or if elements
+                are still inverted or degenerate after *max_iterations*
+                passes.
         """
+        if np.uint8(pv.CellType.TETRA) not in tetrahedra.cells_dict:
+            raise ValueError(
+                "tetrahedra has no TETRA cells to repair; got cell types "
+                f"{sorted(tetrahedra.cells_dict)}."
+            )
         connectivity = tetrahedra.cells_dict[np.uint8(pv.CellType.TETRA)]
 
         def volumes(points: np.ndarray) -> np.ndarray:
@@ -803,13 +809,22 @@ def volumes(points: np.ndarray) -> np.ndarray:
         )
         starts = np.concatenate([edges[:, 0], edges[:, 1]])
         ends = np.concatenate([edges[:, 1], edges[:, 0]])
+        order = np.argsort(starts, kind="stable")
+        sorted_starts = starts[order]
+        sorted_ends = ends[order]
+        split_points = np.searchsorted(sorted_starts, np.arange(len(points) + 1))
+        neighbors = {
+            node: np.unique(sorted_ends[split_points[node] : split_points[node + 1]])
+            for node in np.unique(connectivity)
+        }
 
         for _ in range(max_iterations):
             bad = volumes(points) <= 0.0
             if not np.any(bad):
                 break
+            snapshot = points.copy()
             for node in np.unique(connectivity[bad]):
-                neighbor_points = points[ends[starts == node]]
+                neighbor_points = snapshot[neighbors[node]]
                 if len(neighbor_points):
                     points[node] = neighbor_points.mean(axis=0)
 

src/monai_physio/train_physicsnemo_physics_informed_motion.py

@@ -651,9 +651,14 @@ def _bind_reference_meshes(
             # A per-subject fit carries no cell-quality constraint, so it can
             # flip a handful of elements even though the template did not;
             # repair here rather than only at fit time so meshes fitted before
-            # this check existed still load.
-            mesh = contour_tools.repair_inverted_tetrahedra(mesh)
-            points = np.asarray(mesh.points, dtype=np.float64)
+            # this check existed still load. Repair against self._tets, the
+            # connectivity tet_volumes() below actually uses, rather than
+            # whatever cells the file happens to store -- a mismatch there
+            # would repair the wrong topology and still leave the physics
+            # elements inverted.
+            tet_grid = pv.UnstructuredGrid({pv.CellType.TETRA: self._tets}, mesh.points)
+            repaired = contour_tools.repair_inverted_tetrahedra(tet_grid)
+            points = np.asarray(repaired.points, dtype=np.float64)
             _, nodal = tet_volumes(points, self._tets)
             reference = torch.from_numpy(points).to(device=device, dtype=torch.float32)
             volumes = torch.from_numpy(nodal).to(device=device, dtype=torch.float32)

tests/test_contour_mesh_extraction.py

@@ -407,5 +407,45 @@ def test_anatomy_color_survives(self, contour_tools: ContourTools) -> None:
         )
 
 
+class TestRepairInvertedTetrahedra:
+    """Node relaxation must fix a flippable mesh and give up on one that can't."""
+
+    @staticmethod
+    def _single_tetra(points: np.ndarray) -> pv.UnstructuredGrid:
+        """Build a one-cell tetrahedral mesh from four *points*."""
+        cells = np.array([4, 0, 1, 2, 3])
+        return pv.UnstructuredGrid(cells, [pv.CellType.TETRA], points)
+
+    def test_repairs_an_inverted_tetrahedron(self, contour_tools: ContourTools) -> None:
+        """Swapping two corners inverts the cell; relaxation must restore it."""
+        points = np.array(
+            [[0.0, 0.0, 0.0], [1.0, 1.0, 0.0], [1.0, 0.0, 0.0], [0.0, 0.0, 1.0]]
+        )
+        mesh = self._single_tetra(points)
+
+        repaired = contour_tools.repair_inverted_tetrahedra(mesh)
+
+        corners = repaired.points[repaired.cells_dict[np.uint8(pv.CellType.TETRA)]][0]
+        edges = corners[1:, :] - corners[0:1, :]
+        assert np.linalg.det(edges) / 6.0 > 0.0
+
+    def test_raises_when_unrecoverable(self, contour_tools: ContourTools) -> None:
+        """A degenerate cell with no neighbors to average toward can't be fixed."""
+        points = np.array(
+            [[0.0, 0.0, 0.0], [1.0, 0.0, 0.0], [2.0, 0.0, 0.0], [3.0, 0.0, 0.0]]
+        )
+        mesh = self._single_tetra(points)
+
+        with pytest.raises(ValueError, match="still inverted or degenerate"):
+            contour_tools.repair_inverted_tetrahedra(mesh, max_iterations=2)
+
+    def test_raises_when_no_tetra_cells(self, contour_tools: ContourTools) -> None:
+        """A mesh without TETRA cells fails with a clear message, not a KeyError."""
+        mesh = pv.UnstructuredGrid()
+
+        with pytest.raises(ValueError, match="no TETRA cells"):
+            contour_tools.repair_inverted_tetrahedra(mesh)
+
+
 if __name__ == "__main__":
     pytest.main([__file__, "-v", "-s"])

tests/test_physics_informed_motion.py

@@ -498,3 +498,70 @@ def test_the_epoch_log_separates_the_two_loss_terms() -> None:
     assert "physics=0.000000" in messages[-1], (
         f"An epoch that accumulated nothing should report zero: {messages[-1]}"
     )
+
+
+def test_bind_reference_meshes_repairs_against_template_elements(tmp_path: Any) -> None:
+    """Repair must use ``self._tets``, not whatever cells the file stores.
+
+    A fitted reference file's own connectivity is never read back by
+    ``tet_volumes`` -- only ``self._tets`` is -- so repairing against the
+    file's cells instead would validate the wrong topology. Here the file
+    stores one degenerate, unrecoverable cell that touches only a single
+    node; repairing against it would raise, but ``self._tets`` names a
+    perfectly valid mesh, so binding must succeed unchanged.
+    """
+    import torch
+
+    from monai_physio.physicsnemo_tools import DistributedContext
+    from monai_physio.train_physicsnemo_physics_informed_motion import (
+        TrainPhysicsNeMoPhysicsInformedMotion,
+    )
+
+    points, tets = _grid_mesh(size=3)
+    mismatched = pv.UnstructuredGrid(
+        {pv.CellType.TETRA: np.array([[0, 0, 0, 0]])}, points
+    )
+    mesh_path = tmp_path / "reference.vtu"
+    mismatched.save(mesh_path)
+
+    method = TrainPhysicsNeMoPhysicsInformedMotion()
+    method._tets = tets
+    method._sample_subjects = ["subj0"]
+    method._reference_meshes = {"subj0": mesh_path}
+    context = DistributedContext(
+        device=torch.device("cpu"), rank=0, local_rank=0, world_size=1
+    )
+
+    method._bind_reference_meshes(context, n_points=len(points))
+
+    reference, volumes = method._reference_cache["subj0"]
+    assert np.allclose(reference.numpy(), points, atol=1e-5)
+    assert torch.all(volumes > 0)
+
+
+def test_bind_reference_meshes_tolerates_a_file_with_no_cells(tmp_path: Any) -> None:
+    """A reference file need not carry any cells at all; only its points do."""
+    import torch
+
+    from monai_physio.physicsnemo_tools import DistributedContext
+    from monai_physio.train_physicsnemo_physics_informed_motion import (
+        TrainPhysicsNeMoPhysicsInformedMotion,
+    )
+
+    points, tets = _grid_mesh(size=3)
+    mesh_path = tmp_path / "reference.vtp"
+    pv.PolyData(points).save(mesh_path)
+
+    method = TrainPhysicsNeMoPhysicsInformedMotion()
+    method._tets = tets
+    method._sample_subjects = ["subj0"]
+    method._reference_meshes = {"subj0": mesh_path}
+    context = DistributedContext(
+        device=torch.device("cpu"), rank=0, local_rank=0, world_size=1
+    )
+
+    method._bind_reference_meshes(context, n_points=len(points))
+
+    reference, volumes = method._reference_cache["subj0"]
+    assert np.allclose(reference.numpy(), points, atol=1e-5)
+    assert torch.all(volumes > 0)