Skip to content

GitLab

Commit 6b33aeb8 authored by zhangqha's avatar zhangqha
Browse files

BladeDISC DeePMD code

parents
Pipeline #179 canceled with stages
Hide whitespace changes
Inline Side-by-side
Showing with 6136 additions and 0 deletions
_skbuild/linux-x86_64-3.6/setuptools/lib/deepmd/__init__.py 0 → 100644
View file @ 6b33aeb8
"""Root of the deepmd package, exposes all public classes and submodules."""
try:
from importlib import metadata
except ImportError: # for Python<3.8
import importlib_metadata as metadata
import deepmd.utils.network as network
from . import cluster, descriptor, fit, loss, utils, nvnmd
from .env import set_mkl
from .infer import DeepEval, DeepPotential
from .infer.data_modifier import DipoleChargeModifier
set_mkl()
try:
from ._version import version as __version__
except ImportError:
from .__about__ import __version__
# load third-party plugins
for ep in metadata.entry_points().get('deepmd', []):
ep.load()
__all__ = [
"descriptor",
"fit",
"loss",
"utils",
"cluster",
"network",
"DeepEval",
"DeepPotential",
"DipoleChargeModifier",
"nvnmd",
]
_skbuild/linux-x86_64-3.6/setuptools/lib/deepmd/__main__.py 0 → 100644
View file @ 6b33aeb8
"""Package dp entry point."""
from .entrypoints.main import main
if __name__ == '__main__':
main()
_skbuild/linux-x86_64-3.6/setuptools/lib/deepmd/_version.py 0 → 100644
View file @ 6b33aeb8
# coding: utf-8
# file generated by setuptools_scm
# don't change, don't track in version control
version = '2.1.6.dev0+ng6e3d4a62.d20230313'
_skbuild/linux-x86_64-3.6/setuptools/lib/deepmd/calculator.py 0 → 100644
View file @ 6b33aeb8
"""ASE calculator interface module."""
from pathlib import Path
from typing import TYPE_CHECKING, Dict, List, Optional, Union
from ase.calculators.calculator import (
Calculator, all_changes, PropertyNotImplementedError
)
from deepmd import DeepPotential
if TYPE_CHECKING:
from ase import Atoms
__all__ = ["DP"]
class DP(Calculator):
"""Implementation of ASE deepmd calculator.
Implemented propertie are `energy`, `forces` and `stress`
Parameters
----------
model : Union[str, Path]
path to the model
label : str, optional
calculator label, by default "DP"
type_dict : Dict[str, int], optional
mapping of element types and their numbers, best left None and the calculator
will infer this information from model, by default None
Examples
--------
Compute potential energy
>>> from ase import Atoms
>>> from deepmd.calculator import DP
>>> water = Atoms('H2O',
>>> positions=[(0.7601, 1.9270, 1),
>>> (1.9575, 1, 1),
>>> (1., 1., 1.)],
>>> cell=[100, 100, 100],
>>> calculator=DP(model="frozen_model.pb"))
>>> print(water.get_potential_energy())
>>> print(water.get_forces())
Run BFGS structure optimization
>>> from ase.optimize import BFGS
>>> dyn = BFGS(water)
>>> dyn.run(fmax=1e-6)
>>> print(water.get_positions())
"""
name = "DP"
implemented_properties = ["energy", "free_energy", "forces", "virial", "stress"]
def __init__(
self,
model: Union[str, "Path"],
label: str = "DP",
type_dict: Dict[str, int] = None,
**kwargs
) -> None:
Calculator.__init__(self, label=label, **kwargs)
self.dp = DeepPotential(str(Path(model).resolve()))
if type_dict:
self.type_dict = type_dict
else:
self.type_dict = dict(
zip(self.dp.get_type_map(), range(self.dp.get_ntypes()))
)
def calculate(
self,
atoms: Optional["Atoms"] = None,
properties: List[str] = ["energy", "forces", "virial"],
system_changes: List[str] = all_changes,
):
"""Run calculation with deepmd model.
Parameters
----------
atoms : Optional[Atoms], optional
atoms object to run the calculation on, by default None
properties : List[str], optional
unused, only for function signature compatibility,
by default ["energy", "forces", "stress"]
system_changes : List[str], optional
unused, only for function signature compatibility, by default all_changes
"""
if atoms is not None:
self.atoms = atoms.copy()
coord = self.atoms.get_positions().reshape([1, -1])
if sum(self.atoms.get_pbc()) > 0:
cell = self.atoms.get_cell().reshape([1, -1])
else:
cell = None
symbols = self.atoms.get_chemical_symbols()
atype = [self.type_dict[k] for k in symbols]
e, f, v = self.dp.eval(coords=coord, cells=cell, atom_types=atype)
self.results['energy'] = e[0][0]
# see https://gitlab.com/ase/ase/-/merge_requests/2485
self.results['free_energy'] = e[0][0]
self.results['forces'] = f[0]
self.results['virial'] = v[0].reshape(3, 3)
# convert virial into stress for lattice relaxation
if "stress" in properties:
if sum(atoms.get_pbc()) > 0:
# the usual convention (tensile stress is positive)
# stress = -virial / volume
stress = -0.5 * (v[0].copy() + v[0].copy().T) / atoms.get_volume()
# Voigt notation
self.results['stress'] = stress.flat[[0, 4, 8, 5, 2, 1]]
else:
raise PropertyNotImplementedError
_skbuild/linux-x86_64-3.6/setuptools/lib/deepmd/cluster/__init__.py 0 → 100644
View file @ 6b33aeb8
"""Module that reads node resources, auto detects if running local or on SLURM."""
from .local import get_resource as get_local_res
from .slurm import get_resource as get_slurm_res
import os
from typing import List, Tuple, Optional
__all__ = ["get_resource"]
def get_resource() -> Tuple[str, List[str], Optional[List[int]]]:
"""Get local or slurm resources: nodename, nodelist, and gpus.
Returns
-------
Tuple[str, List[str], Optional[List[int]]]
nodename, nodelist, and gpus
"""
if "SLURM_JOB_NODELIST" in os.environ:
return get_slurm_res()
else:
return get_local_res()
_skbuild/linux-x86_64-3.6/setuptools/lib/deepmd/cluster/local.py 0 → 100644
View file @ 6b33aeb8
"""Get local GPU resources."""
import os
import socket
import subprocess as sp
import sys
from deepmd.env import tf
from typing import List, Tuple, Optional
__all__ = ["get_gpus", "get_resource"]
def get_gpus():
"""Get available IDs of GPU cards at local.
These IDs are valid when used as the TensorFlow device ID.
Returns
-------
Optional[List[int]]
List of available GPU IDs. Otherwise, None.
"""
test_cmd = 'from tensorflow.python.client import device_lib; ' \
'devices = device_lib.list_local_devices(); ' \
'gpus = [d.name for d in devices if d.device_type == "GPU"]; ' \
'print(len(gpus))'
with sp.Popen([sys.executable, "-c", test_cmd], stderr=sp.PIPE, stdout=sp.PIPE) as p:
stdout, stderr = p.communicate()
if p.returncode != 0:
decoded = stderr.decode('UTF-8')
raise RuntimeError('Failed to detect availbe GPUs due to:\n%s' % decoded)
decoded = stdout.decode('UTF-8').strip()
num_gpus = int(decoded)
return list(range(num_gpus)) if num_gpus > 0 else None
def get_resource() -> Tuple[str, List[str], Optional[List[int]]]:
"""Get local resources: nodename, nodelist, and gpus.
Returns
-------
Tuple[str, List[str], Optional[List[int]]]
nodename, nodelist, and gpus
"""
nodename = socket.gethostname()
nodelist = [nodename]
gpus = get_gpus()
return nodename, nodelist, gpus
_skbuild/linux-x86_64-3.6/setuptools/lib/deepmd/cluster/slurm.py 0 → 100644
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/setuptools/lib/deepmd/common.py 0 → 100644
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/setuptools/lib/deepmd/descriptor/__init__.py 0 → 100644
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/setuptools/lib/deepmd/descriptor/descriptor.py 0 → 100644
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/setuptools/lib/deepmd/descriptor/hybrid.py 0 → 100644
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/setuptools/lib/deepmd/descriptor/loc_frame.py 0 → 100644
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/setuptools/lib/deepmd/descriptor/se.py 0 → 100644
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/setuptools/lib/deepmd/descriptor/se_a.py 0 → 100644
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/setuptools/lib/deepmd/descriptor/se_a_ebd.py 0 → 100644
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/setuptools/lib/deepmd/descriptor/se_a_ef.py 0 → 100644
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/setuptools/lib/deepmd/descriptor/se_atten.py 0 → 100644
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/setuptools/lib/deepmd/descriptor/se_r.py 0 → 100644
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/setuptools/lib/deepmd/descriptor/se_t.py 0 → 100644
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/setuptools/lib/deepmd/entrypoints/__init__.py 0 → 100644
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",
]
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