Mesh#

The Mesh class is the central data structure of PhysicsNeMo-Mesh. It is a tensorclass built on TensorDict, representing an n-dimensional simplicial manifold embedded in m-dimensional Euclidean space.

A Mesh stores vertex coordinates (points), cell connectivity (cells), and three TensorDict containers for attaching arbitrary tensor data at the vertex, cell, and global levels. All tensors move together under .to(device) calls, and expensive geometric quantities – centroids, normals, areas, curvature – are computed lazily on first access and cached internally.

Most mesh operations (subdivision, derivatives, transformations) are available both as Mesh methods and as standalone functions in the corresponding submodules. Each pair shares one canonical function, and normal Python descriptor binding supplies the instance as the standalone function’s mesh argument.

To construct a triangle mesh from a surface mesh whose cells are arbitrary polygons – a “polygon soup” (see Tessellation) – use Mesh.from_polygons().

import torch
from physicsnemo.mesh import Mesh

points = torch.tensor([[0.0, 0.0], [1.0, 0.0], [0.5, 1.0]])
cells = torch.tensor([[0, 1, 2]])
mesh = Mesh(points=points, cells=cells)

# Geometric properties (lazily computed, cached)
print(mesh.cell_centroids)   # shape (1, 2)
print(mesh.cell_areas)       # shape (1,)

# Attach data and compute derivatives
mesh.point_data["T"] = torch.tensor([1.0, 2.0, 3.0])
mesh = mesh.compute_point_derivatives(keys="T", method="lsq")
print(mesh.point_data["T_gradient"])  # shape (3, 2)

class physicsnemo.mesh.mesh.Mesh( points: torch.Tensor, cells: torch.Tensor | None = None, point_data: tensordict._td.TensorDict | dict[str, torch.Tensor] | None = None, cell_data: tensordict._td.TensorDict | dict[str, torch.Tensor] | None = None, global_data: tensordict._td.TensorDict | dict[str, torch.Tensor] | None = None, *, _cache: tensordict._td.TensorDict | None = None, batch_size, device=None, names=None, )[source]#

Bases: object

property cell_areas: Tensor#

Compute volumes (areas) of n-simplices.

This works for simplices of any manifold dimension embedded in any spatial dimension. For example: edges in 2D/3D, triangles in 2D/3D/4D, tetrahedra in 3D/4D, etc.

Uses dimension-specific closed-form expressions for n <= 3 (Lagrange identity, scalar triple product, etc.) and falls back to the Gram determinant for higher dimensions. See compute_cell_areas() for details.

cell_areas is always the purely geometric simplex measure. For meshes whose cells represent more than their own geometry (e.g. after cell subsampling), the effective integration measure is provided by physicsnemo.mesh.calculus.measure.

Returns:: Tensor of shape (n_cells,) containing the volume of each cell.
Return type:: torch.Tensor

property cell_centroids: Tensor#

Compute the centroids (geometric centers) of all cells.

The centroid of a cell is computed as the arithmetic mean of its vertex positions. For an n-simplex with vertices (v0, v1, …, vn), the centroid is centroid = (v0 + v1 + ... + vn) / (n + 1).

The result is cached in _cache["cell", "centroids"] for efficiency.

Returns:: Tensor of shape (n_cells, n_spatial_dims) containing the centroid of each cell.
Return type:: torch.Tensor

cell_data_to_point_data( overwrite_keys: bool = False, ) → Mesh[source]#

Convert cell data to point data by averaging.

For each point, computes the average of the cell data values from all cells that contain that point. The resulting point data is added to the mesh’s point_data dictionary. Original cell data is preserved.

Parameters:: overwrite_keys (bool) – If True, silently overwrite any existing point_data keys. If False, raise an error if a key already exists in point_data.
Returns:: New Mesh with converted data added to point_data. Original cell_data is preserved.
Return type:: Mesh
Raises:: ValueError – If a cell_data key already exists in point_data and overwrite_keys=False.

Notes

Cell fields are averaged in floating point, so an integer or boolean cell field is returned as a torch.float64 point field (the per-point mean of integers is generally non-integral and is not truncated). See scatter_aggregate for the underlying dtype-promotion rule.

Examples

>>> mesh = Mesh(points, cells, cell_data={"pressure": cell_pressures})
>>> mesh_with_point_data = mesh.cell_data_to_point_data()
>>> # Now mesh has both cell_data["pressure"] and point_data["pressure"]

property cell_normals: Tensor#

Compute unit normal vectors for codimension-1 cells.

Normal vectors are uniquely defined (up to orientation) only for codimension-1 manifolds, where n_manifold_dims = n_spatial_dims - 1.

Uses dimension-specific closed-form expressions for d=2 (rotation) and d=3 (cross product), falling back to signed minor determinants for higher dimensions. See compute_cell_normals() for details.

Returns:: Tensor of shape (n_cells, n_spatial_dims) containing unit normal vectors.
Return type:: torch.Tensor
Raises:: ValueError – If the mesh is not codimension-1 (n_manifold_dims ≠ n_spatial_dims - 1).

clean( tolerance: float = 1e-12, merge_points: bool = True, remove_duplicate_cells: bool = True, remove_unused_points: bool = True, ) → Mesh[source]#

Clean and repair this mesh.

Performs up to three cleaning operations in sequence:

Merge duplicate points (merge_points): Finds points within tolerance L2 distance using BVH spatial queries and merges them into a single representative. Point data values are averaged across merged groups. Cost: \(O(N \log N)\) where \(N\) is the number of points. This is the most expensive step - on meshes with millions of points it can take tens of seconds.
Remove duplicate cells (remove_duplicate_cells): Sorts vertex indices within each cell and removes cells that share the same vertex set. Cost: \(O(C \log C)\) where \(C\) is the number of cells. Typically fast.
Remove unused points (remove_unused_points): Drops points not referenced by any cell and compacts the point array. Cost: \(O(N + C \cdot V)\) where \(V\) is vertices per cell. Very fast (linear scatter + mask).

This is useful after importing meshes from external sources (VTK, STL, CAD) that may have redundant geometry. For programmatic mesh operations like slice_cells that don’t create duplicates, you can disable the expensive steps and only keep remove_unused_points=True for a large speedup.

Parameters:

tolerance (float, optional) – Absolute L2 distance threshold for merging duplicate points.
merge_points (bool, optional) – Whether to merge spatially-duplicate points (default True).
remove_duplicate_cells (bool, optional) – Whether to remove cells with identical vertex sets (default True).
remove_unused_points (bool, optional) – Whether to drop points not referenced by any cell (default True).

Returns:

Cleaned mesh with same structure but repaired topology.

Return type:

Mesh

Examples

>>> import torch
>>> from physicsnemo.mesh import Mesh
>>> # Mesh with duplicate points
>>> points = torch.tensor([[0., 0.], [1., 0.], [0., 0.], [1., 1.]])
>>> cells = torch.tensor([[0, 1, 3], [2, 1, 3]])
>>> mesh = Mesh(points=points, cells=cells)
>>> cleaned = mesh.clean()
>>> assert cleaned.n_points == 3  # points 0 and 2 merged
>>>
>>> # Fast path: only remove unreferenced points (after slice_cells, etc.)
>>> subset = mesh.slice_cells(torch.tensor([0]))
>>> compacted = subset.clean(
...     merge_points=False,
...     remove_duplicate_cells=False,
...     remove_unused_points=True,
... )

property codimension: int#

Compute the codimension of the mesh.

The codimension is the difference between the spatial dimension and the manifold dimension: codimension = n_spatial_dims - n_manifold_dims.

Returns:: The codimension of the mesh (always non-negative).
Return type:: int

Notes

Edges (1-simplices) in 2D: codimension = 2 - 1 = 1 (codimension-1)
Triangles (2-simplices) in 3D: codimension = 3 - 2 = 1 (codimension-1)
Edges in 3D: codimension = 3 - 1 = 2 (codimension-2)
Points in 2D: codimension = 2 - 0 = 2 (codimension-2)

compute_cell_derivatives( keys: str | tuple[str, ...] | Sequence[str | tuple[str, ...]] | None = None, method: Literal['lsq', 'dec'] = 'lsq', gradient_type: Literal['intrinsic', 'extrinsic', 'both'] = 'intrinsic', ) → Mesh#

Compute gradients of cell_data fields.

Call it as compute_cell_derivatives(mesh, ...) or as mesh.compute_cell_derivatives(...). The bound method supplies mesh automatically.

Parameters:

mesh (Mesh) – Simplicial mesh with cell_data fields to differentiate.
keys (str or tuple[str, ...] or Sequence or None) – Fields to compute gradients of (same format as compute_point_derivatives()).
method ({"lsq"}) – Discretization method for cell-centered data. Currently only "lsq" (weighted least-squares) is implemented. DEC gradients for cell-centered data are not available because the standard DEC exterior derivative maps vertex 0-forms to edge 1-forms; there is no analogous cell-to-cell operator in the primal DEC complex.
gradient_type ({"intrinsic", "extrinsic", "both"}) – Type of gradient to compute.

Returns:

A new Mesh with gradient fields added to cell_data. The original mesh is not modified.

Return type:

Mesh

Raises:

NotImplementedError – If method="dec" is requested.

compute_point_derivatives( keys: str | tuple[str, ...] | Sequence[str | tuple[str, ...]] | None = None, method: Literal['lsq', 'dec'] = 'lsq', gradient_type: Literal['intrinsic', 'extrinsic', 'both'] = 'intrinsic', ) → Mesh#

Compute gradients of point_data fields.

Computes discrete gradients using either DEC or LSQ methods, with support for both intrinsic (tangent space) and extrinsic (ambient space) derivatives.

Call it as compute_point_derivatives(mesh, ...) or as mesh.compute_point_derivatives(...). The bound method supplies mesh automatically.

Parameters:

mesh (Mesh) – Simplicial mesh with point_data fields to differentiate.
keys (str or tuple[str, ...] or Sequence or None) –
Fields to compute gradients of. Options:
- None: All non-cached fields (excludes "_cache" subdictionary).
- str: Single field name (e.g., "pressure").
- tuple: Nested path (e.g., ("flow", "temperature")).
- Sequence: List of the above.
method ({"lsq", "dec"}) –
Discretization method:
- "lsq": Weighted least-squares reconstruction (CFD standard).
- "dec": Discrete Exterior Calculus (differential geometry).
gradient_type ({"intrinsic", "extrinsic", "both"}) –
Type of gradient to compute:
- "intrinsic": Project onto manifold tangent space.
- "extrinsic": Full ambient space gradient.
- "both": Compute and store both.

Returns:

A new Mesh with gradient fields added to point_data. The original mesh is not modified. Field naming convention:

gradient_type="intrinsic" or "extrinsic": "{field}_gradient"
gradient_type="both": "{field}_gradient_intrinsic" and "{field}_gradient_extrinsic"

Return type:

Mesh

Example

>>> import torch
>>> from physicsnemo.mesh.primitives.basic import two_triangles_2d
>>> mesh = two_triangles_2d.load()
>>> mesh.point_data["pressure"] = torch.randn(mesh.n_points)
>>> mesh_with_grad = compute_point_derivatives(mesh, keys="pressure")
>>> grad_p = mesh_with_grad.point_data["pressure_gradient"]

compute_point_normals( weighting: Literal['area', 'unweighted', 'angle', 'angle_area'] = 'angle_area', ) → Tensor[source]#

Compute normal vectors at mesh vertices with specified weighting.

For each point (vertex), computes a normal vector by averaging the normals of all adjacent cells. This provides a smooth approximation of the surface normal at each vertex.

Four weighting schemes are available (following industry conventions from Autodesk Maya and 3ds Max):

“area”: Area-weighted averaging, where larger faces have more influence on the vertex normal. The normal at vertex v is computed as: point_normal_v = normalize(sum(cell_normal * cell_area)). This reduces the influence of small sliver triangles.
“unweighted”: Simple averaging, where each adjacent face contributes equally regardless of size. The normal at vertex v is: point_normal_v = normalize(sum(cell_normal)). This matches PyVista/VTK’s compute_normals behavior.
“angle”: Angle-weighted averaging, where faces are weighted by the interior angle at the vertex. Faces with larger angles at the vertex have more influence. This often provides the most geometrically accurate normals for curved surfaces.
“angle_area” (default): Combined angle and area weighting, where each face’s contribution is weighted by both its area and the angle at the vertex. This is the default in Maya and balances both geometric factors.

Normal vectors are only well-defined for codimension-1 manifolds, where each cell has a unique normal direction. For higher codimensions, normals are ambiguous and this method will raise an error.

Parameters:

weighting ({"area", "unweighted", "angle", "angle_area"}) –

Weighting scheme for averaging adjacent cell normals.

”area”: Weight by cell area (larger faces have more influence).
”unweighted”: Equal weight for all adjacent cells (matches PyVista/VTK).
”angle”: Weight by interior angle at the vertex.
”angle_area”: Weight by both angle and area (Maya default).

Returns:

Tensor of shape (n_points, n_spatial_dims) containing unit normal vectors at each vertex. For isolated points (with no adjacent cells), the normal is a zero vector.

Return type:

torch.Tensor

Raises:

ValueError – If the mesh is not codimension-1 (n_manifold_dims ≠ n_spatial_dims - 1), if an invalid weighting scheme is specified, or if angle-based weighting is requested for 1-simplices (edges) which have no interior angle.

See also

point_normals: Property returning angle-area-weighted normals (canonical default).
cell_normals: Compute cell (face) normals.

Examples

>>> # Triangle mesh in 3D
>>> mesh = create_triangle_mesh_3d()
>>> normals = mesh.compute_point_normals()  # area-weighted (default)
>>> normals_unweighted = mesh.compute_point_normals(weighting="unweighted")
>>> normals_angle = mesh.compute_point_normals(weighting="angle")
>>> # Normals are unit vectors (or zero for isolated points)
>>> assert torch.allclose(normals.norm(dim=-1), torch.ones(mesh.n_points), atol=1e-6)

curl( field: str | tuple[str, ...] | Float[Tensor, 'n 3'], data_source: Literal['points', 'cells'] = 'points', ) → Float[Tensor, 'n 3'][source]#

Curl of a 3D vector point or cell field (LSQ), returned as a tensor.

Accepts a field key (looked up in point_data / cell_data according to data_source) or a raw vector tensor of shape (n, 3), mirroring integrate(). Only defined for n_spatial_dims == 3.

Parameters:

field (str, tuple[str, ...], or torch.Tensor) – Vector field, by data key or by value.
data_source ({"points", "cells"}, optional) – Whether field lives at vertices (default) or at cell centers.

Returns:

Curl vector per entity, shape (n_points, 3) or (n_cells, 3) according to data_source.

Return type:

torch.Tensor

property device: device#: Retrieves the device type of tensor class.

Displace every mesh point by a dense vector field.

Computes points + displacement without changing connectivity, optionally multiplying the displacement by point_weights[..., None]. displacement and point_weights may be raw tensors or keys (including nested tuple keys) in point_data.

Call it as displace(mesh, ...) or as mesh.displace(...). The bound method supplies mesh automatically.

Parameters:

mesh (Mesh) – Mesh whose points are displaced. The source mesh is not modified.
displacement (str, tuple[str, ...], or torch.Tensor) – Dense displacement vectors with shape (mesh.n_points, mesh.n_spatial_dims), or a point-data key resolving to such a tensor. The tensor and mesh.points must have the same float32 or float64 dtype and device.
point_weights (str, tuple[str, ...], torch.Tensor, or None, optional) – Optional bool or floating-point weights with shape (mesh.n_points,), or a point-data key resolving to those point weights. Floating-point weights may be signed or greater than one. Default is None.
implementation ({"torch"} or None, optional) – Backend override. None selects Torch for dense displacement.

Returns:

New mesh with displaced points and unchanged connectivity and attached fields.

Return type:

Mesh

Notes

Attached fields are treated as Lagrangian data and are not pushed forward. Geometry-dependent caches are invalidated and topology caches are retained. The operation does not detect or repair inverted, degenerate, or self-intersecting cells; call validate() explicitly when needed.

divergence( field: str | tuple[str, ...] | Float[Tensor, 'n n_spatial_dims'], method: Literal['lsq', 'dec'] = 'lsq', data_source: Literal['points', 'cells'] = 'points', ) → Float[Tensor, 'n'][source]#

Divergence of a vector point or cell field, returned as a tensor.

Accepts a field key (looked up in point_data / cell_data according to data_source) or a raw vector tensor of shape (n, n_spatial_dims), mirroring integrate().

Parameters:

field (str, tuple[str, ...], or torch.Tensor) – Vector field, by data key or by value.
method ({"lsq", "dec"}) – Discretization (default "lsq"). "dec" is only available for point data (the DEC operators act on vertex forms).
data_source ({"points", "cells"}, optional) – Whether field lives at vertices (default) or at cell centers.

Returns:

Scalar divergence per entity, shape (n_points,) or (n_cells,) according to data_source.

Return type:

torch.Tensor

draw( backend: Literal['matplotlib', 'pyvista', 'auto'] = 'auto', show: bool = True, point_scalars: None | Tensor | str | tuple[str, ...] = None, cell_scalars: None | Tensor | str | tuple[str, ...] = None, cmap: str = 'viridis', vmin: float | None = None, vmax: float | None = None, alpha_points: float = 1.0, alpha_cells: float = 1.0, alpha_edges: float = 1.0, show_edges: bool = True, ax: AxesOrPlotter | None = None, backend_options: dict[str, Any] | None = None, ) → AxesOrPlotter#

Draw a mesh using matplotlib or PyVista backend.

This is the main visualization function for Mesh objects. It automatically selects the appropriate backend based on spatial dimensions, or allows explicit backend specification. Call it as draw(mesh, ...) or as mesh.draw(...). The bound method supplies mesh automatically.

Parameters:

mesh (Mesh) – Mesh object to visualize.
backend ({"auto", "matplotlib", "pyvista"}) –
Visualization backend to use:
- ”auto”: Automatically select based on n_spatial_dims (matplotlib for 0D/1D/2D, PyVista for 3D)
- ”matplotlib”: Force matplotlib backend (supports 3D via mplot3d)
- ”pyvista”: Force PyVista backend (requires n_spatial_dims <= 3)
show (bool) – Whether to display the plot immediately (calls plt.show() or plotter.show()). If False, returns the plotter/axes for further customization before display.
point_scalars (torch.Tensor or str or tuple[str, ...] or None, optional) –
Scalar data to color points. Mutually exclusive with cell_scalars. Can be:
- None: Points use neutral color (black)
- torch.Tensor: Direct scalar values, shape (n_points,) or (n_points, …) where trailing dimensions are L2-normed
- str or tuple[str, …]: Key to lookup in mesh.point_data
cell_scalars (torch.Tensor or str or tuple[str, ...] or None, optional) –
Scalar data to color cells. Mutually exclusive with point_scalars. Can be:
- None: Cells use neutral color (lightblue if no scalars, lightgray if point_scalars active)
- torch.Tensor: Direct scalar values, shape (n_cells,) or (n_cells, …) where trailing dimensions are L2-normed
- str or tuple[str, …]: Key to lookup in mesh.cell_data
cmap (str) – Colormap name for scalar visualization.
vmin (float or None, optional) – Minimum value for colormap normalization. If None, uses data min.
vmax (float or None, optional) – Maximum value for colormap normalization. If None, uses data max.
alpha_points (float) – Opacity for points, range [0, 1].
alpha_cells (float) – Opacity for cells/faces, range [0, 1].
alpha_edges (float) – Opacity for cell edges, range [0, 1].
show_edges (bool) – Whether to draw cell edges.
ax (matplotlib.axes.Axes or pyvista.Plotter, optional) – Existing canvas to draw on. For matplotlib, a matplotlib Axes; for PyVista, a pyvista Plotter. If None, a new figure/plotter is created. Use this to overlay multiple meshes on the same scene.
backend_options (dict[str, Any], optional) – Additional keyword arguments forwarded to the underlying visualization backend (e.g. PyVista’s plotter.add_mesh()).

Returns:

matplotlib backend: matplotlib.axes.Axes object
PyVista backend: pyvista.Plotter object

Return type:

matplotlib.axes.Axes or pyvista.Plotter

Raises:

ValueError – If both point_scalars and cell_scalars are specified, or if n_spatial_dims is not supported by the chosen backend, or if backend selection fails.
ImportError – If the requested backend is not installed.

Examples

>>> # Draw mesh with automatic backend selection
>>> mesh.draw()
>>>
>>> # Color cells by pressure data
>>> mesh.draw(cell_scalars="pressure", cmap="coolwarm")
>>>
>>> # Color points by velocity magnitude (computing norm of vector field)
>>> mesh.draw(point_scalars="velocity")  # velocity is (n_points, 3)
>>>
>>> # Use nested TensorDict key
>>> mesh.draw(cell_scalars=("flow", "temperature"))
>>>
>>> # Customize and display later
>>> import matplotlib.pyplot as plt
>>> ax = mesh.draw(show=False, backend="matplotlib")
>>> ax.set_title("My Mesh")
>>> plt.show()

