deepmd.pt.infer.deep_eval

Attributes

log
_EMBEDDING_DTYPE_TO_TORCH

Classes

DeepEval

PyTorch backend implementation of DeepEval.

Functions

_is_sezm_model_params(→ bool)

Return whether the params describe a SeZM / DPA4 model.

Module Contents

deepmd.pt.infer.deep_eval.log

deepmd.pt.infer.deep_eval._EMBEDDING_DTYPE_TO_TORCH

deepmd.pt.infer.deep_eval._is_sezm_model_params(model_params: dict[str, Any]) → bool

Return whether the params describe a SeZM / DPA4 model.

class deepmd.pt.infer.deep_eval.DeepEval(model_file: str, output_def: deepmd.dpmodel.output_def.ModelOutputDef, *args: Any, auto_batch_size: bool | int | deepmd.pt.utils.auto_batch_size.AutoBatchSize = True, neighbor_list: ase.neighborlist.NewPrimitiveNeighborList | None = None, head: str | int | None = None, no_jit: bool = False, nlist_backend: str = 'auto', **kwargs: Any)

Bases: deepmd.infer.deep_eval.DeepEvalBackend

PyTorch backend implementation of DeepEval.

Parameters:

model_filePath

The name of the frozen model file.

output_defModelOutputDef

The output definition of the model.

*argslist

Positional arguments.

auto_batch_sizebool or int or AutomaticBatchSize, default: True

If True, automatic batch size will be used. If int, it will be used as the initial batch size.

neighbor_listase.neighborlist.NewPrimitiveNeighborList, optional

The ASE neighbor list class to produce the neighbor list. If None, the neighbor list will be built natively in the model.

**kwargsdict

Keyword arguments.

output_def

model_path

neighbor_list = None

modifier = None

rcut

type_map

_has_spin

_has_hessian

_uses_edge_schema

_setup_nlist_backend(nlist_backend: str) → None

Resolve the neighbor-list construction strategy from a user choice.

"native" uses the dense all-pairs builder; "vesin" / "nv" force the O(N) vesin.torch / nvalchemiops cell list (raising if unavailable or the model/inputs are unsupported); "auto" picks the first available O(N) builder (vesin, then nv) and otherwise falls back to the native builder. Results are unchanged either way – only the neighbor-search cost differs.

get_rcut() → float

Get the cutoff radius of this model.

get_ntypes() → int

Get the number of atom types of this model.

get_type_map() → list[str]

Get the type map (element name of the atom types) of this model.

get_dim_fparam() → int

Get the number (dimension) of frame parameters of this DP.

get_dim_aparam() → int

Get the number (dimension) of atomic parameters of this DP.

has_default_fparam() → bool

Check if the model has default frame parameters.

has_chg_spin_ebd() → bool

Check if the model has charge spin embedding.

has_default_chg_spin() → bool

Check if the model has default charge_spin values.

get_intensive() → bool

get_var_name() → str

Get the name of the property.

property model_type: type[deepmd.infer.deep_eval.DeepEval]

The the evaluator of the model type.

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.

get_numb_dos() → int

Get the number of DOS.

get_task_dim() → int

Get the output dimension.

get_has_efield() → bool

Check if the model has efield.

get_ntypes_spin() → int

Get the number of spin atom types of this model. Only used in old implement.

get_has_spin() → bool

Check if the model has spin atom types.

get_use_spin() → list[bool]

Get the per-type spin usage of this model.

get_has_hessian() → bool

Check if the model has hessian.

get_model_branch() → tuple[dict[str, str], dict[str, dict[str, Any]]]

Get the model branch information.

eval(coords: numpy.ndarray, cells: numpy.ndarray | None, atom_types: numpy.ndarray, atomic: bool = False, fparam: numpy.ndarray | None = None, aparam: numpy.ndarray | None = None, charge_spin: numpy.ndarray | None = None, **kwargs: Any) → dict[str, numpy.ndarray]

Evaluate the energy, force and virial by using this DP.

Parameters:

coords

The coordinates of atoms. The array should be of size nframes x natoms x 3

cells

The cell of the region. If None then non-PBC is assumed, otherwise using PBC. The array should be of size nframes x 9

atom_types

The atom types The list should contain natoms ints

atomic

Calculate the atomic energy and virial

fparam

The frame parameter. The array can be of size : - nframes x dim_fparam. - dim_fparam. Then all frames are assumed to be provided with the same fparam.

aparam

The atomic parameter The array can be of size : - nframes x natoms x dim_aparam. - natoms x dim_aparam. Then all frames are assumed to be provided with the same aparam. - dim_aparam. Then all frames and atoms are provided with the same aparam.

**kwargs

