PhysicsNeMo Utils#

The PhysicsNeMo Utils module provides a comprehensive set of utilities that support various aspects of scientific computing, machine learning, and physics simulations. These utilities range from optimization helpers and distributed computing tools to specialized functions for weather/climate modeling and geometry processing. The module is designed to simplify common tasks while maintaining high performance and scalability.

Optimization utils#

The optimization utilities provide tools for capturing and managing training states, gradients, and optimization processes. These are particularly useful when implementing custom training loops or specialized optimization strategies.

class physicsnemo.utils.capture.StaticCaptureEvaluateNoGrad(*args, **kwargs)[source]#

Bases: _StaticCapture

An performance optimization decorator for PyTorch no grad evaluation.

This class should be initialized as a decorator on a function that computes run the forward pass of the model that does not require gradient calculations. This is the recommended method to use for inference and validation methods.

Parameters:

model (physicsnemo.models.Module) – PhysicsNeMo Model
logger (Optional[Logger], optional) – PhysicsNeMo Launch Logger, by default None
use_graphs (bool, optional) – Toggle CUDA graphs if supported by model, by default True
use_amp (bool, optional) – Toggle AMP if supported by mode, by default True
cuda_graph_warmup (int, optional) – Number of warmup steps for cuda graphs, by default 11
amp_type (Union[float16, bfloat16], optional) – Auto casting type for AMP, by default torch.float16
label (Optional[str], optional) – Static capture checkpoint label, by default None

Raises:

ValueError – If the model provided is not a physicsnemo.models.Module. I.e. has no meta data.

Example

>>> # Create model
>>> model = physicsnemo.models.mlp.FullyConnected(2, 64, 2)
>>> input = torch.rand(8, 2)
>>> # Create evaluate function with optimization wrapper
>>> @StaticCaptureEvaluateNoGrad(model=model)
... def eval_step(model, invar):
...     predvar = model(invar)
...     return predvar
...
>>> output = eval_step(model, input)
>>> output.size()
torch.Size([8, 2])

Note

Capturing multiple cuda graphs in a single program can lead to potential invalid CUDA memory access errors on some systems. Prioritize capturing training graphs when this occurs.

class physicsnemo.utils.capture.StaticCaptureTraining(*args, **kwargs)[source]#

Bases: _StaticCapture

A performance optimization decorator for PyTorch training functions.

This class should be initialized as a decorator on a function that computes the forward pass of the neural network and loss function. The user should only call the defind training step function. This will apply optimizations including: AMP and Cuda Graphs.

Parameters:

model (physicsnemo.models.Module) – PhysicsNeMo Model
optim (torch.optim) – Optimizer
logger (Optional[Logger], optional) – PhysicsNeMo Launch Logger, by default None
use_graphs (bool, optional) – Toggle CUDA graphs if supported by model, by default True
use_amp (bool, optional) – Toggle AMP if supported by mode, by default True
cuda_graph_warmup (int, optional) – Number of warmup steps for cuda graphs, by default 11
amp_type (Union[float16, bfloat16], optional) – Auto casting type for AMP, by default torch.float16
gradient_clip_norm (Optional[float], optional) – Threshold for gradient clipping
label (Optional[str], optional) – Static capture checkpoint label, by default None

Raises:

ValueError – If the model provided is not a physicsnemo.models.Module. I.e. has no meta data.

Example

>>> # Create model
>>> model = physicsnemo.models.mlp.FullyConnected(2, 64, 2)
>>> input = torch.rand(8, 2)
>>> output = torch.rand(8, 2)
>>> # Create optimizer
>>> optim = torch.optim.Adam(model.parameters(), lr=0.001)
>>> # Create training step function with optimization wrapper
>>> @StaticCaptureTraining(model=model, optim=optim)
... def training_step(model, invar, outvar):
...     predvar = model(invar)
...     loss = torch.sum(torch.pow(predvar - outvar, 2))
...     return loss
...
>>> # Sample training loop
>>> for i in range(3):
...     loss = training_step(model, input, output)
...

Note

Static captures must be checkpointed when training using the state_dict() if AMP is being used with gradient scaler. By default, this requires static captures to be instantiated in the same order as when they were checkpointed. The label parameter can be used to relax/circumvent this ordering requirement.

Note

Capturing multiple cuda graphs in a single program can lead to potential invalid CUDA memory access errors on some systems. Prioritize capturing training graphs when this occurs.

Neighbor utils#

Utilities for optimized layers for neighbor searches.

physicsnemo.utils.neighbors.radius_search( points: Tensor, queries: Tensor, radius: float, max_points: int | None = None, return_dists: bool = False, return_points: bool = False, backend: Literal['warp', 'torch'] = 'warp', ) → tuple[Tensor][source]#

Performs radius-based neighbor search to find points within a specified radius of query points.

Can use brute-force methods with PyTorch, or an accelerated spatial decomposition method with Warp.

This function does not currently accept a batch index.

This function has differing behavior based on the argument for max_points. If max_points is None, the function will find ALL points within the radius and return a flattened list of indices, (optionally) distances, and (optionally) points. The indices will have a shape of (2, N) where N is the aggregate number of neighbors found for all queries. The 0th index of the output represents the index of the query points, and the 1st index represents the index of the neighbor points within the search space.

If max_points is not None, the function will find the max_points closest points within the radius and return a statically sized array of indices, (optionally) distances, and (optionally) points. The indices will have a shape of (queries.shape[0], max_points). Each row i of the indices will be neighbors of queries[i]. If there are fewer points than max_points, then the unused indices will be set to -1 and the distances and points will be set to 0 for unused points.

Because the shape when max_points=None is dynamic, this function is incompatible with torch.compile in that case. When max_points is set, this function is compatible with torch.compile regardless of backend.

The different backends are not necessarily certain to provide identical output, for two reasons: first, if max_points is lower than the number of neighbors found, the selected points may be stochastic. Second, when max_points is None or max_points is greater than the number of neighbors, the outputs may be ordered differently by the two backends. Do not rely on the exact order of the neighbors in the outputs.

Parameters:

points (torch.Tensor) – The reference point cloud tensor of shape (N, 3) where N is the number of points.
queries (torch.Tensor) – The query points tensor of shape (M, 3) where M is the number of query points.
radius (float) – The search radius. Points within or at this radius of a query point will be considered neighbors.
max_points (int | None, optional) – Maximum number of neighbors to return for each query point. If None, returns all neighbors within radius. Defaults to None. See documentation for details.
return_dists (bool, optional) – If True, returns the distances to the neighbor points. Defaults to False.
return_points (bool, optional) – If True, returns the actual neighbor points in addition to their indices. Defaults to False.
backend (Literal["warp", "torch"], optional) – The backend implementation to use for the search. Either “warp” or “torch”. Defaults to “warp”.

Returns:

A tuple containing:

indices (torch.Tensor): Indices of neighbor points for each query point
counts (torch.Tensor): Number of neighbors found for each query point
distances (torch.Tensor, optional): Distances to neighbor points if return_dists=True
neighbor_points (torch.Tensor, optional): Actual neighbor points if return_points=True

Return type:

tuple

Raises:

ValueError – If backend is not “warp” or “torch”

GraphCast utils#

A collection of utilities specifically designed for working with the GraphCast model, including data processing, graph construction, and specialized loss functions. These utilities are essential for implementing and training GraphCast-based weather prediction models.

class physicsnemo.utils.graphcast.data_utils.StaticData( static_dataset_path: str, latitudes: Tensor, longitudes: Tensor, )[source]#

Bases: object

Class to load static data from netCDF files. Static data includes land-sea mask, geopotential, and latitude-longitude coordinates.

Parameters:

static_dataset_path (str) – Path to directory containing static data.
latitudes (Tensor) – Tensor with shape (lat,) that includes latitudes.
longitudes (Tensor) – Tensor with shape (lon,) that includes longitudes.

get() → Tensor[source]#

Get all static data.

Returns:: Tensor with shape (1, 5, lat, lon) that includes land-sea mask, geopotential, cosine of latitudes, sine and cosine of longitudes.
Return type:: Tensor

get_geop(normalize: bool = True) → Tensor[source]#

Get geopotential from netCDF file.

Parameters:: normalize (bool, optional) – Whether to normalize the geopotential, by default True
Returns:: Normalized geopotential with shape (1, 1, lat, lon).
Return type:: Tensor

get_lat_lon() → Tensor[source]#

Computes cosine of latitudes and sine and cosine of longitudes.

Returns:: Tensor with shape (1, 3, lat, lon) tha includes cosine of latitudes, sine and cosine of longitudes.
Return type:: Tensor

get_lsm() → Tensor[source]#

Get land-sea mask from netCDF file.

Returns:: Land-sea mask with shape (1, 1, lat, lon).
Return type:: Tensor

class physicsnemo.utils.graphcast.graph.Graph( lat_lon_grid: Tensor, mesh_level: int = 6, multimesh: bool = True, khop_neighbors: int = 0, dtype=torch.float32, )[source]#

Bases: object

Graph class for creating the graph2mesh, latent mesh, and mesh2graph graphs.

Parameters:

lat_lon_grid (Tensor) – Tensor with shape (lat, lon, 2) that includes the latitudes and longitudes meshgrid.
mesh_level (int, optional) – Level of the latent mesh, by default 6
multimesh (bool, optional) – If the latent mesh is a multimesh, by default True If True, the latent mesh includes the nodes corresponding to the specified mesh_level`and incorporates the edges from all mesh levels ranging from level 0 up to and including `mesh_level.
khop_neighbors (int, optional) – This option is used to retrieve a list of indices for the k-hop neighbors of all mesh nodes. It is applicable when a graph transformer is used as the processor. If set to 0, this list is not computed. If a message passing processor is used, it is forced to 0. By default 0.
dtype (torch.dtype, optional) – Data type of the graph, by default torch.float

create_g2m_graph(verbose: bool = True) → Tensor[source]#

Create the graph2mesh graph.

Parameters:: verbose (bool, optional) – verbosity, by default True
Returns:: Graph2mesh graph.
Return type:: DGLGraph

create_m2g_graph(verbose: bool = True) → Tensor[source]#

Create the mesh2grid graph.

Parameters:: verbose (bool, optional) – verbosity, by default True
Returns:: Mesh2grid graph.
Return type:: DGLGraph

create_mesh_graph(verbose: bool = True) → Tensor[source]#

Create the multimesh graph.

Parameters:: verbose (bool, optional) – verbosity, by default True
Returns:: Multimesh graph
Return type:: DGLGraph

physicsnemo.utils.graphcast.graph_utils.add_edge_features( graph: dgl.DGLGraph, pos: Tensor, normalize: bool = True, ) → dgl.DGLGraph[source]#

Adds edge features to the graph.

Parameters:

graph (DGLGraph) – The graph to add edge features to.
pos (Tensor) – The node positions.
normalize (bool, optional) – Whether to normalize the edge features, by default True

Returns:

The graph with edge features.

Return type:

DGLGraph

physicsnemo.utils.graphcast.graph_utils.add_node_features( graph: dgl.DGLGraph, pos: Tensor, ) → dgl.DGLGraph[source]#

Adds cosine of latitude, sine and cosine of longitude as the node features to the graph.

Parameters:

graph (DGLGraph) – The graph to add node features to.
pos (Tensor) – The node positions.

Returns:

graph – The graph with node features.

Return type:

DGLGraph

physicsnemo.utils.graphcast.graph_utils.azimuthal_angle(lon: Tensor) → Tensor[source]#

Gives the azimuthal angle of a point on the sphere

Parameters:: lon (Tensor) – Tensor of shape (N, ) containing the longitude of the point
Returns:: Tensor of shape (N, ) containing the azimuthal angle
Return type:: Tensor

physicsnemo.utils.graphcast.graph_utils.cell_to_adj(cells: List[List[int]])[source]#

creates adjancy matrix in COO format from mesh cells

Parameters:: cells (List[List[int]]) – List of cells, each cell is a list of 3 vertices
Returns:: src, dst – List of source and destination vertices
Return type:: List[int], List[int]

physicsnemo.utils.graphcast.graph_utils.create_graph( src: List, dst: List, to_bidirected: bool = True, add_self_loop: bool = False, dtype: dtype = torch.int32, ) → dgl.DGLGraph[source]#

Creates a DGL graph from an adj matrix in COO format.

Parameters:

src (List) – List of source nodes
dst (List) – List of destination nodes
to_bidirected (bool, optional) – Whether to make the graph bidirectional, by default True
add_self_loop (bool, optional) – Whether to add self loop to the graph, by default False
dtype (torch.dtype, optional) – Graph index data type, by default torch.int32

Returns:

The dgl Graph.

Return type:

DGLGraph

physicsnemo.utils.graphcast.graph_utils.create_heterograph( src: List, dst: List, labels: str, dtype: dtype = torch.int32, num_nodes_dict: dict = None, ) → dgl.DGLGraph[source]#

Creates a heterogeneous DGL graph from an adj matrix in COO format.

Parameters:

src (List) – List of source nodes
dst (List) – List of destination nodes
labels (str) – Label of the edge type
dtype (torch.dtype, optional) – Graph index data type, by default torch.int32
num_nodes_dict (dict, optional) – number of nodes for some node types, see dgl.heterograph for more information

Returns:

The dgl Graph.

Return type:

DGLGraph

physicsnemo.utils.graphcast.graph_utils.deg2rad(deg: Tensor) → Tensor[source]#

Converts degrees to radians

Parameters:: deg – Tensor of shape (N, ) containing the degrees
Returns:: Tensor of shape (N, ) containing the radians
Return type:: Tensor

physicsnemo.utils.graphcast.graph_utils.geospatial_rotation( invar: Tensor, theta: Tensor, axis: str, unit: str = 'rad', ) → Tensor[source]#

Rotation using right hand rule

Parameters:

invar (Tensor) – Tensor of shape (N, 3) containing x, y, z coordinates
theta (Tensor) – Tensor of shape (N, ) containing the rotation angle
axis (str) – Axis of rotation
unit (str, optional) – Unit of the theta, by default “rad”

Returns:

Tensor of shape (N, 3) containing the rotated x, y, z coordinates

Return type:

Tensor

physicsnemo.utils.graphcast.graph_utils.get_face_centroids( vertices: List[Tuple[float, float, float]], faces: List[List[int]], ) → List[Tuple[float, float, float]][source]#

Compute the centroids of triangular faces in a graph.

Parameters: vertices (List[Tuple[float, float, float]]): A list of tuples representing the coordinates of the vertices. faces (List[List[int]]): A list of lists, where each inner list contains three indices representing a triangular face.

Returns: List[Tuple[float, float, float]]: A list of tuples representing the centroids of the faces.

physicsnemo.utils.graphcast.graph_utils.latlon2xyz( latlon: Tensor, radius: float = 1, unit: str = 'deg', ) → Tensor[source]#

Converts latlon in degrees to xyz Based on: https://stackoverflow.com/questions/1185408 - The x-axis goes through long,lat (0,0); - The y-axis goes through (0,90); - The z-axis goes through the poles.

Parameters:

latlon (Tensor) – Tensor of shape (N, 2) containing latitudes and longitudes
radius (float, optional) – Radius of the sphere, by default 1
unit (str, optional) – Unit of the latlon, by default “deg”

Returns:

Tensor of shape (N, 3) containing x, y, z coordinates

Return type:

Tensor

physicsnemo.utils.graphcast.graph_utils.max_edge_length( vertices: List[List[float]], source_nodes: List[int], destination_nodes: List[int], ) → float[source]#

Compute the maximum edge length in a graph.

Parameters: vertices (List[List[float]]): A list of tuples representing the coordinates of the vertices. source_nodes (List[int]): A list of indices representing the source nodes of the edges. destination_nodes (List[int]): A list of indices representing the destination nodes of the edges.

Returns: The maximum edge length in the graph (float).

physicsnemo.utils.graphcast.graph_utils.polar_angle(lat: Tensor) → Tensor[source]#

Gives the polar angle of a point on the sphere

Parameters:: lat (Tensor) – Tensor of shape (N, ) containing the latitude of the point
Returns:: Tensor of shape (N, ) containing the polar angle
Return type:: Tensor

physicsnemo.utils.graphcast.graph_utils.rad2deg(rad)[source]#

Converts radians to degrees

Parameters:: rad – Tensor of shape (N, ) containing the radians
Returns:: Tensor of shape (N, ) containing the degrees
Return type:: Tensor

physicsnemo.utils.graphcast.graph_utils.xyz2latlon( xyz: Tensor, radius: float = 1, unit: str = 'deg', ) → Tensor[source]#

Converts xyz to latlon in degrees Based on: https://stackoverflow.com/questions/1185408 - The x-axis goes through long,lat (0,0); - The y-axis goes through (0,90); - The z-axis goes through the poles.

Parameters:

xyz (Tensor) – Tensor of shape (N, 3) containing x, y, z coordinates
radius (float, optional) – Radius of the sphere, by default 1
unit (str, optional) – Unit of the latlon, by default “deg”

Returns:

Tensor of shape (N, 2) containing latitudes and longitudes

Return type:

Tensor

class physicsnemo.utils.graphcast.loss.CellAreaWeightedLossFunction(area)[source]#

Bases: Module

Loss function with cell area weighting.

Parameters:: area (torch.Tensor) – Cell area with shape [H, W].

forward(invar, outvar)[source]#

Implicit forward function which computes the loss given a prediction and the corresponding targets.

Parameters:

invar (torch.Tensor) – prediction of shape [T, C, H, W].
outvar (torch.Tensor) – target values of shape [T, C, H, W].

class physicsnemo.utils.graphcast.loss.CustomCellAreaWeightedLossAutogradFunction(*args, **kwargs)[source]#

Bases: Function

Autograd fuunction for custom loss with cell area weighting.

static backward( ctx, grad_loss: Tensor, )[source]#: Backward method of custom loss function with cell area weighting.

static forward( ctx, invar: Tensor, outvar: Tensor, area: Tensor, )[source]#: Forward of custom loss function with cell area weighting.

class physicsnemo.utils.graphcast.loss.CustomCellAreaWeightedLossFunction(area: Tensor)[source]#

Bases: CellAreaWeightedLossFunction

Custom loss function with cell area weighting.

Parameters:: area (torch.Tensor) – Cell area with shape [H, W].

forward( invar: Tensor, outvar: Tensor, ) → Tensor[source]#

Implicit forward function which computes the loss given a prediction and the corresponding targets.

Parameters:

invar (torch.Tensor) – prediction of shape [T, C, H, W].
outvar (torch.Tensor) – target values of shape [T, C, H, W].

class physicsnemo.utils.graphcast.loss.GraphCastLossFunction( area, channels_list, dataset_metadata_path, time_diff_std_path, )[source]#

Bases: Module

Loss function as specified in GraphCast. :param area: Cell area with shape [H, W]. :type area: torch.Tensor

assign_atmosphere_weights()[source]#: Assigns weights to atmospheric variables

assign_surface_weights()[source]#: Assigns weights to surface variables

assign_variable_weights()[source]#: assigns per-variable per-pressure level weights

calculate_linear_weights(variables)[source]#: Calculate weights for each variable group.

forward(invar, outvar)[source]#: Implicit forward function which computes the loss given a prediction and the corresponding targets. :param invar: prediction of shape [T, C, H, W]. :type invar: torch.Tensor :param outvar: target values of shape [T, C, H, W]. :type outvar: torch.Tensor

get_channel_dict( dataset_metadata_path, channels_list, )[source]#: Gets lists of surface and atmospheric channels

get_time_diff_std( time_diff_std_path, channels_list, )[source]#: Gets the time difference standard deviation

parse_variable(variable_list)[source]#: Parse variable into its letter and numeric parts.

Filesystem utils#

Utilities for handling file operations, caching, and data management across different storage systems. These utilities abstract away the complexity of dealing with different filesystem types and provide consistent interfaces for data access.

class physicsnemo.utils.filesystem.Package(root: str, seperator: str = '/')[source]#

Bases: object

A generic file system abstraction. Can be used to represent local and remote file systems. Remote files are automatically fetched and stored in the $LOCAL_CACHE or $HOME/.cache/physicsnemo folder. The get method can then be used to access files present.