dumps( prefix: str | None = None, copy_existing: bool = False, *, num_threads: int = 0, return_early: bool = False, share_non_tensor: bool = False, robust_key: bool | None = True, archive: bool | None = None, compression: str | int | None = None, ) → Any#

Saves the tensordict to disk.

This function is a proxy to memmap().

classmethod fields()#

Return a tuple describing the fields of this dataclass.

Accepts a dataclass or an instance of one. Tuple elements are of type Field.

free_form_deform( control_displacements: Float[Tensor, '*lattice_resolution n_spatial_dims'], *, origin: Float[Tensor, 'n_spatial_dims'] | Sequence[float] | None = None, extent: Float[Tensor, 'n_spatial_dims'] | Sequence[float] | None = None, basis: Literal['bernstein', 'bspline', 'linear', 'cubic_hermite', 'quintic_hermite'] = 'bernstein', point_weights: str | tuple[str, ...] | Bool[Tensor, 'n_points'] | Float[Tensor, 'n_points'] | None = None, implementation: Literal['torch', 'warp'] | None = None, ) → Mesh#

Deform a mesh with a control-point lattice by free-form deformation.

An n_1 x ... x n_D array of control displacements defines a field over the axis-aligned box [origin, origin + extent]. Mesh points inside the box move with the tensor-product basis interpolation of those values. Points outside the box are unchanged. A lattice of zero displacements is exactly the identity, and a constant lattice translates every point inside the box.

Call it as free_form_deform(mesh, ...) or as mesh.free_form_deform(...). The bound method supplies mesh automatically.

Parameters:

mesh (Mesh) – Mesh whose points are deformed. The source mesh is not modified.
control_displacements (torch.Tensor) – Displacement vectors, not destination coordinates, for every lattice node, with shape (n_1, ..., n_D, mesh.n_spatial_dims) and the same float32 or float64 dtype and device as mesh.points. Each axis needs at least two nodes for "bernstein" and the node-interpolating bases, and four for "bspline".
origin (torch.Tensor, sequence of float, or None, optional) – Minimum corner of the lattice box with shape (mesh.n_spatial_dims,). None uses the minimum corner of the mesh bounds. For repeated GPU calls with an explicit box, create origin and extent once as device tensors. Reuse them to avoid recreating and transferring sequence values. Default is None.
extent (torch.Tensor, sequence of float, or None, optional) – Edge lengths of the lattice box with the same accepted shapes as origin. Every value must be finite and strictly positive. The operation does not validate tensor values at runtime. None sizes the box from origin to the maximum corner of the mesh bounds. Validating a derived extent synchronizes with the device and is not CUDA Graph capture-safe. For capture, pass both origin and extent as device tensors. Every point-coordinate axis must have positive range when the extent is derived. Supply an explicit extent for lower-dimensional geometry embedded in a higher-dimensional space. Default is None.
basis ({"bernstein", "bspline", "linear", "cubic_hermite", "quintic_hermite"}, optional) –
Per-axis basis family:
- "bernstein" provides classic global-support FFD. Every lattice node influences every point inside the box.
- "bspline" uses a uniform cubic B-spline with local four-node-per-axis support and C2 continuity between knot spans. Coefficient index i corresponds to local coordinate (i - 1) / (n - 3). The first and last coefficient planes lie outside the evaluation box.
- "linear" uses upper-node weight \(s(t)=t\) within each lattice cell. It interpolates every node. It is continuous (C0) across cell boundaries, but its slope can jump.
- "cubic_hermite" uses the cubic Hermite blend \(s(t)=3t^2-2t^3\). Its first derivative vanishes at both cell endpoints. This gives C1 continuity across cell boundaries.
- "quintic_hermite" uses the quintic Hermite blend \(s(t)=6t^5-15t^4+10t^3\). Its first and second derivatives vanish at both endpoints. This gives C2 continuity across cell boundaries. Perlin introduced this improved interpolant in “Improving Noise” [1].
The node-interpolating bases use only the two neighboring nodes per axis. Here, t is the local cell coordinate in [0, 1]. The upper-node weight is \(s(t)\), and the lower-node weight is \(1-s(t)\). Default is "bernstein".
point_weights (str, tuple[str, ...], torch.Tensor, or None, optional) – Optional bool or floating mesh-point weights with shape (mesh.n_points,), or a point_data key resolving to those point weights. All weights must match the point device. Floating weights must also match the point dtype. Default is None.
implementation ({"torch", "warp"} or None, optional) – Backend override. None selects Torch on CPU. On CUDA, it selects Warp when available and otherwise Torch.

Returns:

New mesh with deformed points and unchanged connectivity and fields.

Return type:

Mesh

Raises:

TypeError – If tensors, lattice values, or point weights have unsupported types or dtypes.
ValueError – If shapes, devices, lattice parameters, point weights, or basis are invalid.
KeyError – If a point-data key or implementation name is not found.
ImportError – If an explicitly requested backend is unavailable.

Notes

The operation treats attached fields as Lagrangian data and does not push them forward. It invalidates geometry-dependent caches and retains topology caches. The deformation is generally not continuous across the lattice box boundary. To keep the exterior fixed, zero the outermost coefficient plane on every Bernstein or node-interpolating face. For cubic B-splines, zero the first and last three coefficient planes on every axis. origin and extent are non-differentiable lattice parameters. Optimize control_displacements instead. The operation does not detect or repair inverted, degenerate, or self-intersecting cells. Call validate() explicitly when needed.

References

[1] Perlin, K. (2002). “Improving Noise.” ACM Transactions on Graphics, 21(3), 681-682. https://doi.org/10.1145/566654.566636

Creates a TensorDict from a CSV file.

Requires either pandas or pyarrow to be installed.

Parameters:

path (str or Path) – Path to the CSV file.

Keyword Arguments:

auto_batch_size (bool, optional) – If True, the batch size will be computed automatically. Defaults to False.
batch_dims (int, optional) – If auto_batch_size is True, defines how many dimensions the output tensordict should have. Defaults to None.
device (torch.device, optional) – The device for tensor data. Defaults to None.
batch_size (torch.Size, optional) – The batch size. Defaults to [num_rows].
separator (str, optional) – If provided, column names are split on this separator to create nested TensorDicts. Defaults to None.
dtype (torch.dtype, optional) – If provided, all numeric columns are cast to this dtype. Defaults to None.
**kwargs – Additional keyword arguments forwarded to the CSV reader (pandas.read_csv or pyarrow.csv.read_csv).

Returns:

A TensorDict representation of the CSV data.

Examples

>>> td = TensorDict.from_csv("data.csv")
>>> td = TensorDict.from_csv("data.csv", separator=".", dtype=torch.float32)

Creates a TensorDict from a JSON file.

Supports both standard JSON (array of records) and JSON Lines format. For nested JSON objects, use from_dict() instead.

Requires pandas for best results. Falls back to stdlib json for simple cases.

Parameters:

path (str or Path) – Path to the JSON file.

Keyword Arguments:

auto_batch_size (bool, optional) – If True, the batch size will be computed automatically. Defaults to False.
batch_dims (int, optional) – If auto_batch_size is True, defines how many dimensions the output tensordict should have. Defaults to None.
device (torch.device, optional) – The device for tensor data. Defaults to None.
batch_size (torch.Size, optional) – The batch size. Defaults to [num_rows].
separator (str, optional) – If provided, column names are split on this separator to create nested TensorDicts. Defaults to None.
dtype (torch.dtype, optional) – If provided, all numeric columns are cast to this dtype. Defaults to None.
lines (bool, optional) – If True, reads the file as JSON Lines (one JSON object per line). Defaults to False.
**kwargs – Additional keyword arguments forwarded to the JSON reader.

Returns:

A TensorDict representation of the JSON data.

Examples

>>> td = TensorDict.from_json("data.json")
>>> td = TensorDict.from_json("data.jsonl", lines=True)

Converts a pandas DataFrame to a TensorDict.

Numeric columns become tensors, string/object columns become NonTensorData.

Parameters:

dataframe (pd.DataFrame) – The pandas DataFrame to convert.

Keyword Arguments:

auto_batch_size (bool, optional) – If True, the batch size will be computed automatically. Defaults to False.
batch_dims (int, optional) – If auto_batch_size is True, defines how many dimensions the output tensordict should have. Defaults to None.
device (torch.device, optional) – The device for tensor data. Defaults to None.
batch_size (torch.Size, optional) – The batch size. Defaults to [num_rows].
separator (str, optional) – If provided, column names are split on this separator to create nested TensorDicts. For example, with separator=".", a column "obs.x" becomes td["obs", "x"]. Defaults to None.
dtype (torch.dtype, optional) – If provided, all numeric columns are cast to this dtype. Defaults to None.

Returns:

A TensorDict representation of the DataFrame.

Examples

>>> import pandas as pd
>>> df = pd.DataFrame({"a": [1, 2, 3], "b": [4.0, 5.0, 6.0]})
>>> td = TensorDict.from_pandas(df)
>>> print(td)
TensorDict(
    fields={
        a: Tensor(shape=torch.Size([3]), device=cpu, dtype=torch.int64, is_shared=False),
        b: Tensor(shape=torch.Size([3]), device=cpu, dtype=torch.float64, is_shared=False)},
    batch_size=torch.Size([3]),
    device=None,
    is_shared=False)

Creates a TensorDict from a Parquet file.

Requires either pyarrow or pandas to be installed. Prefers pyarrow when available for better performance.

Parameters:

path (str or Path) – Path to the Parquet file.

Keyword Arguments:

auto_batch_size (bool, optional) – If True, the batch size will be computed automatically. Defaults to False.
batch_dims (int, optional) – If auto_batch_size is True, defines how many dimensions the output tensordict should have. Defaults to None.
device (torch.device, optional) – The device for tensor data. Defaults to None.
batch_size (torch.Size, optional) – The batch size. Defaults to [num_rows].
separator (str, optional) – If provided, column names are split on this separator to create nested TensorDicts. Defaults to None.
dtype (torch.dtype, optional) – If provided, all numeric columns are cast to this dtype. Defaults to None.
columns (list of str, optional) – If provided, only read these columns from the file. Defaults to None (all columns).
**kwargs – Additional keyword arguments forwarded to the Parquet reader.

Returns:

A TensorDict representation of the Parquet data.

Examples

>>> td = TensorDict.from_parquet("data.parquet")
>>> td = TensorDict.from_parquet("data.parquet", columns=["obs", "reward"])

Build a triangulated surface Mesh from a polygon soup.

Triangulates a polygon cell-to-vertex incidence (an Adjacency of vertex rings, as produced by VTK-style readers) into the simplex-only Mesh representation, and broadcasts any per-polygon cell_data to the resulting triangles.

Triangulation uses physicsnemo.mesh.tessellation.triangulate(): a vectorized vertex-0 fan for convex polygons and ear clipping for the rare non-convex ones (so unsigned-area-weighted integrals stay correct).

Parameters:

points (torch.Tensor) – Vertex coordinates of shape \((N_\text{points}, D)\).
polygons (Adjacency) – Cell-to-vertex incidence (CSR): polygon p is the vertex ring polygons.indices[polygons.offsets[p] : polygons.offsets[p + 1]].
point_data (TensorDict or dict[str, torch.Tensor], optional) – Per-vertex data, carried through unchanged.
cell_data (TensorDict or dict[str, torch.Tensor], optional) – Per-polygon data; broadcast to each polygon’s triangles via the triangulation’s parent_index.
global_data (TensorDict or dict[str, torch.Tensor], optional) – Mesh-level data, carried through unchanged.
assume_convex (bool, default False) – If True, skip the convexity test and ear-clip fallback and fan-triangulate every polygon (correct only for convex inputs).

Returns:

A triangle mesh (cells of shape \((N_\text{triangles}, 3)\)).

Return type:

Mesh

Notes

Each polygon ring must be a simple, approximately planar polygon with no repeated consecutive vertices; see physicsnemo.mesh.tessellation.triangulate() for the full input contract.

Examples

>>> import torch
>>> from physicsnemo.mesh import Mesh
>>> from physicsnemo.mesh.neighbors import Adjacency
>>> points = torch.tensor([[0.0, 0.0, 0.0], [1.0, 0.0, 0.0],
...                        [1.0, 1.0, 0.0], [0.0, 1.0, 0.0]])
>>> polygons = Adjacency(offsets=torch.tensor([0, 4]),  # one quad
...                      indices=torch.tensor([0, 1, 2, 3]))
>>> mesh = Mesh.from_polygons(
...     points, polygons, cell_data={"p": torch.tensor([2.5])}
... )
>>> mesh.n_cells
2
>>> mesh.cell_data["p"].tolist()
[2.5, 2.5]

from_schema(

*,

batch_size: Sequence[int] | Size | None = None,

storage: str | None = None,

device=None,

**kwargs,

) → TensorDictBase#

Pre-allocate a zero-filled TensorDict from a schema.

Creates a TensorDictBase whose storage backend is selected by storage. Each entry in schema maps a field name to an (element_shape, dtype) pair; the full stored shape is [*batch_size, *element_shape].

Parameters:

schema – Mapping from field name to (element_shape, dtype). element_shape is the per-element shape (excluding batch_size).

Keyword Arguments:

batch_size – Overall batch dimensions prepended to every element shape. Defaults to ().
storage (str or None) –
Backend selector:
- None – plain TensorDict with regular tensors.
- "memmap" – memory-mapped tensors on disk. Pass prefix=<dir> in kwargs.
- "h5" – HDF5 via PersistentTensorDict. Pass filename=<path> in kwargs.
- "zarr" – zarr (requires zarr>=3.0) via PersistentTensorDict. Pass filename=<path or store> in kwargs.
- "shared" – CPU shared-memory tensors.
- "redis" / "dragonfly" – delegates to TensorDictStore.from_schema().
device – Device for the resulting tensors (ignored by some backends).
**kwargs – Backend-specific arguments forwarded to the underlying constructor (e.g. prefix for memmap, filename for h5, host/port for redis).

Returns:

A new TensorDictBase subclass instance with pre-allocated (zero-filled) keys.

Examples

>>> td = TensorDict.from_schema(
...     {"obs": ([84, 84, 3], torch.uint8),
...      "reward": ([], torch.float32)},
...     batch_size=[1000],
... )
>>> td["obs"].shape
torch.Size([1000, 84, 84, 3])

>>> import tempfile
>>> with tempfile.TemporaryDirectory() as d:
...     td_mm = TensorDict.from_schema(
...         {"obs": ([4], torch.float32)},
...         batch_size=[8],
...         storage="memmap",
...         prefix=d,
...     )
...     assert td_mm.is_memmap()

classmethod from_tensordict( tensordict: TensorDictBase, non_tensordict: dict | None = None, safe: bool = True, ) → Any#

Tensor class wrapper to instantiate a new tensor class object.

Parameters:

tensordict (TensorDictBase) – Dictionary of tensor types
non_tensordict (dict) – Dictionary with non-tensor and nested tensor class objects
safe (bool) – Whether to raise an error if the tensordict is not a TensorDictBase instance

property gaussian_curvature_cells: Tensor#

Compute Gaussian curvature at cell centers.

Averages the intrinsic vertex-based Gaussian curvature (angle defect) over each cell’s vertices, giving a cell-centered field consistent with gaussian_curvature_vertices.

The result is cached in _cache["cell", "gaussian_curvature"] for efficiency.

Returns:: Tensor of shape (n_cells,) containing Gaussian curvature at cells.
Return type:: torch.Tensor

Examples

>>> from physicsnemo.mesh.primitives.surfaces import sphere_icosahedral
>>> mesh = sphere_icosahedral.load(subdivisions=2)
>>> K_cells = mesh.gaussian_curvature_cells

property gaussian_curvature_vertices: Tensor#

Compute intrinsic Gaussian curvature at mesh vertices.

Uses the angle-defect method from discrete differential geometry. For a vertex \(v\) with incident cells \(\sigma \ni v\) and interior angle \(\theta_\sigma(v)\) at \(v\) in each \(\sigma\),

\[K(v) = \frac{\Theta(v)}{|{\star}v|}, \quad \Theta(v) = \Theta_n - \sum_{\sigma \ni v} \theta_\sigma(v),\]

where \(\Theta_n\) is the full angle in an \(n\)-dimensional manifold and \(|{\star}v|\) is the dual 0-cell (Voronoi) volume. This is an intrinsic measure of curvature (Theorema Egregium) that works for any codimension, as it depends only on distances within the manifold.

Signed curvature:

Positive: elliptic/convex (sphere-like).
Zero: flat/parabolic (plane-like).
Negative: hyperbolic/saddle (saddle-like).

The result is cached in _cache["point", "gaussian_curvature"] for efficiency.

Returns:: Signed Gaussian curvature, shape (n_points,). Isolated vertices have NaN curvature.
Return type:: torch.Tensor

Notes

Satisfies the discrete Gauss-Bonnet theorem,

\[\sum_v K(v) \, |{\star}v| = 2 \pi \, \chi(M),\]

where the sum is over vertices and \(\chi(M)\) is the Euler characteristic.

Examples

>>> from physicsnemo.mesh.primitives.surfaces import sphere_icosahedral
>>> # Sphere of radius r has K = 1/r^2
>>> sphere = sphere_icosahedral.load(radius=2.0, subdivisions=3)
>>> K = sphere.gaussian_curvature_vertices
>>> # K.mean() approx 0.25 (= 1 / 2.0^2)

get(key: NestedKey, *args, **kwargs)#

Gets the value stored with the input key.

Parameters:

key (str, tuple of str) – key to be queried. If tuple of str it is equivalent to chained calls of getattr.
default – default value if the key is not found in the tensorclass.

Returns:

value stored with the input key

get_boundary_mesh( data_source: Literal['points', 'cells'] = 'cells', data_aggregation: Literal['mean', 'area_weighted', 'inverse_distance'] = 'mean', ) → Mesh[source]#

Extract the boundary surface of this mesh.

Convenience wrapper around get_facet_mesh() that extracts only boundary facets (those appearing in exactly one parent cell).

See get_facet_mesh() for full parameter documentation.

Parameters:

data_source ({"points", "cells"}, optional) – Source of data inheritance. Default: “cells”.
data_aggregation ({"mean", "area_weighted", "inverse_distance"}, optional) – Strategy for aggregating data. Default: “mean”.

Returns:

Boundary mesh containing only boundary facets.

Return type:

Mesh

Notes

For meshes with internal cavities (like volume meshes with voids or drivaerML-style automotive meshes), this returns BOTH the exterior surface and any interior cavity surfaces. All facets that appear in exactly one parent cell are included, regardless of whether they face “outward” or “inward”.

Examples

>>> from physicsnemo.mesh.primitives.procedural import lumpy_ball
>>> from physicsnemo.mesh.primitives.surfaces import sphere_icosahedral
>>> # Extract triangular surface of a volume mesh
>>> vol_mesh = lumpy_ball.load(n_shells=2, subdivisions=1)
>>> surface_mesh = vol_mesh.get_boundary_mesh()
>>> assert surface_mesh.n_manifold_dims == 2  # triangles
>>>
>>> # For a closed watertight sphere
>>> sphere = sphere_icosahedral.load(subdivisions=3)
>>> boundary = sphere.get_boundary_mesh()
>>> assert boundary.n_cells == 0  # no boundary

get_cell_to_cells_adjacency(adjacency_codimension: int = 1)[source]#

Compute cell-to-cells adjacency based on shared facets.

Two cells are considered adjacent if they share a k-codimension facet.

The result is cached in _cache["topology", ...] for efficiency, keyed by adjacency_codimension. Adjacency depends only on topology (cells), not geometry (points), so the cache is preserved through geometric transforms.

Parameters:

adjacency_codimension (int, optional) –

Codimension of shared facets defining adjacency.

1 (default): Cells must share a codimension-1 facet (e.g., triangles sharing an edge, tetrahedra sharing a triangular face)
2: Cells must share a codimension-2 facet (e.g., tetrahedra sharing an edge)
k: Cells must share a codimension-k facet

Returns:

Adjacency where adjacency.to_list()[i] contains all cell indices that share a k-codimension facet with cell i.

Return type:

Adjacency

Examples

>>> from physicsnemo.mesh.primitives.basic import two_triangles_2d
>>> mesh = two_triangles_2d.load()
>>> adj = mesh.get_cell_to_cells_adjacency(adjacency_codimension=1)
>>> # Get cells sharing an edge with cell 0
>>> neighbors_of_cell_0 = adj.to_list()[0]

get_cell_to_points_adjacency()[source]#

Get the vertices (points) that comprise each cell.

This is a simple wrapper around the cells array that returns it in the standard Adjacency format for consistency with other neighbor queries.

The result is cached in _cache["topology", ...] for efficiency.

Returns:: Adjacency where adjacency.to_list()[i] contains all point indices that are vertices of cell i. For simplicial meshes, all cells have the same number of vertices (n_manifold_dims + 1).
Return type:: Adjacency

