deepmd.pt.model.model.spin_model

Module Contents

Classes

SpinModel	A spin model wrapper, with spin input preprocess and output split.
SpinEnergyModel	A spin model for energy.

class deepmd.pt.model.model.spin_model.SpinModel(backbone_model, spin: deepmd.utils.spin.Spin)

Bases: torch.nn.Module

A spin model wrapper, with spin input preprocess and output split.

process_spin_input(coord, atype, spin)

Generate virtual coordinates and types, concat into the input.

process_spin_input_lower(extended_coord, extended_atype, extended_spin, nlist, mapping: torch.Tensor | None = None)

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)

Split the output both real and virtual atoms, and scale the latter. add_mag: whether to add magnetic tensor onto the real tensor.

Default: True. e.g. Ture for forces and False for atomic virials on real atoms.

virtual_scale: whether to scale the magnetic tensor with virtual scale factor.

Default: True. e.g. Ture for forces and False for atomic virials on virtual atoms.

process_spin_output_lower(extended_atype, extended_out_tensor, nloc: int, add_mag: bool = True, virtual_scale: bool = True)

Split the extended output of both real and virtual atoms with switch, and scale the latter. add_mag: whether to add magnetic tensor onto the real tensor.

Default: True. e.g. Ture for forces and False for atomic virials on real atoms.

virtual_scale: whether to scale the magnetic tensor with virtual scale factor.

Default: True. e.g. Ture for forces and False for atomic virials on virtual atoms.

static extend_nlist(extended_atype, nlist)

static concat_switch_virtual(extended_tensor, extended_tensor_virtual, nloc: int)

Concat real and virtual extended tensors, and switch all the local ones to the first nloc * 2 atoms. - [:, :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.

static expand_aparam(aparam, nloc: int)

Expand the atom parameters for virtual atoms if necessary.

get_type_map() → List[str]

Get the type map.

get_rcut()

Get the cut-off radius.

get_dim_fparam()

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

get_dim_aparam()

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).

model_output_type() → List[str]

Get the output type for the model.

get_model_def_script() → str

Get the model definition script.

get_nnei() → int

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

get_nsel() → int

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

has_spin() → bool

Returns whether it has spin input and output.

__getattr__(name)

Get attribute from the wrapped model.

compute_or_load_stat(sampled_func, stat_file_path: deepmd.utils.path.DPPath | None = None)

Compute or load the statistics parameters of the model, such as 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 lazy sampled function to get data frames from different data systems.

stat_file_path

The dictionary of paths to the statistics files.

forward_common(coord, atype, spin, 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_common_lower(extended_coord, extended_atype, extended_spin, nlist, mapping: torch.Tensor | None = None, fparam: torch.Tensor | None = None, aparam: torch.Tensor | None = None, do_atomic_virial: bool = False)

serialize() → dict

classmethod deserialize(data) → SpinModel

class deepmd.pt.model.model.spin_model.SpinEnergyModel(backbone_model, spin: deepmd.utils.spin.Spin)

Bases: SpinModel

A spin model for energy.

model_type = 'ener'

forward(coord, atype, spin, 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_lower(extended_coord, extended_atype, extended_spin, nlist, mapping: torch.Tensor | None = None, fparam: torch.Tensor | None = None, aparam: torch.Tensor | None = None, do_atomic_virial: bool = False)