Skip to content

GitLab

"vscode:/vscode.git/clone" did not exist on "315506c9ad1d6c73a32e8e2e1a2cb10f0513498c"
Commit 56d2b104 authored by zhangqha's avatar zhangqha
Browse files

BladeDISC DeePMD code update

parent 6b33aeb8
Pipeline #180 failed with stages
in 0 seconds
Hide whitespace changes
Inline Side-by-side
Showing with 0 additions and 7302 deletions
_skbuild/linux-x86_64-3.6/cmake-install/deepmd/cluster/slurm.py deleted 100644 → 0
View file @ 6b33aeb8
"""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 deleted 100644 → 0
View file @ 6b33aeb8
"""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 deleted 100644 → 0
View file @ 6b33aeb8
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 deleted 100644 → 0
View file @ 6b33aeb8
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 deleted 100644 → 0
View file @ 6b33aeb8
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 deleted 100644 → 0
View file @ 6b33aeb8
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 deleted 100644 → 0
View file @ 6b33aeb8
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 deleted 100644 → 0
View file @ 6b33aeb8
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 deleted 100644 → 0
View file @ 6b33aeb8
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 deleted 100644 → 0
View file @ 6b33aeb8
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 deleted 100644 → 0
View file @ 6b33aeb8
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 deleted 100644 → 0
View file @ 6b33aeb8
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 deleted 100644 → 0
View file @ 6b33aeb8
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 deleted 100644 → 0
View file @ 6b33aeb8
"""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 deleted 100644 → 0
View file @ 6b33aeb8
"""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 deleted 100644 → 0
View file @ 6b33aeb8
#!/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 deleted 100644 → 0
View file @ 6b33aeb8
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)
_skbuild/linux-x86_64-3.6/cmake-install/deepmd/entrypoints/doc.py deleted 100644 → 0
View file @ 6b33aeb8
"""Module that prints train input arguments docstrings."""
from deepmd.utils.argcheck import gen_doc, gen_json
__all__ = ["doc_train_input"]
def doc_train_input(*, out_type: str = "rst", **kwargs):
"""Print out trining input arguments to console."""
if out_type == "rst":
doc_str = gen_doc(make_anchor=True)
elif out_type == "json":
doc_str = gen_json()
else:
raise RuntimeError("Unsupported out type %s" % out_type)
print(doc_str)
_skbuild/linux-x86_64-3.6/cmake-install/deepmd/entrypoints/freeze.py deleted 100644 → 0
View file @ 6b33aeb8
#!/usr/bin/env python3
"""Script for freezing TF trained graph so it can be used with LAMMPS and i-PI.
References
----------
https://blog.metaflow.fr/tensorflow-how-to-freeze-a-model-and-serve-it-with-a-python-api-d4f3596b3adc
"""
import logging
import google.protobuf.message
from deepmd.env import tf, FITTING_NET_PATTERN
from deepmd.utils.errors import GraphTooLargeError
from deepmd.utils.sess import run_sess
from deepmd.utils.graph import get_pattern_nodes_from_graph_def
from os.path import abspath
# load grad of force module
import deepmd.op
from typing import List, Optional
from deepmd.nvnmd.entrypoints.freeze import save_weight
__all__ = ["freeze"]
log = logging.getLogger(__name__)
def _transfer_fitting_net_trainable_variables(sess, old_graph_def, raw_graph_def):
old_pattern = FITTING_NET_PATTERN
raw_pattern = FITTING_NET_PATTERN\
.replace('idt', 'idt+_\d+')\
.replace('bias', 'bias+_\d+')\
.replace('matrix', 'matrix+_\d+')
old_graph_nodes = get_pattern_nodes_from_graph_def(
old_graph_def,
old_pattern
)
try :
raw_graph_def = tf.graph_util.convert_variables_to_constants(
sess, # The session is used to retrieve the weights
raw_graph_def, # The graph_def is used to retrieve the nodes
[n + '_1' for n in old_graph_nodes], # The output node names are used to select the usefull nodes
)
except AssertionError:
# if there's no additional nodes
return old_graph_def
raw_graph_nodes = get_pattern_nodes_from_graph_def(
raw_graph_def,
raw_pattern
)
for node in old_graph_def.node:
if node.name not in old_graph_nodes.keys():
continue
tensor = tf.make_ndarray(raw_graph_nodes[node.name + '_1'])
node.attr["value"].tensor.tensor_content = tensor.tostring()
return old_graph_def
def _make_node_names(model_type: str, modifier_type: Optional[str] = None) -> List[str]:
"""Get node names based on model type.
Parameters
----------
model_type : str
str type of model
modifier_type : Optional[str], optional
modifier type if any, by default None
Returns
-------
List[str]
list with all node names to freeze
Raises
------
RuntimeError
if unknown model type
"""
nodes = [
"model_type",
"descrpt_attr/rcut",
"descrpt_attr/ntypes",
"model_attr/tmap",
"model_attr/model_type",
"model_attr/model_version",
"train_attr/min_nbor_dist",
"train_attr/training_script",
]
if model_type == "ener":
nodes += [
"o_energy",
"o_force",
"o_virial",
"o_atom_energy",
"o_atom_virial",
"fitting_attr/dfparam",
"fitting_attr/daparam",
]
elif model_type == "wfc":
nodes += [
"o_wfc",
"model_attr/sel_type",
"model_attr/output_dim",
]
elif model_type == "dipole":
nodes += [
"o_dipole",
"o_global_dipole",
"o_force",
"o_virial",
"o_atom_virial",
"o_rmat",
"o_rmat_deriv",
"o_nlist",
"o_rij",
"descrpt_attr/sel",
"descrpt_attr/ndescrpt",
"model_attr/sel_type",
"model_attr/output_dim",
]
elif model_type == "polar":
nodes += [
"o_polar",
"o_global_polar",
"o_force",
"o_virial",
"o_atom_virial",
"model_attr/sel_type",
"model_attr/output_dim",
]
elif model_type == "global_polar":
nodes += [
"o_global_polar",
"model_attr/sel_type",
"model_attr/output_dim",
]
else:
raise RuntimeError(f"unknow model type {model_type}")
if modifier_type == "dipole_charge":
nodes += [
"modifier_attr/type",
"modifier_attr/mdl_name",
"modifier_attr/mdl_charge_map",
"modifier_attr/sys_charge_map",
"modifier_attr/ewald_h",
"modifier_attr/ewald_beta",
"dipole_charge/model_type",
"dipole_charge/descrpt_attr/rcut",
"dipole_charge/descrpt_attr/ntypes",
"dipole_charge/model_attr/tmap",
"dipole_charge/model_attr/model_type",
"dipole_charge/model_attr/model_version",
"o_dm_force",
"dipole_charge/model_attr/sel_type",
"dipole_charge/o_dipole",
"dipole_charge/model_attr/output_dim",
"o_dm_virial",
"o_dm_av",
]
return nodes
def freeze(
*, checkpoint_folder: str, output: str, node_names: Optional[str] = None, nvnmd_weight: Optional[str] = None, **kwargs
):
"""Freeze the graph in supplied folder.
Parameters
----------
checkpoint_folder : str
location of the folder with model
output : str
output file name
node_names : Optional[str], optional
names of nodes to output, by default None
"""
# We retrieve our checkpoint fullpath
checkpoint = tf.train.get_checkpoint_state(checkpoint_folder)
input_checkpoint = checkpoint.model_checkpoint_path
# expand the output file to full path
output_graph = abspath(output)
# Before exporting our graph, we need to precise what is our output node
# This is how TF decides what part of the Graph he has to keep
# and what part it can dump
# NOTE: this variable is plural, because you can have multiple output nodes
# node_names = "energy_test,force_test,virial_test,t_rcut"
# We clear devices to allow TensorFlow to control
# on which device it will load operations
clear_devices = True
# We import the meta graph and retrieve a Saver
try:
# In case paralle training
import horovod.tensorflow as _
except ImportError:
pass
saver = tf.train.import_meta_graph(
f"{input_checkpoint}.meta", clear_devices=clear_devices
)
# We retrieve the protobuf graph definition
graph = tf.get_default_graph()
try:
input_graph_def = graph.as_graph_def()
except google.protobuf.message.DecodeError as e:
raise GraphTooLargeError(
"The graph size exceeds 2 GB, the hard limitation of protobuf."
" Then a DecodeError was raised by protobuf. You should "
"reduce the size of your model."
) from e
nodes = [n.name for n in input_graph_def.node]
# We start a session and restore the graph weights
with tf.Session() as sess:
saver.restore(sess, input_checkpoint)
model_type = run_sess(sess, "model_attr/model_type:0", feed_dict={}).decode("utf-8")
if "modifier_attr/type" in nodes:
modifier_type = run_sess(sess, "modifier_attr/type:0", feed_dict={}).decode(
"utf-8"
)
else:
modifier_type = None
if node_names is None:
output_node_list = _make_node_names(model_type, modifier_type)
different_set = set(output_node_list) - set(nodes)
if different_set:
log.warning(
"The following nodes are not in the graph: %s. "
"Skip freezeing these nodes. You may be freezing "
"a checkpoint generated by an old version." % different_set
)
# use intersection as output list
output_node_list = list(set(output_node_list) & set(nodes))
else:
output_node_list = node_names.split(",")
log.info(f"The following nodes will be frozen: {output_node_list}")
if nvnmd_weight is not None:
save_weight(sess, nvnmd_weight) # nvnmd
# We use a built-in TF helper to export variables to constants
output_graph_def = tf.graph_util.convert_variables_to_constants(
sess, # The session is used to retrieve the weights
input_graph_def, # The graph_def is used to retrieve the nodes
output_node_list, # The output node names are used to select the usefull nodes
)
# If we need to transfer the fitting net variables
output_graph_def = _transfer_fitting_net_trainable_variables(
sess,
output_graph_def,
input_graph_def
)
# Finally we serialize and dump the output graph to the filesystem
with tf.gfile.GFile(output_graph, "wb") as f:
f.write(output_graph_def.SerializeToString())
log.info(f"{len(output_graph_def.node):d} ops in the final graph.")
_skbuild/linux-x86_64-3.6/cmake-install/deepmd/entrypoints/main.py deleted 100644 → 0
View file @ 6b33aeb8
"""DeePMD-Kit entry point module."""
import argparse
import logging
import textwrap
from pathlib import Path
from typing import Dict, List, Optional
from deepmd import __version__
from deepmd.entrypoints import (
compress,
config,
doc_train_input,
freeze,
test,
train_dp,
transfer,
make_model_devi,
convert,
neighbor_stat,
)
from deepmd.loggers import set_log_handles
from deepmd.nvnmd.entrypoints.train import train_nvnmd
__all__ = ["main", "parse_args", "get_ll", "main_parser"]
def get_ll(log_level: str) -> int:
"""Convert string to python logging level.
Parameters
----------
log_level : str
allowed input values are: DEBUG, INFO, WARNING, ERROR, 3, 2, 1, 0
Returns
-------
int
one of python logging module log levels - 10, 20, 30 or 40
"""
if log_level.isdigit():
int_level = (4 - int(log_level)) * 10
else:
int_level = getattr(logging, log_level)
return int_level
class RawTextArgumentDefaultsHelpFormatter(
argparse.RawTextHelpFormatter, argparse.ArgumentDefaultsHelpFormatter
):
"""This formatter is used to print multile-line help message with default value."""
def main_parser() -> argparse.ArgumentParser:
"""DeePMD-Kit commandline options argument parser.
Returns
-------
argparse.ArgumentParser
main parser of DeePMD-kit
"""
parser = argparse.ArgumentParser(
description="DeePMD-kit: A deep learning package for many-body potential energy"
" representation and molecular dynamics",
formatter_class=argparse.ArgumentDefaultsHelpFormatter,
)
subparsers = parser.add_subparsers(title="Valid subcommands", dest="command")
# * logging options parser *********************************************************
# with use of the parent argument this options will be added to every parser
parser_log = argparse.ArgumentParser(
add_help=False, formatter_class=argparse.ArgumentDefaultsHelpFormatter
)
parser_log.add_argument(
"-v",
"--log-level",
choices=["DEBUG", "3", "INFO", "2", "WARNING", "1", "ERROR", "0"],
default="INFO",
help="set verbosity level by string or number, 0=ERROR, 1=WARNING, 2=INFO "
"and 3=DEBUG",
)
parser_log.add_argument(
"-l",
"--log-path",
type=str,
default=None,
help="set log file to log messages to disk, if not specified, the logs will "
"only be output to console",
)
# * mpi logging parser *************************************************************
parser_mpi_log = argparse.ArgumentParser(
add_help=False, formatter_class=argparse.ArgumentDefaultsHelpFormatter
)
parser_mpi_log.add_argument(
"-m",
"--mpi-log",
type=str,
default="master",
choices=("master", "collect", "workers"),
help="Set the manner of logging when running with MPI. 'master' logs only on "
"main process, 'collect' broadcasts logs from workers to master and 'workers' "
"means each process will output its own log",
)
# * config script ******************************************************************
parser_cfig = subparsers.add_parser(
"config",
parents=[parser_log],
help="fast configuration of parameter file for smooth model",
formatter_class=argparse.ArgumentDefaultsHelpFormatter,
)
parser_cfig.add_argument(
"-o", "--output", type=str, default="input.json", help="the output json file"
)
# * transfer script ****************************************************************
parser_transfer = subparsers.add_parser(
"transfer", parents=[parser_log], help="pass parameters to another model"
)
parser_transfer.add_argument(
"-r",
"--raw-model",
default="raw_frozen_model.pb",
type=str,
help="the model receiving parameters",
)
parser_transfer.add_argument(
"-O",
"--old-model",
default="old_frozen_model.pb",
type=str,
help="the model providing parameters",
)
parser_transfer.add_argument(
"-o",
"--output",
default="frozen_model.pb",
type=str,
help="the model after passing parameters",
)
# * config parser ******************************************************************
parser_train = subparsers.add_parser(
"train",
parents=[parser_log, parser_mpi_log],
help="train a model",
formatter_class=RawTextArgumentDefaultsHelpFormatter,
epilog=textwrap.dedent("""\
examples:
dp train input.json
dp train input.json --restart model.ckpt
dp train input.json --init-model model.ckpt
"""),
)
parser_train.add_argument(
"INPUT", help="the input parameter file in json or yaml format"
)
parser_train.add_argument(
"-i",
"--init-model",
type=str,
default=None,
help="Initialize the model by the provided checkpoint.",
)
parser_train.add_argument(
"-r",
"--restart",
type=str,
default=None,
help="Restart the training from the provided checkpoint.",
)
parser_train.add_argument(
"-o",
"--output",
type=str,
default="out.json",
help="The output file of the parameters used in training.",
)
parser_train.add_argument(
"-f",
"--init-frz-model",
type=str,
default=None,
help="Initialize the training from the frozen model.",
)
parser_train.add_argument(
"--skip-neighbor-stat",
action="store_true",
help="Skip calculating neighbor statistics. Sel checking, automatic sel, and model compression will be disabled.",
)
# * freeze script ******************************************************************
parser_frz = subparsers.add_parser(
"freeze",
parents=[parser_log],
help="freeze the model",
formatter_class=RawTextArgumentDefaultsHelpFormatter,
epilog=textwrap.dedent("""\
examples:
dp freeze
dp freeze -o graph.pb
"""),
)
parser_frz.add_argument(
"-c",
"--checkpoint-folder",
type=str,
default=".",
help="path to checkpoint folder",
)
parser_frz.add_argument(
"-o",
"--output",
type=str,
default="frozen_model.pb",
help="name of graph, will output to the checkpoint folder",
)
parser_frz.add_argument(
"-n",
"--node-names",
type=str,
default=None,
help="the frozen nodes, if not set, determined from the model type",
)
parser_frz.add_argument(
"-w",
"--nvnmd-weight",
type=str,
default=None,
help="the name of weight file (.npy), if set, save the model's weight into the file",
)
# * test script ********************************************************************
parser_tst = subparsers.add_parser(
"test",
parents=[parser_log],
help="test the model",
formatter_class=RawTextArgumentDefaultsHelpFormatter,
epilog=textwrap.dedent("""\
examples:
dp test -m graph.pb -s /path/to/system -n 30
"""),
)
parser_tst.add_argument(
"-m",
"--model",
default="frozen_model.pb",
type=str,
help="Frozen model file to import",
)
parser_tst.add_argument(
"-s",
"--system",
default=".",
type=str,
help="The system dir. Recursively detect systems in this directory",
)
parser_tst.add_argument(
"-S", "--set-prefix", default="set", type=str, help="The set prefix"
)
parser_tst.add_argument(
"-n", "--numb-test", default=100, type=int, help="The number of data for test"
)
parser_tst.add_argument(
"-r", "--rand-seed", type=int, default=None, help="The random seed"
)
parser_tst.add_argument(
"--shuffle-test", action="store_true", default=False, help="Shuffle test data"
)
parser_tst.add_argument(
"-d",
"--detail-file",
type=str,
default=None,
help="File where details of energy force and virial accuracy will be written",
)
parser_tst.add_argument(
"-a",
"--atomic",
action="store_true",
default=False,
help="Test the accuracy of atomic label, i.e. energy / tensor (dipole, polar)",
)
# * compress model *****************************************************************
# Compress a model, which including tabulating the embedding-net.
# The table is composed of fifth-order polynomial coefficients and is assembled
# from two sub-tables. The first table takes the step(parameter) as it's uniform
# step, while the second table takes 10 * step as it\s uniform step
#  The range of the first table is automatically detected by deepmd-kit, while the
# second table ranges from the first table's upper boundary(upper) to the
# extrapolate(parameter) * upper.
parser_compress = subparsers.add_parser(
"compress",
parents=[parser_log, parser_mpi_log],
help="compress a model",
formatter_class=RawTextArgumentDefaultsHelpFormatter,
epilog=textwrap.dedent("""\
examples:
dp compress
dp compress -i graph.pb -o compressed.pb
"""),
)
parser_compress.add_argument(
"-i",
"--input",
default="frozen_model.pb",
type=str,
help="The original frozen model, which will be compressed by the code",
)
parser_compress.add_argument(
"-o",
"--output",
default="frozen_model_compressed.pb",
type=str,
help="The compressed model",
)
parser_compress.add_argument(
"-s",
"--step",
default=0.01,
type=float,
help="Model compression uses fifth-order polynomials to interpolate the embedding-net. "
"It introduces two tables with different step size to store the parameters of the polynomials. "
"The first table covers the range of the training data, while the second table is an extrapolation of the training data. "
"The domain of each table is uniformly divided by a given step size. "
"And the step(parameter) denotes the step size of the first table and the second table will "
"use 10 * step as it's step size to save the memory. "
"Usually the value ranges from 0.1 to 0.001. "
"Smaller step means higher accuracy and bigger model size",
)
parser_compress.add_argument(
"-e",
"--extrapolate",
default=5,
type=int,
help="The domain range of the first table is automatically detected by the code: [d_low, d_up]. "
"While the second table ranges from the first table's upper boundary(d_up) to the extrapolate(parameter) * d_up: [d_up, extrapolate * d_up]",
)
parser_compress.add_argument(
"-f",
"--frequency",
default=-1,
type=int,
help="The frequency of tabulation overflow check(Whether the input environment "
"matrix overflow the first or second table range). "
"By default do not check the overflow",
)
parser_compress.add_argument(
"-c",
"--checkpoint-folder",
type=str,
default="model-compression",
help="path to checkpoint folder",
)
parser_compress.add_argument(
"-t",
"--training-script",
type=str,
default=None,
help="The training script of the input frozen model",
)
# * print docs script **************************************************************
parsers_doc = subparsers.add_parser(
"doc-train-input",
parents=[parser_log],
help="print the documentation (in rst format) of input training parameters.",
formatter_class=argparse.ArgumentDefaultsHelpFormatter,
)
parsers_doc.add_argument(
"--out-type",
default="rst",
type=str,
help="The output type"
)
# * make model deviation ***********************************************************
parser_model_devi = subparsers.add_parser(
"model-devi",
parents=[parser_log],
help="calculate model deviation",
formatter_class=RawTextArgumentDefaultsHelpFormatter,
epilog=textwrap.dedent("""\
examples:
dp model-devi -m graph.000.pb graph.001.pb graph.002.pb graph.003.pb -s ./data -o model_devi.out
"""),
)
parser_model_devi.add_argument(
"-m",
"--models",
default=["graph.000.pb", "graph.001.pb", "graph.002.pb", "graph.003.pb"],
nargs="+",
type=str,
help="Frozen models file to import",
)
parser_model_devi.add_argument(
"-s",
"--system",
default=".",
type=str,
help="The system directory. Recursively detect systems in this directory.",
)
parser_model_devi.add_argument(
"-S", "--set-prefix", default="set", type=str, help="The set prefix"
)
parser_model_devi.add_argument(
"-o",
"--output",
default="model_devi.out",
type=str,
help="The output file for results of model deviation"
)
parser_model_devi.add_argument(
"-f",
"--frequency",
default=1,
type=int,
help="The trajectory frequency of the system"
)
# * convert models
parser_transform = subparsers.add_parser(
'convert-from',
parents=[parser_log],
help='convert lower model version to supported version',
formatter_class=RawTextArgumentDefaultsHelpFormatter,
epilog=textwrap.dedent("""\
examples:
dp convert-from 1.0 -i graph.pb -o graph_new.pb
"""),
)
parser_transform.add_argument(
'FROM',
type = str,
choices = ['0.12', '1.0', '1.1', '1.2', '1.3', '2.0'],
help="The original model compatibility",
)
parser_transform.add_argument(
'-i',
"--input-model",
default = "frozen_model.pb",
type=str,
help = "the input model",
)
parser_transform.add_argument(
"-o",
"--output-model",
default = "convert_out.pb",
type=str,
help='the output model',
)
# neighbor_stat
parser_neighbor_stat = subparsers.add_parser(
'neighbor-stat',
parents=[parser_log],
help='Calculate neighbor statistics',
formatter_class=RawTextArgumentDefaultsHelpFormatter,
epilog=textwrap.dedent("""\
examples:
dp neighbor-stat -s data -r 6.0 -t O H
"""),
)
parser_neighbor_stat.add_argument(
"-s",
"--system",
default=".",
type=str,
help="The system dir. Recursively detect systems in this directory",
)
parser_neighbor_stat.add_argument(
"-r",
"--rcut",
type=float,
required=True,
help="cutoff radius",
)
parser_neighbor_stat.add_argument(
"-t",
"--type-map",
type=str,
nargs='+',
required=True,
help="type map",
)
parser_neighbor_stat.add_argument(
"--one-type",
action="store_true",
default=False,
help="treat all types as a single type. Used with se_atten descriptor.",
)
# --version
parser.add_argument('--version', action='version', version='DeePMD-kit v%s' % __version__)
# * train nvnmd script ******************************************************************
parser_train_nvnmd = subparsers.add_parser(
"train-nvnmd",
parents=[parser_log],
help="train nvnmd model",
formatter_class=argparse.ArgumentDefaultsHelpFormatter,
)
parser_train_nvnmd.add_argument(
"INPUT", help="the input parameter file in json format"
)
parser_train_nvnmd.add_argument(
"-s",
"--step",
default="s1",
type=str,
choices=['s1', 's2'],
help="steps to train model of NVNMD: s1 (train CNN), s2 (train QNN)"
)
return parser
def parse_args(args: Optional[List[str]] = None) -> argparse.Namespace:
"""Parse arguments and convert argument strings to objects.
Parameters
----------
args: List[str]
list of command line arguments, main purpose is testing default option None
takes arguments from sys.argv
Returns
-------
argparse.Namespace
the populated namespace
"""
parser = main_parser()
parsed_args = parser.parse_args(args=args)
if parsed_args.command is None:
parser.print_help()
else:
parsed_args.log_level = get_ll(parsed_args.log_level)
return parsed_args
def main():
"""DeePMD-Kit entry point.
Raises
------
RuntimeError
if no command was input
"""
args = parse_args()
# do not set log handles for None, it is useless
# log handles for train will be set separatelly
# when the use of MPI will be determined in `RunOptions`
if args.command not in (None, "train"):
set_log_handles(args.log_level, Path(args.log_path) if args.log_path else None)
dict_args = vars(args)
if args.command == "train":
train_dp(**dict_args)
elif args.command == "freeze":
freeze(**dict_args)
elif args.command == "config":
config(**dict_args)
elif args.command == "test":
test(**dict_args)
elif args.command == "transfer":
transfer(**dict_args)
elif args.command == "compress":
compress(**dict_args)
elif args.command == "doc-train-input":
doc_train_input(**dict_args)
elif args.command == "model-devi":
make_model_devi(**dict_args)
elif args.command == "convert-from":
convert(**dict_args)
elif args.command == "neighbor-stat":
neighbor_stat(**dict_args)
elif args.command == "train-nvnmd": # nvnmd
train_nvnmd(**dict_args)
elif args.command is None:
pass
else:
raise RuntimeError(f"unknown command {args.command}")
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment