deepmd.dpmodel.model.spin_model

Module Contents

Classes

SpinModel

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

class deepmd.dpmodel.model.spin_model.SpinModel(backbone_model, spin: deepmd.utils.spin.Spin)[source]

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

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 extend_nlist(extended_atype, nlist)[source]
static concat_switch_virtual(extended_tensor, extended_tensor_virtual, nloc: int)[source]
get_type_map() List[str][source]

Get the type map.

get_rcut()[source]

Get the cut-off radius.

get_dim_fparam()[source]

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

get_dim_aparam()[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).

model_output_type() List[str][source]

Get the output type for the model.

get_model_def_script() str[source]

Get the model definition script.

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.

static has_spin() bool[source]

Returns whether it has spin input and output.

__getattr__(name)[source]

Get attribute from the wrapped model.

serialize() dict[source]
classmethod deserialize(data) SpinModel[source]
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.