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.dpa4_model
• deepmd.pt_expt.model.edge_transform_output
• deepmd.pt_expt.model.ener_model
• deepmd.pt_expt.model.frozen
• deepmd.pt_expt.model.get_model
• deepmd.pt_expt.model.graph_lower
• deepmd.pt_expt.model.make_model
• deepmd.pt_expt.model.model
• deepmd.pt_expt.model.native_spin_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
• deepmd.pt_expt.model.unimol_pretrain_model

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.
DPA4EnergyModel	Energy model for the DPA4/SeZM descriptor family.
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.
NativeSpinEnergyModel	pt_expt native-spin energy model.
PolarModel	Common methods for DP models.
PropertyModel	Common methods for DP models.
SpinEnergyModel	A spin model for energy.
UniMolPretrainModel	Uni-Mol v1 self-supervised pretraining on the PyTorch-Exportable backend.

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_lower_graph_exportable

forward_lower_graph_exportable_with_comm

forward_lower_canonical_graph

forward_lower_canonical_graph_exportable

compute_fitting_input_stats(sampled_func: collections.abc.Callable[[], list[dict]]) → None

Recompute the input statistics of every learned sub-model’s fitting net.

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.DPA4EnergyModel(*args: Any, **kwargs: Any)

Bases: deepmd.pt_expt.model.ener_model.EnergyModel

Energy model for the DPA4/SeZM descriptor family.

Behaviorally identical to EnergyModel; additionally owns the pt-checkpoint interop: pt’s SeZMModel serialises with a model-level wrapper ({type: "SeZM", atomic_model: <sezm_atomic dict>, bridging_method, bridging_r_*, lora}) whose atomic dict carries pt-only extras – deserialize() recognises those layouts, normalises them to the standard flat dict, and delegates to the generic path.

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

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

_to_public_keys(model_ret: dict[str, torch.Tensor], do_atomic_virial: bool) → dict[str, torch.Tensor]

Rename one call_common result to the public energy-model keys.

The renaming is a property of the output definition, not of the node axis, so it serves the rectangular and the ragged entry alike.

Parameters:

model_retdict[str, torch.Tensor]

A call_common result, in internal <var>_<derivative> keys.

do_atomic_virialbool

Whether the per-atom virial was requested and should be carried.

Returns:

dict[str, torch.Tensor]

The same tensors under the public names.

_translate_eager_call(model_ret: dict[str, torch.Tensor], atype: torch.Tensor, do_atomic_virial: bool = False) → dict[str, torch.Tensor]

Translate internal energy outputs at the public model boundary.

Parameters:

model_retdict[str, torch.Tensor]

Result returned by a call_common entry.

atypetorch.Tensor

Atom types on the same node axis as the atomic outputs.

do_atomic_virialbool, default: False

Whether the per-atom virial was requested.

Returns:

dict[str, torch.Tensor]

Public model outputs.

Notes

Native-spin models override this translation to add force_mag and mask_mag. Keeping the dispatch here lets rectangular and ragged forwards share one implementation without duplicating a spin-specific forward_ragged.

forward_lower_canonical_graph(atype: torch.Tensor, n_node: torch.Tensor, n_local: torch.Tensor, source: torch.Tensor, edge_vec: torch.Tensor, destination_row_ptr: torch.Tensor, source_row_ptr: torch.Tensor, source_order: torch.Tensor, *, do_atomic_virial: bool, spin: torch.Tensor | None = None) → dict[str, torch.Tensor]

Evaluate an eligible compressed canonical deployment graph.

Parameters:

atype

Flat local-plus-halo atom types with shape (N,), int64.

n_node

Per-frame total node counts with shape (nf,), int64.

n_local

Per-frame owned node counts with shape (nf,), int64.

source

Source-node indices with shape (S,), uint32.

edge_vec

Destination-major edge vectors with shape (S, 3), float32.

destination_row_ptr

Destination CSR offsets with shape (N + 1,), int64.

