deepmd.pt_expt.model

deepmd.pt_expt.model#

Submodules#

Classes#

`DipoleModel`	Common methods for DP models.
`DOSModel`	Common methods for DP models.
`LinearEnergyModel`	Common methods for DP models.
`DPZBLModel`	Common methods for DP models.
`EnergyModel`	Common methods for DP models.
`FrozenModel`	Load model from a frozen model file, which cannot be trained.
`BaseModel`	Base class for pt_expt models.
`PolarModel`	Common methods for DP models.
`PropertyModel`	Common methods for DP models.
`SpinEnergyModel`	A spin model for energy.

Functions#

`make_hessian_model`(→ type)	Make a model that can compute Hessian.
`get_model`(→ deepmd.pt_expt.model.model.BaseModel)	Get a model from a config dictionary.

Package Contents#

deepmd.pt_expt.model.make_hessian_model(T_Model: type) → type[source]#

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:

The model that computes hessian.

class deepmd.pt_expt.model.DipoleModel(*args: Any, **kwargs: Any)[source]#

Bases: deepmd.dpmodel.model.dp_model.DPModelCommon, DPDipoleModel_

Common methods for DP models.

This class provides common functionality for DeepPot models, including neighbor selection updates and fitting network access.

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][source]#

Default forward delegates to call().

Subclasses (e.g. EnergyModel) override this with output translation.

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][source]#

translated_output_def() → dict[str, Any][source]#

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[source]#

class deepmd.pt_expt.model.DOSModel(*args: Any, **kwargs: Any)[source]#

Bases: deepmd.dpmodel.model.dp_model.DPModelCommon, DPDOSModel_

Common methods for DP models.

This class provides common functionality for DeepPot models, including neighbor selection updates and fitting network access.

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][source]#

Default forward delegates to call().

Subclasses (e.g. EnergyModel) override this with output translation.

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][source]#

translated_output_def() → dict[str, Any][source]#

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[source]#

class deepmd.pt_expt.model.LinearEnergyModel(*args: Any, **kwargs: Any)[source]#

Bases: deepmd.dpmodel.model.dp_model.DPModelCommon, DPLinearModel_

Common methods for DP models.

This class provides common functionality for DeepPot models, including neighbor selection updates and fitting network access.

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][source]#

Default forward delegates to call().

Subclasses (e.g. EnergyModel) override this with output translation.

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][source]#

translated_output_def() → dict[str, Any][source]#

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[source]#

classmethod update_sel(train_data: deepmd.utils.data_system.DeepmdDataSystem, type_map: list[str] | None, local_jdata: dict) → tuple[dict, float | None][source]#

Update the selection and perform neighbor statistics.

Parameters:

train_dataDeepmdDataSystem: data used to do neighbor statistics
type_maplist[str], optional: The name of each type of atoms
local_jdatadict: The local data refer to the current class

Returns:

dict: The updated local data
float: The minimum distance between two atoms

class deepmd.pt_expt.model.DPZBLModel(*args: Any, **kwargs: Any)[source]#

Bases: deepmd.dpmodel.model.dp_model.DPModelCommon, DPZBLModel_

Common methods for DP models.

This class provides common functionality for DeepPot models, including neighbor selection updates and fitting network access.

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][source]#

Default forward delegates to call().

Subclasses (e.g. EnergyModel) override this with output translation.

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][source]#

translated_output_def() → dict[str, Any][source]#

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[source]#

class deepmd.pt_expt.model.EnergyModel(*args: Any, **kwargs: Any)[source]#

Bases: deepmd.dpmodel.model.dp_model.DPModelCommon, DPEnergyModel_

Common methods for DP models.

This class provides common functionality for DeepPot models, including neighbor selection updates and fitting network access.

_hessian_enabled = False#

enable_hessian() → None[source]#

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][source]#

Default forward delegates to call().

Subclasses (e.g. EnergyModel) override this with output translation.

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][source]#

translated_output_def() → dict[str, Any][source]#

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[source]#

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:

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.

class deepmd.pt_expt.model.FrozenModel(model_file: str, **kwargs: Any)[source]#

Bases: 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_filestr: The path to the frozen model file.

model#

deepmd.pt_expt.model.get_model(data: dict) → deepmd.pt_expt.model.model.BaseModel[source]#

Get a model from a config dictionary.

Parameters:

datadict: The data to construct the model.

class deepmd.pt_expt.model.BaseModel[source]#

Bases: 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.