deepmd.tf.infer

Submodule containing all the implemented potentials.

Submodules

Package Contents

Classes

DipoleChargeModifier

Parameters:

DeepDipole

Deep dipole model.

DeepDOS

Deep density of states model.

DeepEval

TensorFlow backend implementation for DeepEval.

DeepGlobalPolar

Deep Tensor Model.

DeepPolar

Deep polar model.

DeepPot

Potential energy model.

DeepWFC

Deep WFC model.

EwaldRecp

Evaluate the reciprocal part of the Ewald sum.

Functions

DeepPotential(→ deep_eval.DeepEval)

Factory function that forwards to DeepEval (for compatbility).

calc_model_devi(coord, box, atype, models[, fname, ...])

Python interface to calculate model deviation.

deepmd.tf.infer.DeepPotential(*args, **kwargs) deep_eval.DeepEval[source]

Factory function that forwards to DeepEval (for compatbility).

Parameters:
*args

positional arguments

**kwargs

keyword arguments

Returns:
DeepEval

potentials

class deepmd.tf.infer.DipoleChargeModifier(model_name: str, model_charge_map: List[float], sys_charge_map: List[float], ewald_h: float = 1, ewald_beta: float = 1)[source]

Bases: deepmd.tf.infer.deep_dipole.DeepDipoleOld

Parameters:
model_name

The model file for the DeepDipole model

model_charge_map

Gives the amount of charge for the wfcc

sys_charge_map

Gives the amount of charge for the real atoms

ewald_h

Grid spacing of the reciprocal part of Ewald sum. Unit: A

ewald_beta

Splitting parameter of the Ewald sum. Unit: A^{-1}

build_fv_graph() deepmd.tf.env.tf.Tensor[source]

Build the computational graph for the force and virial inference.

_build_fv_graph_inner()[source]
_enrich(dipole, dof=3)[source]
_slice_descrpt_deriv(deriv)[source]
eval(coord: numpy.ndarray, box: numpy.ndarray, atype: numpy.ndarray, eval_fv: bool = True) Tuple[numpy.ndarray, numpy.ndarray, numpy.ndarray][source]

Evaluate the modification.

Parameters:
coord

The coordinates of atoms

box

The simulation region. PBC is assumed

atype

The atom types

eval_fv

Evaluate force and virial

Returns:
tot_e

The energy modification

tot_f

The force modification

tot_v

The virial modification

_eval_fv(coords, cells, atom_types, ext_f)[source]
_extend_system(coord, box, atype, charge)[source]
modify_data(data: dict, data_sys: deepmd.tf.utils.data.DeepmdData) None[source]

Modify data.

Parameters:
data

Internal data of DeepmdData. Be a dict, has the following keys - coord coordinates - box simulation box - type atom types - find_energy tells if data has energy - find_force tells if data has force - find_virial tells if data has virial - energy energy - force force - virial virial

data_sysDeepmdData

The data system.

class deepmd.tf.infer.DeepDipole(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])[source]

Bases: deepmd.infer.deep_tensor.DeepTensor

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

property output_tensor_name: str

The name of the tensor.

class deepmd.tf.infer.DeepDOS(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])[source]

Bases: deepmd.infer.deep_eval.DeepEval

Deep density of states 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.

property output_def: deepmd.dpmodel.output_def.ModelOutputDef

Get the output definition of this model.

property numb_dos: int

Get the number of DOS.

eval(coords: numpy.ndarray, cells: numpy.ndarray | None, atom_types: List[int] | numpy.ndarray, atomic: bool = False, fparam: numpy.ndarray | None = None, aparam: numpy.ndarray | None = None, mixed_type: bool = False, **kwargs: Dict[str, Any]) Tuple[numpy.ndarray, Ellipsis][source]

Evaluate energy, force, and virial. If atomic is True, also return atomic energy and atomic virial.

Parameters:
coordsnp.ndarray

The coordinates of the atoms, in shape (nframes, natoms, 3).

cellsnp.ndarray

The cell vectors of the system, in shape (nframes, 9). If the system is not periodic, set it to None.

atom_typesList[int] or np.ndarray

The types of the atoms. If mixed_type is False, the shape is (natoms,); otherwise, the shape is (nframes, natoms).

atomicbool, optional

Whether to return atomic energy and atomic virial, by default False.

fparamnp.ndarray, optional

The frame parameters, by default None.

aparamnp.ndarray, optional

The atomic parameters, by default None.

mixed_typebool, optional

Whether the atom_types is mixed type, by default False.

**kwargsDict[str, Any]

Keyword arguments.

Returns:
energy

The energy of the system, in shape (nframes,).

force

The force of the system, in shape (nframes, natoms, 3).

virial

The virial of the system, in shape (nframes, 9).

atomic_energy

The atomic energy of the system, in shape (nframes, natoms). Only returned when atomic is True.

atomic_virial

The atomic virial of the system, in shape (nframes, natoms, 9). Only returned when atomic is True.

get_numb_dos() int[source]
class deepmd.tf.infer.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

Get type of model.

:type:str

property model_version: str

Get version of model.

Returns:
str

version of model

property sess: deepmd.tf.env.tf.Session

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.DeepGlobalPolar(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])[source]

Bases: deepmd.infer.deep_tensor.DeepTensor

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.

property output_tensor_name: str

The name of the tensor.

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

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

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

class deepmd.tf.infer.DeepPolar(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])[source]

