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 5101 additions and 0 deletions
_skbuild/linux-x86_64-3.6/setuptools/lib/deepmd/entrypoints/compress.py 0 → 100644
View file @ 6b33aeb8
"""Compress a model, which including tabulating the embedding-net."""
import os
import json
import logging
from typing import Optional
from deepmd.common import j_loader
from deepmd.env import tf, GLOBAL_ENER_FLOAT_PRECISION
from deepmd.utils.argcheck import normalize
from deepmd.utils.compat import update_deepmd_input
from deepmd.utils.errors import GraphTooLargeError, GraphWithoutTensorError
from deepmd.utils.graph import get_tensor_by_name
from .freeze import freeze
from .train import train, get_rcut, get_min_nbor_dist
from .transfer import transfer
__all__ = ["compress"]
log = logging.getLogger(__name__)
def compress(
*,
input: str,
output: str,
extrapolate: int,
step: float,
frequency: str,
checkpoint_folder: str,
training_script: str,
mpi_log: str,
log_path: Optional[str],
log_level: int,
**kwargs
):
"""Compress model.
The table is composed of fifth-order polynomial coefficients and is assembled from
two sub-tables. The first table takes the step parameter as the domain's uniform step size,
while the second table takes 10 * step as it's uniform step size. The range of the
first table is automatically detected by the code, while the second table ranges
from the first table's upper boundary(upper) to the extrapolate(parameter) * upper.
Parameters
----------
input : str
frozen model file to compress
output : str
compressed model filename
extrapolate : int
scale of model extrapolation
step : float
uniform step size of the tabulation's first table
frequency : str
frequency of tabulation overflow check
checkpoint_folder : str
trining checkpoint folder for freezing
training_script : str
training script of the input frozen model
mpi_log : str
mpi logging mode for training
log_path : Optional[str]
if speccified log will be written to this file
log_level : int
logging level
"""
try:
t_jdata = get_tensor_by_name(input, 'train_attr/training_script')
t_min_nbor_dist = get_tensor_by_name(input, 'train_attr/min_nbor_dist')
jdata = json.loads(t_jdata)
except GraphWithoutTensorError as e:
if training_script == None:
raise RuntimeError(
"The input frozen model: %s has no training script or min_nbor_dist information, "
"which is not supported by the model compression interface. "
"Please consider using the --training-script command within the model compression interface to provide the training script of the input frozen model. "
"Note that the input training script must contain the correct path to the training data." % input
) from e
elif not os.path.exists(training_script):
raise RuntimeError(
"The input training script %s (%s) does not exist! Please check the path of the training script. " % (input, os.path.abspath(input))
) from e
else:
log.info("stage 0: compute the min_nbor_dist")
jdata = j_loader(training_script)
jdata = update_deepmd_input(jdata)
t_min_nbor_dist = get_min_nbor_dist(jdata, get_rcut(jdata))
_check_compress_type(input)
tf.constant(t_min_nbor_dist,
name = 'train_attr/min_nbor_dist',
dtype = GLOBAL_ENER_FLOAT_PRECISION)
jdata["model"]["compress"] = {}
jdata["model"]["compress"]["model_file"] = input
jdata["model"]["compress"]["min_nbor_dist"] = t_min_nbor_dist
jdata["model"]["compress"]["table_config"] = [
extrapolate,
step,
10 * step,
int(frequency),
]
jdata["training"]["save_ckpt"] = os.path.join("model-compression", "model.ckpt")
jdata = update_deepmd_input(jdata)
jdata = normalize(jdata)
# check the descriptor info of the input file
# move to the specific Descriptor class
# stage 1: training or refining the model with tabulation
log.info("\n\n")
log.info("stage 1: compress the model")
control_file = "compress.json"
with open(control_file, "w") as fp:
json.dump(jdata, fp, indent=4)
try:
train(
INPUT=control_file,
init_model=None,
restart=None,
init_frz_model=None,
output=control_file,
mpi_log=mpi_log,
log_level=log_level,
log_path=log_path,
is_compress=True,
)
except GraphTooLargeError as e:
raise RuntimeError(
"The uniform step size of the tabulation's first table is %f, "
"which is too small. This leads to a very large graph size, "
"exceeding protobuf's limitation (2 GB). You should try to "
"increase the step size." % step
) from e
# reset the graph, otherwise the size limitation will be only 2 GB / 2 = 1 GB
tf.reset_default_graph()
# stage 2: freeze the model
log.info("\n\n")
log.info("stage 2: freeze the model")
try:
freeze(checkpoint_folder=checkpoint_folder, output=output, node_names=None)
except GraphTooLargeError as e:
raise RuntimeError(
"The uniform step size of the tabulation's first table is %f, "
"which is too small. This leads to a very large graph size, "
"exceeding protobuf's limitation (2 GB). You should try to "
"increase the step size." % step
) from e
def _check_compress_type(model_file):
try:
t_model_type = bytes.decode(get_tensor_by_name(model_file, 'model_type'))
except GraphWithoutTensorError as e:
# Compatible with the upgraded model, which has no 'model_type' info
t_model_type = None
if t_model_type == "compressed_model":
raise RuntimeError("The input frozen model %s has already been compressed! Please do not compress the model repeatedly. " % model_file)
_skbuild/linux-x86_64-3.6/setuptools/lib/deepmd/entrypoints/config.py 0 → 100644
View file @ 6b33aeb8
#!/usr/bin/env python3
"""Quickly create a configuration file for smooth model."""
import json
import yaml
from pathlib import Path
from typing import Any, Dict, List, Tuple
import numpy as np
__all__ = ["config"]
DEFAULT_DATA: Dict[str, Any] = {
"use_smooth": True,
"sel_a": [],
"rcut_smth": -1,
"rcut": -1,
"filter_neuron": [20, 40, 80],
"filter_resnet_dt": False,
"axis_neuron": 8,
"fitting_neuron": [240, 240, 240],
"fitting_resnet_dt": True,
"coord_norm": True,
"type_fitting_net": False,
"systems": [],
"set_prefix": "set",
"stop_batch": -1,
"batch_size": -1,
"start_lr": 0.001,
"decay_steps": -1,
"decay_rate": 0.95,
"start_pref_e": 0.02,
"limit_pref_e": 1,
"start_pref_f": 1000,
"limit_pref_f": 1,
"start_pref_v": 0,
"limit_pref_v": 0,
"seed": 1,
"disp_file": "lcurve.out",
"disp_freq": 1000,
"numb_test": 10,
"save_freq": 10000,
"save_ckpt": "model.ckpt",
"disp_training": True,
"time_training": True,
}
def valid_dir(path: Path):
"""Check if directory is a valid deepmd system directory.
Parameters
----------
path : Path
path to directory
Raises
------
OSError
if `type.raw` is missing on dir or `box.npy` or `coord.npy` are missing in one
of the sets subdirs
"""
if not (path / "type.raw").is_file():
raise OSError
for ii in path.glob("set.*"):
if not (ii / "box.npy").is_file():
raise OSError
if not (ii / "coord.npy").is_file():
raise OSError
def load_systems(dirs: List[Path]) -> Tuple[List[np.ndarray], List[np.ndarray]]:
"""Load systems to memory for disk.
Parameters
----------
dirs : List[Path]
list of system directories paths
Returns
-------
Tuple[List[np.ndarray], List[np.ndarray]]
atoms types and structure cells formated as Nx9 array
"""
all_type = []
all_box = []
for d in dirs:
sys_type = np.loadtxt(d / "type.raw", dtype=int)
sys_box = np.vstack([np.load(s / "box.npy") for s in d.glob("set.*")])
all_type.append(sys_type)
all_box.append(sys_box)
return all_type, all_box
def get_system_names() -> List[Path]:
"""Get system directory paths from stdin.
Returns
-------
List[Path]
list of system directories paths
"""
dirs = input("Enter system path(s) (seperated by space, wild card supported): \n")
system_dirs = []
for dir_str in dirs.split():
found_dirs = Path.cwd().glob(dir_str)
for d in found_dirs:
valid_dir(d)
system_dirs.append(d)
return system_dirs
def get_rcut() -> float:
"""Get rcut from stdin from user.
Returns
-------
float
input rcut lenght converted to float
Raises
------
ValueError
if rcut is smaller than 0.0
"""
dv = 6.0
rcut_input = input(f"Enter rcut (default {dv:.1f} A): \n")
try:
rcut = float(rcut_input)
except ValueError as e:
print(f"invalid rcut: {e} setting to default: {dv:.1f}")
rcut = dv
if rcut <= 0:
raise ValueError("rcut should be > 0")
return rcut
def get_batch_size_rule() -> int:
"""Get minimal batch size from user from stdin.
Returns
-------
int
size of the batch
Raises
------
ValueError
if batch size is <= 0
"""
dv = 32
matom_input = input(
f"Enter the minimal number of atoms in a batch (default {dv:d}: \n"
)
try:
matom = int(matom_input)
except ValueError as e:
print(f"invalid batch size: {e} setting to default: {dv:d}")
matom = dv
if matom <= 0:
raise ValueError("the number should be > 0")
return matom
def get_stop_batch() -> int:
"""Get stop batch from user from stdin.
Returns
-------
int
size of the batch
Raises
------
ValueError
if stop batch is <= 0
"""
dv = 1000000
sb_input = input(f"Enter the stop batch (default {dv:d}): \n")
try:
sb = int(sb_input)
except ValueError as e:
print(f"invalid stop batch: {e} setting to default: {dv:d}")
sb = dv
if sb <= 0:
raise ValueError("the number should be > 0")
return sb
def get_ntypes(all_type: List[np.ndarray]) -> int:
"""Count number of unique elements.
Parameters
----------
all_type : List[np.ndarray]
list with arrays specifying elements of structures
Returns
-------
int
number of unique elements
"""
return len(np.unique(all_type))
def get_max_density(
all_type: List[np.ndarray], all_box: List[np.ndarray]
) -> np.ndarray:
"""Compute maximum density in suppliedd cells.
Parameters
----------
all_type : List[np.ndarray]
list with arrays specifying elements of structures
all_box : List[np.ndarray]
list with arrays specifying cells for all structures
Returns
-------
float
maximum atom density in all supplies structures for each element individually
"""
ntypes = get_ntypes(all_type)
all_max = []
for tt, bb in zip(all_type, all_box):
vv = np.reshape(bb, [-1, 3, 3])
vv = np.linalg.det(vv)
min_v = np.min(vv)
type_count = []
for ii in range(ntypes):
type_count.append(sum(tt == ii))
max_den = type_count / min_v
all_max.append(max_den)
all_max = np.max(all_max, axis=0)
return all_max
def suggest_sel(
all_type: List[np.ndarray],
all_box: List[np.ndarray],
rcut: float,
ratio: float = 1.5,
) -> List[int]:
"""Suggest selection parameter.
Parameters
----------
all_type : List[np.ndarray]
list with arrays specifying elements of structures
all_box : List[np.ndarray]
list with arrays specifying cells for all structures
rcut : float
cutoff radius
ratio : float, optional
safety margin to add to estimated value, by default 1.5
Returns
-------
List[int]
[description]
"""
max_den = get_max_density(all_type, all_box)
return [int(ii) for ii in max_den * 4.0 / 3.0 * np.pi * rcut ** 3 * ratio]
def suggest_batch_size(all_type: List[np.ndarray], min_atom: int) -> List[int]:
"""Get suggestion for batch size.
Parameters
----------
all_type : List[np.ndarray]
list with arrays specifying elements of structures
min_atom : int
minimal number of atoms in batch
Returns
-------
List[int]
suggested batch sizes for each system
"""
bs = []
for ii in all_type:
natoms = len(ii)
tbs = min_atom // natoms
if (min_atom // natoms) * natoms != min_atom:
tbs += 1
bs.append(tbs)
return bs
def suggest_decay(stop_batch: int) -> Tuple[int, float]:
"""Suggest number of decay steps and decay rate.
Parameters
----------
stop_batch : int
stop batch number
Returns
-------
Tuple[int, float]
number of decay steps and decay rate
"""
decay_steps = int(stop_batch // 200)
decay_rate = 0.95
return decay_steps, decay_rate
def config(*, output: str, **kwargs):
"""Auto config file generator.
Parameters
----------
output: str
file to write config file
Raises
------
RuntimeError
if user does not input any systems
ValueError
if output file is of wrong type
"""
all_sys = get_system_names()
if len(all_sys) == 0:
raise RuntimeError("no system specified")
rcut = get_rcut()
matom = get_batch_size_rule()
stop_batch = get_stop_batch()
all_type, all_box = load_systems(all_sys)
sel = suggest_sel(all_type, all_box, rcut, ratio=1.5)
bs = suggest_batch_size(all_type, matom)
decay_steps, decay_rate = suggest_decay(stop_batch)
jdata = DEFAULT_DATA.copy()
jdata["systems"] = [str(ii) for ii in all_sys]
jdata["sel_a"] = sel
jdata["rcut"] = rcut
jdata["rcut_smth"] = rcut - 0.2
jdata["stop_batch"] = stop_batch
jdata["batch_size"] = bs
jdata["decay_steps"] = decay_steps
jdata["decay_rate"] = decay_rate
with open(output, "w") as fp:
if output.endswith("json"):
json.dump(jdata, fp, indent=4)
elif output.endswith(("yml", "yaml")):
yaml.safe_dump(jdata, fp, default_flow_style=False)
else:
raise ValueError("output file must be of type json or yaml")
_skbuild/linux-x86_64-3.6/setuptools/lib/deepmd/entrypoints/convert.py 0 → 100644
View file @ 6b33aeb8
from deepmd.utils.convert import convert_012_to_21, convert_10_to_21, convert_20_to_21, convert_13_to_21, convert_12_to_21
def convert(
*,
FROM: str,
input_model: str,
output_model: str,
**kwargs,
):
if FROM == '0.12':
convert_012_to_21(input_model, output_model)
elif FROM == '1.0':
convert_10_to_21(input_model, output_model)
elif FROM in ['1.1', '1.2']:
# no difference between 1.1 and 1.2
convert_12_to_21(input_model, output_model)
elif FROM == '1.3':
convert_13_to_21(input_model, output_model)
elif FROM == '2.0':
convert_20_to_21(input_model, output_model)
else:
raise RuntimeError('unsupported model version ' + FROM)
_skbuild/linux-x86_64-3.6/setuptools/lib/deepmd/entrypoints/doc.py 0 → 100644
View file @ 6b33aeb8
"""Module that prints train input arguments docstrings."""
from deepmd.utils.argcheck import gen_doc, gen_json
__all__ = ["doc_train_input"]
def doc_train_input(*, out_type: str = "rst", **kwargs):
"""Print out trining input arguments to console."""
if out_type == "rst":
doc_str = gen_doc(make_anchor=True)
elif out_type == "json":
doc_str = gen_json()
else:
raise RuntimeError("Unsupported out type %s" % out_type)
print(doc_str)
_skbuild/linux-x86_64-3.6/setuptools/lib/deepmd/entrypoints/freeze.py 0 → 100644
View file @ 6b33aeb8
#!/usr/bin/env python3
"""Script for freezing TF trained graph so it can be used with LAMMPS and i-PI.
References
----------
https://blog.metaflow.fr/tensorflow-how-to-freeze-a-model-and-serve-it-with-a-python-api-d4f3596b3adc
"""
import logging
import google.protobuf.message
from deepmd.env import tf, FITTING_NET_PATTERN
from deepmd.utils.errors import GraphTooLargeError
from deepmd.utils.sess import run_sess
from deepmd.utils.graph import get_pattern_nodes_from_graph_def
from os.path import abspath
# load grad of force module
import deepmd.op
from typing import List, Optional
from deepmd.nvnmd.entrypoints.freeze import save_weight
__all__ = ["freeze"]
log = logging.getLogger(__name__)
def _transfer_fitting_net_trainable_variables(sess, old_graph_def, raw_graph_def):
old_pattern = FITTING_NET_PATTERN
raw_pattern = FITTING_NET_PATTERN\
.replace('idt', 'idt+_\d+')\
.replace('bias', 'bias+_\d+')\
.replace('matrix', 'matrix+_\d+')
old_graph_nodes = get_pattern_nodes_from_graph_def(
old_graph_def,
old_pattern
)
try :
raw_graph_def = tf.graph_util.convert_variables_to_constants(
sess, # The session is used to retrieve the weights
raw_graph_def, # The graph_def is used to retrieve the nodes
[n + '_1' for n in old_graph_nodes], # The output node names are used to select the usefull nodes
)
except AssertionError:
# if there's no additional nodes
return old_graph_def
raw_graph_nodes = get_pattern_nodes_from_graph_def(
raw_graph_def,
raw_pattern
)
for node in old_graph_def.node:
if node.name not in old_graph_nodes.keys():
continue
tensor = tf.make_ndarray(raw_graph_nodes[node.name + '_1'])
node.attr["value"].tensor.tensor_content = tensor.tostring()
return old_graph_def
def _make_node_names(model_type: str, modifier_type: Optional[str] = None) -> List[str]:
"""Get node names based on model type.
Parameters
----------
model_type : str
str type of model
modifier_type : Optional[str], optional
modifier type if any, by default None
Returns
-------
List[str]
list with all node names to freeze
Raises
------
RuntimeError
if unknown model type
"""
nodes = [
"model_type",
"descrpt_attr/rcut",
"descrpt_attr/ntypes",
"model_attr/tmap",
"model_attr/model_type",
"model_attr/model_version",
"train_attr/min_nbor_dist",
"train_attr/training_script",
]
if model_type == "ener":
nodes += [
"o_energy",
"o_force",
"o_virial",
"o_atom_energy",
"o_atom_virial",
"fitting_attr/dfparam",
"fitting_attr/daparam",
]
elif model_type == "wfc":
nodes += [
"o_wfc",
"model_attr/sel_type",
"model_attr/output_dim",
]
elif model_type == "dipole":
nodes += [
"o_dipole",
"o_global_dipole",
"o_force",
"o_virial",
"o_atom_virial",
"o_rmat",
"o_rmat_deriv",
"o_nlist",
"o_rij",
"descrpt_attr/sel",
"descrpt_attr/ndescrpt",
"model_attr/sel_type",
"model_attr/output_dim",
]
elif model_type == "polar":
nodes += [
"o_polar",
"o_global_polar",
"o_force",
"o_virial",
"o_atom_virial",
"model_attr/sel_type",
"model_attr/output_dim",
]
elif model_type == "global_polar":
nodes += [
"o_global_polar",
"model_attr/sel_type",
"model_attr/output_dim",
]
else:
raise RuntimeError(f"unknow model type {model_type}")
if modifier_type == "dipole_charge":
nodes += [
"modifier_attr/type",
"modifier_attr/mdl_name",
"modifier_attr/mdl_charge_map",
"modifier_attr/sys_charge_map",
"modifier_attr/ewald_h",
"modifier_attr/ewald_beta",
"dipole_charge/model_type",
"dipole_charge/descrpt_attr/rcut",
"dipole_charge/descrpt_attr/ntypes",
"dipole_charge/model_attr/tmap",
"dipole_charge/model_attr/model_type",
"dipole_charge/model_attr/model_version",
"o_dm_force",
"dipole_charge/model_attr/sel_type",
"dipole_charge/o_dipole",
"dipole_charge/model_attr/output_dim",
"o_dm_virial",
"o_dm_av",
]
return nodes
def freeze(
*, checkpoint_folder: str, output: str, node_names: Optional[str] = None, nvnmd_weight: Optional[str] = None, **kwargs
):
"""Freeze the graph in supplied folder.
Parameters
----------
checkpoint_folder : str
location of the folder with model
output : str
output file name
node_names : Optional[str], optional
names of nodes to output, by default None
"""
# We retrieve our checkpoint fullpath
checkpoint = tf.train.get_checkpoint_state(checkpoint_folder)
input_checkpoint = checkpoint.model_checkpoint_path
# expand the output file to full path
output_graph = abspath(output)
# Before exporting our graph, we need to precise what is our output node
# This is how TF decides what part of the Graph he has to keep
# and what part it can dump
# NOTE: this variable is plural, because you can have multiple output nodes
# node_names = "energy_test,force_test,virial_test,t_rcut"
# We clear devices to allow TensorFlow to control
# on which device it will load operations
clear_devices = True
# We import the meta graph and retrieve a Saver
try:
# In case paralle training
import horovod.tensorflow as _
except ImportError:
pass
saver = tf.train.import_meta_graph(
f"{input_checkpoint}.meta", clear_devices=clear_devices
)
# We retrieve the protobuf graph definition
graph = tf.get_default_graph()
try:
input_graph_def = graph.as_graph_def()
except google.protobuf.message.DecodeError as e:
raise GraphTooLargeError(
"The graph size exceeds 2 GB, the hard limitation of protobuf."
" Then a DecodeError was raised by protobuf. You should "
"reduce the size of your model."
) from e
nodes = [n.name for n in input_graph_def.node]
# We start a session and restore the graph weights
with tf.Session() as sess:
saver.restore(sess, input_checkpoint)
model_type = run_sess(sess, "model_attr/model_type:0", feed_dict={}).decode("utf-8")
if "modifier_attr/type" in nodes:
modifier_type = run_sess(sess, "modifier_attr/type:0", feed_dict={}).decode(
"utf-8"
)
else:
modifier_type = None
if node_names is None:
output_node_list = _make_node_names(model_type, modifier_type)
different_set = set(output_node_list) - set(nodes)
if different_set:
log.warning(
"The following nodes are not in the graph: %s. "
"Skip freezeing these nodes. You may be freezing "
"a checkpoint generated by an old version." % different_set
)
# use intersection as output list
output_node_list = list(set(output_node_list) & set(nodes))
else:
output_node_list = node_names.split(",")
log.info(f"The following nodes will be frozen: {output_node_list}")
if nvnmd_weight is not None:
save_weight(sess, nvnmd_weight) # nvnmd
# We use a built-in TF helper to export variables to constants
output_graph_def = tf.graph_util.convert_variables_to_constants(
sess, # The session is used to retrieve the weights
input_graph_def, # The graph_def is used to retrieve the nodes
output_node_list, # The output node names are used to select the usefull nodes
)
# If we need to transfer the fitting net variables
output_graph_def = _transfer_fitting_net_trainable_variables(
sess,
output_graph_def,
input_graph_def
)
# Finally we serialize and dump the output graph to the filesystem
with tf.gfile.GFile(output_graph, "wb") as f:
f.write(output_graph_def.SerializeToString())
log.info(f"{len(output_graph_def.node):d} ops in the final graph.")
_skbuild/linux-x86_64-3.6/setuptools/lib/deepmd/entrypoints/main.py 0 → 100644
View file @ 6b33aeb8
"""DeePMD-Kit entry point module."""
import argparse
import logging
import textwrap
from pathlib import Path
from typing import Dict, List, Optional
from deepmd import __version__
from deepmd.entrypoints import (
compress,
config,
doc_train_input,
freeze,
test,
train_dp,
transfer,
make_model_devi,
convert,
neighbor_stat,
)
from deepmd.loggers import set_log_handles
from deepmd.nvnmd.entrypoints.train import train_nvnmd
__all__ = ["main", "parse_args", "get_ll", "main_parser"]
def get_ll(log_level: str) -> int:
"""Convert string to python logging level.
Parameters
----------
log_level : str
allowed input values are: DEBUG, INFO, WARNING, ERROR, 3, 2, 1, 0
Returns
-------
int
one of python logging module log levels - 10, 20, 30 or 40
"""
if log_level.isdigit():
int_level = (4 - int(log_level)) * 10
else:
int_level = getattr(logging, log_level)
return int_level
class RawTextArgumentDefaultsHelpFormatter(
argparse.RawTextHelpFormatter, argparse.ArgumentDefaultsHelpFormatter
):
"""This formatter is used to print multile-line help message with default value."""
def main_parser() -> argparse.ArgumentParser:
"""DeePMD-Kit commandline options argument parser.
Returns
-------
argparse.ArgumentParser
main parser of DeePMD-kit
"""
parser = argparse.ArgumentParser(
description="DeePMD-kit: A deep learning package for many-body potential energy"
" representation and molecular dynamics",
formatter_class=argparse.ArgumentDefaultsHelpFormatter,
)
subparsers = parser.add_subparsers(title="Valid subcommands", dest="command")
# * logging options parser *********************************************************
# with use of the parent argument this options will be added to every parser
parser_log = argparse.ArgumentParser(
add_help=False, formatter_class=argparse.ArgumentDefaultsHelpFormatter
)
parser_log.add_argument(
"-v",
"--log-level",
choices=["DEBUG", "3", "INFO", "2", "WARNING", "1", "ERROR", "0"],
default="INFO",
help="set verbosity level by string or number, 0=ERROR, 1=WARNING, 2=INFO "
"and 3=DEBUG",
)
parser_log.add_argument(
"-l",
"--log-path",
type=str,
default=None,
help="set log file to log messages to disk, if not specified, the logs will "
"only be output to console",
)
# * mpi logging parser *************************************************************
parser_mpi_log = argparse.ArgumentParser(
add_help=False, formatter_class=argparse.ArgumentDefaultsHelpFormatter
)
parser_mpi_log.add_argument(
"-m",
"--mpi-log",
type=str,
default="master",
choices=("master", "collect", "workers"),
help="Set the manner of logging when running with MPI. 'master' logs only on "
"main process, 'collect' broadcasts logs from workers to master and 'workers' "
"means each process will output its own log",
)
# * config script ******************************************************************
parser_cfig = subparsers.add_parser(
"config",
parents=[parser_log],
help="fast configuration of parameter file for smooth model",
formatter_class=argparse.ArgumentDefaultsHelpFormatter,
)
parser_cfig.add_argument(
"-o", "--output", type=str, default="input.json", help="the output json file"
)
# * transfer script ****************************************************************
parser_transfer = subparsers.add_parser(
"transfer", parents=[parser_log], help="pass parameters to another model"
)
parser_transfer.add_argument(
"-r",
"--raw-model",
default="raw_frozen_model.pb",
type=str,
help="the model receiving parameters",
)
parser_transfer.add_argument(
"-O",
"--old-model",
default="old_frozen_model.pb",
type=str,
help="the model providing parameters",
)
parser_transfer.add_argument(
"-o",
"--output",
default="frozen_model.pb",
type=str,
help="the model after passing parameters",
)
# * config parser ******************************************************************
parser_train = subparsers.add_parser(
"train",
parents=[parser_log, parser_mpi_log],
help="train a model",
formatter_class=RawTextArgumentDefaultsHelpFormatter,
epilog=textwrap.dedent("""\
examples:
dp train input.json
dp train input.json --restart model.ckpt
dp train input.json --init-model model.ckpt
"""),
)
parser_train.add_argument(
"INPUT", help="the input parameter file in json or yaml format"
)
parser_train.add_argument(
"-i",
"--init-model",
type=str,
default=None,
help="Initialize the model by the provided checkpoint.",
)
parser_train.add_argument(
"-r",
"--restart",
type=str,
default=None,
help="Restart the training from the provided checkpoint.",
)
parser_train.add_argument(
"-o",
"--output",
type=str,
default="out.json",
help="The output file of the parameters used in training.",
)
parser_train.add_argument(
"-f",
"--init-frz-model",
type=str,
default=None,
help="Initialize the training from the frozen model.",
)
parser_train.add_argument(
"--skip-neighbor-stat",
action="store_true",
help="Skip calculating neighbor statistics. Sel checking, automatic sel, and model compression will be disabled.",
)
# * freeze script ******************************************************************
parser_frz = subparsers.add_parser(
"freeze",
parents=[parser_log],
help="freeze the model",
formatter_class=RawTextArgumentDefaultsHelpFormatter,
epilog=textwrap.dedent("""\
examples:
dp freeze
dp freeze -o graph.pb
"""),
)
parser_frz.add_argument(
"-c",
"--checkpoint-folder",
type=str,
default=".",
help="path to checkpoint folder",
)
parser_frz.add_argument(
"-o",
"--output",
type=str,
default="frozen_model.pb",
help="name of graph, will output to the checkpoint folder",
)
parser_frz.add_argument(
"-n",
"--node-names",
type=str,
default=None,
help="the frozen nodes, if not set, determined from the model type",
)
parser_frz.add_argument(
"-w",
"--nvnmd-weight",
type=str,
default=None,
help="the name of weight file (.npy), if set, save the model's weight into the file",
)
# * test script ********************************************************************
parser_tst = subparsers.add_parser(
"test",
parents=[parser_log],
help="test the model",
formatter_class=RawTextArgumentDefaultsHelpFormatter,
epilog=textwrap.dedent("""\
examples:
dp test -m graph.pb -s /path/to/system -n 30
"""),
)
parser_tst.add_argument(
"-m",
"--model",
default="frozen_model.pb",
type=str,
help="Frozen model file to import",
)
parser_tst.add_argument(
"-s",
"--system",
default=".",
type=str,
help="The system dir. Recursively detect systems in this directory",
)
parser_tst.add_argument(
"-S", "--set-prefix", default="set", type=str, help="The set prefix"
)
parser_tst.add_argument(
"-n", "--numb-test", default=100, type=int, help="The number of data for test"
)
parser_tst.add_argument(
"-r", "--rand-seed", type=int, default=None, help="The random seed"
)
parser_tst.add_argument(
"--shuffle-test", action="store_true", default=False, help="Shuffle test data"
)
parser_tst.add_argument(
"-d",
"--detail-file",
type=str,
default=None,
help="File where details of energy force and virial accuracy will be written",
)
parser_tst.add_argument(
"-a",
"--atomic",
action="store_true",
default=False,
help="Test the accuracy of atomic label, i.e. energy / tensor (dipole, polar)",
)
# * compress model *****************************************************************
# Compress a model, which including tabulating the embedding-net.
# The table is composed of fifth-order polynomial coefficients and is assembled
# from two sub-tables. The first table takes the step(parameter) as it's uniform
# step, while the second table takes 10 * step as it\s uniform step
#  The range of the first table is automatically detected by deepmd-kit, while the
# second table ranges from the first table's upper boundary(upper) to the
# extrapolate(parameter) * upper.
parser_compress = subparsers.add_parser(
"compress",
parents=[parser_log, parser_mpi_log],
help="compress a model",
formatter_class=RawTextArgumentDefaultsHelpFormatter,
epilog=textwrap.dedent("""\
examples:
dp compress
dp compress -i graph.pb -o compressed.pb
"""),
)
parser_compress.add_argument(
"-i",
"--input",
default="frozen_model.pb",
type=str,
help="The original frozen model, which will be compressed by the code",
)
parser_compress.add_argument(
"-o",
"--output",
default="frozen_model_compressed.pb",
type=str,
help="The compressed model",
)
parser_compress.add_argument(
"-s",
"--step",
default=0.01,
type=float,
help="Model compression uses fifth-order polynomials to interpolate the embedding-net. "
"It introduces two tables with different step size to store the parameters of the polynomials. "
"The first table covers the range of the training data, while the second table is an extrapolation of the training data. "
"The domain of each table is uniformly divided by a given step size. "
"And the step(parameter) denotes the step size of the first table and the second table will "
"use 10 * step as it's step size to save the memory. "
"Usually the value ranges from 0.1 to 0.001. "
"Smaller step means higher accuracy and bigger model size",
)
parser_compress.add_argument(
"-e",
"--extrapolate",
default=5,
type=int,
help="The domain range of the first table is automatically detected by the code: [d_low, d_up]. "
"While the second table ranges from the first table's upper boundary(d_up) to the extrapolate(parameter) * d_up: [d_up, extrapolate * d_up]",
)
parser_compress.add_argument(
"-f",
"--frequency",
default=-1,
type=int,
help="The frequency of tabulation overflow check(Whether the input environment "
"matrix overflow the first or second table range). "
"By default do not check the overflow",
)
parser_compress.add_argument(
"-c",
"--checkpoint-folder",
type=str,
default="model-compression",
help="path to checkpoint folder",
)
parser_compress.add_argument(
"-t",
"--training-script",
type=str,
default=None,
help="The training script of the input frozen model",
)
# * print docs script **************************************************************
parsers_doc = subparsers.add_parser(
"doc-train-input",
parents=[parser_log],
help="print the documentation (in rst format) of input training parameters.",
formatter_class=argparse.ArgumentDefaultsHelpFormatter,
)
parsers_doc.add_argument(
"--out-type",
default="rst",
type=str,
help="The output type"
)
# * make model deviation ***********************************************************
parser_model_devi = subparsers.add_parser(
"model-devi",
parents=[parser_log],
help="calculate model deviation",
formatter_class=RawTextArgumentDefaultsHelpFormatter,
epilog=textwrap.dedent("""\
examples:
dp model-devi -m graph.000.pb graph.001.pb graph.002.pb graph.003.pb -s ./data -o model_devi.out
"""),
)
parser_model_devi.add_argument(
"-m",
"--models",
default=["graph.000.pb", "graph.001.pb", "graph.002.pb", "graph.003.pb"],
nargs="+",
type=str,
help="Frozen models file to import",
)
parser_model_devi.add_argument(
"-s",
"--system",
default=".",
type=str,
help="The system directory. Recursively detect systems in this directory.",
)
parser_model_devi.add_argument(
"-S", "--set-prefix", default="set", type=str, help="The set prefix"
)
parser_model_devi.add_argument(
"-o",
"--output",
default="model_devi.out",
type=str,
help="The output file for results of model deviation"
)
parser_model_devi.add_argument(
"-f",
"--frequency",
default=1,
type=int,
help="The trajectory frequency of the system"
)
# * convert models
parser_transform = subparsers.add_parser(
'convert-from',
parents=[parser_log],
help='convert lower model version to supported version',
formatter_class=RawTextArgumentDefaultsHelpFormatter,
epilog=textwrap.dedent("""\
examples:
dp convert-from 1.0 -i graph.pb -o graph_new.pb
"""),
)
parser_transform.add_argument(
'FROM',
type = str,
choices = ['0.12', '1.0', '1.1', '1.2', '1.3', '2.0'],
help="The original model compatibility",
)
parser_transform.add_argument(
'-i',
"--input-model",
default = "frozen_model.pb",
type=str,
help = "the input model",
)
parser_transform.add_argument(
"-o",
"--output-model",
default = "convert_out.pb",
type=str,
help='the output model',
)
# neighbor_stat
parser_neighbor_stat = subparsers.add_parser(
'neighbor-stat',
parents=[parser_log],
help='Calculate neighbor statistics',
formatter_class=RawTextArgumentDefaultsHelpFormatter,
epilog=textwrap.dedent("""\
examples:
dp neighbor-stat -s data -r 6.0 -t O H
"""),
)
parser_neighbor_stat.add_argument(
"-s",
"--system",
default=".",
type=str,
help="The system dir. Recursively detect systems in this directory",
)
parser_neighbor_stat.add_argument(
"-r",
"--rcut",
type=float,
required=True,
help="cutoff radius",
)
parser_neighbor_stat.add_argument(
"-t",
"--type-map",
type=str,
nargs='+',
required=True,
help="type map",
)
parser_neighbor_stat.add_argument(
"--one-type",
action="store_true",
default=False,
help="treat all types as a single type. Used with se_atten descriptor.",
)
# --version
parser.add_argument('--version', action='version', version='DeePMD-kit v%s' % __version__)
# * train nvnmd script ******************************************************************
parser_train_nvnmd = subparsers.add_parser(
"train-nvnmd",
parents=[parser_log],
help="train nvnmd model",
formatter_class=argparse.ArgumentDefaultsHelpFormatter,
)
parser_train_nvnmd.add_argument(
"INPUT", help="the input parameter file in json format"
)
parser_train_nvnmd.add_argument(
"-s",
"--step",
default="s1",
type=str,
choices=['s1', 's2'],
help="steps to train model of NVNMD: s1 (train CNN), s2 (train QNN)"
)
return parser
def parse_args(args: Optional[List[str]] = None) -> argparse.Namespace:
"""Parse arguments and convert argument strings to objects.
Parameters
----------
args: List[str]
list of command line arguments, main purpose is testing default option None
takes arguments from sys.argv
Returns
-------
argparse.Namespace
the populated namespace
"""
parser = main_parser()
parsed_args = parser.parse_args(args=args)
if parsed_args.command is None:
parser.print_help()
else:
parsed_args.log_level = get_ll(parsed_args.log_level)
return parsed_args
def main():
"""DeePMD-Kit entry point.
Raises
------
RuntimeError
if no command was input
"""
args = parse_args()
# do not set log handles for None, it is useless
# log handles for train will be set separatelly
# when the use of MPI will be determined in `RunOptions`
if args.command not in (None, "train"):
set_log_handles(args.log_level, Path(args.log_path) if args.log_path else None)
dict_args = vars(args)
if args.command == "train":
train_dp(**dict_args)
elif args.command == "freeze":
freeze(**dict_args)
elif args.command == "config":
config(**dict_args)
elif args.command == "test":
test(**dict_args)
elif args.command == "transfer":
transfer(**dict_args)
elif args.command == "compress":
compress(**dict_args)
elif args.command == "doc-train-input":
doc_train_input(**dict_args)
elif args.command == "model-devi":
make_model_devi(**dict_args)
elif args.command == "convert-from":
convert(**dict_args)
elif args.command == "neighbor-stat":
neighbor_stat(**dict_args)
elif args.command == "train-nvnmd": # nvnmd
train_nvnmd(**dict_args)
elif args.command is None:
pass
else:
raise RuntimeError(f"unknown command {args.command}")
_skbuild/linux-x86_64-3.6/setuptools/lib/deepmd/entrypoints/neighbor_stat.py 0 → 100644
View file @ 6b33aeb8
import logging
from typing import List
from deepmd.common import expand_sys_str
from deepmd.utils.data_system import DeepmdDataSystem
from deepmd.utils.neighbor_stat import NeighborStat
log = logging.getLogger(__name__)
def neighbor_stat(
*,
system: str,
rcut: float,
type_map: List[str],
one_type: bool = False,
**kwargs,
):
"""Calculate neighbor statistics.
Parameters
----------
system : str
system to stat
rcut : float
cutoff radius
type_map : list[str]
type map
one_type : bool, optional, default=False
treat all types as a single type
Examples
--------
>>> neighbor_stat(system='.', rcut=6., type_map=["C", "H", "O", "N", "P", "S", "Mg", "Na", "HW", "OW", "mNa", "mCl", "mC", "mH", "mMg", "mN", "mO", "mP"])
min_nbor_dist: 0.6599510670195264
max_nbor_size: [23, 26, 19, 16, 2, 2, 1, 1, 72, 37, 5, 0, 31, 29, 1, 21, 20, 5]
"""
all_sys = expand_sys_str(system)
if not len(all_sys):
raise RuntimeError("Did not find valid system")
data = DeepmdDataSystem(
systems=all_sys,
batch_size=1,
test_size=1,
rcut=rcut,
type_map=type_map,
)
data.get_batch()
nei = NeighborStat(data.get_ntypes(), rcut, one_type=one_type)
min_nbor_dist, max_nbor_size = nei.get_stat(data)
log.info("min_nbor_dist: %f" % min_nbor_dist)
log.info("max_nbor_size: %s" % str(max_nbor_size))
return min_nbor_dist, max_nbor_size
_skbuild/linux-x86_64-3.6/setuptools/lib/deepmd/entrypoints/test.py 0 → 100644
View file @ 6b33aeb8
"""Test trained DeePMD model."""
import logging
from pathlib import Path
from typing import TYPE_CHECKING, List, Dict, Optional, Tuple
import numpy as np
from deepmd import DeepPotential
from deepmd.common import expand_sys_str
from deepmd.utils import random as dp_random
from deepmd.utils.data import DeepmdData
from deepmd.utils.weight_avg import weighted_average
if TYPE_CHECKING:
from deepmd.infer import DeepDipole, DeepPolar, DeepPot, DeepWFC
from deepmd.infer.deep_tensor import DeepTensor
__all__ = ["test"]
log = logging.getLogger(__name__)
def test(
*,
model: str,
system: str,
set_prefix: str,
numb_test: int,
rand_seed: Optional[int],
shuffle_test: bool,
detail_file: str,
atomic: bool,
**kwargs,
):
"""Test model predictions.
Parameters
----------
model : str
path where model is stored
system : str
system directory
set_prefix : str
string prefix of set
numb_test : int
munber of tests to do
rand_seed : Optional[int]
seed for random generator
shuffle_test : bool
whether to shuffle tests
detail_file : Optional[str]
file where test details will be output
atomic : bool
whether per atom quantities should be computed
Raises
------
RuntimeError
if no valid system was found
"""
all_sys = expand_sys_str(system)
if len(all_sys) == 0:
raise RuntimeError("Did not find valid system")
err_coll = []
siz_coll = []
# init random seed
if rand_seed is not None:
dp_random.seed(rand_seed % (2 ** 32))
# init model
dp = DeepPotential(model)
for cc, system in enumerate(all_sys):
log.info("# ---------------output of dp test--------------- ")
log.info(f"# testing system : {system}")
# create data class
tmap = dp.get_type_map() if dp.model_type == "ener" else None
data = DeepmdData(system, set_prefix, shuffle_test=shuffle_test, type_map=tmap)
if dp.model_type == "ener":
err = test_ener(
dp,
data,
system,
numb_test,
detail_file,
atomic,
append_detail=(cc != 0),
)
elif dp.model_type == "dipole":
err = test_dipole(dp, data, numb_test, detail_file, atomic)
elif dp.model_type == "polar":
err = test_polar(dp, data, numb_test, detail_file, atomic=atomic)
elif dp.model_type == "global_polar": # should not appear in this new version
log.warning("Global polar model is not currently supported. Please directly use the polar mode and change loss parameters.")
err = test_polar(dp, data, numb_test, detail_file, atomic=False) # YWolfeee: downward compatibility
log.info("# ----------------------------------------------- ")
err_coll.append(err)
avg_err = weighted_average(err_coll)
if len(all_sys) != len(err_coll):
log.warning("Not all systems are tested! Check if the systems are valid")
if len(all_sys) > 1:
log.info("# ----------weighted average of errors----------- ")
log.info(f"# number of systems : {len(all_sys)}")
if dp.model_type == "ener":
print_ener_sys_avg(avg_err)
elif dp.model_type == "dipole":
print_dipole_sys_avg(avg_err)
elif dp.model_type == "polar":
print_polar_sys_avg(avg_err)
elif dp.model_type == "global_polar":
print_polar_sys_avg(avg_err)
elif dp.model_type == "wfc":
print_wfc_sys_avg(avg_err)
log.info("# ----------------------------------------------- ")
def rmse(diff: np.ndarray) -> np.ndarray:
"""Calculate average root mean square error.
Parameters
----------
diff: np.ndarray
difference
Returns
-------
np.ndarray
array with normalized difference
"""
return np.sqrt(np.average(diff * diff))
def save_txt_file(
fname: Path, data: np.ndarray, header: str = "", append: bool = False
):
"""Save numpy array to test file.
Parameters
----------
fname : str
filename
data : np.ndarray
data to save to disk
header : str, optional
header string to use in file, by default ""
append : bool, optional
if true file will be appended insted of overwriting, by default False
"""
flags = "ab" if append else "w"
with fname.open(flags) as fp:
np.savetxt(fp, data, header=header)
def test_ener(
dp: "DeepPot",
data: DeepmdData,
system: str,
numb_test: int,
detail_file: Optional[str],
has_atom_ener: bool,
append_detail: bool = False,
) -> Tuple[List[np.ndarray], List[int]]:
"""Test energy type model.
Parameters
----------
dp : DeepPot
instance of deep potential
data: DeepmdData
data container object
system : str
system directory
numb_test : int
munber of tests to do
detail_file : Optional[str]
file where test details will be output
has_atom_ener : bool
whether per atom quantities should be computed
append_detail : bool, optional
if true append output detail file, by default False
Returns
-------
Tuple[List[np.ndarray], List[int]]
arrays with results and their shapes
"""
data.add("energy", 1, atomic=False, must=False, high_prec=True)
data.add("force", 3, atomic=True, must=False, high_prec=False)
data.add("virial", 9, atomic=False, must=False, high_prec=False)
if dp.has_efield:
data.add("efield", 3, atomic=True, must=True, high_prec=False)
if has_atom_ener:
data.add("atom_ener", 1, atomic=True, must=True, high_prec=False)
if dp.get_dim_fparam() > 0:
data.add(
"fparam", dp.get_dim_fparam(), atomic=False, must=True, high_prec=False
)
if dp.get_dim_aparam() > 0:
data.add("aparam", dp.get_dim_aparam(), atomic=True, must=True, high_prec=False)
test_data = data.get_test()
natoms = len(test_data["type"][0])
nframes = test_data["box"].shape[0]
numb_test = min(nframes, numb_test)
coord = test_data["coord"][:numb_test].reshape([numb_test, -1])
box = test_data["box"][:numb_test]
if dp.has_efield:
efield = test_data["efield"][:numb_test].reshape([numb_test, -1])
else:
efield = None
if not data.pbc:
box = None
atype = test_data["type"][0]
if dp.get_dim_fparam() > 0:
fparam = test_data["fparam"][:numb_test]
else:
fparam = None
if dp.get_dim_aparam() > 0:
aparam = test_data["aparam"][:numb_test]
else:
aparam = None
ret = dp.eval(
coord,
box,
atype,
fparam=fparam,
aparam=aparam,
atomic=has_atom_ener,
efield=efield,
)
energy = ret[0]
force = ret[1]
virial = ret[2]
energy = energy.reshape([numb_test, 1])
force = force.reshape([numb_test, -1])
virial = virial.reshape([numb_test, 9])
if has_atom_ener:
ae = ret[3]
av = ret[4]
ae = ae.reshape([numb_test, -1])
av = av.reshape([numb_test, -1])
rmse_e = rmse(energy - test_data["energy"][:numb_test].reshape([-1, 1]))
rmse_f = rmse(force - test_data["force"][:numb_test])
rmse_v = rmse(virial - test_data["virial"][:numb_test])
rmse_ea = rmse_e / natoms
rmse_va = rmse_v / natoms
if has_atom_ener:
rmse_ae = rmse(
test_data["atom_ener"][:numb_test].reshape([-1]) - ae.reshape([-1])
)
# print ("# energies: %s" % energy)
log.info(f"# number of test data : {numb_test:d} ")
log.info(f"Energy RMSE : {rmse_e:e} eV")
log.info(f"Energy RMSE/Natoms : {rmse_ea:e} eV")
log.info(f"Force RMSE : {rmse_f:e} eV/A")
if data.pbc:
log.info(f"Virial RMSE : {rmse_v:e} eV")
log.info(f"Virial RMSE/Natoms : {rmse_va:e} eV")
if has_atom_ener:
log.info(f"Atomic ener RMSE : {rmse_ae:e} eV")
if detail_file is not None:
detail_path = Path(detail_file)
pe = np.concatenate(
(
np.reshape(test_data["energy"][:numb_test], [-1, 1]),
np.reshape(energy, [-1, 1]),
),
axis=1,
)
save_txt_file(
detail_path.with_suffix(".e.out"),
pe,
header="%s: data_e pred_e" % system,
append=append_detail,
)
pf = np.concatenate(
(
np.reshape(test_data["force"][:numb_test], [-1, 3]),
np.reshape(force, [-1, 3]),
),
axis=1,
)
save_txt_file(
detail_path.with_suffix(".f.out"),
pf,
header="%s: data_fx data_fy data_fz pred_fx pred_fy pred_fz" % system,
append=append_detail,
)
pv = np.concatenate(
(
np.reshape(test_data["virial"][:numb_test], [-1, 9]),
np.reshape(virial, [-1, 9]),
),
axis=1,
)
save_txt_file(
detail_path.with_suffix(".v.out"),
pv,
header=f"{system}: data_vxx data_vxy data_vxz data_vyx data_vyy "
"data_vyz data_vzx data_vzy data_vzz pred_vxx pred_vxy pred_vxz pred_vyx "
"pred_vyy pred_vyz pred_vzx pred_vzy pred_vzz",
append=append_detail,
)
return {
"rmse_ea" : (rmse_ea, energy.size),
"rmse_f" : (rmse_f, force.size),
"rmse_va" : (rmse_va, virial.size),
}
def print_ener_sys_avg(avg: Dict[str,float]):
"""Print errors summary for energy type potential.
Parameters
----------
avg : np.ndarray
array with summaries
"""
log.info(f"Energy RMSE/Natoms : {avg['rmse_ea']:e} eV")
log.info(f"Force RMSE : {avg['rmse_f']:e} eV/A")
log.info(f"Virial RMSE/Natoms : {avg['rmse_va']:e} eV")
def run_test(dp: "DeepTensor", test_data: dict, numb_test: int):
"""Run tests.
Parameters
----------
dp : DeepTensor
instance of deep potential
test_data : dict
dictionary with test data
numb_test : int
munber of tests to do
Returns
-------
[type]
[description]
"""
nframes = test_data["box"].shape[0]
numb_test = min(nframes, numb_test)
coord = test_data["coord"][:numb_test].reshape([numb_test, -1])
box = test_data["box"][:numb_test]
atype = test_data["type"][0]
prediction = dp.eval(coord, box, atype)
return prediction.reshape([numb_test, -1]), numb_test, atype
def test_wfc(
dp: "DeepWFC",
data: DeepmdData,
numb_test: int,
detail_file: Optional[str],
) -> Tuple[List[np.ndarray], List[int]]:
"""Test energy type model.
Parameters
----------
dp : DeepPot
instance of deep potential
data: DeepmdData
data container object
numb_test : int
munber of tests to do
detail_file : Optional[str]
file where test details will be output
Returns
-------
Tuple[List[np.ndarray], List[int]]
arrays with results and their shapes
"""
data.add(
"wfc", 12, atomic=True, must=True, high_prec=False, type_sel=dp.get_sel_type()
)
test_data = data.get_test()
wfc, numb_test, _ = run_test(dp, test_data, numb_test)
rmse_f = rmse(wfc - test_data["wfc"][:numb_test])
log.info("# number of test data : {numb_test:d} ")
log.info("WFC RMSE : {rmse_f:e} eV/A")
if detail_file is not None:
detail_path = Path(detail_file)
pe = np.concatenate(
(
np.reshape(test_data["wfc"][:numb_test], [-1, 12]),
np.reshape(wfc, [-1, 12]),
),
axis=1,
)
np.savetxt(
detail_path.with_suffix(".out"),
pe,
header="ref_wfc(12 dofs) predicted_wfc(12 dofs)",
)
return {
'rmse' : (rmse_f, wfc.size)
}
def print_wfc_sys_avg(avg):
"""Print errors summary for wfc type potential.
Parameters
----------
avg : np.ndarray
array with summaries
"""
log.info(f"WFC RMSE : {avg['rmse']:e} eV/A")
def test_polar(
dp: "DeepPolar",
data: DeepmdData,
numb_test: int,
detail_file: Optional[str],
*,
atomic: bool,
) -> Tuple[List[np.ndarray], List[int]]:
"""Test energy type model.
Parameters
----------
dp : DeepPot
instance of deep potential
data: DeepmdData
data container object
numb_test : int
munber of tests to do
detail_file : Optional[str]
file where test details will be output
global_polar : bool
wheter to use glovbal version of polar potential
Returns
-------
Tuple[List[np.ndarray], List[int]]
arrays with results and their shapes
"""
data.add(
"polarizability" if not atomic else "atomic_polarizability",
9,
atomic=atomic,
must=True,
high_prec=False,
type_sel=dp.get_sel_type(),
)
test_data = data.get_test()
polar, numb_test, atype = run_test(dp, test_data, numb_test)
sel_type = dp.get_sel_type()
sel_natoms = 0
for ii in sel_type:
sel_natoms += sum(atype == ii)
# YWolfeee: do summation in global polar mode
if not atomic:
polar = np.sum(polar.reshape((polar.shape[0],-1,9)),axis=1)
rmse_f = rmse(polar - test_data["polarizability"][:numb_test])
rmse_fs = rmse_f / np.sqrt(sel_natoms)
rmse_fa = rmse_f / sel_natoms
else:
rmse_f = rmse(polar - test_data["atomic_polarizability"][:numb_test])
log.info(f"# number of test data : {numb_test:d} ")
log.info(f"Polarizability RMSE : {rmse_f:e}")
if not atomic:
log.info(f"Polarizability RMSE/sqrtN : {rmse_fs:e}")
log.info(f"Polarizability RMSE/N : {rmse_fa:e}")
log.info(f"The unit of error is the same as the unit of provided label.")
if detail_file is not None:
detail_path = Path(detail_file)
pe = np.concatenate(
(
np.reshape(test_data["polarizability"][:numb_test], [-1, 9]),
np.reshape(polar, [-1, 9]),
),
axis=1,
)
np.savetxt(
detail_path.with_suffix(".out"),
pe,
header="data_pxx data_pxy data_pxz data_pyx data_pyy data_pyz data_pzx "
"data_pzy data_pzz pred_pxx pred_pxy pred_pxz pred_pyx pred_pyy pred_pyz "
"pred_pzx pred_pzy pred_pzz",
)
return {
"rmse" : (rmse_f, polar.size)
}
def print_polar_sys_avg(avg):
"""Print errors summary for polar type potential.
Parameters
----------
avg : np.ndarray
array with summaries
"""
log.info(f"Polarizability RMSE : {avg['rmse']:e} eV/A")
def test_dipole(
dp: "DeepDipole",
data: DeepmdData,
numb_test: int,
detail_file: Optional[str],
atomic: bool,
) -> Tuple[List[np.ndarray], List[int]]:
"""Test energy type model.
Parameters
----------
dp : DeepPot
instance of deep potential
data: DeepmdData
data container object
numb_test : int
munber of tests to do
detail_file : Optional[str]
file where test details will be output
atomic : bool
whether atomic dipole is provided
Returns
-------
Tuple[List[np.ndarray], List[int]]
arrays with results and their shapes
"""
data.add(
"dipole" if not atomic else "atomic_dipole",
3,
atomic=atomic,
must=True,
high_prec=False,
type_sel=dp.get_sel_type()
)
test_data = data.get_test()
dipole, numb_test, atype = run_test(dp, test_data, numb_test)
sel_type = dp.get_sel_type()
sel_natoms = 0
for ii in sel_type:
sel_natoms += sum(atype == ii)
# do summation in atom dimension
if not atomic:
dipole = np.sum(dipole.reshape((dipole.shape[0], -1, 3)),axis=1)
rmse_f = rmse(dipole - test_data["dipole"][:numb_test])
rmse_fs = rmse_f / np.sqrt(sel_natoms)
rmse_fa = rmse_f / sel_natoms
else:
rmse_f = rmse(dipole - test_data["atomic_dipole"][:numb_test])
log.info(f"# number of test data : {numb_test:d}")
log.info(f"Dipole RMSE : {rmse_f:e}")
if not atomic:
log.info(f"Dipole RMSE/sqrtN : {rmse_fs:e}")
log.info(f"Dipole RMSE/N : {rmse_fa:e}")
log.info(f"The unit of error is the same as the unit of provided label.")
if detail_file is not None:
detail_path = Path(detail_file)
pe = np.concatenate(
(
np.reshape(test_data["dipole"][:numb_test], [-1, 3]),
np.reshape(dipole, [-1, 3]),
),
axis=1,
)
np.savetxt(
detail_path.with_suffix(".out"),
pe,
header="data_x data_y data_z pred_x pred_y pred_z",
)
return {
'rmse' : (rmse_f, dipole.size)
}
def print_dipole_sys_avg(avg):
"""Print errors summary for dipole type potential.
Parameters
----------
avg : np.ndarray
array with summaries
"""
log.info(f"Dipole RMSE : {avg['rmse']:e} eV/A")
_skbuild/linux-x86_64-3.6/setuptools/lib/deepmd/entrypoints/train.py 0 → 100644
View file @ 6b33aeb8
"""DeePMD training entrypoint script.
Can handle local or distributed training.
"""
import json
import logging
import time
import os
from typing import Dict, List, Optional, Any
from deepmd.common import data_requirement, expand_sys_str, j_loader, j_must_have
from deepmd.env import tf, reset_default_tf_session_config, GLOBAL_ENER_FLOAT_PRECISION
from deepmd.infer.data_modifier import DipoleChargeModifier
from deepmd.train.run_options import BUILD, CITATION, WELCOME, RunOptions
from deepmd.train.trainer import DPTrainer
from deepmd.utils import random as dp_random
from deepmd.utils.argcheck import normalize
from deepmd.utils.compat import update_deepmd_input
from deepmd.utils.data_system import DeepmdDataSystem
from deepmd.utils.sess import run_sess
from deepmd.utils.neighbor_stat import NeighborStat
from deepmd.utils.path import DPPath
__all__ = ["train"]
log = logging.getLogger(__name__)
def train(
*,
INPUT: str,
init_model: Optional[str],
restart: Optional[str],
output: str,
init_frz_model: str,
mpi_log: str,
log_level: int,
log_path: Optional[str],
is_compress: bool = False,
skip_neighbor_stat: bool = False,
**kwargs,
):
"""Run DeePMD model training.
Parameters
----------
INPUT : str
json/yaml control file
init_model : Optional[str]
path to checkpoint folder or None
restart : Optional[str]
path to checkpoint folder or None
output : str
path for dump file with arguments
init_frz_model : str
path to frozen model or None
mpi_log : str
mpi logging mode
log_level : int
logging level defined by int 0-3
log_path : Optional[str]
logging file path or None if logs are to be output only to stdout
is_compress: bool
indicates whether in the model compress mode
skip_neighbor_stat : bool, default=False
skip checking neighbor statistics
Raises
------
RuntimeError
if distributed training job nem is wrong
"""
run_opt = RunOptions(
init_model=init_model,
restart=restart,
init_frz_model=init_frz_model,
log_path=log_path,
log_level=log_level,
mpi_log=mpi_log
)
if run_opt.is_distrib and len(run_opt.gpus or []) > 1:
# avoid conflict of visible gpus among multipe tf sessions in one process
reset_default_tf_session_config(cpu_only=True)
# load json database
jdata = j_loader(INPUT)
jdata = update_deepmd_input(jdata, warning=True, dump="input_v2_compat.json")
jdata = normalize(jdata)
if not is_compress and not skip_neighbor_stat:
jdata = update_sel(jdata)
with open(output, "w") as fp:
json.dump(jdata, fp, indent=4)
# save the training script into the graph
# remove white spaces as it is not compressed
tf.constant(json.dumps(jdata, separators=(',', ':')), name='train_attr/training_script', dtype=tf.string)
for message in WELCOME + CITATION + BUILD:
log.info(message)
run_opt.print_resource_summary()
_do_work(jdata, run_opt, is_compress)
def _do_work(jdata: Dict[str, Any], run_opt: RunOptions, is_compress: bool = False):
"""Run serial model training.
Parameters
----------
jdata : Dict[str, Any]
arguments read form json/yaml control file
run_opt : RunOptions
object with run configuration
is_compress : Bool
indicates whether in model compress mode
Raises
------
RuntimeError
If unsupported modifier type is selected for model
"""
# make necessary checks
assert "training" in jdata
# init the model
model = DPTrainer(jdata, run_opt=run_opt, is_compress = is_compress)
rcut = model.model.get_rcut()
type_map = model.model.get_type_map()
if len(type_map) == 0:
ipt_type_map = None
else:
ipt_type_map = type_map
# init random seed of data systems
seed = jdata["training"].get("seed", None)
if seed is not None:
# avoid the same batch sequence among workers
seed += run_opt.my_rank
seed = seed % (2 ** 32)
dp_random.seed(seed)
# setup data modifier
modifier = get_modifier(jdata["model"].get("modifier", None))
# decouple the training data from the model compress process
train_data = None
valid_data = None
if not is_compress:
# init data
train_data = get_data(jdata["training"]["training_data"], rcut, ipt_type_map, modifier)
train_data.print_summary("training")
if jdata["training"].get("validation_data", None) is not None:
valid_data = get_data(jdata["training"]["validation_data"], rcut, train_data.type_map, modifier)
valid_data.print_summary("validation")
# get training info
stop_batch = j_must_have(jdata["training"], "numb_steps")
model.build(train_data, stop_batch)
if not is_compress:
# train the model with the provided systems in a cyclic way
start_time = time.time()
model.train(train_data, valid_data)
end_time = time.time()
log.info("finished training")
log.info(f"wall time: {(end_time - start_time):.3f} s")
else:
model.save_compressed()
log.info("finished compressing")
def get_data(jdata: Dict[str, Any], rcut, type_map, modifier):
systems = j_must_have(jdata, "systems")
if isinstance(systems, str):
systems = expand_sys_str(systems)
help_msg = 'Please check your setting for data systems'
# check length of systems
if len(systems) == 0:
msg = 'cannot find valid a data system'
log.fatal(msg)
raise IOError(msg, help_msg)
# rougly check all items in systems are valid
for ii in systems:
ii = DPPath(ii)
if (not ii.is_dir()):
msg = f'dir {ii} is not a valid dir'
log.fatal(msg)
raise IOError(msg, help_msg)
if (not (ii / 'type.raw').is_file()):
msg = f'dir {ii} is not a valid data system dir'
log.fatal(msg)
raise IOError(msg, help_msg)
batch_size = j_must_have(jdata, "batch_size")
sys_probs = jdata.get("sys_probs", None)
auto_prob = jdata.get("auto_prob", "prob_sys_size")
data = DeepmdDataSystem(
systems=systems,
batch_size=batch_size,
test_size=1, # to satisfy the old api
shuffle_test=True, # to satisfy the old api
rcut=rcut,
type_map=type_map,
modifier=modifier,
trn_all_set=True, # sample from all sets
sys_probs=sys_probs,
auto_prob_style=auto_prob
)
data.add_dict(data_requirement)
return data
def get_modifier(modi_data=None):
modifier: Optional[DipoleChargeModifier]
if modi_data is not None:
if modi_data["type"] == "dipole_charge":
modifier = DipoleChargeModifier(
modi_data["model_name"],
modi_data["model_charge_map"],
modi_data["sys_charge_map"],
modi_data["ewald_h"],
modi_data["ewald_beta"],
)
else:
raise RuntimeError("unknown modifier type " + str(modi_data["type"]))
else:
modifier = None
return modifier
def get_rcut(jdata):
descrpt_data = jdata['model']['descriptor']
rcut_list = []
if descrpt_data['type'] == 'hybrid':
for ii in descrpt_data['list']:
rcut_list.append(ii['rcut'])
else:
rcut_list.append(descrpt_data['rcut'])
return max(rcut_list)
def get_type_map(jdata):
return jdata['model'].get('type_map', None)
def get_nbor_stat(jdata, rcut, one_type: bool = False):
max_rcut = get_rcut(jdata)
type_map = get_type_map(jdata)
if type_map and len(type_map) == 0:
type_map = None
train_data = get_data(jdata["training"]["training_data"], max_rcut, type_map, None)
train_data.get_batch()
data_ntypes = train_data.get_ntypes()
if type_map is not None:
map_ntypes = len(type_map)
else:
map_ntypes = data_ntypes
ntypes = max([map_ntypes, data_ntypes])
neistat = NeighborStat(ntypes, rcut, one_type=one_type)
min_nbor_dist, max_nbor_size = neistat.get_stat(train_data)
# moved from traier.py as duplicated
# TODO: this is a simple fix but we should have a clear
# architecture to call neighbor stat
tf.constant(min_nbor_dist,
name = 'train_attr/min_nbor_dist',
dtype = GLOBAL_ENER_FLOAT_PRECISION)
tf.constant(max_nbor_size,
name = 'train_attr/max_nbor_size',
dtype = tf.int32)
return min_nbor_dist, max_nbor_size
def get_sel(jdata, rcut, one_type: bool = False):
_, max_nbor_size = get_nbor_stat(jdata, rcut, one_type=one_type)
return max_nbor_size
def get_min_nbor_dist(jdata, rcut):
min_nbor_dist, _ = get_nbor_stat(jdata, rcut)
return min_nbor_dist
def parse_auto_sel(sel):
if type(sel) is not str:
return False
words = sel.split(':')
if words[0] == 'auto':
return True
else:
return False
def parse_auto_sel_ratio(sel):
if not parse_auto_sel(sel):
raise RuntimeError(f'invalid auto sel format {sel}')
else:
words = sel.split(':')
if len(words) == 1:
ratio = 1.1
elif len(words) == 2:
ratio = float(words[1])
else:
raise RuntimeError(f'invalid auto sel format {sel}')
return ratio
def wrap_up_4(xx):
return 4 * ((int(xx) + 3) // 4)
def update_one_sel(jdata, descriptor):
if descriptor['type'] == 'loc_frame':
return descriptor
rcut = descriptor['rcut']
tmp_sel = get_sel(jdata, rcut, one_type=descriptor['type'] in ('se_atten',))
sel = descriptor['sel']
if isinstance(sel, int):
# convert to list and finnally convert back to int
sel = [sel]
if parse_auto_sel(descriptor['sel']) :
ratio = parse_auto_sel_ratio(descriptor['sel'])
descriptor['sel'] = sel = [int(wrap_up_4(ii * ratio)) for ii in tmp_sel]
else:
# sel is set by user
for ii, (tt, dd) in enumerate(zip(tmp_sel, sel)):
if dd and tt > dd:
# we may skip warning for sel=0, where the user is likely
# to exclude such type in the descriptor
log.warning(
"sel of type %d is not enough! The expected value is "
"not less than %d, but you set it to %d. The accuracy"
" of your model may get worse." %(ii, tt, dd)
)
if descriptor['type'] in ('se_atten',):
descriptor['sel'] = sel = sum(sel)
return descriptor
def update_sel(jdata):
log.info("Calculate neighbor statistics... (add --skip-neighbor-stat to skip this step)")
descrpt_data = jdata['model']['descriptor']
if descrpt_data['type'] == 'hybrid':
for ii in range(len(descrpt_data['list'])):
descrpt_data['list'][ii] = update_one_sel(jdata, descrpt_data['list'][ii])
else:
descrpt_data = update_one_sel(jdata, descrpt_data)
jdata['model']['descriptor'] = descrpt_data
return jdata
_skbuild/linux-x86_64-3.6/setuptools/lib/deepmd/entrypoints/transfer.py 0 → 100644
View file @ 6b33aeb8
"""Module used for transfering parameters between models."""
from typing import Dict, Optional, Sequence, Tuple
from deepmd.env import tf, TRANSFER_PATTERN
import re
import numpy as np
import logging
__all__ = ["transfer"]
log = logging.getLogger(__name__)
@np.vectorize
def convert_number(number: int) -> float:
binary = bin(number).replace("0b", "").zfill(16)
sign = int(binary[0]) * -2 + 1
exp = int(binary[1:6], 2)
frac = (int(binary[6:], 2) + 2 ** 10) * (2 ** -10)
return sign * (2 ** (exp - 15)) * frac
def convert_matrix(
matrix: np.ndarray, shape: Sequence[int], dtype: Optional[type] = None
) -> np.ndarray:
"""Convert matrix of integers to self defined binary format.
Parameters
----------
matrix : np.ndarray
array of ints
shape : Sequence[int]
shape to cast resulting array to
dtype : Optional[type]
type that finall array will be cast to, If None no casting will take place
Returns
-------
np.ndarray
array cast to required format
"""
conv = convert_number(matrix.flatten()).reshape(shape)
if dtype:
conv = conv.astype(dtype)
return conv
def transfer(*, old_model: str, raw_model: str, output: str, **kwargs):
"""Transfer operation from old fron graph to new prepared raw graph.
Parameters
----------
old_model : str
frozen old graph model
raw_model : str
new model that will accept ops from old model
output : str
new model with transfered parameters will be saved to this location
"""
raw_graph = load_graph(raw_model)
old_graph = load_graph(old_model)
log.info(f"{len(raw_graph.as_graph_def().node)} ops in the raw graph")
log.info(f"{len(old_graph.as_graph_def().node)} ops in the old graph")
new_graph_def = transform_graph(raw_graph, old_graph)
with tf.gfile.GFile(output, mode="wb") as f:
f.write(new_graph_def.SerializeToString())
log.info("the output model is saved in " + output)
def load_graph(graph_name: str) -> tf.Graph:
"""Load graph from passed in path.
Parameters
----------
graph_name : str
path to frozen graph on disk
Returns
-------
tf.Graph
tf graph object
"""
graph_def = tf.GraphDef()
with open(graph_name, "rb") as f:
graph_def.ParseFromString(f.read())
with tf.Graph().as_default() as graph:
tf.import_graph_def(graph_def, name="")
return graph
def transform_graph(raw_graph: tf.Graph, old_graph: tf.Graph) -> tf.Graph:
"""Trasform old graph into new.
Parameters
----------
raw_graph : tf.Graph
graph receiving parameters from the old one
old_graph : tf.Graph
graph providing parameters
Returns
-------
tf.Graph
new graph with parameters transfered form the old one
"""
old_graph_def = old_graph.as_graph_def()
raw_graph_def = raw_graph.as_graph_def()
raw_graph_node = load_transform_node(raw_graph_def)
old_graph_node = load_transform_node(old_graph_def)
for node in raw_graph_def.node:
if node.name not in raw_graph_node.keys():
continue
old_node = old_graph_node[node.name]
raw_node = raw_graph_node[node.name]
cp_attr = CopyNodeAttr(node)
check_dim(raw_graph_node, old_graph_node, node.name)
tensor_shape = [dim.size for dim in raw_node.tensor_shape.dim]
old_graph_dtype = tf.as_dtype(old_node.dtype).as_numpy_dtype
raw_graph_dtype = tf.as_dtype(raw_node.dtype).as_numpy_dtype
log.info(
f"{node.name} is passed from old graph({old_graph_dtype}) "
f"to raw graph({raw_graph_dtype})"
)
if raw_graph_dtype == np.float16:
if old_graph_dtype == np.float64 or old_graph_dtype == np.float32:
if (len(tensor_shape) != 1) or (tensor_shape[0] != 1):
tensor = np.frombuffer(old_node.tensor_content, dtype = old_graph_dtype)
tensor = tensor.astype(raw_graph_dtype)
cp_attr.from_str(tensor)
else:
tensor = load_tensor(old_node, old_graph_dtype, raw_graph_dtype)
cp_attr.from_array(tensor, tf.float16, [1])
elif old_graph_dtype[1] == "float16":
tensor = convertMatrix(np.array(old_node.half_val), tensor_shape)
cp_attr.from_array(tensor, raw_graph_dtype)
elif raw_graph_dtype == np.float64 or raw_graph_dtype == np.float32:
if old_graph_dtype == np.float64 or old_graph_dtype == np.float32:
if (len(tensor_shape) != 1) or (tensor_shape[0] != 1):
tensor = np.frombuffer(old_node.tensor_content, dtype = old_graph_dtype)
tensor = tensor.astype(raw_graph_dtype)
cp_attr.from_str(tensor)
else:
tensor = load_tensor(old_node, old_graph_dtype, raw_graph_dtype)
cp_attr.from_array(tensor, raw_graph_dtype, shape=[1])
elif old_graph_dtype == np.float16:
if (len(tensor_shape) != 1) or (tensor_shape[0] != 1):
tensor = convertMatrix(np.array(old_node.half_val), tensor_shape).astype(raw_graph_dtype)
cp_attr.from_str(tensor)
else:
tensor = convertMatrix(np.array(old_node.half_val), tensor_shape).astype(raw_graph_dtype)
cp_attr.from_array(tensor, raw_graph_dtype)
return raw_graph_def
class CopyNodeAttr:
def __init__(self, node) -> None:
self.node = node
def from_array(
self, tensor: np.ndarray, dtype: type, shape: Optional[Sequence[int]] = None
):
if shape is None:
shape = tensor.shape
self.node.attr["value"].CopyFrom(
tf.AttrValue(tensor=tf.make_tensor_proto(tensor, dtype, shape))
)
def from_str(self, tensor: np.ndarray):
self.node.attr["value"].tensor.tensor_content = tensor.tostring()
def load_tensor(node: tf.Tensor, dtype_old: type, dtype_new: type) -> np.ndarray:
if dtype_old == np.float64:
tensor = np.array(node.double_val).astype(dtype_new)
elif dtype_old == np.float32:
tensor = np.array(node.float_val).astype(dtype_new)
return tensor
def check_dim(raw_graph_node: tf.Tensor, old_graph_node: tf.Tensor, node_name: str):
"""Check if dimensions of tensor in old and new graph is equal.
Parameters
----------
raw_graph_node : tf.Tensor
node of the receiving graph
old_graph_node : tf.Tensor
node of the graph from which will node be extracted
node_name : str
name of the node
Raises
------
RuntimeError
if node dimension do not match
"""
raw_graph_dim = raw_graph_node[node_name].tensor_shape
old_graph_dim = old_graph_node[node_name].tensor_shape
if raw_graph_dim != old_graph_dim:
raise RuntimeError(
f"old graph {old_graph_dim} and raw graph {raw_graph_dim} "
f"has different {node_name} dim"
)
def load_transform_node(graph: tf.Graph) -> Dict[str, tf.Tensor]:
"""Load nodes and their names from graph to dict.
Parameters
----------
graph : tf.Graph
tensforflow graph
Returns
-------
Dict[str, tf.Tensor]
mapping on graph node names and corresponding tensors
"""
transform_node_pattern = re.compile(TRANSFER_PATTERN)
transform_node = {}
for node in graph.node:
if transform_node_pattern.fullmatch(node.name) is not None:
transform_node[node.name] = node.attr["value"].tensor
return transform_node
_skbuild/linux-x86_64-3.6/setuptools/lib/deepmd/env.py 0 → 100644
View file @ 6b33aeb8
"""Module that sets tensorflow working environment and exports inportant constants."""
import logging
import os
import re
import platform
from configparser import ConfigParser
from importlib import reload
from pathlib import Path
from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple
from packaging.version import Version
import numpy as np
if TYPE_CHECKING:
from types import ModuleType
# import tensorflow v1 compatability
try:
import tensorflow.compat.v1 as tf
tf.disable_v2_behavior()
except ImportError:
import tensorflow as tf
try:
import tensorflow.compat.v2 as tfv2
except ImportError:
tfv2 = None
__all__ = [
"GLOBAL_CONFIG",
"GLOBAL_TF_FLOAT_PRECISION",
"GLOBAL_NP_FLOAT_PRECISION",
"GLOBAL_ENER_FLOAT_PRECISION",
"global_float_prec",
"global_cvt_2_tf_float",
"global_cvt_2_ener_float",
"MODEL_VERSION",
"SHARED_LIB_MODULE",
"default_tf_session_config",
"reset_default_tf_session_config",
"op_module",
"op_grads_module",
"TRANSFER_PATTERN",
"FITTING_NET_PATTERN",
"EMBEDDING_NET_PATTERN",
"TYPE_EMBEDDING_PATTERN",
"ATTENTION_LAYER_PATTERN",
"TF_VERSION"
]
SHARED_LIB_MODULE = "op"
# Python library version
try:
tf_py_version = tf.version.VERSION
except AttributeError:
tf_py_version = tf.__version__
EMBEDDING_NET_PATTERN = str(
r"filter_type_\d+/matrix_\d+_\d+|"
r"filter_type_\d+/bias_\d+_\d+|"
r"filter_type_\d+/idt_\d+_\d+|"
r"filter_type_all/matrix_\d+|"
r"filter_type_all/matrix_\d+_\d+|"
r"filter_type_all/matrix_\d+_\d+_\d+|"
r"filter_type_all/bias_\d+|"
r"filter_type_all/bias_\d+_\d+|"
r"filter_type_all/bias_\d+_\d+_\d+|"
r"filter_type_all/idt_\d+|"
r"filter_type_all/idt_\d+_\d+|"
)
FITTING_NET_PATTERN = str(
r"layer_\d+/matrix|"
r"layer_\d+_type_\d+/matrix|"
r"layer_\d+/bias|"
r"layer_\d+_type_\d+/bias|"
r"layer_\d+/idt|"
r"layer_\d+_type_\d+/idt|"
r"final_layer/matrix|"
r"final_layer_type_\d+/matrix|"
r"final_layer/bias|"
r"final_layer_type_\d+/bias|"
)
TYPE_EMBEDDING_PATTERN = str(
r"type_embed_net+/matrix_\d+|"
r"type_embed_net+/bias_\d+|"
r"type_embed_net+/idt_\d+|"
)
ATTENTION_LAYER_PATTERN = str(
r"attention_layer_\d+/c_query/matrix|"
r"attention_layer_\d+/c_query/bias|"
r"attention_layer_\d+/c_key/matrix|"
r"attention_layer_\d+/c_key/bias|"
r"attention_layer_\d+/c_value/matrix|"
r"attention_layer_\d+/c_value/bias|"
r"attention_layer_\d+/c_out/matrix|"
r"attention_layer_\d+/c_out/bias|"
r"attention_layer_\d+/layer_normalization/beta|"
r"attention_layer_\d+/layer_normalization/gamma|"
r"attention_layer_\d+/layer_normalization_\d+/beta|"
r"attention_layer_\d+/layer_normalization_\d+/gamma|"
)
TRANSFER_PATTERN = \
EMBEDDING_NET_PATTERN + \
FITTING_NET_PATTERN + \
TYPE_EMBEDDING_PATTERN + \
str(
r"descrpt_attr/t_avg|"
r"descrpt_attr/t_std|"
r"fitting_attr/t_fparam_avg|"
r"fitting_attr/t_fparam_istd|"
r"fitting_attr/t_aparam_avg|"
r"fitting_attr/t_aparam_istd|"
r"model_attr/t_tab_info|"
r"model_attr/t_tab_data|"
)
def set_env_if_empty(key: str, value: str, verbose: bool = True):
"""Set environment variable only if it is empty.
Parameters
----------
key : str
env variable name
value : str
env variable value
verbose : bool, optional
if True action will be logged, by default True
"""
if os.environ.get(key) is None:
os.environ[key] = value
if verbose:
logging.warn(
f"Environment variable {key} is empty. Use the default value {value}"
)
def set_mkl():
"""Tuning MKL for the best performance.
References
----------
TF overview
https://www.tensorflow.org/guide/performance/overview
Fixing an issue in numpy built by MKL
https://github.com/ContinuumIO/anaconda-issues/issues/11367
https://github.com/numpy/numpy/issues/12374
check whether the numpy is built by mkl, see
https://github.com/numpy/numpy/issues/14751
"""
if "mkl_rt" in np.__config__.get_info("blas_mkl_info").get("libraries", []):
set_env_if_empty("KMP_BLOCKTIME", "0")
set_env_if_empty(
"KMP_AFFINITY", "granularity=fine,verbose,compact,1,0")
reload(np)
def set_tf_default_nthreads():
"""Set TF internal number of threads to default=automatic selection.
Notes
-----
`TF_INTRA_OP_PARALLELISM_THREADS` and `TF_INTER_OP_PARALLELISM_THREADS`
control TF configuration of multithreading.
"""
if "OMP_NUM_THREADS" not in os.environ or \
"TF_INTRA_OP_PARALLELISM_THREADS" not in os.environ or \
"TF_INTER_OP_PARALLELISM_THREADS" not in os.environ:
logging.warning(
"To get the best performance, it is recommended to adjust "
"the number of threads by setting the environment variables "
"OMP_NUM_THREADS, TF_INTRA_OP_PARALLELISM_THREADS, and "
"TF_INTER_OP_PARALLELISM_THREADS.")
set_env_if_empty("TF_INTRA_OP_PARALLELISM_THREADS", "0", verbose=False)
set_env_if_empty("TF_INTER_OP_PARALLELISM_THREADS", "0", verbose=False)
def get_tf_default_nthreads() -> Tuple[int, int]:
"""Get TF paralellism settings.
Returns
-------
Tuple[int, int]
number of `TF_INTRA_OP_PARALLELISM_THREADS` and
`TF_INTER_OP_PARALLELISM_THREADS`
"""
return int(os.environ.get("TF_INTRA_OP_PARALLELISM_THREADS", "0")), int(
os.environ.get("TF_INTER_OP_PARALLELISM_THREADS", "0")
)
def get_tf_session_config() -> Any:
"""Configure tensorflow session.
Returns
-------
Any
session configure object
"""
set_tf_default_nthreads()
intra, inter = get_tf_default_nthreads()
config = tf.ConfigProto(
gpu_options=tf.GPUOptions(allow_growth=True),
intra_op_parallelism_threads=intra, inter_op_parallelism_threads=inter
)
if Version(tf_py_version) >= Version('1.15') and int(os.environ.get("DP_AUTO_PARALLELIZATION", 0)):
config.graph_options.rewrite_options.custom_optimizers.add().name = "dpparallel"
return config
default_tf_session_config = get_tf_session_config()
def reset_default_tf_session_config(cpu_only: bool):
"""Limit tensorflow session to CPU or not.
Parameters
----------
cpu_only : bool
If enabled, no GPU device is visible to the TensorFlow Session.
"""
global default_tf_session_config
if cpu_only:
default_tf_session_config.device_count['GPU'] = 0
else:
if 'GPU' in default_tf_session_config.device_count:
del default_tf_session_config.device_count['GPU']
def get_module(module_name: str) -> "ModuleType":
"""Load force module.
Returns
-------
ModuleType
loaded force module
Raises
------
FileNotFoundError
if module is not found in directory
"""
if platform.system() == "Windows":
ext = ".dll"
prefix = ""
#elif platform.system() == "Darwin":
# ext = ".dylib"
else:
ext = ".so"
prefix = "lib"
module_file = (
(Path(__file__).parent / SHARED_LIB_MODULE / (prefix + module_name))
.with_suffix(ext)
.resolve()
)
if not module_file.is_file():
raise FileNotFoundError(f"module {module_name} does not exist")
else:
try:
module = tf.load_op_library(str(module_file))
except tf.errors.NotFoundError as e:
# check CXX11_ABI_FLAG is compatiblity
# see https://gcc.gnu.org/onlinedocs/libstdc++/manual/using_dual_abi.html
# ABI should be the same
if 'CXX11_ABI_FLAG' in tf.__dict__:
tf_cxx11_abi_flag = tf.CXX11_ABI_FLAG
else:
tf_cxx11_abi_flag = tf.sysconfig.CXX11_ABI_FLAG
if TF_CXX11_ABI_FLAG != tf_cxx11_abi_flag:
raise RuntimeError(
"This deepmd-kit package was compiled with "
"CXX11_ABI_FLAG=%d, but TensorFlow runtime was compiled "
"with CXX11_ABI_FLAG=%d. These two library ABIs are "
"incompatible and thus an error is raised when loading %s. "
"You need to rebuild deepmd-kit against this TensorFlow "
"runtime." % (
TF_CXX11_ABI_FLAG,
tf_cxx11_abi_flag,
module_name,
)) from e
# different versions may cause incompatibility
# see #406, #447, #557, #774, and #796 for example
# throw a message if versions are different
if TF_VERSION != tf_py_version:
raise RuntimeError(
"The version of TensorFlow used to compile this "
"deepmd-kit package is %s, but the version of TensorFlow "
"runtime you are using is %s. These two versions are "
"incompatible and thus an error is raised when loading %s. "
"You need to install TensorFlow %s, or rebuild deepmd-kit "
"against TensorFlow %s.\nIf you are using a wheel from "
"pypi, you may consider to install deepmd-kit execuating "
"`pip install deepmd-kit --no-binary deepmd-kit` "
"instead." % (
TF_VERSION,
tf_py_version,
module_name,
TF_VERSION,
tf_py_version,
)) from e
error_message = (
"This deepmd-kit package is inconsitent with TensorFlow "
"Runtime, thus an error is raised when loading %s. "
"You need to rebuild deepmd-kit against this TensorFlow "
"runtime." % (
module_name,
)
)
if TF_CXX11_ABI_FLAG == 1:
# #1791
error_message += (
"\nWARNING: devtoolset on RHEL6 and RHEL7 does not support _GLIBCXX_USE_CXX11_ABI=1. "
"See https://bugzilla.redhat.com/show_bug.cgi?id=1546704"
)
raise RuntimeError(error_message) from e
return module
def _get_package_constants(
config_file: Path = Path(__file__).parent / "pkg_config/run_config.ini",
) -> Dict[str, str]:
"""Read package constants set at compile time by CMake to dictionary.
Parameters
----------
config_file : str, optional
path to CONFIG file, by default "pkg_config/run_config.ini"
Returns
-------
Dict[str, str]
dictionary with package constants
"""
config = ConfigParser()
config.read(config_file)
return dict(config.items("CONFIG"))
GLOBAL_CONFIG = _get_package_constants()
MODEL_VERSION = GLOBAL_CONFIG["model_version"]
TF_VERSION = GLOBAL_CONFIG["tf_version"]
TF_CXX11_ABI_FLAG = int(GLOBAL_CONFIG["tf_cxx11_abi_flag"])
op_module = get_module("op_abi")
op_grads_module = get_module("op_grads")
# FLOAT_PREC
dp_float_prec = os.environ.get("DP_INTERFACE_PREC", "high").lower()
if dp_float_prec in ("high", ""):
# default is high
GLOBAL_TF_FLOAT_PRECISION = tf.float64
GLOBAL_NP_FLOAT_PRECISION = np.float64
GLOBAL_ENER_FLOAT_PRECISION = np.float64
global_float_prec = "double"
elif dp_float_prec == "low":
GLOBAL_TF_FLOAT_PRECISION = tf.float32
GLOBAL_NP_FLOAT_PRECISION = np.float32
GLOBAL_ENER_FLOAT_PRECISION = np.float64
global_float_prec = "float"
else:
raise RuntimeError(
"Unsupported float precision option: %s. Supported: high,"
"low. Please set precision with environmental variable "
"DP_INTERFACE_PREC." % dp_float_prec
)
def global_cvt_2_tf_float(xx: tf.Tensor) -> tf.Tensor:
"""Cast tensor to globally set TF precision.
Parameters
----------
xx : tf.Tensor
input tensor
Returns
-------
tf.Tensor
output tensor cast to `GLOBAL_TF_FLOAT_PRECISION`
"""
return tf.cast(xx, GLOBAL_TF_FLOAT_PRECISION)
def global_cvt_2_ener_float(xx: tf.Tensor) -> tf.Tensor:
"""Cast tensor to globally set energy precision.
Parameters
----------
xx : tf.Tensor
input tensor
Returns
-------
tf.Tensor
output tensor cast to `GLOBAL_ENER_FLOAT_PRECISION`
"""
return tf.cast(xx, GLOBAL_ENER_FLOAT_PRECISION)
_skbuild/linux-x86_64-3.6/setuptools/lib/deepmd/fit/__init__.py 0 → 100644
View file @ 6b33aeb8
from .ener import EnerFitting
from .wfc import WFCFitting
from .dipole import DipoleFittingSeA
from .polar import PolarFittingSeA
from .polar import GlobalPolarFittingSeA
from .polar import PolarFittingLocFrame
_skbuild/linux-x86_64-3.6/setuptools/lib/deepmd/fit/dipole.py 0 → 100644
View file @ 6b33aeb8
import warnings
import numpy as np
from typing import Tuple, List
from deepmd.env import tf
from deepmd.common import add_data_requirement, get_activation_func, get_precision, cast_precision
from deepmd.utils.network import one_layer, one_layer_rand_seed_shift
from deepmd.utils.graph import get_fitting_net_variables_from_graph_def
from deepmd.descriptor import DescrptSeA
from deepmd.fit.fitting import Fitting
from deepmd.env import global_cvt_2_tf_float
from deepmd.env import GLOBAL_TF_FLOAT_PRECISION
class DipoleFittingSeA (Fitting) :
"""
Fit the atomic dipole with descriptor se_a
Parameters
----------
descrpt : tf.Tensor
The descrptor
neuron : List[int]
Number of neurons in each hidden layer of the fitting net
resnet_dt : bool
Time-step `dt` in the resnet construction:
y = x + dt * \phi (Wx + b)
sel_type : List[int]
The atom types selected to have an atomic dipole prediction. If is None, all atoms are selected.
seed : int
Random seed for initializing the network parameters.
activation_function : str
The activation function in the embedding net. Supported options are |ACTIVATION_FN|
precision : str
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,
descrpt : tf.Tensor,
neuron : List[int] = [120,120,120],
resnet_dt : bool = True,
sel_type : List[int] = None,
seed : int = None,
activation_function : str = 'tanh',
precision : str = 'default',
uniform_seed: bool = False
) -> None:
"""
Constructor
"""
if not isinstance(descrpt, DescrptSeA) :
raise RuntimeError('DipoleFittingSeA only supports DescrptSeA')
self.ntypes = descrpt.get_ntypes()
self.dim_descrpt = descrpt.get_dim_out()
# args = ClassArg()\
# .add('neuron', list, default = [120,120,120], alias = 'n_neuron')\
# .add('resnet_dt', bool, default = True)\
# .add('sel_type', [list,int], default = [ii for ii in range(self.ntypes)], alias = 'dipole_type')\
# .add('seed', int)\
# .add("activation_function", str, default = "tanh")\
# .add('precision', str, default = "default")
# class_data = args.parse(jdata)
self.n_neuron = neuron
self.resnet_dt = resnet_dt
self.sel_type = sel_type
if self.sel_type is None:
self.sel_type = [ii for ii in range(self.ntypes)]
self.seed = seed
self.uniform_seed = uniform_seed
self.seed_shift = one_layer_rand_seed_shift()
self.fitting_activation_fn = get_activation_func(activation_function)
self.fitting_precision = get_precision(precision)
self.dim_rot_mat_1 = descrpt.get_dim_rot_mat_1()
self.dim_rot_mat = self.dim_rot_mat_1 * 3
self.useBN = False
self.fitting_net_variables = None
self.mixed_prec = None
def get_sel_type(self) -> int:
"""
Get selected type
"""
return self.sel_type
def get_out_size(self) -> int:
"""
Get the output size. Should be 3
"""
return 3
@cast_precision
def build (self,
input_d : tf.Tensor,
rot_mat : tf.Tensor,
natoms : tf.Tensor,
reuse : bool = None,
suffix : str = '') -> tf.Tensor:
"""
Build the computational graph for fitting net
Parameters
----------
input_d
The input descriptor
rot_mat
The rotation matrix from the descriptor.
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
reuse
The weights in the networks should be reused when get the variable.
suffix
Name suffix to identify this descriptor
Returns
-------
dipole
The atomic dipole.
"""
start_index = 0
inputs = tf.reshape(input_d, [-1, natoms[0], self.dim_descrpt])
rot_mat = tf.reshape(rot_mat, [-1, natoms[0], self.dim_rot_mat])
count = 0
outs_list = []
for type_i in range(self.ntypes):
# cut-out inputs
inputs_i = tf.slice (inputs,
[ 0, start_index, 0],
[-1, natoms[2+type_i], -1] )
inputs_i = tf.reshape(inputs_i, [-1, self.dim_descrpt])
rot_mat_i = tf.slice (rot_mat,
[ 0, start_index, 0],
[-1, natoms[2+type_i], -1] )
rot_mat_i = tf.reshape(rot_mat_i, [-1, self.dim_rot_mat_1, 3])
start_index += natoms[2+type_i]
if not type_i in self.sel_type :
continue
layer = inputs_i
for ii in range(0,len(self.n_neuron)) :
if ii >= 1 and self.n_neuron[ii] == self.n_neuron[ii-1] :
layer+= one_layer(layer, self.n_neuron[ii], name='layer_'+str(ii)+'_type_'+str(type_i)+suffix, reuse=reuse, seed = self.seed, use_timestep = self.resnet_dt, activation_fn = self.fitting_activation_fn, precision = self.fitting_precision, uniform_seed = self.uniform_seed, initial_variables = self.fitting_net_variables, mixed_prec = self.mixed_prec)
else :
layer = one_layer(layer, self.n_neuron[ii], name='layer_'+str(ii)+'_type_'+str(type_i)+suffix, reuse=reuse, seed = self.seed, activation_fn = self.fitting_activation_fn, precision = self.fitting_precision, uniform_seed = self.uniform_seed, initial_variables = self.fitting_net_variables, mixed_prec = self.mixed_prec)
if (not self.uniform_seed) and (self.seed is not None): self.seed += self.seed_shift
# (nframes x natoms) x naxis
final_layer = one_layer(layer, self.dim_rot_mat_1, activation_fn = None, name='final_layer_type_'+str(type_i)+suffix, reuse=reuse, seed = self.seed, precision = self.fitting_precision, uniform_seed = self.uniform_seed, initial_variables = self.fitting_net_variables, mixed_prec = self.mixed_prec, final_layer = True)
if (not self.uniform_seed) and (self.seed is not None): self.seed += self.seed_shift
# (nframes x natoms) x 1 * naxis
final_layer = tf.reshape(final_layer, [tf.shape(inputs)[0] * natoms[2+type_i], 1, self.dim_rot_mat_1])
# (nframes x natoms) x 1 x 3(coord)
final_layer = tf.matmul(final_layer, rot_mat_i)
# nframes x natoms x 3
final_layer = tf.reshape(final_layer, [tf.shape(inputs)[0], natoms[2+type_i], 3])
# concat the results
outs_list.append(final_layer)
count += 1
outs = tf.concat(outs_list, axis = 1)
tf.summary.histogram('fitting_net_output', outs)
return tf.reshape(outs, [-1])
# return tf.reshape(outs, [tf.shape(inputs)[0] * natoms[0] * 3 // 3])
def init_variables(self,
graph: tf.Graph,
graph_def: tf.GraphDef,
suffix : str = "",
) -> None:
"""
Init the fitting 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
suffix to name scope
"""
self.fitting_net_variables = get_fitting_net_variables_from_graph_def(graph_def, suffix=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.fitting_precision = get_precision(mixed_prec['output_prec'])
_skbuild/linux-x86_64-3.6/setuptools/lib/deepmd/fit/ener.py 0 → 100644
View file @ 6b33aeb8
import warnings
import numpy as np
from typing import Tuple, List
from packaging.version import Version
from deepmd.env import tf
from deepmd.common import add_data_requirement, get_activation_func, get_precision, cast_precision
from deepmd.utils.network import one_layer_rand_seed_shift
from deepmd.utils.network import one_layer as one_layer_deepmd
from deepmd.utils.type_embed import embed_atom_type
from deepmd.utils.graph import get_fitting_net_variables_from_graph_def, load_graph_def, get_tensor_by_name_from_graph
from deepmd.utils.errors import GraphWithoutTensorError
from deepmd.fit.fitting import Fitting
from deepmd.env import global_cvt_2_tf_float
from deepmd.env import GLOBAL_TF_FLOAT_PRECISION, TF_VERSION
from deepmd.nvnmd.utils.config import nvnmd_cfg
from deepmd.nvnmd.fit.ener import one_layer_nvnmd
class EnerFitting (Fitting):
r"""Fitting the energy of the system. The force and the virial can also be trained.
The potential energy :math:`E` is a fitting network function of the descriptor :math:`\mathcal{D}`:
.. math::
E(\mathcal{D}) = \mathcal{L}^{(n)} \circ \mathcal{L}^{(n-1)}
\circ \cdots \circ \mathcal{L}^{(1)} \circ \mathcal{L}^{(0)}
The first :math:`n` hidden layers :math:`\mathcal{L}^{(0)}, \cdots, \mathcal{L}^{(n-1)}` are given by
.. math::
\mathbf{y}=\mathcal{L}(\mathbf{x};\mathbf{w},\mathbf{b})=
\boldsymbol{\phi}(\mathbf{x}^T\mathbf{w}+\mathbf{b})
where :math:`\mathbf{x} \in \mathbb{R}^{N_1}`$` is the input vector and :math:`\mathbf{y} \in \mathbb{R}^{N_2}`
is the output vector. :math:`\mathbf{w} \in \mathbb{R}^{N_1 \times N_2}` and
:math:`\mathbf{b} \in \mathbb{R}^{N_2}`$` are weights and biases, respectively,
both of which are trainable if `trainable[i]` is `True`. :math:`\boldsymbol{\phi}`
is the activation function.
The output layer :math:`\mathcal{L}^{(n)}` is given by
.. math::
\mathbf{y}=\mathcal{L}^{(n)}(\mathbf{x};\mathbf{w},\mathbf{b})=
\mathbf{x}^T\mathbf{w}+\mathbf{b}
where :math:`\mathbf{x} \in \mathbb{R}^{N_{n-1}}`$` is the input vector and :math:`\mathbf{y} \in \mathbb{R}`
is the output scalar. :math:`\mathbf{w} \in \mathbb{R}^{N_{n-1}}` and
:math:`\mathbf{b} \in \mathbb{R}`$` are weights and bias, respectively,
both of which are trainable if `trainable[n]` is `True`.
Parameters
----------
descrpt
The descrptor :math:`\mathcal{D}`
neuron
Number of neurons :math:`N` in each hidden layer of the fitting net
resnet_dt
Time-step `dt` in the resnet construction:
:math:`y = x + dt * \phi (Wx + b)`
numb_fparam
Number of frame parameter
numb_aparam
Number of atomic parameter
rcond
The condition number for the regression of atomic energy.
tot_ener_zero
Force the total energy to zero. Useful for the charge fitting.
trainable
If the weights of fitting net are trainable.
Suppose that we have :math:`N_l` hidden layers in the fitting net,
this list is of length :math:`N_l + 1`, specifying if the hidden layers and the output layer are trainable.
seed
Random seed for initializing the network parameters.
atom_ener
Specifying atomic energy contribution in vacuum. The `set_davg_zero` key in the descrptor should be set.
activation_function
The activation function :math:`\boldsymbol{\phi}` 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,
descrpt : tf.Tensor,
neuron : List[int] = [120,120,120],
resnet_dt : bool = True,
numb_fparam : int = 0,
numb_aparam : int = 0,
rcond : float = 1e-3,
tot_ener_zero : bool = False,
trainable : List[bool] = None,
seed : int = None,
atom_ener : List[float] = [],
activation_function : str = 'tanh',
precision : str = 'default',
uniform_seed: bool = False
) -> None:
"""
Constructor
"""
# model param
self.ntypes = descrpt.get_ntypes()
self.dim_descrpt = descrpt.get_dim_out()
# args = ()\
# .add('numb_fparam', int, default = 0)\
# .add('numb_aparam', int, default = 0)\
# .add('neuron', list, default = [120,120,120], alias = 'n_neuron')\
# .add('resnet_dt', bool, default = True)\
# .add('rcond', float, default = 1e-3) \
# .add('tot_ener_zero', bool, default = False) \
# .add('seed', int) \
# .add('atom_ener', list, default = [])\
# .add("activation_function", str, default = "tanh")\
# .add("precision", str, default = "default")\
# .add("trainable", [list, bool], default = True)
self.numb_fparam = numb_fparam
self.numb_aparam = numb_aparam
self.n_neuron = neuron
self.resnet_dt = resnet_dt
self.rcond = rcond
self.seed = seed
self.uniform_seed = uniform_seed
self.seed_shift = one_layer_rand_seed_shift()
self.tot_ener_zero = tot_ener_zero
self.fitting_activation_fn = get_activation_func(activation_function)
self.fitting_precision = get_precision(precision)
self.trainable = trainable
if self.trainable is None:
self.trainable = [True for ii in range(len(self.n_neuron) + 1)]
if type(self.trainable) is bool:
self.trainable = [self.trainable] * (len(self.n_neuron)+1)
assert(len(self.trainable) == len(self.n_neuron) + 1), 'length of trainable should be that of n_neuron + 1'
self.atom_ener = []
self.atom_ener_v = atom_ener
for at, ae in enumerate(atom_ener):
if ae is not None:
self.atom_ener.append(tf.constant(ae, self.fitting_precision, name = "atom_%d_ener" % at))
else:
self.atom_ener.append(None)
self.useBN = False
self.bias_atom_e = np.zeros(self.ntypes, dtype=np.float64)
# data requirement
if self.numb_fparam > 0 :
add_data_requirement('fparam', self.numb_fparam, atomic=False, must=True, high_prec=False)
self.fparam_avg = None
self.fparam_std = None
self.fparam_inv_std = None
if self.numb_aparam > 0:
add_data_requirement('aparam', self.numb_aparam, atomic=True, must=True, high_prec=False)
self.aparam_avg = None
self.aparam_std = None
self.aparam_inv_std = None
self.fitting_net_variables = None
self.mixed_prec = None
def get_numb_fparam(self) -> int:
"""
Get the number of frame parameters
"""
return self.numb_fparam
def get_numb_aparam(self) -> int:
"""
Get the number of atomic parameters
"""
return self.numb_fparam
def compute_output_stats(self,
all_stat: dict,
mixed_type: bool = False
) -> None:
"""
Compute the ouput statistics
Parameters
----------
all_stat
must have the following components:
all_stat['energy'] of shape n_sys x n_batch x n_frame
can be prepared by model.make_stat_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.
"""
self.bias_atom_e = self._compute_output_stats(all_stat, rcond=self.rcond, mixed_type=mixed_type)
def _compute_output_stats(self, all_stat, rcond=1e-3, mixed_type=False):
data = all_stat['energy']
# data[sys_idx][batch_idx][frame_idx]
sys_ener = np.array([])
for ss in range(len(data)):
sys_data = []
for ii in range(len(data[ss])):
for jj in range(len(data[ss][ii])):
sys_data.append(data[ss][ii][jj])
sys_data = np.concatenate(sys_data)
sys_ener = np.append(sys_ener, np.average(sys_data))
sys_tynatom = np.array([])
if mixed_type:
data = all_stat['real_natoms_vec']
nsys = len(data)
for ss in range(len(data)):
tmp_tynatom = []
for ii in range(len(data[ss])):
for jj in range(len(data[ss][ii])):
tmp_tynatom.append(data[ss][ii][jj].astype(np.float64))
tmp_tynatom = np.average(np.array(tmp_tynatom), axis=0)
sys_tynatom = np.append(sys_tynatom, tmp_tynatom)
else:
data = all_stat['natoms_vec']
nsys = len(data)
for ss in range(len(data)):
sys_tynatom = np.append(sys_tynatom, data[ss][0].astype(np.float64))
sys_tynatom = np.reshape(sys_tynatom, [nsys,-1])
sys_tynatom = sys_tynatom[:,2:]
if len(self.atom_ener) > 0:
# Atomic energies stats are incorrect if atomic energies are assigned.
# In this situation, we directly use these assigned energies instead of computing stats.
# This will make the loss decrease quickly
assigned_atom_ener = np.array(list((ee for ee in self.atom_ener_v if ee is not None)))
assigned_ener_idx = list((ii for ii, ee in enumerate(self.atom_ener_v) if ee is not None))
# np.dot out size: nframe
sys_ener -= np.dot(sys_tynatom[:, assigned_ener_idx], assigned_atom_ener)
sys_tynatom[:, assigned_ener_idx] = 0.
energy_shift,resd,rank,s_value \
= np.linalg.lstsq(sys_tynatom, sys_ener, rcond = rcond)
if len(self.atom_ener) > 0:
for ii in assigned_ener_idx:
energy_shift[ii] = self.atom_ener_v[ii]
return energy_shift
def compute_input_stats(self,
all_stat : dict,
protection : float = 1e-2) -> None:
"""
Compute the input statistics
Parameters
----------
all_stat
if numb_fparam > 0 must have all_stat['fparam']
if numb_aparam > 0 must have all_stat['aparam']
can be prepared by model.make_stat_input
protection
Divided-by-zero protection
"""
# stat fparam
if self.numb_fparam > 0:
cat_data = np.concatenate(all_stat['fparam'], axis = 0)
cat_data = np.reshape(cat_data, [-1, self.numb_fparam])
self.fparam_avg = np.average(cat_data, axis = 0)
self.fparam_std = np.std(cat_data, axis = 0)
for ii in range(self.fparam_std.size):
if self.fparam_std[ii] < protection:
self.fparam_std[ii] = protection
self.fparam_inv_std = 1./self.fparam_std
# stat aparam
if self.numb_aparam > 0:
sys_sumv = []
sys_sumv2 = []
sys_sumn = []
for ss_ in all_stat['aparam'] :
ss = np.reshape(ss_, [-1, self.numb_aparam])
sys_sumv.append(np.sum(ss, axis = 0))
sys_sumv2.append(np.sum(np.multiply(ss, ss), axis = 0))
sys_sumn.append(ss.shape[0])
sumv = np.sum(sys_sumv, axis = 0)
sumv2 = np.sum(sys_sumv2, axis = 0)
sumn = np.sum(sys_sumn)
self.aparam_avg = (sumv)/sumn
self.aparam_std = self._compute_std(sumv2, sumv, sumn)
for ii in range(self.aparam_std.size):
if self.aparam_std[ii] < protection:
self.aparam_std[ii] = protection
self.aparam_inv_std = 1./self.aparam_std
def _compute_std (self, sumv2, sumv, sumn) :
return np.sqrt(sumv2/sumn - np.multiply(sumv/sumn, sumv/sumn))
def _build_lower(
self,
start_index,
natoms,
inputs,
fparam = None,
aparam = None,
bias_atom_e = 0.0,
suffix = '',
reuse = None
):
# cut-out inputs
inputs_i = tf.slice (inputs,
[ 0, start_index, 0],
[-1, natoms, -1] )
inputs_i = tf.reshape(inputs_i, [-1, self.dim_descrpt])
layer = inputs_i
if fparam is not None:
ext_fparam = tf.tile(fparam, [1, natoms])
ext_fparam = tf.reshape(ext_fparam, [-1, self.numb_fparam])
ext_fparam = tf.cast(ext_fparam,self.fitting_precision)
layer = tf.concat([layer, ext_fparam], axis = 1)
if aparam is not None:
ext_aparam = tf.slice(aparam,
[ 0, start_index * self.numb_aparam],
[-1, natoms * self.numb_aparam])
ext_aparam = tf.reshape(ext_aparam, [-1, self.numb_aparam])
ext_aparam = tf.cast(ext_aparam,self.fitting_precision)
layer = tf.concat([layer, ext_aparam], axis = 1)
if nvnmd_cfg.enable:
one_layer = one_layer_nvnmd
else:
one_layer = one_layer_deepmd
for ii in range(0,len(self.n_neuron)) :
if ii >= 1 and self.n_neuron[ii] == self.n_neuron[ii-1] and (not nvnmd_cfg.enable):
layer+= one_layer(
layer,
self.n_neuron[ii],
name='layer_'+str(ii)+suffix,
reuse=reuse,
seed = self.seed,
use_timestep = self.resnet_dt,
activation_fn = self.fitting_activation_fn,
precision = self.fitting_precision,
trainable = self.trainable[ii],
uniform_seed = self.uniform_seed,
initial_variables = self.fitting_net_variables,
mixed_prec = self.mixed_prec)
else :
layer = one_layer(
layer,
self.n_neuron[ii],
name='layer_'+str(ii)+suffix,
reuse=reuse,
seed = self.seed,
activation_fn = self.fitting_activation_fn,
precision = self.fitting_precision,
trainable = self.trainable[ii],
uniform_seed = self.uniform_seed,
initial_variables = self.fitting_net_variables,
mixed_prec = self.mixed_prec)
if (not self.uniform_seed) and (self.seed is not None): self.seed += self.seed_shift
final_layer = one_layer(
layer,
1,
activation_fn = None,
bavg = bias_atom_e,
name='final_layer'+suffix,
reuse=reuse,
seed = self.seed,
precision = self.fitting_precision,
trainable = self.trainable[-1],
uniform_seed = self.uniform_seed,
initial_variables = self.fitting_net_variables,
mixed_prec = self.mixed_prec,
final_layer = True)
if (not self.uniform_seed) and (self.seed is not None): self.seed += self.seed_shift
return final_layer
@cast_precision
def build (self,
inputs : tf.Tensor,
natoms : tf.Tensor,
input_dict : dict = None,
reuse : bool = None,
suffix : str = '',
) -> tf.Tensor:
"""
Build the computational graph for fitting net
Parameters
----------
inputs
The input descriptor
input_dict
Additional dict for inputs.
if numb_fparam > 0, should have input_dict['fparam']
if numb_aparam > 0, should have input_dict['aparam']
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
reuse
The weights in the networks should be reused when get the variable.
suffix
Name suffix to identify this descriptor
Returns
-------
ener
The system energy
"""
if input_dict is None:
input_dict = {}
bias_atom_e = self.bias_atom_e
type_embedding = input_dict.get('type_embedding', None)
atype = input_dict.get('atype', None)
if self.numb_fparam > 0:
if self.fparam_avg is None:
self.fparam_avg = 0.
if self.fparam_inv_std is None:
self.fparam_inv_std = 1.
if self.numb_aparam > 0:
if self.aparam_avg is None:
self.aparam_avg = 0.
if self.aparam_inv_std is None:
self.aparam_inv_std = 1.
with tf.variable_scope('fitting_attr' + suffix, reuse = reuse) :
t_dfparam = tf.constant(self.numb_fparam,
name = 'dfparam',
dtype = tf.int32)
t_daparam = tf.constant(self.numb_aparam,
name = 'daparam',
dtype = tf.int32)
if type_embedding is not None:
self.t_bias_atom_e = tf.get_variable('t_bias_atom_e',
self.bias_atom_e.shape,
dtype=self.fitting_precision,
trainable=False,
initializer=tf.constant_initializer(self.bias_atom_e))
if self.numb_fparam > 0:
t_fparam_avg = tf.get_variable('t_fparam_avg',
self.numb_fparam,
dtype = GLOBAL_TF_FLOAT_PRECISION,
trainable = False,
initializer = tf.constant_initializer(self.fparam_avg))
t_fparam_istd = tf.get_variable('t_fparam_istd',
self.numb_fparam,
dtype = GLOBAL_TF_FLOAT_PRECISION,
trainable = False,
initializer = tf.constant_initializer(self.fparam_inv_std))
if self.numb_aparam > 0:
t_aparam_avg = tf.get_variable('t_aparam_avg',
self.numb_aparam,
dtype = GLOBAL_TF_FLOAT_PRECISION,
trainable = False,
initializer = tf.constant_initializer(self.aparam_avg))
t_aparam_istd = tf.get_variable('t_aparam_istd',
self.numb_aparam,
dtype = GLOBAL_TF_FLOAT_PRECISION,
trainable = False,
initializer = tf.constant_initializer(self.aparam_inv_std))
inputs = tf.reshape(inputs, [-1, natoms[0], self.dim_descrpt])
if len(self.atom_ener):
# only for atom_ener
nframes = input_dict.get('nframes')
if nframes is not None:
# like inputs, but we don't want to add a dependency on inputs
inputs_zero = tf.zeros((nframes, natoms[0], self.dim_descrpt), dtype=self.fitting_precision)
else:
inputs_zero = tf.zeros_like(inputs, dtype=self.fitting_precision)
if bias_atom_e is not None :
assert(len(bias_atom_e) == self.ntypes)
fparam = None
aparam = None
if self.numb_fparam > 0 :
fparam = input_dict['fparam']
fparam = tf.reshape(fparam, [-1, self.numb_fparam])
fparam = (fparam - t_fparam_avg) * t_fparam_istd
if self.numb_aparam > 0 :
aparam = input_dict['aparam']
aparam = tf.reshape(aparam, [-1, self.numb_aparam])
aparam = (aparam - t_aparam_avg) * t_aparam_istd
aparam = tf.reshape(aparam, [-1, self.numb_aparam * natoms[0]])
if type_embedding is not None:
atype_nall = tf.reshape(atype, [-1, natoms[1]])
self.atype_nloc = tf.reshape(tf.slice(atype_nall, [0, 0], [-1, natoms[0]]), [-1]) ## lammps will make error
atype_embed = tf.nn.embedding_lookup(type_embedding, self.atype_nloc)
else:
atype_embed = None
self.atype_embed = atype_embed
if atype_embed is None:
start_index = 0
outs_list = []
for type_i in range(self.ntypes):
if bias_atom_e is None :
type_bias_ae = 0.0
else :
type_bias_ae = bias_atom_e[type_i]
final_layer = self._build_lower(
start_index, natoms[2+type_i],
inputs, fparam, aparam,
bias_atom_e=type_bias_ae, suffix='_type_'+str(type_i)+suffix, reuse=reuse
)
# concat the results
if type_i < len(self.atom_ener) and self.atom_ener[type_i] is not None:
zero_layer = self._build_lower(
start_index, natoms[2+type_i],
inputs_zero, fparam, aparam,
bias_atom_e=type_bias_ae, suffix='_type_'+str(type_i)+suffix, reuse=True
)
final_layer += self.atom_ener[type_i] - zero_layer
final_layer = tf.reshape(final_layer, [tf.shape(inputs)[0], natoms[2+type_i]])
outs_list.append(final_layer)
start_index += natoms[2+type_i]
# concat the results
# concat once may be faster than multiple concat
outs = tf.concat(outs_list, axis = 1)
# with type embedding
else:
if len(self.atom_ener) > 0:
raise RuntimeError("setting atom_ener is not supported by type embedding")
atype_embed = tf.cast(atype_embed, self.fitting_precision)
type_shape = atype_embed.get_shape().as_list()
inputs = tf.concat(
[tf.reshape(inputs,[-1,self.dim_descrpt]),atype_embed],
axis=1
)
self.dim_descrpt = self.dim_descrpt + type_shape[1]
inputs = tf.reshape(inputs, [-1, natoms[0], self.dim_descrpt])
final_layer = self._build_lower(
0, natoms[0],
inputs, fparam, aparam,
bias_atom_e=0.0, suffix=suffix, reuse=reuse
)
outs = tf.reshape(final_layer, [tf.shape(inputs)[0], natoms[0]])
# add bias
self.atom_ener_before = outs
self.add_type = tf.reshape(tf.nn.embedding_lookup(self.t_bias_atom_e, self.atype_nloc), [tf.shape(inputs)[0], natoms[0]])
outs = outs + self.add_type
self.atom_ener_after = outs
if self.tot_ener_zero:
force_tot_ener = 0.0
outs = tf.reshape(outs, [-1, natoms[0]])
outs_mean = tf.reshape(tf.reduce_mean(outs, axis = 1), [-1, 1])
outs_mean = outs_mean - tf.ones_like(outs_mean, dtype = GLOBAL_TF_FLOAT_PRECISION) * (force_tot_ener/global_cvt_2_tf_float(natoms[0]))
outs = outs - outs_mean
outs = tf.reshape(outs, [-1])
tf.summary.histogram('fitting_net_output', outs)
return tf.reshape(outs, [-1])
def init_variables(self,
graph: tf.Graph,
graph_def: tf.GraphDef,
suffix : str = "",
) -> None:
"""
Init the fitting 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
suffix to name scope
"""
self.fitting_net_variables = get_fitting_net_variables_from_graph_def(graph_def, suffix=suffix)
if self.numb_fparam > 0:
self.fparam_avg = get_tensor_by_name_from_graph(graph, 'fitting_attr%s/t_fparam_avg' % suffix)
self.fparam_inv_std = get_tensor_by_name_from_graph(graph, 'fitting_attr%s/t_fparam_istd' % suffix)
if self.numb_aparam > 0:
self.aparam_avg = get_tensor_by_name_from_graph(graph, 'fitting_attr%s/t_aparam_avg' % suffix)
self.aparam_inv_std = get_tensor_by_name_from_graph(graph, 'fitting_attr%s/t_aparam_istd' % suffix)
try:
self.bias_atom_e = get_tensor_by_name_from_graph(graph, 'fitting_attr%s/t_bias_atom_e' % suffix)
except GraphWithoutTensorError:
# model without type_embedding has no t_bias_atom_e
pass
def enable_compression(self,
model_file: str,
suffix: str = ""
) -> None:
"""
Set the fitting net attributes from the frozen model_file when fparam or aparam is not zero
Parameters
----------
model_file : str
The input frozen model file
suffix : str, optional
The suffix of the scope
"""
if self.numb_fparam > 0 or self.numb_aparam > 0:
graph, _ = load_graph_def(model_file)
if self.numb_fparam > 0:
self.fparam_avg = get_tensor_by_name_from_graph(graph, 'fitting_attr%s/t_fparam_avg' % suffix)
self.fparam_inv_std = get_tensor_by_name_from_graph(graph, 'fitting_attr%s/t_fparam_istd' % suffix)
if self.numb_aparam > 0:
self.aparam_avg = get_tensor_by_name_from_graph(graph, 'fitting_attr%s/t_aparam_avg' % suffix)
self.aparam_inv_std = get_tensor_by_name_from_graph(graph, 'fitting_attr%s/t_aparam_istd' % 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.fitting_precision = get_precision(mixed_prec['output_prec'])
_skbuild/linux-x86_64-3.6/setuptools/lib/deepmd/fit/fitting.py 0 → 100644
View file @ 6b33aeb8
from deepmd.env import tf
from deepmd.utils import Plugin, PluginVariant
class Fitting:
@property
def precision(self) -> tf.DType:
"""Precision of fitting network."""
return self.fitting_precision
def init_variables(self,
graph: tf.Graph,
graph_def: tf.GraphDef,
suffix : str = "",
) -> None:
"""
Init the fitting 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
suffix to name scope
Notes
-----
This method is called by others when the fitting supported initialization from the given variables.
"""
raise NotImplementedError(
"Fitting %s doesn't support initialization from the given variables!" % type(self).__name__)
_skbuild/linux-x86_64-3.6/setuptools/lib/deepmd/fit/polar.py 0 → 100644
View file @ 6b33aeb8
import warnings
import numpy as np
from typing import Tuple, List
from deepmd.env import tf
from deepmd.common import add_data_requirement, cast_precision, get_activation_func, get_precision
from deepmd.utils.network import one_layer, one_layer_rand_seed_shift
from deepmd.utils.graph import get_fitting_net_variables_from_graph_def
from deepmd.descriptor import DescrptLocFrame
from deepmd.descriptor import DescrptSeA
from deepmd.fit.fitting import Fitting
from deepmd.env import global_cvt_2_tf_float
from deepmd.env import GLOBAL_TF_FLOAT_PRECISION
class PolarFittingLocFrame () :
"""
Fitting polarizability with local frame descriptor.
.. deprecated:: 2.0.0
This class is not supported any more.
"""
def __init__ (self, jdata, descrpt) :
if not isinstance(descrpt, DescrptLocFrame) :
raise RuntimeError('PolarFittingLocFrame only supports DescrptLocFrame')
self.ntypes = descrpt.get_ntypes()
self.dim_descrpt = descrpt.get_dim_out()
args = ClassArg()\
.add('neuron', list, default = [120,120,120], alias = 'n_neuron')\
.add('resnet_dt', bool, default = True)\
.add('sel_type', [list,int], default = [ii for ii in range(self.ntypes)], alias = 'pol_type')\
.add('seed', int)\
.add("activation_function", str, default = "tanh")\
.add('precision', str, default = "default")
class_data = args.parse(jdata)
self.n_neuron = class_data['neuron']
self.resnet_dt = class_data['resnet_dt']
self.sel_type = class_data['sel_type']
self.seed = class_data['seed']
self.fitting_activation_fn = get_activation_func(class_data["activation_function"])
self.fitting_precision = get_precision(class_data['precision'])
self.useBN = False
def get_sel_type(self):
return self.sel_type
def get_out_size(self):
return 9
def build (self,
input_d,
rot_mat,
natoms,
reuse = None,
suffix = '') :
start_index = 0
inputs = tf.cast(tf.reshape(input_d, [-1, natoms[0], self.dim_descrpt]), self.fitting_precision)
rot_mat = tf.reshape(rot_mat, [-1, 9 * natoms[0]])
count = 0
outs_list = []
for type_i in range(self.ntypes):
# cut-out inputs
inputs_i = tf.slice (inputs,
[ 0, start_index, 0],
[-1, natoms[2+type_i], -1] )
inputs_i = tf.reshape(inputs_i, [-1, self.dim_descrpt])
rot_mat_i = tf.slice (rot_mat,
[ 0, start_index* 9],
[-1, natoms[2+type_i]* 9] )
rot_mat_i = tf.reshape(rot_mat_i, [-1, 3, 3])
start_index += natoms[2+type_i]
if not type_i in self.sel_type :
continue
layer = inputs_i
for ii in range(0,len(self.n_neuron)) :
if ii >= 1 and self.n_neuron[ii] == self.n_neuron[ii-1] :
layer+= one_layer(layer, self.n_neuron[ii], name='layer_'+str(ii)+'_type_'+str(type_i)+suffix, reuse=reuse, seed = self.seed, use_timestep = self.resnet_dt, activation_fn = self.fitting_activation_fn, precision = self.fitting_precision)
else :
layer = one_layer(layer, self.n_neuron[ii], name='layer_'+str(ii)+'_type_'+str(type_i)+suffix, reuse=reuse, seed = self.seed, activation_fn = self.fitting_activation_fn, precision = self.fitting_precision)
# (nframes x natoms) x 9
final_layer = one_layer(layer, 9, activation_fn = None, name='final_layer_type_'+str(type_i)+suffix, reuse=reuse, seed = self.seed, precision = self.fitting_precision, final_layer = True)
# (nframes x natoms) x 3 x 3
final_layer = tf.reshape(final_layer, [tf.shape(inputs)[0] * natoms[2+type_i], 3, 3])
# (nframes x natoms) x 3 x 3
final_layer = final_layer + tf.transpose(final_layer, perm = [0,2,1])
# (nframes x natoms) x 3 x 3(coord)
final_layer = tf.matmul(final_layer, rot_mat_i)
# (nframes x natoms) x 3(coord) x 3(coord)
final_layer = tf.matmul(rot_mat_i, final_layer, transpose_a = True)
# nframes x natoms x 3 x 3
final_layer = tf.reshape(final_layer, [tf.shape(inputs)[0], natoms[2+type_i], 3, 3])
# concat the results
outs_list.append(final_layer)
count += 1
outs = tf.concat(outs_list, axis = 1)
tf.summary.histogram('fitting_net_output', outs)
return tf.cast(tf.reshape(outs, [-1]), GLOBAL_TF_FLOAT_PRECISION)
class PolarFittingSeA (Fitting) :
"""
Fit the atomic polarizability with descriptor se_a
Parameters
----------
descrpt : tf.Tensor
The descrptor
neuron : List[int]
Number of neurons in each hidden layer of the fitting net
resnet_dt : bool
Time-step `dt` in the resnet construction:
y = x + dt * \phi (Wx + b)
sel_type : List[int]
The atom types selected to have an atomic polarizability prediction. If is None, all atoms are selected.
fit_diag : bool
Fit the diagonal part of the rotational invariant polarizability matrix, which will be converted to normal polarizability matrix by contracting with the rotation matrix.
scale : List[float]
The output of the fitting net (polarizability matrix) for type i atom will be scaled by scale[i]
diag_shift : List[float]
The diagonal part of the polarizability matrix of type i will be shifted by diag_shift[i]. The shift operation is carried out after scale.
seed : int
Random seed for initializing the network parameters.
activation_function : str
The activation function in the embedding net. Supported options are |ACTIVATION_FN|
precision : str
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,
descrpt : tf.Tensor,
neuron : List[int] = [120,120,120],
resnet_dt : bool = True,
sel_type : List[int] = None,
fit_diag : bool = True,
scale : List[float] = None,
shift_diag : bool = True, # YWolfeee: will support the user to decide whether to use this function
#diag_shift : List[float] = None, YWolfeee: will not support the user to assign a shift
seed : int = None,
activation_function : str = 'tanh',
precision : str = 'default',
uniform_seed: bool = False
) -> None:
"""
Constructor
"""
if not isinstance(descrpt, DescrptSeA) :
raise RuntimeError('PolarFittingSeA only supports DescrptSeA')
self.ntypes = descrpt.get_ntypes()
self.dim_descrpt = descrpt.get_dim_out()
# args = ClassArg()\
# .add('neuron', list, default = [120,120,120], alias = 'n_neuron')\
# .add('resnet_dt', bool, default = True)\
# .add('fit_diag', bool, default = True)\
# .add('diag_shift', [list,float], default = [0.0 for ii in range(self.ntypes)])\
# .add('scale', [list,float], default = [1.0 for ii in range(self.ntypes)])\
# .add('sel_type', [list,int], default = [ii for ii in range(self.ntypes)], alias = 'pol_type')\
# .add('seed', int)\
# .add("activation_function", str , default = "tanh")\
# .add('precision', str, default = "default")
# class_data = args.parse(jdata)
self.n_neuron = neuron
self.resnet_dt = resnet_dt
self.sel_type = sel_type
self.fit_diag = fit_diag
self.seed = seed
self.uniform_seed = uniform_seed
self.seed_shift = one_layer_rand_seed_shift()
#self.diag_shift = diag_shift
self.shift_diag = shift_diag
self.scale = scale
self.fitting_activation_fn = get_activation_func(activation_function)
self.fitting_precision = get_precision(precision)
if self.sel_type is None:
self.sel_type = [ii for ii in range(self.ntypes)]
if self.scale is None:
self.scale = [1.0 for ii in range(self.ntypes)]
#if self.diag_shift is None:
# self.diag_shift = [0.0 for ii in range(self.ntypes)]
if type(self.sel_type) is not list:
self.sel_type = [self.sel_type]
self.constant_matrix = np.zeros(len(self.sel_type)) # len(sel_type) x 1, store the average diagonal value
#if type(self.diag_shift) is not list:
# self.diag_shift = [self.diag_shift]
if type(self.scale) is not list:
self.scale = [self.scale]
self.dim_rot_mat_1 = descrpt.get_dim_rot_mat_1()
self.dim_rot_mat = self.dim_rot_mat_1 * 3
self.useBN = False
self.fitting_net_variables = None
self.mixed_prec = None
def get_sel_type(self) -> List[int]:
"""
Get selected atom types
"""
return self.sel_type
def get_out_size(self) -> int:
"""
Get the output size. Should be 9
"""
return 9
def compute_input_stats(self,
all_stat,
protection = 1e-2):
"""
Compute the input statistics
Parameters
----------
all_stat
Dictionary of inputs.
can be prepared by model.make_stat_input
protection
Divided-by-zero protection
"""
if not ('polarizability' in all_stat.keys()):
self.avgeig = np.zeros([9])
warnings.warn('no polarizability data, cannot do data stat. use zeros as guess')
return
data = all_stat['polarizability']
all_tmp = []
for ss in range(len(data)):
tmp = np.concatenate(data[ss], axis = 0)
tmp = np.reshape(tmp, [-1, 3, 3])
tmp,_ = np.linalg.eig(tmp)
tmp = np.absolute(tmp)
tmp = np.sort(tmp, axis = 1)
all_tmp.append(tmp)
all_tmp = np.concatenate(all_tmp, axis = 1)
self.avgeig = np.average(all_tmp, axis = 0)
# YWolfeee: support polar normalization, initialize to a more appropriate point
if self.shift_diag:
mean_polar = np.zeros([len(self.sel_type), 9])
sys_matrix, polar_bias = [], []
for ss in range(len(all_stat['type'])):
atom_has_polar = [w for w in all_stat['type'][ss][0] if (w in self.sel_type)] # select atom with polar
if all_stat['find_atomic_polarizability'][ss] > 0.0:
for itype in range(len(self.sel_type)): # Atomic polar mode, should specify the atoms
index_lis = [index for index, w in enumerate(atom_has_polar) \
if atom_has_polar[index] == self.sel_type[itype]] # select index in this type
sys_matrix.append(np.zeros((1,len(self.sel_type))))
sys_matrix[-1][0,itype] = len(index_lis)
polar_bias.append(np.sum(
all_stat['atomic_polarizability'][ss].reshape((-1,9))[index_lis],axis=0).reshape((1,9)))
else: # No atomic polar in this system, so it should have global polar
if not all_stat['find_polarizability'][ss] > 0.0: # This system is jsut a joke?
continue
# Till here, we have global polar
sys_matrix.append(np.zeros((1,len(self.sel_type)))) # add a line in the equations
for itype in range(len(self.sel_type)): # Atomic polar mode, should specify the atoms
index_lis = [index for index, w in enumerate(atom_has_polar) \
if atom_has_polar[index] == self.sel_type[itype]] # select index in this type
sys_matrix[-1][0,itype] = len(index_lis)
# add polar_bias
polar_bias.append(all_stat['polarizability'][ss].reshape((1,9)))
matrix, bias = np.concatenate(sys_matrix,axis=0), np.concatenate(polar_bias,axis=0)
atom_polar,_,_,_ \
= np.linalg.lstsq(matrix, bias, rcond = 1e-3)
for itype in range(len(self.sel_type)):
self.constant_matrix[itype] = np.mean(np.diagonal(atom_polar[itype].reshape((3,3))))
@cast_precision
def build (self,
input_d : tf.Tensor,
rot_mat : tf.Tensor,
natoms : tf.Tensor,
reuse : bool = None,
suffix : str = '') :
"""
Build the computational graph for fitting net
Parameters
----------
input_d
The input descriptor
rot_mat
The rotation matrix from the descriptor.
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
reuse
The weights in the networks should be reused when get the variable.
suffix
Name suffix to identify this descriptor
Returns
-------
atomic_polar
The atomic polarizability
"""
start_index = 0
inputs = tf.reshape(input_d, [-1, self.dim_descrpt * natoms[0]])
rot_mat = tf.reshape(rot_mat, [-1, self.dim_rot_mat * natoms[0]])
count = 0
outs_list = []
for type_i in range(self.ntypes):
# cut-out inputs
inputs_i = tf.slice (inputs,
[ 0, start_index* self.dim_descrpt],
[-1, natoms[2+type_i]* self.dim_descrpt] )
inputs_i = tf.reshape(inputs_i, [-1, self.dim_descrpt])
rot_mat_i = tf.slice (rot_mat,
[ 0, start_index* self.dim_rot_mat],
[-1, natoms[2+type_i]* self.dim_rot_mat] )
rot_mat_i = tf.reshape(rot_mat_i, [-1, self.dim_rot_mat_1, 3])
start_index += natoms[2+type_i]
if not type_i in self.sel_type :
continue
layer = inputs_i
for ii in range(0,len(self.n_neuron)) :
if ii >= 1 and self.n_neuron[ii] == self.n_neuron[ii-1] :
layer+= one_layer(layer, self.n_neuron[ii], name='layer_'+str(ii)+'_type_'+str(type_i)+suffix, reuse=reuse, seed = self.seed, use_timestep = self.resnet_dt, activation_fn = self.fitting_activation_fn, precision = self.fitting_precision, uniform_seed = self.uniform_seed, initial_variables = self.fitting_net_variables, mixed_prec = self.mixed_prec)
else :
layer = one_layer(layer, self.n_neuron[ii], name='layer_'+str(ii)+'_type_'+str(type_i)+suffix, reuse=reuse, seed = self.seed, activation_fn = self.fitting_activation_fn, precision = self.fitting_precision, uniform_seed = self.uniform_seed, initial_variables = self.fitting_net_variables, mixed_prec = self.mixed_prec)
if (not self.uniform_seed) and (self.seed is not None): self.seed += self.seed_shift
if self.fit_diag :
bavg = np.zeros(self.dim_rot_mat_1)
# bavg[0] = self.avgeig[0]
# bavg[1] = self.avgeig[1]
# bavg[2] = self.avgeig[2]
# (nframes x natoms) x naxis
final_layer = one_layer(layer, self.dim_rot_mat_1, activation_fn = None, name='final_layer_type_'+str(type_i)+suffix, reuse=reuse, seed = self.seed, bavg = bavg, precision = self.fitting_precision, uniform_seed = self.uniform_seed, initial_variables = self.fitting_net_variables, mixed_prec = self.mixed_prec, final_layer = True)
if (not self.uniform_seed) and (self.seed is not None): self.seed += self.seed_shift
# (nframes x natoms) x naxis
final_layer = tf.reshape(final_layer, [tf.shape(inputs)[0] * natoms[2+type_i], self.dim_rot_mat_1])
# (nframes x natoms) x naxis x naxis
final_layer = tf.matrix_diag(final_layer)
else :
bavg = np.zeros(self.dim_rot_mat_1*self.dim_rot_mat_1)
# bavg[0*self.dim_rot_mat_1+0] = self.avgeig[0]
# bavg[1*self.dim_rot_mat_1+1] = self.avgeig[1]
# bavg[2*self.dim_rot_mat_1+2] = self.avgeig[2]
# (nframes x natoms) x (naxis x naxis)
final_layer = one_layer(layer, self.dim_rot_mat_1*self.dim_rot_mat_1, activation_fn = None, name='final_layer_type_'+str(type_i)+suffix, reuse=reuse, seed = self.seed, bavg = bavg, precision = self.fitting_precision, uniform_seed = self.uniform_seed, initial_variables = self.fitting_net_variables, mixed_prec = self.mixed_prec, final_layer = True)
if (not self.uniform_seed) and (self.seed is not None): self.seed += self.seed_shift
# (nframes x natoms) x naxis x naxis
final_layer = tf.reshape(final_layer, [tf.shape(inputs)[0] * natoms[2+type_i], self.dim_rot_mat_1, self.dim_rot_mat_1])
# (nframes x natoms) x naxis x naxis
final_layer = final_layer + tf.transpose(final_layer, perm = [0,2,1])
# (nframes x natoms) x naxis x 3(coord)
final_layer = tf.matmul(final_layer, rot_mat_i)
# (nframes x natoms) x 3(coord) x 3(coord)
final_layer = tf.matmul(rot_mat_i, final_layer, transpose_a = True)
# nframes x natoms x 3 x 3
final_layer = tf.reshape(final_layer, [tf.shape(inputs)[0], natoms[2+type_i], 3, 3])
# shift and scale
sel_type_idx = self.sel_type.index(type_i)
final_layer = final_layer * self.scale[sel_type_idx]
final_layer = final_layer + self.constant_matrix[sel_type_idx] * tf.eye(3, batch_shape=[tf.shape(inputs)[0], natoms[2+type_i]], dtype = GLOBAL_TF_FLOAT_PRECISION)
# concat the results
outs_list.append(final_layer)
count += 1
outs = tf.concat(outs_list, axis = 1)
tf.summary.histogram('fitting_net_output', outs)
return tf.reshape(outs, [-1])
def init_variables(self,
graph: tf.Graph,
graph_def: tf.GraphDef,
suffix : str = "",
) -> None:
"""
Init the fitting 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
suffix to name scope
"""
self.fitting_net_variables = get_fitting_net_variables_from_graph_def(graph_def, suffix=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.fitting_precision = get_precision(mixed_prec['output_prec'])
class GlobalPolarFittingSeA () :
"""
Fit the system polarizability with descriptor se_a
Parameters
----------
descrpt : tf.Tensor
The descrptor
neuron : List[int]
Number of neurons in each hidden layer of the fitting net
resnet_dt : bool
Time-step `dt` in the resnet construction:
y = x + dt * \phi (Wx + b)
sel_type : List[int]
The atom types selected to have an atomic polarizability prediction
fit_diag : bool
Fit the diagonal part of the rotational invariant polarizability matrix, which will be converted to normal polarizability matrix by contracting with the rotation matrix.
scale : List[float]
The output of the fitting net (polarizability matrix) for type i atom will be scaled by scale[i]
diag_shift : List[float]
The diagonal part of the polarizability matrix of type i will be shifted by diag_shift[i]. The shift operation is carried out after scale.
seed : int
Random seed for initializing the network parameters.
activation_function : str
The activation function in the embedding net. Supported options are |ACTIVATION_FN|
precision : str
The precision of the embedding net parameters. Supported options are |PRECISION|
"""
def __init__ (self,
descrpt : tf.Tensor,
neuron : List[int] = [120,120,120],
resnet_dt : bool = True,
sel_type : List[int] = None,
fit_diag : bool = True,
scale : List[float] = None,
diag_shift : List[float] = None,
seed : int = None,
activation_function : str = 'tanh',
precision : str = 'default'
) -> None:
"""
Constructor
"""
if not isinstance(descrpt, DescrptSeA) :
raise RuntimeError('GlobalPolarFittingSeA only supports DescrptSeA')
self.ntypes = descrpt.get_ntypes()
self.dim_descrpt = descrpt.get_dim_out()
self.polar_fitting = PolarFittingSeA(descrpt,
neuron,
resnet_dt,
sel_type,
fit_diag,
scale,
diag_shift,
seed,
activation_function,
precision)
def get_sel_type(self) -> int:
"""
Get selected atom types
"""
return self.polar_fitting.get_sel_type()
def get_out_size(self) -> int:
"""
Get the output size. Should be 9
"""
return self.polar_fitting.get_out_size()
def build (self,
input_d,
rot_mat,
natoms,
reuse = None,
suffix = '') -> tf.Tensor:
"""
Build the computational graph for fitting net
Parameters
----------
input_d
The input descriptor
rot_mat
The rotation matrix from the descriptor.
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
reuse
The weights in the networks should be reused when get the variable.
suffix
Name suffix to identify this descriptor
Returns
-------
polar
The system polarizability
"""
inputs = tf.reshape(input_d, [-1, self.dim_descrpt * natoms[0]])
outs = self.polar_fitting.build(input_d, rot_mat, natoms, reuse, suffix)
# nframes x natoms x 9
outs = tf.reshape(outs, [tf.shape(inputs)[0], -1, 9])
outs = tf.reduce_sum(outs, axis = 1)
tf.summary.histogram('fitting_net_output', outs)
return tf.reshape(outs, [-1])
def init_variables(self,
graph: tf.Graph,
graph_def: tf.GraphDef,
suffix : str = "",
) -> None:
"""
Init the fitting 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
suffix to name scope
"""
self.polar_fitting.init_variables(graph=graph, graph_def=graph_def, suffix=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.polar_fitting.enable_mixed_precision(mixed_prec)
_skbuild/linux-x86_64-3.6/setuptools/lib/deepmd/fit/wfc.py 0 → 100644
View file @ 6b33aeb8
import warnings
import numpy as np
from typing import Tuple, List
from deepmd.env import tf
from deepmd.common import ClassArg, add_data_requirement, get_activation_func, get_precision
from deepmd.utils.network import one_layer, one_layer_rand_seed_shift
from deepmd.descriptor import DescrptLocFrame
from deepmd.env import global_cvt_2_tf_float
from deepmd.env import GLOBAL_TF_FLOAT_PRECISION
class WFCFitting () :
"""
Fitting Wannier function centers (WFCs) with local frame descriptor.
.. deprecated:: 2.0.0
This class is not supported any more.
"""
def __init__ (self, jdata, descrpt):
if not isinstance(descrpt, DescrptLocFrame) :
raise RuntimeError('WFC only supports DescrptLocFrame')
self.ntypes = descrpt.get_ntypes()
self.dim_descrpt = descrpt.get_dim_out()
args = ClassArg()\
.add('neuron', list, default = [120,120,120], alias = 'n_neuron')\
.add('resnet_dt', bool, default = True)\
.add('wfc_numb', int, must = True)\
.add('sel_type', [list,int], default = [ii for ii in range(self.ntypes)], alias = 'wfc_type')\
.add('seed', int)\
.add("activation_function", str, default = "tanh")\
.add('precision', str, default = "default")\
.add('uniform_seed', bool, default = False)
class_data = args.parse(jdata)
self.n_neuron = class_data['neuron']
self.resnet_dt = class_data['resnet_dt']
self.wfc_numb = class_data['wfc_numb']
self.sel_type = class_data['sel_type']
self.seed = class_data['seed']
self.uniform_seed = class_data['uniform_seed']
self.seed_shift = one_layer_rand_seed_shift()
self.fitting_activation_fn = get_activation_func(class_data["activation_function"])
self.fitting_precision = get_precision(class_data['precision'])
self.useBN = False
def get_sel_type(self):
return self.sel_type
def get_wfc_numb(self):
return self.wfc_numb
def get_out_size(self):
return self.wfc_numb * 3
def build (self,
input_d,
rot_mat,
natoms,
reuse = None,
suffix = '') :
start_index = 0
inputs = tf.cast(tf.reshape(input_d, [-1, self.dim_descrpt * natoms[0]]), self.fitting_precision)
rot_mat = tf.reshape(rot_mat, [-1, 9 * natoms[0]])
count = 0
for type_i in range(self.ntypes):
# cut-out inputs
inputs_i = tf.slice (inputs,
[ 0, start_index* self.dim_descrpt],
[-1, natoms[2+type_i]* self.dim_descrpt] )
inputs_i = tf.reshape(inputs_i, [-1, self.dim_descrpt])
rot_mat_i = tf.slice (rot_mat,
[ 0, start_index* 9],
[-1, natoms[2+type_i]* 9] )
rot_mat_i = tf.reshape(rot_mat_i, [-1, 3, 3])
start_index += natoms[2+type_i]
if not type_i in self.sel_type :
continue
layer = inputs_i
for ii in range(0,len(self.n_neuron)) :
if ii >= 1 and self.n_neuron[ii] == self.n_neuron[ii-1] :
layer+= one_layer(layer, self.n_neuron[ii], name='layer_'+str(ii)+'_type_'+str(type_i)+suffix, reuse=reuse, seed = self.seed, use_timestep = self.resnet_dt, activation_fn = self.fitting_activation_fn, precision = self.fitting_precision, uniform_seed = self.uniform_seed)
else :
layer = one_layer(layer, self.n_neuron[ii], name='layer_'+str(ii)+'_type_'+str(type_i)+suffix, reuse=reuse, seed = self.seed, activation_fn = self.fitting_activation_fn, precision = self.fitting_precision, uniform_seed = self.uniform_seed)
if (not self.uniform_seed) and (self.seed is not None): self.seed += self.seed_shift
# (nframes x natoms) x (nwfc x 3)
final_layer = one_layer(layer, self.wfc_numb * 3, activation_fn = None, name='final_layer_type_'+str(type_i)+suffix, reuse=reuse, seed = self.seed, precision = self.fitting_precision, uniform_seed = self.uniform_seed)
if (not self.uniform_seed) and (self.seed is not None): self.seed += self.seed_shift
# (nframes x natoms) x nwfc(wc) x 3(coord_local)
final_layer = tf.reshape(final_layer, [tf.shape(inputs)[0] * natoms[2+type_i], self.wfc_numb, 3])
# (nframes x natoms) x nwfc(wc) x 3(coord)
final_layer = tf.matmul(final_layer, rot_mat_i)
# nframes x natoms x nwfc(wc) x 3(coord_local)
final_layer = tf.reshape(final_layer, [tf.shape(inputs)[0], natoms[2+type_i], self.wfc_numb, 3])
# concat the results
if count == 0:
outs = final_layer
else:
outs = tf.concat([outs, final_layer], axis = 1)
count += 1
tf.summary.histogram('fitting_net_output', outs)
return tf.cast(tf.reshape(outs, [-1]), GLOBAL_TF_FLOAT_PRECISION)
_skbuild/linux-x86_64-3.6/setuptools/lib/deepmd/infer/__init__.py 0 → 100644
View file @ 6b33aeb8
"""Submodule containing all the implemented potentials."""
from pathlib import Path
from typing import Union
from .data_modifier import DipoleChargeModifier
from .deep_dipole import DeepDipole
from .deep_eval import DeepEval
from .deep_polar import DeepGlobalPolar, DeepPolar
from .deep_pot import DeepPot
from .deep_wfc import DeepWFC
from .ewald_recp import EwaldRecp
from .model_devi import calc_model_devi
__all__ = [
"DeepPotential",
"DeepDipole",
"DeepEval",
"DeepGlobalPolar",
"DeepPolar",
"DeepPot",
"DeepWFC",
"DipoleChargeModifier",
"EwaldRecp",
"calc_model_devi"
]
def DeepPotential(
model_file: Union[str, Path],
load_prefix: str = "load",
default_tf_graph: bool = False,
) -> Union[DeepDipole, DeepGlobalPolar, DeepPolar, DeepPot, DeepWFC]:
"""Factory function that will inialize appropriate potential read from `model_file`.
Parameters
----------
model_file: str
The name of the frozen model file.
load_prefix: str
The prefix in the load computational graph
default_tf_graph : bool
If uses the default tf graph, otherwise build a new tf graph for evaluation
Returns
-------
Union[DeepDipole, DeepGlobalPolar, DeepPolar, DeepPot, DeepWFC]
one of the available potentials
Raises
------
RuntimeError
if model file does not correspond to any implementd potential
"""
mf = Path(model_file)
model_type = DeepEval(
mf, load_prefix=load_prefix, default_tf_graph=default_tf_graph
).model_type
if model_type == "ener":
dp = DeepPot(mf, load_prefix=load_prefix, default_tf_graph=default_tf_graph)
elif model_type == "dipole":
dp = DeepDipole(mf, load_prefix=load_prefix, default_tf_graph=default_tf_graph)
elif model_type == "polar":
dp = DeepPolar(mf, load_prefix=load_prefix, default_tf_graph=default_tf_graph)
elif model_type == "global_polar":
dp = DeepGlobalPolar(
mf, load_prefix=load_prefix, default_tf_graph=default_tf_graph
)
elif model_type == "wfc":
dp = DeepWFC(mf, load_prefix=load_prefix, default_tf_graph=default_tf_graph)
else:
raise RuntimeError(f"unknow model type {model_type}")
return dp
_skbuild/linux-x86_64-3.6/setuptools/lib/deepmd/infer/data_modifier.py 0 → 100644
View file @ 6b33aeb8
import os
import numpy as np
from typing import Tuple, List
from deepmd.infer.deep_dipole import DeepDipole
from deepmd.infer.ewald_recp import EwaldRecp
from deepmd.env import tf
from deepmd.common import select_idx_map, make_default_mesh
from deepmd.env import GLOBAL_TF_FLOAT_PRECISION
from deepmd.env import GLOBAL_NP_FLOAT_PRECISION
from deepmd.env import GLOBAL_ENER_FLOAT_PRECISION
from deepmd.env import global_cvt_2_tf_float
from deepmd.env import global_cvt_2_ener_float
from deepmd.env import op_module
from deepmd.utils.sess import run_sess
class DipoleChargeModifier(DeepDipole):
"""
Parameters
----------
model_name
The model file for the DeepDipole model
model_charge_map
Gives the amount of charge for the wfcc
sys_charge_map
Gives the amount of charge for the real atoms
ewald_h
Grid spacing of the reciprocal part of Ewald sum. Unit: A
ewald_beta
Splitting parameter of the Ewald sum. Unit: A^{-1}
"""
def __init__(self,
model_name : str,
model_charge_map : List[float],
sys_charge_map : List[float],
ewald_h : float = 1,
ewald_beta : float = 1
) -> None:
"""
Constructor
"""
# the dipole model is loaded with prefix 'dipole_charge'
self.modifier_prefix = 'dipole_charge'
# init dipole model
DeepDipole.__init__(self,
model_name,
load_prefix = self.modifier_prefix,
default_tf_graph = True)
self.model_name = model_name
self.model_charge_map = model_charge_map
self.sys_charge_map = sys_charge_map
self.sel_type = list(self.get_sel_type())
# init ewald recp
self.ewald_h = ewald_h
self.ewald_beta = ewald_beta
self.er = EwaldRecp(self.ewald_h, self.ewald_beta)
# dimension of dipole
self.ext_dim = 3
self.t_ndesc = self.graph.get_tensor_by_name(os.path.join(self.modifier_prefix, 'descrpt_attr/ndescrpt:0'))
self.t_sela = self.graph.get_tensor_by_name(os.path.join(self.modifier_prefix, 'descrpt_attr/sel:0'))
[self.ndescrpt, self.sel_a] = run_sess(self.sess, [self.t_ndesc, self.t_sela])
self.sel_r = [ 0 for ii in range(len(self.sel_a)) ]
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
assert(self.ndescrpt == self.ndescrpt_a + self.ndescrpt_r)
self.force = None
self.ntypes = len(self.sel_a)
def build_fv_graph(self) -> tf.Tensor:
"""
Build the computational graph for the force and virial inference.
"""
with tf.variable_scope('modifier_attr') :
t_mdl_name = tf.constant(self.model_name,
name = 'mdl_name',
dtype = tf.string)
t_modi_type = tf.constant(self.modifier_prefix,
name = 'type',
dtype = tf.string)
t_mdl_charge_map = tf.constant(' '.join([str(ii) for ii in self.model_charge_map]),
name = 'mdl_charge_map',
dtype = tf.string)
t_sys_charge_map = tf.constant(' '.join([str(ii) for ii in self.sys_charge_map]),
name = 'sys_charge_map',
dtype = tf.string)
t_ewald_h = tf.constant(self.ewald_h,
name = 'ewald_h',
dtype = tf.float64)
t_ewald_b = tf.constant(self.ewald_beta,
name = 'ewald_beta',
dtype = tf.float64)
with self.graph.as_default():
return self._build_fv_graph_inner()
def _build_fv_graph_inner(self):
self.t_ef = tf.placeholder(GLOBAL_TF_FLOAT_PRECISION, [None], name = 't_ef')
nf = 10
nfxnas = 64*nf
nfxna = 192*nf
nf = -1
nfxnas = -1
nfxna = -1
self.t_box_reshape = tf.reshape(self.t_box, [-1, 9])
t_nframes = tf.shape(self.t_box_reshape)[0]
# (nframes x natoms) x ndescrpt
self.descrpt = self.graph.get_tensor_by_name(os.path.join(self.modifier_prefix, 'o_rmat:0'))
self.descrpt_deriv = self.graph.get_tensor_by_name(os.path.join(self.modifier_prefix, 'o_rmat_deriv:0'))
self.nlist = self.graph.get_tensor_by_name(os.path.join(self.modifier_prefix, 'o_nlist:0'))
self.rij = self.graph.get_tensor_by_name(os.path.join(self.modifier_prefix, 'o_rij:0'))
# self.descrpt_reshape = tf.reshape(self.descrpt, [nf, 192 * self.ndescrpt])
# self.descrpt_deriv = tf.reshape(self.descrpt_deriv, [nf, 192 * self.ndescrpt * 3])
# nframes x (natoms_sel x 3)
self.t_ef_reshape = tf.reshape(self.t_ef, [t_nframes, -1])
# nframes x (natoms x 3)
self.t_ef_reshape = self._enrich(self.t_ef_reshape, dof = 3)
# (nframes x natoms) x 3
self.t_ef_reshape = tf.reshape(self.t_ef_reshape, [nfxna, 3])
# nframes x (natoms_sel x 3)
self.t_tensor_reshape = tf.reshape(self.t_tensor, [t_nframes, -1])
# nframes x (natoms x 3)
self.t_tensor_reshape = self._enrich(self.t_tensor_reshape, dof = 3)
# (nframes x natoms) x 3
self.t_tensor_reshape = tf.reshape(self.t_tensor_reshape, [nfxna, 3])
# (nframes x natoms) x ndescrpt
[self.t_ef_d] = tf.gradients(self.t_tensor_reshape, self.descrpt, self.t_ef_reshape)
# nframes x (natoms x ndescrpt)
self.t_ef_d = tf.reshape(self.t_ef_d, [nf, self.t_natoms[0] * self.ndescrpt])
# t_ef_d is force (with -1), prod_forc takes deriv, so we need the opposite
self.t_ef_d_oppo = -self.t_ef_d
force = op_module.prod_force_se_a(self.t_ef_d_oppo,
self.descrpt_deriv,
self.nlist,
self.t_natoms,
n_a_sel = self.nnei_a,
n_r_sel = self.nnei_r)
virial, atom_virial \
= op_module.prod_virial_se_a (self.t_ef_d_oppo,
self.descrpt_deriv,
self.rij,
self.nlist,
self.t_natoms,
n_a_sel = self.nnei_a,
n_r_sel = self.nnei_r)
force = tf.identity(force, name='o_dm_force')
virial = tf.identity(virial, name='o_dm_virial')
atom_virial = tf.identity(atom_virial, name='o_dm_av')
return force, virial, atom_virial
def _enrich(self, dipole, dof = 3):
coll = []
sel_start_idx = 0
for type_i in range(self.ntypes):
if type_i in self.sel_type:
di = tf.slice(dipole,
[ 0, sel_start_idx * dof],
[-1, self.t_natoms[2+type_i] * dof])
sel_start_idx += self.t_natoms[2+type_i]
else:
di = tf.zeros([tf.shape(dipole)[0], self.t_natoms[2+type_i] * dof],
dtype = GLOBAL_TF_FLOAT_PRECISION)
coll.append(di)
return tf.concat(coll, axis = 1)
def _slice_descrpt_deriv(self, deriv):
coll = []
start_idx = 0
for type_i in range(self.ntypes):
if type_i in self.sel_type:
di = tf.slice(deriv,
[ 0, start_idx * self.ndescrpt],
[-1, self.t_natoms[2+type_i] * self.ndescrpt])
coll.append(di)
start_idx += self.t_natoms[2+type_i]
return tf.concat(coll, axis = 1)
def eval(self,
coord : np.ndarray,
box : np.ndarray,
atype : np.ndarray,
eval_fv : bool = True
) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
"""
Evaluate the modification
Parameters
----------
coord
The coordinates of atoms
box
The simulation region. PBC is assumed
atype
The atom types
eval_fv
Evaluate force and virial
Returns
-------
tot_e
The energy modification
tot_f
The force modification
tot_v
The virial modification
"""
atype = np.array(atype, dtype=int)
coord, atype, imap = self.sort_input(coord, atype)
# natoms = coord.shape[1] // 3
natoms = atype.size
nframes = coord.shape[0]
box = np.reshape(box, [nframes, 9])
atype = np.reshape(atype, [natoms])
sel_idx_map = select_idx_map(atype, self.sel_type)
nsel = len(sel_idx_map)
# setup charge
charge = np.zeros([natoms])
for ii in range(natoms):
charge[ii] = self.sys_charge_map[atype[ii]]
charge = np.tile(charge, [nframes, 1])
# add wfcc
all_coord, all_charge, dipole = self._extend_system(coord, box, atype, charge)
# print('compute er')
batch_size = 5
tot_e = []
all_f = []
all_v = []
for ii in range(0,nframes,batch_size):
e,f,v = self.er.eval(all_coord[ii:ii+batch_size], all_charge[ii:ii+batch_size], box[ii:ii+batch_size])
tot_e.append(e)
all_f.append(f)
all_v.append(v)
tot_e = np.concatenate(tot_e, axis = 0)
all_f = np.concatenate(all_f, axis = 0)
all_v = np.concatenate(all_v, axis = 0)
# print('finish er')
# reshape
tot_e.reshape([nframes,1])
tot_f = None
tot_v = None
if self.force is None:
self.force, self.virial, self.av = self.build_fv_graph()
if eval_fv:
# compute f
ext_f = all_f[:,natoms*3:]
corr_f = []
corr_v = []
corr_av = []
for ii in range(0,nframes,batch_size):
f, v, av = self._eval_fv(coord[ii:ii+batch_size], box[ii:ii+batch_size], atype, ext_f[ii:ii+batch_size])
corr_f.append(f)
corr_v.append(v)
corr_av.append(av)
corr_f = np.concatenate(corr_f, axis = 0)
corr_v = np.concatenate(corr_v, axis = 0)
corr_av = np.concatenate(corr_av, axis = 0)
tot_f = all_f[:,:natoms*3] + corr_f
for ii in range(nsel):
orig_idx = sel_idx_map[ii]
tot_f[:,orig_idx*3:orig_idx*3+3] += ext_f[:,ii*3:ii*3+3]
tot_f = self.reverse_map(np.reshape(tot_f, [nframes,-1,3]), imap)
# reshape
tot_f = tot_f.reshape([nframes,natoms,3])
# compute v
dipole3 = np.reshape(dipole, [nframes, nsel, 3])
ext_f3 = np.reshape(ext_f, [nframes, nsel, 3])
ext_f3 = np.transpose(ext_f3, [0, 2, 1])
# fd_corr_v = -np.matmul(ext_f3, dipole3).T.reshape([nframes, 9])
# fd_corr_v = -np.matmul(ext_f3, dipole3)
# fd_corr_v = np.transpose(fd_corr_v, [0, 2, 1]).reshape([nframes, 9])
fd_corr_v = -np.matmul(ext_f3, dipole3).reshape([nframes, 9])
# print(all_v, '\n', corr_v, '\n', fd_corr_v)
tot_v = all_v + corr_v + fd_corr_v
# reshape
tot_v = tot_v.reshape([nframes,9])
return tot_e, tot_f, tot_v
def _eval_fv(self, coords, cells, atom_types, ext_f) :
# reshape the inputs
cells = np.reshape(cells, [-1, 9])
nframes = cells.shape[0]
coords = np.reshape(coords, [nframes, -1])
natoms = coords.shape[1] // 3
# sort inputs
coords, atom_types, imap, sel_at, sel_imap = self.sort_input(coords, atom_types, sel_atoms = self.get_sel_type())
# make natoms_vec and default_mesh
natoms_vec = self.make_natoms_vec(atom_types)
assert(natoms_vec[0] == natoms)
default_mesh = make_default_mesh(cells)
# evaluate
tensor = []
feed_dict_test = {}
feed_dict_test[self.t_natoms] = natoms_vec
feed_dict_test[self.t_type ] = np.tile(atom_types, [nframes, 1]).reshape([-1])
feed_dict_test[self.t_coord ] = coords.reshape([-1])
feed_dict_test[self.t_box ] = cells.reshape([-1])
feed_dict_test[self.t_mesh ] = default_mesh.reshape([-1])
feed_dict_test[self.t_ef ] = ext_f.reshape([-1])
# print(run_sess(self.sess, tf.shape(self.t_tensor), feed_dict = feed_dict_test))
fout, vout, avout \
= run_sess(self.sess, [self.force, self.virial, self.av],
feed_dict = feed_dict_test)
# print('fout: ', fout.shape, fout)
fout = self.reverse_map(np.reshape(fout, [nframes,-1,3]), imap)
fout = np.reshape(fout, [nframes, -1])
return fout, vout, avout
def _extend_system(self, coord, box, atype, charge):
natoms = coord.shape[1] // 3
nframes = coord.shape[0]
# sel atoms and setup ref coord
sel_idx_map = select_idx_map(atype, self.sel_type)
nsel = len(sel_idx_map)
coord3 = coord.reshape([nframes, natoms, 3])
ref_coord = coord3[:,sel_idx_map,:]
ref_coord = np.reshape(ref_coord, [nframes, nsel * 3])
batch_size = 8
all_dipole = []
for ii in range(0,nframes,batch_size):
dipole = DeepDipole.eval(self,
coord[ii:ii+batch_size],
box[ii:ii+batch_size],
atype)
all_dipole.append(dipole)
dipole = np.concatenate(all_dipole, axis = 0)
assert(dipole.shape[0] == nframes)
dipole = np.reshape(dipole, [nframes, nsel * 3])
wfcc_coord = ref_coord + dipole
# wfcc_coord = dipole
wfcc_charge = np.zeros([nsel])
for ii in range(nsel):
orig_idx = self.sel_type.index(atype[sel_idx_map[ii]])
wfcc_charge[ii] = self.model_charge_map[orig_idx]
wfcc_charge = np.tile(wfcc_charge, [nframes, 1])
wfcc_coord = np.reshape(wfcc_coord, [nframes, nsel * 3])
wfcc_charge = np.reshape(wfcc_charge, [nframes, nsel])
all_coord = np.concatenate((coord, wfcc_coord), axis = 1)
all_charge = np.concatenate((charge, wfcc_charge), axis = 1)
return all_coord, all_charge, dipole
def modify_data(self,
data : dict) -> None:
"""
Modify data.
Parameters
----------
data
Internal data of DeepmdData.
Be a dict, has the following keys
- coord coordinates
- box simulation box
- type atom types
- find_energy tells if data has energy
- find_force tells if data has force
- find_virial tells if data has virial
- energy energy
- force force
- virial virial
"""
if 'find_energy' not in data and 'find_force' not in data and 'find_virial' not in data:
return
get_nframes=None
coord = data['coord'][:get_nframes,:]
box = data['box'][:get_nframes,:]
atype = data['type'][:get_nframes,:]
atype = atype[0]
nframes = coord.shape[0]
tot_e, tot_f, tot_v = self.eval(coord, box, atype)
# print(tot_f[:,0])
if 'find_energy' in data and data['find_energy'] == 1.0 :
data['energy'] -= tot_e.reshape(data['energy'].shape)
if 'find_force' in data and data['find_force'] == 1.0 :
data['force'] -= tot_f.reshape(data['force'].shape)
if 'find_virial' in data and data['find_virial'] == 1.0 :
data['virial'] -= tot_v.reshape(data['virial'].shape)
_skbuild/linux-x86_64-3.6/setuptools/lib/deepmd/infer/deep_dipole.py 0 → 100644
View file @ 6b33aeb8
from deepmd.infer.deep_tensor import DeepTensor
from typing import TYPE_CHECKING
if TYPE_CHECKING:
from pathlib import Path
class DeepDipole(DeepTensor):
"""Constructor.
Parameters
----------
model_file : Path
The name of the frozen model file.
load_prefix: str
The prefix in the load computational graph
default_tf_graph : bool
If uses the default tf graph, otherwise build a new tf graph for evaluation
Warnings
--------
For developers: `DeepTensor` initializer must be called at the end after
`self.tensors` are modified because it uses the data in `self.tensors` dict.
Do not chanage the order!
"""
def __init__(
self, model_file: "Path", load_prefix: str = "load", default_tf_graph: bool = False
) -> None:
# use this in favor of dict update to move attribute from class to
# instance namespace
self.tensors = dict(
{
# output tensor
"t_tensor": "o_dipole:0",
},
**self.tensors
)
DeepTensor.__init__(
self,
model_file,
load_prefix=load_prefix,
default_tf_graph=default_tf_graph,
)
def get_dim_fparam(self) -> int:
"""Unsupported in this model."""
raise NotImplementedError("This model type does not support this attribute")
def get_dim_aparam(self) -> int:
"""Unsupported in this model."""
raise NotImplementedError("This model type does not support this attribute")
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