deepmd.dpmodel.fitting.dpa4_ener

deepmd.dpmodel.fitting.dpa4_ener#

SeZM (DPA4) GLU energy fitting network, dpmodel implementation.

Mirrors deepmd.pt.model.task.sezm_ener with the array-API call convention and the pt-state_dict-key serialization contract.

Classes#

`GLUFittingNet`	GLU-based fitting network for SeZM.
`SeZMNetworkCollection`	Network collection for SeZM fitting networks.
`SeZMEnergyFittingNet`	SeZM energy fitting with GLU hidden layers.

Functions#

_resolve_auto_neuron(→ list[int])

Resolve SeZM fitting hidden widths, using 0 as the auto-width marker.

Module Contents#

class deepmd.dpmodel.fitting.dpa4_ener.GLUFittingNet(in_dim: int, out_dim: int, neuron: list[int] | None = None, activation_function: str = 'silu', resnet_dt: bool = False, precision: str = DEFAULT_PRECISION, bias_out: bool = False, seed: int | list[int] | None = None, trainable: bool | list[bool] = True, descriptor_dim: int | None = None, dim_case_embd: int = 0, case_film_embd: bool = False)[source]#

Bases: deepmd.dpmodel.NativeOP

GLU-based fitting network for SeZM.

Parameters:

in_dim: Input dimension.
out_dim: Output dimension.
neuron: Hidden layer sizes. Empty list means direct linear projection.
activation_function: Activation function used for GLU gating.
resnet_dt: Reserved for compatibility; not used in GLU layers.
precision: Numerical precision.
bias_out: Whether the output layer uses bias.
seed: Random seed.
trainable: Whether parameters are trainable.
descriptor_dim: Descriptor feature width. Kept for serialization compatibility with the case-FiLM path (not implemented here).
dim_case_embd: Case one-hot width.
case_film_embd: Whether to use case FiLM instead of input concatenation. Not implemented in the dpmodel backend.

in_dim[source]#

out_dim[source]#

neuron[source]#

activation_function = 'silu'[source]#

resnet_dt = False[source]#

precision = 'float64'[source]#

bias_out = False[source]#

descriptor_dim[source]#

dim_case_embd = 0[source]#

case_film_embd = False[source]#

hidden_layers = [][source]#

output_layer[source]#

call_until_last(xx: deepmd.dpmodel.array_api.Array) → deepmd.dpmodel.array_api.Array[source]#: Return activations before the output projection.

call(xx: deepmd.dpmodel.array_api.Array) → deepmd.dpmodel.array_api.Array[source]#: Forward pass for the GLU fitting net.

serialize() → dict[str, Any][source]#: Serialize the network to a dict (pt state_dict key contract).

classmethod deserialize(data: dict) → GLUFittingNet[source]#: Deserialize the network from a dict.

class deepmd.dpmodel.fitting.dpa4_ener.SeZMNetworkCollection(ndim: int, ntypes: int, network_type: str = 'sezm_fitting_network', networks: list[Any] | None = None)[source]#

Network collection for SeZM fitting networks.

Parameters:

ndim: The number of type dimensions.
ntypes: Number of atom types.
network_type: The network type name. Only “sezm_fitting_network” is supported.
networks: The networks to initialize with.

NETWORK_TYPE_MAP: ClassVar[dict[str, type]][source]#

ndim[source]#

ntypes[source]#

network_type[source]#

_networks: list[GLUFittingNet | None][source]#

networks[source]#

_convert_key(key: int | tuple | str) → int[source]#

__getitem__(key: int | tuple | str) → GLUFittingNet[source]#

__setitem__(key: int | tuple | str, value: Any) → None[source]#

serialize() → dict[str, Any][source]#: Serialize the networks to a dict.

classmethod deserialize(data: dict) → SeZMNetworkCollection[source]#: Deserialize the networks from a dict.

deepmd.dpmodel.fitting.dpa4_ener._resolve_auto_neuron(neuron: list[int] | None, *, dim_descrpt: int, numb_fparam: int, numb_aparam: int, dim_case_embd: int, case_film_embd: bool, use_aparam_as_mask: bool) → list[int][source]#: Resolve SeZM fitting hidden widths, using 0 as the auto-width marker.

class deepmd.dpmodel.fitting.dpa4_ener.SeZMEnergyFittingNet(ntypes: int, dim_descrpt: int, neuron: list[int] | None = None, bias_atom_e: deepmd.dpmodel.array_api.Array | None = None, resnet_dt: bool = False, numb_fparam: int = 0, numb_aparam: int = 0, dim_case_embd: int = 0, case_film_embd: bool = False, activation_function: str = 'silu', bias_out: bool = False, precision: str = 'float32', mixed_types: bool = True, seed: int | list[int] | None = None, type_map: list[str] | None = None, default_fparam: list | None = None, **kwargs: Any)[source]#

Bases: deepmd.dpmodel.fitting.invar_fitting.InvarFitting

SeZM energy fitting with GLU hidden layers.

This uses the same configuration keys as the standard energy fitting but replaces hidden MLP layers with GLU blocks.

seed = None[source]#

bias_out = False[source]#

case_film_embd[source]#

_build_glu_fitting_layers() → None[source]#

classmethod deserialize(data: dict) → SeZMEnergyFittingNet[source]#

Deserialize the fitting.

Parameters:

datadict: The serialized data

Returns:

BF: The deserialized fitting

serialize() → dict[source]#: Serialize the fitting to dict.