deepmd.tf.common

Collection of functions and classes used throughout the whole package.

Attributes

PRECISION_DICT
ACTIVATION_FN_DICT

Functions

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.
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(→ None)	Reset all state generated by DeePMD-kit.

Module Contents

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 precision 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.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 specifying 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() → None

Reset all state generated by DeePMD-kit.