FlipBoundaryOrientation: reverse the winding with Mesh.with_cells

CHANGELOG.md

@@ -563,6 +563,9 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - The container build compiles pyg-lib from source without the uv cache, so
   it no longer reuses a pyg-lib wheel built against an older PyTorch, which
   failed to load with an undefined symbol after a base image upgrade.
+- `Mesh.with_cells` and `fix_orientation` keep the effective measure and area
+  of each cell whose vertices are only reordered, as when flipping a winding,
+  instead of rescaling them by float32 roundoff or raising on zero-area cells.
 - `DistributedManager` now initializes on CPU-only hosts, where PyTorch rejected
   its CPU `device_id` and a missing local rank raised `ZeroDivisionError`.
   `cleanup()` no longer raises after a failed initialization.

physicsnemo/datapipes/transforms/mesh/transforms.py

@@ -1000,19 +1000,16 @@ def __call__(self, mesh: Mesh) -> Mesh:
             raise ValueError(
                 "FlipBoundaryOrientation needs cells of at least 2 vertices"
             )
-        cells = torch.cat((mesh.cells[:, :-2], mesh.cells[:, -2:].flip(1)), dim=1)
-        cell_data = mesh.cell_data.copy()  # only the negated fields get new tensors
+        ### Each cell keeps its vertices, so with_cells keeps its measure and drops
+        ### the normals cached for the old winding.
+        flipped = mesh.with_cells(
+            torch.cat((mesh.cells[:, :-2], mesh.cells[:, -2:].flip(1)), dim=1)
+        )
+        cell_data = flipped.cell_data.copy()  # only the negated fields get new tensors
         for key in self.normal_fields:
             if key in cell_data.keys(include_nested=True):
                 cell_data[key] = -cell_data[key]
