From b0e979a7411e08d550434730816d65a9549a156c Mon Sep 17 00:00:00 2001 From: Peter Sharpe Date: Fri, 13 Nov 2026 11:48:51 -0500 Subject: [PATCH 1/4] Keep measures of reordered cells in Mesh.with_cells A row that only reorders its own vertices, such as a winding flip, is the same simplex: it now keeps its effective measure, cached area, and cached centroid bitwise, instead of being rescaled by float32 roundoff in the area ratio or raising on a zero-area cell with a nonzero measure. Rows whose vertex set changes keep the area-ratio transfer. Signed-off-by: Peter Sharpe --- CHANGELOG.md | 3 ++ physicsnemo/mesh/calculus/measure.py | 30 ++++++++--- physicsnemo/mesh/mesh.py | 25 ++++++++- test/mesh/calculus/test_measure.py | 78 ++++++++++++++++++++++++++++ test/mesh/repair/test_orientation.py | 22 ++++++++ test/mesh/utilities/test_cache.py | 27 ++++++++++ 6 files changed, 177 insertions(+), 8 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index c325aeefbb..9a6c6ddadc 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -555,6 +555,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. diff --git a/physicsnemo/mesh/calculus/measure.py b/physicsnemo/mesh/calculus/measure.py index 51e64f64a7..0adeae5e54 100644 --- a/physicsnemo/mesh/calculus/measure.py +++ b/physicsnemo/mesh/calculus/measure.py @@ -201,16 +201,31 @@ 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 rows marked by the boolean mask ``unchanged`` are their source rows + with vertices reordered: they keep the source's measure and cached area + bitwise, so later transfers see the same pair. + """ 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 + 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" ) @@ -218,11 +233,14 @@ def _transfer_cell_measures( # 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: diff --git a/physicsnemo/mesh/mesh.py b/physicsnemo/mesh/mesh.py index e6beec8b90..f7061dcc83 100644 --- a/physicsnemo/mesh/mesh.py +++ b/physicsnemo/mesh/mesh.py @@ -2100,6 +2100,10 @@ def with_cells( The replacement must have exactly the same shape as the current cells. User data is preserved, while all 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. Outside tracing and autograd it + also keeps its cached centroid and, when the mesh has effective + measures, its cached area. Parameters ---------- @@ -2155,9 +2159,26 @@ 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: + sorted_cells = 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( diff --git a/test/mesh/calculus/test_measure.py b/test/mesh/calculus/test_measure.py index 5726a2cfe2..8f7c9e5015 100644 --- a/test/mesh/calculus/test_measure.py +++ b/test/mesh/calculus/test_measure.py @@ -246,3 +246,81 @@ 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): + """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), + cells=torch.randperm(60, generator=g).reshape(20, 3), + ) + sub = mesh.slice_cells(torch.arange(0, 20, 2)) + 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): + mesh = Mesh( + points=torch.tensor([[0.0, 0.0, 0.0], [1.0, 0.0, 0.0], [2.0, 0.0, 0.0]]), + cells=torch.tensor([[0, 1, 2]]), + ) + set_cell_measures(mesh, torch.tensor([1.5])) + + reordered = self._flip(mesh, order) + + assert torch.equal(cell_measures(reordered), torch.tensor([1.5])) + + def test_winding_flip_of_slivers_keeps_measures_bitwise(self): + """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), + cells=torch.arange(600).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)) + + def test_changed_vertex_sets_still_follow_the_area_ratio(self): + 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]] + ), + cells=torch.tensor([[0, 1, 2], [0, 2, 1]]), + ) + set_cell_measures(mesh, torch.tensor([3.0, 5.0])) + + updated = mesh.with_cells(torch.tensor([[0, 2, 1], [0, 3, 1]])) + + assert torch.equal(cell_measures(updated), torch.tensor([3.0, 10.0])) + + def test_changed_vertex_set_of_a_zero_area_cell_raises(self): + 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]] + ), + cells=torch.tensor([[0, 1, 2], [0, 1, 3]]), + ) + set_cell_measures(mesh, torch.tensor([1.0, 1.0])) + + with pytest.raises(ValueError, match="zero-measure cell"): + mesh.with_cells(torch.tensor([[0, 1, 3], [0, 3, 1]])) diff --git a/test/mesh/repair/test_orientation.py b/test/mesh/repair/test_orientation.py index b9f90de832..6ae6c54605 100644 --- a/test/mesh/repair/test_orientation.py +++ b/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 diff --git a/test/mesh/utilities/test_cache.py b/test/mesh/utilities/test_cache.py index a64ab497b0..843beff9e8 100644 --- a/test/mesh/utilities/test_cache.py +++ b/test/mesh/utilities/test_cache.py @@ -465,6 +465,33 @@ def test_rejects_floating_point_cells(self): with pytest.raises(TypeError, match="integer dtype"): mesh.with_cells(mesh.cells.to(torch.float32)) + def test_reordered_rows_keep_cached_centroids(self): + mesh = Mesh( + points=torch.rand(4, 3, generator=torch.Generator().manual_seed(0)), + cells=torch.tensor([[0, 1, 2], [0, 1, 3]]), + ) + centroids = mesh.cell_centroids + + updated = mesh.with_cells(torch.tensor([[1, 2, 0], [0, 2, 3]])) + + cached = updated._cache["cell", "centroids"] + assert torch.equal(cached[0], centroids[0]) + torch.testing.assert_close(cached[1], mesh.points[[0, 2, 3]].mean(dim=0)) + + def test_reordering_writes_no_cache_when_points_require_grad(self): + from physicsnemo.mesh.calculus.measure import set_cell_measures + + mesh = Mesh( + points=torch.rand(3, 3).requires_grad_(), + cells=torch.tensor([[0, 1, 2]]), + ) + set_cell_measures(mesh, torch.ones(1)) + mesh._cache["cell", "centroids"] = torch.zeros(1, 3) + + updated = mesh.with_cells(mesh.cells[:, [0, 2, 1]]) + + assert not updated._cache["cell"].keys() + def test_result_containers_are_independent(self): mesh = self._cached_triangle() From a73e716968b201310f74e942654fecfe4c210b70 Mon Sep 17 00:00:00 2001 From: Peter Sharpe Date: Fri, 13 Nov 2026 12:24:56 -0500 Subject: [PATCH 2/4] Accept unsigned replacement cells in Mesh.with_cells The reordered-row check compared the raw replacement with the mesh's int64 cells, and PyTorch cannot promote uint16, uint32 or uint64 with int64. Compare result.cells, which the constructor has normalized. Signed-off-by: Peter Sharpe --- physicsnemo/mesh/mesh.py | 4 +++- test/mesh/utilities/test_cache.py | 17 +++++++++++++++++ 2 files changed, 20 insertions(+), 1 deletion(-) diff --git a/physicsnemo/mesh/mesh.py b/physicsnemo/mesh/mesh.py index f7061dcc83..6169aa99b2 100644 --- a/physicsnemo/mesh/mesh.py +++ b/physicsnemo/mesh/mesh.py @@ -2172,7 +2172,9 @@ def with_cells( ) unchanged = None if centroids is not None or EFFECTIVE_MEASURE_KEY in self.cell_data: - sorted_cells = cells.sort(dim=1).values + # Use the normalized cells: PyTorch cannot compare uint16, uint32 + # or uint64 with int64. + 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( diff --git a/test/mesh/utilities/test_cache.py b/test/mesh/utilities/test_cache.py index 843beff9e8..1cb5a5d4ea 100644 --- a/test/mesh/utilities/test_cache.py +++ b/test/mesh/utilities/test_cache.py @@ -478,6 +478,23 @@ def test_reordered_rows_keep_cached_centroids(self): assert torch.equal(cached[0], centroids[0]) torch.testing.assert_close(cached[1], mesh.points[[0, 2, 3]].mean(dim=0)) + @pytest.mark.parametrize("dtype", [torch.uint16, torch.uint32, torch.uint64]) + def test_reordering_unsigned_cells_keeps_measures_and_centroids(self, dtype): + from physicsnemo.mesh.calculus.measure import cell_measures, set_cell_measures + + mesh = Mesh( + points=torch.rand(3, 3, generator=torch.Generator().manual_seed(0)), + cells=torch.tensor([[0, 1, 2]]), + ) + set_cell_measures(mesh, torch.tensor([1.5])) + centroids = mesh.cell_centroids + + updated = mesh.with_cells(mesh.cells[:, [0, 2, 1]].to(dtype)) + + assert updated.cells.dtype == torch.int64 + assert torch.equal(cell_measures(updated), torch.tensor([1.5])) + assert torch.equal(updated._cache["cell", "centroids"], centroids) + def test_reordering_writes_no_cache_when_points_require_grad(self): from physicsnemo.mesh.calculus.measure import set_cell_measures From 1e7f581307c11d292ec705a8300975c08334ebab Mon Sep 17 00:00:00 2001 From: Peter Sharpe Date: Fri, 13 Nov 2026 12:27:00 -0500 Subject: [PATCH 3/4] Make the dtype comment in Mesh.with_cells exact Signed-off-by: Peter Sharpe --- physicsnemo/mesh/mesh.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/physicsnemo/mesh/mesh.py b/physicsnemo/mesh/mesh.py index 6169aa99b2..dc562cf02f 100644 --- a/physicsnemo/mesh/mesh.py +++ b/physicsnemo/mesh/mesh.py @@ -2172,8 +2172,8 @@ def with_cells( ) unchanged = None if centroids is not None or EFFECTIVE_MEASURE_KEY in self.cell_data: - # Use the normalized cells: PyTorch cannot compare uint16, uint32 - # or uint64 with int64. + # 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: From 2f72cd5ab9edda0e4e27c83893e15f958e982a2b Mon Sep 17 00:00:00 2001 From: Peter Sharpe Date: Fri, 13 Nov 2026 13:03:30 -0500 Subject: [PATCH 4/4] Document the caches Mesh.with_cells keeps, and widen its tests Returns and Notes now say which caches the result holds without keep, and that they are computed eagerly for all cells. _transfer_cell_measures documents how unchanged indexes rows when parents is given, and says why it caches the old area bits. New tests cover a compiled flip, an empty mesh, keep=("cell", "areas"), and CUDA for the measure tests; the centroid tests use a sentinel and the sliver test checks the cached area, so each fails if its value is recomputed. Signed-off-by: Peter Sharpe --- physicsnemo/mesh/calculus/measure.py | 9 +++-- physicsnemo/mesh/mesh.py | 16 ++++++--- test/mesh/calculus/test_measure.py | 53 ++++++++++++++++------------ test/mesh/mesh/test_compile.py | 35 +++++++++++++++++- test/mesh/utilities/test_cache.py | 42 +++++++++++++++++++--- 5 files changed, 119 insertions(+), 36 deletions(-) diff --git a/physicsnemo/mesh/calculus/measure.py b/physicsnemo/mesh/calculus/measure.py index 0adeae5e54..fbd5d2fc70 100644 --- a/physicsnemo/mesh/calculus/measure.py +++ b/physicsnemo/mesh/calculus/measure.py @@ -208,9 +208,10 @@ def _transfer_cell_measures( ) -> None: """Transfer explicit cell measures through geometric changes/subdivision. - Result rows marked by the boolean mask ``unchanged`` are their source rows - with vertices reordered: they keep the source's measure and cached area - bitwise, so later transfers see the same pair. + 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 @@ -223,6 +224,8 @@ def _transfer_cell_measures( 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()): diff --git a/physicsnemo/mesh/mesh.py b/physicsnemo/mesh/mesh.py index dc562cf02f..fa4bea3613 100644 --- a/physicsnemo/mesh/mesh.py +++ b/physicsnemo/mesh/mesh.py @@ -2098,12 +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. Outside tracing and autograd it - also keeps its cached centroid and, when the mesh has effective - measures, its cached area. + keeps its effective measure bitwise and, as the Notes describe, its + cached centroid and area. Parameters ---------- @@ -2117,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 ------ @@ -2132,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 diff --git a/test/mesh/calculus/test_measure.py b/test/mesh/calculus/test_measure.py index 8f7c9e5015..fdcfdc512f 100644 --- a/test/mesh/calculus/test_measure.py +++ b/test/mesh/calculus/test_measure.py @@ -251,14 +251,14 @@ def test_weights_invariant_under_rigid_and_scaling_transforms(self): 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): + 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), - cells=torch.randperm(60, generator=g).reshape(20, 3), + 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)) + sub = mesh.slice_cells(torch.arange(0, 20, 2, device=device)) scale_measures(sub, 2.0) flipped = self._flip(sub) @@ -267,18 +267,20 @@ def test_winding_flip_keeps_subsampled_measures_and_areas_bitwise(self): 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): + 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]]), - cells=torch.tensor([[0, 1, 2]]), + 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])) + 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])) + assert torch.equal(cell_measures(reordered), torch.tensor([1.5], device=device)) - def test_winding_flip_of_slivers_keeps_measures_bitwise(self): + 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) @@ -287,8 +289,8 @@ def test_winding_flip_of_slivers_keeps_measures_bitwise(self): 200, 3, 1, generator=g ) mesh = Mesh( - points=(base + offsets * direction).reshape(-1, 3), - cells=torch.arange(600).reshape(200, 3), + 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]] @@ -299,28 +301,33 @@ def test_winding_flip_of_slivers_keeps_measures_bitwise(self): 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): + 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]] + [[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]]), + cells=torch.tensor([[0, 1, 2], [0, 2, 1]], device=device), ) - set_cell_measures(mesh, torch.tensor([3.0, 5.0])) + set_cell_measures(mesh, torch.tensor([3.0, 5.0], device=device)) - updated = mesh.with_cells(torch.tensor([[0, 2, 1], [0, 3, 1]])) + 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])) + 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): + 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]] + [[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]]), + cells=torch.tensor([[0, 1, 2], [0, 1, 3]], device=device), ) - set_cell_measures(mesh, torch.tensor([1.0, 1.0])) + 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]])) + mesh.with_cells(torch.tensor([[0, 1, 3], [0, 3, 1]], device=device)) diff --git a/test/mesh/mesh/test_compile.py b/test/mesh/mesh/test_compile.py index 144788cc27..0d7a64ab5c 100644 --- a/test/mesh/mesh/test_compile.py +++ b/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): diff --git a/test/mesh/utilities/test_cache.py b/test/mesh/utilities/test_cache.py index 1cb5a5d4ea..7d0cd38ece 100644 --- a/test/mesh/utilities/test_cache.py +++ b/test/mesh/utilities/test_cache.py @@ -470,12 +470,13 @@ def test_reordered_rows_keep_cached_centroids(self): points=torch.rand(4, 3, generator=torch.Generator().manual_seed(0)), cells=torch.tensor([[0, 1, 2], [0, 1, 3]]), ) - centroids = mesh.cell_centroids + # A sentinel, which no recompute can match + mesh._cache["cell", "centroids"] = torch.full((2, 3), 7.0) updated = mesh.with_cells(torch.tensor([[1, 2, 0], [0, 2, 3]])) cached = updated._cache["cell", "centroids"] - assert torch.equal(cached[0], centroids[0]) + assert torch.equal(cached[0], torch.full((3,), 7.0)) torch.testing.assert_close(cached[1], mesh.points[[0, 2, 3]].mean(dim=0)) @pytest.mark.parametrize("dtype", [torch.uint16, torch.uint32, torch.uint64]) @@ -487,13 +488,46 @@ def test_reordering_unsigned_cells_keeps_measures_and_centroids(self, dtype): cells=torch.tensor([[0, 1, 2]]), ) set_cell_measures(mesh, torch.tensor([1.5])) - centroids = mesh.cell_centroids + sentinel = torch.full((1, 3), 7.0) + mesh._cache["cell", "centroids"] = sentinel updated = mesh.with_cells(mesh.cells[:, [0, 2, 1]].to(dtype)) assert updated.cells.dtype == torch.int64 assert torch.equal(cell_measures(updated), torch.tensor([1.5])) - assert torch.equal(updated._cache["cell", "centroids"], centroids) + assert torch.equal(updated._cache["cell", "centroids"], sentinel) + + def test_reordering_an_empty_mesh(self): + from physicsnemo.mesh.calculus.measure import cell_measures, set_cell_measures + + mesh = Mesh(points=torch.rand(3, 3), cells=torch.empty(0, 3, dtype=torch.int64)) + set_cell_measures(mesh, torch.empty(0)) + _ = mesh.cell_centroids + + updated = mesh.with_cells(mesh.cells[:, [0, 2, 1]]) + + assert cell_measures(updated).shape == (0,) + assert updated._cache["cell", "centroids"].shape == (0, 3) + + def test_keep_retains_the_area_cache_through_the_measure_transfer(self): + """A retained area cache is used as is, even for a row whose vertices change.""" + from physicsnemo.mesh.calculus.measure import cell_measures, set_cell_measures + + 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]] + ), + cells=torch.tensor([[0, 1, 2], [0, 2, 1]]), + ) + set_cell_measures(mesh, torch.tensor([3.0, 5.0])) + areas = mesh.cell_areas + + updated = mesh.with_cells( + torch.tensor([[0, 2, 1], [0, 3, 1]]), keep=("cell", "areas") + ) + + assert torch.equal(updated._cache["cell", "areas"], areas) + assert torch.equal(cell_measures(updated), torch.tensor([3.0, 5.0])) def test_reordering_writes_no_cache_when_points_require_grad(self): from physicsnemo.mesh.calculus.measure import set_cell_measures