Embeddings#

class physicsnemo.nn.module.embedding_layers.FourierEmbedding( num_channels: int, scale: int = 16, amp_mode: bool = False, )[source]#

Bases: Module

Generates Fourier embeddings for timesteps, primarily used in the NCSN++ architecture.

This class generates embeddings by first multiplying input tensor x and internally stored random frequencies, and then concatenating the cosine and sine of the resultant.

Parameters:#

num_channelsint: The number of channels in the embedding. The final embedding size will be 2 * num_channels because of concatenation of cosine and sine results.
scaleint, optional: A scale factor applied to the random frequencies, controlling their range and thereby the frequency of oscillations in the embedding space. By default 16.
amp_modebool, optional: A boolean flag indicating whether mixed-precision (AMP) training is enabled. Defaults to False.

forward(x)[source]#

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

class physicsnemo.nn.module.embedding_layers.OneHotEmbedding(*args, **kwargs)[source]#

Bases: Module

Soft one-hot embedding for normalized timesteps in \([0, 1]\).

The \(i\)-th channel is \(\max(0, 1 - |t(D-1) - i|)\), giving a soft one-hot vector of dimension \(D\). Used for timestep conditioning in ModAFNO when method="onehot".

Parameters:: num_channels (int) – Number of channels (embedding dimension).
Forward:: t (torch.Tensor) – Input tensor, shape B ..., with normalized timesteps in [0, 1].
Outputs:: torch.Tensor – Output tensor, shape B D.

Examples

>>> import torch
>>> from physicsnemo.nn.module.embedding_layers import OneHotEmbedding
>>> emb = OneHotEmbedding(num_channels=64)
>>> t = torch.tensor([[0.0], [0.5], [1.0]])
>>> out = emb(t)
>>> out.shape
torch.Size([3, 64])

forward( t: Float[Tensor, 'B ...'], ) → Float[Tensor, 'B D'][source]#: Forward pass computing soft one-hot embeddings.

class physicsnemo.nn.module.embedding_layers.PositionalEmbedding( num_channels: int, max_positions: int = 10000, endpoint: bool = False, amp_mode: bool = False, learnable: bool = False, freq_embed_dim: int | None = None, mlp_hidden_dim: int | None = None, embed_fn: Literal['cos_sin', 'np_sin_cos'] = 'cos_sin', )[source]#

Bases: Module

A module for generating positional embeddings based on timesteps. This embedding technique is employed in the DDPM++ and ADM architectures.

Parameters:#

num_channelsint

Number of channels for the embedding.

max_positionsint, optional

Maximum number of positions for the embeddings, by default 10000.

endpointbool, optional

If True, the embedding considers the endpoint. By default False.

amp_modebool, optional

A boolean flag indicating whether mixed-precision (AMP) training is enabled. Defaults to False.

learnablebool, optional

A boolean flag indicating whether learnable positional embedding is enabled. Defaults to False.

freq_embed_dim: int, optional

The dimension of the frequency embedding. Defaults to None, in which case it will be set to num_channels.

mlp_hidden_dim: int, optional

The dimension of the hidden layer in the MLP. Defaults to None, in which case it will be set to 2 * num_channels. Only applicable if learnable is True; if learnable is False, this parameter is ignored.

embed_fn: Literal[“cos_sin”, “np_sin_cos”], optional

The function to use for embedding into sin/cos features (allows for swapping the order of sin/cos). Defaults to ‘cos_sin’. Options:

‘cos_sin’: Uses torch to compute frequency embeddings and returns in order (cos, sin)

‘np_sin_cos’: Uses numpy to compute frequency embeddings and returns in order (sin, cos)

forward(x)[source]#

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

class physicsnemo.nn.module.embedding_layers.SinusoidalTimestepEmbedding(*args, **kwargs)[source]#

Bases: Module

Sinusoidal embedding for timesteps (e.g. for modulation / diffusion).

For input \(x\) (timestep) and \(D =\) num_channels (even), the frequencies are \(\omega_k = k\pi\) for \(k = 1, \ldots, D/2\), and the output is the concatenation of cosines and sines:

\[\mathrm{embed}(x) = \big[ \cos(x \omega_1), \ldots, \cos(x \omega_{D/2}), \sin(x \omega_1), \ldots, \sin(x \omega_{D/2}) \big] \in \mathbb{R}^D.\]

This is a simpler scheme than PositionalEmbedding (which uses geometric frequencies and optional learnable MLP) which works well for timestep conditioning in ModAFNO model.

Parameters:: num_channels (int) – Number of output channels. Must be even.
Forward:: x (torch.Tensor) – Input tensor, shape B ... (e.g. \((B,)\) or \((B, 1)\)), containing timesteps.
Outputs:: torch.Tensor – Output tensor, shape B D with \(D =\) num_channels.

Examples

>>> import torch
>>> from physicsnemo.nn.module.embedding_layers import SinusoidalTimestepEmbedding
>>> emb = SinusoidalTimestepEmbedding(num_channels=64)
>>> t = torch.tensor([0.0, 0.5, 1.0])
>>> out = emb(t)
>>> out.shape
torch.Size([3, 64])