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.
SpinEmbedding	Per-atom spin embedding for the native spin scheme.

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)

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

embed_dim

dtype

seed = None

device

padding = True

adam_type_embedding

forward(atype: torch.Tensor) → torch.Tensor

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)

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 node degree for the initial embedding.

channels

Number of channels per (l, m) coefficient.

dtype

Parameter dtype.

lmax

channels

ebed_dim = 1

dtype

device

precision

forward(*, n_nodes: int, edge_cache: deepmd.pt.model.descriptor.sezm_nn.edge_cache.EdgeFeatureCache, radial_feat: torch.Tensor, zonal_coupling: torch.Tensor | None = None, spin_l1_message: torch.Tensor | None = None) → torch.Tensor

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.

zonal_coupling

Optional precomputed zonal coupling with shape (E, D-1). If None, it is gathered from edge_cache.Dt_full.

spin_l1_message

Optional per-edge neighbor-spin l=1 message with shape (E, 3, C) for the native spin scheme (built by SpinEmbedding.edge_l1). It is added to the l=1 rows of the per-edge message, so it shares this module’s source gate, scatter and degree normalization with the geometric message.

Returns:

torch.Tensor

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

serialize() → dict[str, Any]

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

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, use_spin: list[bool] | None = None, dtype: torch.dtype, trainable: bool, seed: int | list[int] | None = None)

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:

