deepmd.infer.deep_tensor

Module Contents

Classes

DeepTensor

Deep Tensor Model.

class deepmd.infer.deep_tensor.DeepTensor(model_file: str, *args: List[Any], auto_batch_size: bool | int | deepmd.utils.batch_size.AutoBatchSize = True, neighbor_list: ase.neighborlist.NewPrimitiveNeighborList | None = None, **kwargs: Dict[str, Any])

Bases: deepmd.infer.deep_eval.DeepEval

Deep Tensor Model.

Parameters:

model_filePath

The name of the frozen model file.

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

abstract property output_tensor_name: str

The name of the tensor.

property output_def: deepmd.dpmodel.output_def.ModelOutputDef

Get the output definition of this model.

eval(coords: numpy.ndarray, cells: numpy.ndarray | None, atom_types: List[int] | numpy.ndarray, atomic: bool = True, fparam: numpy.ndarray | None = None, aparam: numpy.ndarray | None = None, mixed_type: bool = False, **kwargs: dict) → numpy.ndarray

Evaluate the model.

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_typeslist[int] or np.ndarray

The atom types The list should contain natoms ints

atomic

If True (default), return the atomic tensor Otherwise return the global tensor

fparam

Not used in this model

aparam

Not used in this model

efield

Not used in this model

mixed_type

Whether to perform the mixed_type mode. If True, the input data has the mixed_type format (see doc/model/train_se_atten.md), in which frames in a system may have different natoms_vec(s), with the same nloc.

Returns:

tensor

The returned tensor If atomic == False then of size nframes x output_dim else of size nframes x natoms x output_dim

eval_full(coords: numpy.ndarray, cells: numpy.ndarray | None, atom_types: numpy.ndarray, atomic: bool = False, fparam: numpy.ndarray | None = None, aparam: numpy.ndarray | None = None, mixed_type: bool = False, **kwargs: dict) → Tuple[numpy.ndarray, Ellipsis]

Evaluate the model with interface similar to the energy model. Will return global tensor, component-wise force and virial and optionally atomic tensor and atomic virial.

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

Whether to calculate atomic tensor and virial

fparam

Not used in this model

aparam

Not used in this model

mixed_type

Whether to perform the mixed_type mode. If True, the input data has the mixed_type format (see doc/model/train_se_atten.md), in which frames in a system may have different natoms_vec(s), with the same nloc.

Returns:

tensor

The global tensor. shape: [nframes x nout]

force

The component-wise force (negative derivative) on each atom. shape: [nframes x nout x natoms x 3]

virial

The component-wise virial of the tensor. shape: [nframes x nout x 9]

atom_tensor

The atomic tensor. Only returned when atomic == True shape: [nframes x natoms x nout]

atom_virial

The atomic virial. Only returned when atomic == True shape: [nframes x nout x natoms x 9]