-        ### A new Mesh, so no cached normals of the old winding survive.
-        return Mesh(
-            points=mesh.points,
-            cells=cells,
-            point_data=mesh.point_data,
-            cell_data=cell_data,
-            global_data=mesh.global_data,
-        )
+        return flipped.with_data(cell_data=cell_data)
 
     def apply_to_domain(self, domain: DomainMesh) -> DomainMesh:
         """Reverse the named boundaries of a :class:`DomainMesh`.

physicsnemo/mesh/calculus/measure.py

@@ -201,28 +201,49 @@ def lumped_point_measures(mesh: "Mesh") -> torch.Tensor:
 
 
 def _transfer_cell_measures(
-    source: "Mesh", result: "Mesh", parents: torch.Tensor | None = None
+    source: "Mesh",
+    result: "Mesh",
+    parents: torch.Tensor | None = None,
+    unchanged: torch.Tensor | None = None,
 ) -> None:
-    """Transfer explicit cell measures through geometric changes/subdivision."""
+    """Transfer explicit cell measures through geometric changes/subdivision.
+
+    Result row ``i`` comes from source row ``parents[i]``, or from row ``i``
+    without ``parents``. ``unchanged`` is a boolean mask over result rows: a
+    marked row only reorders its source row's vertices, so it keeps that row's
+    measure bitwise and, outside tracing and autograd, its cached area.
+    """
     if EFFECTIVE_MEASURE_KEY not in source.cell_data:
         return
     measures = cell_measures(source)
     old_areas = source.cell_areas
     if parents is not None:
         measures, old_areas = measures[parents], old_areas[parents]
-    if not is_tracing() and bool(((old_areas == 0) & (measures != 0)).any()):
+    new_areas = result.cell_areas
+    degenerate = (old_areas == 0) & (measures != 0)
+    if unchanged is not None:
+        new_areas = torch.where(unchanged, old_areas, new_areas)
+        degenerate = degenerate & ~unchanged
+        # Later transfers divide by the cached area, so cache the bits paired
+        # with the measure rather than a recompute in the new vertex order.
+        if not (is_tracing() or result.points.requires_grad):
+            result._cache["cell", "areas"] = new_areas
+    if not is_tracing() and bool(degenerate.any()):
         raise ValueError(
             "Cannot transform nonzero effective measure on a zero-measure cell; supply replacement measures"
         )
     denominator = torch.where(old_areas != 0, old_areas, torch.ones_like(old_areas))
     # The measures follow the widest dtype involved; the area ratio can overflow
     # float16, so compute in at least float32.
     out_dtype = torch.promote_types(
-        torch.promote_types(measures.dtype, result.cell_areas.dtype), old_areas.dtype
+        torch.promote_types(measures.dtype, new_areas.dtype), old_areas.dtype
     )
     dtype = torch.promote_types(out_dtype, torch.float32)
-    ratio = result.cell_areas.to(dtype) / denominator.to(dtype)
-    set_cell_measures(result, (measures.to(dtype) * ratio).to(out_dtype))
+    ratio = new_areas.to(dtype) / denominator.to(dtype)
+    transferred = (measures.to(dtype) * ratio).to(out_dtype)
+    if unchanged is not None:
+        transferred = torch.where(unchanged, measures.to(out_dtype), transferred)
+    set_cell_measures(result, transferred)
 
 
 def _require_preserved_point_measures(mesh: "Mesh") -> None:

physicsnemo/mesh/mesh.py

@@ -2098,8 +2098,11 @@ def with_cells(
         This operation is for connectivity changes that preserve cell indexing
         and simplex type, such as reversing the winding of selected triangles.
         The replacement must have exactly the same shape as the current cells.
-        User data is preserved, while all caches are cleared by default because
+        User data is preserved, while caches are cleared by default because
         changing connectivity can invalidate cell, point, and topology caches.
+        A row that only reorders its own vertices is the same simplex, so it
+        keeps its effective measure bitwise and, as the Notes describe, its
+        cached centroid and area.
 
         Parameters
         ----------
@@ -2113,7 +2116,8 @@ def with_cells(
         -------
         Mesh
             New mesh with replacement connectivity, preserved points and data,
-            and only the selected caches.
+            and the selected caches. It may also cache cell centroids and
+            areas; see Notes.
 
         Raises
         ------
@@ -2128,6 +2132,12 @@ def with_cells(
         expert operation: the caller is responsible for ensuring every retained
         value remains valid for the replacement connectivity.
 
+        The result also caches cell centroids when this mesh has them cached,
+        and cell areas when it has effective measures, unless ``keep`` retains
+        them. Neither is cached while tracing or when ``points`` requires
+        gradients. Both are computed eagerly for all cells, then each row that
+        only reorders its vertices takes this mesh's values bitwise.
+
         Examples
         --------
         >>> flipped = mesh.with_cells(  # doctest: +SKIP
@@ -2155,9 +2165,28 @@ def with_cells(
             cells=cells,
             keep=keep,
         )
-        from physicsnemo.mesh.calculus.measure import _transfer_cell_measures
+        from physicsnemo.mesh.calculus.measure import (
+            EFFECTIVE_MEASURE_KEY,
+            _transfer_cell_measures,
+        )
 
-        _transfer_cell_measures(self, result)
+        # Like _cached, leave the cache alone while tracing or under autograd.
+        centroids = (
+            None
+            if is_tracing() or self.points.requires_grad
+            else self._cache.get(("cell", "centroids"), None)
+        )
+        unchanged = None
+        if centroids is not None or EFFECTIVE_MEASURE_KEY in self.cell_data:
+            # Use the normalized cells: PyTorch compares uint16, uint32 and
+            # uint64 only with their own dtype.
+            sorted_cells = result.cells.sort(dim=1).values
+            unchanged = (sorted_cells == self.cells.sort(dim=1).values).all(dim=1)
+        if centroids is not None:
+            result._cache["cell", "centroids"] = torch.where(
+                unchanged[:, None], centroids, result.cell_centroids
+            )
+        _transfer_cell_measures(self, result, unchanged=unchanged)
         return result
 
     def with_data(

test/datapipes/transforms/test_flip_boundary_orientation.py

@@ -22,7 +22,9 @@
 from physicsnemo.datapipes.registry import COMPONENT_REGISTRY
 from physicsnemo.datapipes.transforms.mesh import FlipBoundaryOrientation
 from physicsnemo.mesh import DomainMesh, Mesh
+from physicsnemo.mesh.calculus.measure import cell_measures, set_cell_measures
 from physicsnemo.mesh.primitives.basic import single_triangle_3d
+from physicsnemo.mesh.primitives.surfaces import sphere_icosahedral
 
 
 def _surface() -> Mesh:
@@ -57,6 +59,18 @@ def test_mesh_reversed_with_its_stored_normals(self):
         ### The input keeps its own stored normals.
         torch.testing.assert_close(mesh.cell_data["normals"], mesh.cell_normals)
 
+    def test_effective_measures_keep_their_areas(self):
+        """Each cell keeps its measure and the area it came from, bit for bit,
+        though float32 areas differ between the two vertex orders."""
+        mesh = sphere_icosahedral.load(subdivisions=3)
+        set_cell_measures(mesh, mesh.cell_areas * 2.0)
+        out = FlipBoundaryOrientation()(mesh)
+        assert not torch.equal(
+            Mesh(points=mesh.points, cells=out.cells).cell_areas, mesh.cell_areas
+        )
+        assert torch.equal(cell_measures(out), cell_measures(mesh))
+        assert torch.equal(out.cell_areas, mesh.cell_areas)
+
     def test_flip_twice_is_identity(self):
         """Reversing twice restores the cells and the stored normals exactly."""
         mesh = _surface()

test/mesh/calculus/test_measure.py

@@ -246,3 +246,88 @@ def test_weights_invariant_under_rigid_and_scaling_transforms(self):
         torch.testing.assert_close(
             moved.integrate("f"), base_integral / 25.0, rtol=1e-5, atol=1e-7
         )
+
+    @staticmethod
+    def _flip(mesh: Mesh, order: tuple[int, ...] = (0, 2, 1)) -> Mesh:
+        return mesh.with_cells(mesh.cells[:, list(order)])
+
+    def test_winding_flip_keeps_subsampled_measures_and_areas_bitwise(self, device):
+        """A reordered row is the same simplex: its measure and area are kept."""
+        g = torch.Generator().manual_seed(0)
+        mesh = Mesh(
+            points=torch.rand(60, 3, generator=g).to(device),
+            cells=torch.randperm(60, generator=g).reshape(20, 3).to(device),
+        )
+        sub = mesh.slice_cells(torch.arange(0, 20, 2, device=device))
+        scale_measures(sub, 2.0)
+
+        flipped = self._flip(sub)
+
+        assert torch.equal(cell_measures(flipped), cell_measures(sub))
+        assert torch.equal(flipped._cache["cell", "areas"], sub.cell_areas)
+
+    @pytest.mark.parametrize("order", [[0, 2, 1], [1, 2, 0]], ids=["flip", "cyclic"])
+    def test_reordering_a_zero_area_cell_keeps_its_measure(self, order, device):
+        mesh = Mesh(
+            points=torch.tensor(
+                [[0.0, 0.0, 0.0], [1.0, 0.0, 0.0], [2.0, 0.0, 0.0]], device=device
+            ),
+            cells=torch.tensor([[0, 1, 2]], device=device),
+        )
+        set_cell_measures(mesh, torch.tensor([1.5], device=device))
+
+        reordered = self._flip(mesh, order)
+
+        assert torch.equal(cell_measures(reordered), torch.tensor([1.5], device=device))
+
+    def test_winding_flip_of_slivers_keeps_measures_bitwise(self, device):
+        """Float32 areas of slivers depend on the vertex order; measures do not."""
+        g = torch.Generator().manual_seed(0)
+        base = torch.rand(200, 1, 3, generator=g)
+        direction = torch.rand(200, 1, 3, generator=g)
+        offsets = torch.tensor([0.0, 1.0, 2.0])[:, None] + 1e-4 * torch.rand(
+            200, 3, 1, generator=g
+        )
+        mesh = Mesh(
+            points=(base + offsets * direction).reshape(-1, 3).to(device),
+            cells=torch.arange(600, device=device).reshape(200, 3),
+        )
+        set_cell_measures(mesh, mesh.cell_areas * 3.0)
+        flipped_cells = mesh.cells[:, [0, 2, 1]]
+        assert not torch.equal(
+            Mesh(points=mesh.points, cells=flipped_cells).cell_areas, mesh.cell_areas
+        )
+
+        flipped = mesh.with_cells(flipped_cells)
+
+        assert torch.equal(cell_measures(flipped), cell_measures(mesh))
+        assert torch.equal(flipped._cache["cell", "areas"], mesh.cell_areas)
+
+    def test_changed_vertex_sets_still_follow_the_area_ratio(self, device):
+        mesh = Mesh(
+            points=torch.tensor(
+                [[0.0, 0.0, 0.0], [1.0, 0.0, 0.0], [0.0, 1.0, 0.0], [0.0, 2.0, 0.0]],
+                device=device,
+            ),
+            cells=torch.tensor([[0, 1, 2], [0, 2, 1]], device=device),
+        )
+        set_cell_measures(mesh, torch.tensor([3.0, 5.0], device=device))
+
+        updated = mesh.with_cells(torch.tensor([[0, 2, 1], [0, 3, 1]], device=device))
+
+        assert torch.equal(
+            cell_measures(updated), torch.tensor([3.0, 10.0], device=device)
+        )
+
+    def test_changed_vertex_set_of_a_zero_area_cell_raises(self, device):
+        mesh = Mesh(
+            points=torch.tensor(
+                [[0.0, 0.0, 0.0], [1.0, 0.0, 0.0], [2.0, 0.0, 0.0], [0.0, 1.0, 0.0]],
+                device=device,
+            ),
+            cells=torch.tensor([[0, 1, 2], [0, 1, 3]], device=device),
+        )
+        set_cell_measures(mesh, torch.tensor([1.0, 1.0], device=device))
+
+        with pytest.raises(ValueError, match="zero-measure cell"):
+            mesh.with_cells(torch.tensor([[0, 1, 3], [0, 3, 1]], device=device))

test/mesh/mesh/test_compile.py

@@ -47,7 +47,7 @@
 import torch
 
 from physicsnemo.mesh import Mesh
-from physicsnemo.mesh.calculus.measure import set_cell_measures
+from physicsnemo.mesh.calculus.measure import cell_measures, set_cell_measures
 from physicsnemo.mesh.primitives.surfaces import sphere_icosahedral
 
 ### Fixtures ###
@@ -264,6 +264,39 @@ def fn(p: torch.Tensor, c: torch.Tensor, replacement: torch.Tensor) -> torch.Ten
         compiled(points, cells, tetrahedra)
 
 
+def test_with_cells_keeps_measures_under_fullgraph_compile() -> None:
+    """A compiled winding flip keeps effective measures and ignores ``_cache``.
+
+    A cold and a warm input mesh share one graph, and the compiled call leaves
+    the input mesh's cache keys unchanged.
+    """
+    g = torch.Generator().manual_seed(0)
+    cold = Mesh(
+        points=torch.rand(600, 3, generator=g),
+        cells=torch.randperm(600, generator=g).reshape(200, 3),
+    )
+    # Float32 areas of these triangles depend on the vertex order
+    flipped = Mesh(points=cold.points, cells=cold.cells[:, [0, 2, 1]])
+    assert not torch.equal(flipped.cell_areas, cold.cell_areas)
+    warm = Mesh(points=cold.points * 2.0, cells=cold.cells)
+    for mesh in (cold, warm):
+        set_cell_measures(mesh, torch.full((mesh.n_cells,), 0.5))
+    _ = warm.cell_centroids, warm.cell_areas  # eager reads fill the cache
+
+    torch._dynamo.reset()
+    compiled = torch.compile(
+        lambda mesh: cell_measures(mesh.with_cells(mesh.cells[:, [0, 2, 1]])),
+        backend="eager",
+        fullgraph=True,
+        dynamic=True,
+    )
+    with torch._dynamo.config.patch(error_on_recompile=True):
+        for mesh in (cold, warm):
+            keys = set(mesh._cache.keys(True, True))
+            assert torch.equal(compiled(mesh), cell_measures(mesh))
+            assert set(mesh._cache.keys(True, True)) == keys
+
+
 @pytest.mark.parametrize("backend", ["eager", "inductor"])
 @pytest.mark.parametrize("dtype", [torch.int64, torch.uint64])
 def test_integer_coordinates_under_fullgraph_compile(backend, dtype):

test/mesh/repair/test_orientation.py

@@ -89,6 +89,28 @@ def test_closed_surfaces_in_3d(self, device):
         }
         torch.testing.assert_close(oriented.cell_normals, two_spheres.cell_normals)
 
+    def test_keeps_effective_measures_and_areas_bitwise(self, device):
+        """Flips keep each cell's (area, measure) pair, even on a degenerate cell."""
+        from physicsnemo.mesh.calculus.measure import cell_measures, set_cell_measures
+        from physicsnemo.mesh.primitives.surfaces import sphere_icosahedral
+
+        sphere = sphere_icosahedral.load(subdivisions=2, device=device)
+        # Jittered points make float32 areas depend on the vertex order
+        g = torch.Generator().manual_seed(0)
+        jitter = 0.01 * torch.rand(sphere.points.shape, generator=g).to(device)
+        degenerate = torch.tensor([[0, 0, 1]], device=device)
+        sphere = Mesh(
+            points=sphere.points + jitter, cells=torch.cat([sphere.cells, degenerate])
+        )
+        scrambled = _reverse(sphere, _every_third(sphere.n_cells, device))
+        set_cell_measures(scrambled, torch.full((sphere.n_cells,), 0.5, device=device))
+
+        oriented, stats = fix_orientation(scrambled)
+
+        assert stats["n_cells_flipped"] > 0
+        assert torch.equal(cell_measures(oriented), cell_measures(scrambled))
+        assert torch.equal(oriented.cell_areas, scrambled.cell_areas)
+
     def test_sharp_creases(self, device):
         """Faces meeting at right angles are oriented by connectivity, not normals."""
         from physicsnemo.mesh.primitives.surfaces import cube_surface