Other parameters

Returns:

output_dictdict

The output of the evaluation. The keys are the names of the output variables, and the values are the corresponding output arrays.

_get_request_defs(atomic: bool) → list[deepmd.dpmodel.output_def.OutputVariableDef]

Get the requested output definitions.

When atomic is True, all output_def are requested. When atomic is False, only energy (tensor), force, and virial are requested.

Parameters:

atomicbool

Whether to request the atomic output.

Returns:

list[OutputVariableDef]

The requested output definitions.

_eval_func(inner_func: collections.abc.Callable, numb_test: int, natoms: int) → collections.abc.Callable

Wrapper method with auto batch size.

Parameters:

inner_funcCallable

the method to be wrapped

numb_testint

number of tests

natomsint

number of atoms

Returns:

Callable

the wrapper

_get_natoms_and_nframes(coords: numpy.ndarray, atom_types: numpy.ndarray, mixed_type: bool = False) → tuple[int, int]

_eval_model(coords: numpy.ndarray, cells: numpy.ndarray | None, atom_types: numpy.ndarray, fparam: numpy.ndarray | None, aparam: numpy.ndarray | None, request_defs: list[deepmd.dpmodel.output_def.OutputVariableDef], charge_spin: numpy.ndarray | None) → tuple[numpy.ndarray, Ellipsis]

_eval_lower_strategy(coord: torch.Tensor, atype: torch.Tensor, box: torch.Tensor | None, fparam: torch.Tensor | None, aparam: torch.Tensor | None, charge_spin: torch.Tensor | None, do_atomic_virial: bool) → dict[str, torch.Tensor]

Evaluate via the selected O(N) NeighborList strategy.

Uses the selected O(N) builder (vesin or nv). Models that declare the edge-vector contract consume it directly; other energy models keep the historical extended-coordinate contract and fold extended outputs back to local atoms. Returns a dict keyed by backend names, matching the normal model() output so the caller’s extraction is unchanged.

_eval_model_spin(coords: numpy.ndarray, cells: numpy.ndarray | None, atom_types: numpy.ndarray, spins: numpy.ndarray, fparam: numpy.ndarray | None, aparam: numpy.ndarray | None, request_defs: list[deepmd.dpmodel.output_def.OutputVariableDef], charge_spin: numpy.ndarray | None) → tuple[numpy.ndarray, Ellipsis]

_get_output_shape(odef: deepmd.dpmodel.output_def.OutputVariableDef, nframes: int, natoms: int) → list[int]

eval_typeebd() → numpy.ndarray

Evaluate output of type embedding network by using this model.

Returns:

np.ndarray

The output of type embedding network. The shape is [ntypes, o_size] or [ntypes + 1, o_size], where ntypes is the number of types, and o_size is the number of nodes in the output layer. If there are multiple type embedding networks, these outputs will be concatenated along the second axis.

Raises:

KeyError

If the model does not enable type embedding.

See also

deepmd.pt.model.network.network.TypeEmbedNetConsistent

The type embedding network.

get_model_def_script() → dict

Get model definition script.

serialize() → dict[str, Any]

Serialize the loaded model as a model tree.

Most in-tree backends return the lossless, weight-bearing model subtree from the serialized file payload. Backends that cannot recover a lossless tree may override this method to document and implement their narrower behavior.

Returns:

dict

Serialized model tree that can be consumed by Node.deserialize.

get_model_size() → dict

Get model parameter count.

Returns:

dict

A dictionary containing the number of parameters in the model. The keys are ‘descriptor’, ‘fitting_net’, and ‘total’.

get_observed_types() → dict

Get observed types (elements) of the model during data statistics.

Returns:

dict

A dictionary containing the information of observed type in the model: - ‘type_num’: the total number of observed types in this model. - ‘observed_type’: a list of the observed types in this model.

get_model() → deepmd.pt.model.model.model.BaseModel

Get the PyTorch model.

Returns:

BaseModel

The PyTorch model instance.

eval_descriptor(coords: numpy.ndarray, cells: numpy.ndarray | None, atom_types: numpy.ndarray, fparam: numpy.ndarray | None = None, aparam: numpy.ndarray | None = None, **kwargs: Any) → numpy.ndarray

Evaluate descriptors by using this DP.

Deprecated since version Use: eval_embedding() instead, which returns the descriptor together with the atomic and structural features in a single forward pass. This method is a thin wrapper kept for compatibility. For models frozen before forward_embedding existed, it falls back to the descriptor hook baked into that TorchScript module.

Parameters:

coords

The coordinates of atoms. The array should be of size nframes x natoms x 3

cells

The cell of the region. If None then non-PBC is assumed, otherwise using PBC. The array should be of size nframes x 9