Examples

>>> from physicsnemo.mesh.primitives.basic import two_triangles_2d
>>> mesh = two_triangles_2d.load()
>>> adj = mesh.get_cell_to_points_adjacency()
>>> # Get vertices of cell 0
>>> vertices_of_cell_0 = adj.to_list()[0]

get_facet_mesh( manifold_codimension: int = 1, data_source: Literal['points', 'cells'] = 'cells', data_aggregation: Literal['mean', 'area_weighted', 'inverse_distance'] = 'mean', target_counts: list[int] | Literal['boundary', 'shared', 'interior', 'all'] = 'all', ) → Mesh[source]#

Extract k-codimension facet mesh from this n-dimensional mesh.

Extracts all (n-k)-simplices from the current n-simplicial mesh. For example:

Triangle mesh (2-simplices) → edge mesh (1-simplices) [codimension=1, default]
Triangle mesh (2-simplices) → vertex mesh (0-simplices) [codimension=2]
Tetrahedral mesh (3-simplices) → triangular facet mesh (2-simplices) [codimension=1, default]
Tetrahedral mesh (3-simplices) → edge mesh (1-simplices) [codimension=2]

The resulting mesh shares the same vertex positions but has connectivity representing the lower-dimensional simplices. Data can be inherited from either the parent cells or the boundary points.

Parameters:

manifold_codimension (int, optional) –
Codimension of extracted mesh relative to parent.
- 1: Extract (n-1)-facets (default, immediate boundaries of all cells)
- 2: Extract (n-2)-facets (e.g., edges from tets, vertices from triangles)
- k: Extract (n-k)-facets
data_source ({"points", "cells"}, optional) –
Source of data inheritance:
- ”cells”: Facets inherit from parent cells they bound. When multiple cells share a facet, data is aggregated according to data_aggregation.
- ”points”: Facets inherit from their boundary vertices. Data from multiple boundary points is averaged.
data_aggregation ({"mean", "area_weighted", "inverse_distance"}, optional) –
Strategy for aggregating data from multiple sources (only applies when data_source=”cells”):
- ”mean”: Simple arithmetic mean
- ”area_weighted”: Weighted by parent cell areas
- ”inverse_distance”: Weighted by inverse distance from facet centroid to parent cell centroids
target_counts (list[int] | {"boundary", "shared", "interior", "all"}, optional) –
Which facets to keep based on how many parent cells share them:
- ”all”: Keep all unique facets (default)
- ”boundary”: Keep only boundary facets (appearing in exactly 1 cell)
- ”shared”: Keep only shared facets (appearing in 2+ cells)
- ”interior”: Keep only interior facets (appearing in exactly 2 cells)
- list[int]: Keep facets with counts matching any value in the list

Returns:

New Mesh with n_manifold_dims = self.n_manifold_dims - manifold_codimension, embedded in the same spatial dimension. The mesh shares the same points array but has new cells connectivity and aggregated cell_data.

Return type:

Mesh

Raises:

ValueError – If manifold_codimension is too large for this mesh (would result in negative manifold dimension).

Examples

>>> from physicsnemo.mesh.primitives.basic import two_triangles_2d
>>> # Extract edges from a triangle mesh (codimension 1)
>>> triangle_mesh = two_triangles_2d.load()
>>> edge_mesh = triangle_mesh.get_facet_mesh(manifold_codimension=1)
>>> assert edge_mesh.n_manifold_dims == 1  # edges
>>>
>>> # Extract vertices from a triangle mesh (codimension 2)
>>> vertex_mesh = triangle_mesh.get_facet_mesh(manifold_codimension=2)
>>> assert vertex_mesh.n_manifold_dims == 0  # vertices
>>> facet_mesh = triangle_mesh.get_facet_mesh(
...     data_source="cells",
...     data_aggregation="area_weighted"
... )

get_point_to_cells_adjacency()[source]#

Compute the star of each vertex (all cells containing each point).

For each point in the mesh, finds all cells that contain that point. This is the graph-theoretic “star” operation on vertices.

The result is cached in _cache["topology", ...] for efficiency. Adjacency depends only on topology (cells), not geometry (points), so the cache is preserved through geometric transforms.

Returns:: Adjacency where adjacency.to_list()[i] contains all cell indices that contain point i. Isolated points (not in any cells) have empty lists.
Return type:: Adjacency

Examples

>>> from physicsnemo.mesh.primitives.basic import two_triangles_2d
>>> mesh = two_triangles_2d.load()
>>> adj = mesh.get_point_to_cells_adjacency()
>>> # Get cells containing point 0
>>> cells_of_point_0 = adj.to_list()[0]

get_point_to_points_adjacency()[source]#

Compute point-to-point adjacency (graph edges of the mesh).

For each point, finds all other points that share a cell with it. In simplicial meshes, this is equivalent to finding all points connected by an edge.

The result is cached in _cache["topology", ...] for efficiency. Adjacency depends only on topology (cells), not geometry (points), so the cache is preserved through geometric transforms.

Returns:: Adjacency where adjacency.to_list()[i] contains all point indices that share a cell (edge) with point i. Isolated points have empty lists.
Return type:: Adjacency

Examples

>>> from physicsnemo.mesh.primitives.basic import two_triangles_2d
>>> mesh = two_triangles_2d.load()
>>> adj = mesh.get_point_to_points_adjacency()
>>> # Get neighbors of point 0
>>> neighbors_of_point_0 = adj.to_list()[0]

gradient( field: str | tuple[str, ...] | Float[Tensor, 'n ...'], method: Literal['lsq', 'dec'] = 'lsq', gradient_type: Literal['intrinsic', 'extrinsic'] = 'intrinsic', data_source: Literal['points', 'cells'] = 'points', ) → Float[Tensor, 'n n_spatial_dims ...'][source]#

Gradient of a point or cell field, returned as a tensor.

Single-field convenience that returns the gradient tensor directly, accepting a field key (looked up in point_data / cell_data according to data_source) or a raw tensor – mirroring integrate(). (Contrast compute_point_derivatives() / compute_cell_derivatives(), which return a new mesh with the gradient stored under an auto-generated key, and can process several fields at once.)

Parameters:

field (str, tuple[str, ...], or torch.Tensor) – Field, by data key or by value.
method ({"lsq", "dec"}) – Discretization (default "lsq"). "dec" is only available for point data: the DEC exterior derivative maps vertex 0-forms to edge 1-forms, and there is no analogous cell-to-cell operator.
gradient_type ({"intrinsic", "extrinsic"}) – Project onto the tangent space ("intrinsic", default) or use the full ambient-space gradient ("extrinsic").
data_source ({"points", "cells"}, optional) – Whether field lives at vertices (default) or at cell centers.

Returns:

Gradient of shape (n, n_spatial_dims, *field.shape[1:]), where n is n_points or n_cells according to data_source. For a vector field, gradient[i, k, j] is \(\partial field_{i,j} / \partial x_k\).

Return type:

torch.Tensor

integrate( field: str | tuple[str, ...] | Float[Tensor, 'n_cells_or_points ...'], data_source: Literal['cells', 'points'] = 'cells', *, nan_policy: Literal['omit', 'propagate'] = 'omit', ) → Float[Tensor, '...']#

Integrate a field over the mesh domain.

This is the public entry point for ordinary mesh-field integration. It selects P0 cell quadrature or P1 point quadrature from data_source and resolves field from a string key or tensor.

Call it as integrate(mesh, ...) or as mesh.integrate(...). The bound method supplies mesh automatically.

Parameters:

mesh (Mesh) – Simplicial mesh.
field (str, tuple[str, ...], or torch.Tensor) –
Field to integrate.
- str or tuple: looked up in cell_data or point_data according to data_source.
- torch.Tensor: used directly.
data_source ({"cells", "points"}) – Whether field is cell-centered (P0) or vertex-centered (P1).
nan_policy ({"omit", "propagate"}, default "omit") – NaN reduction behavior. "omit" preserves the historical masked-data behavior; "propagate" is appropriate when NaNs should remain visible, such as inside neural-operator reductions.

Returns:

Integral value. Shape matches the trailing dimensions of the field (scalar field -> 0-d tensor, vector field -> 1-d tensor, etc.).

Return type:

torch.Tensor

Raises:

KeyError – If field is a string key not present in the specified data source.
ValueError – If the mesh has no cells, or if a raw tensor has the wrong leading dimension for the specified data_source.

Examples

>>> import torch
>>> from physicsnemo.mesh import Mesh
>>> pts = torch.tensor([[0., 0.], [1., 0.], [0.5, 1.]])
>>> cells = torch.tensor([[0, 1, 2]])
>>> mesh = Mesh(points=pts, cells=cells)
>>> mesh.cell_data["p"] = torch.tensor([3.0])
>>> mesh.integrate("p")  # integrate cell-centered pressure
tensor(1.5000)
>>> mesh.point_data["T"] = torch.tensor([1.0, 2.0, 3.0])
>>> mesh.integrate("T", data_source="points")  # P1 integral
tensor(1.)

integrate_flux( field: str | tuple[str, ...] | Float[Tensor, 'n_cells_or_points n_spatial_dims'], data_source: Literal['cells', 'points'] = 'cells', *, nan_policy: Literal['omit', 'propagate'] = 'omit', ) → Float[Tensor, '']#

Compute the surface flux integral for codimension-1 meshes.

Computes the oriented flux of a vector field through the mesh surface:

\[\int_\Gamma \mathbf{F} \cdot \mathbf{n}\,d\Gamma\]

This is only defined for codimension-1 meshes (surfaces in 3D, curves in 2D) where unique cell normals exist.

Call it as integrate_flux(mesh, ...) or as mesh.integrate_flux(...). The bound method supplies mesh automatically.

For cell data, the flux is:

\[\int_\Gamma \mathbf{F} \cdot \mathbf{n}\,d\Gamma = \sum_c (\mathbf{F}_c \cdot \mathbf{n}_c)\,|\sigma_c|\]

For point data, the P1 vertex-averaged field is dotted with the cell normal (which is constant per cell):

\[\int_\Gamma \mathbf{F} \cdot \mathbf{n}\,d\Gamma = \sum_c \Bigl(\frac{1}{n_v}\sum_{v \in c} \mathbf{F}(v)\Bigr) \cdot \mathbf{n}_c\,|\sigma_c|\]

Parameters:

mesh (Mesh) – Codimension-1 simplicial mesh (i.e. n_manifold_dims == n_spatial_dims - 1).
field (str, tuple[str, ...], or torch.Tensor) – Vector field to integrate. Must have last dimension equal to n_spatial_dims.
data_source ({"cells", "points"}) – Whether field is cell-centered or vertex-centered.
nan_policy ({"omit", "propagate"}, default "omit") – "omit" excludes cells whose normal flux is NaN. "propagate" uses an ordinary sum so any NaN normal-flux contribution remains visible.

Returns:

Scalar flux value (0-d tensor).

Return type:

torch.Tensor

Raises:

KeyError – If field is a string key not present in the specified data source.
ValueError – If the mesh is not codimension-1, if the field leading dimension does not match the expected entity count, or if the field does not have the correct trailing dimension.

Examples

>>> import torch
>>> from physicsnemo.mesh import Mesh
>>> # Unit square boundary in 2D (4 edges forming a closed loop)
>>> pts = torch.tensor([[0., 0.], [1., 0.], [1., 1.], [0., 1.]])
>>> cells = torch.tensor([[0, 1], [1, 2], [2, 3], [3, 0]])
>>> mesh = Mesh(points=pts, cells=cells)
>>> # Constant outward velocity field - flux through closed boundary
>>> mesh.cell_data["v"] = torch.zeros(4, 2)
>>> mesh.integrate_flux("v")
tensor(0.)

integrate_moment( left: str | tuple[str, ...] | Float[Tensor, 'n_cells ...'], right: str | tuple[str, ...] | Float[Tensor, 'n_cells ...'], *, aligned_dims: int = 0, accumulation_dtype: dtype | None = torch.float32, nan_policy: Literal['omit', 'propagate'] = 'omit', ) → Tensor#

Integrate the outer product of two cell-centered fields.

Computes the P0 quadrature moment

\[M = \sum_c |\sigma_c|\, a_c \otimes b_c,\]

where a is left, b is right, and \(|\sigma_c|\) is the cell’s effective measure (its geometric area times any recorded measure weight; see physicsnemo.mesh.calculus.measure). By default the result has shape left.shape[1:] + right.shape[1:]. aligned_dims may designate a common leading subset of the trailing dimensions as independent groups; those axes appear only once in the output rather than participating in the outer product. The implementation evaluates a batched weighted matrix product and never materializes the per-cell outer product.

Call it as integrate_moment(mesh, ...) or as mesh.integrate_moment(...). The bound method supplies mesh automatically.

Parameters:

mesh (Mesh) – Simplicial mesh with at least one cell.
left (str, tuple[str, ...], or torch.Tensor) – Cell-centered fields. String and tuple keys are resolved from mesh.cell_data. Their leading dimensions must equal mesh.n_cells; arbitrary trailing dimensions are supported.
right (str, tuple[str, ...], or torch.Tensor) – Cell-centered fields. String and tuple keys are resolved from mesh.cell_data. Their leading dimensions must equal mesh.n_cells; arbitrary trailing dimensions are supported.
aligned_dims (int, default=0) – Number of leading trailing dimensions shared by left and right and treated as aligned batch/group axes. For example, inputs shaped (N, H, A) and (N, H, B) with aligned_dims=1 produce (H, A, B) instead of (H, A, H, B). The aligned shapes must match exactly.
accumulation_dtype (torch.dtype or None, default torch.float32) – Minimum dtype used by the weighted matrix product. The actual compute dtype is the promotion of both inputs, the cell measures, and this dtype, so the default accumulates reduced-precision inputs in at least FP32 without downcasting FP64 inputs. Pass None to use ordinary input promotion with no additional precision floor, or torch.float64 to request at least FP64 accumulation.
nan_policy ({"omit", "propagate"}, default "omit") – "omit" replaces NaN field contributions with zero before the matrix product. "propagate" leaves them untouched.

Returns:

Weighted outer-product moment with shape aligned_shape + left_event_shape + right_event_shape and the accumulation dtype.

Return type:

torch.Tensor

Raises:

KeyError – If a named field is absent from mesh.cell_data.
TypeError – If aligned_dims is not an integer or accumulation_dtype is not floating-point or complex.
ValueError – If the mesh is empty, a leading dimension is wrong, aligned dimensions are invalid, the fields are on different devices, or nan_policy is invalid.

Notes

nan_policy="omit" is intended for finite data containing NaN masks. As with any matrix product, indeterminate expressions involving infinities (for example 0 * inf) are not treated as missing data.

is_manifold( check_level: Literal['facets', 'edges', 'full'] = 'full', ) → bool#

Check if mesh is a valid topological manifold.

A mesh is a manifold if it locally looks like Euclidean space at every point. This function checks various topological constraints depending on the check level.

Call it as is_manifold(mesh, ...) or as mesh.is_manifold(...). The bound method supplies mesh automatically.

Parameters:

mesh (Mesh) – Input simplicial mesh to check
check_level ({"facets", "edges", "full"}, optional) –
Level of checking to perform:
- ”facets”: Only check codimension-1 facets (each appears 1-2 times)
- ”edges”: Check facets + edge neighborhoods (for 2D/3D meshes)
- ”full”: Complete manifold validation (default)

Returns:

True if mesh passes the specified manifold checks, False otherwise

Return type:

bool

Examples

>>> from physicsnemo.mesh.primitives.surfaces import sphere_icosahedral, cylinder_open
>>> # Valid manifold (sphere)
>>> sphere = sphere_icosahedral.load(subdivisions=3)
>>> assert is_manifold(sphere) == True
>>>
>>> # Manifold with boundary (open cylinder)
>>> cylinder = cylinder_open.load()
>>> assert is_manifold(cylinder) == True  # manifold with boundary is OK

Notes

This function checks topological constraints but does not check for geometric self-intersections (which would require expensive spatial queries).

is_watertight() → bool#

Check if mesh is watertight (has no boundary).

A mesh is watertight if every codimension-1 facet is shared by exactly 2 cells. This means the mesh forms a closed surface/volume with no holes or gaps.

Call it as is_watertight(mesh) or as mesh.is_watertight(). The bound method supplies mesh automatically.

Parameters:: mesh (Mesh) – Input simplicial mesh to check
Returns:: True if mesh is watertight (no boundary facets), False otherwise
Return type:: bool

Examples

>>> from physicsnemo.mesh.primitives.surfaces import sphere_icosahedral, cylinder_open
>>> # Closed sphere is watertight
>>> sphere = sphere_icosahedral.load(subdivisions=3)
>>> assert is_watertight(sphere) == True
>>>
>>> # Open cylinder with holes at ends
>>> cylinder = cylinder_open.load()
>>> assert is_watertight(cylinder) == False

laplacian( field: str | tuple[str, ...] | Float[Tensor, 'n ...'], data_source: Literal['points', 'cells'] = 'points', ) → Float[Tensor, 'n ...'][source]#

Laplace-Beltrami operator on a point field (DEC), returned as a tensor.

Uses the intrinsic cotangent Laplacian (physicsnemo.mesh.calculus.compute_laplacian_points_dec()). Accepts a field key (looked up in point_data) or a raw point tensor, mirroring integrate().

Parameters:

field (str, tuple[str, ...], or torch.Tensor) – Point field, by point_data key or by value.
data_source ({"points", "cells"}, optional) – Only "points" is supported: the cotangent Laplace-Beltrami operator is defined on vertex functions, and there is no DEC Laplacian for cell-centered data. The kwarg exists for signature consistency with gradient() / divergence() / curl(); passing "cells" raises. (For a cell-centered Laplacian, compose mesh.divergence(mesh.gradient(f, gradient_type="extrinsic", data_source="cells"), data_source="cells") explicitly – a double-LSQ discretization with different accuracy properties.)

Returns:

Laplace-Beltrami of the field, same shape as the input field.

Return type:

torch.Tensor

classmethod load(prefix: str | Path, *args, **kwargs) → Any#

Loads a tensordict from disk.

This class method is a proxy to load_memmap().

load_(prefix: str | Path, *args, **kwargs)#

Loads a tensordict from disk within the current tensordict.

This class method is a proxy to load_memmap_().

classmethod load_memmap( prefix: str | Path, device: device | None = None, non_blocking: bool = False, *, out: TensorDictBase | None = None, robust_key: bool | None = True, subpath: NestedKey | None = None, mode: str = 'r', num_threads: int = 0, allow_pickle: bool | None = None, ) → Any#

Loads a memory-mapped tensordict from disk.

Parameters:

prefix (str or Path to folder) – the path to the folder where the saved tensordict should be fetched, or the path to a memmap archive file written through save(..., archive=True) / a ".tdz" prefix (or packed with pack_memmap()). Archives are memory-mapped once and every leaf is exposed as a zero-copy view into the mapping: only the pages of the leaves that are actually accessed are read from disk. Unlike directory-backed tensordicts, in-place writes to the leaves of an archive-loaded tensordict do not propagate to the file.
device (torch.device or equivalent, optional) – if provided, the data will be asynchronously cast to that device. Supports “meta” device, in which case the data isn’t loaded but a set of empty “meta” tensors are created. This is useful to get a sense of the total model size and structure without actually opening any file.
non_blocking (bool, optional) – if True, synchronize won’t be called after loading tensors on device. Defaults to False.
out (TensorDictBase, optional) – optional tensordict where the data should be written.
robust_key (bool, optional) – if True (default), expects robust key encoding was used when saving and decodes filenames accordingly. If False, uses legacy behavior. If None, uses the default robust behavior.
subpath (NestedKey or str path, optional) – the location of a nested tensordict to load, as a nested key (e.g. ("module", "0"), with arbitrary nesting allowed as usual) or as a "/"-separated string path (e.g. "module/0"). Only that subtree is loaded. Works both for directories (equivalent to appending the path to prefix) and for archives.
mode (str, optional) – "r" (default) or "r+". Only relevant when loading an archive: with "r" the archive is mapped copy-on-write and in-place writes to the leaves stay in memory; with "r+" the mapping is shared and in-place writes propagate to the file, like directory-backed tensordicts. "r+" requires uncompressed, aligned tensor payloads (i.e. archives written by tensordict without compression) and is not available for nested-tensor leaves. In-place writes do not update the per-entry CRC-32 stored by the zip format; load_memmap() ignores checksums, but call refresh_archive_checksums() before handing a modified archive to tools that verify them (unzip, unpack_memmap(), …). Directory prefixes are always write-through and ignore this argument.
num_threads (int, optional) – number of threads used to decompress the leaves of a compressed archive (deflate entries are inflated in parallel, which scales nearly linearly). Without compression, loading is a metadata-only operation and this argument has no effect. Defaults to 0 (sequential).
allow_pickle (bool, optional) – whether pickled non-tensor fields may be loaded. Pickle can execute arbitrary code, so pass True only for data from a trusted source and False for untrusted data. During the 0.14 compatibility window, omitting this option loads pickle with a FutureWarning; the default will change to False in 0.15. Saves without a pickle sidecar do not require this option.