source_row_ptr

Source CSR offsets with shape (N + 1,), int64.

source_order

Source-grouped edge positions with shape (S,) and the same dtype as source.

do_atomic_virial

Whether to return the per-node virial.

spin

Per-node magnetic moments with shape (N, 3) for a native-spin model, or None. Ghost rows carry their owner’s moment.

Returns:

dict[str, torch.Tensor]

Public energy-model outputs on the flat node axis.

forward_lower_canonical_graph_exportable(atype: torch.Tensor, n_node: torch.Tensor, n_local: torch.Tensor, source: torch.Tensor, edge_vec: torch.Tensor, destination_row_ptr: torch.Tensor, source_row_ptr: torch.Tensor, source_order: torch.Tensor, *, do_atomic_virial: bool, **make_fx_kwargs: Any) → torch.nn.Module

Trace the compact canonical deployment forward with make_fx.

Parameters:

atype, n_node, n_local, source, edge_vec

Compact graph node and edge tensors.

destination_row_ptr, source_row_ptr, source_order

Compact dual-CSR topology.

do_atomic_virial

Whether the traced output includes per-node virial.

**make_fx_kwargs

Additional arguments passed to make_fx().

Returns:

torch.nn.Module

Traced eight-input compact deployment module.

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:

coord

Atomic coordinates.

atype

Atomic type indices.

box

Simulation-cell vectors, or None for a non-periodic system.

fparam

Optional frame parameters.

aparam

Optional atomic parameters.

do_atomic_virial

Whether to return per-atom virials.

charge_spin

Optional frame-level charge and spin conditioning.

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_ragged(coord: torch.Tensor, atype: torch.Tensor, n_node: 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, spin: torch.Tensor | None = None) → dict[str, torch.Tensor]

Evaluate the energy model over a batch whose node axis is flat.

The counterpart of forward() for frames held concatenated rather than padded to a common atom count. Arguments and results share their meaning with call_common_ragged(), whose per-atom outputs keep the flat axis; only the keys are the public ones.

Parameters:

coordtorch.Tensor

Local coordinates with shape (N, 3), frame-major over n_node.

atypetorch.Tensor

Local atom types with shape (N,).

n_nodetorch.Tensor

Atoms per frame with shape (nf,).

boxtorch.Tensor or None, optional

Simulation cell with shape (nf, 3, 3).

fparamtorch.Tensor or None, optional

Frame parameters with shape (nf, ndf).

aparamtorch.Tensor or None, optional

Atomic parameters with shape (N, nda).

do_atomic_virialbool, default: False

Whether to return per-atom virials.

charge_spintorch.Tensor or None, optional

Frame-level charge and spin conditioning with shape (nf, 2).

spintorch.Tensor or None, optional

Per-atom native spin with shape (N, 3).

Returns:

dict[str, torch.Tensor]

Public energy-model keys; per-atom entries have leading dimension N and per-frame entries nf.

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-coordinate sample used for tracing.

extended_atype

Extended atom-type sample used for tracing.

nlist

Neighbor-list sample used for tracing.

mapping

Extended-to-local mapping sample used for tracing.

fparam

Optional frame-parameter sample.

aparam

Optional atomic-parameter sample.

do_atomic_virial

Whether the traced module returns per-atom virials.

charge_spin

Optional charge/spin conditioning sample.

**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.

forward_lower_graph_exportable(atype: torch.Tensor, n_node: torch.Tensor, n_local: torch.Tensor, edge_index: torch.Tensor, edge_vec: torch.Tensor, edge_mask: torch.Tensor, destination_order: torch.Tensor, destination_row_ptr: torch.Tensor, source_order: torch.Tensor, source_row_ptr: torch.Tensor, fparam: torch.Tensor | None = None, aparam: torch.Tensor | None = None, do_atomic_virial: bool = False, charge_spin: torch.Tensor | None = None, destination_sorted: bool = False, **make_fx_kwargs: Any) → torch.nn.Module