atom_types

The atom types The list should contain natoms ints

fparam

The frame parameter. The array can be of size : - nframes x dim_fparam. - dim_fparam. Then all frames are assumed to be provided with the same fparam.

aparam

The atomic parameter The array can be of size : - nframes x natoms x dim_aparam. - natoms x dim_aparam. Then all frames are assumed to be provided with the same aparam. - dim_aparam. Then all frames and atoms are provided with the same aparam.

Returns:

descriptor

Descriptors.

eval_fitting_last_layer(coords: numpy.ndarray, cells: numpy.ndarray | None, atom_types: numpy.ndarray, fparam: numpy.ndarray | None = None, aparam: numpy.ndarray | None = None, **kwargs: Any) → numpy.ndarray

Evaluate fitting before last layer by using this DP.

Deprecated since version Use: eval_embedding() instead, which returns this activation as the atomic_feature output. This method is a thin wrapper kept for compatibility. For models frozen before forward_embedding existed, it falls back to the fitting-last-layer hook baked into that TorchScript module.

Parameters:

coords

The coordinates of atoms. The array should be of size nframes x natoms x 3

cells

The cell of the region. If None then non-PBC is assumed, otherwise using PBC. The array should be of size nframes x 9

atom_types

The atom types The list should contain natoms ints

fparam

The frame parameter. The array can be of size : - nframes x dim_fparam. - dim_fparam. Then all frames are assumed to be provided with the same fparam.

aparam

The atomic parameter The array can be of size : - nframes x natoms x dim_aparam. - natoms x dim_aparam. Then all frames are assumed to be provided with the same aparam. - dim_aparam. Then all frames and atoms are provided with the same aparam.

Returns:

fitting

Fitting output before last layer.

_eval_legacy_feature(coords: numpy.ndarray, cells: numpy.ndarray | None, atom_types: numpy.ndarray, fparam: numpy.ndarray | None, aparam: numpy.ndarray | None, *, enable_hook: collections.abc.Callable[[bool], None], read_feature: collections.abc.Callable[[], torch.Tensor], **kwargs: Any) → numpy.ndarray

Extract a cached descriptor or fitting feature from a legacy frozen model.

Models frozen before forward_embedding expose these features only through a hook that caches them during a forward pass. The retry loop keeps the cache consistent when the auto batch size splits the forward and hits an out-of-memory condition.

eval_embedding(coords: numpy.ndarray, cells: numpy.ndarray | None, atom_types: numpy.ndarray, fparam: numpy.ndarray | None = None, aparam: numpy.ndarray | None = None, charge_spin: numpy.ndarray | None = None, dtype: str = 'fp32', **kwargs: Any) → tuple[numpy.ndarray, numpy.ndarray, numpy.ndarray]

Evaluate the descriptor, atomic feature, and structural feature.

A single forward pass produces all three embeddings without any force or virial autograd. The descriptor is the per-atom local-environment representation; the atomic feature is the activation after the last fitting hidden layer; the structural feature is the masked atom-sum of the atomic feature, a whole-structure summary. For models with a single shared fitting network, projecting the structural feature through the fitting output layer reproduces the (bias-free) total energy. The output precision is selected by dtype and defaults to float32.

Parameters:

coords

The coordinates of atoms. The array should be of size nframes x natoms x 3

cells

The cell of the region. If None then non-PBC is assumed, otherwise using PBC. The array should be of size nframes x 9

atom_types

The atom types The list should contain natoms ints

fparam

The frame parameter. The array can be of size : - nframes x dim_fparam. - dim_fparam. Then all frames are assumed to be provided with the same fparam.

aparam

The atomic parameter The array can be of size : - nframes x natoms x dim_aparam. - natoms x dim_aparam. Then all frames are assumed to be provided with the same aparam. - dim_aparam. Then all frames and atoms are provided with the same aparam.

charge_spin

The frame-level charge and spin conditions. The array should be of size nframes x 2

dtype

Output dtype: "fp32", "fp64", or "native".

Returns:

descriptor

The per-atom descriptor, of size nframes x natoms x dim_descriptor.

atomic_feature

The per-atom last hidden activation, of size nframes x natoms x dim_hidden.

structural_feature

The per-structure pooled feature, of size nframes x dim_hidden.

Raises:

NotImplementedError

If the loaded model does not support embedding extraction.

_eval_embedding(coords: numpy.ndarray, cells: numpy.ndarray | None, atom_types: numpy.ndarray, fparam: numpy.ndarray | None, aparam: numpy.ndarray | None, charge_spin: numpy.ndarray | None, dtype: str) → tuple[numpy.ndarray, numpy.ndarray, numpy.ndarray]