BladeDISC DeePMD code

_skbuild/linux-x86_64-3.6/cmake-install/deepmd/calculator.py

+"""ASE calculator interface module."""
+
+from pathlib import Path
+from typing import TYPE_CHECKING, Dict, List, Optional, Union
+
+from ase.calculators.calculator import (
+    Calculator, all_changes, PropertyNotImplementedError
+)
+
+from deepmd import DeepPotential
+
+if TYPE_CHECKING:
+    from ase import Atoms
+
+__all__ = ["DP"]
+
+
+class DP(Calculator):
+    """Implementation of ASE deepmd calculator.
+
+    Implemented propertie are `energy`, `forces` and `stress`
+
+    Parameters
+    ----------
+    model : Union[str, Path]
+        path to the model
+    label : str, optional
+        calculator label, by default "DP"
+    type_dict : Dict[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())
+    """
+
+    name = "DP"
+    implemented_properties = ["energy", "free_energy", "forces", "virial", "stress"]
+
+    def __init__(
+        self,
+        model: Union[str, "Path"],
+        label: str = "DP",
+        type_dict: Dict[str, int] = None,
+        **kwargs
+    ) -> None:
+        Calculator.__init__(self, label=label, **kwargs)
+        self.dp = DeepPotential(str(Path(model).resolve()))
+        if type_dict:
+            self.type_dict = type_dict
+        else:
+            self.type_dict = dict(
+                zip(self.dp.get_type_map(), range(self.dp.get_ntypes()))
+            )
+
+    def calculate(
+        self,
+        atoms: Optional["Atoms"] = None,
+        properties: List[str] = ["energy", "forces", "virial"],
+        system_changes: List[str] = all_changes,
+    ):
+        """Run calculation with deepmd model.
+
+        Parameters
+        ----------
+        atoms : Optional[Atoms], optional
+            atoms object to run the calculation on, by default None
+        properties : List[str], optional
+            unused, only for function signature compatibility,
+            by default ["energy", "forces", "stress"]
+        system_changes : List[str], optional
+            unused, only for function signature compatibility, by default all_changes
+        """
+        if atoms is not None:
+            self.atoms = atoms.copy()
+
+        coord = self.atoms.get_positions().reshape([1, -1])
+        if sum(self.atoms.get_pbc()) > 0:
+            cell = self.atoms.get_cell().reshape([1, -1])
+        else:
+            cell = None
+        symbols = self.atoms.get_chemical_symbols()
+        atype = [self.type_dict[k] for k in symbols]
+        e, f, v = self.dp.eval(coords=coord, cells=cell, atom_types=atype)
+        self.results['energy'] = e[0][0]
+        # see https://gitlab.com/ase/ase/-/merge_requests/2485
+        self.results['free_energy'] = e[0][0]
+        self.results['forces'] = f[0]
+        self.results['virial'] = v[0].reshape(3, 3)
+
+        # convert virial into stress for lattice relaxation
+        if "stress" in properties:
+            if sum(atoms.get_pbc()) > 0:
+                # the usual convention (tensile stress is positive)
+                # stress = -virial / volume
+                stress = -0.5 * (v[0].copy() + v[0].copy().T) / atoms.get_volume()
+                # Voigt notation
+                self.results['stress'] = stress.flat[[0, 4, 8, 5, 2, 1]]
+            else:
+                raise PropertyNotImplementedError

_skbuild/linux-x86_64-3.6/cmake-install/deepmd/cluster/__init__.py

+"""Module that reads node resources, auto detects if running local or on SLURM."""
+
+from .local import get_resource as get_local_res
+from .slurm import get_resource as get_slurm_res
+import os
+from typing import List, Tuple, Optional
+
+__all__ = ["get_resource"]
+
+
+def get_resource() -> Tuple[str, List[str], Optional[List[int]]]:
+    """Get local or slurm resources: nodename, nodelist, and gpus.
+
+    Returns
+    -------
+    Tuple[str, List[str], Optional[List[int]]]
+        nodename, nodelist, and gpus
+    """
+    if "SLURM_JOB_NODELIST" in os.environ:
+        return get_slurm_res()
+    else:
+        return get_local_res()

_skbuild/linux-x86_64-3.6/cmake-install/deepmd/cluster/local.py

+"""Get local GPU resources."""
+
+import os
+import socket
+import subprocess as sp
+import sys
+
+from deepmd.env import tf
+from typing import List, Tuple, Optional
+
+
+__all__ = ["get_gpus", "get_resource"]
+
+
+def get_gpus():
+    """Get available IDs of GPU cards at local.
+    These IDs are valid when used as the TensorFlow device ID.
+
+    Returns
+    -------
+    Optional[List[int]]
+        List of available GPU IDs. Otherwise, None.
+    """
+    test_cmd = 'from tensorflow.python.client import device_lib; ' \
+               'devices = device_lib.list_local_devices(); ' \
+               'gpus = [d.name for d in devices if d.device_type == "GPU"]; ' \
+               'print(len(gpus))'
+    with sp.Popen([sys.executable, "-c", test_cmd], stderr=sp.PIPE, stdout=sp.PIPE) as p:
+        stdout, stderr = p.communicate()
+        if p.returncode != 0:
+            decoded = stderr.decode('UTF-8')
+            raise RuntimeError('Failed to detect availbe GPUs due to:\n%s' % decoded)
+        decoded = stdout.decode('UTF-8').strip()
+        num_gpus = int(decoded)
+        return list(range(num_gpus)) if num_gpus > 0 else None
+
+
+def get_resource() -> Tuple[str, List[str], Optional[List[int]]]:
+    """Get local resources: nodename, nodelist, and gpus.
+
+    Returns
+    -------
+    Tuple[str, List[str], Optional[List[int]]]
+        nodename, nodelist, and gpus
+    """
+    nodename = socket.gethostname()
+    nodelist = [nodename]
+    gpus = get_gpus()
+    return nodename, nodelist, gpus

_skbuild/linux-x86_64-3.6/cmake-install/deepmd/cluster/slurm.py

+"""MOdule to get resources on SLURM cluster.
+
+References
+----------
+https://github.com/deepsense-ai/tensorflow_on_slurm ####
+"""
+
+import hostlist
+import os
+
+from deepmd.cluster import local
+from typing import List, Tuple, Optional
+
+__all__ = ["get_resource"]
+
+
+def get_resource() -> Tuple[str, List[str], Optional[List[int]]]:
+    """Get SLURM resources: nodename, nodelist, and gpus.
+
+    Returns
+    -------
+    Tuple[str, List[str], Optional[List[int]]]
+        nodename, nodelist, and gpus
+
+    Raises
+    ------
+    RuntimeError
+        if number of nodes could not be retrieved
+    ValueError
+        list of nodes is not of the same length sa number of nodes
+    ValueError
+        if current nodename is not found in node list
+    """
+    nodelist = hostlist.expand_hostlist(os.environ["SLURM_JOB_NODELIST"])
+    nodename = os.environ["SLURMD_NODENAME"]
+    num_nodes_env = os.getenv("SLURM_JOB_NUM_NODES")
+    if num_nodes_env:
+        num_nodes = int(num_nodes_env)
+    else:
+        raise RuntimeError("Could not get SLURM number of nodes")
+
+    if len(nodelist) != num_nodes:
+        raise ValueError(
+            f"Number of slurm nodes {len(nodelist)} not equal to {num_nodes}"
+        )
+    if nodename not in nodelist:
+        raise ValueError(
+            f"Nodename({nodename}) not in nodelist({nodelist}). This should not happen!"
+        )
+    gpus = local.get_gpus()
+    return nodename, nodelist, gpus

_skbuild/linux-x86_64-3.6/cmake-install/deepmd/common.py

