deepmd_gnn.nequip

Nequip model.

Classes

NequipModel

Nequip model.

Functions

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

Module Contents

deepmd_gnn.nequip._load_observed_type_stat_compat() → tuple[Any, Any, Any]

class deepmd_gnn.nequip.NequipModel(type_map: list[str], sel: int, r_max: float = 6.0, num_layers: int = 4, l_max: int = 2, num_features: int = 32, nonlinearity_type: str = 'gate', parity: bool = True, num_basis: int = 8, BesselBasis_trainable: bool = True, PolynomialCutoff_p: int = 6, invariant_layers: int = 2, invariant_neurons: int = 64, use_sc: bool = True, irreps_edge_sh: str = '0e + 1e', feature_irreps_hidden: str = '32x0o + 32x0e + 32x1o + 32x1e', chemical_embedding_irreps_out: str = '32x0e', conv_to_output_hidden_irreps_out: str = '16x0e', precision: str = 'float32', **kwargs: Any)

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

Nequip 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_layersint

number of interaction blocks, we find 3-5 to work best

l_maxint

the maximum irrep order (rotation order) for the network’s features, l=1 is a good default, l=2 is more accurate but slower

num_featuresint

the multiplicity of the features, 32 is a good default for accurate network, if you want to be more accurate, go larger, if you want to be faster, go lower

nonlinearity_typestr

may be ‘gate’ or ‘norm’, ‘gate’ is recommended

paritybool

whether to include features with odd mirror parityy; often turning parity off gives equally good results but faster networks, so do consider this

num_basisint

number of basis functions used in the radial basis, 8 usually works best

BesselBasis_trainablebool

set true to train the bessel weights

PolynomialCutoff_pint

p-exponent used in polynomial cutoff function, smaller p corresponds to stronger decay with distance

invariant_layersint

number of radial layers, usually 1-3 works best, smaller is faster

invariant_neuronsint

number of hidden neurons in radial function, smaller is faster

use_scbool

use self-connection or not, usually gives big improvement

irreps_edge_shstr

irreps for the chemical embedding of species

feature_irreps_hiddenstr

irreps used for hidden features, here we go up to lmax=1, with even and odd parities; for more accurate but slower networks, use l=2 or higher, smaller number of features is faster

chemical_embedding_irreps_outstr

irreps of the spherical harmonics used for edges. If a single integer, indicates the full SH up to L_max=that_integer

conv_to_output_hidden_irreps_outstr

irreps used in hidden layer of output block

mm_types: list[int]

e0: torch.Tensor

_observed_type: list[str] | None

params

type_map

ntypes

preset_out_bias: dict[str, list]

sel

num_layers = 4

rcut = 6.0

model

property atomic_model: Any

Provide a compatibility view matching wrapped deepmd-kit models.

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.

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.

get_dim_aparam() → int

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

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.

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]

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]

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, box: 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.

boxtorch.Tensor, optional

The box tensor.

serialize() → dict

Serialize the model.

classmethod deserialize(data: dict) → NequipModel

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.