Examples

>>> from tensordict import TensorDict
>>> td = TensorDict.fromkeys(["a", "b", "c", ("nested", "e")], 0)
>>> td.memmap("./saved_td")
>>> td_load = TensorDict.load_memmap("./saved_td")
>>> assert (td == td_load).all()

This method also allows loading nested tensordicts.

Examples

>>> nested = TensorDict.load_memmap("./saved_td/nested")
>>> assert nested["e"] == 0

A tensordict can also be loaded on “meta” device or, alternatively, as a fake tensor.

Examples

>>> import tempfile
>>> td = TensorDict({"a": torch.zeros(()), "b": {"c": torch.zeros(())}})
>>> with tempfile.TemporaryDirectory() as path:
...     td.save(path)
...     td_load = TensorDict.load_memmap(path, device="meta")
...     print("meta:", td_load)
...     from torch._subclasses import FakeTensorMode
...     with FakeTensorMode():
...         td_load = TensorDict.load_memmap(path)
...         print("fake:", td_load)
meta: TensorDict(
    fields={
        a: Tensor(shape=torch.Size([]), device=meta, dtype=torch.float32, is_shared=False),
        b: TensorDict(
            fields={
                c: Tensor(shape=torch.Size([]), device=meta, dtype=torch.float32, is_shared=False)},
            batch_size=torch.Size([]),
            device=meta,
            is_shared=False)},
    batch_size=torch.Size([]),
    device=meta,
    is_shared=False)
fake: TensorDict(
    fields={
        a: FakeTensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, is_shared=False),
        b: TensorDict(
            fields={
                c: FakeTensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, is_shared=False)},
            batch_size=torch.Size([]),
            device=cpu,
            is_shared=False)},
    batch_size=torch.Size([]),
    device=cpu,
    is_shared=False)

load_state_dict( state_dict: dict[str, Any], strict=True, assign=False, from_flatten=None, )#

Loads a state_dict into the tensorclass.

Supports both the new format (logical keys with _metadata) and the legacy format (_tensordict/_non_tensordict wrapper keys).

property mean_curvature_vertices: Tensor#

Compute extrinsic mean curvature at mesh vertices.

Uses the cotangent Laplace-Beltrami operator:: H = (1/2) * ||L @ points|| / voronoi_area

Mean curvature is an extrinsic measure (depends on embedding) and is only defined for codimension-1 manifolds where normal vectors exist.

For 2D surfaces: H = (k1 + k2) / 2 where k1, k2 are principal curvatures

Signed curvature:

Positive: Convex (sphere exterior with outward normals)
Negative: Concave (sphere interior with outward normals)
Zero: Minimal surface (soap film)

The result is cached in _cache["point", "mean_curvature"] for efficiency.

Returns:: Tensor of shape (n_points,) containing signed mean curvature. Isolated vertices have NaN curvature.
Return type:: torch.Tensor
Raises:: ValueError – If mesh is not codimension-1.

Examples

>>> from physicsnemo.mesh.primitives.surfaces import sphere_icosahedral
>>> # Sphere of radius r has H = 1/r
>>> sphere = sphere_icosahedral.load(radius=2.0, subdivisions=3)
>>> H = sphere.mean_curvature_vertices
>>> # H.mean() ≈ 0.5 (= 1/2.0)

memmap( prefix: str | None = None, copy_existing: bool = False, *, num_threads: int = 0, return_early: bool = False, share_non_tensor: bool = False, existsok: bool = True, robust_key: bool | None = True, archive: bool | None = None, compression: str | int | None = None, ) → Any#

Writes all tensors onto a corresponding memory-mapped Tensor in a new tensordict.

Parameters:

prefix (str) – directory prefix where the memory-mapped tensors will be stored. The directory tree structure will mimic the tensordict’s. If prefix ends with ".tdz" (or archive=True is passed), a single-file archive is written instead of a directory: a standard zip file whose entries replicate the memmap directory layout. See archive below.
copy_existing (bool) – If False (default), an exception will be raised if an entry in the tensordict is already a tensor stored on disk with an associated file, but is not saved in the correct location according to prefix. If True, any existing Tensor will be copied to the new location.

Keyword Arguments:

num_threads (int, optional) – the number of threads used to write the memmap tensors. Defaults to 0.
return_early (bool, optional) – if True and num_threads>0, the method will return a future of the tensordict.
share_non_tensor (bool, optional) – if True, the non-tensor data will be shared between the processes and writing operation (such as inplace update or set) on any of the workers within a single node will update the value on all other workers. If the number of non_tensor leaves is high (e.g., sharing large stacks of non-tensor data) this may result in OOM or similar errors. Defaults to False.
existsok (bool, optional) – if False, an exception will be raised if a tensor already exists in the same path. Defaults to True.
robust_key (bool, optional) – if True (default), uses robust key encoding that safely handles keys with path separators and special characters. If False, uses legacy behavior (keys used as-is). If None, uses the default robust behavior.
archive (bool, optional) – if True, prefix designates a single file rather than a directory and the tensordict is written as a memmap archive: a zip file mirroring the memmap directory tree, with tensor payloads stored uncompressed and aligned so that load_memmap() can memory-map the file and expose every leaf as a zero-copy view. If None (default), archive mode is enabled when prefix ends with ".tdz". The result of load_memmap() on an archive behaves like the result of from_consolidated(): all leaves are views into a single storage, and in-place writes do not propagate to the file. Archives and memmap directories are mutually convertible with pack_memmap() / unpack_memmap() (or any zip tool). Note that archives are written sequentially (single data pass) and num_threads has no effect on them.
compression (str or int, optional) – compression for archive entries ("stored", "deflate", "bzip2", "lzma" or a zipfile constant). Defaults to "stored" (uncompressed), which is what enables zero-copy loading. Compressed archives load correctly but leaves are decompressed in memory on access. Only valid in archive mode.

The TensorDict is then locked, meaning that any writing operations that isn’t in-place will throw an exception (eg, rename, set or remove an entry). Once the tensordict is unlocked, the memory-mapped attribute is turned to False, because cross-process identity is not guaranteed anymore.

Returns:: A new tensordict with the tensors stored on disk if return_early=False, otherwise a TensorDictFuture instance.

Note

Serialising in this fashion might be slow with deeply nested tensordicts, so it is not recommended to call this method inside a training loop.

memmap_( prefix: str | None = None, copy_existing: bool = False, *, num_threads: int = 0, return_early: bool = False, share_non_tensor: bool = False, existsok: bool = True, robust_key: bool | None = True, ) → Any#

Writes all tensors onto a corresponding memory-mapped Tensor, in-place.

Parameters:

prefix (str) – directory prefix where the memory-mapped tensors will be stored. The directory tree structure will mimic the tensordict’s.
copy_existing (bool) – If False (default), an exception will be raised if an entry in the tensordict is already a tensor stored on disk with an associated file, but is not saved in the correct location according to prefix. If True, any existing Tensor will be copied to the new location.

Keyword Arguments:

num_threads (int, optional) – the number of threads used to write the memmap tensors. Defaults to 0.
return_early (bool, optional) – if True and num_threads>0, the method will return a future of the tensordict. The resulting tensordict can be queried using future.result().
share_non_tensor (bool, optional) – if True, the non-tensor data will be shared between the processes and writing operation (such as inplace update or set) on any of the workers within a single node will update the value on all other workers. If the number of non-tensor leaves is high (e.g., sharing large stacks of non-tensor data) this may result in OOM or similar errors. Defaults to False.
existsok (bool, optional) – if False, an exception will be raised if a tensor already exists in the same path. Defaults to True.
robust_key (bool, optional) – if True (default), uses robust key encoding that safely handles keys with path separators and special characters. If False, uses legacy behavior (keys used as-is). If None, uses the default robust behavior.

The TensorDict is then locked, meaning that any writing operations that isn’t in-place will throw an exception (eg, rename, set or remove an entry). Once the tensordict is unlocked, the memory-mapped attribute is turned to False, because cross-process identity is not guaranteed anymore.

Returns:: self if return_early=False, otherwise a TensorDictFuture instance.

Note

Serialising in this fashion might be slow with deeply nested tensordicts, so it is not recommended to call this method inside a training loop.

memmap_like( prefix: str | None = None, copy_existing: bool = False, *, existsok: bool = True, num_threads: int = 0, return_early: bool = False, share_non_tensor: bool = False, robust_key: bool | None = True, archive: bool | None = None, ) → Any#

Creates a contentless Memory-mapped tensordict with the same shapes as the original one.

Parameters:

prefix (str) – directory prefix where the memory-mapped tensors will be stored. The directory tree structure will mimic the tensordict’s. If prefix ends with ".tdz" (or archive=True is passed), a preallocated single-file archive is created instead. See archive below.
copy_existing (bool) – If False (default), an exception will be raised if an entry in the tensordict is already a tensor stored on disk with an associated file, but is not saved in the correct location according to prefix. If True, any existing Tensor will be copied to the new location.

Keyword Arguments:

num_threads (int, optional) – the number of threads used to write the memmap tensors. Defaults to 0.
return_early (bool, optional) – if True and num_threads>0, the method will return a future of the tensordict.
share_non_tensor (bool, optional) – if True, the non-tensor data will be shared between the processes and writing operation (such as inplace update or set) on any of the workers within a single node will update the value on all other workers. If the number of non-tensor leaves is high (e.g., sharing large stacks of non-tensor data) this may result in OOM or similar errors. Defaults to False.
existsok (bool, optional) – if False, an exception will be raised if a tensor already exists in the same path. Defaults to True.
robust_key (bool, optional) – if True (default), uses robust key encoding that safely handles keys with path separators and special characters. If False, uses legacy behavior (keys used as-is). If None, uses the default robust behavior.
archive (bool, optional) – if True, prefix designates a single file and a preallocated, zero-filled memmap archive is created and loaded back with load_memmap(prefix, mode="r+"): the returned tensordict writes through to the archive. If None (default), archive mode is enabled when prefix ends with ".tdz". In-place writes leave the zip per-entry checksums stale; call refresh_archive_checksums() before handing the archive to tools that verify them. Nested tensors are not supported in this mode.

The TensorDict is then locked, meaning that any writing operations that isn’t in-place will throw an exception (eg, rename, set or remove an entry). Once the tensordict is unlocked, the memory-mapped attribute is turned to False, because cross-process identity is not guaranteed anymore.

Returns:: A new TensorDict instance with data stored as memory-mapped tensors if return_early=False, otherwise a TensorDictFuture instance.

Note

This is the recommended method to write a set of large buffers on disk, as memmap_() will copy the information, which can be slow for large content.

Examples

>>> td = TensorDict({
...     "a": torch.zeros((3, 64, 64), dtype=torch.uint8),
...     "b": torch.zeros(1, dtype=torch.int64),
... }, batch_size=[]).expand(1_000_000)  # expand does not allocate new memory
>>> buffer = td.memmap_like("/path/to/dataset")

memmap_refresh_(*, allow_pickle: bool | None = None)#

Refreshes the content of the memory-mapped tensordict if it has a saved_path.

This method will raise an exception if no path is associated with it.

Parameters:: allow_pickle (bool, optional) – whether pickled non-tensor fields may be loaded. See load_memmap().

classmethod merge( meshes: Sequence[Mesh], global_data_strategy: Literal['stack'] = 'stack', ) → Mesh[source]#

Merge multiple meshes into a single mesh.

Parameters:

meshes (Sequence[Mesh]) – List of Mesh objects to merge. All constituent tensors across all meshes must reside on the same device.
global_data_strategy ({"stack"}) – Strategy for handling global_data. Currently only “stack” is supported, which stacks global_data fields along a new dimension.

Returns:

A new Mesh object containing all the merged data.

Return type:

Mesh

Raises:

ValueError – If the meshes list is empty, or if meshes have inconsistent dimensions or cell_data keys.
TypeError – If any element in meshes is not a Mesh object.
RuntimeError – If tensors from different meshes reside on different devices.

morph( control_points: Tensor, control_displacements: Tensor, *, radius: float | Tensor, point_weights: str | tuple[str, ...] | Tensor | None = None, kernel: Literal['wendland_c2'] = 'wendland_c2', implementation: Literal['torch', 'warp'] | None = None, ) → Mesh#

Morph a mesh from sparse, compactly supported control displacements.

Control influence uses a Wendland-C2 compact Shepard field with a stationary zero-displacement background. Each control’s influence vanishes smoothly at its support boundary. Where supports overlap, all active controls and the background are blended together; the result is not a simple sum or average. The field is zero outside the union of all supports.

With no point weights, the field at a unique control coordinate is exactly that control’s displacement. Duplicate controls at one coordinate contribute their mean displacement. A control may be anywhere in world coordinates and need not coincide with a mesh point.

Call it as morph(mesh, ...) or as mesh.morph(...). The bound method supplies mesh automatically.

Parameters:

mesh (Mesh) – Mesh whose points are morphed. The source mesh is not modified.
control_points (torch.Tensor) – World-coordinate controls with shape (n_controls, mesh.n_spatial_dims) and the same dtype and device as mesh.points.
control_displacements (torch.Tensor) – Displacement vectors, not destination coordinates, with exactly the same shape, dtype, and device as control_points.
radius (float or torch.Tensor) – Support distance in mesh coordinate units. Supply one scalar for every control or a tensor with shape (n_controls,) that matches the control dtype and device. Every tensor value must remain positive and finite; values are not validated at runtime.
point_weights (str, tuple[str, ...], torch.Tensor, or None, optional) – Optional bool or floating mesh-point weights with shape (mesh.n_points,), or a point_data key resolving to those point weights. These are query-point weights, not per-control values. Default is None.
kernel ({"wendland_c2"}, optional) – Compact radial kernel used to blend control displacements. Default is "wendland_c2".
implementation ({"torch", "warp"} or None, optional) – Backend override. None selects Torch on CPU and Warp on CUDA when Warp is available, otherwise Torch.

Returns:

New mesh with morphed points and unchanged connectivity and attached fields.

Return type:

Mesh

Notes

Attached fields are treated as Lagrangian data and are not pushed forward. Geometry-dependent caches are invalidated and topology caches are retained. Parameterize learned radii to remain positive, for example as torch.nn.functional.softplus(raw_radius) + eps. The operation does not detect or repair inverted, degenerate, or self-intersecting cells; call validate() explicitly when needed.

property n_cells: int#: Number of cells in the mesh.

property n_manifold_dims: int#: Intrinsic dimension of each simplicial cell.

property n_points: int#: Number of points in the mesh.

property n_spatial_dims: int#: Dimension of the ambient coordinate space.

pad( target_n_points: int | None = None, target_n_cells: int | None = None, data_padding_value: float = nan, ) → Mesh[source]#

Pad points and cells arrays to specified sizes.

This is the low-level padding method that performs the actual padding operation. Padding uses null/degenerate elements that don’t affect computations:

Points: Additional points at the last existing point (preserves bounding box)
cells: Degenerate cells with all vertices at the last existing point (zero area)
cell data: NaN-valued padding for all cell data fields (default)

Parameters:

target_n_points (int or None, optional) – Target number of points. If None, no point padding is applied. Must be >= current n_points if specified. Also accepts SymInt for torch.compile.
target_n_cells (int or None, optional) – Target number of cells. If None, no cell padding is applied. Must be >= current n_cells if specified. Also accepts SymInt for torch.compile.
data_padding_value (float) – Value to use for padding data fields. Defaults to NaN for floating and complex fields; integer and boolean fields use 0 when the requested value is NaN, preserving their dtype.

Returns:

A new Mesh with padded arrays. If both targets are None or equal to current sizes, returns self unchanged.

Return type:

Mesh

Raises:

ValueError – If target sizes are less than current sizes.

Examples

>>> mesh = Mesh(points, cells)  # 100 points, 200 cells
>>> padded = mesh.pad(target_n_points=128, target_n_cells=256)
>>> padded.n_points  # 128
>>> padded.n_cells   # 256

pad_to_next_power( power: float = 1.5, data_padding_value: float = nan, ) → Mesh[source]#

Pads points and cells arrays to their next power of power (integer-floored).

This is useful for torch.compile with dynamic=False, where fixed tensor shapes are required. By padding to powers of a base (default 1.5), we can reuse compiled kernels across a reasonable range of mesh sizes while minimizing memory overhead.

This method computes the target sizes as floor(power^n) for the smallest n such that the result is >= the current size, then calls .pad() to perform the actual padding.

Parameters:

power (float) – Base for computing the next power. Must be > 1. Provides a good balance between memory efficiency and compile cache hits.
data_padding_value (float) – Value to use for padding data fields. Defaults to NaN for floating and complex fields; integer and boolean fields use 0 when the requested value is NaN, preserving their dtype.

Returns:

A new Mesh with padded points and cells arrays. The padding uses null elements that don’t affect geometric computations.

Return type:

Mesh

Raises:

ValueError – If power <= 1.

Examples

>>> mesh = Mesh(points, cells)  # 100 points, 200 cells
>>> padded = mesh.pad_to_next_power(power=1.5)
>>> # Points padded to floor(1.5^n) >= 100, cells to floor(1.5^m) >= 200
>>> # For power=1.5: 100 points -> 129 points, 200 cells -> 216 cells
>>> # Padding cells have zero area and don't affect computations

point_data_to_cell_data( overwrite_keys: bool = False, ) → Mesh[source]#

Convert point data to cell data by averaging.

For each cell, computes the average of the point data values from all points (vertices) that define that cell. The resulting cell data is added to the mesh’s cell_data dictionary. Original point data is preserved.

Parameters:: overwrite_keys (bool) – If True, silently overwrite any existing cell_data keys. If False, raise an error if a key already exists in cell_data.
Returns:: New Mesh with converted data added to cell_data. Original point_data is preserved.
Return type:: Mesh
Raises:: ValueError – If a point_data key already exists in cell_data and overwrite_keys=False.

Examples

>>> mesh = Mesh(points, cells, point_data={"temperature": point_temps})
>>> mesh_with_cell_data = mesh.point_data_to_cell_data()
>>> # Now mesh has both point_data["temperature"] and cell_data["temperature"]

property point_normals: Tensor#

Compute weighted normal vectors at mesh vertices.

This property returns the canonical/default point normals. For 2D+ manifolds (surfaces, volumes), angle-area weighting is used, which balances face area and vertex interior angle for high-quality normals. For 1D manifolds (curves), area weighting (i.e. segment-length weighting) is used, since interior angles are not defined for edges.

For explicit weighting control, use compute_point_normals().

The result is cached in _cache["point", "normals"] for efficiency.

Returns:: Tensor of shape (n_points, n_spatial_dims) containing unit normal vectors at each vertex. For isolated points (with no adjacent cells), the normal is a zero vector.
Return type:: torch.Tensor
Raises:: ValueError – If the mesh is not codimension-1 (n_manifold_dims != n_spatial_dims - 1).

See also

compute_point_normals: Compute point normals with explicit weighting choice.
cell_normals: Compute cell (face) normals.

Examples

>>> # Triangle mesh in 3D
>>> mesh = create_triangle_mesh_3d()
>>> normals = mesh.point_normals  # (n_points, 3), angle-area-weighted
>>> # Normals are unit vectors (or zero for isolated points)
>>> assert torch.allclose(normals.norm(dim=-1), torch.ones(mesh.n_points), atol=1e-6)

property quality_metrics: TensorDict#

Compute geometric quality metrics for every cell.

Returns:: Per-cell metrics including normalized aspect ratio, edge-length ratio, minimum and maximum angles, and a combined quality score. A regular simplex has aspect ratio and quality score equal to 1.
Return type:: TensorDict

Examples

>>> from physicsnemo.mesh.primitives.basic import two_triangles_2d
>>> mesh = two_triangles_2d.load()
>>> metrics = mesh.quality_metrics
>>> assert "quality_score" in metrics.keys()

See also

physicsnemo.mesh.validation.compute_quality_metrics: Standalone functional form.

radial_basis_function_deform( control_points: Float[Tensor, 'n_controls n_spatial_dims'], control_displacements: Float[Tensor, 'n_controls n_spatial_dims'], *, kernel: Literal['thin_plate_spline'] = 'thin_plate_spline', polynomial: bool = True, smoothing: float = 0.0, point_weights: str | tuple[str, ...] | Bool[Tensor, 'n_points'] | Float[Tensor, 'n_points'] | None = None, implementation: Literal['torch', 'warp'] | None = None, ) → Mesh#

Deform a mesh with a global thin-plate-spline RBF field.

A thin-plate-spline radial field is fitted to the prescribed sparse control displacements and evaluated at every mesh point. With the default affine polynomial tail, zero smoothing, and a nonsingular control layout, the unweighted field interpolates every control displacement up to solver precision.

Call it as radial_basis_function_deform(mesh, ...) or as mesh.radial_basis_function_deform(...). The bound method supplies mesh automatically.

Parameters:

