deepmd_gnn.mace

Wrapper for MACE models.

Attributes

ELEMENTS
PeriodicTable

Classes

MaceModel

Mace model.

Functions

_is_freeze_command(→ bool)	Return whether the current process is running a DeePMD freeze command.
_is_make_fx_fake_tensor(→ bool)

Module Contents

deepmd_gnn.mace._is_freeze_command() → bool

Return whether the current process is running a DeePMD freeze command.

deepmd_gnn.mace.ELEMENTS = ['H', 'He', 'Li', 'Be', 'B', 'C', 'N', 'O', 'F', 'Ne', 'Na', 'Mg', 'Al', 'Si', 'P', 'S', 'Cl',...

deepmd_gnn.mace.PeriodicTable

deepmd_gnn.mace._is_make_fx_fake_tensor(value: Any) → bool

class deepmd_gnn.mace.MaceModel(type_map: list[str], sel: int, r_max: float = 5.0, num_radial_basis: int = 8, num_cutoff_basis: int = 5, max_ell: int = 3, interaction: str = 'RealAgnosticResidualInteractionBlock', num_interactions: int = 2, hidden_irreps: str = '128x0e + 128x1o', pair_repulsion: bool = False, distance_transform: str = 'None', correlation: int = 3, gate: str = 'silu', MLP_irreps: str = '16x0e', radial_type: str = 'bessel', radial_MLP: list[int] = [64, 64, 64], std: float = 1, avg_num_neighbors: float | None = None, enable_cueq: bool = False, **kwargs: Any)

Bases: deepmd.pt.model.model.model.BaseModel

Mace model.

Parameters

type_maplist[str]

The name of each type of atoms

selint

Maximum number of neighbor atoms

r_maxfloat, optional

distance cutoff (in Ang)

num_radial_basisint, optional

number of radial basis functions

num_cutoff_basisint, optional

number of basis functions for smooth cutoff

max_ellint, optional

highest ell of spherical harmonics

interactionstr, optional

name of interaction block

num_interactionsint, optional

number of interactions

hidden_irrepsstr, optional

hidden irreps

pair_repulsionbool

use amsgrad variant of optimizer

distance_transformstr, optional

distance transform

correlationint

correlation order at each layer

gatestr, optional

non linearity for last readout

MLP_irrepsstr, optional

hidden irreps of the MLP in last readout

radial_typestr, optional

type of radial basis functions

radial_MLPstr, optional

width of the radial MLP

stdfloat, optional

Standard deviation of force components in the training set

mm_types: list[int]

_observed_type: list[str] | None

_use_exportable_edge_index: bool

_use_exportable_border_op: bool

params: dict[str, Any]

_disable_cueq_for_freeze

type_map

ntypes

rcut = 5.0

num_interactions = 2

preset_out_bias: dict[str, list]

sel

avg_num_neighbors

model

atomic_numbers = []

_enable_cueq_for_runtime() → bool

_params_for_serialization() → dict[str, Any]

_sync_freeze_model_def_script() → None

_make_mace_network_from_params(*, enable_cueq: bool, script_model: bool) → torch.nn.Module

_make_raw_mace_network_from_current_state() → torch.nn.Module

_replace_with_scripted_e3nn_model_for_freeze(raw_model: torch.nn.Module | None = None) → torch.nn.Module

_adapt_cueq_state_for_freeze(state_dict: collections.abc.MutableMapping[str, Any], prefix: str) → None

_load_from_state_dict(state_dict: collections.abc.MutableMapping[str, Any], prefix: str, local_metadata: collections.abc.MutableMapping[str, Any], strict: bool, missing_keys: list[str], unexpected_keys: list[str], error_msgs: list[str]) → None

Copy parameters and buffers from state_dict into only this module, but not its descendants.

This is called on every submodule in load_state_dict(). Metadata saved for this module in input state_dict is provided as local_metadata. For state dicts without metadata, local_metadata is empty. Subclasses can achieve class-specific backward compatible loading using the version number at local_metadata.get(“version”, None). Additionally, local_metadata can also contain the key assign_to_params_buffers that indicates whether keys should be assigned their corresponding tensor in the state_dict.

Note

state_dict is not the same object as the input state_dict to load_state_dict(). So it can be modified.

Args:

state_dict (dict): a dict containing parameters and

persistent buffers.

prefix (str): the prefix for parameters and buffers used in this

module

local_metadata (dict): a dict containing the metadata for this module.

See

strict (bool): whether to strictly enforce that the keys in

state_dict with prefix match the names of parameters and buffers in this module

missing_keys (list of str): if strict=True, add missing keys to

this list

unexpected_keys (list of str): if strict=True, add unexpected

keys to this list

error_msgs (list of str): error messages should be added to this

list, and will be reported together in load_state_dict()

property atomic_model: Any

Provide a compatibility view matching wrapped deepmd-kit models.

property descriptor: Any

Provide descriptor metadata expected by deepmd-kit’s pt2 export.

get_descriptor() → Any

Return descriptor metadata expected by DeePMD-kit trainers.

_ensure_compile_trace_model() → torch.nn.Module

_forward_lower_common_compile_trace(nloc: int, extended_coord: torch.Tensor, extended_atype: torch.Tensor, nlist: torch.Tensor, mapping: torch.Tensor | None, fparam: torch.Tensor | None, aparam: torch.Tensor | None, do_atomic_virial: bool, comm_dict: dict[str, torch.Tensor] | None) → dict[str, torch.Tensor]

property observed_type: list[str] | None

Observed element types collected during statistics.

get_observed_type_list() → list[str]

Get observed element types collected during statistics.

compute_or_load_stat(sampled_func, stat_file_path: deepmd.utils.path.DPPath | None = None, preset_observed_type: list[str] | None = None) → None

Compute or load the statistics parameters of the model.

For example, mean and standard deviation of descriptors or the energy bias of the fitting net. When sampled is provided, all the statistics parameters will be calculated (or re-calculated for update), and saved in the stat_file_path`(s). When `sampled is not provided, it will check the existence of `stat_file_path`(s) and load the calculated statistics parameters.

Parameters

sampled_func

The sampled data frames from different data systems.

stat_file_path

The path to the statistics files.

preset_observed_type

Optional observed element types to seed or override self._observed_type. This compatibility parameter is accepted for newer deepmd-kit versions; when provided, it is used directly instead of restoring or collecting observed types from statistics data.

fitting_output_def() → deepmd.dpmodel.output_def.FittingOutputDef

Get the output def of developer implemented atomic models.

atomic_output_def() → deepmd.dpmodel.output_def.FittingOutputDef

Get the atomic output def used by the exportable backend.

get_rcut() → float

Get the cut-off radius.

get_type_map() → list[str]

Get the type map.

get_sel() → list[int]

Return the number of selected atoms for each type.

get_dim_fparam() → int

Get the number (dimension) of frame parameters of this atomic model.

has_default_fparam() → bool

Return whether default frame parameters are available.

get_default_fparam() → torch.Tensor | None

Get the default frame parameters.

get_dim_aparam() → int

Get the number (dimension) of atomic parameters of this atomic model.

has_chg_spin_ebd() → bool

Return whether charge-spin embedding is enabled.

get_dim_chg_spin() → int

Get the dimension of charge-spin input.

has_default_chg_spin() → bool

Return whether default charge-spin values are available.

get_default_chg_spin() → torch.Tensor | None

Get the default charge-spin values.

get_sel_type() → list[int]

Get the selected atom types of this model.

Only atoms with selected atom types have atomic contribution to the result of the model. If returning an empty list, all atom types are selected.

is_aparam_nall() → bool

Check whether the shape of atomic parameters is (nframes, nall, ndim).

If False, the shape is (nframes, nloc, ndim).

mixed_types() → bool

Return whether the model is in mixed-types mode.

If true, the model 1. assumes total number of atoms aligned across frames; 2. uses a neighbor list that does not distinguish different atomic types. If false, the model 1. assumes total number of atoms of each atom type aligned across frames; 2. uses a neighbor list that distinguishes different atomic types.

has_message_passing() → bool

Return whether the descriptor has message passing.

has_message_passing_across_ranks() → bool

Return whether MPI ranks must exchange layer-wise ghost features.

need_sorted_nlist_for_lower() → bool

Return whether lower-interface neighbor lists must be sorted.

forward(coord: torch.Tensor, atype: torch.Tensor, box: torch.Tensor | None = None, fparam: torch.Tensor | None = None, aparam: torch.Tensor | None = None, charge_spin: torch.Tensor | None = None, do_atomic_virial: bool = False) → dict[str, torch.Tensor]

Forward pass of the model.

Parameters

coordtorch.Tensor

The coordinates of atoms.

atypetorch.Tensor

The atomic types of atoms.

boxtorch.Tensor, optional

The box tensor.

fparamtorch.Tensor, optional

The frame parameters.

aparamtorch.Tensor, optional

The atomic parameters.

do_atomic_virialbool, optional

Whether to compute atomic virial.

forward_lower(extended_coord: torch.Tensor, extended_atype: torch.Tensor, nlist: torch.Tensor, mapping: torch.Tensor | None = None, fparam: torch.Tensor | None = None, aparam: torch.Tensor | None = None, do_atomic_virial: bool = False, comm_dict: dict[str, torch.Tensor] | None = None, charge_spin: torch.Tensor | None = None) → dict[str, torch.Tensor]

Forward lower pass of the model.

Parameters

extended_coordtorch.Tensor

The extended coordinates of atoms.

extended_atypetorch.Tensor

The extended atomic types of atoms.

nlisttorch.Tensor

The neighbor list.

mappingtorch.Tensor, optional

The mapping tensor.

fparamtorch.Tensor, optional

The frame parameters.

aparamtorch.Tensor, optional

The atomic parameters.

do_atomic_virialbool, optional

Whether to compute atomic virial.

comm_dictdict[str, torch.Tensor], optional

The communication dictionary.

forward_lower_common(nloc: int, extended_coord: torch.Tensor, extended_atype: torch.Tensor, nlist: torch.Tensor, mapping: torch.Tensor | None = None, box: torch.Tensor | None = None, fparam: torch.Tensor | None = None, aparam: torch.Tensor | None = None, do_atomic_virial: bool = False, comm_dict: dict[str, torch.Tensor] | None = None) → dict[str, torch.Tensor]

Forward lower common pass of the model.

Parameters

extended_coordtorch.Tensor

The extended coordinates of atoms.

extended_atypetorch.Tensor

The extended atomic types of atoms.

nlisttorch.Tensor

The neighbor list.

mappingtorch.Tensor, optional

The mapping tensor.

fparamtorch.Tensor, optional

The frame parameters.

aparamtorch.Tensor, optional

The atomic parameters.

do_atomic_virialbool, optional

Whether to compute atomic virial.

comm_dictdict[str, torch.Tensor], optional

The communication dictionary.

forward_mace_exportable(nf: int, nall: int, edge_index: torch.Tensor, node_attrs: torch.Tensor, positions: torch.Tensor, shifts: torch.Tensor) → torch.Tensor

Evaluate node energies through traceable MACE submodules.

forward_common_lower(extended_coord: torch.Tensor, extended_atype: torch.Tensor, nlist: torch.Tensor, mapping: torch.Tensor | None = None, fparam: torch.Tensor | None = None, aparam: torch.Tensor | None = None, charge_spin: torch.Tensor | None = None, do_atomic_virial: bool = False, comm_dict: dict[str, torch.Tensor] | None = None) → dict[str, torch.Tensor]

Forward lower pass with internal DeePMD output names.

forward_common_lower_exportable(extended_coord: torch.Tensor, extended_atype: torch.Tensor, nlist: torch.Tensor, mapping: torch.Tensor | None = None, fparam: torch.Tensor | None = None, aparam: torch.Tensor | None = None, charge_spin: torch.Tensor | None = None, do_atomic_virial: bool = False, **make_fx_kwargs: object) → torch.nn.Module

Trace forward_common_lower for torch.export serialization.

forward_common_lower_exportable_with_comm(extended_coord: torch.Tensor, extended_atype: torch.Tensor, nlist: torch.Tensor, mapping: torch.Tensor | None, fparam: torch.Tensor | None, aparam: torch.Tensor | None, charge_spin: torch.Tensor | None, send_list: torch.Tensor, send_proc: torch.Tensor, recv_proc: torch.Tensor, send_num: torch.Tensor, recv_num: torch.Tensor, communicator: torch.Tensor, nlocal: torch.Tensor, nghost: torch.Tensor, do_atomic_virial: bool = False, **make_fx_kwargs: object) → torch.nn.Module

Trace forward_common_lower with explicit MPI comm tensors.

communicate_node_features(node_feats: torch.Tensor, nloc: int, nall: int, comm_dict: dict[str, torch.Tensor]) → torch.Tensor

Communicate local node features to ghost slots using DeePMD border_op.

forward_mace_with_comm(nloc: int, nall: int, edge_index: torch.Tensor, node_attrs: torch.Tensor, positions: torch.Tensor, shifts: torch.Tensor, comm_dict: dict[str, torch.Tensor]) → torch.Tensor

Evaluate MACE node energies with layer-wise MPI ghost communication.

serialize() → dict

Serialize the model.

classmethod deserialize(data: dict) → MaceModel

Deserialize the model.

get_nnei() → int

Return the total number of selected neighboring atoms in cut-off radius.

get_nsel() → int

Return the total number of selected neighboring atoms in cut-off radius.

classmethod update_sel(train_data: deepmd.utils.data_system.DeepmdDataSystem, type_map: list[str] | None, local_jdata: dict) → tuple[dict, float | None]

Update the selection and perform neighbor statistics.

Parameters

train_dataDeepmdDataSystem

data used to do neighbor statictics

type_maplist[str], optional

The name of each type of atoms

local_jdatadict

The local data refer to the current class

Returns

dict

The updated local data

float

The minimum distance between two atoms

model_output_type() → list[str]

Get the output type for the model.

translated_output_def() → dict[str, Any]

Get the translated output def for the model.

model_output_def() → deepmd.dpmodel.output_def.ModelOutputDef

Get the output def for the model.

classmethod get_model(model_params: dict) → MaceModel

Get the model by the parameters.

Parameters

model_paramsdict

The model parameters

Returns

BaseBaseModel

The model