deepmd.pt_expt.infer.deep_eval

Classes

DeepEval

PyTorch Exportable backend implementation of DeepEval.

Functions

_reconstruct_model_output_def(...)

Reconstruct ModelOutputDef from stored fitting_output_defs metadata.

Module Contents

deepmd.pt_expt.infer.deep_eval._reconstruct_model_output_def(metadata: dict) → deepmd.dpmodel.output_def.ModelOutputDef

Reconstruct ModelOutputDef from stored fitting_output_defs metadata.

class deepmd.pt_expt.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, **kwargs: Any)

Bases: deepmd.infer.deep_eval.DeepEvalBackend

PyTorch Exportable backend implementation of DeepEval.

Loads a .pte file containing a torch.export-ed model and evaluates it using pre-built neighbor lists.

Parameters:

model_filePath

The name of the .pte model file.

output_defModelOutputDef

The output definition of the model.

*argslist

Positional arguments.

auto_batch_sizebool or int or AutoBatchSize, 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

exported_module

metadata

rcut

type_map

_model_output_def

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.

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_has_efield() → bool

Check if the model has efield.

get_ntypes_spin() → int

Get the number of spin atom types of this model.

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, **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 should be of size nframes x dim_fparam.

aparam

The atomic parameter. The array should be of size nframes x natoms x dim_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.

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

Wrapper method with auto batch size.

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

_build_nlist_native(coords: torch.Tensor, cells: torch.Tensor | None, atom_types: torch.Tensor) → tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]

Build extended coords, atype, nlist, mapping using native nlist.

Parameters:

coordstorch.Tensor

Coordinates, shape (nframes, natoms, 3).

cellstorch.Tensor or None

Cell vectors, shape (nframes, 9). None for non-PBC.

atom_typestorch.Tensor

Atom types, shape (nframes, natoms).

Returns:

extended_coord, extended_atype, nlist, mapping

All as torch.Tensor on the same device as inputs.

_build_nlist_ase(coords: numpy.ndarray, cells: numpy.ndarray | None, atom_types: numpy.ndarray) → tuple[numpy.ndarray, numpy.ndarray, numpy.ndarray, numpy.ndarray]

Build extended coords, atype, nlist, mapping using ASE neighbor list.

Handles multiple frames by building per frame and padding to a common nall.

Parameters:

coordsnp.ndarray

Coordinates, shape (nframes, natoms, 3).

cellsnp.ndarray or None

Cell vectors, shape (nframes, 9). None for non-PBC.

atom_typesnp.ndarray

Atom types, shape (nframes, natoms).

Returns:

extended_coord, extended_atype, nlist, mapping

_build_nlist_ase_single(positions: numpy.ndarray, cell: numpy.ndarray | None, atype: numpy.ndarray) → tuple[numpy.ndarray, numpy.ndarray, numpy.ndarray, numpy.ndarray]

Build extended coords, atype, nlist, mapping for a single frame.

Parameters:

positionsnp.ndarray

Atom positions, shape (natoms, 3).

cellnp.ndarray or None

Cell vector, shape (9,). None for non-PBC.

atypenp.ndarray

Atom types, shape (natoms,).

Returns:

extended_coordnp.ndarray, shape (nall, 3)

extended_atypenp.ndarray, shape (nall,)

nlistnp.ndarray, shape (nloc, nsel)

mappingnp.ndarray, shape (nall,)

_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]) → tuple[numpy.ndarray, Ellipsis]

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

get_model_def_script() → dict

Get model definition script.

get_model() → torch.nn.Module

Get the exported model module.

Returns:

torch.nn.Module

The exported model module.