deepmd.pt_expt.model ==================== .. py:module:: deepmd.pt_expt.model Submodules ---------- .. toctree:: :maxdepth: 1 /autoapi/deepmd/pt_expt/model/dipole_model/index /autoapi/deepmd/pt_expt/model/dos_model/index /autoapi/deepmd/pt_expt/model/dp_linear_model/index /autoapi/deepmd/pt_expt/model/dp_zbl_model/index /autoapi/deepmd/pt_expt/model/ener_model/index /autoapi/deepmd/pt_expt/model/frozen/index /autoapi/deepmd/pt_expt/model/get_model/index /autoapi/deepmd/pt_expt/model/make_model/index /autoapi/deepmd/pt_expt/model/model/index /autoapi/deepmd/pt_expt/model/polar_model/index /autoapi/deepmd/pt_expt/model/property_model/index /autoapi/deepmd/pt_expt/model/spin_ener_model/index /autoapi/deepmd/pt_expt/model/spin_model/index /autoapi/deepmd/pt_expt/model/transform_output/index Classes ------- .. autoapisummary:: deepmd.pt_expt.model.DipoleModel deepmd.pt_expt.model.DOSModel deepmd.pt_expt.model.LinearEnergyModel deepmd.pt_expt.model.DPZBLModel deepmd.pt_expt.model.EnergyModel deepmd.pt_expt.model.FrozenModel deepmd.pt_expt.model.BaseModel deepmd.pt_expt.model.PolarModel deepmd.pt_expt.model.PropertyModel deepmd.pt_expt.model.SpinEnergyModel Functions --------- .. autoapisummary:: deepmd.pt_expt.model.make_hessian_model deepmd.pt_expt.model.get_model Package Contents ---------------- .. py:function:: make_hessian_model(T_Model: type) -> type Make a model that can compute Hessian. With the JAX-mirrored approach, hessian is computed in ``forward_common_atomic`` (in make_model.py) on extended coordinates. This wrapper only needs to override ``atomic_output_def()`` to set ``r_hessian=True``, and ``communicate_extended_output`` in dpmodel naturally maps it from nall to nloc. :Parameters: **T_Model** The model. Should provide the ``atomic_output_def`` method. :Returns: :obj:`The` :obj:`model` :obj:`that` :obj:`computes` hessian. .. .. !! processed by numpydoc !! .. py:class:: DipoleModel(*args: Any, **kwargs: Any) Bases: :py:obj:`deepmd.dpmodel.model.dp_model.DPModelCommon`, :py:obj:`DPDipoleModel_` Common methods for DP models. This class provides common functionality for DeepPot models, including neighbor selection updates and fitting network access. .. !! processed by numpydoc !! .. py:method:: 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, charge_spin: torch.Tensor | None = None) -> dict[str, torch.Tensor] Default forward delegates to call(). Subclasses (e.g. EnergyModel) override this with output translation. .. !! processed by numpydoc !! .. py:method:: 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, charge_spin: torch.Tensor | None = None) -> dict[str, torch.Tensor] .. py:method:: translated_output_def() -> dict[str, Any] .. py:method:: forward_lower_exportable(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, charge_spin: torch.Tensor | None = None, **make_fx_kwargs: Any) -> torch.nn.Module .. py:class:: DOSModel(*args: Any, **kwargs: Any) Bases: :py:obj:`deepmd.dpmodel.model.dp_model.DPModelCommon`, :py:obj:`DPDOSModel_` Common methods for DP models. This class provides common functionality for DeepPot models, including neighbor selection updates and fitting network access. .. !! processed by numpydoc !! .. py:method:: 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, charge_spin: torch.Tensor | None = None) -> dict[str, torch.Tensor] Default forward delegates to call(). Subclasses (e.g. EnergyModel) override this with output translation. .. !! processed by numpydoc !! .. py:method:: 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, charge_spin: torch.Tensor | None = None) -> dict[str, torch.Tensor] .. py:method:: translated_output_def() -> dict[str, Any] .. py:method:: forward_lower_exportable(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, charge_spin: torch.Tensor | None = None, **make_fx_kwargs: Any) -> torch.nn.Module .. py:class:: LinearEnergyModel(*args: Any, **kwargs: Any) Bases: :py:obj:`deepmd.dpmodel.model.dp_model.DPModelCommon`, :py:obj:`DPLinearModel_` Common methods for DP models. This class provides common functionality for DeepPot models, including neighbor selection updates and fitting network access. .. !! processed by numpydoc !! .. py:method:: 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, charge_spin: torch.Tensor | None = None) -> dict[str, torch.Tensor] Default forward delegates to call(). Subclasses (e.g. EnergyModel) override this with output translation. .. !! processed by numpydoc !! .. py:method:: 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, charge_spin: torch.Tensor | None = None) -> dict[str, torch.Tensor] .. py:method:: translated_output_def() -> dict[str, Any] .. py:method:: forward_lower_exportable(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, charge_spin: torch.Tensor | None = None, **make_fx_kwargs: Any) -> torch.nn.Module .. py:method:: update_sel(train_data: deepmd.utils.data_system.DeepmdDataSystem, type_map: list[str] | None, local_jdata: dict) -> tuple[dict, float | None] :classmethod: Update the selection and perform neighbor statistics. :Parameters: **train_data** : :obj:`DeepmdDataSystem` data used to do neighbor statistics **type_map** : :class:`python:list`\[:class:`python:str`], :obj:`optional` The name of each type of atoms **local_jdata** : :class:`python:dict` The local data refer to the current class :Returns: :class:`python:dict` The updated local data :class:`python:float` The minimum distance between two atoms .. !! processed by numpydoc !! .. py:class:: DPZBLModel(*args: Any, **kwargs: Any) Bases: :py:obj:`deepmd.dpmodel.model.dp_model.DPModelCommon`, :py:obj:`DPZBLModel_` Common methods for DP models. This class provides common functionality for DeepPot models, including neighbor selection updates and fitting network access. .. !! processed by numpydoc !! .. py:method:: 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, charge_spin: torch.Tensor | None = None) -> dict[str, torch.Tensor] Default forward delegates to call(). Subclasses (e.g. EnergyModel) override this with output translation. .. !! processed by numpydoc !! .. py:method:: 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, charge_spin: torch.Tensor | None = None) -> dict[str, torch.Tensor] .. py:method:: translated_output_def() -> dict[str, Any] .. py:method:: forward_lower_exportable(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, charge_spin: torch.Tensor | None = None, **make_fx_kwargs: Any) -> torch.nn.Module .. py:class:: EnergyModel(*args: Any, **kwargs: Any) Bases: :py:obj:`deepmd.dpmodel.model.dp_model.DPModelCommon`, :py:obj:`DPEnergyModel_` Common methods for DP models. This class provides common functionality for DeepPot models, including neighbor selection updates and fitting network access. .. !! processed by numpydoc !! .. py:attribute:: _hessian_enabled :value: False .. py:method:: enable_hessian() -> None .. py:method:: 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, charge_spin: torch.Tensor | None = None) -> dict[str, torch.Tensor] Default forward delegates to call(). Subclasses (e.g. EnergyModel) override this with output translation. .. !! processed by numpydoc !! .. py:method:: 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, charge_spin: torch.Tensor | None = None) -> dict[str, torch.Tensor] .. py:method:: translated_output_def() -> dict[str, Any] .. py:method:: forward_lower_exportable(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, charge_spin: torch.Tensor | None = None, **make_fx_kwargs: Any) -> torch.nn.Module Trace ``forward_lower`` into an exportable module. Delegates to ``forward_common_lower_exportable`` for tracing, then translates the internal keys to the ``forward_lower`` convention. :Parameters: **extended_coord, extended_atype, nlist, mapping, fparam, aparam, do_atomic_virial** Sample inputs with representative shapes (used for tracing). **\*\*make_fx_kwargs** Extra keyword arguments forwarded to ``make_fx`` (e.g. ``tracing_mode="symbolic"``). :Returns: :obj:`torch.nn.Module` A traced module whose ``forward`` accepts ``(extended_coord, extended_atype, nlist, mapping, fparam, aparam)`` and returns a dict with the same keys as ``forward_lower``. .. !! processed by numpydoc !! .. py:class:: FrozenModel(model_file: str, **kwargs: Any) Bases: :py:obj:`deepmd.dpmodel.model.frozen.FrozenModel` Load model from a frozen model file, which cannot be trained. The frozen model delegates all operations to the deserialized inner model. ``serialize()`` returns the inner model's data, and ``deserialize()`` dispatches to the appropriate model class via ``BaseModel.deserialize``. :Parameters: **model_file** : :class:`python:str` The path to the frozen model file. .. !! processed by numpydoc !! .. py:attribute:: model .. py:function:: get_model(data: dict) -> deepmd.pt_expt.model.model.BaseModel Get a model from a config dictionary. :Parameters: **data** : :class:`python:dict` The data to construct the model. .. !! processed by numpydoc !! .. py:class:: BaseModel Bases: :py:obj:`make_base_model`\ (\ ) Base class for pt_expt models. Provides the plugin registry so that model classes can be registered with ``@BaseModel.register("ener")`` etc. .. seealso:: :obj:`deepmd.dpmodel.model.base_model.BaseBaseModel` Backend-independent BaseModel class. .. !! processed by numpydoc !! .. py:class:: PolarModel(*args: Any, **kwargs: Any) Bases: :py:obj:`deepmd.dpmodel.model.dp_model.DPModelCommon`, :py:obj:`DPPolarModel_` Common methods for DP models. This class provides common functionality for DeepPot models, including neighbor selection updates and fitting network access. .. !! processed by numpydoc !! .. py:method:: 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, charge_spin: torch.Tensor | None = None) -> dict[str, torch.Tensor] Default forward delegates to call(). Subclasses (e.g. EnergyModel) override this with output translation. .. !! processed by numpydoc !! .. py:method:: 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, charge_spin: torch.Tensor | None = None) -> dict[str, torch.Tensor] .. py:method:: translated_output_def() -> dict[str, Any] .. py:method:: forward_lower_exportable(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, charge_spin: torch.Tensor | None = None, **make_fx_kwargs: Any) -> torch.nn.Module .. py:class:: PropertyModel(*args: Any, **kwargs: Any) Bases: :py:obj:`deepmd.dpmodel.model.dp_model.DPModelCommon`, :py:obj:`DPPropertyModel_` Common methods for DP models. This class provides common functionality for DeepPot models, including neighbor selection updates and fitting network access. .. !! processed by numpydoc !! .. py:method:: get_var_name() -> str Get the name of the property. .. !! processed by numpydoc !! .. py:method:: 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, charge_spin: torch.Tensor | None = None) -> dict[str, torch.Tensor] Default forward delegates to call(). Subclasses (e.g. EnergyModel) override this with output translation. .. !! processed by numpydoc !! .. py:method:: 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, charge_spin: torch.Tensor | None = None) -> dict[str, torch.Tensor] .. py:method:: translated_output_def() -> dict[str, Any] .. py:method:: forward_lower_exportable(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, charge_spin: torch.Tensor | None = None, **make_fx_kwargs: Any) -> torch.nn.Module .. py:class:: SpinEnergyModel(backbone_model: deepmd.dpmodel.atomic_model.dp_atomic_model.DPAtomicModel, spin: deepmd.utils.spin.Spin) Bases: :py:obj:`deepmd.pt_expt.model.spin_model.SpinModel` A spin model for energy. .. !! processed by numpydoc !! .. py:attribute:: model_type :value: 'ener' .. py:method:: translated_output_def() -> dict[str, Any] Get the translated output definition. Maps internal output names to user-facing names, e.g. ``energy`` -> ``atom_energy``, ``energy_redu`` -> ``energy``, ``energy_derv_r`` -> ``force``, ``energy_derv_r_mag`` -> ``force_mag``. .. !! processed by numpydoc !! .. py:method:: forward(coord: torch.Tensor, atype: torch.Tensor, spin: torch.Tensor, box: torch.Tensor | None = None, fparam: torch.Tensor | None = None, aparam: torch.Tensor | None = None, do_atomic_virial: bool = False, charge_spin: torch.Tensor | None = None) -> dict[str, torch.Tensor] .. py:method:: forward_lower(extended_coord: torch.Tensor, extended_atype: torch.Tensor, extended_spin: 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, charge_spin: torch.Tensor | None = None) -> dict[str, torch.Tensor] .. py:method:: forward_lower_exportable(extended_coord: torch.Tensor, extended_atype: torch.Tensor, extended_spin: 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, charge_spin: torch.Tensor | None = None, **make_fx_kwargs: Any) -> torch.nn.Module Trace ``forward_lower`` into an exportable module. Delegates to ``forward_common_lower_exportable`` for tracing, then translates the internal keys to the ``forward_lower`` convention. :Parameters: **extended_coord, extended_atype, extended_spin, nlist, mapping, fparam, aparam, do_atomic_virial** Sample inputs with representative shapes (used for tracing). **\*\*make_fx_kwargs** Extra keyword arguments forwarded to ``make_fx`` (e.g. ``tracing_mode="symbolic"``). :Returns: :obj:`torch.nn.Module` A traced module whose ``forward`` accepts ``(extended_coord, extended_atype, extended_spin, nlist, mapping, fparam, aparam)`` and returns a dict with the same keys as ``forward_lower``. .. !! processed by numpydoc !!