1. Build `r_tilde = [s, s*r_hat]` where `s = edge_env / r` and `r_hat = edge_vec / r`
2. 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`
3. env_agg: aggregate outer product `r_tilde ⊗ g` by destination node
4. D matrix: `D = env_agg^T @ env_agg[:, :, :axis_dim]`
5. 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.

use_spinlist[bool] | None

Per-type spin flags (native spin scheme). When provided, the neighbor spin is appended as extra coordinate channels of the environment matrix, so the inner product D = M^T M additionally yields the neighbor spin-spin invariants. A per-type mask gates the channel, so a non-magnetic neighbor contributes zero and carries zero magnetic force.

dtypetorch.dtype

Parameter dtype.

trainablebool

Whether parameters are trainable.

seedint | list[int] | None

Random seed for reproducibility.

ntypes

n_radial

channels

embed_dim = 64

axis_dim = 8

type_dim = 16

hidden_dim = 64

mlp_bias = False

activation_function = ''

eps

spin_flags = None

dtype

device

precision

coord_dim = 7

rbf_out_dim

rbf_proj_layer1

rbf_proj_layer2

env_type_embed

g_layer1

g_layer2

output_proj

forward(*, edge_cache: deepmd.pt.model.descriptor.sezm_nn.edge_cache.EdgeFeatureCache, atype_flat: torch.Tensor, n_nodes: int, spin: torch.Tensor | None = None) → torch.Tensor

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).

spintorch.Tensor | None

Per-node spin vectors with shape (N, 3) for the native spin scheme. Used only when use_spin is set; the source (neighbor) spin is appended to the environment matrix as an envelope-gated coordinate channel. When None the spin channels are zero-padded so the coordinate dimension stays fixed.

Returns:

torch.Tensor

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

serialize() → dict[str, Any]

classmethod deserialize(data: dict[str, Any]) → EnvironmentInitialEmbedding

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)

Bases: torch.nn.Module

Frame-level charge and spin embedding for scalar type features.

Parameters:

embed_dim

Embedding dimension.

activation_function

Activation function used by the mixing layer.

dtype

Parameter dtype.

seed

Random seed for initialization.

trainable

Whether parameters are trainable.

embed_dim

activation_function = ''

dtype

precision

charge_embedding

spin_embedding

mix_layer

forward(charge_spin: torch.Tensor) → torch.Tensor

Embed frame-level charge and spin.

Parameters:

charge_spin

Frame charge and spin values with shape (nf, 2).

Returns:

torch.Tensor

Mixed condition embedding with shape (nf, embed_dim).

class deepmd.pt.model.descriptor.sezm_nn.embedding.SpinEmbedding(*, ntypes: int, channels: int, use_spin: list[bool], activation_function: str = 'silu', dtype: torch.dtype, seed: int | list[int] | None = None, trainable: bool = True)

Bases: torch.nn.Module

Per-atom spin embedding for the native spin scheme.

The per-atom spin vector s is injected as an equivariant extension of the type embedding, producing two additive contributions to the descriptor node features:

• l = 0 (invariant): a small network of the squared magnitude |s|^2 yields a per-channel scalar added to the scalar type embedding. The squared magnitude is used (rather than |s|) so the feature is smooth at s = 0 and its gradient there vanishes, keeping the magnetic force continuous as a spin crosses zero.

• l = 1 (equivariant): the Cartesian spin vector is mapped to the packed l = 1 coefficients through the SeZM Wigner-D convention (derived from build_cartesian_basis()), then scaled by a per-type per-channel weight. The map is linear in s, so the contribution vanishes at s = 0 and rotates as an l = 1 object under SO(3), i.e. cart_to_l1(R s) = D^1(R) cart_to_l1(s).

Both contributions are gated by a per-type spin mask, so atom types without spin contribute exactly zero regardless of their (nominally zero) input.

Parameters:

ntypes

Number of (real) atom types.

channels

Number of channels per (l, m) coefficient.

use_spin

Per-type boolean flags marking which atom types carry spin.

activation_function

Activation used by the magnitude network.

dtype

Parameter dtype.

seed

Random seed for initialization.

trainable

Whether parameters are trainable.

ntypes

channels

activation_function = ''

dtype

device

precision

mag_layer1

mag_layer2

adam_spin_vec_weight

adam_spin_nbr_weight

forward(spin: torch.Tensor, atype: torch.Tensor) → tuple[torch.Tensor, torch.Tensor]

Compute the l=0 and l=1 spin contributions.

Parameters:

spin

Per-atom spin vectors with shape (N, 3).

atype

Per-atom types with shape (N,).

Returns:

tuple[torch.Tensor, torch.Tensor]

(scalar, vector) where scalar has shape (N, channels) for the l=0 contribution and vector has shape (N, 3, channels) for the packed l=1 contribution (orders m = -1, 0, +1). Both are exactly zero for atom types without spin.

edge_l1(spin: torch.Tensor, atype: torch.Tensor, edge_cache: deepmd.pt.model.descriptor.sezm_nn.edge_cache.EdgeFeatureCache) → torch.Tensor

Build the per-edge neighbor-spin l=1 message for the GIE aggregation.

Each edge carries the packed l = 1 coefficients of the source (neighbor) spin, scaled by a per-source-type per-channel weight and gated by the C^3 envelope. The message is returned per edge; the geometric initial embedding folds it into the l=1 rows and applies the shared source gate, scatter and degree normalization, so a neighbor’s spin direction enters an atom’s l=1 backbone before any interaction block (the spin analogue of the geometric initial embedding).

Parameters:

spin

Per-node spin vectors with shape (N, 3).

atype

Per-node types with shape (N,).

edge_cache

Edge cache providing src and edge_env.

Returns:

torch.Tensor

Per-edge packed l=1 message with shape (E, 3, channels), exactly zero for non-magnetic neighbors.

_build_cart_to_l1_matrix() → torch.Tensor

Build the (3, 3) Cartesian-to-packed-l=1 projection.

The packed l = 1 coefficient of a vector v is obtained by projecting the skew-symmetric matrix [v]_x onto the antisymmetric l = 1 block of build_cartesian_basis(). With packed order m = -1, 0, +1, row d and Cartesian component k give M[d, k] = <[e_k]_x, B[1 + d]>_F, so coeff = M @ v and M @ (R v) = D^1(R) (M @ v).