Presently one can use Package with the following directories: - Package(“/path/to/local/directory”) = local file system - Package(“s3://bucket/path/to/directory”) = object store file system - Package(“http://url/path/to/directory”) = http file system - Package(“ngc://model/<org_id/team_id/model_id>@<version>”) = ngc model file system

Parameters:

root (str) – Root directory for file system
seperator (str, optional) – directory seperator. Defaults to “/”.

get(path: str, recursive: bool = False) → str[source]#

Get a local path to the item at path

path might be a remote file, in which case it is downloaded to a local cache at $LOCAL_CACHE or $HOME/.cache/physicsnemo first.

Diffusion utils#

Tools for working with diffusion models and other generative approaches, including deterministic and stochastic sampling utilities.

physicsnemo.utils.diffusion.deterministic_sampler.deterministic_sampler( net: ~torch.nn.modules.module.Module, latents: ~torch.Tensor, img_lr: ~torch.Tensor, class_labels: ~torch.Tensor | None = None, randn_like: ~typing.Callable = <built-in method randn_like of type object>, patching: ~physicsnemo.utils.patching.GridPatching2D | None = None, mean_hr: ~torch.Tensor | None = None, lead_time_label: ~torch.Tensor | None = None, num_steps: int = 18, sigma_min: float | None = None, sigma_max: float | None = None, rho: float = 7.0, solver: ~typing.Literal['heun', 'euler'] = 'heun', discretization: ~typing.Literal['vp', 've', 'iddpm', 'edm'] = 'edm', schedule: ~typing.Literal['vp', 've', 'linear'] = 'linear', scaling: ~typing.Literal['vp', 'none'] = 'none', epsilon_s: float = 0.001, C_1: float = 0.001, C_2: float = 0.008, M: int = 1000, alpha: float = 1.0, S_churn: int = 0, S_min: float = 0.0, S_max: float = inf, S_noise: float = 1.0, dtype: ~torch.dtype = torch.float64, ) → Tensor[source]#

Generalized sampler, representing the superset of all sampling methods discussed in the paper Elucidating the Design Space of Diffusion-Based Generative Models (EDM).

This function integrates an ODE (probability flow) or SDE over multiple time-steps to generate samples from the diffusion model provided by the argument ‘net’. It can be used to combine multiple choices to design a custom sampler, including multiple integration solver, discretization method, noise schedule, and so on.

Parameters:

net (torch.nn.Module) – The diffusion model to use in the sampling process.
latents (torch.Tensor) – The latent random noise used as the initial condition for the stochastic ODE.
img_lr (torch.Tensor) – Low-resolution input image for conditioning the diffusion process. Passed as a keywork argument to the model net.
class_labels (Optional[torch.Tensor]) – Labels of the classes used as input to a class-conditionned diffusion model. Passed as a keyword argument to the model net. If provided, it must be a tensor containing integer values. Defaults to None, in which case it is ignored.
randn_like (Callable) – Random Number Generator to generate random noise that is added during the stochastic sampling. Must have the same signature as torch.randn_like and return torch.Tensor. Defaults to torch.randn_like.
patching (Optional[GridPatching2D], default=None) – A patching utility for patch-based diffusion. Implements methods to extract patches from an image and batch the patches along dim=0. Should also implement a fuse method to reconstruct the original image from a batch of patches. See GridPatching2D for details. By default None, in which case non-patched diffusion is used.
mean_hr (Optional[Tensor], optional) – Optional tensor containing mean high-resolution images for conditioning. Must have same height and width as img_lr, with shape $(B_{hr}, C_{hr}, H, W)$ where the batch dimension $B_{hr}$ can be either 1, either equal to batch_size, or can be omitted. If $B_{hr} = 1$ or is omitted, mean_hr will be expanded to match the shape of img_lr. By default None.
lead_time_label (Optional[Tensor], optional) – Optional lead time labels. By default None.
num_steps (Optional[int]) – Number of time-steps for the stochastic ODE integration. Defaults to 18.
sigma_min (Optional[float]) – Minimum noise level for the diffusion process. sigma_min, sigma_max, and rho are used to compute the time-step discretization, based on the choice of discretization. For the default choice (discretization='heun'), the noise level schedule is computed as: $\sigma_i = (\sigma_{max}^{1/\rho} + i / (\text{num_steps} - 1) * (\sigma_{min}^{1/\rho} - \sigma_{max}^{1/\rho}))^{\rho}$. For other choices of discretization, see details in the EDM paper. Defaults to None, in which case defaults values depending of the specified discretization are used.
sigma_max (Optional[float]) – Maximum noise level for the diffusion process. See sigma_min for details. Defaults to None, in which case defaults values depending of the specified discretization are used.
rho (float, optional) – Exponent used in the noise schedule. See sigma_min for details. Only used when discretization="heun". Values in the range [5, 10] produce better images. Lower values lead to truncation errors equalized over all time steps. Defaults to 7.
solver (Literal["heun", "euler"]) – The numerical method used to integrate the stochastic ODE. "euler" is 1st order solver, which is faster but produces lower-quality images. "heun" is 2nd order, more expensive, but produces higher-quality images. Defaults to "heun".
discretization (Literal["vp", "ve", "iddpm", "edm"]) – The method to discretize time-steps $t_i$ in the diffusion process. See the EDM paper for details. Defaults to "edm".
schedule (Literal["vp", "ve", "linear"]) – The type of noise level schedule. Defaults to "linear". If schedule="ve", then $\sigma(t) = \sqrt{t}$. If schedule="linear", then $\sigma(t) = t$. If schedule="vp", see EDM paper for details. Defaults to "linear".
scaling (Literal["vp", "none"]) – The type of time-dependent signal scaling $s(t)$, such that $x = s(t) \hat{x}$. See EDM paper for details on the "vp" scaling. Defaults to "none", in which case $s(t)=1$.
epsilon_s (float, optional) – Parameter to compute both the noise level schedule and the time-step discetization. Only used when discretization="vp" or schedule="vp". Ignored in other cases. Defaults to 1e-3.
C_1 (float, optional) – Parameters to compute the time-step discetization. Only used when discretization="iddpm". Defaults to 0.001.
C_2 (float, optional) – Same as for C_1. Only used when discretization="iddpm". Defaults to 0.008.
M (int, optional) – Same as for C_1 and C_2. Only used when discretization="iddpm". Defaults to 1000.
alpha (float, optional) – Controls (i.e. multiplies) the step size $t_{i+1} - \hat{t}_i$ in the stochastic sampler, where $\hat{t}_i$ is the temporarily increased noise level. Defaults to 1.0, which is the recommended value.
S_churn (int, optional) – Controls the amount of stochasticty injected in the SDE in the stochatsic sampler. Larger values of S_churn lead to larger values of $\hat{t}_i$, which in turn lead to injecting more stochasticity in the SDE by Defaults to 0, which means no stochasticity is injected.
S_min (float, optional) – S_min and S_max control the time-step range over which stochasticty is injected in the SDE. Stochasticity is injected through $\hat{t}_i$ for time-steps $t_i$ such that $S_{min} \leq t_i \leq S_{max}$. Defaults to 0.0.
S_max (float, optional) – See S_min. Defaults to float("inf").
S_noise (float, optional) – Controls the amount of stochasticty injected in the SDE in the stochatsic sampler. Added signal noise is proportinal to $\epsilon_i$ where $\epsilon_i \sim \mathcal{N}(0, S_{noise}^2)$. Defaults to 1.0.
dtype (torch.dtype, optional) – Controls the precision used for sampling

Returns:

Generated batch of samples. Same shape as the input latents.

Return type:

torch.Tensor

physicsnemo.utils.diffusion.stochastic_sampler.stochastic_sampler(net: ~torch.nn.modules.module.Module, latents: ~torch.Tensor, img_lr: ~torch.Tensor, class_labels: ~torch.Tensor | None = None, randn_like: ~typing.Callable[[~torch.Tensor], ~torch.Tensor] = <built-in method randn_like of type object>, patching: ~physicsnemo.utils.patching.GridPatching2D | None = None, mean_hr: ~torch.Tensor | None = None, lead_time_label: ~torch.Tensor | None = None, num_steps: int = 18, sigma_min: float = 0.002, sigma_max: float = 800, rho: float = 7, S_churn: float = 0, S_min: float = 0, S_max: float = inf, S_noise: float = 1) → Tensor[source]#

Proposed EDM sampler (Algorithm 2) with minor changes to enable super-resolution and patch-based diffusion.

Parameters:

net (torch.nn.Module) –
The neural network model that generates denoised images from noisy inputs. Expected signature: net(x, x_lr, t_hat, class_labels, lead_time_label=lead_time_label, embedding_selector=embedding_selector).
Inputs:
- x (torch.Tensor): Noisy input of shape $(B, C_{out}, H, W)$
- x_lr (torch.Tensor): Conditioning input of shape $(B, C_{cond}, H, W)$
- t_hat (torch.Tensor): Noise level of shape $(B, 1, 1, 1)$ or scalar
- class_labels (torch.Tensor, optional): Optional class labels
- lead_time_label (torch.Tensor, optional): Optional lead time labels
- embedding_selector (callable, optional): Function to select positional embeddings. Used for patch-based diffusion.
Output:
- denoised (torch.Tensor): Denoised prediction of shape $(B, C_{out}, H, W)$
Required attributes:
- sigma_min (float): Minimum supported noise level for the model
- sigma_max (float): Maximum supported noise level for the model
- round_sigma (callable): Method to convert sigma values to tensor representation
latents (Tensor) – The latent variables (e.g., noise) used as the initial input for the sampler. Has shape $(B, C_{out}, H, W)$.
img_lr (Tensor) – Low-resolution input image for conditioning the super-resolution process. Must have shape $(B, C_{lr}, H, W)$.
class_labels (Optional[Tensor], optional) – Class labels for conditional generation, if required by the model. By default None.
randn_like (Callable[[Tensor], Tensor]) – Function to generate random noise with the same shape as the input tensor. By default torch.randn_like.
patching (Optional[GridPatching2D], default=None) – A patching utility for patch-based diffusion. Implements methods to extract patches from an image and batch the patches along dim=0. Should also implement a fuse method to reconstruct the original image from a batch of patches. See GridPatching2D for details. By default None, in which case non-patched diffusion is used.
mean_hr (Optional[Tensor], optional) – Optional tensor containing mean high-resolution images for conditioning. Must have same height and width as img_lr, with shape $(B_{hr}, C_{hr}, H, W)$ where the batch dimension $B_{hr}$ can be either 1, either equal to batch_size, or can be omitted. If $B_{hr} = 1$ or is omitted, mean_hr will be expanded to match the shape of img_lr. By default None.
lead_time_label (Optional[Tensor], optional) – Optional lead time labels. By default None.
num_steps (int) – Number of time steps for the sampler. By default 18.
sigma_min (float) – Minimum noise level. By default 0.002.
sigma_max (float) – Maximum noise level. By default 800.
rho (float) – Exponent used in the time step discretization. By default 7.
S_churn (float) – Churn parameter controlling the level of noise added in each step. By default 0.
S_min (float) – Minimum time step for applying churn. By default 0.
S_max (float) – Maximum time step for applying churn. By default float("inf").
S_noise (float) – Noise scaling factor applied during the churn step. By default 1.

Returns:

The final denoised image produced by the sampler. Same shape as latents: $(B, C_{out}, H, W)$.

Return type:

Tensor

See also

EDMPrecondSuperResolution: A model wrapper that provides preconditioning for super-resolution diffusion models and implements the required interface for this sampler.

Miscellaneous utility classes and functions.

class physicsnemo.utils.diffusion.utils.EasyDict[source]#

Bases: dict

Convenience class that behaves like a dict but allows access with the attribute syntax.

class physicsnemo.utils.diffusion.utils.InfiniteSampler( dataset: Dataset, rank: int = 0, num_replicas: int = 1, shuffle: bool = True, seed: int = 0, window_size: float = 0.5, start_idx: int = 0, )[source]#

Bases: Sampler[int]

Sampler for torch.utils.data.DataLoader that loops over the dataset indefinitely.

This sampler yields indices indefinitely, optionally shuffling items as it goes. It can also perform distributed sampling when rank and num_replicas are specified.

Parameters:

dataset (torch.utils.data.Dataset) – The dataset to sample from
rank (int, default=0) – The rank of the current process within num_replicas processes
num_replicas (int, default=1) – The number of processes participating in distributed sampling
shuffle (bool, default=True) – Whether to shuffle the indices
seed (int, default=0) – Random seed for reproducibility when shuffling
window_size (float, default=0.5) – Fraction of dataset to use as window for shuffling. Must be between 0 and 1. A larger window means more thorough shuffling but slower iteration.
start_idx (int, default=0) – The initial index to use for the sampler. This is used for resuming training.

class physicsnemo.utils.diffusion.utils.StackedRandomGenerator( device: device, seeds: Sequence[int], )[source]#

Bases: object

Wrapper for torch.Generator that allows specifying a different random seed for each sample in a minibatch.

Parameters:

device (torch.device) – Device to use for the random number generator.
seeds (Sequence[int]) – Sequence (e.g. list or tuple) of random seeds for each sample in the minibatch. Its length defines the batch size of generated samples.

randint(

*args: Any,

size: Size | Sequence[int],

**kwargs: Any,

) → Tensor[source]#

Generate stacked samples from a uniform distribution over the integers.

Parameters:

*args (Any) – Required positional arguments to pass to torch.randint.
size (Sequence[int] | torch.Size) – Size of the output tensor. Accepts any sequence of integers or a torch.Size instance. First dimension must match the number of random seeds.
**kwargs (Any) – Additional keyword arguments to pass to torch.randint.

Returns:

Stacked samples from a uniform distribution over the integers. Shape matches size.

Return type:

torch.Tensor

randn(

size: Size | Sequence[int],

**kwargs: Any,

) → Tensor[source]#

Generate stacked samples from a standard normal distribution. Each sample is generated using a different random seed.

Parameters:

size (Sequence[int] | torch.Size) – Size of the output tensor. Accepts any sequence of integers or a torch.Size instance. First dimension must match the number of random seeds.
**kwargs (Any) – Additional arguments to pass to torch.randn.

Returns:

Stacked samples from a standard normal distribution. Shape matches size.

Return type:

torch.Tensor

randn_like( input: Tensor, ) → Tensor[source]#

Generate stacked samples from a standard normal distribution with the same shape and data type as the input tensor.

Parameters:: input (torch.Tensor) – Input tensor to match the shape, data type, memory layout, and device of.
Returns:: Stacked samples from a standard normal distribution. Shape matches input.shape.
Return type:: torch.Tensor

randt(

nu: int,

size: Size | Sequence[int],

**kwargs: Any,

) → Tensor[source]#

Generate stacked samples from a standard Student-t distribution with nu degrees of freedom. This is useful when sampling from heavy-tailed diffusion models.

Parameters:

nu (int) – Degrees of freedom for the Student-t distribution. Must be > 2.
size (Sequence[int] | torch.Size) – Size of the output tensor. Accepts any sequence of integers or a torch.Size instance. First dimension must match the number of random seeds.
**kwargs (Any) – Additional arguments to pass to torch.randn.

Returns:

Stacked samples from a standard Student-t distribution. Shape matches size.

Return type:

torch.Tensor

physicsnemo.utils.diffusion.utils.assert_shape(tensor, ref_shape)[source]#: Assert that the shape of a tensor matches the given list of integers. None indicates that the size of a dimension is allowed to vary. Performs symbolic assertion when used in torch.jit.trace().

physicsnemo.utils.diffusion.utils.call_func_by_name(

*args,

func_name: str = None,

**kwargs,

) → Any[source]#: Finds the python object with the given name and calls it as a function.

physicsnemo.utils.diffusion.utils.check_ddp_consistency(module, ignore_regex=None)[source]#: Check DistributedDataParallel consistency across processes.

physicsnemo.utils.diffusion.utils.constant( value, shape=None, dtype=None, device=None, memory_format=None, )[source]#: Cached construction of constant tensors

physicsnemo.utils.diffusion.utils.construct_class_by_name(

*args,

class_name: str = None,

**kwargs,

) → Any[source]#: Finds the python class with the given name and constructs it with the given arguments.

physicsnemo.utils.diffusion.utils.convert_datetime_to_cftime( time: datetime, cls=cftime.DatetimeGregorian, ) → cftime.DatetimeGregorian[source]#: Convert a Python datetime object to a cftime DatetimeGregorian object.

physicsnemo.utils.diffusion.utils.copy_files_and_create_dirs( files: List[Tuple[str, str]], ) → None[source]#: Takes in a list of tuples of (src, dst) paths and copies files. Will create all necessary directories.

physicsnemo.utils.diffusion.utils.copy_params_and_buffers(src_module, dst_module, require_all=False)[source]#: Copy parameters and buffers from a source module to target module

physicsnemo.utils.diffusion.utils.ddp_sync(module, sync)[source]#: Context manager for easily enabling/disabling DistributedDataParallel synchronization.

physicsnemo.utils.diffusion.utils.format_time(seconds: int | float) → str[source]#: Convert the seconds to human readable string with days, hours, minutes and seconds.

physicsnemo.utils.diffusion.utils.format_time_brief(seconds: int | float) → str[source]#: Convert the seconds to human readable string with days, hours, minutes and seconds.

physicsnemo.utils.diffusion.utils.get_dtype_and_ctype( type_obj: Any, ) → Tuple[dtype, Any][source]#: Given a type name string (or an object having a __name__ attribute), return matching Numpy and ctypes types that have the same size in bytes.

physicsnemo.utils.diffusion.utils.get_module_dir_by_obj_name(obj_name: str) → str[source]#: Get the directory path of the module containing the given object name.

physicsnemo.utils.diffusion.utils.get_module_from_obj_name( obj_name: str, ) → Tuple[ModuleType, str][source]#: Searches for the underlying module behind the name to some python object. Returns the module and the object name (original name with module part removed).

physicsnemo.utils.diffusion.utils.get_obj_by_name(name: str) → Any[source]#: Finds the python object with the given name.

physicsnemo.utils.diffusion.utils.get_obj_from_module( module: ModuleType, obj_name: str, ) → Any[source]#: Traverses the object name and returns the last (rightmost) python object.

physicsnemo.utils.diffusion.utils.get_top_level_function_name(obj: Any) → str[source]#: Return the fully-qualified name of a top-level function.

physicsnemo.utils.diffusion.utils.is_top_level_function(obj: Any) → bool[source]#: Determine whether the given object is a top-level function, i.e., defined at module scope using ‘def’.

physicsnemo.utils.diffusion.utils.list_dir_recursively_with_ignore( dir_path: str, ignores: List[str] = None, add_base_to_relative: bool = False, ) → List[Tuple[str, str]][source]#: List all files recursively in a given directory while ignoring given file and directory names. Returns list of tuples containing both absolute and relative paths.

physicsnemo.utils.diffusion.utils.named_params_and_buffers(module)[source]#: Get named parameters and buffers of a nn.Module

physicsnemo.utils.diffusion.utils.params_and_buffers(module)[source]#: Get parameters and buffers of a nn.Module

physicsnemo.utils.diffusion.utils.parse_int_list(s)[source]#: Parse a comma separated list of numbers or ranges and return a list of ints. Example: ‘1,2,5-10’ returns [1, 2, 5, 6, 7, 8, 9, 10]

physicsnemo.utils.diffusion.utils.print_module_summary( module, inputs, max_nesting=3, skip_redundant=True, )[source]#: Print summary table of module hierarchy.

physicsnemo.utils.diffusion.utils.profiled_function(fn)[source]#: Function decorator that calls torch.autograd.profiler.record_function().

physicsnemo.utils.diffusion.utils.suppress_tracer_warnings()[source]#: Context manager to temporarily suppress known warnings in torch.jit.trace(). Note: Cannot use catch_warnings because of https://bugs.python.org/issue29672

physicsnemo.utils.diffusion.utils.time_range( start_time: datetime, end_time: datetime, step: timedelta, inclusive: bool = False, )[source]#: Like the Python range iterator, but with datetimes.

physicsnemo.utils.diffusion.utils.tuple_product(t: Tuple) → Any[source]#: Calculate the product of the tuple elements.

Geometry utils#

Utilities for geometric operations, including neighbor search and signed distance field calculations. These are essential for physics simulations and geometric deep learning applications.

Weather / Climate utils#

Specialized utilities for weather and climate modeling, including calculations for solar radiation and atmospheric parameters. These utilities are used extensively in weather prediction models.

physicsnemo.utils.insolation.insolation( dates, lat, lon, scale=1.0, daily=False, enforce_2d=False, clip_zero=True, )[source]#

Calculate the approximate solar insolation for given dates.

For an example reference, see: https://brian-rose.github.io/ClimateLaboratoryBook/courseware/insolation/

Parameters:

dates (np.ndarray)
dates – 1d array: datetime or Timestamp
lat (np.ndarray) – 1d or 2d array of latitudes
lon (np.ndarray) – 1d or 2d array of longitudes (0-360deg). If 2d, must match the shape of lat.
scale (float, optional) – scaling factor (solar constant)
daily (bool, optional) – if True, return the daily max solar radiation (lat and day of year dependent only)
enforce_2d (bool, optional) – if True and lat/lon are 1-d arrays, turns them into 2d meshes.
clip_zero (bool, optional) – if True, set values below 0 to 0

Returns:

np.ndarray

Return type:

insolation (date, lat, lon)

Patching utils#

Patching utilities are particularly useful for patch-based diffusion, also called multi-diffusion. This approach is used to scale diffusion to very large images. The following patching utilities extract patches from 2D images, and typically gather them in the batch dimension. A batch of patches is therefore composed of multiple smaller patches extracted from each sample in the original batch of larger images. Diffusion models can then process these patches independently. These utilities also support fusing operations to reconstruct the entire predicted image from the individual predicted patches.

class physicsnemo.utils.patching.BasePatching2D( img_shape: Tuple[int, int], patch_shape: Tuple[int, int], )[source]#

Bases: ABC

Abstract base class for 2D image patching operations.

This class provides a foundation for implementing various image patching strategies. It handles basic parameter validation and provides default methods for patching and fusing.

It is designed to be extensible to support different patching strategies. Any new patching strategy for 2D images should inherit from this class and implement the abstract methods.

Parameters:

img_shape (Tuple[int, int]) – The height and width of the full input images $(H, W)$.
patch_shape (Tuple[int, int]) – The height and width of the patches to extract $(H_p, W_p)$.

abstract apply(input: Tensor, **kwargs) → Tensor[source]#

Apply the patching operation to a batch of full images.

Parameters:

input (Tensor) – Batch of full input images of shape $(B, C, H, W)$.
**kwargs (dict) – Additional keyword arguments specific to the patching implementation.

Returns:

Patched tensor, shape depends on specific implementation.

Return type:

Tensor

fuse(input: Tensor, **kwargs) → Tensor[source]#

Fuse patches back into a complete image.

Parameters:

input (Tensor) – Input tensor containing patches. Shape depends on specific implementation.
**kwargs (dict) – Additional keyword arguments specific to the fusion implementation.

Returns:

Fused tensor. Shape depends on specific implementation.

Return type:

Tensor

Raises:

NotImplementedError – If the subclass does not implement this method.

global_index( batch_size: int, device: device | str = 'cpu', ) → Tensor[source]#

Returns a tensor containing the global indices for each patch.

Global indices correspond to $(y, x)$ global grid coordinates of each element within the original image (before patching). It is typically used to keep track of the original position of each patch in the original image.

Parameters:

batch_size (int) – The size $B$ of the batch of images to patch.
device (Union[torch.device, str], default="cpu") – Proper device to initialize global_index on.

Returns:

A tensor of shape $(P, 2, H_p, W_p)$, where $P$ is the number of patches to extract (corresponds to self.patch_num for classes that implement this attribute). The y-coordinate is stored in global_index[:, 0, :, :] and the x-coordinate is stored in global_index[:, 1, :, :].

Return type:

Tensor

class physicsnemo.utils.patching.GridPatching2D( img_shape: Tuple[int, int], patch_shape: Tuple[int, int], overlap_pix: int = 0, boundary_pix: int = 0, )[source]#

Bases: BasePatching2D

Class for deterministically extracting patches from 2D images in a grid pattern.

This class provides utilities to extract patches from images in a deterministic manner, with configurable overlap and boundary pixels. The patches are extracted in a grid-like pattern covering the entire image.

Parameters:

img_shape (Tuple[int, int]) – The height and width of the full input images $(H, W)$.
patch_shape (Tuple[int, int]) – The height and width of the patches to extract $(H_p, W_p)$.
overlap_pix (int, optional, default=0) – Number of pixels to overlap between adjacent patches.
boundary_pix (int, optional, default=0) – Number of pixels to crop as boundary from each patch.

patch_num#

Total number of patches $P$ that will be extracted from the image, calculated as $P = P_x * P_y$.

Type:: int

See also

physicsnemo.utils.patching.BasePatching2D: The base class providing the patching interface.
physicsnemo.utils.patching.RandomPatching2D: Alternative patching strategy using random patch locations.

apply( input: Tensor, additional_input: Tensor | None = None, ) → Tensor[source]#

Apply deterministic patching to the input tensor.

Splits the input tensor into patches in a grid-like pattern. Can optionally concatenate additional interpolated data to each patch. Extracted patches are batched along the first dimension of the output. The layout of the output assumes that for any patch index i, out[B * i: B * (i + 1)] corresponds to the same patch extracted from each batch element of input.

Parameters:

input (Tensor) – Batch of full input images of shape $(B, C, H, W)$.
additional_input (Optional[Tensor], optional, default=None) – Additional data to concatenate to each patch. Shape must be $(B, C_{add}, H_{add}, W_{add})$. Will be interpolated to match patch dimensions $(H_p, W_p)$ Note: ``additional_input`` is not patched or decomposed.

Returns:

Tensor containing patches with shape $(P \times B, C [+ C_{add}], H_p, W_p)$. If additional_input is provided, it is channel-wise concatenated to the extracted patches.

Return type:

Tensor

See also

physicsnemo.utils.patching.image_batching(): The underlying function used to perform the patching operation.

fuse( input: Tensor, batch_size: int, ) → Tensor[source]#

Fuse patches back into a complete image.

Reconstructs the original image by stitching together patches, accounting for overlapping regions and boundary pixels. In overlapping regions, values are averaged.

Parameters:

input (Tensor) – Input tensor containing patches with shape $(P \times B, C, H_p, W_p)$. Note: the patch layout along the batch dimension should be the same as the one returned by the method :meth:`~physicsnemo.utils.patching.GridPatching2D.apply`.
batch_size (int) – The original batch size $B$ before patching.

Returns:

Reconstructed image tensor with shape $(B, C, H, W)$.

Return type:

Tensor

See also

physicsnemo.utils.patching.image_fuse(): The underlying function used to perform the fusion operation.

static get_overlap_count( patch_shape: tuple[int, int], img_shape: tuple[int, int], overlap_pix: int, boundary_pix: int, ) → Tensor[source]#

Compute overlap count map for image patch reconstruction.

Calculates how many times each pixel in the padded image is covered by extracted patches, based on the patch size, overlap size, and boundary padding. This is useful for normalizing the reconstructed image after folding overlapping patches.

The overlap count is stored in self._overlap_count.

Parameters:

img_shape (Tuple[int, int]) – The height and width of the full input images $(H, W)$.
patch_shape (Tuple[int, int]) – The height and width of the patches to extract $(H_p, W_p)$.
overlap_pix (int) – The number of overlapping pixels between adjacent patches.
boundary_pix (int) – The number of pixels to crop as a boundary from each patch.

Returns:

Tensor indicating how many times each pixel in the original input is visited (or covered) by patches. Shape is $(1, 1, H_{pad}, W_{pad})$, where $H_{pad}$ and $W_{pad}$ are the padded image dimensions. Those are computed as $H_{pad} = (H_p - \text{overlap_pix} - \text{boundary_pix}) \times (P_H - 1) + H_p + \text{boundary_pix}$, where $P_H$ is the number of patches along the height of the image (and similarly for $W_{pad}$).

Return type:

Tensor

class physicsnemo.utils.patching.RandomPatching2D( img_shape: Tuple[int, int], patch_shape: Tuple[int, int], patch_num: int, )[source]#

Bases: BasePatching2D

Class for randomly extracting patches from 2D images.

This class provides utilities to randomly extract patches from a batch of full images represented as 4D tensors. It maintains a list of random patch indices that can be reset as needed.

Parameters:

img_shape (Tuple[int, int]) – The height and width $(H, W)$ of the full input images.
patch_shape (Tuple[int, int]) – The height and width $(H_p, W_p)$ of the patches to extract.
patch_num (int) – The number of patches $P$ to extract.

patch_indices#

The indices of the patches to extract from the images. These indices correspond to the $(y, x)$ coordinates of the upper left corner of each patch.

Type:: List[Tuple[int, int]]

See also

physicsnemo.utils.patching.BasePatching2D: The base class providing the patching interface.
physicsnemo.utils.patching.GridPatching2D: Alternative patching strategy using deterministic patch locations.

apply( input: Tensor, additional_input: Tensor | None = None, ) → Tensor[source]#

Applies the patching operation by extracting patches specified by self.patch_indices from the input Tensor. Extracted patches are batched along the first dimension of the output. The layout of the output assumes that for any patch index i, out[B * i: B * (i + 1)] corresponds to the same patch extracted from each batch element of input.

Parameters:

input (Tensor) – The input tensor representing the full image with shape $(B, C, H, W)$.
additional_input (Optional[Tensor], optional) – Its shape should be $(B, C_{add}, H_{add}, W_{add})$. Must have same batch size as input. Bilinear interpolation is used to interpolate additional_input onto a 2D grid of shape $(H_p, W_p)$. It is then channel-wise concatenated to the extracted patches. Note: ``additional_input`` is not patched or decomposed.

Returns:

A tensor of shape $(P \times B, C [+ C_{add}], H_p, W_p)$. If additional_input is provided, it is channel-wise concatenated to the extracted patches.

Return type:

Tensor

See also

physicsnemo.utils.patching.image_batching(): The underlying function used to perform the patching operation.

get_patch_indices() → List[Tuple[int, int]][source]#

Get the current list of patch starting indices.

These are the upper-left coordinates of each extracted patch from the full image.

Returns:: A list of (row, column) tuples representing patch starting positions.
Return type:: List[Tuple[int, int]]

property patch_num: int#

Get the number of patches to extract.

Returns:: The number of patches $P$ to extract.
Return type:: int

reset_patch_indices() → None[source]#: Generate new random indices for the patches to extract. These are the starting indices of the patches to extract (upper left corner).

set_patch_num(value: int) → None[source]#

Set the number of patches to extract and reset patch indices. This is the only way to modify the patch_num attribute.

Parameters:: value (int) – The new number of patches $P$ to extract.

physicsnemo.utils.patching.image_batching( input: Tensor, patch_shape_y: int, patch_shape_x: int, overlap_pix: int, boundary_pix: int, input_interp: Tensor | None = None, ) → Tensor[source]#

Splits a full image into a batch of patched images.

This function takes a full image and splits it into patches, adding padding where necessary. It can also concatenate additional interpolated data to each patch if provided.

Parameters:

input (Tensor) – The input tensor representing a batch of full image with shape $(B, C, H, W)$.
patch_shape_y (int) – The height $H_p$ of each image patch.
patch_shape_x (int) – The width $W_p$ of each image patch.
overlap_pix (int) – The number of overlapping pixels between adjacent patches.
boundary_pix (int) – The number of pixels to crop as a boundary from each patch.
input_interp (Optional[Tensor], optional) – Optional additional data to concatenate to each patch with shape $(B, C_{add}, H_{add}, W_{add})$. Note: ``additional_input`` is not patched or decomposed.

Returns:

A tensor containing the image patches, with shape $(P \times B, C [+ C_{add}], H_p, W_p)$. If additional_input is provided, it is channel-wise concatenated to the extracted patches.

Return type:

Tensor

physicsnemo.utils.patching.image_fuse( input: Tensor, img_shape_y: int, img_shape_x: int, batch_size: int, overlap_pix: int, boundary_pix: int, overlap_count: Tensor | None = None, ) → Tensor[source]#

Reconstructs a full image from a batch of patched images. Reverts the patching operation performed by image_batching().

It assumes that the patches are extracted in a grid-like pattern, and that their layout along the batch dimension is the same as the one returned by image_batching().

This function takes a batch of image patches and reconstructs the full image by stitching the patches together. The function accounts for overlapping and boundary pixels, ensuring that overlapping areas are averaged. Note: a simple unweighted average between overlapping patches is used to fuse the patches.

Parameters:

input (Tensor) – The input tensor containing the image patches with shape $(P \times B, C, H_p, W_p)$.
img_shape_y (int) – The height $H$ of the original full image.
img_shape_x (int) – The width $W$ of the original full image.
batch_size (int) – The original batch size $B$ before patching.
overlap_pix (int) – The number of overlapping pixels between adjacent patches.
boundary_pix (int) – The number of pixels to crop as a boundary from each patch.
overlap_count (Tensor, optional, default=None) – A tensor of shape $(1, 1, H, W)$ containing the number of overlaps for each pixel (i.e. the number of patches that cover each pixel). This is typically computed by get_overlap_count(). If not provided, it will be computed internally.

Returns:

The reconstructed full image tensor with shape $(B, C, H, W)$.

Return type:

Tensor

Domino utils#

Utilities for working with the Domino model, including data processing and grid construction. These utilities are essential for implementing and training Domino-based models.

Utilities for data processing and training with the DoMINO model architecture.

This module provides essential utilities for computational fluid dynamics data processing, mesh manipulation, field normalization, and geometric computations. It supports both CPU (NumPy) and GPU (CuPy) operations with automatic fallbacks.

physicsnemo.utils.domino.utils.area_weighted_shuffle_array( arr: ndarray | ndarray, n_points: int, area: ndarray | ndarray, area_factor: float = 1.0, ) → tuple[ndarray | ndarray, ndarray | ndarray][source]#

Perform area-weighted random sampling from array.

This function samples points from an array with probability proportional to their associated area weights. This is particularly useful in CFD applications where larger cells or surface elements should have higher sampling probability.

Parameters:

arr – Input array to sample from, shape (n_points, …).
n_points – Number of points to sample. If greater than arr.shape[0], samples all available points.
area – Area weights for each point, shape (n_points,). Larger values indicate higher sampling probability.
area_factor – Exponent applied to area weights to control sampling bias. Values > 1.0 increase bias toward larger areas, values < 1.0 reduce bias. Defaults to 1.0 (linear weighting).

Returns:

Sampled array subset weighted by area
Indices of the selected points

Return type:

Tuple containing

Note

For GPU arrays (CuPy), the sampling is performed on CPU due to memory efficiency considerations. The Alias method could be implemented for future GPU acceleration.

Examples

>>> import numpy as np
>>> np.random.seed(42)  # For reproducible results
>>> mesh_data = np.array([[1.0], [2.0], [3.0], [4.0]])
>>> cell_areas = np.array([0.1, 0.1, 0.1, 10.0])  # Last point has much larger area
>>> subset, indices = area_weighted_shuffle_array(mesh_data, 2, cell_areas)
>>> subset.shape
(2, 1)
>>> indices.shape
(2,)
>>> # The point with large area (index 3) should likely be selected
>>> len(set(indices)) <= 2  # At most 2 unique indices
True
>>> # Use higher area_factor for stronger bias toward large areas
>>> subset_biased, _ = area_weighted_shuffle_array(mesh_data, 2, cell_areas, area_factor=2.0)

physicsnemo.utils.domino.utils.array_type( array: ndarray | ndarray, ) → type[<module 'numpy' from '/usr/local/lib/python3.12/dist-packages/numpy/__init__.py'>] | type[<module 'cupy' from '/usr/local/lib/python3.12/dist-packages/cupy/__init__.py'>][source]#

Determine the array module (NumPy or CuPy) for the given array.

This function enables array-agnostic code by returning the appropriate array module that can be used for operations on the input array.

Parameters:: array – Input array that can be either NumPy or CuPy array.
Returns:: The array module (numpy or cupy) corresponding to the input array type.

Examples

>>> import numpy as np
>>> arr = np.array([1, 2, 3])
>>> xp = array_type(arr)
>>> result = xp.sum(arr)  # Uses numpy.sum

physicsnemo.utils.domino.utils.calculate_center_of_mass( centers: ndarray | ndarray, sizes: ndarray | ndarray, ) → ndarray | ndarray[source]#

Calculate the center of mass for a collection of elements.

Computes the volume-weighted centroid of mesh elements, commonly used in computational fluid dynamics for mesh analysis and load balancing.

Parameters:

centers – Array of shape (n_elements, 3) containing the centroid coordinates of each element.
sizes – Array of shape (n_elements,) containing the volume or area of each element used as weights.

Returns:

Array of shape (1, 3) containing the x, y, z coordinates of the center of mass.

Raises:

ValueError – If centers and sizes have incompatible shapes.

Examples

>>> import numpy as np
>>> centers = np.array([[0.0, 0.0, 0.0], [1.0, 1.0, 1.0], [2.0, 2.0, 2.0]])
>>> sizes = np.array([1.0, 2.0, 3.0])
>>> com = calculate_center_of_mass(centers, sizes)
>>> np.allclose(com, [[4.0/3.0, 4.0/3.0, 4.0/3.0]])
True

physicsnemo.utils.domino.utils.calculate_normal_positional_encoding( coordinates_a: ndarray | ndarray, coordinates_b: ndarray | ndarray | None = None, cell_dimensions: Sequence[float] = (1.0, 1.0, 1.0), ) → ndarray | ndarray[source]#

Calculate sinusoidal positional encoding for 3D coordinates.

This function computes transformer-style positional encodings for 3D spatial coordinates, enabling neural networks to understand spatial relationships. The encoding uses sinusoidal functions at different frequencies to create unique representations for each spatial position.

Parameters:

coordinates_a – Primary coordinates array of shape (n_points, 3).
coordinates_b – Optional secondary coordinates for computing relative positions. If provided, the encoding is computed for (coordinates_a - coordinates_b).
cell_dimensions – Characteristic length scales for x, y, z dimensions used for normalization. Defaults to unit dimensions.

Returns:

Array of shape (n_points, 12) containing positional encodings with 4 encoding dimensions per spatial axis (x, y, z).

Examples

>>> import numpy as np
>>> coords = np.array([[0.0, 0.0, 0.0], [1.0, 1.0, 1.0]])
>>> cell_size = [0.1, 0.1, 0.1]
>>> encoding = calculate_normal_positional_encoding(coords, cell_dimensions=cell_size)
>>> encoding.shape
(2, 12)
>>> # Relative positioning example
>>> coords_b = np.array([[0.5, 0.5, 0.5], [0.5, 0.5, 0.5]])
>>> encoding_rel = calculate_normal_positional_encoding(coords, coords_b, cell_size)
>>> encoding_rel.shape
(2, 12)

physicsnemo.utils.domino.utils.calculate_pos_encoding( nx: ndarray | ndarray, d: int = 8, ) → list[ndarray | ndarray][source]#

Calculate sinusoidal positional encoding for transformer architectures.

This function computes positional encodings using alternating sine and cosine functions at different frequencies. These encodings help neural networks understand positional relationships in sequences or spatial data.

Parameters:

nx – Input positions/coordinates to encode.
d – Encoding dimensionality. Must be even number. Defaults to 8.

Returns:

List of d arrays containing alternating sine and cosine encodings. Each pair (sin, cos) uses progressively lower frequencies.

Examples

>>> import numpy as np
>>> positions = np.array([0.0, 1.0, 2.0])
>>> encodings = calculate_pos_encoding(positions, d=4)
>>> len(encodings)
4
>>> all(enc.shape == (3,) for enc in encodings)
True

physicsnemo.utils.domino.utils.combine_dict( old_dict: dict[Any, Any], new_dict: dict[Any, Any], ) → dict[Any, Any][source]#

Combine two dictionaries by adding values for matching keys.

This function performs element-wise addition of dictionary values for keys that exist in both dictionaries. It’s commonly used for accumulating statistics or metrics across multiple iterations.

Parameters:

old_dict – Base dictionary to update.
new_dict – Dictionary with values to add to old_dict.

Returns:

Updated old_dict with combined values.

Note

This function modifies old_dict in place and returns it. Values must support the + operator.

Examples

>>> stats1 = {"loss": 0.5, "accuracy": 0.8}
>>> stats2 = {"loss": 0.3, "accuracy": 0.1}
>>> combined = combine_dict(stats1, stats2)
>>> combined["loss"]
0.8
>>> combined["accuracy"]
0.9

physicsnemo.utils.domino.utils.convert_point_data_to_cell_data( input_data: vtkDataSet, ) → vtkDataSet[source]#

Convert point-based field data to cell-based field data.

This function transforms field variables defined at mesh vertices (nodes) to values defined at cell centers. This conversion is often needed when switching between different numerical methods or visualization requirements.

Parameters:: input_data – VTK dataset with point data to be converted.
Returns:: VTK dataset with the same geometry but field data moved from points to cells. Values are typically averaged from the surrounding points.

physicsnemo.utils.domino.utils.convert_to_tet_mesh( polydata: vtkPolyData, ) → vtkUnstructuredGrid[source]#

Convert surface polydata to a tetrahedral volumetric mesh.

This function performs tetrahedralization of a surface mesh, creating a 3D volumetric mesh suitable for finite element analysis. The process fills the interior of the surface with tetrahedral elements.

Parameters:: polydata – VTK polydata representing a closed surface mesh.
Returns:: VTK unstructured grid containing tetrahedral elements filling the volume enclosed by the input surface.
Raises:: RuntimeError – If tetrahedralization fails (e.g., non-manifold surface).

physicsnemo.utils.domino.utils.create_directory(filepath: str | Path) → None[source]#

Create directory and all necessary parent directories.

This function creates a directory at the specified path, including any necessary parent directories. It’s equivalent to mkdir -p in Unix systems.

Parameters:: filepath – Path to the directory to create. Can be string or Path object.

physicsnemo.utils.domino.utils.create_grid( max_coords: ndarray | ndarray, min_coords: ndarray | ndarray, resolution: ndarray | ndarray, ) → ndarray | ndarray[source]#

Create a 3D regular grid from coordinate bounds and resolution.

This function generates a regular 3D grid spanning from min_coords to max_coords with the specified resolution in each dimension. The resulting grid is commonly used for interpolation, visualization, and regular sampling.

Parameters:

max_coords – Maximum coordinates [x_max, y_max, z_max] for the grid bounds.
min_coords – Minimum coordinates [x_min, y_min, z_min] for the grid bounds.
resolution – Number of grid points [nx, ny, nz] in each dimension.

Returns:

Grid array of shape (nx, ny, nz, 3) containing 3D coordinates for each grid point. The last dimension contains [x, y, z] coordinates.

Examples

>>> import numpy as np
>>> min_bounds = np.array([0.0, 0.0, 0.0])
>>> max_bounds = np.array([1.0, 1.0, 1.0])
>>> grid_res = np.array([2, 2, 2])
>>> grid = create_grid(max_bounds, min_bounds, grid_res)
>>> grid.shape
(2, 2, 2, 3)
>>> np.allclose(grid[0, 0, 0], [0.0, 0.0, 0.0])
True
>>> np.allclose(grid[1, 1, 1], [1.0, 1.0, 1.0])
True

physicsnemo.utils.domino.utils.dict_to_device( state_dict: dict[str, Any], device: Any, exclude_keys: list[str] | None = None, ) → dict[str, Any][source]#

Move dictionary values to specified device (GPU/CPU).

This function transfers PyTorch tensors in a dictionary to the specified compute device while preserving the dictionary structure. It’s commonly used for moving model parameters and data between CPU and GPU.

Parameters:

state_dict – Dictionary containing tensors and other values.
device – Target device (e.g., torch.device(‘cuda:0’)).
exclude_keys – List of keys to skip during device transfer. Defaults to [“filename”] if None.

Returns:

New dictionary with tensors moved to the specified device. Non-tensor values and excluded keys are preserved as-is.

Examples

>>> import torch
>>> data = {"weights": torch.randn(10, 10), "filename": "model.pt"}
>>> gpu_data = dict_to_device(data, torch.device('cuda:0'))

physicsnemo.utils.domino.utils.extract_surface_triangles( tetrahedral_mesh: vtkUnstructuredGrid, ) → list[int][source]#

Extract surface triangle indices from a tetrahedral mesh.

This function identifies the boundary faces of a 3D tetrahedral mesh and returns the vertex indices that form triangular faces on the surface. This is essential for visualization and boundary condition application.

Parameters:: tetrahedral_mesh – VTK unstructured grid containing tetrahedral elements.
Returns:: List of vertex indices forming surface triangles. Every three consecutive indices define one triangle.
Raises:: NotImplementedError – If the surface contains non-triangular faces.

physicsnemo.utils.domino.utils.get_fields( data_attributes: vtkDataSetAttributes, variable_names: list[str], ) → list[ndarray][source]#

Extract multiple field variables from VTK data attributes.

This function extracts field variables from VTK data attributes (either point data or cell data) and returns them as a list of NumPy arrays. It handles both point and cell data seamlessly.

Parameters:

data_attributes – VTK data attributes object (point data or cell data).
variable_names – List of variable names to extract.

Returns:

List of NumPy arrays, one for each requested variable. Each array has shape (n_points/n_cells, n_components) where n_components depends on the variable (1 for scalars, 3 for vectors, etc.).

Raises:

ValueError – If a requested variable is not found in the data attributes.

physicsnemo.utils.domino.utils.get_fields_from_cell( cell_data: vtkCellData, variable_names: list[str], ) → ndarray[source]#

Extract field variables from VTK cell data.

This function extracts multiple field variables from VTK cell data and organizes them into a structured NumPy array. Each variable becomes a column in the output array.

Parameters:

cell_data – VTK cell data object containing field variables.
variable_names – List of variable names to extract from the cell data.

Returns:

NumPy array of shape (n_cells, n_variables) containing the extracted field data. Variables are ordered according to the input list.

Raises:

ValueError – If a requested variable name is not found in the cell data.

physicsnemo.utils.domino.utils.get_filenames( filepath: str | Path, exclude_dirs: bool = False, ) → list[str][source]#

Get list of filenames in a directory with optional directory filtering.

This function returns all items in a directory, with options to exclude subdirectories. It handles special cases like .zarr directories which are treated as files even when exclude_dirs is True.

Parameters:

filepath – Path to the directory to list. Can be string or Path object.
exclude_dirs – If True, exclude subdirectories from results. Exception: .zarr directories are always included as they represent data files in array storage format.

Returns:

List of filenames/directory names found in the specified directory.

Raises:

FileNotFoundError – If the specified directory does not exist.

physicsnemo.utils.domino.utils.get_node_to_elem( polydata: vtkDataSet, ) → vtkDataSet[source]#

Convert point data to cell data for VTK dataset.

This function transforms field variables defined at mesh vertices to values defined at cell centers using VTK’s built-in conversion filter.

Parameters:: polydata – VTK dataset with point data to be converted.
Returns:: VTK dataset with field data moved from points to cells.

physicsnemo.utils.domino.utils.get_surface_data( polydata: vtkPolyData, variable_names: list[str], ) → tuple[ndarray, list[ndarray], list[tuple[int, int]]][source]#

Extract surface mesh data including vertices, fields, and edge connectivity.

This function extracts comprehensive surface mesh information including vertex coordinates, field data at vertices, and edge connectivity information. It’s commonly used for processing CFD surface results and boundary conditions.

Parameters:

polydata – VTK polydata representing a surface mesh.
variable_names – List of field variable names to extract from the mesh.

Returns:

Vertex coordinates as NumPy array of shape (n_vertices, 3)
List of field arrays, one per variable
List of edge tuples representing mesh connectivity

Return type:

Tuple containing

Raises:

ValueError – If a requested variable is not found or polygon data is missing.

physicsnemo.utils.domino.utils.get_vertices( polydata: vtkPolyData, ) → ndarray[source]#

Extract vertex coordinates from VTK polydata object.

This function converts VTK polydata to a NumPy array containing the 3D coordinates of all vertices in the mesh.

Parameters:: polydata – VTK polydata object containing mesh geometry.
Returns:: NumPy array of shape (n_points, 3) containing [x, y, z] coordinates for each vertex.

physicsnemo.utils.domino.utils.get_volume_data( polydata: vtkPolyData, variable_names: list[str], ) → tuple[ndarray, list[ndarray]][source]#

Extract vertices and field data from 3D volumetric mesh.

This function extracts both geometric information (vertex coordinates) and field data from a 3D volumetric mesh. It’s commonly used for processing finite element analysis results.

Parameters:

polydata – VTK polydata representing a 3D volumetric mesh.
variable_names – List of field variable names to extract.

Returns:

Vertex coordinates as NumPy array of shape (n_vertices, 3)
List of field arrays, one per variable

Return type:

Tuple containing

physicsnemo.utils.domino.utils.mean_std_sampling( field: ndarray | ndarray, mean: ndarray | ndarray, std: ndarray | ndarray, tolerance: float = 3.0, ) → list[int][source]#

Identify outlier points based on statistical distance from mean.

This function identifies data points that are statistical outliers by checking if they fall outside mean ± tolerance*std for any field component. It’s useful for data cleaning and identifying regions of interest in CFD data.

Parameters:

field – Input field array of shape (n_points, n_components).
mean – Mean values for each field component, shape (n_components,).
std – Standard deviation for each component, shape (n_components,).
tolerance – Number of standard deviations to use as outlier threshold. Defaults to 3.0 (99.7% of normal distribution).

Returns:

List of indices identifying outlier points that exceed the statistical threshold.

Examples

>>> import numpy as np
>>> # Create test data with outliers
>>> field = np.array([[1.0], [2.0], [3.0], [10.0]])  # 10.0 is outlier
>>> field_mean = np.array([2.0])
>>> field_std = np.array([1.0])
>>> outliers = mean_std_sampling(field, field_mean, field_std, 2.0)
>>> 3 in outliers  # Index 3 (value 10.0) should be detected as outlier
True

physicsnemo.utils.domino.utils.nd_interpolator( coordinates: ndarray | ndarray, field: ndarray | ndarray, grid: ndarray | ndarray, k: int = 2, ) → ndarray | ndarray[source]#

Perform n-dimensional interpolation using k-nearest neighbors.

This function interpolates field values from scattered points to a regular grid using k-nearest neighbor averaging. It’s useful for reconstructing fields on regular grids from irregular measurement points.

Parameters:

coordinates – Array of shape (n_points, n_dims) containing source point coordinates.
field – Array of shape (n_points, n_fields) containing field values at source points.
grid – Array of shape (n_field_points, n_dims) containing target grid points for interpolation.
k – Number of nearest neighbors to use for interpolation.

Returns:

Interpolated field values at grid points using k-nearest neighbor averaging.

Note

This function currently uses SciPy’s KDTree which only supports CPU arrays. A future enhancement could add CuML support for GPU acceleration.

Examples

>>> import numpy as np
>>> # Simple 2D interpolation example
>>> coords = np.array([[0.0, 0.0], [1.0, 0.0], [0.0, 1.0], [1.0, 1.0]])
>>> field_vals = np.array([[1.0], [2.0], [3.0], [4.0]])
>>> grid_points = np.array([[0.5, 0.5]])
>>> result = nd_interpolator([coords], field_vals, grid_points)
>>> result.shape[0] == 1  # One grid point
True

Normalize field values to the range [-1, 1].

Applies min-max normalization to scale field values to a symmetric range around zero. This is commonly used in machine learning preprocessing to ensure numerical stability and faster convergence.

Parameters:

field – Input field array to be normalized.
max_val – Maximum values for normalization, can be scalar or array. If None, computed from the field data.
min_val – Minimum values for normalization, can be scalar or array. If None, computed from the field data.

Returns:

Normalized field with values in the range [-1, 1].

Raises:

ZeroDivisionError – If max_val equals min_val (zero range).

Examples

>>> import numpy as np
>>> field = np.array([1.0, 2.0, 3.0, 4.0, 5.0])
>>> normalized = normalize(field, 5.0, 1.0)
>>> np.allclose(normalized, [-1.0, -0.5, 0.0, 0.5, 1.0])
True
>>> # Auto-compute min/max
>>> normalized_auto = normalize(field)
>>> np.allclose(normalized_auto, [-1.0, -0.5, 0.0, 0.5, 1.0])
True

physicsnemo.utils.domino.utils.pad( arr: ndarray | ndarray, n_points: int, pad_value: float = 0.0, ) → ndarray | ndarray[source]#

Pad 2D array with constant values to reach target size.

This function extends a 2D array by adding rows filled with a constant value. It’s commonly used to standardize array sizes in batch processing for machine learning applications.

Parameters:

arr – Input array of shape (n_points, n_features) to be padded.
n_points – Target number of points (rows) after padding.
pad_value – Constant value used for padding. Defaults to 0.0.

Returns:

Padded array of shape (n_points, n_features). If n_points <= arr.shape[0], returns the original array unchanged.

Examples

>>> import numpy as np
>>> arr = np.array([[1.0, 2.0], [3.0, 4.0]])
>>> padded = pad(arr, 4, -1.0)
>>> padded.shape
(4, 2)
>>> np.array_equal(padded[:2], arr)
True
>>> bool(np.all(padded[2:] == -1.0))
True
>>> # No padding needed
>>> same = pad(arr, 2)
>>> np.array_equal(same, arr)
True

physicsnemo.utils.domino.utils.pad_inp( arr: ndarray | ndarray, n_points: int, pad_value: float = 0.0, ) → ndarray | ndarray[source]#

Pad 3D array with constant values to reach target size.

This function extends a 3D array by adding entries along the first dimension filled with a constant value. Used for standardizing 3D tensor sizes in batch processing workflows.

Parameters:

arr – Input array of shape (n_points, height, width) to be padded.
n_points – Target number of points along first dimension after padding.
pad_value – Constant value used for padding. Defaults to 0.0.

Returns:

Padded array of shape (n_points, height, width). If n_points <= arr.shape[0], returns the original array unchanged.

Examples

>>> import numpy as np
>>> arr = np.array([[[1.0, 2.0]], [[3.0, 4.0]]])
>>> padded = pad_inp(arr, 4, 0.0)
>>> padded.shape
(4, 1, 2)
>>> np.array_equal(padded[:2], arr)
True
>>> bool(np.all(padded[2:] == 0.0))
True

physicsnemo.utils.domino.utils.shuffle_array( arr: ndarray | ndarray, n_points: int, ) → tuple[ndarray | ndarray, ndarray | ndarray][source]#

Randomly sample points from array without replacement.

This function performs random sampling from the input array, selecting n_points points without replacement. It’s commonly used for creating training subsets and data augmentation in machine learning workflows.

Parameters:

arr – Input array to sample from, shape (n_points, …).
n_points – Number of points to sample. If greater than arr.shape[0], all points are returned.

Returns:

Sampled array subset
Indices of the selected points

Return type:

Tuple containing

Examples

>>> import numpy as np
>>> np.random.seed(42)  # For reproducible results
>>> data = np.array([[1, 2], [3, 4], [5, 6], [7, 8]])
>>> subset, indices = shuffle_array(data, 2)
>>> subset.shape
(2, 2)
>>> indices.shape
(2,)
>>> len(np.unique(indices)) == 2  # No duplicates
True

physicsnemo.utils.domino.utils.shuffle_array_without_sampling( arr: ndarray | ndarray, ) → tuple[ndarray | ndarray, ndarray | ndarray][source]#

Shuffle array order without changing the number of elements.

This function reorders all elements in the array randomly while preserving all data points. It’s useful for randomizing data order before training while maintaining the complete dataset.

Parameters:

arr – Input array to shuffle, shape (n_points, …).

Returns:

Shuffled array with same shape as input
Permutation indices used for shuffling

Return type:

Tuple containing

Examples

>>> import numpy as np
>>> np.random.seed(42)  # For reproducible results
>>> data = np.array([[1], [2], [3], [4]])
>>> shuffled, indices = shuffle_array_without_sampling(data)
>>> shuffled.shape
(4, 1)
>>> indices.shape
(4,)
>>> set(indices) == set(range(4))  # All original indices present
True

physicsnemo.utils.domino.utils.solution_weighted_shuffle_array( arr: ndarray | ndarray, n_points: int, solution: ndarray | ndarray, scaling_factor: float = 1.0, ) → tuple[ndarray | ndarray, ndarray | ndarray][source]#

Perform solution-weighted random sampling from array.

This function samples points from an array with probability proportional to their associated solution weights. This is particularly useful in CFD applications where larger cells or surface elements should have higher sampling probability.

Parameters:

arr – Input array to sample from, shape (n_points, …).
n_points – Number of points to sample. If greater than arr.shape[0], samples all available points.
solution – Solution weights for each point, shape (n_points,). Larger values indicate higher sampling probability.
scaling_factor – Exponent applied to solution weights to control sampling bias. Values > 1.0 increase bias toward larger solution fields, values < 1.0 reduce bias. Defaults to 1.0 (linear weighting).

Returns:

Sampled array subset weighted by solution fields
Indices of the selected points

Return type:

Tuple containing

Note

For GPU arrays (CuPy), the sampling is performed on CPU due to memory efficiency considerations. The Alias method could be implemented for future GPU acceleration.

Examples

>>> import numpy as np
>>> np.random.seed(42)  # For reproducible results
>>> mesh_data = np.array([[1.0], [2.0], [3.0], [4.0]])
>>> solution = np.array([0.1, 0.1, 0.1, 10.0])  # Last point has much larger solution field
>>> subset, indices = solution_weighted_shuffle_array(mesh_data, 2, solution)
>>> subset.shape
(2, 1)
>>> indices.shape
(2,)
>>> # The point with large area (index 3) should likely be selected
>>> len(set(indices)) <= 2  # At most 2 unique indices
True
>>> # Use higher scaling_factor for stronger bias toward large solution fields
>>> subset_biased, _ = solution_weighted_shuffle_array(mesh_data, 2, solution, scaling_factor=2.0)

Standardize field values to have zero mean and unit variance.

Applies z-score normalization to center the data around zero with standard deviation of one. This is preferred over min-max normalization when the data follows a normal distribution.

Parameters:

field – Input field array to be standardized.
mean – Mean values for standardization. If None, computed from field data.
std – Standard deviation values for standardization. If None, computed from field data.

Returns:

Standardized field with approximately zero mean and unit variance.

Raises:

ZeroDivisionError – If std contains zeros.

Examples

>>> import numpy as np
>>> field = np.array([1.0, 2.0, 3.0, 4.0, 5.0])
>>> standardized = standardize(field, 3.0, np.sqrt(2.5))
>>> np.allclose(standardized, [-1.265, -0.632, 0.0, 0.632, 1.265], atol=1e-3)
True
>>> # Auto-compute mean/std
>>> standardized_auto = standardize(field)
>>> np.allclose(np.mean(standardized_auto), 0.0)
True
>>> np.allclose(np.std(standardized_auto, ddof=0), 1.0)
True

physicsnemo.utils.domino.utils.unnormalize( normalized_field: ndarray | ndarray, max_val: ndarray | ndarray, min_val: ndarray | ndarray, ) → ndarray | ndarray[source]#

Reverse the normalization process to recover original field values.

Transforms normalized values from the range [-1, 1] back to their original physical range using the stored min/max values.

Parameters:

normalized_field – Field values in the normalized range [-1, 1].
max_val – Maximum values used in the original normalization.
min_val – Minimum values used in the original normalization.

Returns:

Field values restored to their original physical range.

Examples

>>> import numpy as np
>>> normalized = np.array([-1.0, -0.5, 0.0, 0.5, 1.0])
>>> original = unnormalize(normalized, 5.0, 1.0)
>>> np.allclose(original, [1.0, 2.0, 3.0, 4.0, 5.0])
True

physicsnemo.utils.domino.utils.unstandardize( standardized_field: ndarray | ndarray, mean: ndarray | ndarray, std: ndarray | ndarray, ) → ndarray | ndarray[source]#

Reverse the standardization process to recover original field values.

Transforms standardized values (zero mean, unit variance) back to their original distribution using the stored mean and standard deviation.

Parameters:

standardized_field – Field values with zero mean and unit variance.
mean – Mean values used in the original standardization.
std – Standard deviation values used in the original standardization.

Returns:

Field values restored to their original distribution.

Examples

>>> import numpy as np
>>> standardized = np.array([-1.265, -0.632, 0.0, 0.632, 1.265])
>>> original = unstandardize(standardized, 3.0, np.sqrt(2.5))
>>> np.allclose(original, [1.0, 2.0, 3.0, 4.0, 5.0], atol=1e-3)
True

physicsnemo.utils.domino.utils.write_to_vtp( polydata: vtkPolyData, filename: str, ) → None[source]#

Write VTK polydata to a VTP (VTK PolyData) file format.

VTP files are XML-based and store polygonal data including points, polygons, and associated field data. This format is commonly used for surface meshes in computational fluid dynamics visualization.

Parameters:

polydata – VTK polydata object containing mesh geometry and fields.
filename – Output filename with .vtp extension. Directory will be created if it doesn’t exist.

Raises:

RuntimeError – If writing fails due to file permissions or disk space.

physicsnemo.utils.domino.utils.write_to_vtu( unstructured_grid: vtkUnstructuredGrid, filename: str, ) → None[source]#

Write VTK unstructured grid to a VTU (VTK Unstructured Grid) file format.

VTU files store 3D volumetric meshes with arbitrary cell types including tetrahedra, hexahedra, and pyramids. This format is essential for storing finite element analysis results.

Parameters:

unstructured_grid – VTK unstructured grid object containing volumetric mesh geometry and field data.
filename – Output filename with .vtu extension. Directory will be created if it doesn’t exist.

Raises:

RuntimeError – If writing fails due to file permissions or disk space.

CorrDiff utils#

Utilities for working with the CorrDiff model, particularly for the diffusion and regression steps.

class physicsnemo.utils.corrdiff.utils.NetCDFWriter( f, lat, lon, input_channels, output_channels, has_lead_time=False, )[source]#

Bases: object

NetCDF Writer

write_input(channel_name, time_index, val)[source]#: Write input data to NetCDF file.

write_prediction( channel_name, time_index, ensemble_index, val, )[source]#: Write prediction data to NetCDF file.

write_time(time_index, time)[source]#: Write time information to NetCDF file.

write_truth(channel_name, time_index, val)[source]#: Write ground truth data to NetCDF file.

physicsnemo.utils.corrdiff.utils.diffusion_step( net: Module, sampler_fn: callable, img_shape: tuple, img_out_channels: int, rank_batches: list, img_lr: Tensor, rank: int, device: device, mean_hr: Tensor = None, lead_time_label: Tensor = None, distribution: Literal['normal', 'student_t'] = 'normal', nu: int | None = None, ) → Tensor[source]#

Generate images using diffusion techniques as described in the relevant paper.

This function applies a diffusion model to generate high-resolution images based on low-resolution inputs. It supports optional conditioning on high-resolution mean predictions and lead time labels.

For each low-resolution sample in img_lr, the function generates multiple high-resolution samples, with different random seeds, specified in rank_batches. The function then concatenates these high-resolution samples across the batch dimension.

Parameters:

net (torch.nn.Module) – The diffusion model network.
sampler_fn (callable) – Function used to sample images from the diffusion model.
img_shape (tuple) – Shape of the images, (height, width).
img_out_channels (int) – Number of output channels for the image.
rank_batches (list) – List of batches of seeds to process.
img_lr (torch.Tensor) – Low-resolution input image with shape (seed_batch_size, channels_lr, height, width).
rank (int, optional) – Rank of the current process for distributed processing.
device (torch.device, optional) – Device to perform computations.
mean_hr (torch.Tensor, optional) – High-resolution mean tensor to be used as an additional input, with shape (1, channels_hr, height, width). Default is None.
lead_time_label (torch.Tensor, optional) – Lead time label tensor for temporal conditioning, with shape (batch_size, lead_time_dims). Default is None.
distribution (Literal["normal", "student_t"], optional, default="normal") – Distribution to use for the latent state.
nu (int, optional, default=None) – Degrees of freedom for the Student-t distribution. Only used if distribution is “student_t”. Must be provided and greater than 2 in this case.

Returns:

Generated images concatenated across batches with shape (seed_batch_size * len(rank_batches), out_channels, height, width).

Return type:

torch.Tensor

physicsnemo.utils.corrdiff.utils.get_time_from_range(times_range, time_format='%Y-%m-%dT%H:%M:%S')[source]#

Generates a list of times within a given range.

Parameters:

times_range – A list containing start time, end time, and optional interval (hours).
time_format – The format of the input times (default: “%Y-%m-%dT%H:%M:%S”).

Returns:

A list of times within the specified range.

physicsnemo.utils.corrdiff.utils.regression_step( net: Module, img_lr: Tensor, latents_shape: Size, lead_time_label: Tensor | None = None, ) → Tensor[source]#

Perform a regression step to produce ensemble mean prediction.

This function takes a low-resolution input and performs a regression step to produce an ensemble mean prediction. It processes a single instance and then replicates the results across the batch dimension if needed.

Parameters:

net (torch.nn.Module) – U-Net model for regression.
img_lr (torch.Tensor) – Low-resolution input to the network with shape (1, channels, height, width). Must have a batch dimension of 1.
latents_shape (torch.Size) – Shape of the latent representation with format (batch_size, out_channels, image_shape_y, image_shape_x).
lead_time_label (Optional[torch.Tensor], optional) – Lead time label tensor for lead time conditioning, with shape (1, lead_time_dims). Default is None.

Returns:

Predicted ensemble mean at the next time step with shape matching latents_shape.

Return type:

torch.Tensor

Raises:

ValueError – If img_lr has a batch size greater than 1.

Profiling utils#

Utilities for profiling the performance of a model.