deepmd.pt_expt.model

Submodules

• deepmd.pt_expt.model.dipole_model
• deepmd.pt_expt.model.dos_model
• deepmd.pt_expt.model.dp_linear_model
• deepmd.pt_expt.model.dp_zbl_model
• deepmd.pt_expt.model.ener_model
• deepmd.pt_expt.model.frozen
• deepmd.pt_expt.model.get_model
• deepmd.pt_expt.model.make_model
• deepmd.pt_expt.model.model
• deepmd.pt_expt.model.polar_model
• deepmd.pt_expt.model.property_model
• deepmd.pt_expt.model.spin_ener_model
• deepmd.pt_expt.model.spin_model
• deepmd.pt_expt.model.transform_output

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

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)

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]

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]

translated_output_def() → dict[str, Any]

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

class deepmd.pt_expt.model.DOSModel(*args: Any, **kwargs: Any)

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]

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]

translated_output_def() → dict[str, Any]

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

class deepmd.pt_expt.model.LinearEnergyModel(*args: Any, **kwargs: Any)

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]

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]

translated_output_def() → dict[str, Any]

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

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

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)

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]

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]

translated_output_def() → dict[str, Any]

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

class deepmd.pt_expt.model.EnergyModel(*args: Any, **kwargs: Any)

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

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, neighbor_list: deepmd.dpmodel.utils.neighbor_list.NeighborList | None = None) → dict[str, torch.Tensor]

Evaluate the energy model.

Most arguments share the meaning of call_common().

Parameters:

neighbor_list

The neighbor-list construction strategy forwarded to call_common(). None uses the default all-pairs builder (DefaultNeighborList), reproducing the historical behavior; an alternative strategy (e.g. the vesin O(N) cell list) may be injected to accelerate neighbor-list construction without changing the model outputs.

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]

translated_output_def() → dict[str, Any]

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:

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)

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

Get a model from a config dictionary.

Parameters:

datadict

The data to construct the model.

class deepmd.pt_expt.model.BaseModel

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.

See also

deepmd.dpmodel.model.base_model.BaseBaseModel

Backend-independent BaseModel class.

_SEZM_MODEL_TYPES

_SEZM_ATOMIC_TYPES

classmethod deserialize(data: dict[str, Any]) → BaseModel

Deserialize the model.

Parameters:

datadict

The serialized data

Returns:

BaseModel

The deserialized model

static _unwrap_pt_sezm_model(data: dict[str, Any]) → dict[str, Any]

Unwrap pt’s SeZMModel serialization to the inner atomic dict.

static _normalize_pt_sezm_atomic(data: dict[str, Any]) → dict[str, Any]

Convert a pt sezm_atomic dict to a standard atomic dict.

Strips the pt-only dens head state (dens_fitting / active_mode / the dens_force_rmsd @variable) and rewrites the type/@version so the generic dpmodel atomic-model deserialize accepts it. A non-energy active mode or a populated dens head is rejected because pt_expt only implements the energy path.

class deepmd.pt_expt.model.PolarModel(*args: Any, **kwargs: Any)

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

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]

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]

translated_output_def() → dict[str, Any]

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

class deepmd.pt_expt.model.PropertyModel(*args: Any, **kwargs: Any)

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

Common methods for DP models.

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

get_var_name() → str

Get the name of the property.

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.

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]

translated_output_def() → dict[str, Any]

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

class deepmd.pt_expt.model.SpinEnergyModel(backbone_model: deepmd.dpmodel.atomic_model.dp_atomic_model.DPAtomicModel, spin: deepmd.utils.spin.Spin)

Bases: deepmd.pt_expt.model.spin_model.SpinModel

A spin model for energy.

model_type = 'ener'

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.

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]

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]

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:

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.