deepmd package

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

class deepmd.DeepEval(model_file: Path, load_prefix: str = 'load', default_tf_graph: bool = False, auto_batch_size: bool | int | AutoBatchSize = False)[source]

Bases: object

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.

Attributes:
model_type

Get type of model.

model_version

Get version of model.

sess

Get TF session.

Methods

eval_typeebd()

Evaluate output of type embedding network by using this model.

make_natoms_vec(atom_types[, mixed_type])

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.
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.utils.type_embed.TypeEmbedNet

The type embedding network.

Examples

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

>>> from deepmd.infer import DeepPotential
>>> dp = DeepPotential('graph.pb')
>>> dp.eval_typeebd()
load_prefix: str
make_natoms_vec(atom_types: ndarray, mixed_type: bool = False) → 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

property model_type: str

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

deepmd.DeepPotential(model_file: str | Path, load_prefix: str = 'load', default_tf_graph: bool = False) → DeepDipole | DeepGlobalPolar | DeepPolar | DeepPot | DeepWFC[source]

Factory function that will inialize appropriate potential read from model_file.

Parameters:
model_filestr

The name of the frozen model file.

load_prefixstr

The prefix in the load computational graph

default_tf_graphbool

If uses the default tf graph, otherwise build a new tf graph for evaluation

Returns:
Union[DeepDipole, DeepGlobalPolar, DeepPolar, DeepPot, DeepWFC]

one of the available potentials

Raises:
RuntimeError

if model file does not correspond to any implementd potential

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

Bases: DeepDipole

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.

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)

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

load_prefix: str
modify_data(data: dict) → 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

Subpackages

Submodules

deepmd.calculator module

ASE calculator interface module.

class deepmd.calculator.DP(model: str | Path, label: str = 'DP', type_dict: Dict[str, int] | None = 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

Examples

Compute potential energy

>>> from ase import Atoms
>>> from deepmd.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: Atoms | None = 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: List[str] = ['energy', 'free_energy', 'forces', 'virial', 'stress']

Properties calculator can handle (energy, forces, …)

name = 'DP'

deepmd.common module

Collection of functions and classes used throughout the whole package.

deepmd.common.add_data_requirement(key: str, ndof: int, atomic: bool = False, must: bool = False, high_prec: bool = False, type_sel: bool | None = None, repeat: int = 1, default: float = 0.0, dtype: dtype | None = None)[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

deepmd.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.common.clear_session()[source]

Reset all state generated by DeePMD-kit.

deepmd.common.expand_sys_str(root_dir: 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.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.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.common.get_activation_func(activation_fn: _ACTIVATION | None) → Callable[[Tensor], Tensor] | None[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.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.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.common.j_loader(filename: 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.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.common.make_default_mesh(test_box: ndarray, cell_size: float = 3.0) → ndarray[source]

Get number of cells of size=`cell_size` fit into average box.

Parameters:
test_boxnp.ndarray

numpy array with cells of shape Nx9

cell_sizefloat, optional

length of one cell, by default 3.0

Returns:
np.ndarray

mesh for supplied boxes, how many cells fit in each direction

deepmd.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.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.env module

Module that sets tensorflow working environment and exports inportant constants.

deepmd.env.GLOBAL_ENER_FLOAT_PRECISION

alias of float64

deepmd.env.GLOBAL_NP_FLOAT_PRECISION

alias of float64

deepmd.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.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.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.lmp module

Register entry points for lammps-wheel.

deepmd.lmp.get_env(paths: List[str | None]) → str[source]

Get the environment variable from given paths.

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

Get library path from a module.

Parameters:
modulestr

The module name.

Returns:
list[str]

The library path.

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

Get the directory of the deepmd-kit OP library.