Trace forward_common_lower_graph into an exportable module with public output keys.

Delegates to forward_common_lower_graph_exportable for tracing, then translates the internal keys to the forward_lower convention.

Parameters:

atype

(N,) flat local-plus-halo atom types, N == sum(n_node).

n_node

(nf,) per-frame total node counts.

n_local

(nf,) per-frame owned node counts.

edge_index

(2, E) [src, dst] edge endpoints (flat local indices).

edge_vec

(E, 3) neighbor-minus-center edge vectors (sample for tracing).

edge_mask

(E,) valid-edge mask (sample for tracing).

destination_order

(E,) destination-grouped edge permutation.

source_order

(E,) source-grouped edge permutation.

destination_row_ptr, source_row_ptr

(N + 1,) destination/source CSR offsets.

destination_sorted

Static export-time assertion that the payload is destination-major and destination_order is identity.

fparam, aparam, do_atomic_virial, charge_spin

As in forward_lower.

**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 (atype, n_node, n_local, edge_index, edge_vec, edge_mask, destination_order, destination_row_ptr, source_order, source_row_ptr, fparam, aparam, charge_spin) and returns a dict with the public keys: atom_energy, energy, force, virial, atom_virial (the last only when do_atomic_virial). Unlike the dense forward_lower_exportable() (which emits extended_force / extended_virial), the graph path emits force and atom_virial directly on the local-plus-halo node axis.

forward_lower_graph_exportable_with_comm(atype: torch.Tensor, n_node: torch.Tensor, n_local: torch.Tensor, edge_index: torch.Tensor, edge_vec: torch.Tensor, edge_mask: torch.Tensor, destination_order: torch.Tensor, destination_row_ptr: torch.Tensor, source_order: torch.Tensor, source_row_ptr: torch.Tensor, fparam: torch.Tensor | None, aparam: torch.Tensor | None, charge_spin: torch.Tensor | None, send_list: torch.Tensor, send_proc: torch.Tensor, recv_proc: torch.Tensor, send_num: torch.Tensor, recv_num: torch.Tensor, communicator: torch.Tensor, nlocal: torch.Tensor, nghost: torch.Tensor, do_atomic_virial: bool = False, **make_fx_kwargs: Any) → torch.nn.Module

Trace forward_common_lower_graph with comm_dict tensors as additional positional inputs – the with-comm counterpart of forward_lower_graph_exportable() for message-passing graph descriptors (dpa2’s repformer block drives cross-rank ghost refresh via deepmd_export::border_op, see _exchange_ghosts_graph()).

Mirrors the dense forward_common_lower_exportable_with_comm (pt_expt/model/make_model.py): packs the 8 trailing positional comm tensors into a comm_dict inside the traced function. Also derives n_local (the per-frame OWNED node count, reshaped to (1,); single-frame – LAMMPS always drives inference with nf=1) from the scalar nlocal tensor, so the differentiated reduction excludes ghost (not-owned) nodes (see forward_common_lower_graph()’s n_local parameter). Unlike the plain-graph export path (which traces forward_common_lower_graph_exportable and then wraps a SECOND make_fx trace around the key-translation closure), this method traces ONCE: the comm-dict packing, n_local derivation, the forward_common_lower_graph call and the key translation all live in a single traced fn – following the dense with-comm precedent, which is also a single trace.

Parameters:

atype, n_node, edge_index, edge_vec, edge_mask, fparam, aparam, charge_spin, do_atomic_virial

As in forward_lower_graph_exportable().

destination_order

Destination-major edge permutation used by fused graph operators.

destination_row_ptr

Destination CSR row pointers.

source_order

Source-major edge permutation used by force assembly.

source_row_ptr

Source CSR row pointers.

send_list, send_proc, recv_proc, send_num, recv_num, communicator, nlocal, nghost

The 8 comm tensors (see _make_comm_sample_inputs in serialization.py), packed into comm_dict inside the traced function.