Bases: deepmd.infer.deep_tensor.DeepTensor

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

property output_tensor_name: str

The name of the tensor.

class deepmd.tf.infer.DeepPot(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])[source]

Bases: deepmd.infer.deep_eval.DeepEval

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

Examples

>>> from deepmd.infer import DeepPot
>>> import numpy as np
>>> dp = DeepPot("graph.pb")
>>> coord = np.array([[1, 0, 0], [0, 0, 1.5], [1, 0, 3]]).reshape([1, -1])
>>> cell = np.diag(10 * np.ones(3)).reshape([1, -1])
>>> atype = [1, 0, 1]
>>> e, f, v = dp.eval(coord, cell, atype)

where e, f and v are predicted energy, force and virial of the system, respectively.

property output_def: deepmd.dpmodel.output_def.ModelOutputDef

Get the output definition of this model.

property output_def_mag: deepmd.dpmodel.output_def.ModelOutputDef

Get the output definition of this model with magnetic parts.

eval(coords: numpy.ndarray, cells: numpy.ndarray | None, atom_types: List[int] | numpy.ndarray, atomic: bool = False, fparam: numpy.ndarray | None = None, aparam: numpy.ndarray | None = None, mixed_type: bool = False, **kwargs: Dict[str, Any]) Tuple[numpy.ndarray, Ellipsis][source]

Evaluate energy, force, and virial. If atomic is True, also return atomic energy and atomic virial.

Parameters:
coordsnp.ndarray

The coordinates of the atoms, in shape (nframes, natoms, 3).

cellsnp.ndarray

The cell vectors of the system, in shape (nframes, 9). If the system is not periodic, set it to None.

atom_typesList[int] or np.ndarray

The types of the atoms. If mixed_type is False, the shape is (natoms,); otherwise, the shape is (nframes, natoms).

atomicbool, optional

Whether to return atomic energy and atomic virial, by default False.

fparamnp.ndarray, optional

The frame parameters, by default None.

aparamnp.ndarray, optional

The atomic parameters, by default None.

mixed_typebool, optional

Whether the atom_types is mixed type, by default False.

**kwargsDict[str, Any]

Keyword arguments.

Returns:
energy

The energy of the system, in shape (nframes,).

force

The force of the system, in shape (nframes, natoms, 3).

virial

The virial of the system, in shape (nframes, 9).

atomic_energy

The atomic energy of the system, in shape (nframes, natoms). Only returned when atomic is True.

atomic_virial

The atomic virial of the system, in shape (nframes, natoms, 9). Only returned when atomic is True.

class deepmd.tf.infer.DeepWFC(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])[source]

Bases: deepmd.infer.deep_tensor.DeepTensor

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

property output_tensor_name: str

The name of the tensor.

class deepmd.tf.infer.EwaldRecp(hh, beta)[source]

Evaluate the reciprocal part of the Ewald sum.

eval(coord: numpy.ndarray, charge: numpy.ndarray, box: numpy.ndarray) Tuple[numpy.ndarray, numpy.ndarray, numpy.ndarray][source]

Evaluate.

Parameters:
coord

The coordinates of atoms

charge

The atomic charge

box

The simulation region. PBC is assumed

Returns:
e

The energy

f

The force

v

The virial

deepmd.tf.infer.calc_model_devi(coord, box, atype, models, fname=None, frequency=1, mixed_type=False, fparam: numpy.ndarray | None = None, aparam: numpy.ndarray | None = None, real_data: dict | None = None, atomic: bool = False, relative: float | None = None, relative_v: float | None = None)[source]

Python interface to calculate model deviation.

Parameters:
coordnumpy.ndarray, n_frames x n_atoms x 3

Coordinates of system to calculate

boxnumpy.ndarray or None, n_frames x 3 x 3

Box to specify periodic boundary condition. If None, no pbc will be used

atypenumpy.ndarray, n_atoms x 1

Atom types

modelslist of DeepPot models

Models used to evaluate deviation

fnamestr or None

File to dump results, default None

frequencyint

Steps between frames (if the system is given by molecular dynamics engine), default 1

mixed_typebool

Whether the input atype is in mixed_type format or not

fparamnumpy.ndarray

frame specific parameters

aparamnumpy.ndarray

atomic specific parameters

real_datadict, optional

real data to calculate RMS real error

atomicbool, default: False

If True, calculate the force model deviation of each atom.

relativefloat, default: None

If given, calculate the relative model deviation of force. The value is the level parameter for computing the relative model deviation of the force.

relative_vfloat, default: None

If given, calculate the relative model deviation of virial. The value is the level parameter for computing the relative model deviation of the virial.

Returns:
model_devinumpy.ndarray, n_frames x 8

Model deviation results. The first column is index of steps, the other 7 columns are max_devi_v, min_devi_v, avg_devi_v, max_devi_f, min_devi_f, avg_devi_f, devi_e.

Examples

>>> from deepmd.tf.infer import calc_model_devi
>>> from deepmd.tf.infer import DeepPot as DP
>>> import numpy as np
>>> coord = np.array([[1, 0, 0], [0, 0, 1.5], [1, 0, 3]]).reshape([1, -1])
>>> cell = np.diag(10 * np.ones(3)).reshape([1, -1])
>>> atype = [1, 0, 1]
>>> graphs = [DP("graph.000.pb"), DP("graph.001.pb")]
>>> model_devi = calc_model_devi(coord, cell, atype, graphs)