deepmd.dpmodel.model#
The model that takes the coordinates, cell and atom types as input and predicts some property. The models are automatically generated from atomic models by the deepmd.dpmodel.make_model method.
The make_model method does the reduction, auto-differentiation (dummy for dpmodels) and communication of the atomic properties according to output variable definition deepmd.dpmodel.OutputVariableDef.
All models should be inherited from deepmd.dpmodel.model.base_model.BaseModel
. Models generated by make_model have already done it.
Submodules#
- deepmd.dpmodel.model.base_model
- deepmd.dpmodel.model.dipole_model
- deepmd.dpmodel.model.dos_model
- deepmd.dpmodel.model.dp_model
- deepmd.dpmodel.model.dp_zbl_model
- deepmd.dpmodel.model.ener_model
- deepmd.dpmodel.model.make_model
- deepmd.dpmodel.model.model
- deepmd.dpmodel.model.polar_model
- deepmd.dpmodel.model.property_model
- deepmd.dpmodel.model.spin_model
- deepmd.dpmodel.model.transform_output
Classes#
The unit operation of a native model. | |
The unit operation of a native model. | |
A spin model wrapper, with spin input preprocess and output split. |
Functions#
| Make a model as a derived class of an atomic model. |
Package Contents#
- class deepmd.dpmodel.model.DPModelCommon[source]#
- class deepmd.dpmodel.model.EnergyModel(*args, **kwargs)[source]#
Bases:
deepmd.dpmodel.model.dp_model.DPModelCommon
,DPEnergyModel_
The unit operation of a native model.
- deepmd.dpmodel.model.make_model(T_AtomicModel: type[deepmd.dpmodel.atomic_model.base_atomic_model.BaseAtomicModel])[source]#
Make a model as a derived class of an atomic model.
The model provide two interfaces.
1. the call_lower, that takes extended coordinates, atyps and neighbor list, and outputs the atomic and property and derivatives (if required) on the extended region.
2. the call, that takes coordinates, atypes and cell and predicts the atomic and reduced property, and derivatives (if required) on the local region.
- Parameters:
- T_AtomicModel
The atomic model.
- Returns:
CM
The model.
- class deepmd.dpmodel.model.PropertyModel(*args, **kwargs)[source]#
Bases:
deepmd.dpmodel.model.dp_model.DPModelCommon
,DPPropertyModel_
The unit operation of a native model.
- class deepmd.dpmodel.model.SpinModel(backbone_model, spin: deepmd.utils.spin.Spin)[source]#
Bases:
deepmd.dpmodel.common.NativeOP
A spin model wrapper, with spin input preprocess and output split.
- backbone_model#
- spin#
- ntypes_real#
- virtual_scale_mask#
- spin_mask#
- process_spin_input(coord, atype, spin)[source]#
Generate virtual coordinates and types, concat into the input.
- process_spin_input_lower(extended_coord: numpy.ndarray, extended_atype: numpy.ndarray, extended_spin: numpy.ndarray, nlist: numpy.ndarray, mapping: numpy.ndarray | None = None)[source]#
Add extended_spin into extended_coord to generate virtual atoms, and extend nlist and mapping. Note that the final extended_coord_updated with shape [nframes, nall + nall, 3] has the following order: - [:, :nloc]: original nloc real atoms. - [:, nloc: nloc + nloc]: virtual atoms corresponding to nloc real atoms. - [:, nloc + nloc: nloc + nall]: ghost real atoms. - [:, nloc + nall: nall + nall]: virtual atoms corresponding to ghost real atoms.
- process_spin_output(atype, out_tensor, add_mag: bool = True, virtual_scale: bool = True)[source]#
Split the output both real and virtual atoms, and scale the latter.
- process_spin_output_lower(extended_atype, extended_out_tensor, nloc: int, add_mag: bool = True, virtual_scale: bool = True)[source]#
Split the extended output of both real and virtual atoms with switch, and scale the latter.
- static expand_aparam(aparam, nloc: int)[source]#
Expand the atom parameters for virtual atoms if necessary.
- 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).
- get_nnei() int [source]#
Returns the total number of selected neighboring atoms in the cut-off radius.
- get_nsel() int [source]#
Returns the total number of selected neighboring atoms in the cut-off radius.
- call(coord, atype, spin, box: numpy.ndarray | None = None, fparam: numpy.ndarray | None = None, aparam: numpy.ndarray | None = None, do_atomic_virial: bool = False) dict[str, numpy.ndarray] [source]#
Return model prediction.
- Parameters:
- coord
The coordinates of the atoms. shape: nf x (nloc x 3)
- atype
The type of atoms. shape: nf x nloc
- spin
The spins of the atoms. shape: nf x (nloc x 3)
- box
The simulation box. shape: nf x 9
- fparam
frame parameter. nf x ndf
- aparam
atomic parameter. nf x nloc x nda
- do_atomic_virial
If calculate the atomic virial.
- Returns:
ret_dict
The result dict of type dict[str,np.ndarray]. The keys are defined by the ModelOutputDef.
- call_lower(extended_coord: numpy.ndarray, extended_atype: numpy.ndarray, extended_spin: numpy.ndarray, nlist: numpy.ndarray, mapping: numpy.ndarray | None = None, fparam: numpy.ndarray | None = None, aparam: numpy.ndarray | None = None, do_atomic_virial: bool = False)[source]#
Return model prediction. Lower interface that takes extended atomic coordinates, types and spins, nlist, and mapping as input, and returns the predictions on the extended region. The predictions are not reduced.
- Parameters:
- extended_coord
coordinates in extended region. nf x (nall x 3).
- extended_atype
atomic type in extended region. nf x nall.
- extended_spin
spins in extended region. nf x (nall x 3).
- nlist
neighbor list. nf x nloc x nsel.
- mapping
maps the extended indices to local indices. nf x nall.
- fparam
frame parameter. nf x ndf
- aparam
atomic parameter. nf x nloc x nda
- do_atomic_virial
whether calculate atomic virial
- Returns:
result_dict
the result dict, defined by the FittingOutputDef.
- forward_lower#