Runtime device contract: ALL 8 stay on CPU, symmetric with the dense with-comm artifact – they are consumed only by the opaque border_op whose HOST code dereferences their data_ptr (send_list carries raw host pointers) and reads nlocal/nghost via cheap host .item() calls. Deriving the in-graph owned count from a device-placed nlocal instead (the previous design) made every per-layer border_op forward AND custom backward pull the scalars back with synchronizing D2H reads (4 * nlayers per MD step). The C++ run_model_graph_with_comm implements this placement.

n_local

(1,) int64 ON THE MODEL DEVICE: the per-frame OWNED node count consumed IN-GRAPH by the owned-node energy mask (it becomes a device kernel operand after move_to_device_pass, like n_node; a CPU tensor fed there is read as a device pointer – CUDA illegal memory access). Carries the same value as the nlocal comm tensor; the two inputs exist precisely to separate the device-compute role from the host-MPI-control role.

**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 (atype, n_node, n_local, edge_index, edge_vec, edge_mask, destination_order, destination_row_ptr, source_order, source_row_ptr, fparam, aparam, charge_spin, send_list, send_proc, recv_proc, send_num, recv_num, communicator, nlocal, nghost) and returns a dict with the SAME public keys as forward_lower_graph_exportable() (atom_energy, energy, force, virial, atom_virial when do_atomic_virial).

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. Deserialization is pure registry dispatch: descriptor-family-specific wire formats (e.g. the pt SeZM checkpoint layouts) are owned by the family’s registered model class (see deepmd.pt_expt.model.dpa4_model.DPA4EnergyModel), never by this base.

See also

deepmd.dpmodel.model.base_model.BaseBaseModel

Backend-independent BaseModel class.

class deepmd.pt_expt.model.NativeSpinEnergyModel(*args: Any, spin: deepmd.utils.spin.Spin, **kwargs: Any)

Bases: make_native_spin_model(EnergyModel)

pt_expt native-spin energy model.

make_native_spin_model applied to THIS backend’s EnergyModel (construction, output defs, serialization all come from the factory; the deserialize closure rebuilds through the pt_expt model class, so a registry round trip yields a real torch.nn.Module), plus two torch-specific overrides:

• forward(): the pt_expt call_common produces REAL autograd energy_derv_r/energy_derv_r_mag/energy_derv_c_redu tensors, unlike the dpmodel factory’s energy-only call (which is restricted to force/force_mag/virial as None placeholders because dpmodel has no autograd).

• forward_lower_graph_exportable(): the graph-spin .pt2 positional ABI (spin at index 10) over the inherited forward_common_lower_graph_exportable.

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]

Return native-spin model predictions with public output keys.

Parameters:

coord

The coordinates of the atoms. shape: nf x (nloc x 3)

atype

The type of atoms. shape: nf x nloc

spin

The per-local-atom spin. 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.

charge_spin

Frame-level charge/spin conditioning, shape nf x 2. Forwarded to the backbone like every other conditioning input; native spin and charge-spin FiLM combine freely (DPA4 declares both capabilities).

Returns:

ret_dict

The result dict with keys atom_energy, energy, force, force_mag, virial, mask_mag, and (when do_atomic_virial) atom_virial. force and force_mag are real autograd tensors (-dE/dcoord and -dE/dspin), NOT placeholders.

forward_ragged(coord: torch.Tensor, atype: torch.Tensor, n_node: 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]

Return native-spin predictions on a flat, frame-segmented node axis.

forward_lower_graph_exportable(atype: torch.Tensor, n_node: torch.Tensor, n_local: torch.Tensor, edge_index: torch.Tensor, edge_vec: torch.Tensor, edge_mask: torch.Tensor, destination_order: torch.Tensor, destination_row_ptr: torch.Tensor, source_order: torch.Tensor, source_row_ptr: torch.Tensor, spin: torch.Tensor, fparam: torch.Tensor | None = None, aparam: torch.Tensor | None = None, charge_spin: torch.Tensor | None = None, do_atomic_virial: bool = False, destination_sorted: bool = False, **make_fx_kwargs: Any) → torch.nn.Module

