`deepmd.tf.infer.deep_eval`

Module Contents

Classes

`DeepEval`	TensorFlow backend implementation for DeepEval.
`DeepEvalOld`	Common methods for DeepPot, DeepWFC, DeepPolar, ...

class deepmd.tf.infer.deep_eval.DeepEval(model_file: pathlib.Path, output_def: deepmd.dpmodel.output_def.ModelOutputDef, *args: list, load_prefix: str = 'load', default_tf_graph: bool = False, auto_batch_size: bool | int | deepmd.tf.utils.batch_size.AutoBatchSize = False, input_map: dict | None = None, neighbor_list=None, **kwargs: dict)[source]

Bases: deepmd.infer.deep_eval.DeepEvalBackend

TensorFlow backend implementation for DeepEval.

Parameters:

model_filePath: The name of the frozen model file.
output_defModelOutputDef: The output definition of the model.
*argslist: Positional arguments.
load_prefix: str: The prefix in the load computational graph
default_tf_graphbool: If uses the default tf graph, otherwise build a new tf graph for evaluation
auto_batch_sizebool or int or AutomaticBatchSize, default: False: If True, automatic batch size will be used. If int, it will be used as the initial batch size.
input_mapdict, optional: The input map for tf.import_graph_def. Only work with default tf graph
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.

property model_type: deepmd.infer.deep_eval.DeepEval[source]

Get type of model.

:type:str

property model_version: str[source]

Get version of model.

Returns:

str: version of model

property sess: deepmd.tf.env.tf.Session[source]: Get TF session.

_init_tensors()[source]

_init_attr()[source]

_graph_compatable() → bool[source]

Check the model compatability.

Returns:

bool: If the model stored in the graph file is compatable with the current code

_get_tensor(tensor_name: str) → deepmd.tf.env.tf.Tensor[source]

Get TF graph tensor.

Parameters:

tensor_namestr: name of tensor to get

Returns:

tf.Tensor: loaded tensor

static _load_graph(frozen_graph_filename: pathlib.Path, prefix: str = 'load', default_tf_graph: bool = False, input_map: dict | None = None)[source]

static sort_input(coord: numpy.ndarray, atom_type: numpy.ndarray, sel_atoms: List[int] | None = None)[source]

Sort atoms in the system according their types.

Parameters:

coord: The coordinates of atoms. Should be of shape [nframes, natoms, 3]
atom_type: The type of atoms Should be of shape [natoms]
sel_atoms: The selected atoms by type

Returns:

coord_out: The coordinates after sorting
atom_type_out: The atom types after sorting
idx_map: The index mapping from the input to the output. For example coord_out = coord[:,idx_map,:]
sel_atom_type: Only output if sel_atoms is not None The sorted selected atom types
sel_idx_map: Only output if sel_atoms is not None The index mapping from the selected atoms to sorted selected atoms.

static reverse_map(vec: numpy.ndarray, imap: List[int]) → numpy.ndarray[source]

Reverse mapping of a vector according to the index map.

Parameters:

vec: Input vector. Be of shape [nframes, natoms, -1]
imap: Index map. Be of shape [natoms]

Returns:

vec_out: Reverse mapped vector.

make_natoms_vec(atom_types: numpy.ndarray) → numpy.ndarray[source]

Make the natom vector used by deepmd-kit.

Parameters:

atom_types: The type of atoms

Returns:

natoms: The number of atoms. This tensor has the length of Ntypes + 2 natoms[0]: number of local atoms natoms[1]: total number of atoms held by this processor natoms[i]: 2 <= i < Ntypes+2, number of type i atoms

eval_typeebd() → numpy.ndarray[source]

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], where ntypes is the number of types, and o_size is the number of nodes in the output layer.

Raises:

KeyError: If the model does not enable type embedding.

See also

deepmd.tf.utils.type_embed.TypeEmbedNet: The type embedding network.

Examples

Get the output of type embedding network of graph.pb:

>>> from deepmd.tf.infer import DeepPotential
>>> dp = DeepPotential("graph.pb")
>>> dp.eval_typeebd()

build_neighbor_list(coords: numpy.ndarray, cell: numpy.ndarray | None, atype: numpy.ndarray, imap: numpy.ndarray, neighbor_list)[source]

Make the mesh with neighbor list for a single frame.

Parameters:

coordsnp.ndarray: The coordinates of atoms. Should be of shape [natoms, 3]
cellOptional[np.ndarray]: The cell of the system. Should be of shape [3, 3]
atypenp.ndarray: The type of atoms. Should be of shape [natoms]
imapnp.ndarray: The index map of atoms. Should be of shape [natoms]
neighbor_listase.neighborlist.NewPrimitiveNeighborList: ASE neighbor list. The following method or attribute will be used/set: bothways, self_interaction, update, build, first_neigh, pair_second, offset_vec.

Returns:

natoms_vecnp.ndarray: The number of atoms. This tensor has the length of Ntypes + 2 natoms[0]: nloc natoms[1]: nall natoms[i]: 2 <= i < Ntypes+2, number of type i atoms for nloc
coordsnp.ndarray: The coordinates of atoms, including ghost atoms. Should be of shape [nframes, nall, 3]
atypenp.ndarray: The type of atoms, including ghost atoms. Should be of shape [nall]
meshnp.ndarray: The mesh in nei_mode=4.
imapnp.ndarray: The index map of atoms. Should be of shape [nall]
ghost_mapnp.ndarray: The index map of ghost atoms. Should be of shape [nghost]

