deepmd.tf package

Root of the deepmd package, exposes all public classes and submodules.

class deepmd.tf.DeepEval(model_file: str, *args, **kwargs)[source]

Bases: 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.

Attributes
model_type

Get type of model.

model_version

Get version of model.

sess

Get TF session.

Methods

build_neighbor_list(coords, cell, atype, ...)

Make the mesh with neighbor list for a single frame.

eval(coords, cells, atom_types[, atomic, ...])

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

eval_descriptor(coords, cells, atom_types[, ...])

Evaluate descriptors by using this DP.

eval_typeebd()

Evaluate output of type embedding network by using this model.

get_dim_aparam()

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

get_dim_fparam()

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

get_has_efield()

Check if the model has efield.

get_ntypes()

Get the number of atom types of this model.

get_ntypes_spin()

Get the number of spin atom types of this model.

get_numb_dos()

Get the number of DOS.

get_rcut()

Get the cut-off radius of this model.

get_sel_type()

Get the selected atom types of this model.

get_type_map()

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

make_natoms_vec(atom_types)

Make the natom vector used by deepmd-kit.

reverse_map(vec, imap)

Reverse mapping of a vector according to the index map.

sort_input(coord, atom_type[, sel_atoms])

Sort atoms in the system according their types.
build_neighbor_list(coords: ndarray, cell: Optional[ndarray], atype: ndarray, imap: 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]