Trace the graph-spin lower into an exportable module.

THIS METHOD OWNS the positional .pt2 ABI for graph-spin models (mirrored verbatim by the C++ / serialization seams): spin sits at index 10, right after the shared NeighborGraph CSR block and before the conditional fparam/aparam tail; the conditional charge_spin tail follows at index 13 (combined native-spin + charge-spin FiLM models, review 3638047227; None otherwise). forward_lower_graph_exportable_with_comm() extends this SAME prefix with the 8 comm tensors for multi-rank.

Two-layer make_fx trace, mirroring forward_lower_graph_exportable(): the inner layer (the inherited forward_common_lower_graph_exportable()) traces forward_common_lower_graph with spin as a SECOND autograd leaf next to edge_vec (charge_spin is a frame-level conditioning input, not a leaf); this outer layer re-traces with the PUBLIC positional ABI above and translates the internal fitting keys to the public output keys via the shared _translate_eager_call() (the same single owner as the eager forward(), so mask_mag is emitted here too).

Parameters:

atype

(N,) flat local-plus-halo atom types, N == sum(n_node).

n_node

(nf,) per-frame total node counts.

n_local

(nf,) per-frame owned node counts.

edge_index

(2, E) [src, dst] edge endpoints (flat local indices).

edge_vec

(E, 3) neighbor-minus-center edge vectors (sample for tracing).

edge_mask

(E,) valid-edge mask (sample for tracing).

destination_order

(E,) destination-grouped edge permutation.

destination_row_ptr

(N + 1,) destination CSR offsets.

source_order

(E,) source-grouped edge permutation.

source_row_ptr

(N + 1,) source CSR offsets.

spin

(N, 3) per-node native spin (sample for tracing). ALWAYS present in this ABI, unlike the energy model’s optional charge_spin.

fparam

Frame parameter, (nf, ndf), or None when dim_fparam == 0.

aparam

Atomic parameter, (N, nda), or None when dim_aparam == 0.

charge_spin

Frame-level charge/spin FiLM conditioning, (nf, dim_chg_spin), or None when the descriptor has no add_chg_spin_ebd (conditional tail, slot 13).

do_atomic_virial

Whether to also return atom_virial.

destination_sorted

Static export-time assertion that the payload is destination-major and destination_order is identity.

**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 (atype, n_node, n_local, edge_index, edge_vec, edge_mask, destination_order, destination_row_ptr, source_order, source_row_ptr, spin, fparam, aparam, charge_spin) and returns a dict with the public keys atom_energy, energy, force, force_mag, virial, mask_mag, and (when do_atomic_virial) atom_virial.

forward_lower_canonical_graph_exportable(atype: torch.Tensor, n_node: torch.Tensor, n_local: torch.Tensor, source: torch.Tensor, edge_vec: torch.Tensor, destination_row_ptr: torch.Tensor, source_row_ptr: torch.Tensor, source_order: torch.Tensor, spin: torch.Tensor, *, do_atomic_virial: bool, **make_fx_kwargs: Any) → torch.nn.Module

Trace the compact canonical spin lower into an exportable module.

THIS METHOD OWNS the positional .pt2 ABI for compact canonical spin models (mirrored verbatim by the C++ / serialization seams): spin sits at index 8, directly after the dual-CSR block, which is the same placement rule the graph ABI uses at index 10. The compact ABI has no conditional tail, so index 8 is the last slot.

The magnetic force is a value here rather than an autograd result: the fused operator emits the on-site half from the node kernel and the neighbour half per edge, and the deployment path reduces the second onto source nodes. mask_mag is derived from the types by the same single owner the eager translation uses.

Parameters:

atype, n_node, n_local, source, edge_vec

Compact graph node and edge tensors.

