deepmd_gnn.mace

deepmd_gnn.mace#

Wrapper for MACE models.

Attributes#

`ELEMENTS`
`PeriodicTable`

Classes#

MaceModel

Mace model.

Functions#

_load_observed_type_stat_compat(→ tuple[Any, Any, Any])

Module Contents#

deepmd_gnn.mace._load_observed_type_stat_compat() → tuple[Any, Any, Any][source]#

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

deepmd_gnn.mace.PeriodicTable[source]#

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, **kwargs: Any)[source]#

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][source]#

_observed_type: list[str] | None[source]#

params[source]#

type_map[source]#

ntypes[source]#

rcut = 5.0[source]#

num_interactions = 2[source]#

preset_out_bias: dict[str, list][source]#

sel[source]#

avg_num_neighbors[source]#

model[source]#

atomic_numbers = [][source]#

property atomic_model: Any[source]#: Provide a compatibility view matching wrapped deepmd-kit models.

property observed_type: list[str] | None[source]#: Observed element types collected during statistics.

get_observed_type_list() → list[str][source]#: 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[source]#

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[source]#: Get the output def of developer implemented atomic models.

get_rcut() → float[source]#: Get the cut-off radius.

get_type_map() → list[str][source]#: Get the type map.

get_sel() → list[int][source]#: Return the number of selected atoms for each type.

get_dim_fparam() → int[source]#: Get the number (dimension) of frame parameters of this atomic model.

get_dim_aparam() → int[source]#: Get the number (dimension) of atomic parameters of this atomic model.

get_sel_type() → list[int][source]#

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[source]#

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

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

mixed_types() → bool[source]#

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[source]#: Return whether the descriptor has message passing.

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

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) → dict[str, torch.Tensor][source]#

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, 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][source]#

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.

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

classmethod deserialize(data: dict) → MaceModel[source]#: Deserialize the model.

get_nnei() → int[source]#: Return the total number of selected neighboring atoms in cut-off radius.

get_nsel() → int[source]#: 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][source]#

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][source]#: Get the output type for the model.

translated_output_def() → dict[str, Any][source]#: Get the translated output def for the model.

model_output_def() → deepmd.dpmodel.output_def.ModelOutputDef[source]#: Get the output def for the model.

classmethod get_model(model_params: dict) → MaceModel[source]#

Get the model by the parameters.

Parameters#

model_paramsdict: The model parameters

Returns#

BaseBaseModel: The model