deepmd.tf.common

Collection of functions and classes used throughout the whole package.

Module Contents

Functions

add_data_requirement(key, ndof[, atomic, must, ...])	Specify data requirements for training.
expand_sys_str(→ List[str])	Recursively iterate over directories taking those that contain type.raw file.
get_np_precision(→ numpy.dtype)	Get numpy precision constant from string.
j_loader(→ Dict[str, Any])	Load yaml or json settings file.
j_must_have(→ _DICT_VAL)	Assert that supplied dictionary conaines specified key.
make_default_mesh(→ numpy.ndarray)	Make mesh.
select_idx_map(→ numpy.ndarray)	Build map of indices for element supplied element types from all atoms list.
gelu(→ deepmd.tf.env.tf.Tensor)	Gaussian Error Linear Unit.
gelu_tf(→ deepmd.tf.env.tf.Tensor)	Gaussian Error Linear Unit.
get_activation_func(...)	Get activation function callable based on string name.
get_precision(→ Any)	Convert str to TF DType constant.
safe_cast_tensor(→ deepmd.tf.env.tf.Tensor)	Convert a Tensor from a precision to another precision.
cast_precision(→ Callable)	A decorator that casts and casts back the input
clear_session()	Reset all state generated by DeePMD-kit.

Attributes

data_requirement
PRECISION_DICT
ACTIVATION_FN_DICT

deepmd.tf.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: numpy.dtype | None = None, output_natoms_for_type_sel: bool = False)

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

deepmd.tf.common.expand_sys_str(root_dir: str | pathlib.Path) → List[str]

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.get_np_precision(precision: _PRECISION) → numpy.dtype

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.j_loader(filename: str | pathlib.Path) → Dict[str, Any]

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

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) → numpy.ndarray

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.select_idx_map(atom_types: numpy.ndarray, select_types: numpy.ndarray) → numpy.ndarray

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.common.PRECISION_DICT

deepmd.tf.common.gelu(x: deepmd.tf.env.tf.Tensor) → deepmd.tf.env.tf.Tensor

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: deepmd.tf.env.tf.Tensor) → deepmd.tf.env.tf.Tensor

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

deepmd.tf.common.get_activation_func(activation_fn: deepmd.common._ACTIVATION | None) → Callable[[deepmd.tf.env.tf.Tensor], deepmd.tf.env.tf.Tensor]

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_precision(precision: deepmd.common._PRECISION) → Any

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.safe_cast_tensor(input: deepmd.tf.env.tf.Tensor, from_precision: deepmd.tf.env.tf.DType, to_precision: deepmd.tf.env.tf.DType) → deepmd.tf.env.tf.Tensor

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.cast_precision(func: Callable) → Callable

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()

Reset all state generated by DeePMD-kit.