get_ntypes() → int[source]: Get the number of atom types of this model.

get_ntypes_spin() → int[source]: Get the number of spin atom types of this model.

get_rcut() → float[source]: Get the cut-off radius of this model.

get_type_map() → List[str][source]: Get the type map (element name of the atom types) of this model.

get_sel_type() → numpy.ndarray | None[source]

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_dim_fparam() → int[source]: Get the number (dimension) of frame parameters of this DP.

get_dim_aparam() → int[source]: Get the number (dimension) of atomic parameters of this DP.

_eval_func(inner_func: Callable, numb_test: int, natoms: int) → Callable[source]

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: List[int] | numpy.ndarray) → Tuple[int, int][source]

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

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.
efield: The external field on atoms. The array should be of size nframes x natoms x 3
**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.

_prepare_feed_dict(coords, cells, atom_types, fparam=None, aparam=None, efield=None)[source]

_eval_inner(coords, cells, atom_types, fparam=None, aparam=None, efield=None, **kwargs)[source]

_get_output_shape(odef, nframes, natoms)[source]

eval_descriptor(coords: numpy.ndarray, cells: numpy.ndarray, atom_types: numpy.ndarray, fparam: numpy.ndarray | None = None, aparam: numpy.ndarray | None = None, efield: numpy.ndarray | None = None) → numpy.ndarray[source]

Evaluate descriptors 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
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.
efield: The external field on atoms. The array should be of size nframes x natoms x 3

Returns:

descriptor: Descriptors.

_eval_descriptor_inner(coords: numpy.ndarray, cells: numpy.ndarray, atom_types: numpy.ndarray, fparam: numpy.ndarray | None = None, aparam: numpy.ndarray | None = None, efield: numpy.ndarray | None = None) → numpy.ndarray[source]

get_numb_dos() → int[source]: Get the number of DOS.

get_has_efield() → bool[source]: Check if the model has efield.

class deepmd.tf.infer.deep_eval.DeepEvalOld(model_file: pathlib.Path, load_prefix: str = 'load', default_tf_graph: bool = False, auto_batch_size: bool | int | deepmd.tf.utils.batch_size.AutoBatchSize = False, input_map: dict | None = None, neighbor_list=None)[source]

Common methods for DeepPot, DeepWFC, DeepPolar, …

Parameters:

model_filePath: The name of the frozen model file.
load_prefix: str: The prefix in the load computational graph
default_tf_graphbool: If uses the default tf graph, otherwise build a new tf graph for evaluation
auto_batch_sizebool or int or AutomaticBatchSize, default: False: If True, automatic batch size will be used. If int, it will be used as the initial batch size.
input_mapdict, optional: The input map for tf.import_graph_def. Only work with default tf graph
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.

property model_type: str[source]

Get type of model.

:type:str

property model_version: str[source]

Get version of model.

Returns:

str: version of model

property sess: deepmd.tf.env.tf.Session[source]: Get TF session.

load_prefix: str[source]

_graph_compatable() → bool[source]

Check the model compatability.

Returns:

bool: If the model stored in the graph file is compatable with the current code

_get_tensor(tensor_name: str, attr_name: str | None = None) → deepmd.tf.env.tf.Tensor[source]

Get TF graph tensor and assign it to class namespace.

Parameters:

tensor_namestr: name of tensor to get
attr_nameOptional[str], optional: if specified, class attribute with this name will be created and tensor will be assigned to it, by default None

Returns:

tf.Tensor: loaded tensor

static _load_graph(frozen_graph_filename: pathlib.Path, prefix: str = 'load', default_tf_graph: bool = False, input_map: dict | None = None)[source]

static sort_input(coord: numpy.ndarray, atom_type: numpy.ndarray, sel_atoms: List[int] | None = None, mixed_type: bool = False)[source]

Sort atoms in the system according their types.

Parameters:

coord: The coordinates of atoms. Should be of shape [nframes, natoms, 3]
atom_type: The type of atoms Should be of shape [natoms]
sel_atoms: The selected atoms by type
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:

coord_out: The coordinates after sorting
atom_type_out: The atom types after sorting
idx_map: The index mapping from the input to the output. For example coord_out = coord[:,idx_map,:]
sel_atom_type: Only output if sel_atoms is not None The sorted selected atom types
sel_idx_map: Only output if sel_atoms is not None The index mapping from the selected atoms to sorted selected atoms.

static reverse_map(vec: numpy.ndarray, imap: List[int]) → numpy.ndarray[source]

Reverse mapping of a vector according to the index map.

Parameters:

vec: Input vector. Be of shape [nframes, natoms, -1]
imap: Index map. Be of shape [natoms]

Returns:

vec_out: Reverse mapped vector.

make_natoms_vec(atom_types: numpy.ndarray, mixed_type: bool = False) → numpy.ndarray[source]

Make the natom vector used by deepmd-kit.

Parameters:

atom_types: The type of atoms
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:

natoms: The number of atoms. This tensor has the length of Ntypes + 2 natoms[0]: number of local atoms natoms[1]: total number of atoms held by this processor natoms[i]: 2 <= i < Ntypes+2, number of type i atoms

eval_typeebd() → numpy.ndarray[source]

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], where ntypes is the number of types, and o_size is the number of nodes in the output layer.

Raises:

KeyError: If the model does not enable type embedding.

deepmd.tf.infer.deep_eval

Module Contents

Classes

`deepmd.tf.infer.deep_eval`