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 and 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.core.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.core.Module. I.e. has no meta data.

Example

>>> # Create model
>>> import physicsnemo
>>> 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.core.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.core.Module. I.e. has no meta data.

Example

>>> # Create model
>>> import physicsnemo
>>> 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.

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.

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.

Diffusion utils#

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

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.

param net:: The diffusion model to use in the sampling process.
type net:: torch.nn.Module
param latents:: The latent random noise used as the initial condition for the stochastic ODE.
type latents:: torch.Tensor
param img_lr:: Low-resolution input image for conditioning the diffusion process. Passed as a keywork argument to the model net.
type img_lr:: torch.Tensor
param class_labels:: 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.
type class_labels:: Optional[torch.Tensor]
param randn_like:: 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.
type randn_like:: Callable
param patching:: 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.
type patching:: Optional[GridPatching2D], default=None
param mean_hr:: 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.
type mean_hr:: Optional[Tensor], optional
param lead_time_label:: Lead-time labels to pass to the model, shape (batch_size,). If not provided, the model is called without a lead-time label input.
type lead_time_label:: Optional[Tensor], optional
param num_steps:: Number of time-steps for the stochastic ODE integration. Defaults to 18.
type num_steps:: Optional[int]
param sigma_min:: 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.
type sigma_min:: Optional[float]
param sigma_max:: 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.
type sigma_max:: Optional[float]
param rho:: 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.
type rho:: float, optional
param solver:: 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".
type solver:: Literal[“heun”, “euler”]
param discretization:: The method to discretize time-steps \(t_i\) in the diffusion process. See the EDM paper for details. Defaults to "edm".
type discretization:: Literal[“vp”, “ve”, “iddpm”, “edm”]
param schedule:: 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".
type schedule:: Literal[“vp”, “ve”, “linear”]
param scaling:: 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\).
type scaling:: Literal[“vp”, “none”]
param epsilon_s:: 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.
type epsilon_s:: float, optional
param C_1:: Parameters to compute the time-step discetization. Only used when discretization="iddpm". Defaults to 0.001.
type C_1:: float, optional
param C_2:: Same as for C_1. Only used when discretization="iddpm". Defaults to 0.008.
type C_2:: float, optional
param M:: Same as for C_1 and C_2. Only used when discretization="iddpm". Defaults to 1000.
type M:: int, optional
param alpha:: 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.
type alpha:: float, optional
param S_churn:: 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.
type S_churn:: int, optional
param S_min:: 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.
type S_min:: float, optional
param S_max:: See S_min. Defaults to float("inf").
type S_max:: float, optional
param S_noise:: 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.
type S_noise:: float, optional
param dtype:: Controls the precision used for sampling
type dtype:: torch.dtype, optional
returns:: Generated batch of samples. Same shape as the input latents.
rtype:: torch.Tensor

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

param net:

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

type net:

torch.nn.Module

param latents:

The latent variables (e.g., noise) used as the initial input for the sampler. Has shape \((B, C_{out}, H, W)\).

type latents:

Tensor

param img_lr:

Low-resolution input image for conditioning the super-resolution process. Must have shape \((B, C_{lr}, H, W)\).

type img_lr:

Tensor

param class_labels:

Class labels for conditional generation, if required by the model. By default None.

type class_labels:

Optional[Tensor], optional

param randn_like:

Function to generate random noise with the same shape as the input tensor. By default torch.randn_like.

type randn_like:

Callable[[Tensor], Tensor]

param patching:

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.

type patching:

Optional[GridPatching2D], default=None

param mean_hr:

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.

type mean_hr:

Optional[Tensor], optional

param lead_time_label:

Optional lead time labels. By default None.

type lead_time_label:

Optional[Tensor], optional

param num_steps:

Number of time steps for the sampler. By default 18.

type num_steps:

int

param sigma_min:

Minimum noise level. By default 0.002.

type sigma_min:

float

param sigma_max:

Maximum noise level. By default 800.

type sigma_max:

float

param rho:

Exponent used in the time step discretization. By default 7.

type rho:

float

param S_churn:

Churn parameter controlling the level of noise added in each step. By default 0.

type S_churn:

float

param S_min:

Minimum time step for applying churn. By default 0.

type S_min:

float

param S_max:

Maximum time step for applying churn. By default float("inf").

type S_max:

float

param S_noise:

Noise scaling factor applied during the churn step. By default 1.

type S_noise:

float

returns:

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

rtype:

Tensor

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 that are 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.

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.

CorrDiff utils#

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

Profiling utils#

Utilities for profiling the performance of a model.