deepmd.dpmodel.descriptor.dpa4c_nn.readout

deepmd.dpmodel.descriptor.dpa4c_nn.readout#

Fixed invariant readout for degree-wise DPA4C moments.

Classes#

InvariantReadout

Contract degree-wise moments into fixed O(3)-invariant features.

Module Contents#

class deepmd.dpmodel.descriptor.dpa4c_nn.readout.InvariantReadout(channels: int, lmax: int, *, precision: str = DEFAULT_PRECISION, trainable: bool = True, seed: int | list[int] | None = None)[source]#

Bases: deepmd.dpmodel.NativeOP

Contract degree-wise moments into fixed O(3)-invariant features.

The readout carries no learned nonlinearity. Degrees one and two first pass through full-width residual channel maps, after which three fixed contractions are emitted: the exact channel Gram of every non-scalar degree, a Cartesian bispectrum over every O(3)-even degree triple, and the projected quartic |Q_b v_a|^2. The scalar moments are prepended unchanged.

Parameters:

channels: Scalar degree-zero channel width.
lmax: Maximum angular degree.
precision: Parameter precision.
trainable: Whether the channel-alignment and probe projections are trainable.
seed: Random seed reserved for the readout.

degree_channels[source]#

bispectrum_ranks[source]#

channels[source]#

lmax[source]#

precision = ''[source]#

trainable = True[source]#

degree_offsets = (0,)[source]#

channel_alignment = [][source]#

gram_index[source]#

gram_scale[source]#

gram_offsets = (0,)[source]#

degree_triples[source]#

coupling_offsets[source]#

probe_offsets[source]#

bispectrum_coupling[source]#

probe_index[source]#

probe_scale[source]#

probe_projections = [][source]#

bispectrum_dim[source]#

quartic_dim[source]#

call(moments: Any) → Any[source]#

Build invariant features from flat degree-wise moments.

The output concatenates the scalar moments, the exact aligned Grams in degree order, the bispectrum in degree-triple order, and the projected quartic.

Parameters:

moments: Flat moment tensor with shape (N, S), where S = sum((2 * l + 1) * degree_channels[l]).

Returns:

Any: Invariant features with shape (N, get_dim_out()).

build_grams(blocks: list[Any], xp: Any) → list[Any][source]#

Build Frobenius-isometric exact channel Grams.

Parameters:

blocks: Aligned degree blocks. Entry l has shape (N, 2 * l + 1, degree_channels[l]).
xp: Array namespace associated with blocks.

Returns:

list[Any]: Upper-triangular Gram blocks for degrees one through lmax.

build_bispectrum(projected: list[Any], xp: Any) → tuple[list[Any], Any][source]#

Contract projected moments with fixed Cartesian Gaunt tensors.

Parameters:

projected: Probe-projected degree blocks. Entry l - 1 has shape (N, 2 * l + 1, bispectrum_ranks[l - 1]).
xp: Array namespace associated with projected.

Returns:

bispectrum: Independent cubic contractions grouped by degree triple.
quartic: Projected |Q_b v_a|^2 values with shape (N, K_2 * K_1).

contract_vector_tensor(vector: Any, packed_tensor: Any, xp: Any) → tuple[Any, Any][source]#

Contract the 112 triple and reuse Q_b v_a for the quartic.

Both outputs are built from the same matrix-vector intermediate \(Q_bv_a\), so the quartic costs one extra reduction rather than a second contraction.

Parameters:

vector: Degree-one probes with shape (N, 3, K_1).
packed_tensor: Degree-two probes with shape (N, 5, K_2).
xp: Array namespace associated with the probes.

Returns:

bispectrum_112: Full ordered cubic contractions with shape (N, K_1 * K_1 * K_2).
quartic: Values \(|Q_bv_a|^2\) with shape (N, K_2 * K_1).

contract_bispectrum(coupling: Any, value_1: Any, value_2: Any, value_3: Any, xp: Any) → Any[source]#

Contract one angular coupling without backend-specific einsum.

Parameters:

coupling: Cartesian Gaunt tensor with shape (2 * l1 + 1, 2 * l2 + 1, 2 * l3 + 1).
value_1: First degree block with shape (N, 2 * l1 + 1, K_1).
value_2: Second degree block with shape (N, 2 * l2 + 1, K_2).
value_3: Third degree block with shape (N, 2 * l3 + 1, K_3).
xp: Array namespace associated with the degree blocks.

Returns:

Any: Ordered contractions with shape (N, K_1 * K_2 * K_3).

get_dim_out() → int[source]#

Return the geometric output width.

Returns:

int: Scalar moments, aligned exact Grams, bispectrum probes, and the projected quartic.

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

Serialize the invariant readout.

Returns:

dict[str, Any]: Versioned readout configuration and trainable projections.

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

Deserialize an InvariantReadout.

Parameters:

data: Versioned dictionary produced by serialize().

Returns:

InvariantReadout: Reconstructed readout with restored projections.