eval(coords: ndarray, cells: ndarray, atom_types: ndarray, atomic: bool = False, fparam: Optional[ndarray] = None, aparam: Optional[ndarray] = None, efield: Optional[ndarray] = None) → Dict[str, 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

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.

eval_descriptor(coords: ndarray, cells: ndarray, atom_types: ndarray, fparam: Optional[ndarray] = None, aparam: Optional[ndarray] = None, efield: Optional[ndarray] = None) → 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_typeebd() → 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()
get_dim_aparam() → int[source]

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

get_dim_fparam() → int[source]

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

get_has_efield() → bool[source]

Check if the model has efield.

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_numb_dos() → int[source]

Get the number of DOS.

get_rcut() → float[source]

Get the cut-off radius of this model.

get_sel_type() → Optional[ndarray][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_type_map() → List[str][source]

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

make_natoms_vec(atom_types: ndarray) → 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

property model_type: DeepEvalWrapper

Get type of model.

:type:str

property model_version: str

Get version of model.

Returns
str

version of model

static reverse_map(vec: ndarray, imap: List[int]) → 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.

property sess: Session

Get TF session.

static sort_input(coord: ndarray, atom_type: ndarray, sel_atoms: Optional[List[int]] = 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.

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

Factory function that forwards to DeepEval (for compatbility).

Parameters
*args

positional arguments

**kwargs

keyword arguments

Returns
DeepEval

potentials

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

Bases: 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}

Attributes
model_type

Get type of model.

model_version

Get version of model.

sess

Get TF session.

Methods

build_fv_graph()

Build the computational graph for the force and virial inference.

build_neighbor_list(coords, cell, atype, ...)

Make the mesh with neighbor list for a single frame.

eval(coord, box, atype[, eval_fv])

Evaluate the modification.

eval_full(coords, cells, atom_types[, ...])

Evaluate the model with interface similar to the energy model.

eval_typeebd()

Evaluate output of type embedding network by using this model.

get_dim_aparam()

Unsupported in this model.

get_dim_fparam()

Unsupported in this model.

get_ntypes()

Get the number of atom types of this model.

get_rcut()

Get the cut-off radius of this model.

get_sel_type()

Get the selected atom types of this model.

get_type_map()

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

make_natoms_vec(atom_types[, mixed_type])

Make the natom vector used by deepmd-kit.

modify_data(data, data_sys)

Modify data.

reverse_map(vec, imap)

Reverse mapping of a vector according to the index map.

sort_input(coord, atom_type[, sel_atoms, ...])

Sort atoms in the system according their types.
build_fv_graph() → Tensor[source]

Build the computational graph for the force and virial inference.

eval(coord: ndarray, box: ndarray, atype: ndarray, eval_fv: bool = True) → Tuple[ndarray, ndarray, 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

modify_data(data: dict, data_sys: 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.

Subpackages

Submodules

deepmd.tf.calculator module

class deepmd.tf.calculator.DP(model: Union[str, Path], label: str = 'DP', type_dict: Optional[Dict[str, int]] = None, neighbor_list=None, **kwargs)[source]

Bases: Calculator

Implementation of ASE deepmd calculator.

Implemented propertie are energy, forces and stress

Parameters
modelUnion[str, Path]

path to the model

labelstr, optional

calculator label, by default “DP”

type_dictDict[str, int], optional

mapping of element types and their numbers, best left None and the calculator will infer this information from model, by default None

neighbor_listase.neighborlist.NeighborList, optional

The neighbor list object. If None, then build the native neighbor list.

Examples

Compute potential energy

>>> from ase import Atoms
>>> from deepmd.tf.calculator import DP
>>> water = Atoms('H2O',
>>>             positions=[(0.7601, 1.9270, 1),
>>>                        (1.9575, 1, 1),
>>>                        (1., 1., 1.)],
>>>             cell=[100, 100, 100],
>>>             calculator=DP(model="frozen_model.pb"))
>>> print(water.get_potential_energy())
>>> print(water.get_forces())

Run BFGS structure optimization

>>> from ase.optimize import BFGS
>>> dyn = BFGS(water)
>>> dyn.run(fmax=1e-6)
>>> print(water.get_positions())
Attributes
directory
label

Methods

band_structure()

Create band-structure object for plotting.

calculate([atoms, properties, system_changes])

Run calculation with deepmd model.

calculate_numerical_forces(atoms[, d])

Calculate numerical forces using finite difference.

calculate_numerical_stress(atoms[, d, voigt])

Calculate numerical stress using finite difference.

calculate_properties(atoms, properties)

This method is experimental; currently for internal use.

check_state(atoms[, tol])

Check for any system changes since last calculation.

get_magnetic_moments([atoms])

Calculate magnetic moments projected onto atoms.

get_property(name[, atoms, allow_calculation])

Get the named property.

get_stresses([atoms])

the calculator should return intensive stresses, i.e., such that stresses.sum(axis=0) == stress

read(label)

Read atoms, parameters and calculated properties from output file.

reset()

Clear all information from old calculation.

set(**kwargs)

Set parameters like set(key1=value1, key2=value2, ...).

set_label(label)

Set label and convert label to directory and prefix.

calculation_required

export_properties

get_atoms

get_charges

get_default_parameters

get_dipole_moment

get_forces

get_magnetic_moment

get_potential_energies

get_potential_energy

get_stress

read_atoms

todict
calculate(atoms: Optional[Atoms] = None, properties: List[str] = ['energy', 'forces', 'virial'], system_changes: List[str] = ['positions', 'numbers', 'cell', 'pbc', 'initial_charges', 'initial_magmoms'])[source]

Run calculation with deepmd model.

Parameters
atomsOptional[Atoms], optional

atoms object to run the calculation on, by default None

propertiesList[str], optional

unused, only for function signature compatibility, by default [“energy”, “forces”, “stress”]

system_changesList[str], optional

unused, only for function signature compatibility, by default all_changes

implemented_properties: ClassVar[List[str]] = ['energy', 'free_energy', 'forces', 'virial', 'stress']

Properties calculator can handle (energy, forces, …)

name = 'DP'

deepmd.tf.common module

Collection of functions and classes used throughout the whole package.

deepmd.tf.common.add_data_requirement(key: str, ndof: int, atomic: bool = False, must: bool = False, high_prec: bool = False, type_sel: Optional[bool] = None, repeat: int = 1, default: float = 0.0, dtype: Optional[dtype] = None, output_natoms_for_type_sel: bool = False)[source]

Specify data requirements for training.

Parameters
keystr

type of data stored in corresponding *.npy file e.g. forces or energy

ndofint

number of the degrees of freedom, this is tied to atomic parameter e.g. forces have atomic=True and ndof=3

atomicbool, optional

specifies whwther the ndof keyworrd applies to per atom quantity or not, by default False

mustbool, optional

specifi if the *.npy data file must exist, by default False

high_precbool, optional

if true load data to np.float64 else np.float32, by default False

type_selbool, optional

select only certain type of atoms, by default None

repeatint, optional

if specify repaeat data repeat times, by default 1

defaultfloat, optional, default=0.

default value of data

dtypenp.dtype, optional

the dtype of data, overwrites high_prec if provided

output_natoms_for_type_selbool, optional

if True and type_sel is True, the atomic dimension will be natoms instead of nsel

deepmd.tf.common.cast_precision(func: Callable) → Callable[source]

A decorator that casts and casts back the input and output tensor of a method.

The decorator should be used in a classmethod.

The decorator will do the following thing: (1) It casts input Tensors from GLOBAL_TF_FLOAT_PRECISION to precision defined by property precision. (2) It casts output Tensors from precision to GLOBAL_TF_FLOAT_PRECISION. (3) It checks inputs and outputs and only casts when input or output is a Tensor and its dtype matches GLOBAL_TF_FLOAT_PRECISION and precision, respectively. If it does not match (e.g. it is an integer), the decorator will do nothing on it.

Returns
Callable

a decorator that casts and casts back the input and output tensor of a method

Examples

>>> class A:
...     @property
...     def precision(self):
...         return tf.float32
...
...     @cast_precision
...     def f(x: tf.Tensor, y: tf.Tensor) -> tf.Tensor:
...         return x**2 + y
deepmd.tf.common.clear_session()[source]

Reset all state generated by DeePMD-kit.

deepmd.tf.common.expand_sys_str(root_dir: Union[str, Path]) → List[str][source]

Recursively iterate over directories taking those that contain type.raw file.

Parameters
root_dirUnion[str, Path]

starting directory

Returns
List[str]

list of string pointing to system directories

deepmd.tf.common.gelu(x: Tensor) → Tensor[source]

Gaussian Error Linear Unit.

This is a smoother version of the RELU, implemented by custom operator.

Parameters
xtf.Tensor

float Tensor to perform activation

Returns
tf.Tensor

x with the GELU activation applied

References

Original paper https://arxiv.org/abs/1606.08415

deepmd.tf.common.gelu_tf(x: Tensor) → Tensor[source]

Gaussian Error Linear Unit.

This is a smoother version of the RELU, implemented by TF.

Parameters
xtf.Tensor

float Tensor to perform activation

Returns
tf.Tensor

x with the GELU activation applied

References

Original paper https://arxiv.org/abs/1606.08415

deepmd.tf.common.get_activation_func(activation_fn: Optional[_ACTIVATION]) → Optional[Callable[[Tensor], Tensor]][source]

Get activation function callable based on string name.

Parameters
activation_fn_ACTIVATION

one of the defined activation functions

Returns
Callable[[tf.Tensor], tf.Tensor]

correspondingg TF callable

Raises
RuntimeError

if unknown activation function is specified

deepmd.tf.common.get_np_precision(precision: _PRECISION) → dtype[source]

Get numpy precision constant from string.

Parameters
precision_PRECISION

string name of numpy constant or default

Returns
np.dtype

numpy presicion constant

Raises
RuntimeError

if string is invalid

deepmd.tf.common.get_precision(precision: _PRECISION) → Any[source]

Convert str to TF DType constant.

Parameters
precision_PRECISION

one of the allowed precisions

Returns
tf.python.framework.dtypes.DType

appropriate TF constant

Raises
RuntimeError

if supplied precision string does not have acorresponding TF constant

deepmd.tf.common.j_loader(filename: Union[str, Path]) → Dict[str, Any][source]

Load yaml or json settings file.

Parameters
filenameUnion[str, Path]

path to file

Returns
Dict[str, Any]

loaded dictionary

Raises
TypeError

if the supplied file is of unsupported type

deepmd.tf.common.j_must_have(jdata: Dict[str, _DICT_VAL], key: str, deprecated_key: List[str] = []) → _DICT_VAL[source]

Assert that supplied dictionary conaines specified key.

Returns
_DICT_VAL

value that was store unde supplied key

Raises
RuntimeError

if the key is not present

deepmd.tf.common.make_default_mesh(pbc: bool, mixed_type: bool) → ndarray[source]

Make mesh.

Only the size of mesh matters, not the values: * 6 for PBC, no mixed types * 0 for no PBC, no mixed types * 7 for PBC, mixed types * 1 for no PBC, mixed types

Parameters
pbcbool

if True, the mesh will be made for periodic boundary conditions

mixed_typebool

if True, the mesh will be made for mixed types

Returns
np.ndarray

mesh

deepmd.tf.common.safe_cast_tensor(input: Tensor, from_precision: DType, to_precision: DType) → Tensor[source]

Convert a Tensor from a precision to another precision.

If input is not a Tensor or without the specific precision, the method will not cast it.

Parameters
inputtf.Tensor

input tensor

from_precisiontf.DType

Tensor data type that is casted from

to_precisiontf.DType

Tensor data type that casts to

Returns
tf.Tensor

casted Tensor

deepmd.tf.common.select_idx_map(atom_types: ndarray, select_types: ndarray) → ndarray[source]

Build map of indices for element supplied element types from all atoms list.

Parameters
atom_typesnp.ndarray

array specifing type for each atoms as integer

select_typesnp.ndarray

types of atoms you want to find indices for

Returns
np.ndarray

indices of types of atoms defined by select_types in atom_types array

Warning

select_types array will be sorted before finding indices in atom_types

deepmd.tf.env module

Module that sets tensorflow working environment and exports inportant constants.

deepmd.tf.env.GLOBAL_ENER_FLOAT_PRECISION

alias of float64

deepmd.tf.env.GLOBAL_NP_FLOAT_PRECISION

alias of float64

deepmd.tf.env.global_cvt_2_ener_float(xx: Tensor) → Tensor[source]

Cast tensor to globally set energy precision.

Parameters
xxtf.Tensor

input tensor

Returns
tf.Tensor

output tensor cast to GLOBAL_ENER_FLOAT_PRECISION

deepmd.tf.env.global_cvt_2_tf_float(xx: Tensor) → Tensor[source]

Cast tensor to globally set TF precision.

Parameters
xxtf.Tensor

input tensor

Returns
tf.Tensor

output tensor cast to GLOBAL_TF_FLOAT_PRECISION

deepmd.tf.env.reset_default_tf_session_config(cpu_only: bool)[source]

Limit tensorflow session to CPU or not.

Parameters
cpu_onlybool

If enabled, no GPU device is visible to the TensorFlow Session.

deepmd.tf.lmp module

Register entry points for lammps-wheel.

deepmd.tf.lmp.get_env(paths: List[Optional[str]]) → str[source]

Get the environment variable from given paths.

deepmd.tf.lmp.get_library_path(module: str, filename: str) → List[str][source]

Get library path from a module.

Parameters
modulestr

The module name.

filenamestr

The library filename pattern.

Returns
list[str]

The library path.

deepmd.tf.lmp.get_op_dir() → str[source]

Get the directory of the deepmd-kit OP library.