+"""Collection of functions and classes used throughout the whole package."""
+
+import json
+import warnings
+import tensorflow
+from functools import wraps
+from pathlib import Path
+from typing import (
+    TYPE_CHECKING,
+    Any,
+    Callable,
+    Dict,
+    List,
+    Optional,
+    Tuple,
+    TypeVar,
+    Union,
+)
+
+import numpy as np
+import yaml
+
+from deepmd.env import op_module, tf
+from tensorflow.python.framework import tensor_util
+from deepmd.env import GLOBAL_TF_FLOAT_PRECISION, GLOBAL_NP_FLOAT_PRECISION
+from deepmd.utils.sess import run_sess
+from deepmd.utils.errors import GraphWithoutTensorError
+from deepmd.utils.path import DPPath
+
+if TYPE_CHECKING:
+    _DICT_VAL = TypeVar("_DICT_VAL")
+    _OBJ = TypeVar("_OBJ")
+    try:
+        from typing import Literal  # python >3.6
+    except ImportError:
+        from typing_extensions import Literal  # type: ignore
+    _ACTIVATION = Literal["relu", "relu6", "softplus", "sigmoid", "tanh", "gelu", "gelu_tf"]
+    _PRECISION = Literal["default", "float16", "float32", "float64"]
+
+# define constants
+PRECISION_DICT = {
+    "default": GLOBAL_TF_FLOAT_PRECISION,
+    "float16": tf.float16,
+    "float32": tf.float32,
+    "float64": tf.float64,
+}
+
+
+def gelu(x: tf.Tensor) -> tf.Tensor:
+    """Gaussian Error Linear Unit.
+
+    This is a smoother version of the RELU, implemented by custom operator.
+
+    Parameters
+    ----------
+    x : tf.Tensor
+        float Tensor to perform activation
+
+    Returns
+    -------
+    tf.Tensor
+        `x` with the GELU activation applied
+
+    References
+    ----------
+    Original paper
+    https://arxiv.org/abs/1606.08415
+    """
+    return op_module.gelu_custom(x)
+
+
+def gelu_tf(x: tf.Tensor) -> tf.Tensor:
+    """Gaussian Error Linear Unit.
+
+    This is a smoother version of the RELU, implemented by TF.
+
+    Parameters
+    ----------
+    x : tf.Tensor
+        float Tensor to perform activation
+
+    Returns
+    -------
+    tf.Tensor
+        `x` with the GELU activation applied
+
+    References
+    ----------
+    Original paper
+    https://arxiv.org/abs/1606.08415
+    """
+    def gelu_wrapper(x):
+        try:
+            return tensorflow.nn.gelu(x, approximate=True)
+        except AttributeError:
+            warnings.warn("TensorFlow does not provide an implementation of gelu, please upgrade your TensorFlow version. Fallback to the custom gelu operator.")
+            return op_module.gelu_custom(x)
+    return (lambda x: gelu_wrapper(x))(x)
+
+# TODO this is not a good way to do things. This is some global variable to which
+# TODO anyone can write and there is no good way to keep track of the changes
+data_requirement = {}
+
+ACTIVATION_FN_DICT = {
+    "relu": tf.nn.relu,
+    "relu6": tf.nn.relu6,
+    "softplus": tf.nn.softplus,
+    "sigmoid": tf.sigmoid,
+    "tanh": tf.nn.tanh,
+    "gelu": gelu,
+    "gelu_tf": gelu_tf,
+}
+
+
+def add_data_requirement(
+    key: str,
+    ndof: int,
+    atomic: bool = False,
+    must: bool = False,
+    high_prec: bool = False,
+    type_sel: bool = None,
+    repeat: int = 1,
+    default: float = 0.,
+):
+    """Specify data requirements for training.
+
+    Parameters
+    ----------
+    key : str
+        type of data stored in corresponding `*.npy` file e.g. `forces` or `energy`
+    ndof : int
+        number of the degrees of freedom, this is tied to `atomic` parameter e.g. forces
+        have `atomic=True` and `ndof=3`
+    atomic : bool, optional
+        specifies whwther the `ndof` keyworrd applies to per atom quantity or not,
+        by default False
+    must : bool, optional
+        specifi if the `*.npy` data file must exist, by default False
+    high_prec : bool, optional
+        if tru load data to `np.float64` else `np.float32`, by default False
+    type_sel : bool, optional
+        select only certain type of atoms, by default None
+    repeat : int, optional
+        if specify repaeat data `repeat` times, by default 1
+    default : float, optional, default=0.
+        default value of data
+    """
+    data_requirement[key] = {
+        "ndof": ndof,
+        "atomic": atomic,
+        "must": must,
+        "high_prec": high_prec,
+        "type_sel": type_sel,
+        "repeat": repeat,
+        "default": default,
+    }
+
+
+def select_idx_map(
+    atom_types: np.ndarray, select_types: np.ndarray
+) -> np.ndarray:
+    """Build map of indices for element supplied element types from all atoms list.
+
+    Parameters
+    ----------
+    atom_types : np.ndarray
+        array specifing type for each atoms as integer
+    select_types : np.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
+
+    Warnings
+    --------
+    `select_types` array will be sorted before finding indices in `atom_types`
+    """
+    sort_select_types = np.sort(select_types)
+    idx_map = np.array([], dtype=int)
+    for ii in sort_select_types:
+        idx_map = np.append(idx_map, np.where(atom_types == ii))
+    return idx_map
+
+
+# TODO not really sure if the docstring is right the purpose of this is a bit unclear
+def make_default_mesh(
+    test_box: np.ndarray, cell_size: float = 3.0
+) -> np.ndarray:
+    """Get number of cells of size=`cell_size` fit into average box.
+
+    Parameters
+    ----------
+    test_box : np.ndarray
+        numpy array with cells of shape Nx9
+    cell_size : float, optional
+        length of one cell, by default 3.0
+
+    Returns
+    -------
+    np.ndarray
+        mesh for supplied boxes, how many cells fit in each direction
+    """
+    cell_lengths = np.linalg.norm(test_box.reshape([-1, 3, 3]), axis=2)
+    avg_cell_lengths = np.average(cell_lengths, axis=0)
+    ncell = (avg_cell_lengths / cell_size).astype(np.int32)
+    ncell[ncell < 2] = 2
+    default_mesh = np.zeros(6, dtype=np.int32)
+    default_mesh[3:6] = ncell
+    return default_mesh
+
+
+# TODO not an ideal approach, every class uses this to parse arguments on its own, json
+# TODO should be parsed once and the parsed result passed to all objects that need it
+class ClassArg:
+    """Class that take care of input json/yaml parsing.
+
+    The rules for parsing are defined by the `add` method, than `parse` is called to
+    process the supplied dict
+
+    Attributes
+    ----------
+    arg_dict: Dict[str, Any]
+        dictionary containing parsing rules
+    alias_map: Dict[str, Any]
+        dictionary with keyword aliases
+    """
+
+    def __init__(self) -> None:
+        self.arg_dict = {}
+        self.alias_map = {}
+
+    def add(
+        self,
+        key: str,
+        types_: Union[type, List[type]],
+        alias: Optional[Union[str, List[str]]] = None,
+        default: Any = None,
+        must: bool = False,
+    ) -> "ClassArg":
+        """Add key to be parsed.
+
+        Parameters
+        ----------
+        key : str
+            key name
+        types_ : Union[type, List[type]]
+            list of allowed key types
+        alias : Optional[Union[str, List[str]]], optional
+            alias for the key, by default None
+        default : Any, optional
+            default value for the key, by default None
+        must : bool, optional
+            if the key is mandatory, by default False
+
+        Returns
+        -------
+        ClassArg
+            instance with added key
+        """
+        if not isinstance(types_, list):
+            types = [types_]
+        else:
+            types = types_
+        if alias is not None:
+            if not isinstance(alias, list):
+                alias_ = [alias]
+            else:
+                alias_ = alias
+        else:
+            alias_ = []
+
+        self.arg_dict[key] = {
+            "types": types,
+            "alias": alias_,
+            "value": default,
+            "must": must,
+        }
+        for ii in alias_:
+            self.alias_map[ii] = key
+
+        return self
+
+    def _add_single(self, key: str, data: Any):
+        vtype = type(data)
+        if data is None:
+            return data
+        if not (vtype in self.arg_dict[key]["types"]):
+            for tp in self.arg_dict[key]["types"]:
+                try:
+                    vv = tp(data)
+                except TypeError:
+                    pass
+                else:
+                    break
+            else:
+                raise TypeError(
+                    f"cannot convert provided key {key} to type(s) "
+                    f'{self.arg_dict[key]["types"]} '
+                )
+        else:
+            vv = data
+        self.arg_dict[key]["value"] = vv
+
+    def _check_must(self):
+        for kk in self.arg_dict:
+            if self.arg_dict[kk]["must"] and self.arg_dict[kk]["value"] is None:
+                raise RuntimeError(f"key {kk} must be provided")
+
+    def parse(self, jdata: Dict[str, Any]) -> Dict[str, Any]:
+        """Parse input dictionary, use the rules defined by add method.
+
+        Parameters
+        ----------
+        jdata : Dict[str, Any]
+            loaded json/yaml data
+
+        Returns
+        -------
+        Dict[str, Any]
+            parsed dictionary
+        """
+        for kk in jdata.keys():
+            if kk in self.arg_dict:
+                key = kk
+                self._add_single(key, jdata[kk])
+            else:
+                if kk in self.alias_map:
+                    key = self.alias_map[kk]
+                    self._add_single(key, jdata[kk])
+        self._check_must()
+        return self.get_dict()
+
+    def get_dict(self) -> Dict[str, Any]:
+        """Get dictionary built from rules defined by add method.
+
+        Returns
+        -------
+        Dict[str, Any]
+            settings dictionary with default values
+        """
+        ret = {}
+        for kk in self.arg_dict.keys():
+            ret[kk] = self.arg_dict[kk]["value"]
+        return ret
+
+
+# TODO maybe rename this to j_deprecated and only warn about deprecated keys,
+# TODO if the deprecated_key argument is left empty function puppose is only custom
+# TODO error since dict[key] already raises KeyError when the key is missing
+def 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
+    """
+    if key not in jdata.keys():
+        for ii in deprecated_key:
+            if ii in jdata.keys():
+                warnings.warn(f"the key {ii} is deprecated, please use {key} instead")
+                return jdata[ii]
+        else:
+            raise RuntimeError(f"json database must provide key {key}")
+    else:
+        return jdata[key]
+
+
+def j_loader(filename: Union[str, Path]) -> Dict[str, Any]:
+    """Load yaml or json settings file.
+
+    Parameters
+    ----------
+    filename : Union[str, Path]
+        path to file
+
+    Returns
+    -------
+    Dict[str, Any]
+        loaded dictionary
+
+    Raises
+    ------
+    TypeError
+        if the supplied file is of unsupported type
+    """
+    filepath = Path(filename)
+    if filepath.suffix.endswith("json"):
+        with filepath.open() as fp:
+            return json.load(fp)
+    elif filepath.suffix.endswith(("yml", "yaml")):
+        with filepath.open() as fp:
+            return yaml.safe_load(fp)
+    else:
+        raise TypeError("config file must be json, or yaml/yml")
+
+
+def get_activation_func(
+    activation_fn: Union["_ACTIVATION", None],
+) -> Union[Callable[[tf.Tensor], tf.Tensor], None]:
+    """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
+    """
+    if activation_fn is None or activation_fn in ['none', 'None']:
+        return None
+    if activation_fn not in ACTIVATION_FN_DICT:
+        raise RuntimeError(f"{activation_fn} is not a valid activation function")
+    return ACTIVATION_FN_DICT[activation_fn]
+
+
+def get_precision(precision: "_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
+    """
+    if precision not in PRECISION_DICT:
+        raise RuntimeError(f"{precision} is not a valid precision")
+    return PRECISION_DICT[precision]
+
+
+# TODO port completely to pathlib when all callers are ported
+def expand_sys_str(root_dir: Union[str, Path]) -> List[str]:
+    """Recursively iterate over directories taking those that contain `type.raw` file.
+
+    Parameters
+    ----------
+    root_dir : Union[str, Path]
+        starting directory
+
+    Returns
+    -------
+    List[str]
+        list of string pointing to system directories
+    """
+    root_dir = DPPath(root_dir)
+    matches = [str(d) for d in root_dir.rglob("*") if (d / "type.raw").is_file()]
+    if (root_dir / "type.raw").is_file():
+        matches.append(str(root_dir))
+    return matches
+
+
+def get_np_precision(precision: "_PRECISION") -> np.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
+    """
+    if precision == "default":
+        return GLOBAL_NP_FLOAT_PRECISION
+    elif precision == "float16":
+        return np.float16
+    elif precision == "float32":
+        return np.float32
+    elif precision == "float64":
+        return np.float64
+    else:
+        raise RuntimeError(f"{precision} is not a valid precision")
+
+
+def safe_cast_tensor(input: tf.Tensor,
+                     from_precision: tf.DType,
+                     to_precision: tf.DType) -> 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
+    ----------
+    input: tf.Tensor
+        input tensor
+    precision : tf.DType
+        Tensor data type that casts to
+
+    Returns
+    -------
+    tf.Tensor
+        casted Tensor
+    """
+    if tensor_util.is_tensor(input) and input.dtype == from_precision:
+        return tf.cast(input, to_precision)
+    return input
+
+
+def 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
+    """
+    @wraps(func)
+    def wrapper(self, *args, **kwargs):
+        # only convert tensors
+        returned_tensor = func(
+            self,
+            *[safe_cast_tensor(vv, GLOBAL_TF_FLOAT_PRECISION, self.precision) for vv in args],
+            **{kk: safe_cast_tensor(vv, GLOBAL_TF_FLOAT_PRECISION, self.precision) for kk, vv in kwargs.items()},
+        )
+        if isinstance(returned_tensor, tuple):
+            return tuple((safe_cast_tensor(vv, self.precision, GLOBAL_TF_FLOAT_PRECISION) for vv in returned_tensor))
+        else:
+            return safe_cast_tensor(returned_tensor, self.precision, GLOBAL_TF_FLOAT_PRECISION)
+    return wrapper

_skbuild/linux-x86_64-3.6/cmake-install/deepmd/descriptor/__init__.py

+from .descriptor import Descriptor
+from .hybrid import DescrptHybrid
+from .se_a import DescrptSeA
+from .se_r import DescrptSeR
+from .se_t import DescrptSeT
+from .se_a_ebd import DescrptSeAEbd
+from .se_a_ef import DescrptSeAEf
+from .se_a_ef import DescrptSeAEfLower
+from .loc_frame import DescrptLocFrame
+from .se_atten import DescrptSeAtten

_skbuild/linux-x86_64-3.6/cmake-install/deepmd/descriptor/descriptor.py

+from abc import ABC, abstractmethod
+from typing import Any, Dict, List, Tuple
+
+import numpy as np
+from deepmd.env import tf
+from deepmd.utils import Plugin, PluginVariant
+
+
+class Descriptor(PluginVariant):
+    r"""The abstract class for descriptors. All specific descriptors should
+    be based on this class.
+
+    The descriptor :math:`\mathcal{D}` describes the environment of an atom,
+    which should be a function of coordinates and types of its neighbour atoms.
+
+    Examples
+    --------
+    >>> descript = Descriptor(type="se_e2_a", rcut=6., rcut_smth=0.5, sel=[50])
+    >>> type(descript)
+    <class 'deepmd.descriptor.se_a.DescrptSeA'>
+
+    Notes
+    -----
+    Only methods and attributes defined in this class are generally public,
+    that can be called by other classes.
+    """
+
+    __plugins = Plugin()
+
+    @staticmethod
+    def register(key: str) -> "Descriptor":
+        """Regiester a descriptor plugin.
+
+        Parameters
+        ----------
+        key : str
+            the key of a descriptor
+
+        Returns
+        -------
+        Descriptor
+            the regiestered descriptor
+
+        Examples
+        --------
+        >>> @Descriptor.register("some_descrpt")
+            class SomeDescript(Descriptor):
+                pass
+        """
+        return Descriptor.__plugins.register(key)
+
+    def __new__(cls, *args, **kwargs):
+        if cls is Descriptor:
+            try:
+                descrpt_type = kwargs['type']
+            except KeyError:
+                raise KeyError('the type of descriptor should be set by `type`')
+            if descrpt_type in Descriptor.__plugins.plugins:
+                cls = Descriptor.__plugins.plugins[descrpt_type]
+            else:
+                raise RuntimeError('Unknown descriptor type: ' + descrpt_type)
+        return super().__new__(cls)
+
+    @abstractmethod
+    def get_rcut(self) -> float:
+        """
+        Returns the cut-off radius.
+
+        Returns
+        -------
+        float
+            the cut-off radius
+
+        Notes
+        -----
+        This method must be implemented, as it's called by other classes.
+        """
+
+    @abstractmethod
+    def get_ntypes(self) -> int:
+        """
+        Returns the number of atom types.
+
+        Returns
+        -------
+        int
+            the number of atom types
+
+        Notes
+        -----
+        This method must be implemented, as it's called by other classes.
+        """
+
+    @abstractmethod
+    def get_dim_out(self) -> int:
+        """
+        Returns the output dimension of this descriptor.
+
+        Returns
+        -------
+        int
+            the output dimension of this descriptor
+
+        Notes
+        -----
+        This method must be implemented, as it's called by other classes.
+        """
+
+    def get_dim_rot_mat_1(self) -> int:
+        """
+        Returns the first dimension of the rotation matrix. The rotation is of shape
+        dim_1 x 3
+
+        Returns
+        -------
+        int
+            the first dimension of the rotation matrix
+        """
+        # TODO: I think this method should be implemented as it's called by dipole and
+        # polar fitting network. However, currently not all descriptors have this
+        # method.
+        raise NotImplementedError
+
+    def get_nlist(self) -> Tuple[tf.Tensor, tf.Tensor, List[int], List[int]]:
+        """
+        Returns neighbor information.
+
+        Returns
+        -------
+        nlist : tf.Tensor
+            Neighbor list
+        rij : tf.Tensor
+            The relative distance between the neighbor and the center atom.
+        sel_a : list[int]
+            The number of neighbors with full information
+        sel_r : list[int]
+            The number of neighbors with only radial information
+        """
+        # TODO: I think this method should be implemented as it's called by energy
+        # model. However, se_ar and hybrid doesn't have this method.
+        raise NotImplementedError
+
+    @abstractmethod
+    def compute_input_stats(self,
+                            data_coord: List[np.ndarray],
+                            data_box: List[np.ndarray],
+                            data_atype: List[np.ndarray],
+                            natoms_vec: List[np.ndarray],
+                            mesh: List[np.ndarray],
+                            input_dict: Dict[str, List[np.ndarray]]
+                            ) -> None:
+        """
+        Compute the statisitcs (avg and std) of the training data. The input will be
+        normalized by the statistics.
+
+        Parameters
+        ----------
+        data_coord : list[np.ndarray]
+            The coordinates. Can be generated by
+            :meth:`deepmd.model.model_stat.make_stat_input`
+        data_box : list[np.ndarray]
+            The box. Can be generated by
+            :meth:`deepmd.model.model_stat.make_stat_input`
+        data_atype : list[np.ndarray]
+            The atom types. Can be generated by :meth:`deepmd.model.model_stat.make_stat_input`
+        natoms_vec : list[np.ndarray]
+            The vector for the number of atoms of the system and different types of
+            atoms. Can be generated by :meth:`deepmd.model.model_stat.make_stat_input`
+        mesh : list[np.ndarray]
+            The mesh for neighbor searching. Can be generated by
+            :meth:`deepmd.model.model_stat.make_stat_input`
+        input_dict : dict[str, list[np.ndarray]]
+            Dictionary for additional input
+
+        Notes
+        -----
+        This method must be implemented, as it's called by other classes.
+        """
+
+    @abstractmethod
+    def build(self,
+              coord_: tf.Tensor,
+              atype_: tf.Tensor,
+              natoms: tf.Tensor,
+              box_: tf.Tensor,
+              mesh: tf.Tensor,
+              input_dict: Dict[str, Any],
+              reuse: bool = None,
+              suffix: str = '',
+              ) -> tf.Tensor:
+        """
+        Build the computational graph for the descriptor.
+
+        Parameters
+        ----------
+        coord_ : tf.Tensor
+            The coordinate of atoms
+        atype_ : tf.Tensor
+            The type of atoms
+        natoms : tf.Tensor
+            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
+        box : tf.Tensor
+            The box of frames
+        mesh : tf.Tensor
+            For historical reasons, only the length of the Tensor matters.
+            if size of mesh == 6, pbc is assumed.
+            if size of mesh == 0, no-pbc is assumed.
+        input_dict : dict[str, Any]
+            Dictionary for additional inputs
+        reuse : bool, optional
+            The weights in the networks should be reused when get the variable.
+        suffix : str, optional
+            Name suffix to identify this descriptor
+
+        Returns
+        -------
+        descriptor: tf.Tensor
+            The output descriptor
+
+        Notes
+        -----
+        This method must be implemented, as it's called by other classes.
+        """
+
+    def enable_compression(self,
+                           min_nbor_dist: float,
+                           model_file: str = 'frozon_model.pb',
+                           table_extrapolate: float = 5.,
+                           table_stride_1: float = 0.01,
+                           table_stride_2: float = 0.1,
+                           check_frequency: int = -1,
+                           suffix: str = "",
+                           ) -> None:
+        """
+        Reveive the statisitcs (distance, max_nbor_size and env_mat_range) of the
+        training data.
+
+        Parameters
+        ----------
+        min_nbor_dist : float
+                The nearest distance between atoms
+        model_file : str, default: 'frozon_model.pb'
+                The original frozen model, which will be compressed by the program
+        table_extrapolate : float, default: 5.
+                The scale of model extrapolation
+        table_stride_1 : float, default: 0.01
+                The uniform stride of the first table
+        table_stride_2 : float, default: 0.1
+                The uniform stride of the second table
+        check_frequency : int, default: -1
+                The overflow check frequency
+        suffix : str, optional
+                The suffix of the scope
+
+        Notes
+        -----
+        This method is called by others when the descriptor supported compression.
+        """
+        raise NotImplementedError(
+            "Descriptor %s doesn't support compression!" % type(self).__name__)
+
+    def enable_mixed_precision(self, mixed_prec: dict = None) -> None:
+        """
+        Reveive the mixed precision setting.
+
+        Parameters
+        ----------
+        mixed_prec
+                The mixed precision setting used in the embedding net
+
+        Notes
+        -----
+        This method is called by others when the descriptor supported compression.
+        """
+        raise NotImplementedError(
+            "Descriptor %s doesn't support mixed precision training!"
+            % type(self).__name__
+        )
+
+    @abstractmethod
+    def prod_force_virial(self,
+                          atom_ener: tf.Tensor,
+                          natoms: tf.Tensor
+                          ) -> Tuple[tf.Tensor, tf.Tensor, tf.Tensor]:
+        """
+        Compute force and virial.
+
+        Parameters
+        ----------
+        atom_ener : tf.Tensor
+            The atomic energy
+        natoms : tf.Tensor
+            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
+
+        Returns
+        -------
+        force : tf.Tensor
+            The force on atoms
+        virial : tf.Tensor
+            The total virial
+        atom_virial : tf.Tensor
+            The atomic virial
+        """
+
+    def get_feed_dict(self,
+                      coord_: tf.Tensor,
+                      atype_: tf.Tensor,
+                      natoms: tf.Tensor,
+                      box: tf.Tensor,
+                      mesh: tf.Tensor
+                      ) -> Dict[str, tf.Tensor]:
+        """
+        Generate the feed_dict for current descriptor
+
+        Parameters
+        ----------
+        coord_ : tf.Tensor
+            The coordinate of atoms
+        atype_ : tf.Tensor
+            The type of atoms
+        natoms : tf.Tensor
+            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
+        box : tf.Tensor
+            The box. Can be generated by deepmd.model.make_stat_input
+        mesh : tf.Tensor
+            For historical reasons, only the length of the Tensor matters.
+            if size of mesh == 6, pbc is assumed. 
+            if size of mesh == 0, no-pbc is assumed. 
+
+        Returns
+        -------
+        feed_dict : dict[str, tf.Tensor]
+            The output feed_dict of current descriptor
+        """
+        feed_dict = {
+            't_coord:0'  :coord_,
+            't_type:0'   :atype_,
+            't_natoms:0' :natoms,
+            't_box:0'    :box,
+            't_mesh:0'   :mesh
+        }
+        return feed_dict
+
+    def init_variables(self,
+                       graph: tf.Graph,
+                       graph_def: tf.GraphDef,
+                       suffix : str = "",
+    ) -> None:
+        """
+        Init the embedding net variables with the given dict
+
+        Parameters
+        ----------
+        graph : tf.Graph
+            The input frozen model graph
+        graph_def : tf.GraphDef
+            The input frozen model graph_def
+        suffix : str, optional
+            The suffix of the scope
+        
+        Notes
+        -----
+        This method is called by others when the descriptor supported initialization from the given variables.
+        """
+        raise NotImplementedError(
+            "Descriptor %s doesn't support initialization from the given variables!" % type(self).__name__)
+
+    def get_tensor_names(self, suffix : str = "") -> Tuple[str]:
+        """Get names of tensors.
+        
+        Parameters
+        ----------
+        suffix : str
+            The suffix of the scope
+
+        Returns
+        -------
+        Tuple[str]
+            Names of tensors
+        """
+        raise NotImplementedError("Descriptor %s doesn't support this property!" % type(self).__name__)
+
+    def pass_tensors_from_frz_model(self,
+                                    *tensors : tf.Tensor,
+    ) -> None:
+        """
+        Pass the descrpt_reshape tensor as well as descrpt_deriv tensor from the frz graph_def
+
+        Parameters
+        ----------
+        *tensors : tf.Tensor
+            passed tensors
+        
+        Notes
+        -----
+        The number of parameters in the method must be equal to the numbers of returns in
+        :meth:`get_tensor_names`.
+        """
+        raise NotImplementedError("Descriptor %s doesn't support this method!" % type(self).__name__)

_skbuild/linux-x86_64-3.6/cmake-install/deepmd/descriptor/hybrid.py

+import numpy as np
+from typing import Tuple, List
+
+from deepmd.env import tf
+from deepmd.common import ClassArg
+from deepmd.env import op_module
+from deepmd.env import GLOBAL_TF_FLOAT_PRECISION
+from deepmd.env import GLOBAL_NP_FLOAT_PRECISION
+# from deepmd.descriptor import DescrptLocFrame
+# from deepmd.descriptor import DescrptSeA
+# from deepmd.descriptor import DescrptSeT
+# from deepmd.descriptor import DescrptSeAEbd
+# from deepmd.descriptor import DescrptSeAEf
+# from deepmd.descriptor import DescrptSeR
+from .descriptor import Descriptor
+from .se_a import DescrptSeA
+from .se_r import DescrptSeR
+from .se_t import DescrptSeT
+from .se_a_ebd import DescrptSeAEbd
+from .se_a_ef import DescrptSeAEf
+from .loc_frame import DescrptLocFrame
+
+@Descriptor.register("hybrid")
+class DescrptHybrid (Descriptor):
+    """Concate a list of descriptors to form a new descriptor.
+
+    Parameters
+    ----------
+    list : list
+            Build a descriptor from the concatenation of the list of descriptors.
+    """
+    def __init__ (self, 
+                  list : list
+    ) -> None :
+        """
+        Constructor
+        """
+        # warning: list is conflict with built-in list
+        descrpt_list = list
+        if descrpt_list == [] or descrpt_list is None:
+            raise RuntimeError('cannot build descriptor from an empty list of descriptors.')
+        formatted_descript_list = []
+        for ii in descrpt_list:
+            if isinstance(ii, Descriptor):
+                formatted_descript_list.append(ii)
+            elif isinstance(ii, dict):
+                formatted_descript_list.append(Descriptor(**ii))
+            else:
+                raise NotImplementedError
+        # args = ClassArg()\
+        #        .add('list', list, must = True)
+        # class_data = args.parse(jdata)
+        # dict_list = class_data['list']
+        self.descrpt_list = formatted_descript_list
+        self.numb_descrpt = len(self.descrpt_list)
+        for ii in range(1, self.numb_descrpt):
+            assert(self.descrpt_list[ii].get_ntypes() == 
+                   self.descrpt_list[ 0].get_ntypes()), \
+                   f'number of atom types in {ii}th descrptor does not match others'
+
+
+    def get_rcut (self) -> float:
+        """
+        Returns the cut-off radius
+        """
+        all_rcut = [ii.get_rcut() for ii in self.descrpt_list]
+        return np.max(all_rcut)
+
+
+    def get_ntypes (self) -> int:
+        """
+        Returns the number of atom types
+        """
+        return self.descrpt_list[0].get_ntypes()
+
+
+    def get_dim_out (self) -> int:
+        """
+        Returns the output dimension of this descriptor
+        """
+        all_dim_out = [ii.get_dim_out() for ii in self.descrpt_list]
+        return sum(all_dim_out)
+
+
+    def get_nlist_i(self, 
+                    ii : int
+    ) -> Tuple[tf.Tensor, tf.Tensor, List[int], List[int]]:
+        """Get the neighbor information of the ii-th descriptor
+
+        Parameters
+        ----------
+        ii : int
+                The index of the descriptor
+
+        Returns
+        -------
+        nlist
+                Neighbor list
+        rij
+                The relative distance between the neighbor and the center atom.
+        sel_a
+                The number of neighbors with full information
+        sel_r
+                The number of neighbors with only radial information
+        """
+        return self.descrpt_list[ii].nlist, self.descrpt_list[ii].rij, self.descrpt_list[ii].sel_a, self.descrpt_list[ii].sel_r
+    
+
+    def compute_input_stats (self,
+                             data_coord : list, 
+                             data_box : list, 
+                             data_atype : list, 
+                             natoms_vec : list,
+                             mesh : list, 
+                             input_dict : dict
+    ) -> None :
+        """
+        Compute the statisitcs (avg and std) of the training data. The input will be normalized by the statistics.
+        
+        Parameters
+        ----------
+        data_coord
+                The coordinates. Can be generated by deepmd.model.make_stat_input
+        data_box
+                The box. Can be generated by deepmd.model.make_stat_input
+        data_atype
+                The atom types. Can be generated by deepmd.model.make_stat_input
+        natoms_vec
+                The vector for the number of atoms of the system and different types of atoms. Can be generated by deepmd.model.make_stat_input
+        mesh
+                The mesh for neighbor searching. Can be generated by deepmd.model.make_stat_input
+        input_dict
+                Dictionary for additional input
+        """
+        for ii in self.descrpt_list:
+            ii.compute_input_stats(data_coord, data_box, data_atype, natoms_vec, mesh, input_dict)
+    
+
+    def build (self, 
+               coord_ : tf.Tensor, 
+               atype_ : tf.Tensor,
+               natoms : tf.Tensor,
+               box_ : tf.Tensor, 
+               mesh : tf.Tensor,
+               input_dict : dict, 
+               reuse : bool = None,
+               suffix : str = ''
+    ) -> tf.Tensor:
+        """
+        Build the computational graph for the descriptor
+
+        Parameters
+        ----------
+        coord_
+                The coordinate of atoms
+        atype_
+                The type of atoms
+        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
+        mesh
+                For historical reasons, only the length of the Tensor matters.
+                if size of mesh == 6, pbc is assumed. 
+                if size of mesh == 0, no-pbc is assumed. 
+        input_dict
+                Dictionary for additional inputs
+        reuse
+                The weights in the networks should be reused when get the variable.
+        suffix
+                Name suffix to identify this descriptor
+
+        Returns
+        -------
+        descriptor
+                The output descriptor
+        """
+        with tf.variable_scope('descrpt_attr' + suffix, reuse = reuse) :
+            t_rcut = tf.constant(self.get_rcut(), 
+                                 name = 'rcut', 
+                                 dtype = GLOBAL_TF_FLOAT_PRECISION)
+            t_ntypes = tf.constant(self.get_ntypes(), 
+                                   name = 'ntypes', 
+                                   dtype = tf.int32)
+        all_dout = []
+        for idx,ii in enumerate(self.descrpt_list):
+            dout = ii.build(coord_, atype_, natoms, box_, mesh, input_dict, suffix=suffix+f'_{idx}', reuse=reuse)
+            dout = tf.reshape(dout, [-1, ii.get_dim_out()])
+            all_dout.append(dout)
+        dout = tf.concat(all_dout, axis = 1)
+        dout = tf.reshape(dout, [-1, natoms[0], self.get_dim_out()])
+        return dout
+        
+
+    def prod_force_virial(self, 
+                          atom_ener : tf.Tensor, 
+                          natoms : tf.Tensor
+    ) -> Tuple[tf.Tensor, tf.Tensor, tf.Tensor]:
+        """
+        Compute force and virial
+
+        Parameters
+        ----------
+        atom_ener
+                The atomic energy
+        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
+
+        Returns
+        -------
+        force
+                The force on atoms
+        virial
+                The total virial
+        atom_virial
+                The atomic virial
+        """
+        for idx,ii in enumerate(self.descrpt_list):
+            ff, vv, av = ii.prod_force_virial(atom_ener, natoms)
+            if idx == 0:
+                force = ff
+                virial = vv
+                atom_virial = av
+            else:
+                force += ff
+                virial += vv
+                atom_virial += av
+        return force, virial, atom_virial
+
+    def enable_compression(self,
+                           min_nbor_dist: float,
+                           model_file: str = 'frozon_model.pb',
+                           table_extrapolate: float = 5.,
+                           table_stride_1: float = 0.01,
+                           table_stride_2: float = 0.1,
+                           check_frequency: int = -1,
+                           suffix: str = ""
+                           ) -> None:
+        """
+        Reveive the statisitcs (distance, max_nbor_size and env_mat_range) of the
+        training data.
+
+        Parameters
+        ----------
+        min_nbor_dist : float
+                The nearest distance between atoms
+        model_file : str, default: 'frozon_model.pb'
+                The original frozen model, which will be compressed by the program
+        table_extrapolate : float, default: 5.
+                The scale of model extrapolation
+        table_stride_1 : float, default: 0.01
+                The uniform stride of the first table
+        table_stride_2 : float, default: 0.1
+                The uniform stride of the second table
+        check_frequency : int, default: -1
+                The overflow check frequency
+        suffix : str, optional
+                The suffix of the scope
+        """
+        for idx, ii in enumerate(self.descrpt_list):
+            ii.enable_compression(min_nbor_dist, model_file, table_extrapolate, table_stride_1, table_stride_2, check_frequency, suffix=f"{suffix}_{idx}")
+
+
+    def enable_mixed_precision(self, mixed_prec : dict = None) -> None:
+        """
+        Reveive the mixed precision setting.
+
+        Parameters
+        ----------
+        mixed_prec
+                The mixed precision setting used in the embedding net
+        """
+        for idx, ii in enumerate(self.descrpt_list):
+            ii.enable_mixed_precision(mixed_prec)
+
+
+    def init_variables(self,
+                       graph: tf.Graph,
+                       graph_def: tf.GraphDef,
+                       suffix : str = "",
+    ) -> None:
+        """
+        Init the embedding net variables with the given dict
+
+        Parameters
+        ----------
+        graph : tf.Graph
+            The input frozen model graph
+        graph_def : tf.GraphDef
+            The input frozen model graph_def
+        suffix : str, optional
+            The suffix of the scope
+        """
+        for idx, ii in enumerate(self.descrpt_list):
+            ii.init_variables(graph, graph_def, suffix=f"{suffix}_{idx}")
+
+    def get_tensor_names(self, suffix : str = "") -> Tuple[str]:
+        """Get names of tensors.
+        
+        Parameters
+        ----------
+        suffix : str
+            The suffix of the scope
+
+        Returns
+        -------
+        Tuple[str]
+            Names of tensors
+        """
+        tensor_names = []
+        for idx, ii in enumerate(self.descrpt_list):
+            tensor_names.extend(ii.get_tensor_names(suffix=f"{suffix}_{idx}"))
+        return tuple(tensor_names)
+
+    def pass_tensors_from_frz_model(self,
+                                    *tensors : tf.Tensor,
+    ) -> None:
+        """
+        Pass the descrpt_reshape tensor as well as descrpt_deriv tensor from the frz graph_def
+
+        Parameters
+        ----------
+        *tensors : tf.Tensor
+            passed tensors
+        """
+        jj = 0
+        for ii in self.descrpt_list:
+            n_tensors = len(ii.get_tensor_names())
+            ii.pass_tensors_from_frz_model(*tensors[jj:jj+n_tensors])
+            jj += n_tensors

_skbuild/linux-x86_64-3.6/cmake-install/deepmd/descriptor/loc_frame.py

+import numpy as np
+from typing import Tuple, List
+
+from deepmd.env import tf
+from deepmd.env import GLOBAL_TF_FLOAT_PRECISION
+from deepmd.env import GLOBAL_NP_FLOAT_PRECISION
+from deepmd.env import op_module
+from deepmd.env import default_tf_session_config
+from deepmd.utils.sess import run_sess
+from .descriptor import Descriptor
+from deepmd.utils.graph import get_tensor_by_name_from_graph
+
+@Descriptor.register("loc_frame")
+class DescrptLocFrame (Descriptor) :
+    """Defines a local frame at each atom, and the compute the descriptor as local
+    coordinates under this frame.
+
+    Parameters
+    ----------
+    rcut
+            The cut-off radius
+    sel_a : list[str]
+            The length of the list should be the same as the number of atom types in the system. 
+            `sel_a[i]` gives the selected number of type-i neighbors. 
+            The full relative coordinates of the neighbors are used by the descriptor.
+    sel_r : list[str]
+            The length of the list should be the same as the number of atom types in the system. 
+            `sel_r[i]` gives the selected number of type-i neighbors. 
+            Only relative distance of the neighbors are used by the descriptor.
+            sel_a[i] + sel_r[i] is recommended to be larger than the maximally possible number of type-i neighbors in the cut-off radius.        
+    axis_rule: list[int]
+            The length should be 6 times of the number of types. 
+            - axis_rule[i*6+0]: class of the atom defining the first axis of type-i atom. 0 for neighbors with full coordinates and 1 for neighbors only with relative distance.\n\n\
+            - axis_rule[i*6+1]: type of the atom defining the first axis of type-i atom.\n\n\
+            - axis_rule[i*6+2]: index of the axis atom defining the first axis. Note that the neighbors with the same class and type are sorted according to their relative distance.\n\n\
+            - axis_rule[i*6+3]: class of the atom defining the first axis of type-i atom. 0 for neighbors with full coordinates and 1 for neighbors only with relative distance.\n\n\
+            - axis_rule[i*6+4]: type of the atom defining the second axis of type-i atom.\n\n\
+            - axis_rule[i*6+5]: class of the atom defining the second axis of type-i atom. 0 for neighbors with full coordinates and 1 for neighbors only with relative distance.    
+    """
+    def __init__(self, 
+                 rcut: float,
+                 sel_a : List[int],
+                 sel_r : List[int],
+                 axis_rule : List[int]
+    ) -> None:
+        """
+        Constructor    
+        """
+        # args = ClassArg()\
+        #        .add('sel_a',    list,   must = True) \
+        #        .add('sel_r',    list,   must = True) \
+        #        .add('rcut',     float,  default = 6.0) \
+        #        .add('axis_rule',list,   must = True)
+        # class_data = args.parse(jdata)
+        self.sel_a = sel_a
+        self.sel_r = sel_r
+        self.axis_rule = axis_rule
+        self.rcut_r = rcut
+        # ntypes and rcut_a === -1
+        self.ntypes = len(self.sel_a)
+        assert(self.ntypes == len(self.sel_r))
+        self.rcut_a = -1
+        # numb of neighbors and numb of descrptors
+        self.nnei_a = np.cumsum(self.sel_a)[-1]
+        self.nnei_r = np.cumsum(self.sel_r)[-1]
+        self.nnei = self.nnei_a + self.nnei_r
+        self.ndescrpt_a = self.nnei_a * 4
+        self.ndescrpt_r = self.nnei_r * 1
+        self.ndescrpt = self.ndescrpt_a + self.ndescrpt_r
+        self.davg = None
+        self.dstd = None
+
+        self.place_holders = {}
+        avg_zero = np.zeros([self.ntypes,self.ndescrpt]).astype(GLOBAL_NP_FLOAT_PRECISION)
+        std_ones = np.ones ([self.ntypes,self.ndescrpt]).astype(GLOBAL_NP_FLOAT_PRECISION)
+        sub_graph = tf.Graph()
+        with sub_graph.as_default():
+            name_pfx = 'd_lf_'
+            for ii in ['coord', 'box']:
+                self.place_holders[ii] = tf.placeholder(GLOBAL_NP_FLOAT_PRECISION, [None, None], name = name_pfx+'t_'+ii)
+            self.place_holders['type'] = tf.placeholder(tf.int32, [None, None], name=name_pfx+'t_type')
+            self.place_holders['natoms_vec'] = tf.placeholder(tf.int32, [self.ntypes+2], name=name_pfx+'t_natoms')
+            self.place_holders['default_mesh'] = tf.placeholder(tf.int32, [None], name=name_pfx+'t_mesh')
+            self.stat_descrpt, descrpt_deriv, rij, nlist, axis, rot_mat \
+                = op_module.descrpt (self.place_holders['coord'],
+                                     self.place_holders['type'],
+                                     self.place_holders['natoms_vec'],
+                                     self.place_holders['box'],
+                                     self.place_holders['default_mesh'],
+                                     tf.constant(avg_zero),
+                                     tf.constant(std_ones),
+                                     rcut_a = self.rcut_a,
+                                     rcut_r = self.rcut_r,
+                                     sel_a = self.sel_a,
+                                     sel_r = self.sel_r,
+                                     axis_rule = self.axis_rule)
+        self.sub_sess = tf.Session(graph = sub_graph, config=default_tf_session_config)
+
+
+    def get_rcut (self) -> float:
+        """
+        Returns the cut-off radisu
+        """
+        return self.rcut_r
+
+    def get_ntypes (self) -> int:
+        """
+        Returns the number of atom types
+        """
+        return self.ntypes
+
+    def get_dim_out (self) -> int:
+        """
+        Returns the output dimension of this descriptor
+        """
+        return self.ndescrpt
+
+    def get_nlist (self) -> Tuple[tf.Tensor, tf.Tensor, List[int], List[int]]:
+        """
+        Returns
+        -------
+        nlist
+                Neighbor list
+        rij
+                The relative distance between the neighbor and the center atom.
+        sel_a
+                The number of neighbors with full information
+        sel_r
+                The number of neighbors with only radial information
+        """
+        return self.nlist, self.rij, self.sel_a, self.sel_r
+
+    def compute_input_stats (self,
+                             data_coord : list, 
+                             data_box : list, 
+                             data_atype : list, 
+                             natoms_vec : list,
+                             mesh : list, 
+                             input_dict : dict
+    ) -> None :
+        """
+        Compute the statisitcs (avg and std) of the training data. The input will be normalized by the statistics.
+        
+        Parameters
+        ----------
+        data_coord
+                The coordinates. Can be generated by deepmd.model.make_stat_input
+        data_box
+                The box. Can be generated by deepmd.model.make_stat_input
+        data_atype
+                The atom types. Can be generated by deepmd.model.make_stat_input
+        natoms_vec
+                The vector for the number of atoms of the system and different types of atoms. Can be generated by deepmd.model.make_stat_input
+        mesh
+                The mesh for neighbor searching. Can be generated by deepmd.model.make_stat_input
+        input_dict
+                Dictionary for additional input
+        """
+        all_davg = []
+        all_dstd = []
+        if True:
+            sumv = []
+            sumn = []
+            sumv2 = []
+            for cc,bb,tt,nn,mm in zip(data_coord,data_box,data_atype,natoms_vec,mesh) :
+                sysv,sysv2,sysn \
+                    = self._compute_dstats_sys_nonsmth(cc,bb,tt,nn,mm)
+                sumv.append(sysv)
+                sumn.append(sysn)
+                sumv2.append(sysv2)
+            sumv = np.sum(sumv, axis = 0)
+            sumn = np.sum(sumn, axis = 0)
+            sumv2 = np.sum(sumv2, axis = 0)
+            for type_i in range(self.ntypes) :
+                davg = sumv[type_i] /  sumn[type_i]
+                dstd = self._compute_std(sumv2[type_i], sumv[type_i], sumn[type_i])
+                for ii in range (len(dstd)) :
+                    if (np.abs(dstd[ii]) < 1e-2) :
+                        dstd[ii] = 1e-2            
+                all_davg.append(davg)
+                all_dstd.append(dstd)
+        self.davg = np.array(all_davg)
+        self.dstd = np.array(all_dstd)        
+
+        
+    def build (self, 
+               coord_ : tf.Tensor, 
+               atype_ : tf.Tensor,
+               natoms : tf.Tensor,
+               box_ : tf.Tensor, 
+               mesh : tf.Tensor,
+               input_dict : dict, 
+               reuse : bool = None,
+               suffix : str = ''
+    ) -> tf.Tensor:
+        """
+        Build the computational graph for the descriptor
+
+        Parameters
+        ----------
+        coord_
+                The coordinate of atoms
+        atype_
+                The type of atoms
+        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
+        mesh
+                For historical reasons, only the length of the Tensor matters.
+                if size of mesh == 6, pbc is assumed. 
+                if size of mesh == 0, no-pbc is assumed. 
+        input_dict
+                Dictionary for additional inputs
+        reuse
+                The weights in the networks should be reused when get the variable.
+        suffix
+                Name suffix to identify this descriptor
+
+        Returns
+        -------
+        descriptor
+                The output descriptor
+        """
+        davg = self.davg
+        dstd = self.dstd
+        with tf.variable_scope('descrpt_attr' + suffix, reuse = reuse) :
+            if davg is None:
+                davg = np.zeros([self.ntypes, self.ndescrpt]) 
+            if dstd is None:
+                dstd = np.ones ([self.ntypes, self.ndescrpt])
+            t_rcut = tf.constant(np.max([self.rcut_r, self.rcut_a]), 
+                                 name = 'rcut', 
+                                 dtype = GLOBAL_TF_FLOAT_PRECISION)
+            t_ntypes = tf.constant(self.ntypes, 
+                                   name = 'ntypes', 
+                                   dtype = tf.int32)
+            self.t_avg = tf.get_variable('t_avg', 
+                                         davg.shape, 
+                                         dtype = GLOBAL_TF_FLOAT_PRECISION,
+                                         trainable = False,
+                                         initializer = tf.constant_initializer(davg))
+            self.t_std = tf.get_variable('t_std', 
+                                         dstd.shape, 
+                                         dtype = GLOBAL_TF_FLOAT_PRECISION,
+                                         trainable = False,
+                                         initializer = tf.constant_initializer(dstd))
+
+        coord = tf.reshape (coord_, [-1, natoms[1] * 3])
+        box   = tf.reshape (box_, [-1, 9])
+        atype = tf.reshape (atype_, [-1, natoms[1]])
+
+        self.descrpt, self.descrpt_deriv, self.rij, self.nlist, self.axis, self.rot_mat \
+            = op_module.descrpt (coord,
+                                 atype,
+                                 natoms,
+                                 box,                                    
+                                 mesh,
+                                 self.t_avg,
+                                 self.t_std,
+                                 rcut_a = self.rcut_a,
+                                 rcut_r = self.rcut_r,
+                                 sel_a = self.sel_a,
+                                 sel_r = self.sel_r,
+                                 axis_rule = self.axis_rule)
+        self.descrpt = tf.reshape(self.descrpt, [-1, self.ndescrpt])
+        tf.summary.histogram('descrpt', self.descrpt)
+        tf.summary.histogram('rij', self.rij)
+        tf.summary.histogram('nlist', self.nlist)
+
+        return self.descrpt
+
+    def get_rot_mat(self) -> tf.Tensor:
+        """
+        Get rotational matrix
+        """
+        return self.rot_mat
+
+    def prod_force_virial(self, 
+                          atom_ener : tf.Tensor, 
+                          natoms : tf.Tensor
+    ) -> Tuple[tf.Tensor, tf.Tensor, tf.Tensor]:
+        """
+        Compute force and virial
+
+        Parameters
+        ----------
+        atom_ener
+                The atomic energy
+        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
+
+        Returns
+        -------
+        force
+                The force on atoms
+        virial
+                The total virial
+        atom_virial
+                The atomic virial
+        """
+        [net_deriv] = tf.gradients (atom_ener, self.descrpt)
+        tf.summary.histogram('net_derivative', net_deriv)
+        net_deriv_reshape = tf.reshape (net_deriv, [np.cast['int64'](-1), natoms[0] * np.cast['int64'](self.ndescrpt)])
+        force = op_module.prod_force (net_deriv_reshape,
+                                      self.descrpt_deriv,
+                                      self.nlist,
+                                      self.axis,
+                                      natoms,
+                                      n_a_sel = self.nnei_a,
+                                      n_r_sel = self.nnei_r)
+        virial, atom_virial \
+            = op_module.prod_virial (net_deriv_reshape,
+                                     self.descrpt_deriv,
+                                     self.rij,
+                                     self.nlist,
+                                     self.axis,
+                                     natoms,
+                                     n_a_sel = self.nnei_a,
+                                     n_r_sel = self.nnei_r)
+        tf.summary.histogram('force', force)
+        tf.summary.histogram('virial', virial)
+        tf.summary.histogram('atom_virial', atom_virial)
+
+        return force, virial, atom_virial
+
+
+    def _compute_dstats_sys_nonsmth (self,
+                                    data_coord, 
+                                    data_box, 
+                                    data_atype, 
+                                    natoms_vec,
+                                    mesh) :    
+        dd_all \
+            = run_sess(self.sub_sess, self.stat_descrpt, 
+                                feed_dict = {
+                                    self.place_holders['coord']: data_coord,
+                                    self.place_holders['type']: data_atype,
+                                    self.place_holders['natoms_vec']: natoms_vec,
+                                    self.place_holders['box']: data_box,
+                                    self.place_holders['default_mesh']: mesh,
+                                })
+        natoms = natoms_vec
+        dd_all = np.reshape(dd_all, [-1, self.ndescrpt * natoms[0]])
+        start_index = 0
+        sysv = []
+        sysn = []
+        sysv2 = []
+        for type_i in range(self.ntypes):
+            end_index = start_index + self.ndescrpt * natoms[2+type_i]
+            dd = dd_all[:, start_index:end_index]
+            dd = np.reshape(dd, [-1, self.ndescrpt])
+            start_index = end_index        
+            # compute
+            sumv = np.sum(dd, axis = 0)
+            sumn = dd.shape[0]
+            sumv2 = np.sum(np.multiply(dd,dd), axis = 0)            
+            sysv.append(sumv)
+            sysn.append(sumn)
+            sysv2.append(sumv2)
+        return sysv, sysv2, sysn
+
+
+    def _compute_std (self,sumv2, sumv, sumn) :
+        return np.sqrt(sumv2/sumn - np.multiply(sumv/sumn, sumv/sumn))
+
+    def init_variables(self,
+                       graph: tf.Graph,
+                       graph_def: tf.GraphDef,
+                       suffix : str = "",
+    ) -> None:
+        """
+        Init the embedding net variables with the given dict
+
+        Parameters
+        ----------
+        graph : tf.Graph
+            The input frozen model graph
+        graph_def : tf.GraphDef
+            The input frozen model graph_def
+        suffix : str, optional
+            The suffix of the scope
+        """
+        self.davg = get_tensor_by_name_from_graph(graph, 'descrpt_attr%s/t_avg' % suffix)
+        self.dstd = get_tensor_by_name_from_graph(graph, 'descrpt_attr%s/t_std' % suffix)

_skbuild/linux-x86_64-3.6/cmake-install/deepmd/descriptor/se.py

+from typing import Tuple, List
+
+from deepmd.env import tf
+from deepmd.utils.graph import get_embedding_net_variables_from_graph_def, get_tensor_by_name_from_graph
+from .descriptor import Descriptor
+
+
+class DescrptSe (Descriptor):
+    """A base class for smooth version of descriptors.
+    
+    Notes
+    -----
+    All of these descriptors have an environmental matrix and an
+    embedding network (:meth:`deepmd.utils.network.embedding_net`), so
+    they can share some similiar methods without defining them twice.
+
+    Attributes
+    ----------
+    embedding_net_variables : dict
+        initial embedding network variables
+    descrpt_reshape : tf.Tensor
+        the reshaped descriptor
+    descrpt_deriv : tf.Tensor
+        the descriptor derivative
+    rij : tf.Tensor
+        distances between two atoms
+    nlist : tf.Tensor
+        the neighbor list
+    
+    """
+    def _identity_tensors(self, suffix : str = "") -> None:
+        """Identify tensors which are expected to be stored and restored.
+        
+        Notes
+        -----
+        These tensors will be indentitied:
+            self.descrpt_reshape : o_rmat
+            self.descrpt_deriv : o_rmat_deriv
+            self.rij : o_rij
+            self.nlist : o_nlist
+        Thus, this method should be called during building the descriptor and
+        after these tensors are initialized.
+
+        Parameters
+        ----------
+        suffix : str
+            The suffix of the scope
+        """
+        self.descrpt_reshape = tf.identity(self.descrpt_reshape, name = 'o_rmat' + suffix)
+        self.descrpt_deriv = tf.identity(self.descrpt_deriv, name = 'o_rmat_deriv' + suffix)
+        self.rij = tf.identity(self.rij, name = 'o_rij' + suffix)
+        self.nlist = tf.identity(self.nlist, name = 'o_nlist' + suffix)
+
+    def get_tensor_names(self, suffix : str = "") -> Tuple[str]:
+        """Get names of tensors.
+        
+        Parameters
+        ----------
+        suffix : str
+            The suffix of the scope
+
+        Returns
+        -------
+        Tuple[str]
+            Names of tensors
+        """
+        return (f'o_rmat{suffix}:0', f'o_rmat_deriv{suffix}:0', f'o_rij{suffix}:0', f'o_nlist{suffix}:0')
+
+    def pass_tensors_from_frz_model(self,
+                                    descrpt_reshape : tf.Tensor,
+                                    descrpt_deriv   : tf.Tensor,
+                                    rij             : tf.Tensor,
+                                    nlist           : tf.Tensor
+    ):
+        """
+        Pass the descrpt_reshape tensor as well as descrpt_deriv tensor from the frz graph_def
+
+        Parameters
+        ----------
+        descrpt_reshape
+                The passed descrpt_reshape tensor
+        descrpt_deriv
+                The passed descrpt_deriv tensor
+        rij
+                The passed rij tensor
+        nlist
+                The passed nlist tensor
+        """
+        self.rij = rij
+        self.nlist = nlist
+        self.descrpt_deriv = descrpt_deriv
+        self.descrpt_reshape = descrpt_reshape
+
+    def init_variables(self,
+                       graph: tf.Graph,
+                       graph_def: tf.GraphDef,
+                       suffix : str = "",
+    ) -> None:
+        """
+        Init the embedding net variables with the given dict
+
+        Parameters
+        ----------
+        graph : tf.Graph
+            The input frozen model graph
+        graph_def : tf.GraphDef
+            The input frozen model graph_def
+        suffix : str, optional
+            The suffix of the scope
+        """
+        self.embedding_net_variables = get_embedding_net_variables_from_graph_def(graph_def, suffix = suffix)
+        self.davg = get_tensor_by_name_from_graph(graph, 'descrpt_attr%s/t_avg' % suffix)
+        self.dstd = get_tensor_by_name_from_graph(graph, 'descrpt_attr%s/t_std' % suffix)
+
+    @property
+    def precision(self) -> tf.DType:
+        """Precision of filter network."""
+        return self.filter_precision

_skbuild/linux-x86_64-3.6/cmake-install/deepmd/descriptor/se_a.py

+import math
+import numpy as np
+from typing import Tuple, List, Dict, Any
+
+from deepmd.env import tf
+from deepmd.common import get_activation_func, get_precision, cast_precision
+from deepmd.env import GLOBAL_TF_FLOAT_PRECISION
+from deepmd.env import GLOBAL_NP_FLOAT_PRECISION
+from deepmd.env import op_module
+from deepmd.env import default_tf_session_config
+from deepmd.utils.network import embedding_net, embedding_net_rand_seed_shift
+from deepmd.utils.tabulate import DPTabulate
+from deepmd.utils.type_embed import embed_atom_type
+from deepmd.utils.sess import run_sess
+from deepmd.utils.graph import load_graph_def, get_tensor_by_name_from_graph, get_tensor_by_name
+from deepmd.utils.errors import GraphWithoutTensorError
+from .descriptor import Descriptor
+from .se import DescrptSe
+
+from deepmd.nvnmd.descriptor.se_a import descrpt2r4, build_davg_dstd, build_op_descriptor, filter_lower_R42GR, filter_GR2D
+from deepmd.nvnmd.utils.config import nvnmd_cfg 
+
+@Descriptor.register("se_e2_a")
+@Descriptor.register("se_a")
+class DescrptSeA (DescrptSe):
+    r"""DeepPot-SE constructed from all information (both angular and radial) of
+    atomic configurations. The embedding takes the distance between atoms as input.
+
+    The descriptor :math:`\mathcal{D}^i \in \mathcal{R}^{M_1 \times M_2}` is given by [1]_
+
+    .. math::
+        \mathcal{D}^i = (\mathcal{G}^i)^T \mathcal{R}^i (\mathcal{R}^i)^T \mathcal{G}^i_<
+
+    where :math:`\mathcal{R}^i \in \mathbb{R}^{N \times 4}` is the coordinate
+    matrix, and each row of :math:`\mathcal{R}^i` can be constructed as follows
+
+    .. math::
+        (\mathcal{R}^i)_j = [
+        \begin{array}{c}
+            s(r_{ji}) & \frac{s(r_{ji})x_{ji}}{r_{ji}} & \frac{s(r_{ji})y_{ji}}{r_{ji}} & \frac{s(r_{ji})z_{ji}}{r_{ji}}
+        \end{array}
+        ]
+
+    where :math:`\mathbf{R}_{ji}=\mathbf{R}_j-\mathbf{R}_i = (x_{ji}, y_{ji}, z_{ji})` is 
+    the relative coordinate and :math:`r_{ji}=\lVert \mathbf{R}_{ji} \lVert` is its norm.
+    The switching function :math:`s(r)` is defined as:
+
+    .. math::
+        s(r)=
+        \begin{cases}
+        \frac{1}{r}, & r<r_s \\
+        \frac{1}{r} \{ {(\frac{r - r_s}{ r_c - r_s})}^3 (-6 {(\frac{r - r_s}{ r_c - r_s})}^2 +15 \frac{r - r_s}{ r_c - r_s} -10) +1 \}, & r_s \leq r<r_c \\
+        0, & r \geq r_c
+        \end{cases}
+
+    Each row of the embedding matrix  :math:`\mathcal{G}^i \in \mathbb{R}^{N \times M_1}` consists of outputs
+    of a embedding network :math:`\mathcal{N}` of :math:`s(r_{ji})`:
+
+    .. math::
+        (\mathcal{G}^i)_j = \mathcal{N}(s(r_{ji}))
+
+    :math:`\mathcal{G}^i_< \in \mathbb{R}^{N \times M_2}` takes first :math:`M_2`$` columns of
+    :math:`\mathcal{G}^i`$`. The equation of embedding network :math:`\mathcal{N}` can be found at
+    :meth:`deepmd.utils.network.embedding_net`.
+
+    Parameters
+    ----------
+    rcut
+            The cut-off radius :math:`r_c`
+    rcut_smth
+            From where the environment matrix should be smoothed :math:`r_s`
+    sel : list[str]
+            sel[i] specifies the maxmum number of type i atoms in the cut-off radius
+    neuron : list[int]
+            Number of neurons in each hidden layers of the embedding net :math:`\mathcal{N}`
+    axis_neuron
+            Number of the axis neuron :math:`M_2` (number of columns of the sub-matrix of the embedding matrix)
+    resnet_dt
+            Time-step `dt` in the resnet construction:
+            y = x + dt * \phi (Wx + b)
+    trainable
+            If the weights of embedding net are trainable.
+    seed
+            Random seed for initializing the network parameters.
+    type_one_side
+            Try to build N_types embedding nets. Otherwise, building N_types^2 embedding nets
+    exclude_types : List[List[int]]
+            The excluded pairs of types which have no interaction with each other.
+            For example, `[[0, 1]]` means no interaction between type 0 and type 1.
+    set_davg_zero
+            Set the shift of embedding net input to zero.
+    activation_function
+            The activation function in the embedding net. Supported options are |ACTIVATION_FN|
+    precision
+            The precision of the embedding net parameters. Supported options are |PRECISION|
+    uniform_seed
+            Only for the purpose of backward compatibility, retrieves the old behavior of using the random seed
+    
+    References
+    ----------
+    .. [1] Linfeng Zhang, Jiequn Han, Han Wang, Wissam A. Saidi, Roberto Car, and E. Weinan. 2018.
+       End-to-end symmetry preserving inter-atomic potential energy model for finite and extended
+       systems. In Proceedings of the 32nd International Conference on Neural Information Processing
+       Systems (NIPS'18). Curran Associates Inc., Red Hook, NY, USA, 4441–4451.
+    """
+    def __init__ (self, 
+                  rcut: float,
+                  rcut_smth: float,
+                  sel: List[str],
+                  neuron: List[int] = [24,48,96],
+                  axis_neuron: int = 8,
+                  resnet_dt: bool = False,
+                  trainable: bool = True,
+                  seed: int = None,
+                  type_one_side: bool = True,
+                  exclude_types: List[List[int]] = [],
+                  set_davg_zero: bool = False,
+                  activation_function: str = 'tanh',
+                  precision: str = 'default',
+                  uniform_seed: bool = False
+    ) -> None:
+        """
+        Constructor
+        """
+        if rcut < rcut_smth:
+            raise RuntimeError("rcut_smth (%f) should be no more than rcut (%f)!" % (rcut_smth, rcut))
+        self.sel_a = sel
+        self.rcut_r = rcut
+        self.rcut_r_smth = rcut_smth
+        self.filter_neuron = neuron
+        self.n_axis_neuron = axis_neuron
+        self.filter_resnet_dt = resnet_dt
+        self.seed = seed
+        self.uniform_seed = uniform_seed
+        self.seed_shift = embedding_net_rand_seed_shift(self.filter_neuron)
+        self.trainable = trainable
+        self.compress_activation_fn = get_activation_func(activation_function)
+        self.filter_activation_fn = get_activation_func(activation_function)
+        self.filter_precision = get_precision(precision)
+        self.exclude_types = set()
+        for tt in exclude_types:
+            assert(len(tt) == 2)
+            self.exclude_types.add((tt[0], tt[1]))
+            self.exclude_types.add((tt[1], tt[0]))
+        self.set_davg_zero = set_davg_zero
+        self.type_one_side = type_one_side
+
+        # descrpt config
+        self.sel_r = [ 0 for ii in range(len(self.sel_a)) ]
+        self.ntypes = len(self.sel_a)
+        assert(self.ntypes == len(self.sel_r))
+        self.rcut_a = -1
+        # numb of neighbors and numb of descrptors
+        self.nnei_a = np.cumsum(self.sel_a)[-1]
+        self.nnei_r = np.cumsum(self.sel_r)[-1]
+        self.nnei = self.nnei_a + self.nnei_r
+        self.ndescrpt_a = self.nnei_a * 4
+        self.ndescrpt_r = self.nnei_r * 1
+        self.ndescrpt = self.ndescrpt_a + self.ndescrpt_r
+        self.useBN = False
+        self.dstd = None
+        self.davg = None
+        self.compress = False
+        self.embedding_net_variables = None
+        self.mixed_prec = None
+        self.place_holders = {}
+        self.nei_type = np.repeat(np.arange(self.ntypes), self.sel_a)  # like a mask
+
+        avg_zero = np.zeros([self.ntypes,self.ndescrpt]).astype(GLOBAL_NP_FLOAT_PRECISION)
+        std_ones = np.ones ([self.ntypes,self.ndescrpt]).astype(GLOBAL_NP_FLOAT_PRECISION)
+        sub_graph = tf.Graph()
+        with sub_graph.as_default():
+            name_pfx = 'd_sea_'
+            for ii in ['coord', 'box']:
+                self.place_holders[ii] = tf.placeholder(GLOBAL_NP_FLOAT_PRECISION, [None, None], name = name_pfx+'t_'+ii)
+            self.place_holders['type'] = tf.placeholder(tf.int32, [None, None], name=name_pfx+'t_type')
+            self.place_holders['natoms_vec'] = tf.placeholder(tf.int32, [self.ntypes+2], name=name_pfx+'t_natoms')
+            self.place_holders['default_mesh'] = tf.placeholder(tf.int32, [None], name=name_pfx+'t_mesh')
+            self.stat_descrpt, descrpt_deriv, rij, nlist \
+                = op_module.prod_env_mat_a(self.place_holders['coord'],
+                                         self.place_holders['type'],
+                                         self.place_holders['natoms_vec'],
+                                         self.place_holders['box'],
+                                         self.place_holders['default_mesh'],
+                                         tf.constant(avg_zero),
+                                         tf.constant(std_ones),
+                                         rcut_a = self.rcut_a,
+                                         rcut_r = self.rcut_r,
+                                         rcut_r_smth = self.rcut_r_smth,
+                                         sel_a = self.sel_a,
+                                         sel_r = self.sel_r)
+        self.sub_sess = tf.Session(graph = sub_graph, config=default_tf_session_config)
+        self.original_sel = None
+
+
+    def get_rcut (self) -> float:
+        """
+        Returns the cut-off radius
+        """
+        return self.rcut_r
+
+    def get_ntypes (self) -> int:
+        """
+        Returns the number of atom types
+        """
+        return self.ntypes
+
+    def get_dim_out (self) -> int:
+        """
+        Returns the output dimension of this descriptor
+        """
+        return self.filter_neuron[-1] * self.n_axis_neuron
+
+    def get_dim_rot_mat_1 (self) -> int:
+        """
+        Returns the first dimension of the rotation matrix. The rotation is of shape dim_1 x 3
+        """
+        return self.filter_neuron[-1]
+
+    def get_nlist (self) -> Tuple[tf.Tensor, tf.Tensor, List[int], List[int]]:
+        """
+        Returns
+        -------
+        nlist
+                Neighbor list
+        rij
+                The relative distance between the neighbor and the center atom.
+        sel_a
+                The number of neighbors with full information
+        sel_r
+                The number of neighbors with only radial information
+        """
+        return self.nlist, self.rij, self.sel_a, self.sel_r
+
+    def compute_input_stats (self,
+                             data_coord : list, 
+                             data_box : list, 
+                             data_atype : list, 
+                             natoms_vec : list,
+                             mesh : list, 
+                             input_dict : dict
+    ) -> None :
+        """
+        Compute the statisitcs (avg and std) of the training data. The input will be normalized by the statistics.
+        
+        Parameters
+        ----------
+        data_coord
+                The coordinates. Can be generated by deepmd.model.make_stat_input
+        data_box
+                The box. Can be generated by deepmd.model.make_stat_input
+        data_atype
+                The atom types. Can be generated by deepmd.model.make_stat_input
+        natoms_vec
+                The vector for the number of atoms of the system and different types of atoms. Can be generated by deepmd.model.make_stat_input
+        mesh
+                The mesh for neighbor searching. Can be generated by deepmd.model.make_stat_input
+        input_dict
+                Dictionary for additional input
+        """
+        all_davg = []
+        all_dstd = []
+        if True:
+            sumr = []
+            suma = []
+            sumn = []
+            sumr2 = []
+            suma2 = []
+            for cc,bb,tt,nn,mm in zip(data_coord,data_box,data_atype,natoms_vec,mesh) :
+                sysr,sysr2,sysa,sysa2,sysn \
+                    = self._compute_dstats_sys_smth(cc,bb,tt,nn,mm)
+                sumr.append(sysr)
+                suma.append(sysa)
+                sumn.append(sysn)
+                sumr2.append(sysr2)
+                suma2.append(sysa2)
+            sumr = np.sum(sumr, axis = 0)
+            suma = np.sum(suma, axis = 0)
+            sumn = np.sum(sumn, axis = 0)
+            sumr2 = np.sum(sumr2, axis = 0)
+            suma2 = np.sum(suma2, axis = 0)
+            for type_i in range(self.ntypes) :
+                davgunit = [sumr[type_i]/(sumn[type_i]+1e-15), 0, 0, 0]
+                dstdunit = [self._compute_std(sumr2[type_i], sumr[type_i], sumn[type_i]), 
+                            self._compute_std(suma2[type_i], suma[type_i], sumn[type_i]), 
+                            self._compute_std(suma2[type_i], suma[type_i], sumn[type_i]), 
+                            self._compute_std(suma2[type_i], suma[type_i], sumn[type_i])
+                            ]
+                davg = np.tile(davgunit, self.ndescrpt // 4)
+                dstd = np.tile(dstdunit, self.ndescrpt // 4)
+                all_davg.append(davg)
+                all_dstd.append(dstd)
+
+        if not self.set_davg_zero:
+            self.davg = np.array(all_davg)
+        self.dstd = np.array(all_dstd)
+
+    def enable_compression(self,
+                           min_nbor_dist : float,
+                           model_file : str = 'frozon_model.pb',
+                           table_extrapolate : float = 5,
+                           table_stride_1 : float = 0.01,
+                           table_stride_2 : float = 0.1,
+                           check_frequency : int = -1,
+                           suffix : str = "",
+    ) -> None:
+        """
+        Reveive the statisitcs (distance, max_nbor_size and env_mat_range) of the training data.
+        
+        Parameters
+        ----------
+        min_nbor_dist
+                The nearest distance between atoms
+        model_file
+                The original frozen model, which will be compressed by the program
+        table_extrapolate
+                The scale of model extrapolation
+        table_stride_1
+                The uniform stride of the first table
+        table_stride_2
+                The uniform stride of the second table
+        check_frequency
+                The overflow check frequency
+        suffix : str, optional
+                The suffix of the scope
+        """
+        # do some checks before the mocel compression process
+        assert (
+            not self.filter_resnet_dt
+        ), "Model compression error: descriptor resnet_dt must be false!"
+        for tt in self.exclude_types:
+            if (tt[0] not in range(self.ntypes)) or (tt[1] not in range(self.ntypes)):
+                raise RuntimeError("exclude types" + str(tt) + " must within the number of atomic types " + str(self.ntypes) + "!")
+        if (self.ntypes * self.ntypes - len(self.exclude_types) == 0):
+            raise RuntimeError("empty embedding-net are not supported in model compression!")
+
+        for ii in range(len(self.filter_neuron) - 1):
+            if self.filter_neuron[ii] * 2 != self.filter_neuron[ii + 1]:
+                raise NotImplementedError(
+                    "Model Compression error: descriptor neuron [%s] is not supported by model compression! "
+                    "The size of the next layer of the neural network must be twice the size of the previous layer." 
+                    % ','.join([str(item) for item in self.filter_neuron])
+                )
+
+        self.compress = True
+        self.table = DPTabulate(
+            self, self.filter_neuron, model_file, self.type_one_side, self.exclude_types, self.compress_activation_fn, suffix=suffix)
+        self.table_config = [table_extrapolate, table_stride_1, table_stride_2, check_frequency]
+        self.lower, self.upper \
+            = self.table.build(min_nbor_dist, 
+                               table_extrapolate, 
+                               table_stride_1, 
+                               table_stride_2)
+        
+        graph, _ = load_graph_def(model_file)
+        self.davg = get_tensor_by_name_from_graph(graph, 'descrpt_attr%s/t_avg' % suffix)
+        self.dstd = get_tensor_by_name_from_graph(graph, 'descrpt_attr%s/t_std' % suffix)
+
+
+    def enable_mixed_precision(self, mixed_prec : dict = None) -> None:
+        """
+        Reveive the mixed precision setting.
+
+        Parameters
+        ----------
+        mixed_prec
+                The mixed precision setting used in the embedding net
+        """
+        self.mixed_prec = mixed_prec
+        self.filter_precision = get_precision(mixed_prec['output_prec'])
+
+
+    def build (self, 
+               coord_ : tf.Tensor, 
+               atype_ : tf.Tensor,
+               natoms : tf.Tensor,
+               box_ : tf.Tensor, 
+               mesh : tf.Tensor,
+               input_dict : dict, 
+               reuse : bool = None,
+               suffix : str = ''
+    ) -> tf.Tensor:
+        """
+        Build the computational graph for the descriptor
+
+        Parameters
+        ----------
+        coord_
+                The coordinate of atoms
+        atype_
+                The type of atoms
+        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
+        mesh
+                For historical reasons, only the length of the Tensor matters.
+                if size of mesh == 6, pbc is assumed. 
+                if size of mesh == 0, no-pbc is assumed. 
+        input_dict
+                Dictionary for additional inputs
+        reuse
+                The weights in the networks should be reused when get the variable.
+        suffix
+                Name suffix to identify this descriptor
+
+        Returns
+        -------
+        descriptor
+                The output descriptor
+        """
+        davg = self.davg
+        dstd = self.dstd
+        if nvnmd_cfg.enable and nvnmd_cfg.restore_descriptor: davg, dstd = build_davg_dstd()
+        with tf.variable_scope('descrpt_attr' + suffix, reuse = reuse) :
+            if davg is None:
+                davg = np.zeros([self.ntypes, self.ndescrpt]) 
+            if dstd is None:
+                dstd = np.ones ([self.ntypes, self.ndescrpt])
+            t_rcut = tf.constant(np.max([self.rcut_r, self.rcut_a]), 
+                                 name = 'rcut', 
+                                 dtype = GLOBAL_TF_FLOAT_PRECISION)
+            t_ntypes = tf.constant(self.ntypes, 
+                                   name = 'ntypes', 
+                                   dtype = tf.int32)
+            t_ndescrpt = tf.constant(self.ndescrpt, 
+                                     name = 'ndescrpt', 
+                                     dtype = tf.int32)            
+            t_sel = tf.constant(self.sel_a, 
+                                name = 'sel', 
+                                dtype = tf.int32)            
+            t_original_sel = tf.constant(self.original_sel if self.original_sel is not None else self.sel_a,
+                name = 'original_sel',
+                dtype = tf.int32)
+            self.t_avg = tf.get_variable('t_avg', 
+                                         davg.shape, 
+                                         dtype = GLOBAL_TF_FLOAT_PRECISION,
+                                         trainable = False,
+                                         initializer = tf.constant_initializer(davg))
+            self.t_std = tf.get_variable('t_std', 
+                                         dstd.shape, 
+                                         dtype = GLOBAL_TF_FLOAT_PRECISION,
+                                         trainable = False,
+                                         initializer = tf.constant_initializer(dstd))
+
+        with tf.control_dependencies([t_sel, t_original_sel]):
+            coord = tf.reshape (coord_, [-1, natoms[1] * 3])
+        box   = tf.reshape (box_, [-1, 9])
+        atype = tf.reshape (atype_, [-1, natoms[1]])
+
+        op_descriptor = build_op_descriptor() if nvnmd_cfg.enable else op_module.prod_env_mat_a
+        self.descrpt, self.descrpt_deriv, self.rij, self.nlist \
+            = op_descriptor           (coord,
+                                       atype,
+                                       natoms,
+                                       box,
+                                       mesh,
+                                       self.t_avg,
+                                       self.t_std,
+                                       rcut_a = self.rcut_a,
+                                       rcut_r = self.rcut_r,
+                                       rcut_r_smth = self.rcut_r_smth,
+                                       sel_a = self.sel_a,
+                                       sel_r = self.sel_r)
+        # only used when tensorboard was set as true
+        tf.summary.histogram('descrpt', self.descrpt)
+        tf.summary.histogram('rij', self.rij)
+        tf.summary.histogram('nlist', self.nlist)
+
+        self.descrpt_reshape = tf.reshape(self.descrpt, [-1, self.ndescrpt])
+        self._identity_tensors(suffix=suffix)
+
+        self.dout, self.qmat = self._pass_filter(self.descrpt_reshape, 
+                                                 atype,
+                                                 natoms, 
+                                                 input_dict,
+                                                 suffix = suffix, 
+                                                 reuse = reuse, 
+                                                 trainable = self.trainable)
+
+        # only used when tensorboard was set as true
+        tf.summary.histogram('embedding_net_output', self.dout)
+        return self.dout
+    
+    def get_rot_mat(self) -> tf.Tensor:
+        """
+        Get rotational matrix
+        """
+        return self.qmat
+
+    def prod_force_virial(self, 
+                          atom_ener : tf.Tensor, 
+                          natoms : tf.Tensor
+    ) -> Tuple[tf.Tensor, tf.Tensor, tf.Tensor]:
+        """
+        Compute force and virial
+
+        Parameters
+        ----------
+        atom_ener
+                The atomic energy
+        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
+
+        Returns
+        -------
+        force
+                The force on atoms
+        virial
+                The total virial
+        atom_virial
+                The atomic virial
+        """
+        [net_deriv] = tf.gradients (atom_ener, self.descrpt_reshape)
+        tf.summary.histogram('net_derivative', net_deriv)
+        net_deriv_reshape = tf.reshape (net_deriv, [np.cast['int64'](-1), natoms[0] * np.cast['int64'](self.ndescrpt)])        
+        force \
+            = op_module.prod_force_se_a (net_deriv_reshape,
+                                          self.descrpt_deriv,
+                                          self.nlist,
+                                          natoms,
+                                          n_a_sel = self.nnei_a,
+                                          n_r_sel = self.nnei_r)
+        virial, atom_virial \
+            = op_module.prod_virial_se_a (net_deriv_reshape,
+                                           self.descrpt_deriv,
+                                           self.rij,
+                                           self.nlist,
+                                           natoms,
+                                           n_a_sel = self.nnei_a,
+                                           n_r_sel = self.nnei_r)
+        tf.summary.histogram('force', force)
+        tf.summary.histogram('virial', virial)
+        tf.summary.histogram('atom_virial', atom_virial)
+        
+        return force, virial, atom_virial
+        
+
+    def _pass_filter(self, 
+                     inputs,
+                     atype,
+                     natoms,
+                     input_dict,
+                     reuse = None,
+                     suffix = '', 
+                     trainable = True) :
+        if input_dict is not None:
+            type_embedding = input_dict.get('type_embedding', None)
+        else:
+            type_embedding = None
+        start_index = 0
+        inputs = tf.reshape(inputs, [-1, natoms[0], self.ndescrpt])
+        output = []
+        output_qmat = []
+        if not (self.type_one_side and len(self.exclude_types) == 0) and type_embedding is None:
+            for type_i in range(self.ntypes):
+                inputs_i = tf.slice (inputs,
+                                     [ 0, start_index, 0],
+                                     [-1, natoms[2+type_i], -1] )
+                inputs_i = tf.reshape(inputs_i, [-1, self.ndescrpt])
+                if self.type_one_side:
+                    # reuse NN parameters for all types to support type_one_side along with exclude_types
+                    reuse = tf.AUTO_REUSE
+                    filter_name = 'filter_type_all'+suffix
+                else:
+                    filter_name = 'filter_type_'+str(type_i)+suffix
+                layer, qmat = self._filter(inputs_i, type_i, name=filter_name, natoms=natoms, reuse=reuse, trainable = trainable, activation_fn = self.filter_activation_fn)
+                layer = tf.reshape(layer, [tf.shape(inputs)[0], natoms[2+type_i], self.get_dim_out()])
+                qmat  = tf.reshape(qmat,  [tf.shape(inputs)[0], natoms[2+type_i], self.get_dim_rot_mat_1() * 3])
+                output.append(layer)
+                output_qmat.append(qmat)
+                start_index += natoms[2+type_i]
+        else :
+            inputs_i = inputs
+            inputs_i = tf.reshape(inputs_i, [-1, self.ndescrpt])
+            type_i = -1
+            if nvnmd_cfg.enable and nvnmd_cfg.quantize_descriptor: 
+                inputs_i = descrpt2r4(inputs_i, natoms)
+            layer, qmat = self._filter(inputs_i, type_i, name='filter_type_all'+suffix, natoms=natoms, reuse=reuse, trainable = trainable, activation_fn = self.filter_activation_fn, type_embedding=type_embedding)
+            layer = tf.reshape(layer, [tf.shape(inputs)[0], natoms[0], self.get_dim_out()])
+            qmat  = tf.reshape(qmat,  [tf.shape(inputs)[0], natoms[0], self.get_dim_rot_mat_1() * 3])
+            output.append(layer)
+            output_qmat.append(qmat)
+        output = tf.concat(output, axis = 1)
+        output_qmat = tf.concat(output_qmat, axis = 1)
+        return output, output_qmat
+
+
+    def _compute_dstats_sys_smth (self,
+                                 data_coord, 
+                                 data_box, 
+                                 data_atype,                             
+                                 natoms_vec,
+                                 mesh) :    
+        dd_all \
+            = run_sess(self.sub_sess, self.stat_descrpt, 
+                                feed_dict = {
+                                    self.place_holders['coord']: data_coord,
+                                    self.place_holders['type']: data_atype,
+                                    self.place_holders['natoms_vec']: natoms_vec,
+                                    self.place_holders['box']: data_box,
+                                    self.place_holders['default_mesh']: mesh,
+                                })
+        natoms = natoms_vec
+        dd_all = np.reshape(dd_all, [-1, self.ndescrpt * natoms[0]])
+        start_index = 0
+        sysr = []
+        sysa = []
+        sysn = []
+        sysr2 = []
+        sysa2 = []
+        for type_i in range(self.ntypes):
+            end_index = start_index + self.ndescrpt * natoms[2+type_i]
+            dd = dd_all[:, start_index:end_index]
+            dd = np.reshape(dd, [-1, self.ndescrpt])
+            start_index = end_index        
+            # compute
+            dd = np.reshape (dd, [-1, 4])
+            ddr = dd[:,:1]
+            dda = dd[:,1:]
+            sumr = np.sum(ddr)
+            suma = np.sum(dda) / 3.
+            sumn = dd.shape[0]
+            sumr2 = np.sum(np.multiply(ddr, ddr))
+            suma2 = np.sum(np.multiply(dda, dda)) / 3.
+            sysr.append(sumr)
+            sysa.append(suma)
+            sysn.append(sumn)
+            sysr2.append(sumr2)
+            sysa2.append(suma2)
+        return sysr, sysr2, sysa, sysa2, sysn
+
+
+    def _compute_std (self,sumv2, sumv, sumn) :
+        if sumn == 0:
+            return 1. / self.rcut_r
+        val = np.sqrt(sumv2/sumn - np.multiply(sumv/sumn, sumv/sumn))
+        if np.abs(val) < 1e-2:
+            val = 1e-2
+        return val
+
+
+    def _concat_type_embedding(
+            self,
+            xyz_scatter,
+            nframes,
+            natoms,
+            type_embedding,
+    ):
+        '''Concatenate `type_embedding` of neighbors and `xyz_scatter`.
+        If not self.type_one_side, concatenate `type_embedding` of center atoms as well.
+
+        Parameters
+        ----------
+        xyz_scatter:
+                shape is [nframes*natoms[0]*self.nnei, 1]
+        nframes:
+                shape is []
+        natoms:
+                shape is [1+1+self.ntypes]
+        type_embedding:
+                shape is [self.ntypes, Y] where Y=jdata['type_embedding']['neuron'][-1]
+
+        Returns
+        -------
+            embedding:
+                environment of each atom represented by embedding.
+        '''
+        te_out_dim = type_embedding.get_shape().as_list()[-1]
+        self.t_nei_type = tf.constant(self.nei_type, dtype=tf.int32)
+        nei_embed = tf.nn.embedding_lookup(type_embedding,tf.cast(self.t_nei_type,dtype=tf.int32))  # shape is [self.nnei, 1+te_out_dim]
+        nei_embed = tf.tile(nei_embed,(nframes*natoms[0],1))  # shape is [nframes*natoms[0]*self.nnei, te_out_dim]
+        nei_embed = tf.reshape(nei_embed,[-1,te_out_dim])
+        embedding_input = tf.concat([xyz_scatter,nei_embed],1)  # shape is [nframes*natoms[0]*self.nnei, 1+te_out_dim]
+        if not self.type_one_side:
+            atm_embed = embed_atom_type(self.ntypes, natoms, type_embedding)  # shape is [natoms[0], te_out_dim]
+            atm_embed = tf.tile(atm_embed,(nframes,self.nnei))  # shape is [nframes*natoms[0], self.nnei*te_out_dim]
+            atm_embed = tf.reshape(atm_embed,[-1,te_out_dim])  # shape is [nframes*natoms[0]*self.nnei, te_out_dim]
+            embedding_input = tf.concat([embedding_input,atm_embed],1)  # shape is [nframes*natoms[0]*self.nnei, 1+te_out_dim+te_out_dim]
+        return embedding_input
+
+
+    def _filter_lower(
+            self,
+            type_i,
+            type_input,
+            start_index,
+            incrs_index,
+            inputs,
+            nframes,
+            natoms,
+            type_embedding=None,
+            is_exclude = False,
+            activation_fn = None,
+            bavg = 0.0,
+            stddev = 1.0,
+            trainable = True,
+            suffix = '',
+    ):
+        """
+        input env matrix, returns R.G
+        """
+        outputs_size = [1] + self.filter_neuron
+        # cut-out inputs
+        # with natom x (nei_type_i x 4)  
+        inputs_i = tf.slice (inputs,
+                             [ 0, start_index* 4],
+                             [-1, incrs_index* 4] )
+        shape_i = inputs_i.get_shape().as_list()
+        natom = tf.shape(inputs_i)[0]
+        # with (natom x nei_type_i) x 4
+        inputs_reshape = tf.reshape(inputs_i, [-1, 4])
+        # with (natom x nei_type_i) x 1
+        xyz_scatter = tf.reshape(tf.slice(inputs_reshape, [0,0],[-1,1]),[-1,1])
+        if type_embedding is not None:
+            xyz_scatter = self._concat_type_embedding(
+                xyz_scatter, nframes, natoms, type_embedding)
+            if self.compress:
+                raise RuntimeError('compression of type embedded descriptor is not supported at the moment')
+        # natom x 4 x outputs_size
+        if nvnmd_cfg.enable:
+          return filter_lower_R42GR(
+            type_i, type_input, inputs_i, is_exclude,
+            activation_fn, bavg, stddev, trainable, 
+            suffix, self.seed, self.seed_shift, self.uniform_seed,
+            self.filter_neuron, self.filter_precision, self.filter_resnet_dt,
+            self.embedding_net_variables
+          )
+        if self.compress and (not is_exclude):
+            if self.type_one_side:
+                net = 'filter_-1_net_' + str(type_i)
+            else:
+                net = 'filter_' + str(type_input) + '_net_' + str(type_i)
+            info = [self.lower[net], self.upper[net], self.upper[net] * self.table_config[0], self.table_config[1], self.table_config[2], self.table_config[3]]
+            return op_module.tabulate_fusion_se_a(tf.cast(self.table.data[net], self.filter_precision), info, xyz_scatter, tf.reshape(inputs_i, [natom, shape_i[1]//4, 4]), last_layer_size = outputs_size[-1])  
+        else:
+          if (not is_exclude):
+              # with (natom x nei_type_i) x out_size
+              xyz_scatter = embedding_net(
+                  xyz_scatter, 
+                  self.filter_neuron, 
+                  self.filter_precision, 
+                  activation_fn = activation_fn, 
+                  resnet_dt = self.filter_resnet_dt,
+                  name_suffix = suffix,
+                  stddev = stddev,
+                  bavg = bavg,
+                  seed = self.seed,
+                  trainable = trainable, 
+                  uniform_seed = self.uniform_seed,
+                  initial_variables = self.embedding_net_variables,
+                  mixed_prec = self.mixed_prec)
+              if (not self.uniform_seed) and (self.seed is not None): self.seed += self.seed_shift
+          else:
+            # we can safely return the final xyz_scatter filled with zero directly
+            return tf.cast(tf.fill((natom, 4, outputs_size[-1]), 0.), self.filter_precision)
+          # natom x nei_type_i x out_size
+          xyz_scatter = tf.reshape(xyz_scatter, (-1, shape_i[1]//4, outputs_size[-1]))  
+          # When using tf.reshape(inputs_i, [-1, shape_i[1]//4, 4]) below
+          # [588 24] -> [588 6 4] correct
+          # but if sel is zero
+          # [588 0] -> [147 0 4] incorrect; the correct one is [588 0 4]
+          # So we need to explicitly assign the shape to tf.shape(inputs_i)[0] instead of -1
+          # natom x 4 x outputs_size
+          return tf.matmul(tf.reshape(inputs_i, [natom, shape_i[1]//4, 4]), xyz_scatter, transpose_a = True)
+
+
+    @cast_precision
+    def _filter(
+            self, 
+            inputs, 
+            type_input,
+            natoms,
+            type_embedding = None,
+            activation_fn=tf.nn.tanh, 
+            stddev=1.0,
+            bavg=0.0,
+            name='linear', 
+            reuse=None,
+            trainable = True):
+        nframes = tf.shape(tf.reshape(inputs, [-1, natoms[0], self.ndescrpt]))[0]
+        # natom x (nei x 4)
+        shape = inputs.get_shape().as_list()
+        outputs_size = [1] + self.filter_neuron
+        outputs_size_2 = self.n_axis_neuron
+        all_excluded = all([(type_input, type_i) in self.exclude_types for type_i in range(self.ntypes)])
+        if all_excluded:
+            # all types are excluded so result and qmat should be zeros
+            # we can safaly return a zero matrix...
+            # See also https://stackoverflow.com/a/34725458/9567349
+            # result: natom x outputs_size x outputs_size_2
+            # qmat: natom x outputs_size x 3
+            natom = tf.shape(inputs)[0]
+            result = tf.cast(tf.fill((natom, outputs_size_2, outputs_size[-1]), 0.), GLOBAL_TF_FLOAT_PRECISION)
+            qmat = tf.cast(tf.fill((natom, outputs_size[-1], 3), 0.), GLOBAL_TF_FLOAT_PRECISION)
+            return result, qmat
+            
+        with tf.variable_scope(name, reuse=reuse):
+          start_index = 0
+          type_i = 0
+          # natom x 4 x outputs_size
+          if type_embedding is None:
+              rets = []
+              for type_i in range(self.ntypes):
+                  ret = self._filter_lower(
+                      type_i, type_input,
+                      start_index, self.sel_a[type_i],
+                      inputs,
+                      nframes,
+                      natoms,
+                      type_embedding = type_embedding,
+                      is_exclude = (type_input, type_i) in self.exclude_types,
+                      activation_fn = activation_fn,
+                      stddev = stddev,
+                      bavg = bavg,
+                      trainable = trainable,
+                      suffix = "_"+str(type_i))
+                  if (type_input, type_i) not in self.exclude_types:
+                      # add zero is meaningless; skip
+                      rets.append(ret)
+                  start_index += self.sel_a[type_i]
+              # faster to use accumulate_n than multiple add
+              xyz_scatter_1 = tf.accumulate_n(rets)
+          else :
+              xyz_scatter_1 = self._filter_lower(
+                  type_i, type_input,
+                  start_index, np.cumsum(self.sel_a)[-1],
+                  inputs,
+                  nframes,
+                  natoms,
+                  type_embedding = type_embedding,
+                  is_exclude = False,
+                  activation_fn = activation_fn,
+                  stddev = stddev,
+                  bavg = bavg,
+                  trainable = trainable)
+          if nvnmd_cfg.enable: return filter_GR2D(xyz_scatter_1)
+          # natom x nei x outputs_size
+          # xyz_scatter = tf.concat(xyz_scatter_total, axis=1)
+          # natom x nei x 4
+          # inputs_reshape = tf.reshape(inputs, [-1, shape[1]//4, 4])
+          # natom x 4 x outputs_size
+          # xyz_scatter_1 = tf.matmul(inputs_reshape, xyz_scatter, transpose_a = True)
+          if self.original_sel is None:
+              # shape[1] = nnei * 4
+              nnei = shape[1] / 4
+          else:
+              nnei = tf.cast(tf.Variable(np.sum(self.original_sel), dtype=tf.int32, trainable=False, name="nnei"), self.filter_precision)
+          xyz_scatter_1 = xyz_scatter_1 / nnei
+          # natom x 4 x outputs_size_2
+          xyz_scatter_2 = tf.slice(xyz_scatter_1, [0,0,0],[-1,-1,outputs_size_2])
+          # # natom x 3 x outputs_size_2
+          # qmat = tf.slice(xyz_scatter_2, [0,1,0], [-1, 3, -1])
+          # natom x 3 x outputs_size_1
+          qmat = tf.slice(xyz_scatter_1, [0,1,0], [-1, 3, -1])
+          # natom x outputs_size_1 x 3
+          qmat = tf.transpose(qmat, perm = [0, 2, 1])
+          # natom x outputs_size x outputs_size_2
+          result = tf.matmul(xyz_scatter_1, xyz_scatter_2, transpose_a = True)
+          # natom x (outputs_size x outputs_size_2)
+          result = tf.reshape(result, [-1, outputs_size_2 * outputs_size[-1]])
+
+        return result, qmat
+
+    def init_variables(self,
+                       graph: tf.Graph,
+                       graph_def: tf.GraphDef,
+                       suffix : str = "",
+    ) -> None:
+        """
+        Init the embedding net variables with the given dict
+
+        Parameters
+        ----------
+        graph : tf.Graph
+            The input frozen model graph
+        graph_def : tf.GraphDef
+            The input frozen model graph_def
+        suffix : str, optional
+            The suffix of the scope
+        """
+        super().init_variables(graph=graph, graph_def=graph_def, suffix=suffix)
+        try:
+            self.original_sel = get_tensor_by_name_from_graph(graph, 'descrpt_attr%s/original_sel' % suffix)
+        except GraphWithoutTensorError:
+            # original_sel is not restored in old graphs, assume sel never changed before
+            pass
+        # check sel == original sel?
+        try:
+            sel = get_tensor_by_name_from_graph(graph, 'descrpt_attr%s/sel' % suffix)
+        except GraphWithoutTensorError:
+            # sel is not restored in old graphs
+            pass
+        else:
+            if not np.array_equal(np.array(self.sel_a), sel):
+                if not self.set_davg_zero:
+                    raise RuntimeError("Adjusting sel is only supported when `set_davg_zero` is true!")
+                # as set_davg_zero, self.davg is safely zero
+                self.davg = np.zeros([self.ntypes, self.ndescrpt]).astype(GLOBAL_NP_FLOAT_PRECISION)
+                new_dstd = np.ones([self.ntypes, self.ndescrpt]).astype(GLOBAL_NP_FLOAT_PRECISION)
+                # shape of davg and dstd is (ntypes, ndescrpt), ndescrpt = 4*sel
+                n_descpt = np.array(self.sel_a) * 4
+                n_descpt_old = np.array(sel) * 4
+                end_index = np.cumsum(n_descpt)
+                end_index_old = np.cumsum(n_descpt_old)
+                start_index = np.roll(end_index, 1)
+                start_index[0] = 0
+                start_index_old = np.roll(end_index_old, 1)
+                start_index_old[0] = 0
+
+                for nn, oo, ii, jj in zip(n_descpt, n_descpt_old, start_index, start_index_old):
+                    if nn < oo:
+                        # new size is smaller, copy part of std
+                        new_dstd[:, ii:ii+nn] = self.dstd[:, jj:jj+nn]
+                    else:
+                        # new size is larger, copy all, the rest remains 1
+                        new_dstd[:, ii:ii+oo] = self.dstd[:, jj:jj+oo]
+                self.dstd = new_dstd
+                if self.original_sel is None:
+                    self.original_sel = sel

_skbuild/linux-x86_64-3.6/cmake-install/deepmd/descriptor/se_a_ebd.py

+import numpy as np
+from typing import Tuple, List
+
+from deepmd.env import tf
+from deepmd.common import ClassArg, get_activation_func, get_precision, add_data_requirement
+from deepmd.utils.network import one_layer
+from deepmd.env import GLOBAL_TF_FLOAT_PRECISION
+from deepmd.env import GLOBAL_NP_FLOAT_PRECISION
+from deepmd.env import op_module
+from deepmd.env import default_tf_session_config
+from deepmd.utils.network import embedding_net
+from .se_a import DescrptSeA
+from .descriptor import Descriptor
+
+@Descriptor.register("se_a_tpe")
+@Descriptor.register("se_a_ebd")
+class DescrptSeAEbd (DescrptSeA):
+    """DeepPot-SE descriptor with type embedding approach.
+
+    Parameters
+    ----------
+    rcut
+            The cut-off radius
+    rcut_smth
+            From where the environment matrix should be smoothed
+    sel : list[str]
+            sel[i] specifies the maxmum number of type i atoms in the cut-off radius
+    neuron : list[int]
+            Number of neurons in each hidden layers of the embedding net
+    axis_neuron
+            Number of the axis neuron (number of columns of the sub-matrix of the embedding matrix)
+    resnet_dt
+            Time-step `dt` in the resnet construction:
+            y = x + dt * \phi (Wx + b)
+    trainable
+            If the weights of embedding net are trainable.
+    seed
+            Random seed for initializing the network parameters.
+    type_one_side
+            Try to build N_types embedding nets. Otherwise, building N_types^2 embedding nets
+    type_nchanl
+            Number of channels for type representation
+    type_nlayer
+            Number of hidden layers for the type embedding net (skip connected).
+    numb_aparam
+            Number of atomic parameters. If >0 it will be embedded with atom types.
+    set_davg_zero
+            Set the shift of embedding net input to zero.
+    activation_function
+            The activation function in the embedding net. Supported options are {0}
+    precision
+            The precision of the embedding net parameters. Supported options are {1}
+    exclude_types : List[List[int]]
+            The excluded pairs of types which have no interaction with each other.
+            For example, `[[0, 1]]` means no interaction between type 0 and type 1.
+    """
+    def __init__ (self, 
+                  rcut: float,
+                  rcut_smth: float,
+                  sel: List[str],
+                  neuron: List[int] = [24,48,96],
+                  axis_neuron: int = 8,
+                  resnet_dt: bool = False,
+                  trainable: bool = True,
+                  seed: int = None,
+                  type_one_side: bool = True,
+                  type_nchanl : int = 2,
+                  type_nlayer : int = 1,
+                  numb_aparam : int = 0,
+                  set_davg_zero: bool = False,
+                  activation_function: str = 'tanh',
+                  precision: str = 'default',
+                  exclude_types: List[List[int]] = [],
+    ) -> None:
+        """
+        Constructor
+        """
+        # args = ClassArg()\
+        #        .add('type_nchanl',      int,    default = 4) \
+        #        .add('type_nlayer',      int,    default = 2) \
+        #        .add('type_one_side',    bool,   default = True) \
+        #        .add('numb_aparam',      int,    default = 0)
+        # class_data = args.parse(jdata)
+        DescrptSeA.__init__(self, 
+                            rcut,
+                            rcut_smth,
+                            sel,
+                            neuron = neuron,
+                            axis_neuron = axis_neuron,
+                            resnet_dt = resnet_dt,
+                            trainable = trainable,
+                            seed = seed,
+                            type_one_side = type_one_side,
+                            set_davg_zero = set_davg_zero,
+                            activation_function = activation_function,
+                            precision = precision
+        )
+        self.type_nchanl = type_nchanl
+        self.type_nlayer = type_nlayer
+        self.type_one_side = type_one_side
+        self.numb_aparam = numb_aparam
+        if self.numb_aparam > 0:
+            add_data_requirement('aparam', 3, atomic=True, must=True, high_prec=False)
+
+
+    def build (self, 
+               coord_ : tf.Tensor, 
+               atype_ : tf.Tensor,
+               natoms : tf.Tensor,
+               box_ : tf.Tensor, 
+               mesh : tf.Tensor,
+               input_dict : dict, 
+               reuse : bool = None,
+               suffix : str = ''
+    ) -> tf.Tensor:
+        """
+        Build the computational graph for the descriptor
+
+        Parameters
+        ----------
+        coord_
+                The coordinate of atoms
+        atype_
+                The type of atoms
+        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
+        mesh
+                For historical reasons, only the length of the Tensor matters.
+                if size of mesh == 6, pbc is assumed. 
+                if size of mesh == 0, no-pbc is assumed. 
+        input_dict
+                Dictionary for additional inputs
+        reuse
+                The weights in the networks should be reused when get the variable.
+        suffix
+                Name suffix to identify this descriptor
+
+        Returns
+        -------
+        descriptor
+                The output descriptor
+        """
+        nei_type = np.array([])
+        for ii in range(self.ntypes):
+            nei_type = np.append(nei_type, ii * np.ones(self.sel_a[ii]))
+        self.nei_type = tf.get_variable('t_nei_type', 
+                                        [self.nnei],
+                                        dtype = GLOBAL_TF_FLOAT_PRECISION,
+                                        trainable = False,
+                                        initializer = tf.constant_initializer(nei_type))
+        self.dout = DescrptSeA.build(self, coord_, atype_, natoms, box_, mesh, input_dict, suffix = suffix, reuse = reuse)
+        tf.summary.histogram('embedding_net_output', self.dout)
+
+        return self.dout
+
+
+    def _type_embed(self, 
+                    atype,
+                    ndim = 1,
+                    reuse = None, 
+                    suffix = '',
+                    trainable = True):
+        ebd_type = tf.cast(atype, self.filter_precision)
+        ebd_type = ebd_type / float(self.ntypes)
+        ebd_type = tf.reshape(ebd_type, [-1, ndim])
+        for ii in range(self.type_nlayer):
+            name = 'type_embed_layer_' + str(ii)
+            ebd_type = one_layer(ebd_type,
+                                 self.type_nchanl,
+                                 activation_fn = self.filter_activation_fn,
+                                 precision = self.filter_precision,
+                                 name = name, 
+                                 reuse = reuse,
+                                 seed = self.seed + ii,
+                                 trainable = trainable)
+        name = 'type_embed_layer_' + str(self.type_nlayer)
+        ebd_type = one_layer(ebd_type,
+                             self.type_nchanl,
+                             activation_fn = None,
+                             precision = self.filter_precision,
+                             name = name, 
+                             reuse = reuse,
+                             seed = self.seed + ii,
+                             trainable = trainable)
+        ebd_type = tf.reshape(ebd_type, [tf.shape(atype)[0], self.type_nchanl])
+        return ebd_type            
+
+
+    def _embedding_net(self, 
+                       inputs,
+                       natoms,
+                       filter_neuron,
+                       activation_fn=tf.nn.tanh, 
+                       stddev=1.0,
+                       bavg=0.0,
+                       name='linear', 
+                       reuse=None,
+                       seed=None, 
+                       trainable = True):
+        '''
+        inputs:  nf x na x (nei x 4)
+        outputs: nf x na x nei x output_size
+        '''
+        # natom x (nei x 4)
+        inputs = tf.reshape(inputs, [-1, self.ndescrpt])
+        shape = inputs.get_shape().as_list()
+        outputs_size = [1] + filter_neuron
+        with tf.variable_scope(name, reuse=reuse):
+            xyz_scatter_total = []
+            # with natom x (nei x 4)  
+            inputs_i = inputs
+            shape_i = inputs_i.get_shape().as_list()
+            # with (natom x nei) x 4  
+            inputs_reshape = tf.reshape(inputs_i, [-1, 4])
+            # with (natom x nei) x 1
+            xyz_scatter = tf.reshape(tf.slice(inputs_reshape, [0,0],[-1,1]),[-1,1])
+            # with (natom x nei) x out_size
+            xyz_scatter = embedding_net(xyz_scatter, 
+                                        self.filter_neuron, 
+                                        self.filter_precision, 
+                                        activation_fn = activation_fn, 
+                                        resnet_dt = self.filter_resnet_dt,
+                                        stddev = stddev,
+                                        bavg = bavg,
+                                        seed = seed,
+                                        trainable = trainable)
+            # natom x nei x out_size
+            xyz_scatter = tf.reshape(xyz_scatter, (-1, shape_i[1]//4, outputs_size[-1]))
+            xyz_scatter_total.append(xyz_scatter)
+        # natom x nei x outputs_size
+        xyz_scatter = tf.concat(xyz_scatter_total, axis=1)
+        # nf x natom x nei x outputs_size
+        xyz_scatter = tf.reshape(xyz_scatter, [tf.shape(inputs)[0], natoms[0], self.nnei, outputs_size[-1]])
+        return xyz_scatter
+
+    
+    def _type_embedding_net_two_sides(self, 
+                                      mat_g, 
+                                      atype,
+                                      natoms,
+                                      name = '',
+                                      reuse = None,
+                                      seed = None,
+                                      trainable = True):
+        outputs_size = self.filter_neuron[-1]
+        nframes = tf.shape(mat_g)[0]
+        # (nf x natom x nei) x (outputs_size x chnl x chnl)
+        mat_g = tf.reshape(mat_g, [nframes * natoms[0] * self.nnei, outputs_size])
+        mat_g = one_layer(mat_g, 
+                          outputs_size * self.type_nchanl * self.type_nchanl, 
+                          activation_fn = None,
+                          precision = self.filter_precision,
+                          name = name+'_amplify',
+                          reuse = reuse,
+                          seed = self.seed,
+                          trainable = trainable)        
+        # nf x natom x nei x outputs_size x chnl x chnl
+        mat_g = tf.reshape(mat_g, [nframes, natoms[0], self.nnei, outputs_size, self.type_nchanl, self.type_nchanl])
+        # nf x natom x outputs_size x chnl x nei x chnl
+        mat_g = tf.transpose(mat_g, perm = [0, 1, 3, 4, 2, 5])
+        # nf x natom x outputs_size x chnl x (nei x chnl)
+        mat_g = tf.reshape(mat_g, [nframes, natoms[0], outputs_size, self.type_nchanl, self.nnei * self.type_nchanl])
+        
+        # nei x nchnl
+        ebd_nei_type = self._type_embed(self.nei_type, 
+                                        reuse = reuse,
+                                        trainable = True,
+                                        suffix = '')
+        # (nei x nchnl)
+        ebd_nei_type = tf.reshape(ebd_nei_type, [self.nnei * self.type_nchanl])
+        # (nframes x natom) x nchnl
+        ebd_atm_type = self._type_embed(atype,
+                                        reuse = True,
+                                        trainable = True,
+                                        suffix = '')    
+        # (nframes x natom x nchnl)
+        ebd_atm_type = tf.reshape(ebd_atm_type, [nframes * natoms[0] * self.type_nchanl])
+
+        # nf x natom x outputs_size x chnl x (nei x chnl)
+        mat_g = tf.multiply(mat_g, ebd_nei_type)
+        # nf x natom x outputs_size x chnl x nei x chnl
+        mat_g = tf.reshape(mat_g, [nframes, natoms[0], outputs_size, self.type_nchanl, self.nnei, self.type_nchanl])
+        # nf x natom x outputs_size x chnl x nei 
+        mat_g = tf.reduce_mean(mat_g, axis = 5)
+        # outputs_size x nei x nf x natom x chnl
+        mat_g = tf.transpose(mat_g, perm = [2, 4, 0, 1, 3])
+        # outputs_size x nei x (nf x natom x chnl)
+        mat_g = tf.reshape(mat_g, [outputs_size, self.nnei, nframes * natoms[0] * self.type_nchanl])
+        # outputs_size x nei x (nf x natom x chnl)
+        mat_g = tf.multiply(mat_g, ebd_atm_type)
+        # outputs_size x nei x nf x natom x chnl
+        mat_g = tf.reshape(mat_g, [outputs_size, self.nnei, nframes, natoms[0], self.type_nchanl])
+        # outputs_size x nei x nf x natom
+        mat_g = tf.reduce_mean(mat_g, axis = 4)
+        # nf x natom x nei x outputs_size
+        mat_g = tf.transpose(mat_g, perm = [2, 3, 1, 0])        
+        # (nf x natom) x nei x outputs_size
+        mat_g = tf.reshape(mat_g, [nframes * natoms[0], self.nnei, outputs_size])
+        return mat_g
+
+
+    def _type_embedding_net_one_side(self, 
+                                     mat_g, 
+                                     atype,
+                                     natoms,
+                                     name = '',
+                                     reuse = None,
+                                     seed = None,
+                                     trainable = True):
+        outputs_size = self.filter_neuron[-1]
+        nframes = tf.shape(mat_g)[0]
+        # (nf x natom x nei) x (outputs_size x chnl x chnl)
+        mat_g = tf.reshape(mat_g, [nframes * natoms[0] * self.nnei, outputs_size])
+        mat_g = one_layer(mat_g, 
+                          outputs_size * self.type_nchanl, 
+                          activation_fn = None,
+                          precision = self.filter_precision,
+                          name = name+'_amplify',
+                          reuse = reuse,
+                          seed = self.seed,
+                          trainable = trainable)        
+        # nf x natom x nei x outputs_size x chnl
+        mat_g = tf.reshape(mat_g, [nframes, natoms[0], self.nnei, outputs_size, self.type_nchanl])
+        # nf x natom x outputs_size x nei x chnl
+        mat_g = tf.transpose(mat_g, perm = [0, 1, 3, 2, 4])
+        # nf x natom x outputs_size x (nei x chnl)
+        mat_g = tf.reshape(mat_g, [nframes, natoms[0], outputs_size, self.nnei * self.type_nchanl])
+
+        # nei x nchnl
+        ebd_nei_type = self._type_embed(self.nei_type, 
+                                        reuse = reuse,
+                                        trainable = True,
+                                        suffix = '')
+        # (nei x nchnl)
+        ebd_nei_type = tf.reshape(ebd_nei_type, [self.nnei * self.type_nchanl])
+
+        # nf x natom x outputs_size x (nei x chnl)
+        mat_g = tf.multiply(mat_g, ebd_nei_type)
+        # nf x natom x outputs_size x nei x chnl
+        mat_g = tf.reshape(mat_g, [nframes, natoms[0], outputs_size, self.nnei, self.type_nchanl])
+        # nf x natom x outputs_size x nei 
+        mat_g = tf.reduce_mean(mat_g, axis = 4)
+        # nf x natom x nei x outputs_size
+        mat_g = tf.transpose(mat_g, perm = [0, 1, 3, 2])
+        # (nf x natom) x nei x outputs_size
+        mat_g = tf.reshape(mat_g, [nframes * natoms[0], self.nnei, outputs_size])
+        return mat_g
+
+
+    def _type_embedding_net_one_side_aparam(self, 
+                                            mat_g, 
+                                            atype,
+                                            natoms,
+                                            aparam,
+                                            name = '',
+                                            reuse = None,
+                                            seed = None,
+                                            trainable = True):
+        outputs_size = self.filter_neuron[-1]
+        nframes = tf.shape(mat_g)[0]
+        # (nf x natom x nei) x (outputs_size x chnl x chnl)
+        mat_g = tf.reshape(mat_g, [nframes * natoms[0] * self.nnei, outputs_size])
+        mat_g = one_layer(mat_g, 
+                          outputs_size * self.type_nchanl, 
+                          activation_fn = None,
+                          precision = self.filter_precision,
+                          name = name+'_amplify',
+                          reuse = reuse,
+                          seed = self.seed,
+                          trainable = trainable)        
+        # nf x natom x nei x outputs_size x chnl
+        mat_g = tf.reshape(mat_g, [nframes, natoms[0], self.nnei, outputs_size, self.type_nchanl])
+        # outputs_size x nf x natom x nei x chnl
+        mat_g = tf.transpose(mat_g, perm = [3, 0, 1, 2, 4])
+        # outputs_size x (nf x natom x nei x chnl)
+        mat_g = tf.reshape(mat_g, [outputs_size, nframes * natoms[0] * self.nnei * self.type_nchanl])        
+        # nf x natom x nnei        
+        embed_type = tf.tile(tf.reshape(self.nei_type, [1, self.nnei]),
+                             [nframes * natoms[0], 1])
+        # (nf x natom x nnei) x 1
+        embed_type = tf.reshape(embed_type, [nframes * natoms[0] * self.nnei, 1])        
+        # nf x (natom x naparam)
+        aparam = tf.reshape(aparam, [nframes, -1])
+        # nf x natom x nnei x naparam        
+        embed_aparam = op_module.map_aparam(aparam, self.nlist, natoms, n_a_sel = self.nnei_a, n_r_sel = self.nnei_r)
+        # (nf x natom x nnei) x naparam
+        embed_aparam = tf.reshape(embed_aparam, [nframes * natoms[0] * self.nnei, self.numb_aparam])
+        # (nf x natom x nnei) x (naparam+1)
+        embed_input = tf.concat((embed_type, embed_aparam), axis = 1)
+        
+        # (nf x natom x nnei) x nchnl
+        ebd_nei_type = self._type_embed(embed_input, 
+                                        ndim = self.numb_aparam + 1,
+                                        reuse = reuse,
+                                        trainable = True,
+                                        suffix = '')
+        # (nf x natom x nei x nchnl)
+        ebd_nei_type = tf.reshape(ebd_nei_type, [nframes * natoms[0] * self.nnei * self.type_nchanl])
+
+        # outputs_size x (nf x natom x nei x chnl)
+        mat_g = tf.multiply(mat_g, ebd_nei_type)
+        # outputs_size x nf x natom x nei x chnl
+        mat_g = tf.reshape(mat_g, [outputs_size, nframes, natoms[0], self.nnei, self.type_nchanl])
+        # outputs_size x nf x natom x nei 
+        mat_g = tf.reduce_mean(mat_g, axis = 4)
+        # nf x natom x nei x outputs_size
+        mat_g = tf.transpose(mat_g, perm = [1, 2, 3, 0])
+        # (nf x natom) x nei x outputs_size
+        mat_g = tf.reshape(mat_g, [nframes * natoms[0], self.nnei, outputs_size])
+        return mat_g
+
+
+    def _pass_filter(self, 
+                     inputs,
+                     atype,
+                     natoms,
+                     input_dict,
+                     reuse = None,
+                     suffix = '', 
+                     trainable = True) :
+        # nf x na x ndescrpt
+        # nf x na x (nnei x 4)
+        inputs = tf.reshape(inputs, [-1, natoms[0], self.ndescrpt])
+        layer, qmat = self._ebd_filter(tf.cast(inputs, self.filter_precision), 
+                                       atype,
+                                       natoms,
+                                       input_dict,
+                                       name='filter_type_all'+suffix, 
+                                       reuse=reuse, 
+                                       seed = self.seed, 
+                                       trainable = trainable, 
+                                       activation_fn = self.filter_activation_fn)
+        output      = tf.reshape(layer, [tf.shape(inputs)[0], natoms[0], self.get_dim_out()])
+        output_qmat = tf.reshape(qmat,  [tf.shape(inputs)[0], natoms[0], self.get_dim_rot_mat_1() * 3])
+        return output, output_qmat
+
+
+    def _ebd_filter(self, 
+                    inputs, 
+                    atype,
+                    natoms,
+                    input_dict,
+                    activation_fn=tf.nn.tanh, 
+                    stddev=1.0,
+                    bavg=0.0,
+                    name='linear', 
+                    reuse=None,
+                    seed=None, 
+                    trainable = True):
+        outputs_size = self.filter_neuron[-1]
+        outputs_size_2 = self.n_axis_neuron
+        # nf x natom x (nei x 4)
+        nframes = tf.shape(inputs)[0]
+        shape = tf.reshape(inputs, [-1, self.ndescrpt]).get_shape().as_list()
+        
+        # nf x natom x nei x outputs_size        
+        mat_g = self._embedding_net(inputs,
+                                    natoms,
+                                    self.filter_neuron,
+                                    activation_fn = activation_fn, 
+                                    stddev = stddev,
+                                    bavg = bavg,
+                                    name = name, 
+                                    reuse = reuse,
+                                    seed = seed,
+                                    trainable = trainable)
+        # nf x natom x nei x outputs_size        
+        mat_g = tf.reshape(mat_g, [nframes, natoms[0], self.nnei, outputs_size])
+        
+        # (nf x natom) x nei x outputs_size
+        if self.type_one_side:
+            if self.numb_aparam > 0:
+                aparam = input_dict['aparam']
+                xyz_scatter \
+                    = self._type_embedding_net_one_side_aparam(mat_g, 
+                                                               atype,
+                                                               natoms, 
+                                                               aparam,
+                                                               name = name,
+                                                               reuse = reuse, 
+                                                               seed = seed,
+                                                               trainable = trainable)
+            else:
+                xyz_scatter \
+                    = self._type_embedding_net_one_side(mat_g, 
+                                                        atype,
+                                                        natoms, 
+                                                        name = name,
+                                                        reuse = reuse, 
+                                                        seed = seed,
+                                                        trainable = trainable)
+        else:
+            xyz_scatter \
+                = self._type_embedding_net_two_sides(mat_g, 
+                                                     atype,
+                                                     natoms, 
+                                                     name = name,
+                                                     reuse = reuse, 
+                                                     seed = seed,
+                                                     trainable = trainable)
+        
+        # natom x nei x 4
+        inputs_reshape = tf.reshape(inputs, [-1, shape[1]//4, 4])
+        # natom x 4 x outputs_size
+        xyz_scatter_1 = tf.matmul(inputs_reshape, xyz_scatter, transpose_a = True)
+        xyz_scatter_1 = xyz_scatter_1 * (4.0 / shape[1])
+        # natom x 4 x outputs_size_2
+        xyz_scatter_2 = tf.slice(xyz_scatter_1, [0,0,0],[-1,-1,outputs_size_2])
+        # # natom x 3 x outputs_size_2
+        # qmat = tf.slice(xyz_scatter_2, [0,1,0], [-1, 3, -1])
+        # natom x 3 x outputs_size_1
+        qmat = tf.slice(xyz_scatter_1, [0,1,0], [-1, 3, -1])
+        # natom x outputs_size_2 x 3
+        qmat = tf.transpose(qmat, perm = [0, 2, 1])
+        # natom x outputs_size x outputs_size_2
+        result = tf.matmul(xyz_scatter_1, xyz_scatter_2, transpose_a = True)
+        # natom x (outputs_size x outputs_size_2)
+        result = tf.reshape(result, [-1, outputs_size_2 * outputs_size])
+
+        return result, qmat
+

_skbuild/linux-x86_64-3.6/cmake-install/deepmd/descriptor/se_a_ef.py

+import numpy as np
+from typing import Tuple, List
+
+from deepmd.env import tf
+from deepmd.common import add_data_requirement
+from deepmd.utils.sess import run_sess
+from deepmd.env import GLOBAL_TF_FLOAT_PRECISION
+from deepmd.env import GLOBAL_NP_FLOAT_PRECISION
+from deepmd.env import op_module
+from deepmd.env import default_tf_session_config
+from .se_a import DescrptSeA
+from .descriptor import Descriptor
+
+@Descriptor.register("se_a_ef")
+class DescrptSeAEf (Descriptor):
+    """
+
+    Parameters
+    ----------
+    rcut
+            The cut-off radius
+    rcut_smth
+            From where the environment matrix should be smoothed
+    sel : list[str]
+            sel[i] specifies the maxmum number of type i atoms in the cut-off radius
+    neuron : list[int]
+            Number of neurons in each hidden layers of the embedding net
+    axis_neuron
+            Number of the axis neuron (number of columns of the sub-matrix of the embedding matrix)
+    resnet_dt
+            Time-step `dt` in the resnet construction:
+            y = x + dt * \phi (Wx + b)
+    trainable
+            If the weights of embedding net are trainable.
+    seed
+            Random seed for initializing the network parameters.
+    type_one_side
+            Try to build N_types embedding nets. Otherwise, building N_types^2 embedding nets
+    exclude_types : List[List[int]]
+            The excluded pairs of types which have no interaction with each other.
+            For example, `[[0, 1]]` means no interaction between type 0 and type 1.
+    set_davg_zero
+            Set the shift of embedding net input to zero.
+    activation_function
+            The activation function in the embedding net. Supported options are |ACTIVATION_FN|
+    precision
+            The precision of the embedding net parameters. Supported options are |PRECISION|
+    uniform_seed
+            Only for the purpose of backward compatibility, retrieves the old behavior of using the random seed
+    """
+    def __init__(self,
+                 rcut: float,
+                 rcut_smth: float,
+                 sel: List[str],
+                 neuron: List[int] = [24,48,96],
+                 axis_neuron: int = 8,
+                 resnet_dt: bool = False,
+                 trainable: bool = True,
+                 seed: int = None,
+                 type_one_side: bool = True,
+                 exclude_types: List[List[int]] = [],
+                 set_davg_zero: bool = False,
+                 activation_function: str = 'tanh',
+                 precision: str = 'default',
+                 uniform_seed = False
+    ) -> None:        
+        """
+        Constructor
+        """
+        self.descrpt_para = DescrptSeAEfLower(
+            op_module.descrpt_se_a_ef_para, 
+            rcut,
+            rcut_smth,
+            sel,
+            neuron,
+            axis_neuron,
+            resnet_dt,
+            trainable,
+            seed,
+            type_one_side,
+            exclude_types,
+            set_davg_zero,
+            activation_function,
+            precision,
+            uniform_seed,
+        )
+        self.descrpt_vert = DescrptSeAEfLower(
+            op_module.descrpt_se_a_ef_vert,
+            rcut,
+            rcut_smth,
+            sel,
+            neuron,
+            axis_neuron,
+            resnet_dt,
+            trainable,
+            seed,
+            type_one_side,
+            exclude_types,
+            set_davg_zero,
+            activation_function,
+            precision,
+            uniform_seed,
+        )
+        
+    def get_rcut (self) -> float:
+        """
+        Returns the cut-off radisu
+        """
+        return self.descrpt_vert.rcut_r
+
+    def get_ntypes (self) -> int:
+        """
+        Returns the number of atom types
+        """
+        return self.descrpt_vert.ntypes
+
+    def get_dim_out (self) -> int:
+        """
+        Returns the output dimension of this descriptor
+        """
+        return self.descrpt_vert.get_dim_out() + self.descrpt_para.get_dim_out()
+
+    def get_dim_rot_mat_1 (self) -> int:
+        """
+        Returns the first dimension of the rotation matrix. The rotation is of shape dim_1 x 3
+        """
+        return self.descrpt_vert.filter_neuron[-1]
+
+    def get_rot_mat(self) -> tf.Tensor:
+        """
+        Get rotational matrix
+        """
+        return self.qmat
+
+
+    def get_nlist (self)  -> Tuple[tf.Tensor, tf.Tensor, List[int], List[int]]:
+        """
+        Returns
+        -------
+        nlist
+                Neighbor list
+        rij
+                The relative distance between the neighbor and the center atom.
+        sel_a
+                The number of neighbors with full information
+        sel_r
+                The number of neighbors with only radial information
+        """
+        return \
+            self.descrpt_vert.nlist, \
+            self.descrpt_vert.rij, \
+            self.descrpt_vert.sel_a, \
+            self.descrpt_vert.sel_r
+
+    def compute_input_stats (self,
+                             data_coord : list, 
+                             data_box : list, 
+                             data_atype : list, 
+                             natoms_vec : list,
+                             mesh : list, 
+                             input_dict : dict
+    ) -> None :
+        """
+        Compute the statisitcs (avg and std) of the training data. The input will be normalized by the statistics.
+        
+        Parameters
+        ----------
+        data_coord
+                The coordinates. Can be generated by deepmd.model.make_stat_input
+        data_box
+                The box. Can be generated by deepmd.model.make_stat_input
+        data_atype
+                The atom types. Can be generated by deepmd.model.make_stat_input
+        natoms_vec
+                The vector for the number of atoms of the system and different types of atoms. Can be generated by deepmd.model.make_stat_input
+        mesh
+                The mesh for neighbor searching. Can be generated by deepmd.model.make_stat_input
+        input_dict
+                Dictionary for additional input
+        """
+        self.descrpt_vert.compute_input_stats(data_coord, data_box, data_atype, natoms_vec, mesh, input_dict)
+        self.descrpt_para.compute_input_stats(data_coord, data_box, data_atype, natoms_vec, mesh, input_dict)
+
+    def build (self, 
+               coord_ : tf.Tensor, 
+               atype_ : tf.Tensor,
+               natoms : tf.Tensor,
+               box_ : tf.Tensor, 
+               mesh : tf.Tensor,
+               input_dict : dict, 
+               reuse : bool = None,
+               suffix : str = ''
+    ) -> tf.Tensor:
+        """
+        Build the computational graph for the descriptor
+
+        Parameters
+        ----------
+        coord_
+                The coordinate of atoms
+        atype_
+                The type of atoms
+        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
+        mesh
+                For historical reasons, only the length of the Tensor matters.
+                if size of mesh == 6, pbc is assumed. 
+                if size of mesh == 0, no-pbc is assumed. 
+        input_dict
+                Dictionary for additional inputs. Should have 'efield'.
+        reuse
+                The weights in the networks should be reused when get the variable.
+        suffix
+                Name suffix to identify this descriptor
+
+        Returns
+        -------
+        descriptor
+                The output descriptor
+        """
+        self.dout_vert = self.descrpt_vert.build(coord_, atype_, natoms, box_, mesh, input_dict)
+        self.dout_para = self.descrpt_para.build(coord_, atype_, natoms, box_, mesh, input_dict, reuse = True)
+        coord = tf.reshape(coord_, [-1, natoms[1] * 3])
+        nframes = tf.shape(coord)[0]
+        self.dout_vert = tf.reshape(self.dout_vert, [nframes * natoms[0], self.descrpt_vert.get_dim_out()])
+        self.dout_para = tf.reshape(self.dout_para, [nframes * natoms[0], self.descrpt_para.get_dim_out()])
+        self.dout = tf.concat([self.dout_vert, self.dout_para], axis = 1)
+        self.dout = tf.reshape(self.dout, [nframes, natoms[0], self.get_dim_out()])
+        self.qmat = self.descrpt_vert.qmat + self.descrpt_para.qmat
+
+        tf.summary.histogram('embedding_net_output', self.dout)
+
+        return self.dout
+
+    def prod_force_virial(self, 
+                          atom_ener : tf.Tensor, 
+                          natoms : tf.Tensor
+    ) -> Tuple[tf.Tensor, tf.Tensor, tf.Tensor]:
+        """
+        Compute force and virial
+
+        Parameters
+        ----------
+        atom_ener
+                The atomic energy
+        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
+
+        Returns
+        -------
+        force
+                The force on atoms
+        virial
+                The total virial
+        atom_virial
+                The atomic virial
+        """
+        f_vert, v_vert, av_vert \
+            = self.descrpt_vert.prod_force_virial(atom_ener, natoms)
+        f_para, v_para, av_para \
+            = self.descrpt_para.prod_force_virial(atom_ener, natoms)
+        force = f_vert + f_para
+        virial = v_vert + v_para
+        atom_vir = av_vert + av_para
+        return force, virial, atom_vir
+
+
+class DescrptSeAEfLower (DescrptSeA):
+    """
+    Helper class for implementing DescrptSeAEf
+    """
+    def __init__ (self, 
+                  op,
+                  rcut: float,
+                  rcut_smth: float,
+                  sel: List[str],
+                  neuron: List[int] = [24,48,96],
+                  axis_neuron: int = 8,
+                  resnet_dt: bool = False,
+                  trainable: bool = True,
+                  seed: int = None,
+                  type_one_side: bool = True,
+                  exclude_types: List[List[int]] = [],
+                  set_davg_zero: bool = False,
+                  activation_function: str = 'tanh',
+                  precision: str = 'default',
+                  uniform_seed : bool = False,
+    ) -> None:
+        DescrptSeA.__init__(
+            self, 
+            rcut,
+            rcut_smth,
+            sel,
+            neuron,
+            axis_neuron,
+            resnet_dt,
+            trainable,
+            seed,
+            type_one_side,
+            exclude_types,
+            set_davg_zero,
+            activation_function,
+            precision,
+            uniform_seed
+        )
+        # DescrptSeA.__init__(self, **jdata)
+        # args = ClassArg()\
+        #        .add('sel',      list,   must = True) \
+        #        .add('rcut',     float,  default = 6.0) \
+        #        .add('rcut_smth',float,  default = 5.5) \
+        #        .add('neuron',   list,   default = [10, 20, 40]) \
+        #        .add('axis_neuron', int, default = 4, alias = 'n_axis_neuron') \
+        #        .add('resnet_dt',bool,   default = False) \
+        #        .add('trainable',bool,   default = True) \
+        #        .add('seed',     int) 
+        # class_data = args.parse(jdata)
+        # self.sel_a = class_data['sel']
+        # self.rcut_r = class_data['rcut']
+        # self.rcut_r_smth = class_data['rcut_smth']
+        # self.filter_neuron = class_data['neuron']
+        # self.n_axis_neuron = class_data['axis_neuron']
+        # self.filter_resnet_dt = class_data['resnet_dt']
+        # self.seed = class_data['seed']
+        # self.trainable = class_data['trainable']
+        self.sel_a = sel
+        self.rcut_r = rcut
+        self.rcut_r_smth = rcut_smth
+        self.filter_neuron = neuron
+        self.n_axis_neuron = axis_neuron
+        self.filter_resnet_dt = resnet_dt
+        self.seed = seed
+        self.trainable = trainable
+        self.op = op
+
+        # descrpt config
+        self.sel_r = [ 0 for ii in range(len(self.sel_a)) ]
+        self.ntypes = len(self.sel_a)
+        assert(self.ntypes == len(self.sel_r))
+        self.rcut_a = -1
+        # numb of neighbors and numb of descrptors
+        self.nnei_a = np.cumsum(self.sel_a)[-1]
+        self.nnei_r = np.cumsum(self.sel_r)[-1]
+        self.nnei = self.nnei_a + self.nnei_r
+        self.ndescrpt_a = self.nnei_a * 4
+        self.ndescrpt_r = self.nnei_r * 1
+        self.ndescrpt = self.ndescrpt_a + self.ndescrpt_r
+        self.useBN = False
+        self.dstd = None
+        self.davg = None
+
+        add_data_requirement('efield', 3, atomic=True, must=True, high_prec=False)
+
+        self.place_holders = {}
+        avg_zero = np.zeros([self.ntypes,self.ndescrpt]).astype(GLOBAL_NP_FLOAT_PRECISION)
+        std_ones = np.ones ([self.ntypes,self.ndescrpt]).astype(GLOBAL_NP_FLOAT_PRECISION)
+        sub_graph = tf.Graph()
+        with sub_graph.as_default():
+            name_pfx = 'd_sea_ef_'
+            for ii in ['coord', 'box']:
+                self.place_holders[ii] = tf.placeholder(GLOBAL_NP_FLOAT_PRECISION, [None, None], name = name_pfx+'t_'+ii)
+            self.place_holders['type'] = tf.placeholder(tf.int32, [None, None], name=name_pfx+'t_type')
+            self.place_holders['natoms_vec'] = tf.placeholder(tf.int32, [self.ntypes+2], name=name_pfx+'t_natoms')
+            self.place_holders['default_mesh'] = tf.placeholder(tf.int32, [None], name=name_pfx+'t_mesh')
+            self.place_holders['efield'] = tf.placeholder(GLOBAL_NP_FLOAT_PRECISION, [None, None], name=name_pfx+'t_efield')
+            self.stat_descrpt, descrpt_deriv, rij, nlist \
+                = self.op(self.place_holders['coord'],
+                          self.place_holders['type'],
+                          self.place_holders['natoms_vec'],
+                          self.place_holders['box'],
+                          self.place_holders['default_mesh'],
+                          self.place_holders['efield'],
+                          tf.constant(avg_zero),
+                          tf.constant(std_ones),
+                          rcut_a = self.rcut_a,
+                          rcut_r = self.rcut_r,
+                          rcut_r_smth = self.rcut_r_smth,
+                          sel_a = self.sel_a,
+                          sel_r = self.sel_r)
+        self.sub_sess = tf.Session(graph = sub_graph, config=default_tf_session_config)
+
+
+
+    def compute_input_stats (self,
+                             data_coord, 
+                             data_box, 
+                             data_atype, 
+                             natoms_vec,
+                             mesh, 
+                             input_dict) :
+        data_efield = input_dict['efield']
+        all_davg = []
+        all_dstd = []
+        if True:
+            sumr = []
+            suma = []
+            sumn = []
+            sumr2 = []
+            suma2 = []
+            for cc,bb,tt,nn,mm,ee in zip(data_coord,data_box,data_atype,natoms_vec,mesh,data_efield) :
+                sysr,sysr2,sysa,sysa2,sysn \
+                    = self._compute_dstats_sys_smth(cc,bb,tt,nn,mm,ee)
+                sumr.append(sysr)
+                suma.append(sysa)
+                sumn.append(sysn)
+                sumr2.append(sysr2)
+                suma2.append(sysa2)
+            sumr = np.sum(sumr, axis = 0)
+            suma = np.sum(suma, axis = 0)
+            sumn = np.sum(sumn, axis = 0)
+            sumr2 = np.sum(sumr2, axis = 0)
+            suma2 = np.sum(suma2, axis = 0)
+            for type_i in range(self.ntypes) :
+                davgunit = [sumr[type_i]/sumn[type_i], 0, 0, 0]
+                dstdunit = [self._compute_std(sumr2[type_i], sumr[type_i], sumn[type_i]), 
+                            self._compute_std(suma2[type_i], suma[type_i], sumn[type_i]), 
+                            self._compute_std(suma2[type_i], suma[type_i], sumn[type_i]), 
+                            self._compute_std(suma2[type_i], suma[type_i], sumn[type_i])
+                            ]
+                davg = np.tile(davgunit, self.ndescrpt // 4)
+                dstd = np.tile(dstdunit, self.ndescrpt // 4)
+                all_davg.append(davg)
+                all_dstd.append(dstd)
+
+        self.davg = np.array(all_davg)
+        self.dstd = np.array(all_dstd)
+
+    def _normalize_3d(self, a):
+        na = tf.norm(a, axis = 1)
+        na = tf.tile(tf.reshape(na, [-1,1]), tf.constant([1, 3]))
+        return tf.divide(a, na)
+
+    def build (self, 
+               coord_, 
+               atype_,
+               natoms,
+               box_, 
+               mesh,
+               input_dict,
+               suffix = '', 
+               reuse = None):
+        efield = input_dict['efield']
+        davg = self.davg
+        dstd = self.dstd
+        with tf.variable_scope('descrpt_attr' + suffix, reuse = reuse) :
+            if davg is None:
+                davg = np.zeros([self.ntypes, self.ndescrpt]) 
+            if dstd is None:
+                dstd = np.ones ([self.ntypes, self.ndescrpt])
+            t_rcut = tf.constant(np.max([self.rcut_r, self.rcut_a]), 
+                                 name = 'rcut', 
+                                 dtype = GLOBAL_TF_FLOAT_PRECISION)
+            t_ntypes = tf.constant(self.ntypes, 
+                                   name = 'ntypes', 
+                                   dtype = tf.int32)
+            t_ndescrpt = tf.constant(self.ndescrpt, 
+                                     name = 'ndescrpt', 
+                                     dtype = tf.int32)            
+            t_sel = tf.constant(self.sel_a, 
+                                name = 'sel', 
+                                dtype = tf.int32)            
+            self.t_avg = tf.get_variable('t_avg', 
+                                         davg.shape, 
+                                         dtype = GLOBAL_TF_FLOAT_PRECISION,
+                                         trainable = False,
+                                         initializer = tf.constant_initializer(davg))
+            self.t_std = tf.get_variable('t_std', 
+                                         dstd.shape, 
+                                         dtype = GLOBAL_TF_FLOAT_PRECISION,
+                                         trainable = False,
+                                         initializer = tf.constant_initializer(dstd))
+
+        coord = tf.reshape (coord_, [-1, natoms[1] * 3])
+        box   = tf.reshape (box_, [-1, 9])
+        atype = tf.reshape (atype_, [-1, natoms[1]])
+        efield = tf.reshape(efield, [-1, 3])
+        efield = self._normalize_3d(efield)
+        efield = tf.reshape(efield, [-1, natoms[0] * 3])
+
+        self.descrpt, self.descrpt_deriv, self.rij, self.nlist \
+            = self.op (coord,
+                       atype,
+                       natoms,
+                       box,
+                       mesh,
+                       efield,
+                       self.t_avg,
+                       self.t_std,
+                       rcut_a = self.rcut_a,
+                       rcut_r = self.rcut_r,
+                       rcut_r_smth = self.rcut_r_smth,
+                       sel_a = self.sel_a,
+                       sel_r = self.sel_r)
+
+        self.descrpt_reshape = tf.reshape(self.descrpt, [-1, self.ndescrpt])
+        self.descrpt_reshape = tf.identity(self.descrpt_reshape, name = 'o_rmat')
+        self.descrpt_deriv = tf.identity(self.descrpt_deriv, name = 'o_rmat_deriv')
+        self.rij = tf.identity(self.rij, name = 'o_rij')
+        self.nlist = tf.identity(self.nlist, name = 'o_nlist')
+
+        # only used when tensorboard was set as true
+        tf.summary.histogram('descrpt', self.descrpt)
+        tf.summary.histogram('rij', self.rij)
+        tf.summary.histogram('nlist', self.nlist)
+
+        self.dout, self.qmat = self._pass_filter(self.descrpt_reshape, atype, natoms, input_dict, suffix = suffix, reuse = reuse, trainable = self.trainable)
+        tf.summary.histogram('embedding_net_output', self.dout)
+
+        return self.dout
+
+
+
+    def _compute_dstats_sys_smth (self,
+                                  data_coord, 
+                                  data_box, 
+                                  data_atype,                             
+                                  natoms_vec,
+                                  mesh,
+                                  data_efield) :
+        dd_all \
+            = run_sess(self.sub_sess, self.stat_descrpt, 
+                                feed_dict = {
+                                    self.place_holders['coord']: data_coord,
+                                    self.place_holders['type']: data_atype,
+                                    self.place_holders['natoms_vec']: natoms_vec,
+                                    self.place_holders['box']: data_box,
+                                    self.place_holders['default_mesh']: mesh,
+                                    self.place_holders['efield']: data_efield,
+                                })
+        natoms = natoms_vec
+        dd_all = np.reshape(dd_all, [-1, self.ndescrpt * natoms[0]])
+        start_index = 0
+        sysr = []
+        sysa = []
+        sysn = []
+        sysr2 = []
+        sysa2 = []
+        for type_i in range(self.ntypes):
+            end_index = start_index + self.ndescrpt * natoms[2+type_i]
+            dd = dd_all[:, start_index:end_index]
+            dd = np.reshape(dd, [-1, self.ndescrpt])
+            start_index = end_index        
+            # compute
+            dd = np.reshape (dd, [-1, 4])
+            ddr = dd[:,:1]
+            dda = dd[:,1:]
+            sumr = np.sum(ddr)
+            suma = np.sum(dda) / 3.
+            sumn = dd.shape[0]
+            sumr2 = np.sum(np.multiply(ddr, ddr))
+            suma2 = np.sum(np.multiply(dda, dda)) / 3.
+            sysr.append(sumr)
+            sysa.append(suma)
+            sysn.append(sumn)
+            sysr2.append(sumr2)
+            sysa2.append(suma2)
+        return sysr, sysr2, sysa, sysa2, sysn
+
+

_skbuild/linux-x86_64-3.6/cmake-install/deepmd/descriptor/se_atten.py

+import math
+import numpy as np
+from typing import Tuple, List, Dict, Any
+from packaging.version import Version
+
+from deepmd.env import tf
+from deepmd.common import get_activation_func, get_precision, cast_precision
+from deepmd.env import GLOBAL_TF_FLOAT_PRECISION
+from deepmd.env import TF_VERSION
+from deepmd.env import GLOBAL_NP_FLOAT_PRECISION
+from deepmd.env import op_module
+from deepmd.env import default_tf_session_config
+from deepmd.utils.network import one_layer, embedding_net, embedding_net_rand_seed_shift
+from deepmd.utils.tabulate import DPTabulate
+from deepmd.utils.type_embed import embed_atom_type
+from deepmd.utils.sess import run_sess
+from deepmd.utils.graph import load_graph_def, get_tensor_by_name_from_graph, get_tensor_by_name
+from deepmd.utils.graph import get_attention_layer_variables_from_graph_def
+from deepmd.utils.errors import GraphWithoutTensorError
+from .descriptor import Descriptor
+from .se_a import DescrptSeA
+
+
+@Descriptor.register("se_atten")
+class DescrptSeAtten(DescrptSeA):
+    """
+    Parameters
+    ----------
+    rcut
+            The cut-off radius :math:`r_c`
+    rcut_smth
+            From where the environment matrix should be smoothed :math:`r_s`
+    sel : list[str]
+            sel[i] specifies the maxmum number of type i atoms in the cut-off radius
+    neuron : list[int]
+            Number of neurons in each hidden layers of the embedding net :math:`\mathcal{N}`
+    axis_neuron
+            Number of the axis neuron :math:`M_2` (number of columns of the sub-matrix of the embedding matrix)
+    resnet_dt
+            Time-step `dt` in the resnet construction:
+            y = x + dt * \phi (Wx + b)
+    trainable
+            If the weights of embedding net are trainable.
+    seed
+            Random seed for initializing the network parameters.
+    type_one_side
+            Try to build N_types embedding nets. Otherwise, building N_types^2 embedding nets
+    exclude_types : List[List[int]]
+            The excluded pairs of types which have no interaction with each other.
+            For example, `[[0, 1]]` means no interaction between type 0 and type 1.
+    set_davg_zero
+            Set the shift of embedding net input to zero.
+    activation_function
+            The activation function in the embedding net. Supported options are |ACTIVATION_FN|
+    precision
+            The precision of the embedding net parameters. Supported options are |PRECISION|
+    uniform_seed
+            Only for the purpose of backward compatibility, retrieves the old behavior of using the random seed
+    attn
+            The length of hidden vector during scale-dot attention computation.
+    attn_layer
+            The number of layers in attention mechanism.
+    attn_dotr
+            Whether to dot the relative coordinates on the attention weights as a gated scheme.
+    attn_mask
+            Whether to mask the diagonal in the attention weights.
+    """
+
+    def __init__(self,
+                 rcut: float,
+                 rcut_smth: float,
+                 sel: int,
+                 ntypes: int,
+                 neuron: List[int] = [24, 48, 96],
+                 axis_neuron: int = 8,
+                 resnet_dt: bool = False,
+                 trainable: bool = True,
+                 seed: int = None,
+                 type_one_side: bool = True,
+                 exclude_types: List[List[int]] = [],
+                 set_davg_zero: bool = False,
+                 activation_function: str = 'tanh',
+                 precision: str = 'default',
+                 uniform_seed: bool = False,
+                 attn: int = 128,
+                 attn_layer: int = 2,
+                 attn_dotr: bool = True,
+                 attn_mask: bool = False
+                 ) -> None:
+        DescrptSeA.__init__(self,
+                            rcut,
+                            rcut_smth,
+                            [sel],
+                            neuron=neuron,
+                            axis_neuron=axis_neuron,
+                            resnet_dt=resnet_dt,
+                            trainable=trainable,
+                            seed=seed,
+                            type_one_side=type_one_side,
+                            exclude_types=exclude_types,
+                            set_davg_zero=set_davg_zero,
+                            activation_function=activation_function,
+                            precision=precision,
+                            uniform_seed=uniform_seed
+                            )
+        """
+        Constructor
+        """
+        assert (Version(TF_VERSION) > Version('2')), "se_atten only support tensorflow version 2.0 or higher."
+        self.ntypes = ntypes
+        self.att_n = attn
+        self.attn_layer = attn_layer
+        self.attn_mask = attn_mask
+        self.attn_dotr = attn_dotr
+
+        # descrpt config
+        self.sel_all_a = [sel]
+        self.sel_all_r = [0]
+        avg_zero = np.zeros([self.ntypes, self.ndescrpt]).astype(GLOBAL_NP_FLOAT_PRECISION)
+        std_ones = np.ones([self.ntypes, self.ndescrpt]).astype(GLOBAL_NP_FLOAT_PRECISION)
+        self.beta = np.zeros([self.attn_layer, self.filter_neuron[-1]]).astype(GLOBAL_NP_FLOAT_PRECISION)
+        self.gamma = np.ones([self.attn_layer, self.filter_neuron[-1]]).astype(GLOBAL_NP_FLOAT_PRECISION)
+        self.attention_layer_variables = None
+        sub_graph = tf.Graph()
+        with sub_graph.as_default():
+            name_pfx = 'd_sea_'
+            for ii in ['coord', 'box']:
+                self.place_holders[ii] = tf.placeholder(GLOBAL_NP_FLOAT_PRECISION, [None, None],
+                                                        name=name_pfx + 't_' + ii)
+            self.place_holders['type'] = tf.placeholder(tf.int32, [None, None], name=name_pfx + 't_type')
+            self.place_holders['natoms_vec'] = tf.placeholder(tf.int32, [self.ntypes + 2], name=name_pfx + 't_natoms')
+            self.place_holders['default_mesh'] = tf.placeholder(tf.int32, [None], name=name_pfx + 't_mesh')
+            self.stat_descrpt, self.descrpt_deriv_t, self.rij_t, self.nlist_t, self.nei_type_vec_t, self.nmask_t \
+                = op_module.prod_env_mat_a_mix(self.place_holders['coord'],
+                                               self.place_holders['type'],
+                                               self.place_holders['natoms_vec'],
+                                               self.place_holders['box'],
+                                               self.place_holders['default_mesh'],
+                                               tf.constant(avg_zero),
+                                               tf.constant(std_ones),
+                                               rcut_a=self.rcut_a,
+                                               rcut_r=self.rcut_r,
+                                               rcut_r_smth=self.rcut_r_smth,
+                                               sel_a=self.sel_all_a,
+                                               sel_r=self.sel_all_r)
+        self.sub_sess = tf.Session(graph=sub_graph, config=default_tf_session_config)
+
+    def compute_input_stats(self,
+                            data_coord: list,
+                            data_box: list,
+                            data_atype: list,
+                            natoms_vec: list,
+                            mesh: list,
+                            input_dict: dict,
+                            mixed_type: bool = False,
+                            real_natoms_vec: list = None
+                            ) -> None:
+        """
+        Compute the statisitcs (avg and std) of the training data. The input will be normalized by the statistics.
+        
+        Parameters
+        ----------
+        data_coord
+                The coordinates. Can be generated by deepmd.model.make_stat_input
+        data_box
+                The box. Can be generated by deepmd.model.make_stat_input
+        data_atype
+                The atom types. Can be generated by deepmd.model.make_stat_input
+        natoms_vec
+                The vector for the number of atoms of the system and different types of atoms.
+                If mixed_type is True, this para is blank. See real_natoms_vec.
+        mesh
+                The mesh for neighbor searching. Can be generated by deepmd.model.make_stat_input
+        input_dict
+                Dictionary for additional input
+        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.
+        real_natoms_vec
+                If mixed_type is True, it takes in the real natoms_vec for each frame.
+        """
+        all_davg = []
+        all_dstd = []
+        if True:
+            sumr = []
+            suma = []
+            sumn = []
+            sumr2 = []
+            suma2 = []
+            if mixed_type:
+                sys_num = 0
+                for cc, bb, tt, nn, mm, r_n in zip(data_coord, data_box, data_atype, natoms_vec, mesh, real_natoms_vec):
+                    sysr, sysr2, sysa, sysa2, sysn \
+                        = self._compute_dstats_sys_smth(cc, bb, tt, nn, mm, mixed_type, r_n)
+                    sys_num += 1
+                    sumr.append(sysr)
+                    suma.append(sysa)
+                    sumn.append(sysn)
+                    sumr2.append(sysr2)
+                    suma2.append(sysa2)
+            else:
+                for cc, bb, tt, nn, mm in zip(data_coord, data_box, data_atype, natoms_vec, mesh):
+                    sysr, sysr2, sysa, sysa2, sysn \
+                        = self._compute_dstats_sys_smth(cc, bb, tt, nn, mm)
+                    sumr.append(sysr)
+                    suma.append(sysa)
+                    sumn.append(sysn)
+                    sumr2.append(sysr2)
+                    suma2.append(sysa2)
+            sumr = np.sum(sumr, axis=0)
+            suma = np.sum(suma, axis=0)
+            sumn = np.sum(sumn, axis=0)
+            sumr2 = np.sum(sumr2, axis=0)
+            suma2 = np.sum(suma2, axis=0)
+            for type_i in range(self.ntypes):
+                davgunit = [sumr[type_i] / (sumn[type_i] + 1e-15), 0, 0, 0]
+                dstdunit = [self._compute_std(sumr2[type_i], sumr[type_i], sumn[type_i]),
+                            self._compute_std(suma2[type_i], suma[type_i], sumn[type_i]),
+                            self._compute_std(suma2[type_i], suma[type_i], sumn[type_i]),
+                            self._compute_std(suma2[type_i], suma[type_i], sumn[type_i])
+                            ]
+                davg = np.tile(davgunit, self.ndescrpt // 4)
+                dstd = np.tile(dstdunit, self.ndescrpt // 4)
+                all_davg.append(davg)
+                all_dstd.append(dstd)
+
+        if not self.set_davg_zero:
+            self.davg = np.array(all_davg)
+        self.dstd = np.array(all_dstd)
+
+    def build(self,
+              coord_: tf.Tensor,
+              atype_: tf.Tensor,
+              natoms: tf.Tensor,
+              box_: tf.Tensor,
+              mesh: tf.Tensor,
+              input_dict: dict,
+              reuse: bool = None,
+              suffix: str = ''
+              ) -> tf.Tensor:
+        """
+        Build the computational graph for the descriptor
+
+        Parameters
+        ----------
+        coord_
+                The coordinate of atoms
+        atype_
+                The type of atoms
+        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
+        mesh
+                For historical reasons, only the length of the Tensor matters.
+                if size of mesh == 6, pbc is assumed. 
+                if size of mesh == 0, no-pbc is assumed. 
+        input_dict
+                Dictionary for additional inputs
+        reuse
+                The weights in the networks should be reused when get the variable.
+        suffix
+                Name suffix to identify this descriptor
+
+        Returns
+        -------
+        descriptor
+                The output descriptor
+        """
+        davg = self.davg
+        dstd = self.dstd
+        with tf.variable_scope('descrpt_attr' + suffix, reuse=reuse):
+            if davg is None:
+                davg = np.zeros([self.ntypes, self.ndescrpt])
+            if dstd is None:
+                dstd = np.ones([self.ntypes, self.ndescrpt])
+            t_rcut = tf.constant(np.max([self.rcut_r, self.rcut_a]),
+                                 name='rcut',
+                                 dtype=GLOBAL_TF_FLOAT_PRECISION)
+            t_ntypes = tf.constant(self.ntypes,
+                                   name='ntypes',
+                                   dtype=tf.int32)
+            t_ndescrpt = tf.constant(self.ndescrpt,
+                                     name='ndescrpt',
+                                     dtype=tf.int32)
+            t_sel = tf.constant(self.sel_a,
+                                name='sel',
+                                dtype=tf.int32)
+            t_original_sel = tf.constant(self.original_sel if self.original_sel is not None else self.sel_a,
+                                         name='original_sel',
+                                         dtype=tf.int32)
+            self.t_avg = tf.get_variable('t_avg',
+                                         davg.shape,
+                                         dtype=GLOBAL_TF_FLOAT_PRECISION,
+                                         trainable=False,
+                                         initializer=tf.constant_initializer(davg))
+            self.t_std = tf.get_variable('t_std',
+                                         dstd.shape,
+                                         dtype=GLOBAL_TF_FLOAT_PRECISION,
+                                         trainable=False,
+                                         initializer=tf.constant_initializer(dstd))
+
+        with tf.control_dependencies([t_sel, t_original_sel]):
+            coord = tf.reshape(coord_, [-1, natoms[1] * 3])
+        box = tf.reshape(box_, [-1, 9])
+        atype = tf.reshape(atype_, [-1, natoms[1]])
+        self.attn_weight = [None for i in range(self.attn_layer)]
+        self.angular_weight = [None for i in range(self.attn_layer)]
+        self.attn_weight_final = [None for i in range(self.attn_layer)]
+
+        self.descrpt, self.descrpt_deriv, self.rij, self.nlist, self.nei_type_vec, self.nmask \
+            = op_module.prod_env_mat_a_mix(coord,
+                                           atype,
+                                           natoms,
+                                           box,
+                                           mesh,
+                                           self.t_avg,
+                                           self.t_std,
+                                           rcut_a=self.rcut_a,
+                                           rcut_r=self.rcut_r,
+                                           rcut_r_smth=self.rcut_r_smth,
+                                           sel_a=self.sel_all_a,
+                                           sel_r=self.sel_all_r)
+        self.nei_type_vec = tf.reshape(self.nei_type_vec, [-1])
+        self.nmask = tf.cast(tf.reshape(self.nmask, [-1, 1, self.sel_all_a[0]]), GLOBAL_TF_FLOAT_PRECISION)
+        self.negative_mask = -(2 << 32) * (1.0 - self.nmask)
+        # only used when tensorboard was set as true
+        tf.summary.histogram('descrpt', self.descrpt)
+        tf.summary.histogram('rij', self.rij)
+        tf.summary.histogram('nlist', self.nlist)
+
+        self.descrpt_reshape = tf.reshape(self.descrpt, [-1, self.ndescrpt])
+        self.atype_nloc = tf.reshape(tf.slice(atype, [0, 0], [-1, natoms[0]]),
+                                     [-1])  ## lammps will have error without this
+        self._identity_tensors(suffix=suffix)
+
+        self.dout, self.qmat = self._pass_filter(self.descrpt_reshape,
+                                                 self.atype_nloc,
+                                                 natoms,
+                                                 input_dict,
+                                                 suffix=suffix,
+                                                 reuse=reuse,
+                                                 trainable=self.trainable)
+
+        # only used when tensorboard was set as true
+        tf.summary.histogram('embedding_net_output', self.dout)
+        return self.dout
+
+    def _pass_filter(self,
+                     inputs,
+                     atype,
+                     natoms,
+                     input_dict,
+                     reuse=None,
+                     suffix='',
+                     trainable=True):
+        assert (input_dict is not None and input_dict.get('type_embedding', None) is not None), \
+            'se_atten desctiptor must use type_embedding'
+        type_embedding = input_dict.get('type_embedding', None)
+        inputs = tf.reshape(inputs, [-1, natoms[0], self.ndescrpt])
+        output = []
+        output_qmat = []
+        inputs_i = inputs
+        inputs_i = tf.reshape(inputs_i, [-1, self.ndescrpt])
+        type_i = -1
+        layer, qmat = self._filter(inputs_i, type_i, natoms, name='filter_type_all' + suffix, suffix=suffix,
+                                   reuse=reuse, trainable=trainable, activation_fn=self.filter_activation_fn,
+                                   type_embedding=type_embedding, atype=atype)
+        layer = tf.reshape(layer, [tf.shape(inputs)[0], natoms[0], self.get_dim_out()])
+        qmat = tf.reshape(qmat, [tf.shape(inputs)[0], natoms[0], self.get_dim_rot_mat_1() * 3])
+        output.append(layer)
+        output_qmat.append(qmat)
+        output = tf.concat(output, axis=1)
+        output_qmat = tf.concat(output_qmat, axis=1)
+        return output, output_qmat
+
+    def _compute_dstats_sys_smth(self,
+                                 data_coord,
+                                 data_box,
+                                 data_atype,
+                                 natoms_vec,
+                                 mesh,
+                                 mixed_type=False,
+                                 real_natoms_vec=None):
+        dd_all, descrpt_deriv_t, rij_t, nlist_t, nei_type_vec_t, nmask_t \
+            = run_sess(self.sub_sess, [self.stat_descrpt, self.descrpt_deriv_t, self.rij_t, self.nlist_t, self.nei_type_vec_t, self.nmask_t],
+                       feed_dict={
+                           self.place_holders['coord']: data_coord,
+                           self.place_holders['type']: data_atype,
+                           self.place_holders['natoms_vec']: natoms_vec,
+                           self.place_holders['box']: data_box,
+                           self.place_holders['default_mesh']: mesh,
+                       })
+        if mixed_type:
+            nframes = dd_all.shape[0]
+            sysr = [0. for i in range(self.ntypes)]
+            sysa = [0. for i in range(self.ntypes)]
+            sysn = [0 for i in range(self.ntypes)]
+            sysr2 = [0. for i in range(self.ntypes)]
+            sysa2 = [0. for i in range(self.ntypes)]
+            for ff in range(nframes):
+                natoms = real_natoms_vec[ff]
+                dd_ff = np.reshape(dd_all[ff], [-1, self.ndescrpt * natoms[0]])
+                start_index = 0
+                for type_i in range(self.ntypes):
+                    end_index = start_index + self.ndescrpt * natoms[2 + type_i]  # center atom split
+                    dd = dd_ff[:, start_index:end_index]
+                    dd = np.reshape(dd, [-1, self.ndescrpt])  # nframes * typen_atoms , nnei * 4
+                    start_index = end_index
+                    # compute
+                    dd = np.reshape(dd, [-1, 4])  # nframes * typen_atoms * nnei, 4
+                    ddr = dd[:, :1]
+                    dda = dd[:, 1:]
+                    sumr = np.sum(ddr)
+                    suma = np.sum(dda) / 3.
+                    sumn = dd.shape[0]
+                    sumr2 = np.sum(np.multiply(ddr, ddr))
+                    suma2 = np.sum(np.multiply(dda, dda)) / 3.
+                    sysr[type_i] += sumr
+                    sysa[type_i] += suma
+                    sysn[type_i] += sumn
+                    sysr2[type_i] += sumr2
+                    sysa2[type_i] += suma2
+        else:
+            natoms = natoms_vec
+            dd_all = np.reshape(dd_all, [-1, self.ndescrpt * natoms[0]])
+            start_index = 0
+            sysr = []
+            sysa = []
+            sysn = []
+            sysr2 = []
+            sysa2 = []
+            for type_i in range(self.ntypes):
+                end_index = start_index + self.ndescrpt * natoms[2 + type_i]  # center atom split
+                dd = dd_all[:, start_index:end_index]
+                dd = np.reshape(dd, [-1, self.ndescrpt])  # nframes * typen_atoms , nnei * 4
+                start_index = end_index
+                # compute
+                dd = np.reshape(dd, [-1, 4])  # nframes * typen_atoms * nnei, 4
+                ddr = dd[:, :1]
+                dda = dd[:, 1:]
+                sumr = np.sum(ddr)
+                suma = np.sum(dda) / 3.
+                sumn = dd.shape[0]
+                sumr2 = np.sum(np.multiply(ddr, ddr))
+                suma2 = np.sum(np.multiply(dda, dda)) / 3.
+                sysr.append(sumr)
+                sysa.append(suma)
+                sysn.append(sumn)
+                sysr2.append(sumr2)
+                sysa2.append(suma2)
+        return sysr, sysr2, sysa, sysa2, sysn
+
+    def _lookup_type_embedding(
+            self,
+            xyz_scatter,
+            natype,
+            type_embedding,
+    ):
+        '''Concatenate `type_embedding` of neighbors and `xyz_scatter`.
+        If not self.type_one_side, concatenate `type_embedding` of center atoms as well.
+
+        Parameters
+        ----------
+        xyz_scatter:
+                shape is [nframes*natoms[0]*self.nnei, 1]
+        nframes:
+                shape is []
+        natoms:
+                shape is [1+1+self.ntypes]
+        type_embedding:
+                shape is [self.ntypes, Y] where Y=jdata['type_embedding']['neuron'][-1]
+
+        Returns
+        -------
+            embedding:
+                environment of each atom represented by embedding.
+        '''
+        te_out_dim = type_embedding.get_shape().as_list()[-1]
+        self.test_type_embedding = type_embedding
+        self.test_nei_embed = tf.nn.embedding_lookup(type_embedding,
+                                                     self.nei_type_vec)  # shape is [self.nnei, 1+te_out_dim]
+        # nei_embed = tf.tile(nei_embed, (nframes * natoms[0], 1))  # shape is [nframes*natoms[0]*self.nnei, te_out_dim]
+        nei_embed = tf.reshape(self.test_nei_embed, [-1, te_out_dim])
+        self.embedding_input = tf.concat([xyz_scatter, nei_embed],
+                                         1)  # shape is [nframes*natoms[0]*self.nnei, 1+te_out_dim]
+        if not self.type_one_side:
+            self.atm_embed = tf.nn.embedding_lookup(type_embedding, natype)  # shape is [nframes*natoms[0], te_out_dim]
+            self.atm_embed = tf.tile(self.atm_embed,
+                                     [1, self.nnei])  # shape is [nframes*natoms[0], self.nnei*te_out_dim]
+            self.atm_embed = tf.reshape(self.atm_embed,
+                                        [-1, te_out_dim])  # shape is [nframes*natoms[0]*self.nnei, te_out_dim]
+            self.embedding_input_2 = tf.concat([self.embedding_input, self.atm_embed],
+                                               1)  # shape is [nframes*natoms[0]*self.nnei, 1+te_out_dim+te_out_dim]
+            return self.embedding_input_2
+        return self.embedding_input
+
+    def _feedforward(self, input_xyz, d_in, d_mid):
+        residual = input_xyz
+        input_xyz = tf.nn.relu(one_layer(
+            input_xyz,
+            d_mid,
+            name='c_ffn1',
+            reuse=tf.AUTO_REUSE,
+            seed=self.seed,
+            activation_fn=None,
+            precision=self.filter_precision,
+            trainable=True,
+            uniform_seed=self.uniform_seed,
+            initial_variables=self.attention_layer_variables))
+        input_xyz = one_layer(
+            input_xyz,
+            d_in,
+            name='c_ffn2',
+            reuse=tf.AUTO_REUSE,
+            seed=self.seed,
+            activation_fn=None,
+            precision=self.filter_precision,
+            trainable=True,
+            uniform_seed=self.uniform_seed,
+            initial_variables=self.attention_layer_variables)
+        input_xyz += residual
+        input_xyz = tf.keras.layers.LayerNormalization()(input_xyz)
+        return input_xyz
+
+    def _scaled_dot_attn(self, Q, K, V, temperature, input_r, dotr=False, do_mask=False, layer=0, save_weights=True):
+        attn = tf.matmul(Q / temperature, K, transpose_b=True)
+        attn *= self.nmask
+        attn += self.negative_mask
+        attn = tf.nn.softmax(attn, axis=-1)
+        attn *= tf.reshape(self.nmask, [-1, attn.shape[-1], 1])
+        if save_weights:
+            self.attn_weight[layer] = attn[0]  # atom 0
+        if dotr:
+            angular_weight = tf.matmul(input_r, input_r, transpose_b=True)  # normalized
+            attn *= angular_weight
+            if save_weights:
+                self.angular_weight[layer] = angular_weight[0]  # atom 0
+                self.attn_weight_final[layer] = attn[0]  # atom 0
+        if do_mask:
+            nei = int(attn.shape[-1])
+            mask = tf.cast(tf.ones((nei, nei)) - tf.eye(nei), self.filter_precision)
+            attn *= mask
+        output = tf.matmul(attn, V)
+        return output
+
+    def _attention_layers(
+            self,
+            input_xyz,
+            layer_num,
+            shape_i,
+            outputs_size,
+            input_r,
+            dotr=False,
+            do_mask=False,
+            trainable=True,
+            suffix=''
+    ):
+        sd_k = tf.sqrt(tf.cast(1., dtype=self.filter_precision))
+        for i in range(layer_num):
+            name = 'attention_layer_{}{}'.format(i, suffix)
+            with tf.variable_scope(name, reuse=tf.AUTO_REUSE):
+                # input_xyz_in = tf.nn.l2_normalize(input_xyz, -1)
+                Q_c = one_layer(
+                    input_xyz,
+                    self.att_n,
+                    name='c_query',
+                    scope=name+'/',
+                    reuse=tf.AUTO_REUSE,
+                    seed=self.seed,
+                    activation_fn=None,
+                    precision=self.filter_precision,
+                    trainable=trainable,
+                    uniform_seed=self.uniform_seed,
+                    initial_variables=self.attention_layer_variables)
+                K_c = one_layer(
+                    input_xyz,
+                    self.att_n,
+                    name='c_key',
+                    scope=name+'/',
+                    reuse=tf.AUTO_REUSE,
+                    seed=self.seed,
+                    activation_fn=None,
+                    precision=self.filter_precision,
+                    trainable=trainable,
+                    uniform_seed=self.uniform_seed,
+                    initial_variables=self.attention_layer_variables)
+                V_c = one_layer(
+                    input_xyz,
+                    self.att_n,
+                    name='c_value',
+                    scope=name+'/',
+                    reuse=tf.AUTO_REUSE,
+                    seed=self.seed,
+                    activation_fn=None,
+                    precision=self.filter_precision,
+                    trainable=trainable,
+                    uniform_seed=self.uniform_seed,
+                    initial_variables=self.attention_layer_variables)
+                # # natom x nei_type_i x out_size
+                # xyz_scatter = tf.reshape(xyz_scatter, (-1, shape_i[1] // 4, outputs_size[-1]))
+                # natom x nei_type_i x att_n
+                Q_c = tf.nn.l2_normalize(tf.reshape(Q_c, (-1, shape_i[1] // 4, self.att_n)), -1)
+                K_c = tf.nn.l2_normalize(tf.reshape(K_c, (-1, shape_i[1] // 4, self.att_n)), -1)
+                V_c = tf.nn.l2_normalize(tf.reshape(V_c, (-1, shape_i[1] // 4, self.att_n)), -1)
+
+                input_att = self._scaled_dot_attn(Q_c, K_c, V_c, sd_k, input_r, dotr=dotr, do_mask=do_mask, layer=i)
+                input_att = tf.reshape(input_att, (-1, self.att_n))
+
+                # (natom x nei_type_i) x out_size
+                input_xyz += one_layer(
+                    input_att,
+                    outputs_size[-1],
+                    name='c_out',
+                    scope=name+'/',
+                    reuse=tf.AUTO_REUSE,
+                    seed=self.seed,
+                    activation_fn=None,
+                    precision=self.filter_precision,
+                    trainable=trainable,
+                    uniform_seed=self.uniform_seed,
+                    initial_variables=self.attention_layer_variables)
+                input_xyz = tf.keras.layers.LayerNormalization(beta_initializer=tf.constant_initializer(self.beta[i]),
+                                                gamma_initializer=tf.constant_initializer(self.gamma[i]))(input_xyz)
+                # input_xyz = self._feedforward(input_xyz, outputs_size[-1], self.att_n)
+        return input_xyz
+
+    def _filter_lower(
+            self,
+            type_i,
+            type_input,
+            start_index,
+            incrs_index,
+            inputs,
+            type_embedding=None,
+            atype=None,
+            is_exclude=False,
+            activation_fn=None,
+            bavg=0.0,
+            stddev=1.0,
+            trainable=True,
+            suffix='',
+            name='filter_',
+            reuse=None
+    ):
+        """
+        input env matrix, returns R.G
+        """
+        outputs_size = [1] + self.filter_neuron
+        # cut-out inputs
+        # with natom x (nei_type_i x 4)  
+        inputs_i = tf.slice(inputs,
+                            [0, start_index * 4],
+                            [-1, incrs_index * 4])
+        shape_i = inputs_i.get_shape().as_list()
+        natom = tf.shape(inputs_i)[0]
+        # with (natom x nei_type_i) x 4
+        inputs_reshape = tf.reshape(inputs_i, [-1, 4])
+        # with (natom x nei_type_i) x 1
+        xyz_scatter = tf.reshape(tf.slice(inputs_reshape, [0, 0], [-1, 1]), [-1, 1])
+        assert atype is not None, 'atype must exist!!'
+        type_embedding = tf.cast(type_embedding, self.filter_precision)
+        xyz_scatter = self._lookup_type_embedding(
+            xyz_scatter, atype, type_embedding)
+        if self.compress:
+            raise RuntimeError('compression of attention descriptor is not supported at the moment')
+        # natom x 4 x outputs_size
+        if (not is_exclude):
+            with tf.variable_scope(name, reuse=reuse):
+                # with (natom x nei_type_i) x out_size
+                xyz_scatter = embedding_net(
+                    xyz_scatter,
+                    self.filter_neuron,
+                    self.filter_precision,
+                    activation_fn=activation_fn,
+                    resnet_dt=self.filter_resnet_dt,
+                    name_suffix=suffix,
+                    stddev=stddev,
+                    bavg=bavg,
+                    seed=self.seed,
+                    trainable=trainable,
+                    uniform_seed=self.uniform_seed,
+                    initial_variables=self.embedding_net_variables,
+                    mixed_prec=self.mixed_prec)
+                if (not self.uniform_seed) and (self.seed is not None): self.seed += self.seed_shift
+            input_r = tf.slice(tf.reshape(inputs_i, (-1, shape_i[1] // 4, 4)), [0, 0, 1], [-1, -1, 3])
+            input_r = tf.nn.l2_normalize(input_r, -1)
+            # natom x nei_type_i x out_size
+            xyz_scatter_att = tf.reshape(
+                self._attention_layers(xyz_scatter, self.attn_layer, shape_i, outputs_size, input_r,
+                                       dotr=self.attn_dotr, do_mask=self.attn_mask, trainable=trainable, suffix=suffix),
+                (-1, shape_i[1] // 4, outputs_size[-1]))
+            # xyz_scatter = tf.reshape(xyz_scatter, (-1, shape_i[1] // 4, outputs_size[-1]))
+        else:
+            # we can safely return the final xyz_scatter filled with zero directly
+            return tf.cast(tf.fill((natom, 4, outputs_size[-1]), 0.), self.filter_precision)
+        # When using tf.reshape(inputs_i, [-1, shape_i[1]//4, 4]) below
+        # [588 24] -> [588 6 4] correct
+        # but if sel is zero
+        # [588 0] -> [147 0 4] incorrect; the correct one is [588 0 4]
+        # So we need to explicitly assign the shape to tf.shape(inputs_i)[0] instead of -1
+        return tf.matmul(tf.reshape(inputs_i, [natom, shape_i[1] // 4, 4]), xyz_scatter_att, transpose_a=True)
+
+    @cast_precision
+    def _filter(
+            self,
+            inputs,
+            type_input,
+            natoms,
+            type_embedding=None,
+            atype=None,
+            activation_fn=tf.nn.tanh,
+            stddev=1.0,
+            bavg=0.0,
+            suffix='',
+            name='linear',
+            reuse=None,
+            trainable=True):
+        nframes = tf.shape(tf.reshape(inputs, [-1, natoms[0], self.ndescrpt]))[0]
+        # natom x (nei x 4)
+        shape = inputs.get_shape().as_list()
+        outputs_size = [1] + self.filter_neuron
+        outputs_size_2 = self.n_axis_neuron
+        all_excluded = all([(type_input, type_i) in self.exclude_types for type_i in range(self.ntypes)])
+        if all_excluded:
+            # all types are excluded so result and qmat should be zeros
+            # we can safaly return a zero matrix...
+            # See also https://stackoverflow.com/a/34725458/9567349
+            # result: natom x outputs_size x outputs_size_2
+            # qmat: natom x outputs_size x 3
+            natom = tf.shape(inputs)[0]
+            result = tf.cast(tf.fill((natom, outputs_size_2, outputs_size[-1]), 0.), GLOBAL_TF_FLOAT_PRECISION)
+            qmat = tf.cast(tf.fill((natom, outputs_size[-1], 3), 0.), GLOBAL_TF_FLOAT_PRECISION)
+            return result, qmat
+
+        start_index = 0
+        type_i = 0
+        # natom x 4 x outputs_size
+        xyz_scatter_1 = self._filter_lower(
+            type_i, type_input,
+            start_index, np.cumsum(self.sel_a)[-1],
+            inputs,
+            type_embedding=type_embedding,
+            is_exclude=False,
+            activation_fn=activation_fn,
+            stddev=stddev,
+            bavg=bavg,
+            trainable=trainable,
+            suffix=suffix,
+            name=name,
+            reuse=reuse,
+            atype=atype)
+        # natom x nei x outputs_size
+        # xyz_scatter = tf.concat(xyz_scatter_total, axis=1)
+        # natom x nei x 4
+        # inputs_reshape = tf.reshape(inputs, [-1, shape[1]//4, 4])
+        # natom x 4 x outputs_size
+        # xyz_scatter_1 = tf.matmul(inputs_reshape, xyz_scatter, transpose_a = True)
+        if self.original_sel is None:
+            # shape[1] = nnei * 4
+            nnei = shape[1] / 4
+        else:
+            nnei = tf.cast(tf.Variable(np.sum(self.original_sel), dtype=tf.int32, trainable=False, name="nnei"),
+                           self.filter_precision)
+        xyz_scatter_1 = xyz_scatter_1 / nnei
+        # natom x 4 x outputs_size_2
+        xyz_scatter_2 = tf.slice(xyz_scatter_1, [0, 0, 0], [-1, -1, outputs_size_2])
+        # # natom x 3 x outputs_size_2
+        # qmat = tf.slice(xyz_scatter_2, [0,1,0], [-1, 3, -1])
+        # natom x 3 x outputs_size_1
+        qmat = tf.slice(xyz_scatter_1, [0, 1, 0], [-1, 3, -1])
+        # natom x outputs_size_1 x 3
+        qmat = tf.transpose(qmat, perm=[0, 2, 1])
+        # natom x outputs_size x outputs_size_2
+        result = tf.matmul(xyz_scatter_1, xyz_scatter_2, transpose_a=True)
+        # natom x (outputs_size x outputs_size_2)
+        result = tf.reshape(result, [-1, outputs_size_2 * outputs_size[-1]])
+
+        return result, qmat
+
+    def init_variables(self,
+                       graph: tf.Graph,
+                       graph_def: tf.GraphDef,
+                       suffix: str = "",
+                       ) -> None:
+        """
+        Init the embedding net variables with the given dict
+
+        Parameters
+        ----------
+        graph : tf.Graph
+            The input frozen model graph
+        graph_def : tf.GraphDef
+            The input frozen model graph_def
+        suffix : str, optional
+            The suffix of the scope
+        """
+        super().init_variables(graph=graph, graph_def=graph_def, suffix=suffix)
+        self.attention_layer_variables = get_attention_layer_variables_from_graph_def(graph_def, suffix=suffix)
+        if self.attn_layer > 0:
+            self.beta[0] = self.attention_layer_variables['attention_layer_0{}/layer_normalization/beta'.format(suffix)]
+            self.gamma[0] = self.attention_layer_variables['attention_layer_0{}/layer_normalization/gamma'.format(suffix)]
+            for i in range(1, self.attn_layer):
+                self.beta[i] = self.attention_layer_variables[
+                    'attention_layer_{}{}/layer_normalization_{}/beta'.format(i, suffix, i)]
+                self.gamma[i] = self.attention_layer_variables[
+                    'attention_layer_{}{}/layer_normalization_{}/gamma'.format(i, suffix, i)]

_skbuild/linux-x86_64-3.6/cmake-install/deepmd/descriptor/se_r.py

+import numpy as np
+from typing import Tuple, List
+
+from deepmd.env import tf
+from deepmd.common import get_activation_func, get_precision, cast_precision
+from deepmd.env import GLOBAL_TF_FLOAT_PRECISION
+from deepmd.env import GLOBAL_NP_FLOAT_PRECISION
+from deepmd.env import op_module
+from deepmd.env import default_tf_session_config
+from deepmd.utils.tabulate import DPTabulate
+from deepmd.utils.graph import load_graph_def, get_tensor_by_name_from_graph
+from deepmd.utils.network import embedding_net, embedding_net_rand_seed_shift
+from deepmd.utils.sess import run_sess
+from .descriptor import Descriptor
+from .se import DescrptSe
+
+
+@Descriptor.register("se_e2_r")
+@Descriptor.register("se_r")
+class DescrptSeR (DescrptSe):
+    """DeepPot-SE constructed from radial information of atomic configurations.
+    
+    The embedding takes the distance between atoms as input.
+
+    Parameters
+    ----------
+    rcut
+            The cut-off radius
+    rcut_smth
+            From where the environment matrix should be smoothed
+    sel : list[str]
+            sel[i] specifies the maxmum number of type i atoms in the cut-off radius
+    neuron : list[int]
+            Number of neurons in each hidden layers of the embedding net
+    resnet_dt
+            Time-step `dt` in the resnet construction:
+            y = x + dt * \phi (Wx + b)
+    trainable
+            If the weights of embedding net are trainable.
+    seed
+            Random seed for initializing the network parameters.
+    type_one_side
+            Try to build N_types embedding nets. Otherwise, building N_types^2 embedding nets
+    exclude_types : List[List[int]]
+            The excluded pairs of types which have no interaction with each other.
+            For example, `[[0, 1]]` means no interaction between type 0 and type 1.
+    activation_function
+            The activation function in the embedding net. Supported options are |ACTIVATION_FN|
+    precision
+            The precision of the embedding net parameters. Supported options are |PRECISION|
+    uniform_seed
+            Only for the purpose of backward compatibility, retrieves the old behavior of using the random seed
+    """
+    def __init__ (self, 
+                  rcut: float,
+                  rcut_smth: float,
+                  sel: List[str],
+                  neuron: List[int] = [24,48,96],
+                  resnet_dt: bool = False,
+                  trainable: bool = True,
+                  seed: int = None,
+                  type_one_side: bool = True,
+                  exclude_types: List[List[int]] = [],
+                  set_davg_zero: bool = False,
+                  activation_function: str = 'tanh',
+                  precision: str = 'default',
+                  uniform_seed: bool = False
+    ) -> None:
+        """
+        Constructor
+        """
+        # args = ClassArg()\
+        #        .add('sel',      list,   must = True) \
+        #        .add('rcut',     float,  default = 6.0) \
+        #        .add('rcut_smth',float,  default = 0.5) \
+        #        .add('neuron',   list,   default = [10, 20, 40]) \
+        #        .add('resnet_dt',bool,   default = False) \
+        #        .add('trainable',bool,   default = True) \
+        #        .add('seed',     int) \
+        #        .add('type_one_side', bool, default = False) \
+        #        .add('exclude_types', list, default = []) \
+        #        .add('set_davg_zero', bool, default = False) \
+        #        .add("activation_function", str, default = "tanh") \
+        #        .add("precision",           str, default = "default")
+        # class_data = args.parse(jdata)
+        if rcut < rcut_smth:
+            raise RuntimeError("rcut_smth (%f) should be no more than rcut (%f)!" % (rcut_smth, rcut))
+        self.sel_r = sel
+        self.rcut = rcut
+        self.rcut_smth = rcut_smth
+        self.filter_neuron = neuron
+        self.filter_resnet_dt = resnet_dt
+        self.seed = seed        
+        self.uniform_seed = uniform_seed
+        self.seed_shift = embedding_net_rand_seed_shift(self.filter_neuron)
+        self.trainable = trainable
+        self.filter_activation_fn = get_activation_func(activation_function) 
+        self.filter_precision = get_precision(precision)  
+        exclude_types = exclude_types
+        self.exclude_types = set()
+        for tt in exclude_types:
+            assert(len(tt) == 2)
+            self.exclude_types.add((tt[0], tt[1]))
+            self.exclude_types.add((tt[1], tt[0]))
+        self.set_davg_zero = set_davg_zero
+        self.type_one_side = type_one_side
+
+        # descrpt config
+        self.sel_a = [ 0 for ii in range(len(self.sel_r)) ]
+        self.ntypes = len(self.sel_r)
+        # numb of neighbors and numb of descrptors
+        self.nnei_a = np.cumsum(self.sel_a)[-1]
+        self.nnei_r = np.cumsum(self.sel_r)[-1]
+        self.nnei = self.nnei_a + self.nnei_r
+        self.ndescrpt_a = self.nnei_a * 4
+        self.ndescrpt_r = self.nnei_r * 1
+        self.ndescrpt = self.nnei_r
+        self.useBN = False
+        self.davg = None
+        self.dstd = None
+        self.compress=False
+        self.embedding_net_variables = None
+
+        self.place_holders = {}
+        avg_zero = np.zeros([self.ntypes,self.ndescrpt]).astype(GLOBAL_NP_FLOAT_PRECISION)
+        std_ones = np.ones ([self.ntypes,self.ndescrpt]).astype(GLOBAL_NP_FLOAT_PRECISION)
+        sub_graph = tf.Graph()
+        with sub_graph.as_default():
+            name_pfx = 'd_ser_'
+            for ii in ['coord', 'box']:
+                self.place_holders[ii] = tf.placeholder(GLOBAL_NP_FLOAT_PRECISION, [None, None], name = name_pfx+'t_'+ii)
+            self.place_holders['type'] = tf.placeholder(tf.int32, [None, None], name=name_pfx+'t_type')
+            self.place_holders['natoms_vec'] = tf.placeholder(tf.int32, [self.ntypes+2], name=name_pfx+'t_natoms')
+            self.place_holders['default_mesh'] = tf.placeholder(tf.int32, [None], name=name_pfx+'t_mesh')
+            self.stat_descrpt, descrpt_deriv, rij, nlist \
+                = op_module.prod_env_mat_r(self.place_holders['coord'],
+                                         self.place_holders['type'],
+                                         self.place_holders['natoms_vec'],
+                                         self.place_holders['box'],
+                                         self.place_holders['default_mesh'],
+                                         tf.constant(avg_zero),
+                                         tf.constant(std_ones),
+                                         rcut = self.rcut,
+                                         rcut_smth = self.rcut_smth,
+                                         sel = self.sel_r)
+            self.sub_sess = tf.Session(graph = sub_graph, config=default_tf_session_config)
+
+
+    def get_rcut (self) :
+        """
+        Returns the cut-off radisu
+        """
+        return self.rcut
+
+    def get_ntypes (self) :
+        """
+        Returns the number of atom types
+        """
+        return self.ntypes
+
+    def get_dim_out (self) :
+        """
+        Returns the output dimension of this descriptor
+        """
+        return self.filter_neuron[-1]
+
+    def get_nlist (self) :
+        """
+        Returns
+        -------
+        nlist
+                Neighbor list
+        rij
+                The relative distance between the neighbor and the center atom.
+        sel_a
+                The number of neighbors with full information
+        sel_r
+                The number of neighbors with only radial information
+        """
+        return self.nlist, self.rij, self.sel_a, self.sel_r
+
+    def compute_input_stats (self,
+                             data_coord, 
+                             data_box, 
+                             data_atype, 
+                             natoms_vec,
+                             mesh, 
+                             input_dict) :    
+        """
+        Compute the statisitcs (avg and std) of the training data. The input will be normalized by the statistics.
+        
+        Parameters
+        ----------
+        data_coord
+                The coordinates. Can be generated by deepmd.model.make_stat_input
+        data_box
+                The box. Can be generated by deepmd.model.make_stat_input
+        data_atype
+                The atom types. Can be generated by deepmd.model.make_stat_input
+        natoms_vec
+                The vector for the number of atoms of the system and different types of atoms. Can be generated by deepmd.model.make_stat_input
+        mesh
+                The mesh for neighbor searching. Can be generated by deepmd.model.make_stat_input
+        input_dict
+                Dictionary for additional input
+        """
+        all_davg = []
+        all_dstd = []
+        sumr = []
+        sumn = []
+        sumr2 = []
+        for cc,bb,tt,nn,mm in zip(data_coord,data_box,data_atype,natoms_vec,mesh) :
+            sysr,sysr2,sysn \
+                = self._compute_dstats_sys_se_r(cc,bb,tt,nn,mm)
+            sumr.append(sysr)
+            sumn.append(sysn)
+            sumr2.append(sysr2)
+        sumr = np.sum(sumr, axis = 0)
+        sumn = np.sum(sumn, axis = 0)
+        sumr2 = np.sum(sumr2, axis = 0)
+        for type_i in range(self.ntypes) :
+            davgunit = [sumr[type_i]/sumn[type_i]]
+            dstdunit = [self._compute_std(sumr2[type_i], sumr[type_i], sumn[type_i])]
+            davg = np.tile(davgunit, self.ndescrpt // 1)
+            dstd = np.tile(dstdunit, self.ndescrpt // 1)
+            all_davg.append(davg)
+            all_dstd.append(dstd)
+
+        if not self.set_davg_zero:
+            self.davg = np.array(all_davg)
+        self.dstd = np.array(all_dstd)
+
+    def enable_compression(self,
+                           min_nbor_dist : float,
+                           model_file : str = 'frozon_model.pb',
+                           table_extrapolate : float = 5,
+                           table_stride_1 : float = 0.01,
+                           table_stride_2 : float = 0.1,
+                           check_frequency : int = -1,
+                           suffix : str = "",
+    ) -> None:
+        """
+        Reveive the statisitcs (distance, max_nbor_size and env_mat_range) of the training data.
+
+        Parameters
+        ----------
+        min_nbor_dist
+                The nearest distance between atoms
+        model_file
+                The original frozen model, which will be compressed by the program
+        table_extrapolate
+                The scale of model extrapolation
+        table_stride_1
+                The uniform stride of the first table
+        table_stride_2
+                The uniform stride of the second table
+        check_frequency
+                The overflow check frequency
+        suffix : str, optional
+                The suffix of the scope
+        """
+        assert (
+            not self.filter_resnet_dt
+        ), "Model compression error: descriptor resnet_dt must be false!"
+
+        for ii in range(len(self.filter_neuron) - 1):
+            if self.filter_neuron[ii] * 2 != self.filter_neuron[ii + 1]:
+                raise NotImplementedError(
+                    "Model Compression error: descriptor neuron [%s] is not supported by model compression! "
+                    "The size of the next layer of the neural network must be twice the size of the previous layer." 
+                    % ','.join([str(item) for item in self.filter_neuron])
+                )
+
+        self.compress = True
+        self.table = DPTabulate(
+            self, self.filter_neuron, model_file, activation_fn = self.filter_activation_fn, suffix=suffix)
+        self.table_config = [table_extrapolate, table_stride_1, table_stride_2, check_frequency]
+        self.lower, self.upper \
+            = self.table.build(min_nbor_dist, 
+                               table_extrapolate, 
+                               table_stride_1, 
+                               table_stride_2)
+        
+        graph, _ = load_graph_def(model_file)
+        self.davg = get_tensor_by_name_from_graph(graph, 'descrpt_attr%s/t_avg' % suffix)
+        self.dstd = get_tensor_by_name_from_graph(graph, 'descrpt_attr%s/t_std' % suffix)
+
+    def build (self, 
+               coord_ : tf.Tensor, 
+               atype_ : tf.Tensor,
+               natoms : tf.Tensor,
+               box_ : tf.Tensor, 
+               mesh : tf.Tensor,
+               input_dict : dict, 
+               reuse : bool = None,
+               suffix : str = ''
+    ) -> tf.Tensor:
+        """
+        Build the computational graph for the descriptor
+
+        Parameters
+        ----------
+        coord_
+                The coordinate of atoms
+        atype_
+                The type of atoms
+        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
+        mesh
+                For historical reasons, only the length of the Tensor matters.
+                if size of mesh == 6, pbc is assumed. 
+                if size of mesh == 0, no-pbc is assumed. 
+        input_dict
+                Dictionary for additional inputs
+        reuse
+                The weights in the networks should be reused when get the variable.
+        suffix
+                Name suffix to identify this descriptor
+
+        Returns
+        -------
+        descriptor
+                The output descriptor
+        """
+        davg = self.davg
+        dstd = self.dstd
+        with tf.variable_scope('descrpt_attr' + suffix, reuse = reuse) :
+            if davg is None:
+                davg = np.zeros([self.ntypes, self.ndescrpt]) 
+            if dstd is None:
+                dstd = np.ones ([self.ntypes, self.ndescrpt])
+            t_rcut = tf.constant(self.rcut, 
+                                 name = 'rcut', 
+                                 dtype = GLOBAL_TF_FLOAT_PRECISION)
+            t_ntypes = tf.constant(self.ntypes, 
+                                   name = 'ntypes', 
+                                   dtype = tf.int32)
+            t_ndescrpt = tf.constant(self.ndescrpt, 
+                                     name = 'ndescrpt', 
+                                     dtype = tf.int32)            
+            t_sel = tf.constant(self.sel_a, 
+                                name = 'sel', 
+                                dtype = tf.int32)            
+            self.t_avg = tf.get_variable('t_avg', 
+                                         davg.shape, 
+                                         dtype = GLOBAL_TF_FLOAT_PRECISION,
+                                         trainable = False,
+                                         initializer = tf.constant_initializer(davg))
+            self.t_std = tf.get_variable('t_std', 
+                                         dstd.shape, 
+                                         dtype = GLOBAL_TF_FLOAT_PRECISION,
+                                         trainable = False,
+                                         initializer = tf.constant_initializer(dstd))
+
+        coord = tf.reshape (coord_, [-1, natoms[1] * 3])
+        box   = tf.reshape (box_, [-1, 9])
+        atype = tf.reshape (atype_, [-1, natoms[1]])
+
+        self.descrpt, self.descrpt_deriv, self.rij, self.nlist \
+            = op_module.prod_env_mat_r(coord,
+                                      atype,
+                                      natoms,
+                                      box,
+                                      mesh,
+                                      self.t_avg,
+                                      self.t_std,
+                                      rcut = self.rcut,
+                                      rcut_smth = self.rcut_smth,
+                                      sel = self.sel_r)
+
+        self.descrpt_reshape = tf.reshape(self.descrpt, [-1, self.ndescrpt])
+        self._identity_tensors(suffix=suffix)
+
+        # only used when tensorboard was set as true
+        tf.summary.histogram('descrpt', self.descrpt)
+        tf.summary.histogram('rij', self.rij)
+        tf.summary.histogram('nlist', self.nlist)
+
+        self.dout = self._pass_filter(self.descrpt_reshape, natoms, suffix = suffix, reuse = reuse, trainable = self.trainable)
+        tf.summary.histogram('embedding_net_output', self.dout)
+
+        return self.dout
+
+
+    def prod_force_virial(self, 
+                          atom_ener : tf.Tensor, 
+                          natoms : tf.Tensor
+    ) -> Tuple[tf.Tensor, tf.Tensor, tf.Tensor]:
+        """
+        Compute force and virial
+
+        Parameters
+        ----------
+        atom_ener
+                The atomic energy
+        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
+
+        Returns
+        -------
+        force
+                The force on atoms
+        virial
+                The total virial
+        atom_virial
+                The atomic virial
+        """
+        [net_deriv] = tf.gradients (atom_ener, self.descrpt_reshape)
+        tf.summary.histogram('net_derivative', net_deriv)
+        net_deriv_reshape = tf.reshape (net_deriv, [np.cast['int64'](-1), natoms[0] * np.cast['int64'](self.ndescrpt)])        
+        force \
+            = op_module.prod_force_se_r (net_deriv_reshape,
+                                         self.descrpt_deriv,
+                                         self.nlist,
+                                         natoms)
+        virial, atom_virial \
+            = op_module.prod_virial_se_r (net_deriv_reshape,
+                                          self.descrpt_deriv,
+                                          self.rij,
+                                          self.nlist,
+                                          natoms)
+        tf.summary.histogram('force', force)
+        tf.summary.histogram('virial', virial)
+        tf.summary.histogram('atom_virial', atom_virial)
+
+        return force, virial, atom_virial
+    
+
+    def _pass_filter(self, 
+                     inputs,
+                     natoms,
+                     reuse = None,
+                     suffix = '', 
+                     trainable = True) :
+        start_index = 0
+        inputs = tf.reshape(inputs, [-1, natoms[0], self.ndescrpt])
+        output = []
+        if not (self.type_one_side and len(self.exclude_types) == 0):
+            for type_i in range(self.ntypes):
+                inputs_i = tf.slice (inputs,
+                                     [ 0, start_index, 0],
+                                     [-1, natoms[2+type_i], -1] )
+                inputs_i = tf.reshape(inputs_i, [-1, self.ndescrpt])
+                if self.type_one_side:
+                    # reuse NN parameters for all types to support type_one_side along with exclude_types
+                    reuse = tf.AUTO_REUSE
+                    filter_name = 'filter_type_all'+suffix
+                else:
+                    filter_name = 'filter_type_'+str(type_i)+suffix
+                layer = self._filter_r(inputs_i, type_i, name=filter_name, natoms=natoms, reuse=reuse, trainable = trainable, activation_fn = self.filter_activation_fn)
+                layer = tf.reshape(layer, [tf.shape(inputs)[0], natoms[2+type_i], self.get_dim_out()])
+                output.append(layer)
+                start_index += natoms[2+type_i]
+        else :
+            inputs_i = inputs
+            inputs_i = tf.reshape(inputs_i, [-1, self.ndescrpt])
+            type_i = -1
+            layer = self._filter_r(inputs_i, type_i, name='filter_type_all'+suffix, natoms=natoms, reuse=reuse, trainable = trainable, activation_fn = self.filter_activation_fn)
+            layer = tf.reshape(layer, [tf.shape(inputs)[0], natoms[0], self.get_dim_out()])
+            output.append(layer)
+        output = tf.concat(output, axis = 1)
+        return output
+
+    def _compute_dstats_sys_se_r (self,
+                                  data_coord, 
+                                  data_box, 
+                                  data_atype,                             
+                                  natoms_vec,
+                                  mesh) :    
+        dd_all \
+            = run_sess(self.sub_sess, self.stat_descrpt, 
+                                feed_dict = {
+                                    self.place_holders['coord']: data_coord,
+                                    self.place_holders['type']: data_atype,
+                                    self.place_holders['natoms_vec']: natoms_vec,
+                                    self.place_holders['box']: data_box,
+                                    self.place_holders['default_mesh']: mesh,
+                                })
+        natoms = natoms_vec
+        dd_all = np.reshape(dd_all, [-1, self.ndescrpt * natoms[0]])
+        start_index = 0
+        sysr = []
+        sysn = []
+        sysr2 = []
+        for type_i in range(self.ntypes):
+            end_index = start_index + self.ndescrpt * natoms[2+type_i]
+            dd = dd_all[:, start_index:end_index]
+            dd = np.reshape(dd, [-1, self.ndescrpt])
+            start_index = end_index        
+            # compute
+            dd = np.reshape (dd, [-1, 1])
+            ddr = dd[:,:1]
+            sumr = np.sum(ddr)
+            sumn = dd.shape[0]
+            sumr2 = np.sum(np.multiply(ddr, ddr))
+            sysr.append(sumr)
+            sysn.append(sumn)
+            sysr2.append(sumr2)
+        return sysr, sysr2, sysn
+
+
+    def _compute_std (self,sumv2, sumv, sumn) :
+        val = np.sqrt(sumv2/sumn - np.multiply(sumv/sumn, sumv/sumn))
+        if np.abs(val) < 1e-2:
+            val = 1e-2
+        return val
+
+    @cast_precision
+    def _filter_r(self, 
+                  inputs, 
+                  type_input,
+                  natoms,
+                  activation_fn=tf.nn.tanh, 
+                  stddev=1.0,
+                  bavg=0.0,
+                  name='linear', 
+                  reuse=None,
+                  trainable = True):
+        # natom x nei
+        outputs_size = [1] + self.filter_neuron
+        with tf.variable_scope(name, reuse=reuse):
+            start_index = 0
+            xyz_scatter_total = []
+            for type_i in range(self.ntypes):
+                # cut-out inputs
+                # with natom x nei_type_i
+                inputs_i = tf.slice (inputs,
+                                     [ 0, start_index       ],
+                                     [-1, self.sel_r[type_i]] )
+                start_index += self.sel_r[type_i]
+                shape_i = inputs_i.get_shape().as_list()
+                # with (natom x nei_type_i) x 1
+                xyz_scatter = tf.reshape(inputs_i, [-1, 1])
+                if self.compress and ((type_input, type_i) not in self.exclude_types):
+                    net = 'filter_' + str(type_input) + '_net_' + str(type_i)
+                    info = [self.lower[net], self.upper[net], self.upper[net] * self.table_config[0], self.table_config[1], self.table_config[2], self.table_config[3]]
+                    xyz_scatter = op_module.tabulate_fusion_se_r(tf.cast(self.table.data[net], self.filter_precision), info, inputs_i, last_layer_size = outputs_size[-1]) 
+                elif (type_input, type_i) not in self.exclude_types:
+                    xyz_scatter = embedding_net(xyz_scatter, 
+                                                self.filter_neuron, 
+                                                self.filter_precision, 
+                                                activation_fn = activation_fn, 
+                                                resnet_dt = self.filter_resnet_dt,
+                                                name_suffix = "_"+str(type_i),
+                                                stddev = stddev,
+                                                bavg = bavg,
+                                                seed = self.seed,
+                                                trainable = trainable, 
+                                                uniform_seed = self.uniform_seed,
+                                                initial_variables = self.embedding_net_variables,
+                                                )
+                    if (not self.uniform_seed) and (self.seed is not None): self.seed += self.seed_shift
+                    # natom x nei_type_i x out_size
+                    xyz_scatter = tf.reshape(xyz_scatter, (-1, shape_i[1], outputs_size[-1]))
+                else:
+                    natom = tf.shape(inputs)[0]
+                    xyz_scatter = tf.cast(tf.fill((natom, shape_i[1], outputs_size[-1]), 0.), self.filter_precision)
+                xyz_scatter_total.append(xyz_scatter)
+
+            # natom x nei x outputs_size
+            xyz_scatter = tf.concat(xyz_scatter_total, axis=1)
+            # natom x outputs_size
+            # 
+            res_rescale = 1./5.
+            result = tf.reduce_mean(xyz_scatter, axis = 1) * res_rescale
+
+        return result

_skbuild/linux-x86_64-3.6/cmake-install/deepmd/descriptor/se_t.py

+import numpy as np
+from typing import Tuple, List
+
+from deepmd.env import tf
+from deepmd.common import get_activation_func, get_precision, cast_precision
+from deepmd.env import GLOBAL_TF_FLOAT_PRECISION
+from deepmd.env import GLOBAL_NP_FLOAT_PRECISION
+from deepmd.env import op_module
+from deepmd.env import default_tf_session_config
+from deepmd.utils.network import embedding_net, embedding_net_rand_seed_shift
+from deepmd.utils.sess import run_sess
+from deepmd.utils.graph import load_graph_def, get_tensor_by_name_from_graph
+from deepmd.utils.tabulate import DPTabulate
+from .descriptor import Descriptor
+from .se import DescrptSe
+
+@Descriptor.register("se_e3")
+@Descriptor.register("se_at")
+@Descriptor.register("se_a_3be")
+class DescrptSeT (DescrptSe):
+    """DeepPot-SE constructed from all information (both angular and radial) of atomic
+    configurations.
+    
+    The embedding takes angles between two neighboring atoms as input.
+
+    Parameters
+    ----------
+    rcut
+            The cut-off radius
+    rcut_smth
+            From where the environment matrix should be smoothed
+    sel : list[str]
+            sel[i] specifies the maxmum number of type i atoms in the cut-off radius
+    neuron : list[int]
+            Number of neurons in each hidden layers of the embedding net
+    resnet_dt
+            Time-step `dt` in the resnet construction:
+            y = x + dt * \phi (Wx + b)
+    trainable
+            If the weights of embedding net are trainable.
+    seed
+            Random seed for initializing the network parameters.
+    set_davg_zero
+            Set the shift of embedding net input to zero.
+    activation_function
+            The activation function in the embedding net. Supported options are |ACTIVATION_FN|
+    precision
+            The precision of the embedding net parameters. Supported options are |PRECISION|
+    uniform_seed
+            Only for the purpose of backward compatibility, retrieves the old behavior of using the random seed
+    """
+    def __init__ (self, 
+                  rcut: float,
+                  rcut_smth: float,
+                  sel: List[str],
+                  neuron: List[int] = [24,48,96],
+                  resnet_dt: bool = False,
+                  trainable: bool = True,
+                  seed: int = None,
+                  set_davg_zero: bool = False,
+                  activation_function: str = 'tanh',
+                  precision: str = 'default',
+                  uniform_seed: bool = False
+    ) -> None:
+        """
+        Constructor
+        """
+        if rcut < rcut_smth:
+            raise RuntimeError("rcut_smth (%f) should be no more than rcut (%f)!" % (rcut_smth, rcut))
+        self.sel_a = sel
+        self.rcut_r = rcut
+        self.rcut_r_smth = rcut_smth
+        self.filter_neuron = neuron
+        self.filter_resnet_dt = resnet_dt
+        self.seed = seed
+        self.uniform_seed = uniform_seed
+        self.seed_shift = embedding_net_rand_seed_shift(self.filter_neuron)
+        self.trainable = trainable
+        self.filter_activation_fn = get_activation_func(activation_function)
+        self.filter_precision = get_precision(precision)
+        # self.exclude_types = set()
+        # for tt in exclude_types:
+        #     assert(len(tt) == 2)
+        #     self.exclude_types.add((tt[0], tt[1]))
+        #     self.exclude_types.add((tt[1], tt[0]))
+        self.set_davg_zero = set_davg_zero
+
+        # descrpt config
+        self.sel_r = [ 0 for ii in range(len(self.sel_a)) ]
+        self.ntypes = len(self.sel_a)
+        assert(self.ntypes == len(self.sel_r))
+        self.rcut_a = -1
+        # numb of neighbors and numb of descrptors
+        self.nnei_a = np.cumsum(self.sel_a)[-1]
+        self.nnei_r = np.cumsum(self.sel_r)[-1]
+        self.nnei = self.nnei_a + self.nnei_r
+        self.ndescrpt_a = self.nnei_a * 4
+        self.ndescrpt_r = self.nnei_r * 1
+        self.ndescrpt = self.ndescrpt_a + self.ndescrpt_r
+        self.useBN = False
+        self.dstd = None
+        self.davg = None
+        self.compress = False
+        self.embedding_net_variables = None
+
+        self.place_holders = {}
+        avg_zero = np.zeros([self.ntypes,self.ndescrpt]).astype(GLOBAL_NP_FLOAT_PRECISION)
+        std_ones = np.ones ([self.ntypes,self.ndescrpt]).astype(GLOBAL_NP_FLOAT_PRECISION)
+        sub_graph = tf.Graph()
+        with sub_graph.as_default():
+            name_pfx = 'd_sea_'
+            for ii in ['coord', 'box']:
+                self.place_holders[ii] = tf.placeholder(GLOBAL_NP_FLOAT_PRECISION, [None, None], name = name_pfx+'t_'+ii)
+            self.place_holders['type'] = tf.placeholder(tf.int32, [None, None], name=name_pfx+'t_type')
+            self.place_holders['natoms_vec'] = tf.placeholder(tf.int32, [self.ntypes+2], name=name_pfx+'t_natoms')
+            self.place_holders['default_mesh'] = tf.placeholder(tf.int32, [None], name=name_pfx+'t_mesh')
+            self.stat_descrpt, descrpt_deriv, rij, nlist \
+                = op_module.prod_env_mat_a(self.place_holders['coord'],
+                                         self.place_holders['type'],
+                                         self.place_holders['natoms_vec'],
+                                         self.place_holders['box'],
+                                         self.place_holders['default_mesh'],
+                                         tf.constant(avg_zero),
+                                         tf.constant(std_ones),
+                                         rcut_a = self.rcut_a,
+                                         rcut_r = self.rcut_r,
+                                         rcut_r_smth = self.rcut_r_smth,
+                                         sel_a = self.sel_a,
+                                         sel_r = self.sel_r)
+        self.sub_sess = tf.Session(graph = sub_graph, config=default_tf_session_config)
+
+
+    def get_rcut (self) -> float:
+        """
+        Returns the cut-off radisu
+        """
+        return self.rcut_r
+
+    def get_ntypes (self) -> int:
+        """
+        Returns the number of atom types
+        """
+        return self.ntypes
+
+    def get_dim_out (self) -> int:
+        """
+        Returns the output dimension of this descriptor
+        """
+        return self.filter_neuron[-1]
+
+    def get_nlist (self) -> Tuple[tf.Tensor, tf.Tensor, List[int], List[int]]:
+        """
+        Returns
+        -------
+        nlist
+                Neighbor list
+        rij
+                The relative distance between the neighbor and the center atom.
+        sel_a
+                The number of neighbors with full information
+        sel_r
+                The number of neighbors with only radial information
+        """
+        return self.nlist, self.rij, self.sel_a, self.sel_r
+
+    def compute_input_stats (self,
+                             data_coord : list, 
+                             data_box : list, 
+                             data_atype : list, 
+                             natoms_vec : list,
+                             mesh : list, 
+                             input_dict : dict
+    ) -> None :
+        """
+        Compute the statisitcs (avg and std) of the training data. The input will be normalized by the statistics.
+        
+        Parameters
+        ----------
+        data_coord
+                The coordinates. Can be generated by deepmd.model.make_stat_input
+        data_box
+                The box. Can be generated by deepmd.model.make_stat_input
+        data_atype
+                The atom types. Can be generated by deepmd.model.make_stat_input
+        natoms_vec
+                The vector for the number of atoms of the system and different types of atoms. Can be generated by deepmd.model.make_stat_input
+        mesh
+                The mesh for neighbor searching. Can be generated by deepmd.model.make_stat_input
+        input_dict
+                Dictionary for additional input
+        """
+        all_davg = []
+        all_dstd = []
+        if True:
+            sumr = []
+            suma = []
+            sumn = []
+            sumr2 = []
+            suma2 = []
+            for cc,bb,tt,nn,mm in zip(data_coord,data_box,data_atype,natoms_vec,mesh) :
+                sysr,sysr2,sysa,sysa2,sysn \
+                    = self._compute_dstats_sys_smth(cc,bb,tt,nn,mm)
+                sumr.append(sysr)
+                suma.append(sysa)
+                sumn.append(sysn)
+                sumr2.append(sysr2)
+                suma2.append(sysa2)
+            sumr = np.sum(sumr, axis = 0)
+            suma = np.sum(suma, axis = 0)
+            sumn = np.sum(sumn, axis = 0)
+            sumr2 = np.sum(sumr2, axis = 0)
+            suma2 = np.sum(suma2, axis = 0)
+            for type_i in range(self.ntypes) :
+                davgunit = [sumr[type_i]/sumn[type_i], 0, 0, 0]
+                dstdunit = [self._compute_std(sumr2[type_i], sumr[type_i], sumn[type_i]), 
+                            self._compute_std(suma2[type_i], suma[type_i], sumn[type_i]), 
+                            self._compute_std(suma2[type_i], suma[type_i], sumn[type_i]), 
+                            self._compute_std(suma2[type_i], suma[type_i], sumn[type_i])
+                            ]
+                davg = np.tile(davgunit, self.ndescrpt // 4)
+                dstd = np.tile(dstdunit, self.ndescrpt // 4)
+                all_davg.append(davg)
+                all_dstd.append(dstd)
+
+        if not self.set_davg_zero:
+            self.davg = np.array(all_davg)
+        self.dstd = np.array(all_dstd)
+
+
+    def enable_compression(self,
+                           min_nbor_dist : float,
+                           model_file : str = 'frozon_model.pb',
+                           table_extrapolate : float = 5,
+                           table_stride_1 : float = 0.01,
+                           table_stride_2 : float = 0.1,
+                           check_frequency : int = -1,
+                           suffix : str = "",
+    ) -> None:
+        """
+        Reveive the statisitcs (distance, max_nbor_size and env_mat_range) of the training data.
+
+        Parameters
+        ----------
+        min_nbor_dist
+                The nearest distance between atoms
+        model_file
+                The original frozen model, which will be compressed by the program
+        table_extrapolate
+                The scale of model extrapolation
+        table_stride_1
+                The uniform stride of the first table
+        table_stride_2
+                The uniform stride of the second table
+        check_frequency
+                The overflow check frequency
+        suffix : str, optional
+                The suffix of the scope
+        """
+        assert (
+            not self.filter_resnet_dt
+        ), "Model compression error: descriptor resnet_dt must be false!"
+
+        for ii in range(len(self.filter_neuron) - 1):
+            if self.filter_neuron[ii] * 2 != self.filter_neuron[ii + 1]:
+                raise NotImplementedError(
+                    "Model Compression error: descriptor neuron [%s] is not supported by model compression! "
+                    "The size of the next layer of the neural network must be twice the size of the previous layer." 
+                    % ','.join([str(item) for item in self.filter_neuron])
+                )
+
+        self.compress = True
+        self.table = DPTabulate(
+            self, self.filter_neuron, model_file, activation_fn = self.filter_activation_fn, suffix=suffix)
+        self.table_config = [table_extrapolate, table_stride_1 * 10, table_stride_2 * 10, check_frequency]
+        self.lower, self.upper \
+            = self.table.build(min_nbor_dist, 
+                               table_extrapolate, 
+                               table_stride_1 * 10, 
+                               table_stride_2 * 10)
+        
+        graph, _ = load_graph_def(model_file)
+        self.davg = get_tensor_by_name_from_graph(graph, 'descrpt_attr%s/t_avg' % suffix)
+        self.dstd = get_tensor_by_name_from_graph(graph, 'descrpt_attr%s/t_std' % suffix)
+
+
+    def build (self, 
+               coord_ : tf.Tensor, 
+               atype_ : tf.Tensor,
+               natoms : tf.Tensor,
+               box_ : tf.Tensor, 
+               mesh : tf.Tensor,
+               input_dict : dict, 
+               reuse : bool = None,
+               suffix : str = ''
+    ) -> tf.Tensor:
+        """
+        Build the computational graph for the descriptor
+
+        Parameters
+        ----------
+        coord_
+                The coordinate of atoms
+        atype_
+                The type of atoms
+        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
+        mesh
+                For historical reasons, only the length of the Tensor matters.
+                if size of mesh == 6, pbc is assumed. 
+                if size of mesh == 0, no-pbc is assumed. 
+        input_dict
+                Dictionary for additional inputs
+        reuse
+                The weights in the networks should be reused when get the variable.
+        suffix
+                Name suffix to identify this descriptor
+
+        Returns
+        -------
+        descriptor
+                The output descriptor
+        """
+        davg = self.davg
+        dstd = self.dstd
+        with tf.variable_scope('descrpt_attr' + suffix, reuse = reuse) :
+            if davg is None:
+                davg = np.zeros([self.ntypes, self.ndescrpt]) 
+            if dstd is None:
+                dstd = np.ones ([self.ntypes, self.ndescrpt])
+            t_rcut = tf.constant(np.max([self.rcut_r, self.rcut_a]), 
+                                 name = 'rcut', 
+                                 dtype = GLOBAL_TF_FLOAT_PRECISION)
+            t_ntypes = tf.constant(self.ntypes, 
+                                   name = 'ntypes', 
+                                   dtype = tf.int32)
+            t_ndescrpt = tf.constant(self.ndescrpt, 
+                                     name = 'ndescrpt', 
+                                     dtype = tf.int32)            
+            t_sel = tf.constant(self.sel_a, 
+                                name = 'sel', 
+                                dtype = tf.int32)            
+            self.t_avg = tf.get_variable('t_avg', 
+                                         davg.shape, 
+                                         dtype = GLOBAL_TF_FLOAT_PRECISION,
+                                         trainable = False,
+                                         initializer = tf.constant_initializer(davg))
+            self.t_std = tf.get_variable('t_std', 
+                                         dstd.shape, 
+                                         dtype = GLOBAL_TF_FLOAT_PRECISION,
+                                         trainable = False,
+                                         initializer = tf.constant_initializer(dstd))
+
+        coord = tf.reshape (coord_, [-1, natoms[1] * 3])
+        box   = tf.reshape (box_, [-1, 9])
+        atype = tf.reshape (atype_, [-1, natoms[1]])
+
+        self.descrpt, self.descrpt_deriv, self.rij, self.nlist \
+            = op_module.prod_env_mat_a (coord,
+                                       atype,
+                                       natoms,
+                                       box,
+                                       mesh,
+                                       self.t_avg,
+                                       self.t_std,
+                                       rcut_a = self.rcut_a,
+                                       rcut_r = self.rcut_r,
+                                       rcut_r_smth = self.rcut_r_smth,
+                                       sel_a = self.sel_a,
+                                       sel_r = self.sel_r)
+
+        self.descrpt_reshape = tf.reshape(self.descrpt, [-1, self.ndescrpt])
+        self._identity_tensors(suffix=suffix)
+
+        self.dout, self.qmat = self._pass_filter(self.descrpt_reshape, 
+                                                 atype,
+                                                 natoms, 
+                                                 input_dict,
+                                                 suffix = suffix, 
+                                                 reuse = reuse, 
+                                                 trainable = self.trainable)
+
+        return self.dout
+
+
+    def prod_force_virial(self, 
+                          atom_ener : tf.Tensor, 
+                          natoms : tf.Tensor
+    ) -> Tuple[tf.Tensor, tf.Tensor, tf.Tensor]:
+        """
+        Compute force and virial
+
+        Parameters
+        ----------
+        atom_ener
+                The atomic energy
+        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
+
+        Returns
+        -------
+        force
+                The force on atoms
+        virial
+                The total virial
+        atom_virial
+                The atomic virial
+        """
+        [net_deriv] = tf.gradients (atom_ener, self.descrpt_reshape)
+        net_deriv_reshape = tf.reshape (net_deriv, [np.cast['int64'](-1), natoms[0] * np.cast['int64'](self.ndescrpt)])        
+        force \
+            = op_module.prod_force_se_a (net_deriv_reshape,
+                                          self.descrpt_deriv,
+                                          self.nlist,
+                                          natoms,
+                                          n_a_sel = self.nnei_a,
+                                          n_r_sel = self.nnei_r)
+        virial, atom_virial \
+            = op_module.prod_virial_se_a (net_deriv_reshape,
+                                           self.descrpt_deriv,
+                                           self.rij,
+                                           self.nlist,
+                                           natoms,
+                                           n_a_sel = self.nnei_a,
+                                           n_r_sel = self.nnei_r)
+        return force, virial, atom_virial
+        
+
+    def _pass_filter(self, 
+                     inputs,
+                     atype,
+                     natoms,
+                     input_dict,
+                     reuse = None,
+                     suffix = '', 
+                     trainable = True) :
+        start_index = 0
+        inputs = tf.reshape(inputs, [-1, natoms[0], self.ndescrpt])
+        output = []
+        output_qmat = []
+        inputs_i = inputs
+        inputs_i = tf.reshape(inputs_i, [-1, self.ndescrpt])
+        type_i = -1
+        layer, qmat = self._filter(inputs_i, type_i, name='filter_type_all'+suffix, natoms=natoms, reuse=reuse, trainable = trainable, activation_fn = self.filter_activation_fn)
+        layer = tf.reshape(layer, [tf.shape(inputs)[0], natoms[0], self.get_dim_out()])
+        # qmat  = tf.reshape(qmat,  [tf.shape(inputs)[0], natoms[0] * self.get_dim_rot_mat_1() * 3])
+        output.append(layer)
+        # output_qmat.append(qmat)
+        output = tf.concat(output, axis = 1)
+        # output_qmat = tf.concat(output_qmat, axis = 1)
+        return output, None
+
+
+    def _compute_dstats_sys_smth (self,
+                                 data_coord, 
+                                 data_box, 
+                                 data_atype,                             
+                                 natoms_vec,
+                                 mesh) :    
+        dd_all \
+            = run_sess(self.sub_sess, self.stat_descrpt, 
+                                feed_dict = {
+                                    self.place_holders['coord']: data_coord,
+                                    self.place_holders['type']: data_atype,
+                                    self.place_holders['natoms_vec']: natoms_vec,
+                                    self.place_holders['box']: data_box,
+                                    self.place_holders['default_mesh']: mesh,
+                                })
+        natoms = natoms_vec
+        dd_all = np.reshape(dd_all, [-1, self.ndescrpt * natoms[0]])
+        start_index = 0
+        sysr = []
+        sysa = []
+        sysn = []
+        sysr2 = []
+        sysa2 = []
+        for type_i in range(self.ntypes):
+            end_index = start_index + self.ndescrpt * natoms[2+type_i]
+            dd = dd_all[:, start_index:end_index]
+            dd = np.reshape(dd, [-1, self.ndescrpt])
+            start_index = end_index        
+            # compute
+            dd = np.reshape (dd, [-1, 4])
+            ddr = dd[:,:1]
+            dda = dd[:,1:]
+            sumr = np.sum(ddr)
+            suma = np.sum(dda) / 3.
+            sumn = dd.shape[0]
+            sumr2 = np.sum(np.multiply(ddr, ddr))
+            suma2 = np.sum(np.multiply(dda, dda)) / 3.
+            sysr.append(sumr)
+            sysa.append(suma)
+            sysn.append(sumn)
+            sysr2.append(sumr2)
+            sysa2.append(suma2)
+        return sysr, sysr2, sysa, sysa2, sysn
+
+
+    def _compute_std (self,sumv2, sumv, sumn) :
+        val = np.sqrt(sumv2/sumn - np.multiply(sumv/sumn, sumv/sumn))
+        if np.abs(val) < 1e-2:
+            val = 1e-2
+        return val
+
+    @cast_precision
+    def _filter(self, 
+                inputs, 
+                type_input,
+                natoms,
+                activation_fn=tf.nn.tanh, 
+                stddev=1.0,
+                bavg=0.0,
+                name='linear', 
+                reuse=None,
+                trainable = True):
+        # natom x (nei x 4)
+        shape = inputs.get_shape().as_list()
+        outputs_size = [1] + self.filter_neuron
+        with tf.variable_scope(name, reuse=reuse):
+            start_index_i = 0
+            result = None
+            for type_i in range(self.ntypes):
+                # cut-out inputs
+                # with natom x (nei_type_i x 4)  
+                inputs_i = tf.slice (inputs,
+                                     [ 0, start_index_i      *4],
+                                     [-1, self.sel_a[type_i] *4] )
+                start_index_j = start_index_i
+                start_index_i += self.sel_a[type_i]
+                nei_type_i = self.sel_a[type_i]
+                shape_i = inputs_i.get_shape().as_list()
+                assert(shape_i[1] == nei_type_i * 4)
+                # with natom x nei_type_i x 4
+                env_i = tf.reshape(inputs_i, [-1, nei_type_i, 4])
+                # with natom x nei_type_i x 3
+                env_i = tf.slice(env_i, [0, 0, 1], [-1, -1, -1])
+                for type_j in range(type_i, self.ntypes):
+                    # with natom x (nei_type_j x 4)  
+                    inputs_j = tf.slice (inputs,
+                                         [ 0, start_index_j      *4],
+                                         [-1, self.sel_a[type_j] *4] )
+                    start_index_j += self.sel_a[type_j]
+                    nei_type_j = self.sel_a[type_j]
+                    shape_j = inputs_j.get_shape().as_list()
+                    assert(shape_j[1] == nei_type_j * 4)
+                    # with natom x nei_type_j x 4
+                    env_j = tf.reshape(inputs_j, [-1, nei_type_j, 4])
+                    # with natom x nei_type_i x 3
+                    env_j = tf.slice(env_j, [0, 0, 1], [-1, -1, -1])
+                    # with natom x nei_type_i x nei_type_j
+                    env_ij = tf.einsum('ijm,ikm->ijk', env_i, env_j)
+                    # with (natom x nei_type_i x nei_type_j)
+                    ebd_env_ij = tf.reshape(env_ij, [-1, 1])
+                    if self.compress:
+                        net = 'filter_' + str(type_i) + '_net_' + str(type_j)
+                        info = [self.lower[net], self.upper[net], self.upper[net] * self.table_config[0], self.table_config[1], self.table_config[2], self.table_config[3]]
+                        res_ij = op_module.tabulate_fusion_se_t(tf.cast(self.table.data[net], self.filter_precision), info, ebd_env_ij, env_ij, last_layer_size = outputs_size[-1]) 
+                    else:
+                        # with (natom x nei_type_i x nei_type_j) x out_size
+                        ebd_env_ij = embedding_net(ebd_env_ij, 
+                                                   self.filter_neuron, 
+                                                   self.filter_precision, 
+                                                   activation_fn = activation_fn, 
+                                                   resnet_dt = self.filter_resnet_dt,
+                                                   name_suffix = f"_{type_i}_{type_j}",
+                                                   stddev = stddev,
+                                                   bavg = bavg,
+                                                   seed = self.seed,
+                                                   trainable = trainable, 
+                                                   uniform_seed = self.uniform_seed,
+                                                   initial_variables = self.embedding_net_variables,
+                                                   )
+                        if (not self.uniform_seed) and (self.seed is not None): self.seed += self.seed_shift
+                        # with natom x nei_type_i x nei_type_j x out_size
+                        ebd_env_ij = tf.reshape(ebd_env_ij, [-1, nei_type_i, nei_type_j, outputs_size[-1]])
+                        # with natom x out_size
+                        res_ij = tf.einsum('ijk,ijkm->im', env_ij, ebd_env_ij)
+                    res_ij = res_ij * (1.0 / float(nei_type_i) / float(nei_type_j))
+                    if result is None:
+                        result = res_ij
+                    else:
+                        result += res_ij
+        return result, None
+

_skbuild/linux-x86_64-3.6/cmake-install/deepmd/entrypoints/__init__.py

+"""Submodule that contains all the DeePMD-Kit entry point scripts."""
+
+from .compress import compress
+from .config import config
+from .doc import doc_train_input
+from .freeze import freeze
+from .test import test
+# import `train` as `train_dp` to avoid the conflict of the
+# module name `train` and the function name `train`
+from .train import train as train_dp
+from .transfer import transfer
+from ..infer.model_devi import make_model_devi
+from .convert import convert
+from .neighbor_stat import neighbor_stat
+
+__all__ = [
+    "config",
+    "doc_train_input",
+    "freeze",
+    "test",
+    "train_dp",
+    "transfer",
+    "compress",
+    "doc_train_input",
+    "make_model_devi",
+    "convert",
+    "neighbor_stat",
+]

_skbuild/linux-x86_64-3.6/cmake-install/deepmd/entrypoints/compress.py

+"""Compress a model, which including tabulating the embedding-net."""
+
+import os
+import json
+import logging
+from typing import Optional
+
+from deepmd.common import j_loader
+from deepmd.env import tf, GLOBAL_ENER_FLOAT_PRECISION
+from deepmd.utils.argcheck import normalize
+from deepmd.utils.compat import update_deepmd_input
+from deepmd.utils.errors import GraphTooLargeError, GraphWithoutTensorError
+from deepmd.utils.graph import get_tensor_by_name
+
+from .freeze import freeze
+from .train import train, get_rcut, get_min_nbor_dist
+from .transfer import transfer
+
+__all__ = ["compress"]
+
+log = logging.getLogger(__name__)
+
+
+def compress(
+    *,
+    input: str,
+    output: str,
+    extrapolate: int,
+    step: float,
+    frequency: str,
+    checkpoint_folder: str,
+    training_script: str,
+    mpi_log: str,
+    log_path: Optional[str],
+    log_level: int,
+    **kwargs
+):
+    """Compress model.
+
+    The table is composed of fifth-order polynomial coefficients and is assembled from
+    two sub-tables. The first table takes the step parameter as the domain's uniform step size,
+    while the second table takes 10 * step as it's uniform step size. The range of the
+    first table is automatically detected by the code, while the second table ranges
+    from the first table's upper boundary(upper) to the extrapolate(parameter) * upper.
+
+    Parameters
+    ----------
+    input : str
+        frozen model file to compress
+    output : str
+        compressed model filename
+    extrapolate : int
+        scale of model extrapolation
+    step : float
+        uniform step size of the tabulation's first table
+    frequency : str
+        frequency of tabulation overflow check
+    checkpoint_folder : str
+        trining checkpoint folder for freezing
+    training_script : str
+        training script of the input frozen model
+    mpi_log : str
+        mpi logging mode for training
+    log_path : Optional[str]
+        if speccified log will be written to this file
+    log_level : int
+        logging level
+    """
+    try:
+        t_jdata = get_tensor_by_name(input, 'train_attr/training_script')
+        t_min_nbor_dist = get_tensor_by_name(input, 'train_attr/min_nbor_dist')
+        jdata = json.loads(t_jdata)
+    except GraphWithoutTensorError as e:
+        if training_script == None:
+            raise RuntimeError(
+                "The input frozen model: %s has no training script or min_nbor_dist information, "
+                "which is not supported by the model compression interface. "
+                "Please consider using the --training-script command within the model compression interface to provide the training script of the input frozen model. "
+                "Note that the input training script must contain the correct path to the training data." % input
+            ) from e
+        elif not os.path.exists(training_script):
+            raise RuntimeError(
+                "The input training script %s (%s) does not exist! Please check the path of the training script. " % (input, os.path.abspath(input))
+            ) from e
+        else:
+            log.info("stage 0: compute the min_nbor_dist")
+            jdata = j_loader(training_script)
+            jdata = update_deepmd_input(jdata)
+            t_min_nbor_dist = get_min_nbor_dist(jdata, get_rcut(jdata))
+
+    _check_compress_type(input)
+
+    tf.constant(t_min_nbor_dist,
+        name = 'train_attr/min_nbor_dist',
+        dtype = GLOBAL_ENER_FLOAT_PRECISION)
+    jdata["model"]["compress"] = {}
+    jdata["model"]["compress"]["model_file"] = input
+    jdata["model"]["compress"]["min_nbor_dist"] = t_min_nbor_dist
+    jdata["model"]["compress"]["table_config"] = [
+        extrapolate,
+        step,
+        10 * step,
+        int(frequency),
+    ]
+    jdata["training"]["save_ckpt"] = os.path.join("model-compression", "model.ckpt")
+    jdata = update_deepmd_input(jdata)
+    jdata = normalize(jdata)
+
+    # check the descriptor info of the input file
+    # move to the specific Descriptor class
+
+    # stage 1: training or refining the model with tabulation
+    log.info("\n\n")
+    log.info("stage 1: compress the model")
+    control_file = "compress.json"
+    with open(control_file, "w") as fp:
+        json.dump(jdata, fp, indent=4)
+    try:
+        train(
+            INPUT=control_file,
+            init_model=None,
+            restart=None,
+            init_frz_model=None,
+            output=control_file,
+            mpi_log=mpi_log,
+            log_level=log_level,
+            log_path=log_path,
+            is_compress=True,
+        )
+    except GraphTooLargeError as e:
+        raise RuntimeError(
+            "The uniform step size of the tabulation's first table is %f, " 
+            "which is too small. This leads to a very large graph size, "
+            "exceeding protobuf's limitation (2 GB). You should try to "
+            "increase the step size." % step
+        ) from e
+
+    # reset the graph, otherwise the size limitation will be only 2 GB / 2 = 1 GB
+    tf.reset_default_graph()
+
+    # stage 2: freeze the model
+    log.info("\n\n")
+    log.info("stage 2: freeze the model")
+    try:
+        freeze(checkpoint_folder=checkpoint_folder, output=output, node_names=None)
+    except GraphTooLargeError as e:
+        raise RuntimeError(
+            "The uniform step size of the tabulation's first table is %f, " 
+            "which is too small. This leads to a very large graph size, "
+            "exceeding protobuf's limitation (2 GB). You should try to "
+            "increase the step size." % step
+        ) from e
+
+def _check_compress_type(model_file):
+    try:
+        t_model_type = bytes.decode(get_tensor_by_name(model_file, 'model_type'))
+    except GraphWithoutTensorError as e:
+        # Compatible with the upgraded model, which has no 'model_type' info
+        t_model_type = None
+    
+    if t_model_type == "compressed_model":
+        raise RuntimeError("The input frozen model %s has already been compressed! Please do not compress the model repeatedly. " % model_file)

_skbuild/linux-x86_64-3.6/cmake-install/deepmd/entrypoints/config.py

+#!/usr/bin/env python3
+"""Quickly create a configuration file for smooth model."""
+
+import json
+import yaml
+from pathlib import Path
+from typing import Any, Dict, List, Tuple
+
+import numpy as np
+
+__all__ = ["config"]
+
+
+DEFAULT_DATA: Dict[str, Any] = {
+    "use_smooth": True,
+    "sel_a": [],
+    "rcut_smth": -1,
+    "rcut": -1,
+    "filter_neuron": [20, 40, 80],
+    "filter_resnet_dt": False,
+    "axis_neuron": 8,
+    "fitting_neuron": [240, 240, 240],
+    "fitting_resnet_dt": True,
+    "coord_norm": True,
+    "type_fitting_net": False,
+    "systems": [],
+    "set_prefix": "set",
+    "stop_batch": -1,
+    "batch_size": -1,
+    "start_lr": 0.001,
+    "decay_steps": -1,
+    "decay_rate": 0.95,
+    "start_pref_e": 0.02,
+    "limit_pref_e": 1,
+    "start_pref_f": 1000,
+    "limit_pref_f": 1,
+    "start_pref_v": 0,
+    "limit_pref_v": 0,
+    "seed": 1,
+    "disp_file": "lcurve.out",
+    "disp_freq": 1000,
+    "numb_test": 10,
+    "save_freq": 10000,
+    "save_ckpt": "model.ckpt",
+    "disp_training": True,
+    "time_training": True,
+}
+
+
+def valid_dir(path: Path):
+    """Check if directory is a valid deepmd system directory.
+
+    Parameters
+    ----------
+    path : Path
+        path to directory
+
+    Raises
+    ------
+    OSError
+        if `type.raw` is missing on dir or `box.npy` or `coord.npy` are missing in one
+        of the sets subdirs
+    """
+    if not (path / "type.raw").is_file():
+        raise OSError
+    for ii in path.glob("set.*"):
+        if not (ii / "box.npy").is_file():
+            raise OSError
+        if not (ii / "coord.npy").is_file():
+            raise OSError
+
+
+def load_systems(dirs: List[Path]) -> Tuple[List[np.ndarray], List[np.ndarray]]:
+    """Load systems to memory for disk.
+
+    Parameters
+    ----------
+    dirs : List[Path]
+        list of system directories paths
+
+    Returns
+    -------
+    Tuple[List[np.ndarray], List[np.ndarray]]
+        atoms types and structure cells formated as Nx9 array
+    """
+    all_type = []
+    all_box = []
+    for d in dirs:
+        sys_type = np.loadtxt(d / "type.raw", dtype=int)
+        sys_box = np.vstack([np.load(s / "box.npy") for s in d.glob("set.*")])
+        all_type.append(sys_type)
+        all_box.append(sys_box)
+    return all_type, all_box
+
+
+def get_system_names() -> List[Path]:
+    """Get system directory paths from stdin.
+
+    Returns
+    -------
+    List[Path]
+        list of system directories paths
+    """
+    dirs = input("Enter system path(s) (seperated by space, wild card supported): \n")
+    system_dirs = []
+    for dir_str in dirs.split():
+        found_dirs = Path.cwd().glob(dir_str)
+        for d in found_dirs:
+            valid_dir(d)
+            system_dirs.append(d)
+
+    return system_dirs
+
+
+def get_rcut() -> float:
+    """Get rcut from stdin from user.
+
+    Returns
+    -------
+    float
+        input rcut lenght converted to float
+
+    Raises
+    ------
+    ValueError
+        if rcut is smaller than 0.0
+    """
+    dv = 6.0
+    rcut_input = input(f"Enter rcut (default {dv:.1f} A): \n")
+    try:
+        rcut = float(rcut_input)
+    except ValueError as e:
+        print(f"invalid rcut: {e} setting to default: {dv:.1f}")
+        rcut = dv
+    if rcut <= 0:
+        raise ValueError("rcut should be > 0")
+    return rcut
+
+
+def get_batch_size_rule() -> int:
+    """Get minimal batch size from user from stdin.
+
+    Returns
+    -------
+    int
+        size of the batch
+
+    Raises
+    ------
+    ValueError
+        if batch size is <= 0
+    """
+    dv = 32
+    matom_input = input(
+        f"Enter the minimal number of atoms in a batch (default {dv:d}: \n"
+    )
+    try:
+        matom = int(matom_input)
+    except ValueError as e:
+        print(f"invalid batch size: {e} setting to default: {dv:d}")
+        matom = dv
+    if matom <= 0:
+        raise ValueError("the number should be > 0")
+    return matom
+
+
+def get_stop_batch() -> int:
+    """Get stop batch from user from stdin.
+
+    Returns
+    -------
+    int
+        size of the batch
+
+    Raises
+    ------
+    ValueError
+        if stop batch is <= 0
+    """
+    dv = 1000000
+    sb_input = input(f"Enter the stop batch (default {dv:d}): \n")
+    try:
+        sb = int(sb_input)
+    except ValueError as e:
+        print(f"invalid stop batch: {e} setting to default: {dv:d}")
+        sb = dv
+    if sb <= 0:
+        raise ValueError("the number should be > 0")
+    return sb
+
+
+def get_ntypes(all_type: List[np.ndarray]) -> int:
+    """Count number of unique elements.
+
+    Parameters
+    ----------
+    all_type : List[np.ndarray]
+        list with arrays specifying elements of structures
+
+    Returns
+    -------
+    int
+        number of unique elements
+    """
+    return len(np.unique(all_type))
+
+
+def get_max_density(
+    all_type: List[np.ndarray], all_box: List[np.ndarray]
+) -> np.ndarray:
+    """Compute maximum density in suppliedd cells.
+
+    Parameters
+    ----------
+    all_type : List[np.ndarray]
+        list with arrays specifying elements of structures
+    all_box : List[np.ndarray]
+        list with arrays specifying cells for all structures
+
+    Returns
+    -------
+    float
+        maximum atom density in all supplies structures for each element individually
+    """
+    ntypes = get_ntypes(all_type)
+    all_max = []
+    for tt, bb in zip(all_type, all_box):
+        vv = np.reshape(bb, [-1, 3, 3])
+        vv = np.linalg.det(vv)
+        min_v = np.min(vv)
+        type_count = []
+        for ii in range(ntypes):
+            type_count.append(sum(tt == ii))
+        max_den = type_count / min_v
+        all_max.append(max_den)
+    all_max = np.max(all_max, axis=0)
+    return all_max
+
+
+def suggest_sel(
+    all_type: List[np.ndarray],
+    all_box: List[np.ndarray],
+    rcut: float,
+    ratio: float = 1.5,
+) -> List[int]:
+    """Suggest selection parameter.
+
+    Parameters
+    ----------
+    all_type : List[np.ndarray]
+        list with arrays specifying elements of structures
+    all_box : List[np.ndarray]
+        list with arrays specifying cells for all structures
+    rcut : float
+        cutoff radius
+    ratio : float, optional
+        safety margin to add to estimated value, by default 1.5
+
+    Returns
+    -------
+    List[int]
+        [description]
+    """
+    max_den = get_max_density(all_type, all_box)
+    return [int(ii) for ii in max_den * 4.0 / 3.0 * np.pi * rcut ** 3 * ratio]
+
+
+def suggest_batch_size(all_type: List[np.ndarray], min_atom: int) -> List[int]:
+    """Get suggestion for batch size.
+
+    Parameters
+    ----------
+    all_type : List[np.ndarray]
+        list with arrays specifying elements of structures
+    min_atom : int
+        minimal number of atoms in batch
+
+    Returns
+    -------
+    List[int]
+        suggested batch sizes for each system
+    """
+    bs = []
+    for ii in all_type:
+        natoms = len(ii)
+        tbs = min_atom // natoms
+        if (min_atom // natoms) * natoms != min_atom:
+            tbs += 1
+        bs.append(tbs)
+    return bs
+
+
+def suggest_decay(stop_batch: int) -> Tuple[int, float]:
+    """Suggest number of decay steps and decay rate.
+
+    Parameters
+    ----------
+    stop_batch : int
+        stop batch number
+
+    Returns
+    -------
+    Tuple[int, float]
+        number of decay steps and decay rate
+    """
+    decay_steps = int(stop_batch // 200)
+    decay_rate = 0.95
+    return decay_steps, decay_rate
+
+
+def config(*, output: str, **kwargs):
+    """Auto config file generator.
+
+    Parameters
+    ----------
+    output: str
+        file to write config file
+
+    Raises
+    ------
+    RuntimeError
+        if user does not input any systems
+    ValueError
+        if output file is of wrong type
+    """
+    all_sys = get_system_names()
+    if len(all_sys) == 0:
+        raise RuntimeError("no system specified")
+    rcut = get_rcut()
+    matom = get_batch_size_rule()
+    stop_batch = get_stop_batch()
+
+    all_type, all_box = load_systems(all_sys)
+    sel = suggest_sel(all_type, all_box, rcut, ratio=1.5)
+    bs = suggest_batch_size(all_type, matom)
+    decay_steps, decay_rate = suggest_decay(stop_batch)
+
+    jdata = DEFAULT_DATA.copy()
+    jdata["systems"] = [str(ii) for ii in all_sys]
+    jdata["sel_a"] = sel
+    jdata["rcut"] = rcut
+    jdata["rcut_smth"] = rcut - 0.2
+    jdata["stop_batch"] = stop_batch
+    jdata["batch_size"] = bs
+    jdata["decay_steps"] = decay_steps
+    jdata["decay_rate"] = decay_rate
+
+    with open(output, "w") as fp:
+        if output.endswith("json"):
+            json.dump(jdata, fp, indent=4)
+        elif output.endswith(("yml", "yaml")):
+            yaml.safe_dump(jdata, fp, default_flow_style=False)
+        else:
+            raise ValueError("output file must be of type json or yaml")

_skbuild/linux-x86_64-3.6/cmake-install/deepmd/entrypoints/convert.py

+from deepmd.utils.convert import convert_012_to_21, convert_10_to_21, convert_20_to_21, convert_13_to_21, convert_12_to_21 
+
+def convert(
+    *,
+    FROM: str,
+    input_model: str,
+    output_model: str,
+    **kwargs,
+):
+    if FROM == '0.12':
+        convert_012_to_21(input_model, output_model)
+    elif FROM == '1.0':
+        convert_10_to_21(input_model, output_model)
+    elif FROM in ['1.1', '1.2']:
+        # no difference between 1.1 and 1.2
+        convert_12_to_21(input_model, output_model)
+    elif FROM == '1.3':
+        convert_13_to_21(input_model, output_model)
+    elif FROM == '2.0':
+        convert_20_to_21(input_model, output_model)
+    else:
+        raise RuntimeError('unsupported model version ' + FROM)