mesh (Mesh) – Mesh whose points are deformed. The source mesh is not modified.
control_points (torch.Tensor) – World-coordinate controls with shape (n_controls, mesh.n_spatial_dims) and the same dtype and device as mesh.points.
control_displacements (torch.Tensor) – Displacement vectors, not destination coordinates, with exactly the same shape, dtype, and device as control_points.
kernel ({"thin_plate_spline"}, optional) – Radial kernel used by the interpolant. Default is "thin_plate_spline".
polynomial (bool, optional) – Add the standard affine polynomial tail and side constraints. This reproduces affine displacement fields. The controls must affinely span the coordinate space, and the augmented system must be nonsingular. Default is True.
smoothing (float, optional) – Nonnegative diagonal regularization added to the radial system. Zero gives exact interpolation for a nonsingular control layout up to solver precision. Positive values relax interpolation accuracy. Default is 0.0.
point_weights (str, tuple[str, ...], torch.Tensor, or None, optional) – Optional bool or floating mesh-point weights with shape (mesh.n_points,), or a point_data key resolving to those weights. Values scale or mask the fitted field after interpolation. Bool weights must be on the same device as the mesh points. Floating weights must have the same dtype and device as the mesh points.
implementation ({"torch", "warp"} or None, optional) – Evaluation-backend override. Both backends use PyTorch for the dense coefficient solve. None selects Torch on CPU and Warp on CUDA when Warp is available, otherwise Torch.

Returns:

New mesh with deformed points and unchanged connectivity and attached fields.

Return type:

Mesh

Raises:

TypeError – If control tensors or Python arguments have unsupported types, or if tensor dtypes are unsupported or mismatched.
ValueError – If tensor shapes, devices, control layout, point weights, or RBF options are invalid.
KeyError – If a point-data key is missing or implementation does not name a registered backend.
ImportError – If an explicitly requested backend is unavailable.
RuntimeError – If runtime validation or coefficient fitting fails, including for a singular system or during CUDA Graph capture.

Notes

The field has global support. Unlike compact Shepard morphing, every control generally influences every mesh point. Attached fields are treated as Lagrangian data and are not pushed forward. Geometry-dependent caches are invalidated and topology caches are retained. The operation does not detect or repair inverted, degenerate, or self-intersecting cells. Call validate() explicitly when needed. Coefficient fitting is not supported inside CUDA Graph capture because the singular-system check requires host interaction.

remesh( n_clusters: int, *, max_iterations: int = 4, transfer_point_data: PointDataSelection = False, resolution_field: ResolutionField = None, ) → Mesh#

Remesh a triangle surface with point-data and resolution controls.

Warp performs integration-mass-weighted centroidal clustering, projects cluster centers back to the source surface with a bounding volume hierarchy, and reconstructs compact triangle connectivity. A direct positive tensor or an attached point-data field can specify relative local linear resolution.

Parameters:

mesh (Mesh) – Input triangle surface. Only 2D triangle manifolds embedded in 3D are supported.
n_clusters (int) – Target output vertex count. Cleanup can produce slightly fewer vertices. Must be between 3 and the input point count, inclusive.
max_iterations (int, optional) – Maximum centroid-relaxation iterations. Default is 4. Values must be non-negative.
transfer_point_data (bool, str, tuple, list, or None, optional) – Point-data fields to interpolate onto the output vertices. False or None transfers no fields. True transfers every point-data leaf. A string or tuple selects one key or nested key path. A list selects several keys or paths. Selected fields must contain real floating-point tensors. Default is False.
resolution_field (str, tuple, torch.Tensor, or None, optional) – Positive scalar tensor with shape (n_points,), or a key or nested key path resolving to one in mesh.point_data. Values specify relative linear resolution. A value twice another requests approximately half the local edge spacing. The fixed n_clusters budget and source geometry limit the realized spacing. The field must use a real floating-point dtype on the mesh device. Direct tensor entries correspond to mesh.points order and are not attached to or transferred with the output mesh. Only relative values matter. Default is None for uniform remeshing.

Returns:

Remeshed surface on the input device. Selected point data is barycentrically interpolated from the original source surface. Cell data and unselected point data are discarded. Global data is preserved.

Return type:

Mesh

Raises:

TypeError – If counts, tuning parameters, point coordinates, a field selection, or a selected field has an invalid type.
ValueError – If a count is out of range or geometry, connectivity, or a selected field is invalid.
KeyError – If a requested point-data key or path does not exist.
NotImplementedError – If mesh is not a 2D triangle surface embedded in 3D.
ImportError – If Warp is unavailable.
RuntimeError – If cleanup cannot reconstruct a nonempty manifold triangle surface or point-data transfer provenance is unavailable.

Notes

Remeshing, topology, projection choices, and resolution control are intentionally non-differentiable. Transferred fields remain differentiable with respect to their source values because the final barycentric interpolation uses PyTorch. Warp computes geometry in centered and scaled coordinates in float32, then restores the input point dtype and coordinate frame. For the 2D squared-distance CVT objective, the implementation converts linear resolution r to integration density r**4. Ideal local point density therefore scales approximately as r**2. These relationships guide allocation but do not guarantee exact edge lengths or local point counts. Because clustering uses spatial distance rather than mesh connectivity, sheets or thin features separated by less than the mean cluster spacing can be assigned to a common cluster and welded together. Projection can map distinct cluster centroids to the same surface position. Output vertices are compacted by connectivity but are not welded by position. Backend-specific tuning remains available through physicsnemo.nn.functional.remeshing(). These advanced parameters may change as the implementation evolves.

Rotate the mesh about an axis by a specified angle.

Call it as rotate(mesh, ...) or as mesh.rotate(...). The bound method supplies mesh automatically.

Parameters:

mesh (Mesh) – Input mesh to rotate.
angle (float) – Rotation angle in radians (counterclockwise, right-hand rule).
axis (Float[torch.Tensor, " n_spatial_dims"] or Sequence[float] or {"x", "y", "z"} or None) – Rotation axis vector. None for 2D, shape \((3,)\) for 3D. String literals "x", "y", "z" are converted to unit vectors (1,0,0), (0,1,0), (0,0,1) respectively.
center (Float[torch.Tensor, " n_spatial_dims"] or Sequence[float] or None) – Center point for rotation. If None, rotates about the origin.
transform_point_data (bool or TensorDict) – Controls transformation of point_data fields. See transform() for full semantics.
transform_cell_data (bool or TensorDict) – Same semantics as transform_point_data, for cell_data.
transform_global_data (bool or TensorDict) – Same semantics as transform_point_data, for global_data.

Returns:

New Mesh with rotated geometry.

Return type:

Mesh

Notes

Cache Handling:

areas: Unchanged (rotation preserves volumes)

centroids: Rotated

normals: Rotated

sample_data_at_points( query_points: Tensor, data_source: Literal['cells', 'points'] = 'cells', multiple_cells_strategy: Literal['mean', 'nan'] = 'mean', project_onto_nearest_cell: bool = False, tolerance: float = 1e-06, bvh: Any = None, ) → TensorDict[source]#

Extract or interpolate mesh data at specified query points.

This method retrieves mesh data at arbitrary spatial locations. Note that “sample” here means “extract/query at specific points” - NOT random sampling. For random point sampling, see sample_random_points_on_cells().

Containment queries are BVH-accelerated (O(n_queries * log(n_cells))).

Parameters:

query_points (torch.Tensor) – Query point locations, shape (n_queries, n_spatial_dims).
data_source ({"cells", "points"}, optional) –
How to retrieve data:
- ”cells”: Use cell data directly (no interpolation)
- ”points”: Interpolate point data using barycentric coordinates
multiple_cells_strategy ({"mean", "nan"}, optional) –
How to handle query points in multiple cells:
- ”mean”: Return arithmetic mean of values from all containing cells
- ”nan”: Return NaN for ambiguous points
project_onto_nearest_cell (bool, optional) – If True, snaps each query point to the centroid of the nearest cell before containment testing. Useful for codimension != 0 manifolds.
tolerance (float, optional) – Tolerance for considering a point inside a cell.
bvh (BVH or None, optional) – Pre-built Bounding Volume Hierarchy. If None (default), one is built automatically. For repeated queries, pre-build with BVH.from_mesh(mesh) and pass it here to avoid redundant work.

Returns:

Data for each query point. Values are NaN for query points outside the mesh.

Return type:

TensorDict

Examples

>>> import torch
>>> from physicsnemo.mesh.primitives.basic import two_triangles_2d
>>> mesh = two_triangles_2d.load()
>>> mesh.cell_data["pressure"] = torch.tensor([1.0, 2.0])
>>> query_pts = torch.tensor([[0.3, 0.3], [0.8, 0.5]])
>>> data = mesh.sample_data_at_points(query_pts, data_source="cells")

sample_random_points_on_cells( cell_indices: Sequence[int] | Tensor | None = None, alpha: float = 1.0, ) → Tensor[source]#

Sample random points on specified cells of the mesh.

Uses a Dirichlet distribution to generate barycentric coordinates, which are then used to compute random points as weighted combinations of cell vertices. The concentration parameter alpha controls the distribution of samples within each cell (simplex).

This is a convenience method that delegates to physicsnemo.mesh.sampling.sample_random_points_on_cells.

Parameters:

cell_indices (Sequence[int] or torch.Tensor or None, optional) – Indices of cells to sample from. Can be a Sequence or tensor. Allows repeated indices to sample multiple points from the same cell. If None, samples one point from each cell (equivalent to arange(n_cells)). Shape: (n_samples,) where n_samples is the number of points to sample.
alpha (float, optional) –
Concentration parameter for the Dirichlet distribution. Controls how samples are distributed within each cell:
- alpha = 1.0: Uniform distribution over the simplex (default)
- alpha > 1.0: Concentrates samples toward the center of each cell
- alpha < 1.0: Concentrates samples toward vertices and edges

Returns:

Random points on cells, shape (n_samples, n_spatial_dims). Each point lies within its corresponding cell. If cell_indices is None, n_samples = n_cells.

Return type:

torch.Tensor

Raises:

NotImplementedError – If alpha != 1.0 and torch.compile is being used. This is due to a PyTorch limitation with Gamma distributions under torch.compile.
IndexError – If any cell_indices are out of bounds.

Examples

>>> import torch
>>> from physicsnemo.mesh.primitives.basic import two_triangles_2d
>>> mesh = two_triangles_2d.load()
>>> # Sample one point from each cell uniformly
>>> points = mesh.sample_random_points_on_cells()
>>> assert points.shape == (mesh.n_cells, mesh.n_spatial_dims)

save( prefix: str | None = None, copy_existing: bool = False, *, num_threads: int = 0, return_early: bool = False, share_non_tensor: bool = False, robust_key: bool | None = True, archive: bool | None = None, compression: str | int | None = None, ) → Any#

Saves the tensordict to disk.

This function is a proxy to memmap().

Scale the mesh by specified factor(s).

Call it as scale(mesh, ...) or as mesh.scale(...). The bound method supplies mesh automatically.

Parameters:

mesh (Mesh) – Input mesh to scale.
factor (float or Float[torch.Tensor, " n_spatial_dims"] or Sequence[float]) – Scale factor(s). Scalar for uniform, vector for non-uniform.
center (Float[torch.Tensor, " n_spatial_dims"] or Sequence[float] or None) – Center point for scaling. If None, scales about the origin.
transform_point_data (bool or TensorDict) – Controls transformation of point_data fields. See transform() for full semantics.
transform_cell_data (bool or TensorDict) – Same semantics as transform_point_data, for cell_data.
transform_global_data (bool or TensorDict) – Same semantics as transform_point_data, for global_data.
assume_invertible (bool or None) –
Controls cache propagation:
- True: Assume all factors are non-zero, propagate caches (compile-safe)
- False: Assume some factor is zero, skip cache propagation (compile-safe)
- None: Check determinant at runtime (may cause graph breaks under torch.compile)

Returns:

New Mesh with scaled geometry.

Return type:

Mesh

Notes

Cache Handling:

areas: Scaled correctly. For non-isotropic transforms of codimension-1
embedded manifolds, per-element scaling is computed using normals.

centroids: Scaled

normals: Transformed by inverse-transpose (direction adjusted, magnitude normalized)

select( *keys, inplace: bool = False, strict: bool = True, as_tensordict: bool = False, )#: TensorClass-specific select that supports as_tensordict.

set( key: NestedKey, value: Any, inplace: bool = False, non_blocking: bool = False, )#

Sets a new key-value pair.

Parameters:

key (str, tuple of str) – name of the key to be set. If tuple of str it is equivalent to chained calls of getattr followed by a final setattr.
value (Any) – value to be stored in the tensorclass
inplace (bool, optional) – if True, set will tentatively try to update the value in-place. If False or if the key isn’t present, the value will be simply written at its destination.

Returns:

self

Project mesh vertices onto the nearest locations of a target surface.

For each source vertex \(x_i\), Shrinkwrap selects the closest point \(p_i\) on target and returns

