deepmd.dpmodel.model.spin_model =============================== .. py:module:: deepmd.dpmodel.model.spin_model Classes ------- .. autoapisummary:: deepmd.dpmodel.model.spin_model.SpinModel Module Contents --------------- .. py:class:: SpinModel(backbone_model, spin: deepmd.utils.spin.Spin) Bases: :py:obj:`deepmd.dpmodel.common.NativeOP` A spin model wrapper, with spin input preprocess and output split. .. !! processed by numpydoc !! .. py:attribute:: backbone_model .. py:attribute:: spin .. py:attribute:: ntypes_real .. py:attribute:: virtual_scale_mask .. py:attribute:: spin_mask .. py:method:: process_spin_input(coord, atype, spin) Generate virtual coordinates and types, concat into the input. .. !! processed by numpydoc !! .. py:method:: process_spin_input_lower(extended_coord: numpy.ndarray, extended_atype: numpy.ndarray, extended_spin: numpy.ndarray, nlist: numpy.ndarray, mapping: Optional[numpy.ndarray] = 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. .. !! processed by numpydoc !! .. py:method:: 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. .. !! processed by numpydoc !! .. py:method:: 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. .. !! processed by numpydoc !! .. py:method:: extend_nlist(extended_atype, nlist) :staticmethod: .. py:method:: concat_switch_virtual(extended_tensor, extended_tensor_virtual, nloc: int) :staticmethod: .. py:method:: expand_aparam(aparam, nloc: int) :staticmethod: Expand the atom parameters for virtual atoms if necessary. .. !! processed by numpydoc !! .. py:method:: get_type_map() -> list[str] Get the type map. .. !! processed by numpydoc !! .. py:method:: get_ntypes() Returns the number of element types. .. !! processed by numpydoc !! .. py:method:: get_rcut() Get the cut-off radius. .. !! processed by numpydoc !! .. py:method:: get_dim_fparam() Get the number (dimension) of frame parameters of this atomic model. .. !! processed by numpydoc !! .. py:method:: get_dim_aparam() Get the number (dimension) of atomic parameters of this atomic model. .. !! processed by numpydoc !! .. py:method:: 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. .. !! processed by numpydoc !! .. py:method:: is_aparam_nall() -> bool Check whether the shape of atomic parameters is (nframes, nall, ndim). If False, the shape is (nframes, nloc, ndim). .. !! processed by numpydoc !! .. py:method:: model_output_type() -> list[str] Get the output type for the model. .. !! processed by numpydoc !! .. py:method:: get_model_def_script() -> str Get the model definition script. .. !! processed by numpydoc !! .. py:method:: get_min_nbor_dist() -> Optional[float] Get the minimum neighbor distance. .. !! processed by numpydoc !! .. py:method:: get_nnei() -> int Returns the total number of selected neighboring atoms in the cut-off radius. .. !! processed by numpydoc !! .. py:method:: get_nsel() -> int Returns the total number of selected neighboring atoms in the cut-off radius. .. !! processed by numpydoc !! .. py:method:: has_spin() -> bool :staticmethod: Returns whether it has spin input and output. .. !! processed by numpydoc !! .. py:method:: model_output_def() Get the output def for the model. .. !! processed by numpydoc !! .. py:method:: __getattr__(name) Get attribute from the wrapped model. .. !! processed by numpydoc !! .. py:method:: serialize() -> dict .. py:method:: deserialize(data) -> SpinModel :classmethod: .. py:method:: call(coord, atype, spin, box: Optional[numpy.ndarray] = None, fparam: Optional[numpy.ndarray] = None, aparam: Optional[numpy.ndarray] = None, do_atomic_virial: bool = False) -> dict[str, numpy.ndarray] 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: :obj:`ret_dict` The result dict of type dict[str,np.ndarray]. The keys are defined by the `ModelOutputDef`. .. !! processed by numpydoc !! .. py:method:: call_lower(extended_coord: numpy.ndarray, extended_atype: numpy.ndarray, extended_spin: numpy.ndarray, nlist: numpy.ndarray, mapping: Optional[numpy.ndarray] = None, fparam: Optional[numpy.ndarray] = None, aparam: Optional[numpy.ndarray] = None, do_atomic_virial: bool = False) 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: :obj:`result_dict` the result dict, defined by the `FittingOutputDef`. .. !! processed by numpydoc !! .. py:attribute:: forward_lower