deepmd.pt.model.descriptor.sezm_nn.embedding

deepmd.pt.model.descriptor.sezm_nn.embedding#

Embedding layers for the SeZM descriptor.

This module defines the type embedding, geometric initial embedding, and environment-seed embedding used to initialize SeZM node features.

Classes#

`SeZMTypeEmbedding`	Minimal SeZM type embedding with Adam-routed parameter naming.
`GeometricInitialEmbedding`	Geometric initial embedding that adds zonal (m=0) rotated features.
`EnvironmentInitialEmbedding`	Environment matrix initial embedding for l=0 features.
`ChargeSpinEmbedding`	Frame-level charge and spin embedding for scalar type features.

Module Contents#

class deepmd.pt.model.descriptor.sezm_nn.embedding.SeZMTypeEmbedding(*, ntypes: int, embed_dim: int, dtype: torch.dtype, seed: int | list[int] | None = None, trainable: bool, padding: bool = True)[source]#

Bases: torch.nn.Module

Minimal SeZM type embedding with Adam-routed parameter naming.

Parameters:

ntypes: Number of atom types.
embed_dim: Embedding dimension.
dtype: Parameter dtype.
seed: Random seed for initialization.
trainable: Whether parameters are trainable.
padding: Whether to append one all-zero padding row.

Notes

The parameter is named with adam_ prefix so HybridMuon routes it to Adam.

ntypes[source]#

embed_dim[source]#

dtype[source]#

seed = None[source]#

device[source]#

padding = True[source]#

adam_type_embedding[source]#

forward(atype: torch.Tensor) → torch.Tensor[source]#

Gather type embeddings.

Parameters:

atype: Atom types with shape (…,). Valid type range is [0, ntypes-1].

Returns:

torch.Tensor: Type embeddings with shape (…, embed_dim).

class deepmd.pt.model.descriptor.sezm_nn.embedding.GeometricInitialEmbedding(*, lmax: int, channels: int, dtype: torch.dtype)[source]#

Bases: torch.nn.Module

Geometric initial embedding that adds zonal (m=0) rotated features.

This module rotates pre-computed radial features for each degree l >= 1 using the zonal (m=0) column of the cached inverse Wigner-D blocks (local->global). The l=0 component is not computed here since it comes from type embedding.

Parameters:

lmax: Maximum degree, should match l_schedule[0].
channels: Number of channels per (l, m) coefficient.
dtype: Parameter dtype.

lmax[source]#

channels[source]#

ebed_dim = 1[source]#

dtype[source]#

device[source]#

precision[source]#

forward(*, n_nodes: int, edge_cache: deepmd.pt.model.descriptor.sezm_nn.edge_cache.EdgeFeatureCache, radial_feat: torch.Tensor) → torch.Tensor[source]#

Parameters:

n_nodes: Number of nodes (nf*nloc).
edge_cache: Per-edge cache containing geometry, weights, and Wigner-D blocks.
radial_feat: Per-edge radial features with shape (E, lmax, C) for l=1..lmax.

Returns:

torch.Tensor: Initial features to add with shape (N, D, C). l=0 is guaranteed zero.

serialize() → dict[str, Any][source]#

classmethod deserialize(data: dict[str, Any]) → GeometricInitialEmbedding[source]#

class deepmd.pt.model.descriptor.sezm_nn.embedding.EnvironmentInitialEmbedding(*, ntypes: int, n_radial: int, channels: int, embed_dim: int = 64, axis_dim: int = 8, type_dim: int = 16, hidden_dim: int = 64, mlp_bias: bool = False, activation_function: str = 'silu', eps: float = 1e-07, dtype: torch.dtype, trainable: bool, seed: int | list[int] | None = None)[source]#

Bases: torch.nn.Module

Environment matrix initial embedding for l=0 features.

Computes an initial embedding based on the 4D environment matrix:

[s, s * rx, s * ry, s * rz]

Combined with independent type embeddings (individual type embedding), providing physical inductive bias for l=0 features.

The computation follows the environment matrix approach where:

Build `r_tilde = [s, s*r_hat]` where `s = edge_env / r` and `r_hat = edge_vec / r`
G network: `g = G(rbf_proj(edge_rbf), type_src, type_dst)` produces per-edge features
   - Uses independent `env_type_embed` instead of projecting from main type embedding
   - Uses `rbf_proj` to project edge_rbf to `rbf_out_dim`
env_agg: aggregate outer product `r_tilde ⊗ g` by destination node
D matrix: `D = env_agg^T @ env_agg[:, :, :axis_dim]`
Output: projection of flattened D matrix into FiLM logits

Parameters:

ntypesint: Number of atom types.
n_radialint: Number of radial basis functions.
channelsint: Output channel dimension per FiLM branch (final output is 2*channels).
embed_dimint: G network output dimension (filter width).
axis_dimint: D matrix axis dimension (must be < embed_dim).
type_dimint: Dimension for independent type embeddings in env_seed.
hidden_dimint: Hidden layer size for G network.
mlp_biasbool: Whether to enable bias terms in env-seed MLP layers (rbf_proj_layer1/2 and g_layer1/2).
activation_functionstr: Activation function for G network hidden layer.
epsfloat: Small epsilon for numerical stability.
dtypetorch.dtype: Parameter dtype.
trainablebool: Whether parameters are trainable.
seedint | list[int] | None: Random seed for reproducibility.

ntypes[source]#

n_radial[source]#

channels[source]#

embed_dim = 64[source]#

axis_dim = 8[source]#

type_dim = 16[source]#

hidden_dim = 64[source]#

mlp_bias = False[source]#

activation_function = ''[source]#

eps[source]#

dtype[source]#

device[source]#

precision[source]#

rbf_out_dim[source]#

rbf_proj_layer1[source]#

rbf_proj_layer2[source]#

env_type_embed[source]#

g_layer1[source]#

g_layer2[source]#

output_proj[source]#

forward(*, edge_cache: deepmd.pt.model.descriptor.sezm_nn.edge_cache.EdgeFeatureCache, atype_flat: torch.Tensor, n_nodes: int) → torch.Tensor[source]#

Compute environment FiLM logits for l=0 conditioning.

Parameters:

edge_cacheEdgeFeatureCache: Edge cache containing src, dst, edge_vec, edge_rbf, edge_env.
atype_flattorch.Tensor: Flattened atom types with shape (N,), where N = nf * nloc.
n_nodesint: Number of nodes (N = nf * nloc).

Returns:

torch.Tensor: FiLM logits with shape (N, 2*channels).

serialize() → dict[str, Any][source]#

classmethod deserialize(data: dict[str, Any]) → EnvironmentInitialEmbedding[source]#: Deserialize from dictionary.

class deepmd.pt.model.descriptor.sezm_nn.embedding.ChargeSpinEmbedding(*, embed_dim: int, activation_function: str, dtype: torch.dtype, seed: int | list[int] | None = None, trainable: bool)[source]#