\[x'_i = x_i + w_i\left(p_i + \delta n_i - x_i\right),\]

where \(w_i\) is an optional point weight, \(\delta\) is offset, and \(n_i\) is the oriented unit normal of the selected target face. Source connectivity and attached data are retained.

Call it as shrinkwrap(mesh, target, ...) or as mesh.shrinkwrap(target, ...). Neither source nor target is modified.

Parameters:

mesh (Mesh) – Source mesh whose 3D point coordinates are projected. Its connectivity may have any supported topology.
target (Mesh) – Nonempty triangle surface embedded in 3D. Source and target points must share their float32 or float64 dtype and device. Target connectivity may use any non-bool integer dtype supported by Mesh and is normalized to int64 for this operation.
offset (float or torch.Tensor, optional) – Signed scalar distance from the target along the selected face normal. A scalar tensor must match the mesh point dtype and device and may require gradients. The value must be finite. Positive values follow target face winding. Default is 0.0.
max_distance (float or None, optional) – Positive finite nearest-surface search radius in mesh coordinate units. Vertices without a target strictly closer than this distance remain unchanged. Values that round to zero in the mesh point dtype are rejected. None performs an unbounded search. Default is None.
point_weights (str, tuple[str, ...], torch.Tensor, or None, optional) – Optional bool or floating per-source-point weights, or a key/path in point_data resolving to one. Zero leaves a vertex unchanged and one applies the full projection. Floating values are not clamped. Default is None.
implementation ({"torch", "warp"} or None, optional) – Nearest-face backend. None selects Torch on CPU and Warp on CUDA when available.

Returns:

New mesh with projected points, preserved connectivity and data, invalidated geometry caches, and retained topology caches.

Return type:

Mesh

Notes

Nearest-face selection is discrete. Away from target face, edge, vertex, and distance-cutoff transitions, gradients propagate through source points, selected target vertices, floating point weights, and tensor-valued offset. Degenerate target faces are ignored. Target orientation must be consistent when using a nonzero offset.

Shrinkwrap does not detect or prevent inverted or self-intersecting source cells. Partial point weights interpolate toward the projected surface and generally do not place the resulting vertex exactly on it. Shrinkwrap is not supported inside CUDA Graph capture with either backend. Float64 targets use the Torch search because Warp searches in float32. Safe float32 coordinates are searched unchanged. Warp falls back to Torch for unsafe coordinate magnitudes or face geometry.

Returns a new Mesh with a subset of the cells.

Parameters:: indices (int or slice or torch.Tensor) – Indices or mask to select cells.
Returns:: New Mesh with subset of cells.
Return type:: Mesh

Notes

Slicing shares unsliced data with the source by reference rather than copying: the returned mesh shares points, point_data, and global_data with this mesh, and the no-op selections None / Ellipsis return this mesh itself. Mutating any shared field on the result therefore also mutates the source; clone first if you need an independent copy.

Returns a new Mesh with a subset of the points.

This method filters points and automatically updates cells to maintain consistency. Cells that reference any removed points are also removed, and the remaining cells have their indices remapped to the new point numbering.

Parameters:

indices (int or slice or Ellipsis or None or torch.Tensor or Sequence) –

Indices or mask to select points. Supports:

int: Single point index
slice: Python slice object
Ellipsis or None: Keep all points (returns self)
torch.Tensor: Integer indices or boolean mask
Sequence[int | bool]: List/tuple of indices or boolean mask

Returns:

New Mesh with subset of points. Cells that reference any removed points are also removed, and remaining cell indices are remapped.

Return type:

Mesh

Notes

The no-op selections None / Ellipsis return this mesh itself, and global_data is shared with the source by reference rather than copied. Mutating shared data on the result therefore also mutates the source; clone first if you need an independent copy.

Examples

>>> import torch
>>> from physicsnemo.mesh import Mesh
>>> # Create a mesh with 4 points and 2 triangular cells
>>> points = torch.tensor([[0., 0.], [1., 0.], [1., 1.], [0., 1.]])
>>> cells = torch.tensor([[0, 1, 2], [0, 2, 3]])
>>> mesh = Mesh(points=points, cells=cells)
>>> # Keep only points 0 and 2 - both cells are removed (they need points 1 or 3)
>>> sliced = mesh.slice_points([0, 2])
>>> sliced.n_points, sliced.n_cells
(2, 0)
>>> # Keep points 0, 1, 2 - first cell is preserved with remapped indices
>>> sliced = mesh.slice_points([0, 1, 2])
>>> sliced.n_points, sliced.n_cells
(3, 1)
>>> sliced.cells.tolist()
[[0, 1, 2]]

sobolev_deform( displacement: str | tuple[str, ...] | Float[torch.Tensor, 'n_points n_spatial_dims'], *, length_scale: float, fixed_points: str | tuple[str, ...] | Bool[torch.Tensor, ' n_points'] | None = None, max_iterations: int = 128, tolerance: float | None = None, implementation: Literal['torch', 'warp'] | None = None, ) → Mesh#

Deform a mesh with a uniform-mass P1 Sobolev displacement.

Filters a dense per-vertex displacement by solving

\[(M + \ell^2 K)u = M d\]

and returns a new mesh with points mesh.points + u. Here \(M=\bar m I\) is a uniform vertex mass matrix scaled by the mean positive lumped P1 mass. \(K\) is the P1 stiffness matrix, and \(\ell\) is length_scale in mesh coordinate units. The uniform mass makes the filter self-adjoint for PyTorch vertex tensors. Connectivity and attached fields are unchanged.

Call it as sobolev_deform(mesh, ...) or as mesh.sobolev_deform(...). The bound method supplies mesh automatically.

Parameters:

mesh (Mesh) – Simplicial mesh whose points are deformed. The source is not modified.
displacement (str, tuple[str, ...], or torch.Tensor) – Raw displacement with shape (mesh.n_points, mesh.n_spatial_dims), or a point_data key resolving to one. It must match the point dtype and device.
length_scale (float) – Nonnegative physical smoothing length. Zero applies the raw displacement directly at unfixed points.
fixed_points (str, tuple[str, ...], torch.Tensor, or None, optional) – Optional bool mask with shape (mesh.n_points,), or a point-data key resolving to one. True entries receive zero displacement. Default is None.
max_iterations (int, optional) – Maximum PCG iterations. Default is 128.
tolerance (float or None, optional) – Positive relative residual tolerance. None selects a dtype-dependent default. Default is None.
implementation ({"torch", "warp"} or None, optional) – Backend override. None selects Torch on CPU. On CUDA, it selects Warp for segments, triangles, and tetrahedra when available. It otherwise selects Torch, with a one-time RuntimeWarning when Warp is unavailable. The Warp backend requires CUDA tensors.

Returns:

New mesh with Sobolev-filtered points, unchanged connectivity, and unchanged attached fields.

Return type:

Mesh

Notes

Unfixed mesh boundaries use the natural homogeneous Neumann condition. Constant displacements are retained when no points are fixed. Isolated points receive their raw displacement.

Both backends participate in autograd through the source points and the raw displacement. Their reverse-mode derivatives solve the adjoint Helmholtz system, which makes the operation suitable for smooth vertex-based optimization. The Warp backend evaluates the geometry vector-Jacobian product analytically. A forward or adjoint solve that does not reach tolerance within max_iterations raises a RuntimeError. Warp supports segment, triangle, and tetrahedron cells. Higher-dimensional simplices use Torch by default. Warp CUDA results and point gradients may vary at roundoff between runs. CUDA Graph capture is not supported because P1 operator assembly and solver diagnostics are not capture-safe.

Geometry caches are invalidated and topology caches are retained. At positive length scales, cells must be finite, nondegenerate simplices. The operation does not detect inverted or self-intersecting output cells.

state_dict( destination=None, prefix='', keep_vars=False, flatten=True, ) → dict[str, Any]#

Returns a state_dict with logical keys, matching TensorDictBase conventions.

Tensor fields appear as data keys. Non-tensor fields (strings, ints, etc.) and the tensorclass type are stored in _metadata. This replaces the legacy _tensordict/_non_tensordict wrapper format.

property statistics: Mapping[str, int | float | tuple[float, float, float, float]]#

Compute summary statistics for the mesh.

Returns:: Mesh counts and distributions of edge lengths, cell measures, aspect ratios, and quality scores.
Return type:: Mapping

Examples

>>> from physicsnemo.mesh.primitives.basic import two_triangles_2d
>>> mesh = two_triangles_2d.load()
>>> stats = mesh.statistics
>>> assert "n_points" in stats and "n_cells" in stats

See also

physicsnemo.mesh.validation.compute_mesh_statistics: Standalone functional form, including a configurable degeneracy tolerance.

strip_caches() → Mesh[source]#

Return a new mesh with all 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:

Accurate benchmarking (prevents false performance benefits from caching)
Reducing memory usage
Forcing recomputation of cached values

Returns:: A new mesh with the same geometry and data, but without cached values.
Return type:: Mesh

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

subdivide( levels: int = 1, filter: Literal['linear', 'butterfly', 'loop'] = 'linear', ) → Mesh[source]#

Subdivide the mesh using iterative application of subdivision schemes.

Subdivision refines the mesh by splitting each n-simplex into 2^n child simplices. Multiple subdivision schemes are supported, each with different geometric and smoothness properties.

This method applies the chosen subdivision scheme iteratively for the specified number of levels. Each level independently subdivides the current mesh.

Parameters:

levels (int, optional) –
Number of subdivision iterations to perform. Each level increases mesh resolution exponentially:
- 0: No subdivision (returns original mesh)
- 1: Each cell splits into 2^n children
- 2: Each cell splits into 4^n children
- k: Each cell splits into (2^k)^n children
filter ({"linear", "butterfly", "loop"}, optional) –
Subdivision scheme to use:
- ”linear”: Simple midpoint subdivision (interpolating). New vertices at exact edge midpoints. Works for any dimension. Preserves original vertices.
- ”butterfly”: Weighted stencil subdivision (interpolating). New vertices use weighted neighbor stencils for smoother results. Currently only supports 2D manifolds (triangular meshes). Preserves original vertices.
- ”loop”: Valence-based subdivision (approximating). Both old and new vertices are repositioned for C² smoothness. Currently only supports 2D manifolds (triangular meshes). Original vertices move to new positions.

Returns:

Subdivided mesh with refined geometry and connectivity.

Manifold and spatial dimensions are preserved
Point data is interpolated to new vertices
Cell data is propagated from parents to children
Global data is preserved unchanged

Return type:

Mesh

Raises:

ValueError – If levels < 0 or if filter is not one of the supported schemes.
NotImplementedError – If butterfly/loop filter used with non-2D manifold.

Notes

Multi-level subdivision is achieved by iterative application. For levels=3, this is equivalent to calling subdivide(levels=1) three times in sequence. This is the standard approach for all subdivision schemes.

Examples

>>> from physicsnemo.mesh.primitives.basic import two_triangles_2d
>>> # Linear subdivision of triangular mesh
>>> mesh = two_triangles_2d.load()
>>> refined = mesh.subdivide(levels=2, filter="linear")
>>> # Each triangle splits into 4, twice: 2 -> 8 -> 32 triangles
>>> assert refined.n_cells == mesh.n_cells * 16

to_dual_graph() → Mesh[1, ...][source]#

Return a 1D Mesh representing the cell-adjacency (dual) graph.

Points are the cell centroids of this mesh. Cells are \((E, 2)\) line segments connecting pairs of cells that share a codimension-1 facet (e.g., cells sharing an edge in 2D or a face in 3D). The parent mesh’s cell_data becomes the point_data of the returned Mesh, since each dual-graph node corresponds to a parent cell.

Returns:: A 1D mesh (n_manifold_dims == 1) whose points are cell centroids and whose cells encode the cell-neighbor adjacency.
Return type:: Mesh[1, …]

Examples

>>> import torch
>>> from physicsnemo.mesh import Mesh
>>> # Two triangles sharing an edge
>>> points = torch.tensor([[0., 0.], [1., 0.], [0.5, 1.], [1.5, 1.]])
>>> cells = torch.tensor([[0, 1, 2], [1, 3, 2]])
>>> mesh = Mesh(points=points, cells=cells)
>>> dual = mesh.to_dual_graph()
>>> assert isinstance(dual, Mesh[1, ...])
>>> assert dual.n_cells == 1  # 1 shared edge -> 1 dual edge

to_edge_graph() → Mesh[1, ...][source]#

Return a 1D Mesh whose cells are the unique edges of this mesh.

Each edge (pair of vertices connected in a cell) appears exactly once. The resulting Mesh has the same points array, with cells of shape \((E, 2)\) where \(E\) is the number of unique edges.

Cell data from the parent mesh is aggregated onto edges via the facet extraction pipeline (mean aggregation by default).

Returns:: A 1D mesh (n_manifold_dims == 1) with edge cells.
Return type:: Mesh[1, …]

Examples

>>> import torch
>>> from physicsnemo.mesh import Mesh
>>> points = torch.tensor([[0., 0.], [1., 0.], [0.5, 1.]])
>>> cells = torch.tensor([[0, 1, 2]])
>>> mesh = Mesh(points=points, cells=cells)
>>> edge_graph = mesh.to_edge_graph()
>>> assert isinstance(edge_graph, Mesh[1, ...])
>>> assert edge_graph.n_cells == 3  # triangle has 3 edges

to_point_cloud( point_source: Literal['vertices', 'cell_centroids'] = 'vertices', ) → Mesh[0, ...][source]#

Return a 0D Mesh (point cloud) with no cell connectivity.

Parameters:

point_source ({"vertices", "cell_centroids"}) –

What becomes the points of the returned Mesh:

"vertices" (default): Uses mesh vertices as points, preserving point_data.
"cell_centroids": Uses cell centroids as points, mapping cell_data to point_data.

Returns:

A 0D mesh (n_manifold_dims == 0) with no cells.

Return type:

Mesh[0, …]

Examples

>>> import torch
>>> from physicsnemo.mesh import Mesh
>>> points = torch.tensor([[0., 0.], [1., 0.], [0.5, 1.]])
>>> cells = torch.tensor([[0, 1, 2]])
>>> mesh = Mesh(points=points, cells=cells)
>>> pc = mesh.to_point_cloud()
>>> assert isinstance(pc, Mesh[0, ...])
>>> assert pc.n_points == 3

to_tensordict( *, retain_none: bool | None = None, ) → TensorDict#

Convert the tensorclass into a regular TensorDict.

Makes a copy of all entries. Memmap and shared memory tensors are converted to regular tensors.

Parameters:: retain_none (bool) – if True, the None values will be written in the tensordict. Otherwise they will be discrarded. Default: True.
Returns:: A new TensorDict object containing the same values as the tensorclass.

transform( matrix: Float[Tensor, 'new_n_spatial_dims n_spatial_dims'], transform_point_data: bool | TensorDict = False, transform_cell_data: bool | TensorDict = False, transform_global_data: bool | TensorDict = False, assume_invertible: bool | None = None, ) → Mesh#

Apply a linear transformation to the mesh.

Call it as transform(mesh, ...) or as mesh.transform(...). The bound method supplies mesh automatically.

Parameters:

mesh (Mesh) – Input mesh to transform.
matrix (Float[torch.Tensor, "new_n_spatial_dims n_spatial_dims"]) – Transformation matrix, shape \((S', S)\).
transform_point_data (bool or TensorDict) – Controls transformation of point_data fields. True transforms all compatible fields; False transforms none; a TensorDict (or dict) with scalar bool leaves selectively transforms only the named fields.
transform_cell_data (bool or TensorDict) – Same semantics as transform_point_data, for cell_data.
transform_global_data (bool or TensorDict) – Same semantics as transform_point_data, for global_data.
assume_invertible (bool or None) –
Controls cache propagation for square matrices:
- True: assume matrix is invertible and propagate caches (compile-safe). This is a promise, not a check. If matrix is in fact singular, the inverse-transpose step silently yields non-finite values instead of raising, so the propagated normals and areas caches – and anything derived from them, such as the sum of cell_areas – come back as NaN. Use False or None unless you know the matrix is non-singular.
- False: assume matrix is singular and skip cache propagation (compile-safe). Caches are dropped and recomputed lazily on demand, which is always correct, just slower.
- None (default): test abs(det(matrix)) > 1e-10 at runtime and take one of the branches above. Safe for singular input, but the test reads a device scalar back to the host, which synchronizes on CUDA and may cause graph breaks under torch.compile.

Returns:

New Mesh with transformed geometry and appropriately updated caches.

Return type:

Mesh

Notes

Cache Handling:

areas: For square invertible matrices:

Full-dimensional meshes: scaled by |det|

Codimension-1 manifolds: per-element scaling using |det| * ||M^{-T} n||

Higher codimension: invalidated

centroids: Always transformed

normals: For square invertible matrices, transformed by inverse-transpose

translate( offset: Float[Tensor, 'n_spatial_dims'] | Sequence[float], ) → Mesh#

Apply a translation to the mesh.

Translation only affects point positions and centroids. Vector/tensor fields are unchanged by translation (they represent directions, not positions).

Call it as translate(mesh, ...) or as mesh.translate(...). The bound method supplies mesh automatically.

Parameters:

mesh (Mesh) – Input mesh to translate.
offset (Float[torch.Tensor, " n_spatial_dims"] or Sequence[float]) – Translation vector, shape \((S,)\).

Returns:

New Mesh with translated geometry.

Return type:

Mesh

Notes

Cache Handling:

areas: Unchanged

centroids: Translated

normals: Unchanged

unbind(dim: int)#

Returns a tuple of indexed tensorclass instances unbound along the indicated dimension.

Resulting tensorclass instances will share the storage of the initial tensorclass instance.

validate( check_degenerate_cells: bool = True, check_duplicate_vertices: bool = True, check_inverted_cells: bool = False, check_out_of_bounds: bool = True, check_manifoldness: bool = False, tolerance: float | None = None, raise_on_error: bool = False, *, check_self_intersection: bool = False, ) → Mapping[str, bool | int | Tensor]#

Validate mesh integrity and detect common errors.

Performs a comprehensive set of checks to ensure mesh is well-formed and suitable for geometric computations. Call it as validate(mesh, ...) or as mesh.validate(...). The bound method supplies mesh automatically.

Parameters:

mesh (Mesh) – Mesh to validate
check_degenerate_cells (bool) – Check for zero/negative area cells
check_duplicate_vertices (bool) – Check for coincident vertices within tolerance
check_inverted_cells (bool) – Check for cells with negative orientation (expensive)
check_out_of_bounds (bool) – Check that cell indices are valid
check_manifoldness (bool) – Check manifold topology (2D only, expensive)
tolerance (float | None) – Tolerance for geometric checks (areas, distances). If None (default), uses a dtype-aware epsilon via safe_eps().
raise_on_error (bool) – If True, raise ValueError on first error. If False, return dict with all validation results.
check_self_intersection (bool) – Request a self-intersection check. This option is keyword-only and not yet implemented; passing True raises NotImplementedError.

Returns:

Dictionary with validation results:

”valid”: bool, True if all enabled checks passed

”n_degenerate_cells”: int, number of degenerate cells found

”degenerate_cell_indices”: Tensor of indices (if any found)

”n_duplicate_vertices”: int, number of duplicate vertex pairs

”duplicate_vertex_pairs”: Tensor of index pairs (if any found)

”n_out_of_bounds_cells”: int, cells with invalid indices

”out_of_bounds_cell_indices”: Tensor of cell indices (if any)

”n_inverted_cells”: int (if check enabled)

”inverted_cell_indices”: Tensor (if check enabled and any found)

”is_manifold”: bool (if check enabled, 2D only)

”non_manifold_edges”: Tensor of edge indices (if check enabled)

Return type:

Mapping[str, bool | int | torch.Tensor]

Raises:

ValueError – If raise_on_error=True and validation fails
NotImplementedError – If check_self_intersection=True because that check is not yet implemented.

Examples

>>> from physicsnemo.mesh.primitives.basic import two_triangles_2d
>>> mesh = two_triangles_2d.load()
>>> report = validate(mesh)
>>> assert report["valid"] == True

Return a new mesh with selected field-data containers replaced.

Geometry and geometric/topological caches are preserved because points and cells do not change. Any data argument left as None is retained; pass an empty dictionary to clear that data association. The source mesh is not modified.

Parameters:

point_data (TensorDict or dict, optional) – Replacement per-point data. None retains the current data.
cell_data (TensorDict or dict, optional) – Replacement per-cell data. None retains the current data.
global_data (TensorDict or dict, optional) – Replacement mesh-level data. None retains the current data.

Returns:

New mesh sharing the immutable geometry tensors and cached geometry values, with independent TensorDict containers for data and cache entries.

Return type:

Mesh

Notes

The TensorDict containers are shallow-copied. Their tensor leaves are shared, matching PyTorch’s usual view-like replacement semantics and avoiding an unexpected copy of potentially large fields. Clone a field explicitly before passing it when independent tensor storage is required.

Examples

>>> updated = mesh.with_data(
...     point_data={"pressure": predicted_pressure},
... )
>>> cleared = updated.with_data(cell_data={})

DomainMesh#

The DomainMesh class groups an interior mesh with named boundary meshes and domain-level data. Operations such as morph() and radial_basis_function_deform() apply one consistent geometry change to every component and return a new domain.

class physicsnemo.mesh.domain_mesh.DomainMesh( interior: physicsnemo.mesh.mesh.Mesh, boundaries: dict[str, physicsnemo.mesh.mesh.Mesh] | tensordict._td.TensorDict | None = None, global_data: dict[str, torch.Tensor] | tensordict._td.TensorDict | None = None, *, batch_size, device=None, names=None, )[source]#

Bases: object

all_meshes() → Iterator[tuple[str, Mesh]][source]#

Iterate over all meshes in the domain.

Yields the interior mesh first (keyed "interior"), then each boundary mesh in sorted key order.

Yields:: tuple[str, Mesh] – (name, mesh) pairs. The first pair is always ("interior", self.interior).

Examples

>>> for name, mesh in dm.all_meshes():
...     print(f"{name}: {mesh.n_points} points")
interior: 100 points
inlet: 3 points
no_slip: 3 points

apply_to_meshes( fn: Callable[[Mesh], Mesh], *, interior: bool = True, boundaries: bool = True, ) → DomainMesh[source]#

Apply a Mesh-to-Mesh function to meshes in the domain.

By default, fn is called on the interior and on each boundary mesh. Use the keyword flags to apply selectively. Components that are skipped are cloned unchanged. Domain-level global_data is always cloned unchanged.

All built-in operations (translate, rotate, subdivide, clean, etc.) delegate here.

This is distinct from the inherited tensorclass apply(), which recursively maps a Tensor -> Tensor callable across every leaf tensor. Use apply() for tensor-level transforms (e.g. dtype casting) and apply_to_meshes() for mesh-level transforms.

Parameters:

fn (Callable[[Mesh], Mesh]) – A function that takes a Mesh and returns a Mesh.
interior (bool) – If True (default), apply fn to the interior mesh.
boundaries (bool) – If True (default), apply fn to every boundary mesh.

Returns:

New domain with the transformed meshes.

Return type:

DomainMesh

Examples

Convert every mesh to a point cloud (drop connectivity):

>>> dm_cloud = dm.apply_to_meshes(lambda m: Mesh(points=m.points))

Subdivide only the boundaries (e.g. to match a finer interior):

>>> dm2 = dm.apply_to_meshes(
...     lambda m: m.subdivide(levels=1), boundaries=True, interior=False
... )

property boundary_names: list[str]#

Sorted list of boundary condition names.

Returns:: The keys of boundaries, sorted alphabetically.
Return type:: list[str]

cell_data_to_point_data( overwrite_keys: bool = False, ) → DomainMesh[source]#

Convert cell data to point data on all meshes in the domain.

Delegates to Mesh.cell_data_to_point_data() for each mesh.

Parameters:: overwrite_keys (bool) – If True, silently overwrite existing point_data keys.
Returns:: New domain with converted data on all meshes.
Return type:: DomainMesh

clean( tolerance: float = 1e-12, merge_points: bool = True, remove_duplicate_cells: bool = True, remove_unused_points: bool = True, ) → DomainMesh[source]#

Clean and repair all meshes in the domain.

Delegates to Mesh.clean() for each mesh independently.

Parameters:

tolerance (float, optional) – L2 distance threshold for merging duplicate points.
merge_points (bool, optional) – Whether to merge spatially-duplicate points.
remove_duplicate_cells (bool, optional) – Whether to remove cells with identical vertex sets.
remove_unused_points (bool, optional) – Whether to drop points not referenced by any cell.

Returns:

New domain with cleaned meshes.

Return type:

DomainMesh

compute_cell_derivatives( keys: str | tuple[str, ...] | list[str | tuple[str, ...]] | None = None, method: Literal['lsq', 'dec'] = 'lsq', gradient_type: Literal['intrinsic', 'extrinsic', 'both'] = 'intrinsic', ) → DomainMesh[source]#

Compute gradients of cell_data fields on all meshes.

Delegates to Mesh.compute_cell_derivatives() for each mesh.

Parameters:

keys (str or tuple or list or None, optional) – Fields to differentiate. None for all non-cached fields.
method ({"lsq", "dec"}, optional) – Discretization method.
gradient_type ({"intrinsic", "extrinsic", "both"}, optional) – Type of gradient to compute.

Returns:

Domain with gradient fields added to each mesh’s cell_data.

Return type:

DomainMesh

compute_point_derivatives( keys: str | tuple[str, ...] | list[str | tuple[str, ...]] | None = None, method: Literal['lsq', 'dec'] = 'lsq', gradient_type: Literal['intrinsic', 'extrinsic', 'both'] = 'intrinsic', ) → DomainMesh[source]#

Compute gradients of point_data fields on all meshes.

Delegates to Mesh.compute_point_derivatives() for each mesh.

Parameters:

keys (str or tuple or list or None, optional) – Fields to differentiate. None for all non-cached fields.
method ({"lsq", "dec"}, optional) – Discretization method.
gradient_type ({"intrinsic", "extrinsic", "both"}, optional) – Type of gradient to compute.

Returns:

Domain with gradient fields added to each mesh’s point_data.

Return type:

DomainMesh

property device: device#: Retrieves the device type of tensor class.

draw( *, backend: Literal['matplotlib', 'pyvista', 'auto'] = 'auto', show: bool = True, point_scalars: None | Tensor | str | tuple[str, ...] = None, cell_scalars: None | Tensor | str | tuple[str, ...] = None, cmap: str = 'viridis', vmin: float | None = None, vmax: float | None = None, alpha_points: float = 1.0, alpha_cells: float = 1.0, alpha_edges: float = 1.0, show_edges: bool = False, boundary_kwargs: dict[str, Any] | None = None, ax: matplotlib.axes.Axes | pyvista.Plotter | None = None, backend_options: dict[str, Any] | None = None, ) → matplotlib.axes.Axes | pyvista.Plotter[source]#

Draw the domain: interior with optional scalar coloring, boundaries overlaid.

Renders the interior as the primary visual layer, then overlays every boundary on the same canvas. The interior parameter set mirrors Mesh.draw() exactly, with two intentional changes: the call is keyword-only, and show_edges defaults to False (rather than True) because dense interior meshes are typically more readable without edges. Both matplotlib and PyVista backends are supported.

Parameters:

backend – Forwarded to Mesh.draw() for the interior mesh. See Mesh.draw() for full descriptions.
show – Forwarded to Mesh.draw() for the interior mesh. See Mesh.draw() for full descriptions.
point_scalars – Forwarded to Mesh.draw() for the interior mesh. See Mesh.draw() for full descriptions.
cell_scalars – Forwarded to Mesh.draw() for the interior mesh. See Mesh.draw() for full descriptions.
cmap – Forwarded to Mesh.draw() for the interior mesh. See Mesh.draw() for full descriptions.
vmin – Forwarded to Mesh.draw() for the interior mesh. See Mesh.draw() for full descriptions.
vmax – Forwarded to Mesh.draw() for the interior mesh. See Mesh.draw() for full descriptions.
alpha_points – Forwarded to Mesh.draw() for the interior mesh. See Mesh.draw() for full descriptions.
alpha_cells – Forwarded to Mesh.draw() for the interior mesh. See Mesh.draw() for full descriptions.
alpha_edges – Forwarded to Mesh.draw() for the interior mesh. See Mesh.draw() for full descriptions.
show_edges – Forwarded to Mesh.draw() for the interior mesh. See Mesh.draw() for full descriptions.
ax – Forwarded to Mesh.draw() for the interior mesh. See Mesh.draw() for full descriptions.
backend_options – Forwarded to Mesh.draw() for the interior mesh. See Mesh.draw() for full descriptions.
boundary_kwargs (dict, optional) –
Keyword arguments forwarded to Mesh.draw() for every boundary mesh. Defaults are tuned for unobtrusive overlay:
- alpha_points = 0 (boundary vertices are not scattered).
- alpha_cells = 0.3 when boundaries are 2-D surfaces, 1.0 when they are 1-D curves. Auto-detected from the first boundary’s Mesh.n_manifold_dims.
- show_edges = False.
User-supplied keys override these defaults. To color individual boundaries by their own scalar fields, compose Mesh.draw() calls directly (see Examples).

Returns:

The canvas, for further customization when show=False.

Return type:

matplotlib.axes.Axes or pyvista.Plotter

Examples

Default visualization with pressure coloring on the interior:

>>> dm.draw(point_scalars="p", cmap="RdBu_r", vmin=-200, vmax=200)

Translucent boundaries with edges visible:

>>> dm.draw(
...     point_scalars="p",
...     boundary_kwargs={"alpha_cells": 0.5, "show_edges": True},
... )

Customize and display later by setting axis limits on the returned canvas:

>>> ax = dm.draw(point_scalars="p", show=False)
>>> ax.set_xlim(-2, 4); ax.set_ylim(-3, 3)

Per-boundary scalar coloring (manual composition - color the no-slip wall by its own shear field while the interior shows pressure):

>>> ax = dm.interior.draw(point_scalars="p", show=False)
>>> dm.boundaries["wall"].draw(
...     ax=ax, cell_scalars="shear", cmap="hot", show=False,
... )
>>> for name in dm.boundary_names:
...     if name == "wall":
...         continue
...     dm.boundaries[name].draw(
...         ax=ax, alpha_cells=0.3, alpha_points=0,
...         show_edges=False, show=False,
...     )
>>> import matplotlib.pyplot as plt; plt.show()

dumps( prefix: str | None = None, copy_existing: bool = False, *, num_threads: int = 0, return_early: bool = False, share_non_tensor: bool = False, robust_key: bool | None = True, archive: bool | None = None, compression: str | int | None = None, ) → Any#

Saves the tensordict to disk.

This function is a proxy to memmap().

classmethod fields()#

Return a tuple describing the fields of this dataclass.

Accepts a dataclass or an instance of one. Tuple elements are of type Field.

free_form_deform( control_displacements: Float[Tensor, '*lattice_resolution n_spatial_dims'], *, origin: Float[Tensor, 'n_spatial_dims'] | Sequence[float] | None = None, extent: Float[Tensor, 'n_spatial_dims'] | Sequence[float] | None = None, basis: Literal['bernstein', 'bspline', 'linear', 'cubic_hermite', 'quintic_hermite'] = 'bernstein', point_weights: str | tuple[str, ...] | None = None, implementation: Literal['torch', 'warp'] | None = None, ) → DomainMesh[source]#

Deform the interior and all boundaries with one lattice field.

Every component uses the same control lattice, box, basis, and backend. With point_weights=None, coincident interior and boundary points receive the same motion. When supplied, point_weights is a common Mesh.point_data key (or nested tuple key) resolved independently on each component. Raw point-weight tensors are rejected because component point counts differ.

Parameters:

control_displacements (torch.Tensor) – Displacement vectors, not destination coordinates, for every lattice node, with shape (n_1, ..., n_D, n_spatial_dims) and the same float32 or float64 dtype and device as every component’s points. Each axis needs at least two nodes for "bernstein" and the node-interpolating bases, and four for "bspline".
origin (torch.Tensor, sequence of float, or None, optional) – Minimum corner of the lattice box with shape (n_spatial_dims,). None uses the minimum corner of the combined component bounds. For repeated GPU calls with an explicit box, create origin and extent once as device tensors. Reuse them to avoid recreating and transferring sequence values.
extent (torch.Tensor, sequence of float, or None, optional) – Edge lengths of the lattice box. Every value must be finite and strictly positive. The operation does not validate tensor values at runtime. None sizes the box from origin to the maximum corner of the combined component bounds. Validating a derived extent synchronizes with the device and is not CUDA Graph capture-safe. For capture, pass both origin and extent as device tensors. Every coordinate axis must have positive range when the extent is derived. Otherwise, supply an explicit extent.
basis ({"bernstein", "bspline", "linear", "cubic_hermite", "quintic_hermite"}, optional) – Per-axis basis family. "bernstein" provides global support. "bspline" uses local four-node-per-axis support. B-spline coefficient index i corresponds to local coordinate (i - 1) / (n - 3). The first and last coefficient planes lie outside the evaluation box. "linear", "cubic_hermite", and "quintic_hermite" use two neighboring nodes per axis. The resulting fields are C0, C1, and C2 across cell boundaries, respectively. See free_form_deform() for their polynomial weights and literature reference. Default is "bernstein".
point_weights (str, tuple[str, ...], or None) – Optional point-data key present in every component and resolved independently on each component. Each resolved tensor must have shape (component.n_points,) and match the component point device. All components must use one common bool or floating dtype. Floating weights must also match the point dtype. Raw tensors are not accepted.
implementation ({"torch", "warp"} or None) – Backend override. Automatic dispatch uses Torch on CPU. On CUDA, it uses Warp when available and otherwise Torch.

Returns:

New domain with deformed component meshes and unchanged domain data.

Return type:

DomainMesh

Raises:

TypeError – If tensors, lattice values, or point weights have unsupported types or dtypes.
ValueError – If component layouts, lattice parameters, point weights, or basis are invalid.
KeyError – If a point-data key or implementation name is not found.
ImportError – If an explicitly requested backend is unavailable.

Notes

The operation retains connectivity and attached mesh and domain data. It treats attached vector and tensor fields as Lagrangian data and does not push them forward. It invalidates geometry caches and retains topology caches on each component. Points outside the lattice box are unchanged. To keep the exterior fixed, zero the outermost coefficient plane on every Bernstein or node-interpolating face. For cubic B-splines, zero the first and last three coefficient planes on every axis. origin and extent are non-differentiable lattice parameters. The operation does not automatically detect inverted, degenerate, or self-intersecting cells. Validate each component mesh explicitly with Mesh.validate() when required.

Creates a TensorDict from a CSV file.

Requires either pandas or pyarrow to be installed.

Parameters:

path (str or Path) – Path to the CSV file.

Keyword Arguments:

auto_batch_size (bool, optional) – If True, the batch size will be computed automatically. Defaults to False.
batch_dims (int, optional) – If auto_batch_size is True, defines how many dimensions the output tensordict should have. Defaults to None.
device (torch.device, optional) – The device for tensor data. Defaults to None.
batch_size (torch.Size, optional) – The batch size. Defaults to [num_rows].
separator (str, optional) – If provided, column names are split on this separator to create nested TensorDicts. Defaults to None.
dtype (torch.dtype, optional) – If provided, all numeric columns are cast to this dtype. Defaults to None.
**kwargs – Additional keyword arguments forwarded to the CSV reader (pandas.read_csv or pyarrow.csv.read_csv).

Returns:

A TensorDict representation of the CSV data.

Examples

>>> td = TensorDict.from_csv("data.csv")
>>> td = TensorDict.from_csv("data.csv", separator=".", dtype=torch.float32)

Creates a TensorDict from a JSON file.

Supports both standard JSON (array of records) and JSON Lines format. For nested JSON objects, use from_dict() instead.

Requires pandas for best results. Falls back to stdlib json for simple cases.

Parameters:

path (str or Path) – Path to the JSON file.

Keyword Arguments:

auto_batch_size (bool, optional) – If True, the batch size will be computed automatically. Defaults to False.
batch_dims (int, optional) – If auto_batch_size is True, defines how many dimensions the output tensordict should have. Defaults to None.
device (torch.device, optional) – The device for tensor data. Defaults to None.
batch_size (torch.Size, optional) – The batch size. Defaults to [num_rows].
separator (str, optional) – If provided, column names are split on this separator to create nested TensorDicts. Defaults to None.
dtype (torch.dtype, optional) – If provided, all numeric columns are cast to this dtype. Defaults to None.
lines (bool, optional) – If True, reads the file as JSON Lines (one JSON object per line). Defaults to False.
**kwargs – Additional keyword arguments forwarded to the JSON reader.

Returns:

A TensorDict representation of the JSON data.

Examples

>>> td = TensorDict.from_json("data.json")
>>> td = TensorDict.from_json("data.jsonl", lines=True)

Converts a pandas DataFrame to a TensorDict.

Numeric columns become tensors, string/object columns become NonTensorData.

Parameters:

dataframe (pd.DataFrame) – The pandas DataFrame to convert.

Keyword Arguments:

auto_batch_size (bool, optional) – If True, the batch size will be computed automatically. Defaults to False.
batch_dims (int, optional) – If auto_batch_size is True, defines how many dimensions the output tensordict should have. Defaults to None.
device (torch.device, optional) – The device for tensor data. Defaults to None.
batch_size (torch.Size, optional) – The batch size. Defaults to [num_rows].
separator (str, optional) – If provided, column names are split on this separator to create nested TensorDicts. For example, with separator=".", a column "obs.x" becomes td["obs", "x"]. Defaults to None.
dtype (torch.dtype, optional) – If provided, all numeric columns are cast to this dtype. Defaults to None.

Returns:

A TensorDict representation of the DataFrame.

Examples

>>> import pandas as pd
>>> df = pd.DataFrame({"a": [1, 2, 3], "b": [4.0, 5.0, 6.0]})
>>> td = TensorDict.from_pandas(df)
>>> print(td)
TensorDict(
    fields={
        a: Tensor(shape=torch.Size([3]), device=cpu, dtype=torch.int64, is_shared=False),
        b: Tensor(shape=torch.Size([3]), device=cpu, dtype=torch.float64, is_shared=False)},
    batch_size=torch.Size([3]),
    device=None,
    is_shared=False)

Creates a TensorDict from a Parquet file.

Requires either pyarrow or pandas to be installed. Prefers pyarrow when available for better performance.

Parameters:

path (str or Path) – Path to the Parquet file.

Keyword Arguments:

auto_batch_size (bool, optional) – If True, the batch size will be computed automatically. Defaults to False.
batch_dims (int, optional) – If auto_batch_size is True, defines how many dimensions the output tensordict should have. Defaults to None.
device (torch.device, optional) – The device for tensor data. Defaults to None.
batch_size (torch.Size, optional) – The batch size. Defaults to [num_rows].
separator (str, optional) – If provided, column names are split on this separator to create nested TensorDicts. Defaults to None.
dtype (torch.dtype, optional) – If provided, all numeric columns are cast to this dtype. Defaults to None.
columns (list of str, optional) – If provided, only read these columns from the file. Defaults to None (all columns).
**kwargs – Additional keyword arguments forwarded to the Parquet reader.

Returns:

A TensorDict representation of the Parquet data.

Examples

>>> td = TensorDict.from_parquet("data.parquet")
>>> td = TensorDict.from_parquet("data.parquet", columns=["obs", "reward"])

from_schema(

*,

batch_size: Sequence[int] | Size | None = None,

storage: str | None = None,

device=None,

**kwargs,

) → TensorDictBase#

Pre-allocate a zero-filled TensorDict from a schema.

Creates a TensorDictBase whose storage backend is selected by storage. Each entry in schema maps a field name to an (element_shape, dtype) pair; the full stored shape is [*batch_size, *element_shape].

Parameters:

schema – Mapping from field name to (element_shape, dtype). element_shape is the per-element shape (excluding batch_size).

Keyword Arguments:

batch_size – Overall batch dimensions prepended to every element shape. Defaults to ().
storage (str or None) –
Backend selector:
- None – plain TensorDict with regular tensors.
- "memmap" – memory-mapped tensors on disk. Pass prefix=<dir> in kwargs.
- "h5" – HDF5 via PersistentTensorDict. Pass filename=<path> in kwargs.
- "zarr" – zarr (requires zarr>=3.0) via PersistentTensorDict. Pass filename=<path or store> in kwargs.
- "shared" – CPU shared-memory tensors.
- "redis" / "dragonfly" – delegates to TensorDictStore.from_schema().
device – Device for the resulting tensors (ignored by some backends).
**kwargs – Backend-specific arguments forwarded to the underlying constructor (e.g. prefix for memmap, filename for h5, host/port for redis).

Returns:

A new TensorDictBase subclass instance with pre-allocated (zero-filled) keys.

Examples

>>> td = TensorDict.from_schema(
...     {"obs": ([84, 84, 3], torch.uint8),
...      "reward": ([], torch.float32)},
...     batch_size=[1000],
... )
>>> td["obs"].shape
torch.Size([1000, 84, 84, 3])

>>> import tempfile
>>> with tempfile.TemporaryDirectory() as d:
...     td_mm = TensorDict.from_schema(
...         {"obs": ([4], torch.float32)},
...         batch_size=[8],
...         storage="memmap",
...         prefix=d,
...     )
...     assert td_mm.is_memmap()

classmethod from_tensordict( tensordict: TensorDictBase, non_tensordict: dict | None = None, safe: bool = True, ) → Any#

Tensor class wrapper to instantiate a new tensor class object.

Parameters:

tensordict (TensorDictBase) – Dictionary of tensor types
non_tensordict (dict) – Dictionary with non-tensor and nested tensor class objects
safe (bool) – Whether to raise an error if the tensordict is not a TensorDictBase instance

get(key: NestedKey, *args, **kwargs)#

Gets the value stored with the input key.

Parameters:

key (str, tuple of str) – key to be queried. If tuple of str it is equivalent to chained calls of getattr.
default – default value if the key is not found in the tensorclass.

Returns:

value stored with the input key

is_boundary_watertight(tolerance: float = 1e-06) → bool[source]#

Check whether the merged boundary meshes form a watertight surface.

Merges all boundary meshes via merge_boundaries(), deduplicates coincident vertices with Mesh.clean(), and calls Mesh.is_watertight() on the result. The clean step is necessary because independently-meshed boundary patches share physical vertices that become duplicated during merge - and float32 round-off from any prior transform may prevent an exact-match merge.

Parameters:: tolerance (float, optional) – L2 distance threshold for merging coincident boundary vertices before the topology check. The default 1e-6 is deliberately looser than Mesh.clean()’s 1e-12 so it absorbs float32 round-off (~1e-7 relative) on the duplicated vertices that merge_boundaries produces from independently-meshed patches. For coordinates that span much smaller or much larger than ~1, pass an explicit value (e.g. 1e-6 * max_extent of the bbox).
Returns:: True if the merged boundary surface is watertight (every codimension-1 facet is shared by exactly 2 cells), False otherwise. Returns False if there are no boundary meshes.
Return type:: bool

Notes

This is not free to compute: the Mesh.clean() step performs a BVH-based duplicate-point merge that scales as \(O(N \log N)\) in the total boundary vertex count \(N\), and dominates the runtime. Callers that need the result repeatedly should cache it.

classmethod load(

prefix: str | Path,

*args,

**kwargs,

) → Any#

Loads a tensordict from disk.

This class method is a proxy to load_memmap().

load_(prefix: str | Path, *args, **kwargs)#

Loads a tensordict from disk within the current tensordict.

This class method is a proxy to load_memmap_().

classmethod load_memmap( prefix: str | Path, device: device | None = None, non_blocking: bool = False, *, out: TensorDictBase | None = None, robust_key: bool | None = True, subpath: NestedKey | None = None, mode: str = 'r', num_threads: int = 0, allow_pickle: bool | None = None, ) → Any#

Loads a memory-mapped tensordict from disk.

Parameters:

prefix (str or Path to folder) – the path to the folder where the saved tensordict should be fetched, or the path to a memmap archive file written through save(..., archive=True) / a ".tdz" prefix (or packed with pack_memmap()). Archives are memory-mapped once and every leaf is exposed as a zero-copy view into the mapping: only the pages of the leaves that are actually accessed are read from disk. Unlike directory-backed tensordicts, in-place writes to the leaves of an archive-loaded tensordict do not propagate to the file.
device (torch.device or equivalent, optional) – if provided, the data will be asynchronously cast to that device. Supports “meta” device, in which case the data isn’t loaded but a set of empty “meta” tensors are created. This is useful to get a sense of the total model size and structure without actually opening any file.
non_blocking (bool, optional) – if True, synchronize won’t be called after loading tensors on device. Defaults to False.
out (TensorDictBase, optional) – optional tensordict where the data should be written.
robust_key (bool, optional) – if True (default), expects robust key encoding was used when saving and decodes filenames accordingly. If False, uses legacy behavior. If None, uses the default robust behavior.
subpath (NestedKey or str path, optional) – the location of a nested tensordict to load, as a nested key (e.g. ("module", "0"), with arbitrary nesting allowed as usual) or as a "/"-separated string path (e.g. "module/0"). Only that subtree is loaded. Works both for directories (equivalent to appending the path to prefix) and for archives.
mode (str, optional) – "r" (default) or "r+". Only relevant when loading an archive: with "r" the archive is mapped copy-on-write and in-place writes to the leaves stay in memory; with "r+" the mapping is shared and in-place writes propagate to the file, like directory-backed tensordicts. "r+" requires uncompressed, aligned tensor payloads (i.e. archives written by tensordict without compression) and is not available for nested-tensor leaves. In-place writes do not update the per-entry CRC-32 stored by the zip format; load_memmap() ignores checksums, but call refresh_archive_checksums() before handing a modified archive to tools that verify them (unzip, unpack_memmap(), …). Directory prefixes are always write-through and ignore this argument.
num_threads (int, optional) – number of threads used to decompress the leaves of a compressed archive (deflate entries are inflated in parallel, which scales nearly linearly). Without compression, loading is a metadata-only operation and this argument has no effect. Defaults to 0 (sequential).
allow_pickle (bool, optional) – whether pickled non-tensor fields may be loaded. Pickle can execute arbitrary code, so pass True only for data from a trusted source and False for untrusted data. During the 0.14 compatibility window, omitting this option loads pickle with a FutureWarning; the default will change to False in 0.15. Saves without a pickle sidecar do not require this option.

Examples

>>> from tensordict import TensorDict
>>> td = TensorDict.fromkeys(["a", "b", "c", ("nested", "e")], 0)
>>> td.memmap("./saved_td")
>>> td_load = TensorDict.load_memmap("./saved_td")
>>> assert (td == td_load).all()

This method also allows loading nested tensordicts.

Examples

>>> nested = TensorDict.load_memmap("./saved_td/nested")
>>> assert nested["e"] == 0

A tensordict can also be loaded on “meta” device or, alternatively, as a fake tensor.

Examples

>>> import tempfile
>>> td = TensorDict({"a": torch.zeros(()), "b": {"c": torch.zeros(())}})
>>> with tempfile.TemporaryDirectory() as path:
...     td.save(path)
...     td_load = TensorDict.load_memmap(path, device="meta")
...     print("meta:", td_load)
...     from torch._subclasses import FakeTensorMode
...     with FakeTensorMode():
...         td_load = TensorDict.load_memmap(path)
...         print("fake:", td_load)
meta: TensorDict(
    fields={
        a: Tensor(shape=torch.Size([]), device=meta, dtype=torch.float32, is_shared=False),
        b: TensorDict(
            fields={
                c: Tensor(shape=torch.Size([]), device=meta, dtype=torch.float32, is_shared=False)},
            batch_size=torch.Size([]),
            device=meta,
            is_shared=False)},
    batch_size=torch.Size([]),
    device=meta,
    is_shared=False)
fake: TensorDict(
    fields={
        a: FakeTensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, is_shared=False),
        b: TensorDict(
            fields={
                c: FakeTensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, is_shared=False)},
            batch_size=torch.Size([]),
            device=cpu,
            is_shared=False)},
    batch_size=torch.Size([]),
    device=cpu,
    is_shared=False)

load_state_dict( state_dict: dict[str, Any], strict=True, assign=False, from_flatten=None, )#

Loads a state_dict into the tensorclass.

Supports both the new format (logical keys with _metadata) and the legacy format (_tensordict/_non_tensordict wrapper keys).

memmap( prefix: str | None = None, copy_existing: bool = False, *, num_threads: int = 0, return_early: bool = False, share_non_tensor: bool = False, existsok: bool = True, robust_key: bool | None = True, archive: bool | None = None, compression: str | int | None = None, ) → Any#

Writes all tensors onto a corresponding memory-mapped Tensor in a new tensordict.

Parameters:

prefix (str) – directory prefix where the memory-mapped tensors will be stored. The directory tree structure will mimic the tensordict’s. If prefix ends with ".tdz" (or archive=True is passed), a single-file archive is written instead of a directory: a standard zip file whose entries replicate the memmap directory layout. See archive below.
copy_existing (bool) – If False (default), an exception will be raised if an entry in the tensordict is already a tensor stored on disk with an associated file, but is not saved in the correct location according to prefix. If True, any existing Tensor will be copied to the new location.

Keyword Arguments:

num_threads (int, optional) – the number of threads used to write the memmap tensors. Defaults to 0.
return_early (bool, optional) – if True and num_threads>0, the method will return a future of the tensordict.
share_non_tensor (bool, optional) – if True, the non-tensor data will be shared between the processes and writing operation (such as inplace update or set) on any of the workers within a single node will update the value on all other workers. If the number of non_tensor leaves is high (e.g., sharing large stacks of non-tensor data) this may result in OOM or similar errors. Defaults to False.
existsok (bool, optional) – if False, an exception will be raised if a tensor already exists in the same path. Defaults to True.
robust_key (bool, optional) – if True (default), uses robust key encoding that safely handles keys with path separators and special characters. If False, uses legacy behavior (keys used as-is). If None, uses the default robust behavior.
archive (bool, optional) – if True, prefix designates a single file rather than a directory and the tensordict is written as a memmap archive: a zip file mirroring the memmap directory tree, with tensor payloads stored uncompressed and aligned so that load_memmap() can memory-map the file and expose every leaf as a zero-copy view. If None (default), archive mode is enabled when prefix ends with ".tdz". The result of load_memmap() on an archive behaves like the result of from_consolidated(): all leaves are views into a single storage, and in-place writes do not propagate to the file. Archives and memmap directories are mutually convertible with pack_memmap() / unpack_memmap() (or any zip tool). Note that archives are written sequentially (single data pass) and num_threads has no effect on them.
compression (str or int, optional) – compression for archive entries ("stored", "deflate", "bzip2", "lzma" or a zipfile constant). Defaults to "stored" (uncompressed), which is what enables zero-copy loading. Compressed archives load correctly but leaves are decompressed in memory on access. Only valid in archive mode.

The TensorDict is then locked, meaning that any writing operations that isn’t in-place will throw an exception (eg, rename, set or remove an entry). Once the tensordict is unlocked, the memory-mapped attribute is turned to False, because cross-process identity is not guaranteed anymore.

Returns:: A new tensordict with the tensors stored on disk if return_early=False, otherwise a TensorDictFuture instance.

Note

Serialising in this fashion might be slow with deeply nested tensordicts, so it is not recommended to call this method inside a training loop.

memmap_( prefix: str | None = None, copy_existing: bool = False, *, num_threads: int = 0, return_early: bool = False, share_non_tensor: bool = False, existsok: bool = True, robust_key: bool | None = True, ) → Any#

Writes all tensors onto a corresponding memory-mapped Tensor, in-place.

Parameters:

prefix (str) – directory prefix where the memory-mapped tensors will be stored. The directory tree structure will mimic the tensordict’s.
copy_existing (bool) – If False (default), an exception will be raised if an entry in the tensordict is already a tensor stored on disk with an associated file, but is not saved in the correct location according to prefix. If True, any existing Tensor will be copied to the new location.

Keyword Arguments:

num_threads (int, optional) – the number of threads used to write the memmap tensors. Defaults to 0.
return_early (bool, optional) – if True and num_threads>0, the method will return a future of the tensordict. The resulting tensordict can be queried using future.result().
share_non_tensor (bool, optional) – if True, the non-tensor data will be shared between the processes and writing operation (such as inplace update or set) on any of the workers within a single node will update the value on all other workers. If the number of non-tensor leaves is high (e.g., sharing large stacks of non-tensor data) this may result in OOM or similar errors. Defaults to False.
existsok (bool, optional) – if False, an exception will be raised if a tensor already exists in the same path. Defaults to True.
robust_key (bool, optional) – if True (default), uses robust key encoding that safely handles keys with path separators and special characters. If False, uses legacy behavior (keys used as-is). If None, uses the default robust behavior.

The TensorDict is then locked, meaning that any writing operations that isn’t in-place will throw an exception (eg, rename, set or remove an entry). Once the tensordict is unlocked, the memory-mapped attribute is turned to False, because cross-process identity is not guaranteed anymore.

Returns:: self if return_early=False, otherwise a TensorDictFuture instance.

Note

Serialising in this fashion might be slow with deeply nested tensordicts, so it is not recommended to call this method inside a training loop.

memmap_like( prefix: str | None = None, copy_existing: bool = False, *, existsok: bool = True, num_threads: int = 0, return_early: bool = False, share_non_tensor: bool = False, robust_key: bool | None = True, archive: bool | None = None, ) → Any#

Creates a contentless Memory-mapped tensordict with the same shapes as the original one.

Parameters:

prefix (str) – directory prefix where the memory-mapped tensors will be stored. The directory tree structure will mimic the tensordict’s. If prefix ends with ".tdz" (or archive=True is passed), a preallocated single-file archive is created instead. See archive below.
copy_existing (bool) – If False (default), an exception will be raised if an entry in the tensordict is already a tensor stored on disk with an associated file, but is not saved in the correct location according to prefix. If True, any existing Tensor will be copied to the new location.

Keyword Arguments:

num_threads (int, optional) – the number of threads used to write the memmap tensors. Defaults to 0.
return_early (bool, optional) – if True and num_threads>0, the method will return a future of the tensordict.
share_non_tensor (bool, optional) – if True, the non-tensor data will be shared between the processes and writing operation (such as inplace update or set) on any of the workers within a single node will update the value on all other workers. If the number of non-tensor leaves is high (e.g., sharing large stacks of non-tensor data) this may result in OOM or similar errors. Defaults to False.
existsok (bool, optional) – if False, an exception will be raised if a tensor already exists in the same path. Defaults to True.
robust_key (bool, optional) – if True (default), uses robust key encoding that safely handles keys with path separators and special characters. If False, uses legacy behavior (keys used as-is). If None, uses the default robust behavior.
archive (bool, optional) – if True, prefix designates a single file and a preallocated, zero-filled memmap archive is created and loaded back with load_memmap(prefix, mode="r+"): the returned tensordict writes through to the archive. If None (default), archive mode is enabled when prefix ends with ".tdz". In-place writes leave the zip per-entry checksums stale; call refresh_archive_checksums() before handing the archive to tools that verify them. Nested tensors are not supported in this mode.

The TensorDict is then locked, meaning that any writing operations that isn’t in-place will throw an exception (eg, rename, set or remove an entry). Once the tensordict is unlocked, the memory-mapped attribute is turned to False, because cross-process identity is not guaranteed anymore.

Returns:: A new TensorDict instance with data stored as memory-mapped tensors if return_early=False, otherwise a TensorDictFuture instance.

Note

This is the recommended method to write a set of large buffers on disk, as memmap_() will copy the information, which can be slow for large content.

Examples

>>> td = TensorDict({
...     "a": torch.zeros((3, 64, 64), dtype=torch.uint8),
...     "b": torch.zeros(1, dtype=torch.int64),
... }, batch_size=[]).expand(1_000_000)  # expand does not allocate new memory
>>> buffer = td.memmap_like("/path/to/dataset")

memmap_refresh_(*, allow_pickle: bool | None = None)#

Refreshes the content of the memory-mapped tensordict if it has a saved_path.

This method will raise an exception if no path is associated with it.

Parameters:: allow_pickle (bool, optional) – whether pickled non-tensor fields may be loaded. See load_memmap().

merge_boundaries( preserve_data: bool = False, ) → Mesh[source]#

Merge all boundary meshes into a single Mesh.

Produces a mesh containing the concatenated points and cells from every boundary. By default, point_data and cell_data are stripped before merging because boundaries typically carry heterogeneous fields (different keys per boundary), which Mesh.merge() cannot concatenate.

Parameters:: preserve_data (bool) – If False (default), strip point_data and cell_data from each boundary before merging - the safe choice for the typical CFD case where each boundary carries its own field set. If True, delegate directly to Mesh.merge(), which preserves data but requires that all boundaries share the same cell_data keys and have point_data that can be concatenated. Use this when every boundary has a consistent set of fields.
Returns:: A single mesh containing the concatenated points and cells from every boundary. Data fields are included only if preserve_data is True.
Return type:: Mesh
Raises:: ValueError – If there are no boundary meshes to merge, if boundary meshes have incompatible manifold dimensions, or (when preserve_data=True) if their data keys are inconsistent.

morph( control_points: Tensor, control_displacements: Tensor, *, radius: float | Tensor, point_weights: str | tuple[str, ...] | None = None, kernel: Literal['wendland_c2'] = 'wendland_c2', implementation: Literal['torch', 'warp'] | None = None, ) → DomainMesh[source]#

Morph the interior and all boundaries with one world-space field.

The same control coordinates, displacements, radii, and backend are used for every component, so coincident interior/boundary points receive the same motion when point_weights is None. When supplied, point_weights is a common Mesh.point_data key (or nested tuple key) resolved on each component independently; raw point-weight tensors are intentionally rejected because component point counts differ. A common key does not require equal values: coincident component points remain coincident only when their resolved point weights also match.

Parameters:

control_points (torch.Tensor) – World-coordinate controls with shape (n_controls, n_spatial_dims) and the same float32 or float64 dtype and device as every component’s points.
control_displacements (torch.Tensor) – Displacement vectors, not destination coordinates, with the same shape, dtype, and device as control_points.
radius (float or torch.Tensor) – Support distance in domain coordinate units. Supply a scalar or one radius per control. A tensor radius must match the control dtype and device; every value must remain positive and finite but is not validated at runtime.
point_weights (str, tuple[str, ...], or None) – Optional point-data key present in every component and resolved independently on each component. Resolved tensors must have one common dtype; floating-point weights match the component point dtype. Raw tensors are not accepted.
kernel ({"wendland_c2"}, optional) – Compact radial kernel used to blend control displacements. Default is "wendland_c2".
implementation ({"torch", "warp"} or None) – Backend override. Auto dispatch uses Torch on CPU and Warp on CUDA when Warp is available, otherwise Torch.

Returns:

New domain with morphed component meshes and unchanged domain data.

Return type:

DomainMesh

Notes

Connectivity and attached mesh and domain data are retained. Attached vector and tensor fields are treated as Lagrangian data and are not pushed forward. Geometry caches are invalidated and topology caches are retained on each component. Parameterize learned radii to remain positive, for example as torch.nn.functional.softplus(raw_radius) + eps. Morphing does not automatically detect inverted, degenerate, or self-intersecting cells. Use each component mesh’s Mesh.validate() method explicitly when required.

property n_boundaries: int#

Number of boundary meshes.

Returns:: The number of entries in boundaries.
Return type:: int

point_data_to_cell_data( overwrite_keys: bool = False, ) → DomainMesh[source]#

Convert point data to cell data on all meshes in the domain.

Delegates to Mesh.point_data_to_cell_data() for each mesh.

Parameters:: overwrite_keys (bool) – If True, silently overwrite existing cell_data keys.
Returns:: New domain with converted data on all meshes.
Return type:: DomainMesh

radial_basis_function_deform( control_points: Float[Tensor, 'n_controls n_spatial_dims'], control_displacements: Float[Tensor, 'n_controls n_spatial_dims'], *, kernel: Literal['thin_plate_spline'] = 'thin_plate_spline', polynomial: bool = True, smoothing: float = 0.0, point_weights: str | tuple[str, ...] | None = None, implementation: Literal['torch', 'warp'] | None = None, ) → DomainMesh[source]#

Deform every component with one global thin-plate-spline RBF field.

The same controls, fitted coefficients, kernel, and evaluation backend are shared across the interior and every boundary. With no point weights, coincident component points therefore receive identical motion. When supplied, point_weights is a common Mesh.point_data key (or nested tuple key) resolved independently on each component. Each resolved weight scales the fitted field at its point. Coincident component points therefore receive identical motion only when their resolved weights match. Raw weight tensors are rejected because component point counts differ.

Parameters:

control_points (torch.Tensor) – World-coordinate controls with shape (n_controls, n_spatial_dims) and the same float32 or float64 dtype and device as every component’s points.
control_displacements (torch.Tensor) – Displacement vectors, not destination coordinates, with the same shape, dtype, and device as control_points.
kernel ({"thin_plate_spline"}, optional) – Radial kernel used by the interpolant. Default is "thin_plate_spline".
polynomial (bool, optional) – Add the standard affine polynomial tail and side constraints. When controls are present, this requires at least D + 1 distinct controls that span the ambient affine basis and form a nonsingular augmented system. Default is True.
smoothing (float, optional) – Nonnegative diagonal regularization. With a nonsingular control layout, zero interpolates the control displacements up to solver precision. Positive values relax interpolation accuracy. Default is 0.0.
point_weights (str, tuple[str, ...], or None) – Optional point-data key present in every component. Resolved tensors must use one common bool or floating dtype. Floating weights must match component point dtypes. Every resolved tensor must be on the same device as its component’s points. Raw tensors are not accepted.
implementation ({"torch", "warp"} or None) – Field-evaluation backend. Both paths use PyTorch for the coefficient solve. Automatic dispatch uses Torch on CPU. On CUDA, it uses Warp when available and otherwise Torch.

Returns:

New domain with deformed component meshes and unchanged domain data.

Return type:

DomainMesh

Raises:

TypeError – If control tensors or Python arguments have unsupported types, or if tensor dtypes are unsupported or mismatched.
ValueError – If component data, tensor shapes, devices, control layout, point weights, or RBF options are invalid.
KeyError – If a point-data key is missing or implementation does not name a registered backend.
ImportError – If an explicitly requested backend is unavailable.
RuntimeError – If runtime validation or coefficient fitting fails, including for a singular system or during CUDA Graph capture.

Notes

The thin-plate-spline field has global support. Connectivity and attached mesh and domain data are retained. Attached vector and tensor fields are treated as Lagrangian data and are not pushed forward. Geometry caches are invalidated and topology caches are retained on each component. The operation does not detect inverted, degenerate, or self-intersecting cells. Use each component mesh’s Mesh.validate() method explicitly when required. Coefficient fitting is not supported inside CUDA Graph capture because the singular-system check requires host interaction.

Rotate all meshes in the domain about an axis.

Builds a rotation matrix and delegates to transform(). Center handling uses translate-rotate-translate at the domain level, so domain-level global_data vectors are correctly rotated but not translated (vectors are translation-invariant).

Parameters:

angle (float) – Rotation angle in radians.
axis (torch.Tensor or Sequence[float] or {"x", "y", "z"}, optional) – Rotation axis vector, shape \((D_s,)\). Use None for 2D.
center (torch.Tensor or Sequence[float], optional) – Center point for rotation, shape \((D_s,)\).
transform_point_data (bool or TensorDict) – Controls transformation of point_data fields. True transforms all compatible fields; a TensorDict (or dict) with scalar bool leaves selects specific fields.
transform_cell_data (bool or TensorDict) – Same semantics, for cell_data.
transform_global_data (bool or TensorDict) – Same semantics, for each mesh’s global_data and the domain-level global_data.

Returns:

New domain with rotated geometry.

Return type:

DomainMesh

save( prefix: str | None = None, copy_existing: bool = False, *, num_threads: int = 0, return_early: bool = False, share_non_tensor: bool = False, robust_key: bool | None = True, archive: bool | None = None, compression: str | int | None = None, ) → Any#

Saves the tensordict to disk.

This function is a proxy to memmap().

scale( factor: float | Float[Tensor, 'n_spatial_dims'], center: Float[Tensor, 'n_spatial_dims'] | Sequence[float] | None = None, transform_point_data: bool | TensorDict = False, transform_cell_data: bool | TensorDict = False, transform_global_data: bool | TensorDict = False, assume_invertible: bool | None = None, ) → DomainMesh[source]#

Scale all meshes in the domain by specified factor(s).

Builds a scale matrix and delegates to transform(). Center handling uses translate-scale-translate at the domain level.

Parameters:

factor (float or torch.Tensor) – Scale factor (scalar) or per-dimension factors, shape \((D_s,)\).
center (torch.Tensor or Sequence[float], optional) – Center point for scaling, shape \((D_s,)\).
transform_point_data (bool or TensorDict) – Controls transformation of point_data fields. True transforms all compatible fields; a TensorDict (or dict) with scalar bool leaves selects specific fields.
transform_cell_data (bool or TensorDict) – Same semantics, for cell_data.
transform_global_data (bool or TensorDict) – Same semantics, for each mesh’s global_data and the domain-level global_data.
assume_invertible (bool or None, optional) – Controls cache propagation. See Mesh.scale().

Returns:

New domain with scaled geometry.

Return type:

DomainMesh

select( *keys, inplace: bool = False, strict: bool = True, as_tensordict: bool = False, )#: TensorClass-specific select that supports as_tensordict.

set( key: NestedKey, value: Any, inplace: bool = False, non_blocking: bool = False, )#

Sets a new key-value pair.

Parameters:

key (str, tuple of str) – name of the key to be set. If tuple of str it is equivalent to chained calls of getattr followed by a final setattr.
value (Any) – value to be stored in the tensorclass
inplace (bool, optional) – if True, set will tentatively try to update the value in-place. If False or if the key isn’t present, the value will be simply written at its destination.

Returns:

self

state_dict( destination=None, prefix='', keep_vars=False, flatten=True, ) → dict[str, Any]#

Returns a state_dict with logical keys, matching TensorDictBase conventions.

Tensor fields appear as data keys. Non-tensor fields (strings, ints, etc.) and the tensorclass type are stored in _metadata. This replaces the legacy _tensordict/_non_tensordict wrapper format.

strip_caches() → DomainMesh[source]#

Remove cached geometry from all meshes in the domain.

Delegates to Mesh.strip_caches() for each mesh.

Returns:: New domain with all cached values cleared.
Return type:: DomainMesh

subdivide( levels: int = 1, filter: Literal['linear', 'butterfly', 'loop'] = 'linear', ) → DomainMesh[source]#

Subdivide all meshes in the domain.

Delegates to Mesh.subdivide() for each mesh.

Parameters:

levels (int, optional) – Number of subdivision iterations.
filter ({"linear", "butterfly", "loop"}, optional) – Subdivision scheme. See Mesh.subdivide().

Returns:

New domain with subdivided meshes.

Return type:

DomainMesh

to_tensordict( *, retain_none: bool | None = None, ) → TensorDict#

Convert the tensorclass into a regular TensorDict.

Makes a copy of all entries. Memmap and shared memory tensors are converted to regular tensors.

Parameters:: retain_none (bool) – if True, the None values will be written in the tensordict. Otherwise they will be discrarded. Default: True.
Returns:: A new TensorDict object containing the same values as the tensorclass.

transform( matrix: Float[Tensor, 'new_n_spatial_dims n_spatial_dims'], transform_point_data: bool | TensorDict = False, transform_cell_data: bool | TensorDict = False, transform_global_data: bool | TensorDict = False, assume_invertible: bool | None = None, ) → DomainMesh[source]#

Apply a linear transformation to all meshes in the domain.

This is the single point of contact for domain-level global_data transformation. Both rotate() and scale() delegate here after building their matrix.

Parameters:

matrix (torch.Tensor) – Transformation matrix, shape \((S', S)\).
transform_point_data (bool or TensorDict) – Controls transformation of point_data fields. True transforms all compatible fields; a TensorDict (or dict) with scalar bool leaves selects specific fields.
transform_cell_data (bool or TensorDict) – Same semantics, for cell_data.
transform_global_data (bool or TensorDict) – Same semantics, for each mesh’s global_data and the domain-level global_data.
assume_invertible (bool or None, optional) – Controls cache propagation. See Mesh.transform().

Returns:

New domain with transformed geometry.

Return type:

DomainMesh

translate( offset: Float[Tensor, 'n_spatial_dims'] | Sequence[float], ) → DomainMesh[source]#

Translate all meshes in the domain by a constant offset.

Delegates to Mesh.translate() for each mesh.

Parameters:: offset (torch.Tensor or Sequence[float]) – Translation vector, shape \((S,)\) where \(S\) is n_spatial_dims.
Returns:: New domain with translated geometry.
Return type:: DomainMesh

unbind(dim: int)#

Returns a tuple of indexed tensorclass instances unbound along the indicated dimension.

Resulting tensorclass instances will share the storage of the initial tensorclass instance.

validate( check_degenerate_cells: bool = True, check_duplicate_vertices: bool = True, check_inverted_cells: bool = False, check_out_of_bounds: bool = True, check_manifoldness: bool = False, tolerance: float | None = None, raise_on_error: bool = False, *, check_self_intersection: bool = False, ) → dict[str, Any][source]#

Validate all meshes in the domain and aggregate results.

Delegates to Mesh.validate() for the interior and each boundary mesh, then aggregates the results into a domain-level report.

Parameters:

check_degenerate_cells (bool, optional) – Check for zero/negative area cells.
check_duplicate_vertices (bool, optional) – Check for coincident vertices.
check_inverted_cells (bool, optional) – Check for negative orientation.
check_out_of_bounds (bool, optional) – Check cell indices are valid.
check_manifoldness (bool, optional) – Check manifold topology.
tolerance (float | None, optional) – Tolerance for geometric checks. If None (default), each mesh uses a dtype-aware epsilon.
raise_on_error (bool, optional) – Raise ValueError on first error vs return report.
check_self_intersection (bool, optional) – Request self-intersection checks for every component. This option is keyword-only and not yet implemented; passing True raises NotImplementedError.

Returns:

Aggregated validation report with keys:

"interior": validation report for the interior mesh (Mapping[str, bool | int | torch.Tensor], see Mesh.validate()).
"boundaries": dict[str, Mapping[str, ...]] of per-boundary reports.
"valid": bool, True only if all meshes pass validation.

Return type:

dict[str, Any]

Raises:

NotImplementedError – If check_self_intersection=True because component-level self-intersection checking is not yet implemented.