PhysicsNeMo-Mesh: improve cache handling APIs - #1867
Conversation
70e0961 to
6c198f7
Compare
peterdsharpe
changed the title
Greptile SummaryThis PR introduces
Important Files Changed
Reviews (1): Last reviewed commit: "Add cache-aware Mesh cell replacement" | Re-trigger Greptile |
![greptile-apps[bot]](https://avatars.githubusercontent.com/in/867647?s=60&v=4){kind=small}
| cache = self._cache.select(*keep, strict=False).copy() | ||
| device = self.points.device | ||
| for category, batch_size in ( | ||
| ("cell", torch.Size([self.n_cells])), | ||
| ("point", torch.Size([self.n_points])), | ||
| ("topology", torch.Size([])), | ||
| ): | ||
| if category not in cache: | ||
| cache[category] = TensorDict( | ||
| {}, | ||
| batch_size=batch_size, | ||
| device=device, | ||
| ) | ||
| return cache |
There was a problem hiding this comment.
keep silently produces wrong results
_cache_with_only unpacks keep directly with *keep, so two easy mistakes produce silent wrong behavior rather than an error:
-
A bare string —
keep="topology"— iterates as characters, callingselect("t", "o", "p", …, strict=False). None of those keys exist, so the result is an empty cache (all caches stripped) instead of the topology cache being retained. -
A bare nested-key tuple —
keep=("cell", "areas")— unpacks as two top-level selects:select("cell", "areas", …), retaining the entire cell sub-TensorDict plus a silently-ignored missing"areas"top-level key, rather than retaining only("cell", "areas").
Both cases are accepted without error because str and tuple both satisfy Sequence[str | tuple]. The upstream public methods strip_caches, with_points, and with_cells all expose the same parameter, so any caller who passes keep="topology" or keep=("cell", "areas") instead of keep=["topology"] or keep=[("cell", "areas")] gets silently incorrect cache state. A short guard at the top of _cache_with_only (e.g. if isinstance(keep, (str, bytes)): raise TypeError(...)) would catch the string case.
| """Return an independent cache container containing only ``keep``. | ||
|
|
||
| Tensor leaves are intentionally shared, but every retained nested | ||
| ``TensorDict`` container is shallow-copied so populating a cache on a | ||
| derived mesh cannot mutate the source mesh's cache structure. | ||
| """ | ||
| cache = self._cache.select(*keep, strict=False).copy() |
There was a problem hiding this comment.
keep parameter from silently stripping all caches instead of retaining the named category.
| """Return an independent cache container containing only ``keep``. | |
| Tensor leaves are intentionally shared, but every retained nested | |
| ``TensorDict`` container is shallow-copied so populating a cache on a | |
| derived mesh cannot mutate the source mesh's cache structure. | |
| """ | |
| cache = self._cache.select(*keep, strict=False).copy() | |
| """Return an independent cache container containing only ``keep``. | |
| Tensor leaves are intentionally shared, but every retained nested | |
| ``TensorDict`` container is shallow-copied so populating a cache on a | |
| derived mesh cannot mutate the source mesh's cache structure. | |
| """ | |
| if isinstance(keep, (str, bytes)): | |
| raise TypeError( | |
| f"`keep` must be a sequence of strings or nested-key tuples, " | |
| f"not a bare string {keep!r}. " | |
| f"Did you mean keep=[{keep!r}]?" | |
| ) | |
| cache = self._cache.select(*keep, strict=False).copy() |
Note: If this suggestion doesn't match your team's coding style, reply to this and let me know. I'll remember it for next time!
| if (v := mesh._cache.get(("cell", "centroids"), None)) is not None: | ||
| new_cache["cell", "centroids"] = v + offset | ||
| translated_mesh._cache["cell", "centroids"] = v + offset | ||
|
|
||
| from physicsnemo.mesh.mesh import Mesh | ||
|
|
||
| return Mesh( | ||
| points=new_points, | ||
| cells=mesh.cells, | ||
| point_data=mesh.point_data, | ||
| cell_data=mesh.cell_data, | ||
| global_data=mesh.global_data, | ||
| _cache=new_cache, | ||
| ) | ||
| return translated_mesh | ||
|
|
||
|
|
||
| def rotate( |
There was a problem hiding this comment.
with_points bypasses immutability convention
translated_mesh._cache["cell", "centroids"] = v + offset directly mutates the _cache TensorDict container of an already-constructed Mesh. While safe here (the container is independent after _new_with_structure), the cache state visible inside translated_mesh differs between the moment with_points returns and the moment translate returns. Any intermediate inspection would see stale centroid values. A cleaner pattern would be to build the corrected cache before calling _new_with_structure, keeping the public API the only write path to the cache.
| ) | ||
| return cleaned | ||
|
|
||
| def strip_caches(self) -> "Mesh": | ||
| r"""Return a new mesh with all cached values removed. | ||
| def strip_caches( | ||
| self, | ||
| keep: Sequence[str | tuple[str, ...]] = (), | ||
| ) -> "Mesh": | ||
| r"""Return a new mesh with cached values removed. | ||
|
|
||
| Cached values (stored under the ``_cache`` key in data TensorDicts) are | ||
| computed lazily for expensive operations like normals, areas, and curvature. | ||
| This method creates a new mesh without these cached values, which is useful | ||
| for: | ||
| Cached values stored in the separate :attr:`_cache` field are computed | ||
| lazily for expensive operations like normals, areas, and curvature. This | ||
| method creates a new mesh without these cached values, except for keys | ||
| explicitly listed in ``keep``. This is useful for: | ||
|
|
||
| - Accurate benchmarking (prevents false performance benefits from caching) | ||
| - Reducing memory usage | ||
| - Forcing recomputation of cached values | ||
|
|
||
| Parameters | ||
| ---------- | ||
| keep : sequence of str or tuple[str, ...], optional | ||
| Cache keys to retain. Strings select a complete top-level cache such as | ||
| ``"topology"``; tuples select a nested entry such as | ||
| ``("cell", "areas")``. Missing keys are ignored. | ||
|
|
||
| Returns | ||
| ------- | ||
| Mesh | ||
| A new mesh with the same geometry and data, but without cached values. | ||
| A new mesh with the same geometry and data, retaining only the requested | ||
| cached values. | ||
|
|
||
| Examples | ||
| -------- | ||
| >>> from physicsnemo.mesh.primitives.surfaces import sphere_icosahedral | ||
| >>> mesh = sphere_icosahedral.load(subdivisions=2) | ||
| >>> _ = mesh.cell_normals # Triggers caching | ||
| >>> mesh_clean = mesh.strip_caches() # Remove cached normals | ||
| >>> mesh_with_areas = mesh.strip_caches(keep=[("cell", "areas")]) | ||
| """ | ||
| return Mesh( | ||
| return self._new_with_structure( | ||
| points=self.points, | ||
| cells=self.cells, | ||
| point_data=self.point_data, | ||
| cell_data=self.cell_data, | ||
| global_data=self.global_data, | ||
| keep=keep, | ||
| ) | ||
|
|
||
|
|
There was a problem hiding this comment.
strip_caches() not accurately described
The PR description states "the existing no-argument strip_caches() behavior remains unchanged," but there is a subtle difference. The old implementation passed point_data=self.point_data, cell_data=self.cell_data, global_data=self.global_data directly to the Mesh constructor, which stored those TensorDicts as-is (shared containers). The new path through _new_with_structure calls .copy() on each, producing independent containers. Code that relied on the aliasing would silently behave differently. The change is an improvement in correctness, but the PR description should note it rather than claiming no-arg parity.
1 participant
PhysicsNeMo Pull Request
Description
Improve the Mesh cache-handling API and route repository consumers through it so cache validity is explicit, reusable, and consistent.
Mesh.strip_caches(keep=...)andDomainMesh.strip_caches(keep=...)to retain selected cache entries or complete cache categories.Mesh.with_points(...)for coordinate changes that preserve point indexing and connectivity. It invalidates geometry-dependent caches while retaining topology caches by default.Mesh.with_cells(...)for connectivity changes that preserve cell indexing and simplex type. It clears cell, point, and topology caches by default, with an explicit expertkeep=override.with_dataso valid caches survive.The private structural helper is named
_new_with_structureto avoid TensorClass generatedreplaceand dataclass__replace__APIs. Those generic replacements are documented as unsafe for changingpointsorcellsbecause they do not invalidate caches.The existing no-argument
strip_caches()behavior remains unchanged. Part of #1851.Validation
Additional checks cover full-graph
torch.compileuse ofwith_cellsand DrivAer ground-boundary winding/normals. The GLOBE model file retains its existing symbolic-shapeF821suppressions; all newly changed lines pass Ruff.Checklist
Dependencies
None.
Review Process
Draft; no reviewers requested yet.