destination_row_ptr, source_row_ptr, source_order

Compact dual-CSR topology.

spin

(N, 3) per-node moment. Ghost rows carry their owner’s moment.

do_atomic_virial

Whether the traced output includes per-node virial.

**make_fx_kwargs

Additional arguments passed to make_fx().

Returns:

torch.nn.Module

Traced nine-input compact deployment module whose output dict adds force_mag and mask_mag to the spin-free key set.

forward_lower_graph_exportable_with_comm(atype: torch.Tensor, n_node: torch.Tensor, n_local: torch.Tensor, edge_index: torch.Tensor, edge_vec: torch.Tensor, edge_mask: torch.Tensor, destination_order: torch.Tensor, destination_row_ptr: torch.Tensor, source_order: torch.Tensor, source_row_ptr: torch.Tensor, spin: torch.Tensor, fparam: torch.Tensor | None, aparam: torch.Tensor | None, charge_spin: torch.Tensor | None, send_list: torch.Tensor, send_proc: torch.Tensor, recv_proc: torch.Tensor, send_num: torch.Tensor, recv_num: torch.Tensor, communicator: torch.Tensor, nlocal: torch.Tensor, nghost: torch.Tensor, do_atomic_virial: bool = False, **make_fx_kwargs: Any) → torch.nn.Module

Trace the multi-rank graph-spin lower into an exportable module.

The with-comm counterpart of forward_lower_graph_exportable(): same positional prefix, spin still at index 10, then the 8 comm tensors appended after the conditional fparam/aparam/charge_spin tail (indices 14-21) – exactly how forward_lower_graph_exportable_with_comm() appends them for the energy model.

spin is the EXTENDED per-node spin (N, 3): ghost rows carry their owner’s spin, delivered by the LAMMPS sp forward-comm before the call, so the descriptor’s spin embedding sees the same value on every rank that holds the node. The per-block ghost FEATURE refresh rides deepmd_export::border_op exactly as in the energy model; spin needs no border exchange of its own because it is an input, not a derived feature.

Single make_fx trace (the energy with-comm precedent), unlike the two-layer trace of the non-comm spin path: the comm-dict packing, the forward_common_lower_graph call with spin as a second autograd leaf, and the public-key translation all live in one traced fn.

Parameters:

atype, n_node, n_local, edge_index, edge_vec, edge_mask, destination_order, destination_row_ptr, source_order, source_row_ptr, spin, fparam, aparam, charge_spin

As in forward_lower_graph_exportable().

send_list, send_proc, recv_proc, send_num, recv_num, communicator, nlocal, nghost

The 8 comm tensors, packed into comm_dict inside the traced function. Same runtime device contract as the energy model’s: ALL 8 stay on CPU (host control metadata for border_op), while the device-side owned count is the separate n_local input at slot 2.

do_atomic_virial

Whether to also return atom_virial.

**make_fx_kwargs

Extra keyword arguments forwarded to make_fx.

Returns:

torch.nn.Module

A traced module accepting the 22-input ABI above and returning the same public keys as forward_lower_graph_exportable() (atom_energy, energy, force, force_mag, virial, mask_mag, plus atom_virial when requested).

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.

get_task_dim() → int

Get the output dimension of the property.

get_intensive() → bool

Whether the property is intensive.

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, charge_spin

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, charge_spin) and returns a dict with the same keys as forward_lower.

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

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

Uni-Mol v1 self-supervised pretraining on the PyTorch-Exportable backend.

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]

Evaluate the pretraining heads on a frame.

Raises:

ValueError

If a periodic cell is supplied; Uni-Mol is molecular.

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, comm_dict: dict[str, torch.Tensor] | None = None, charge_spin: torch.Tensor | None = None) → dict[str, torch.Tensor]

Evaluate the pretraining heads on an extended frame.

translated_output_def() → dict[str, deepmd.dpmodel.output_def.OutputVariableDef]

The head outputs, passed through under their own names.