Commit 56d2b104 authored by zhangqha's avatar zhangqha
Browse files

BladeDISC DeePMD code update

parent 6b33aeb8
Pipeline #180 failed with stages
in 0 seconds
import numpy as np
from .deep_pot import DeepPot
from ..utils.data import DeepmdData
from ..utils.batch_size import AutoBatchSize
from deepmd.common import expand_sys_str
def calc_model_devi_f(fs: np.ndarray):
'''
Parameters
----------
fs : numpy.ndarray
size of `n_models x n_frames x n_atoms x 3`
'''
fs_devi = np.linalg.norm(np.std(fs, axis=0), axis=-1)
max_devi_f = np.max(fs_devi, axis=-1)
min_devi_f = np.min(fs_devi, axis=-1)
avg_devi_f = np.mean(fs_devi, axis=-1)
return max_devi_f, min_devi_f, avg_devi_f
def calc_model_devi_e(es: np.ndarray):
'''
Parameters
----------
es : numpy.ndarray
size of `n_models x n_frames x n_atoms
'''
es_devi = np.std(es, axis=0)
max_devi_e = np.max(es_devi, axis=1)
min_devi_e = np.min(es_devi, axis=1)
avg_devi_e = np.mean(es_devi, axis=1)
return max_devi_e, min_devi_e, avg_devi_e
def calc_model_devi_v(vs: np.ndarray):
'''
Parameters
----------
vs : numpy.ndarray
size of `n_models x n_frames x 9`
'''
vs_devi = np.std(vs, axis=0)
max_devi_v = np.max(vs_devi, axis=-1)
min_devi_v = np.min(vs_devi, axis=-1)
avg_devi_v = np.linalg.norm(vs_devi, axis=-1) / 3
return max_devi_v, min_devi_v, avg_devi_v
def write_model_devi_out(devi: np.ndarray, fname: str, header: str=""):
'''
Parameters
----------
devi : numpy.ndarray
the first column is the steps index
fname : str
the file name to dump
header : str, default=""
the header to dump
'''
assert devi.shape[1] == 7
header = "%s\n%10s" % (header, "step")
for item in 'vf':
header += "%19s%19s%19s" % (f"max_devi_{item}", f"min_devi_{item}", f"avg_devi_{item}")
with open(fname, "ab") as fp:
np.savetxt(fp,
devi,
fmt=['%12d'] + ['%19.6e' for _ in range(6)],
delimiter='',
header=header)
return devi
def _check_tmaps(tmaps, ref_tmap=None):
'''
Check whether type maps are identical
'''
assert isinstance(tmaps, list)
if ref_tmap is None:
ref_tmap = tmaps[0]
assert isinstance(ref_tmap, list)
flag = True
for tmap in tmaps:
if tmap != ref_tmap:
flag = False
break
return flag
def calc_model_devi(coord,
box,
atype,
models,
fname=None,
frequency=1,
):
'''
Python interface to calculate model deviation
Parameters
-----------
coord : numpy.ndarray, `n_frames x n_atoms x 3`
Coordinates of system to calculate
box : numpy.ndarray or None, `n_frames x 3 x 3`
Box to specify periodic boundary condition. If None, no pbc will be used
atype : numpy.ndarray, `n_atoms x 1`
Atom types
models : list of DeepPot models
Models used to evaluate deviation
fname : str or None
File to dump results, default None
frequency : int
Steps between frames (if the system is given by molecular dynamics engine), default 1
Returns
-------
model_devi : numpy.ndarray, `n_frames x 7`
Model deviation results. The first column is index of steps, the other 6 columns are
max_devi_v, min_devi_v, avg_devi_v, max_devi_f, min_devi_f, avg_devi_f.
Examples
--------
>>> from deepmd.infer import calc_model_devi
>>> from deepmd.infer import DeepPot as DP
>>> import numpy as np
>>> coord = np.array([[1,0,0], [0,0,1.5], [1,0,3]]).reshape([1, -1])
>>> cell = np.diag(10 * np.ones(3)).reshape([1, -1])
>>> atype = [1,0,1]
>>> graphs = [DP("graph.000.pb"), DP("graph.001.pb")]
>>> model_devi = calc_model_devi(coord, cell, atype, graphs)
'''
if box is not None:
nopbc = True
else:
nopbc = False
forces = []
virials = []
for dp in models:
ret = dp.eval(
coord,
box,
atype,
)
forces.append(ret[1])
virials.append(ret[2] / len(atype))
forces = np.array(forces)
virials = np.array(virials)
devi = [np.arange(coord.shape[0]) * frequency]
devi += list(calc_model_devi_v(virials))
devi += list(calc_model_devi_f(forces))
devi = np.vstack(devi).T
if fname:
write_model_devi_out(devi, fname)
return devi
def make_model_devi(
*,
models: list,
system: str,
set_prefix: str,
output: str,
frequency: int,
**kwargs
):
'''
Make model deviation calculation
Parameters
----------
models: list
A list of paths of models to use for making model deviation
system: str
The path of system to make model deviation calculation
set_prefix: str
The set prefix of the system
output: str
The output file for model deviation results
frequency: int
The number of steps that elapse between writing coordinates
in a trajectory by a MD engine (such as Gromacs / Lammps).
This paramter is used to determine the index in the output file.
'''
auto_batch_size = AutoBatchSize()
# init models
dp_models = [DeepPot(model, auto_batch_size=auto_batch_size) for model in models]
# check type maps
tmaps = [dp.get_type_map() for dp in dp_models]
if _check_tmaps(tmaps):
tmap = tmaps[0]
else:
raise RuntimeError("The models does not have the same type map.")
all_sys = expand_sys_str(system)
if len(all_sys) == 0:
raise RuntimeError("Did not find valid system")
devis_coll = []
for system in all_sys:
# create data-system
dp_data = DeepmdData(system, set_prefix, shuffle_test=False, type_map=tmap)
data_sets = [dp_data._load_set(set_name) for set_name in dp_data.dirs]
nframes_tot = 0
devis = []
for data in data_sets:
coord = data["coord"]
box = data["box"]
atype = data["type"][0]
if not dp_data.pbc:
box = None
devi = calc_model_devi(coord, box, atype, dp_models)
nframes_tot += coord.shape[0]
devis.append(devi)
devis = np.vstack(devis)
devis[:, 0] = np.arange(nframes_tot) * frequency
write_model_devi_out(devis, output, header=system)
devis_coll.append(devis)
return devis_coll
"""Module taking care of logging duties."""
from .loggers import set_log_handles
__all__ = ["set_log_handles"]
"""Logger initialization for package."""
import logging
import os
from typing import TYPE_CHECKING, Optional
if TYPE_CHECKING:
from pathlib import Path
from mpi4py import MPI
_MPI_APPEND_MODE = MPI.MODE_CREATE | MPI.MODE_APPEND
logging.getLogger(__name__)
__all__ = ["set_log_handles"]
# logger formater
FFORMATTER = logging.Formatter(
"[%(asctime)s] %(app_name)s %(levelname)-7s %(name)-45s %(message)s"
)
CFORMATTER = logging.Formatter(
# "%(app_name)s %(levelname)-7s |-> %(name)-45s %(message)s"
"%(app_name)s %(levelname)-7s %(message)s"
)
FFORMATTER_MPI = logging.Formatter(
"[%(asctime)s] %(app_name)s rank:%(rank)-2s %(levelname)-7s %(name)-45s %(message)s"
)
CFORMATTER_MPI = logging.Formatter(
# "%(app_name)s rank:%(rank)-2s %(levelname)-7s |-> %(name)-45s %(message)s"
"%(app_name)s rank:%(rank)-2s %(levelname)-7s %(message)s"
)
class _AppFilter(logging.Filter):
"""Add field `app_name` to log messages."""
def filter(self, record):
record.app_name = "DEEPMD"
return True
class _MPIRankFilter(logging.Filter):
"""Add MPI rank number to log messages, adds field `rank`."""
def __init__(self, rank: int) -> None:
super().__init__(name="MPI_rank_id")
self.mpi_rank = str(rank)
def filter(self, record):
record.rank = self.mpi_rank
return True
class _MPIMasterFilter(logging.Filter):
"""Filter that lets through only messages emited from rank==0."""
def __init__(self, rank: int) -> None:
super().__init__(name="MPI_master_log")
self.mpi_rank = rank
def filter(self, record):
if self.mpi_rank == 0:
return True
else:
return False
class _MPIFileStream:
"""Wrap MPI.File` so it has the same API as python file streams.
Parameters
----------
filename : Path
disk location of the file stream
MPI : MPI
MPI communicator object
mode : str, optional
file write mode, by default _MPI_APPEND_MODE
"""
def __init__(
self, filename: "Path", MPI: "MPI", mode: str = "_MPI_APPEND_MODE"
) -> None:
self.stream = MPI.File.Open(MPI.COMM_WORLD, filename, mode)
self.stream.Set_atomicity(True)
self.name = "MPIfilestream"
def write(self, msg: str):
"""Write to MPI shared file stream.
Parameters
----------
msg : str
message to write
"""
b = bytearray()
b.extend(map(ord, msg))
self.stream.Write_shared(b)
def close(self):
"""Synchronize and close MPI file stream."""
self.stream.Sync()
self.stream.Close()
class _MPIHandler(logging.FileHandler):
"""Emulate `logging.FileHandler` with MPI shared File that all ranks can write to.
Parameters
----------
filename : Path
file path
MPI : MPI
MPI communicator object
mode : str, optional
file access mode, by default "_MPI_APPEND_MODE"
"""
def __init__(
self,
filename: "Path",
MPI: "MPI",
mode: str = "_MPI_APPEND_MODE",
) -> None:
self.MPI = MPI
super().__init__(filename, mode=mode, encoding=None, delay=False)
def _open(self):
return _MPIFileStream(self.baseFilename, self.MPI, self.mode)
def setStream(self, stream):
"""Stream canot be reasigned in MPI mode."""
raise NotImplementedError("Unable to do for MPI file handler!")
def set_log_handles(
level: int,
log_path: Optional["Path"] = None,
mpi_log: Optional[str] = None
):
"""Set desired level for package loggers and add file handlers.
Parameters
----------
level: int
logging level
log_path: Optional[str]
path to log file, if None logs will be send only to console. If the parent
directory does not exist it will be automatically created, by default None
mpi_log : Optional[str], optional
mpi log type. Has three options. `master` will output logs to file and console
only from rank==0. `collect` will write messages from all ranks to one file
opened under rank==0 and to console. `workers` will open one log file for each
worker designated by its rank, console behaviour is the same as for `collect`.
If this argument is specified, package 'mpi4py' must be already installed.
by default None
Raises
------
RuntimeError
If the argument `mpi_log` is specified, package `mpi4py` is not installed.
References
----------
https://groups.google.com/g/mpi4py/c/SaNzc8bdj6U
https://stackoverflow.com/questions/35869137/avoid-tensorflow-print-on-standard-error
https://stackoverflow.com/questions/56085015/suppress-openmp-debug-messages-when-running-tensorflow-on-cpu
Notes
-----
Logging levels:
+---------+--------------+----------------+----------------+----------------+
| | our notation | python logging | tensorflow cpp | OpenMP |
+=========+==============+================+================+================+
| debug | 10 | 10 | 0 | 1/on/true/yes |
+---------+--------------+----------------+----------------+----------------+
| info | 20 | 20 | 1 | 0/off/false/no |
+---------+--------------+----------------+----------------+----------------+
| warning | 30 | 30 | 2 | 0/off/false/no |
+---------+--------------+----------------+----------------+----------------+
| error | 40 | 40 | 3 | 0/off/false/no |
+---------+--------------+----------------+----------------+----------------+
"""
# silence logging for OpenMP when running on CPU if level is any other than debug
if level <= 10:
os.environ["KMP_WARNINGS"] = "FALSE"
# set TF cpp internal logging level
os.environ['TF_CPP_MIN_LOG_LEVEL'] = str(int((level / 10) - 1))
# get root logger
root_log = logging.getLogger()
# remove all old handlers
root_log.setLevel(level)
for hdlr in root_log.handlers[:]:
root_log.removeHandler(hdlr)
# check if arguments are present
MPI = None
if mpi_log:
try:
from mpi4py import MPI
except ImportError as e:
raise RuntimeError("You cannot specify 'mpi_log' when mpi4py not installed") from e
# * add console handler ************************************************************
ch = logging.StreamHandler()
if MPI:
rank = MPI.COMM_WORLD.Get_rank()
if mpi_log == "master":
ch.setFormatter(CFORMATTER)
ch.addFilter(_MPIMasterFilter(rank))
else:
ch.setFormatter(CFORMATTER_MPI)
ch.addFilter(_MPIRankFilter(rank))
else:
ch.setFormatter(CFORMATTER)
ch.setLevel(level)
ch.addFilter(_AppFilter())
root_log.addHandler(ch)
# * add file handler ***************************************************************
if log_path:
# create directory
log_path.parent.mkdir(exist_ok=True, parents=True)
fh = None
if mpi_log == "master":
rank = MPI.COMM_WORLD.Get_rank()
if rank == 0:
fh = logging.FileHandler(log_path, mode="w")
fh.addFilter(_MPIMasterFilter(rank))
fh.setFormatter(FFORMATTER)
elif mpi_log == "collect":
rank = MPI.COMM_WORLD.Get_rank()
fh = _MPIHandler(log_path, MPI, mode=MPI.MODE_WRONLY | MPI.MODE_CREATE)
fh.addFilter(_MPIRankFilter(rank))
fh.setFormatter(FFORMATTER_MPI)
elif mpi_log == "workers":
rank = MPI.COMM_WORLD.Get_rank()
# if file has suffix than inser rank number before suffix
# e.g deepmd.log -> deepmd_<rank>.log
# if no suffix is present, insert rank as suffix
# e.g. deepmdlog -> deepmdlog.<rank>
if log_path.suffix:
worker_log = (log_path.parent / f"{log_path.stem}_{rank}").with_suffix(
log_path.suffix
)
else:
worker_log = log_path.with_suffix(f".{rank}")
fh = logging.FileHandler(worker_log, mode="w")
fh.setFormatter(FFORMATTER)
else:
fh = logging.FileHandler(log_path, mode="w")
fh.setFormatter(FFORMATTER)
if fh:
fh.setLevel(level)
fh.addFilter(_AppFilter())
root_log.addHandler(fh)
from .ener import EnerStdLoss
from .ener import EnerDipoleLoss
from .tensor import TensorLoss
import numpy as np
from deepmd.env import tf
from deepmd.common import ClassArg, add_data_requirement
from deepmd.env import global_cvt_2_tf_float
from deepmd.env import global_cvt_2_ener_float
from deepmd.utils.sess import run_sess
from .loss import Loss
class EnerStdLoss (Loss) :
"""
Standard loss function for DP models
Parameters
----------
enable_atom_ener_coeff : bool
if true, the energy will be computed as \sum_i c_i E_i
"""
def __init__ (self,
starter_learning_rate : float,
start_pref_e : float = 0.02,
limit_pref_e : float = 1.00,
start_pref_f : float = 1000,
limit_pref_f : float = 1.00,
start_pref_v : float = 0.0,
limit_pref_v : float = 0.0,
start_pref_ae : float = 0.0,
limit_pref_ae : float = 0.0,
start_pref_pf : float = 0.0,
limit_pref_pf : float = 0.0,
relative_f : float = None,
enable_atom_ener_coeff: bool=False,
) -> None:
self.starter_learning_rate = starter_learning_rate
self.start_pref_e = start_pref_e
self.limit_pref_e = limit_pref_e
self.start_pref_f = start_pref_f
self.limit_pref_f = limit_pref_f
self.start_pref_v = start_pref_v
self.limit_pref_v = limit_pref_v
self.start_pref_ae = start_pref_ae
self.limit_pref_ae = limit_pref_ae
self.start_pref_pf = start_pref_pf
self.limit_pref_pf = limit_pref_pf
self.relative_f = relative_f
self.enable_atom_ener_coeff = enable_atom_ener_coeff
self.has_e = (self.start_pref_e != 0.0 or self.limit_pref_e != 0.0)
self.has_f = (self.start_pref_f != 0.0 or self.limit_pref_f != 0.0)
self.has_v = (self.start_pref_v != 0.0 or self.limit_pref_v != 0.0)
self.has_ae = (self.start_pref_ae != 0.0 or self.limit_pref_ae != 0.0)
self.has_pf = (self.start_pref_pf != 0.0 or self.limit_pref_pf != 0.0)
# data required
add_data_requirement('energy', 1, atomic=False, must=False, high_prec=True)
add_data_requirement('force', 3, atomic=True, must=False, high_prec=False)
add_data_requirement('virial', 9, atomic=False, must=False, high_prec=False)
add_data_requirement('atom_ener', 1, atomic=True, must=False, high_prec=False)
add_data_requirement('atom_pref', 1, atomic=True, must=False, high_prec=False, repeat=3)
if self.enable_atom_ener_coeff:
add_data_requirement('atom_ener_coeff', 1, atomic=True, must=False, high_prec=False, default=1.)
def build (self,
learning_rate,
natoms,
model_dict,
label_dict,
suffix):
energy = model_dict['energy']
force = model_dict['force']
virial = model_dict['virial']
atom_ener = model_dict['atom_ener']
energy_hat = label_dict['energy']
force_hat = label_dict['force']
virial_hat = label_dict['virial']
atom_ener_hat = label_dict['atom_ener']
atom_pref = label_dict['atom_pref']
find_energy = label_dict['find_energy']
find_force = label_dict['find_force']
find_virial = label_dict['find_virial']
find_atom_ener = label_dict['find_atom_ener']
find_atom_pref = label_dict['find_atom_pref']
if self.enable_atom_ener_coeff:
# when ener_coeff (\nu) is defined, the energy is defined as
# E = \sum_i \nu_i E_i
# instead of the sum of atomic energies.
#
# A case is that we want to train reaction energy
# A + B -> C + D
# E = - E(A) - E(B) + E(C) + E(D)
# A, B, C, D could be put far away from each other
atom_ener_coeff = label_dict['atom_ener_coeff']
atom_ener_coeff = tf.reshape(atom_ener_coeff, tf.shape(atom_ener))
energy = tf.reduce_sum(atom_ener_coeff * atom_ener, 1)
l2_ener_loss = tf.reduce_mean( tf.square(energy - energy_hat), name='l2_'+suffix)
force_reshape = tf.reshape (force, [-1])
force_hat_reshape = tf.reshape (force_hat, [-1])
atom_pref_reshape = tf.reshape (atom_pref, [-1])
diff_f = force_hat_reshape - force_reshape
if self.relative_f is not None:
force_hat_3 = tf.reshape(force_hat, [-1, 3])
norm_f = tf.reshape(tf.norm(force_hat_3, axis = 1), [-1, 1]) + self.relative_f
diff_f_3 = tf.reshape(diff_f, [-1, 3])
diff_f_3 = diff_f_3 / norm_f
diff_f = tf.reshape(diff_f_3, [-1])
l2_force_loss = tf.reduce_mean(tf.square(diff_f), name = "l2_force_" + suffix)
l2_pref_force_loss = tf.reduce_mean(tf.multiply(tf.square(diff_f), atom_pref_reshape), name = "l2_pref_force_" + suffix)
virial_reshape = tf.reshape (virial, [-1])
virial_hat_reshape = tf.reshape (virial_hat, [-1])
l2_virial_loss = tf.reduce_mean (tf.square(virial_hat_reshape - virial_reshape), name = "l2_virial_" + suffix)
atom_ener_reshape = tf.reshape (atom_ener, [-1])
atom_ener_hat_reshape = tf.reshape (atom_ener_hat, [-1])
l2_atom_ener_loss = tf.reduce_mean (tf.square(atom_ener_hat_reshape - atom_ener_reshape), name = "l2_atom_ener_" + suffix)
atom_norm = 1./ global_cvt_2_tf_float(natoms[0])
atom_norm_ener = 1./ global_cvt_2_ener_float(natoms[0])
pref_e = global_cvt_2_ener_float(find_energy * (self.limit_pref_e + (self.start_pref_e - self.limit_pref_e) * learning_rate / self.starter_learning_rate) )
pref_f = global_cvt_2_tf_float(find_force * (self.limit_pref_f + (self.start_pref_f - self.limit_pref_f) * learning_rate / self.starter_learning_rate) )
pref_v = global_cvt_2_tf_float(find_virial * (self.limit_pref_v + (self.start_pref_v - self.limit_pref_v) * learning_rate / self.starter_learning_rate) )
pref_ae= global_cvt_2_tf_float(find_atom_ener * (self.limit_pref_ae+ (self.start_pref_ae-self.limit_pref_ae) * learning_rate / self.starter_learning_rate) )
pref_pf= global_cvt_2_tf_float(find_atom_pref * (self.limit_pref_pf+ (self.start_pref_pf-self.limit_pref_pf) * learning_rate / self.starter_learning_rate) )
l2_loss = 0
more_loss = {}
if self.has_e :
l2_loss += atom_norm_ener * (pref_e * l2_ener_loss)
more_loss['l2_ener_loss'] = l2_ener_loss
if self.has_f :
l2_loss += global_cvt_2_ener_float(pref_f * l2_force_loss)
more_loss['l2_force_loss'] = l2_force_loss
if self.has_v :
l2_loss += global_cvt_2_ener_float(atom_norm * (pref_v * l2_virial_loss))
more_loss['l2_virial_loss'] = l2_virial_loss
if self.has_ae :
l2_loss += global_cvt_2_ener_float(pref_ae * l2_atom_ener_loss)
more_loss['l2_atom_ener_loss'] = l2_atom_ener_loss
if self.has_pf :
l2_loss += global_cvt_2_ener_float(pref_pf * l2_pref_force_loss)
more_loss['l2_pref_force_loss'] = l2_pref_force_loss
# only used when tensorboard was set as true
self.l2_loss_summary = tf.summary.scalar('l2_loss', tf.sqrt(l2_loss))
self.l2_loss_ener_summary = tf.summary.scalar('l2_ener_loss', global_cvt_2_tf_float(tf.sqrt(l2_ener_loss)) / global_cvt_2_tf_float(natoms[0]))
self.l2_loss_force_summary = tf.summary.scalar('l2_force_loss', tf.sqrt(l2_force_loss))
self.l2_loss_virial_summary = tf.summary.scalar('l2_virial_loss', tf.sqrt(l2_virial_loss) / global_cvt_2_tf_float(natoms[0]))
self.l2_l = l2_loss
self.l2_more = more_loss
return l2_loss, more_loss
def eval(self, sess, feed_dict, natoms):
placeholder = self.l2_l
run_data = [
self.l2_l,
self.l2_more['l2_ener_loss'] if self.has_e else placeholder,
self.l2_more['l2_force_loss'] if self.has_f else placeholder,
self.l2_more['l2_virial_loss'] if self.has_v else placeholder,
self.l2_more['l2_atom_ener_loss'] if self.has_ae else placeholder,
self.l2_more['l2_pref_force_loss'] if self.has_pf else placeholder,
]
error, error_e, error_f, error_v, error_ae, error_pf = run_sess(sess, run_data, feed_dict=feed_dict)
results = {"natoms": natoms[0], "rmse": np.sqrt(error)}
if self.has_e:
results["rmse_e"] = np.sqrt(error_e) / natoms[0]
if self.has_ae:
results["rmse_ae"] = np.sqrt(error_ae)
if self.has_f:
results["rmse_f"] = np.sqrt(error_f)
if self.has_v:
results["rmse_v"] = np.sqrt(error_v) / natoms[0]
if self.has_pf:
results["rmse_pf"] = np.sqrt(error_pf)
return results
def print_header(self): # depreciated
prop_fmt = ' %11s %11s'
print_str = ''
print_str += prop_fmt % ('rmse_tst', 'rmse_trn')
if self.has_e :
print_str += prop_fmt % ('rmse_e_tst', 'rmse_e_trn')
if self.has_ae :
print_str += prop_fmt % ('rmse_ae_tst', 'rmse_ae_trn')
if self.has_f :
print_str += prop_fmt % ('rmse_f_tst', 'rmse_f_trn')
if self.has_v :
print_str += prop_fmt % ('rmse_v_tst', 'rmse_v_trn')
if self.has_pf :
print_str += prop_fmt % ('rmse_pf_tst', 'rmse_pf_trn')
return print_str
def print_on_training(self,
tb_writer,
cur_batch,
sess,
natoms,
feed_dict_test,
feed_dict_batch): # depreciated
run_data = [
self.l2_l,
self.l2_more['l2_ener_loss'],
self.l2_more['l2_force_loss'],
self.l2_more['l2_virial_loss'],
self.l2_more['l2_atom_ener_loss'],
self.l2_more['l2_pref_force_loss']
]
# first train data
train_out = run_sess(sess, run_data, feed_dict=feed_dict_batch)
error_train, error_e_train, error_f_train, error_v_train, error_ae_train, error_pf_train = train_out
# than test data, if tensorboard log writter is present, commpute summary
# and write tensorboard logs
if tb_writer:
summary_merged_op = tf.summary.merge([self.l2_loss_summary, self.l2_loss_ener_summary, self.l2_loss_force_summary, self.l2_loss_virial_summary])
run_data.insert(0, summary_merged_op)
test_out = run_sess(sess, run_data, feed_dict=feed_dict_test)
if tb_writer:
summary = test_out.pop(0)
tb_writer.add_summary(summary, cur_batch)
error_test, error_e_test, error_f_test, error_v_test, error_ae_test, error_pf_test = test_out
print_str = ""
prop_fmt = " %11.2e %11.2e"
print_str += prop_fmt % (np.sqrt(error_test), np.sqrt(error_train))
if self.has_e :
print_str += prop_fmt % (np.sqrt(error_e_test) / natoms[0], np.sqrt(error_e_train) / natoms[0])
if self.has_ae :
print_str += prop_fmt % (np.sqrt(error_ae_test), np.sqrt(error_ae_train))
if self.has_f :
print_str += prop_fmt % (np.sqrt(error_f_test), np.sqrt(error_f_train))
if self.has_v :
print_str += prop_fmt % (np.sqrt(error_v_test) / natoms[0], np.sqrt(error_v_train) / natoms[0])
if self.has_pf:
print_str += prop_fmt % (np.sqrt(error_pf_test), np.sqrt(error_pf_train))
return print_str
class EnerDipoleLoss (Loss) :
def __init__ (self,
starter_learning_rate : float,
start_pref_e : float = 0.1,
limit_pref_e : float = 1.0,
start_pref_ed : float = 1.0,
limit_pref_ed : float = 1.0
) -> None :
self.starter_learning_rate = kwarg['starter_learning_rate']
args = ClassArg()\
.add('start_pref_e', float, must = True, default = 0.1) \
.add('limit_pref_e', float, must = True, default = 1.00)\
.add('start_pref_ed', float, must = True, default = 1.00)\
.add('limit_pref_ed', float, must = True, default = 1.00)
class_data = args.parse(jdata)
self.start_pref_e = class_data['start_pref_e']
self.limit_pref_e = class_data['limit_pref_e']
self.start_pref_ed = class_data['start_pref_ed']
self.limit_pref_ed = class_data['limit_pref_ed']
# data required
add_data_requirement('energy', 1, atomic=False, must=True, high_prec=True)
add_data_requirement('energy_dipole', 3, atomic=False, must=True, high_prec=False)
def build (self,
learning_rate,
natoms,
model_dict,
label_dict,
suffix):
coord = model_dict['coord']
energy = model_dict['energy']
atom_ener = model_dict['atom_ener']
nframes = tf.shape(atom_ener)[0]
natoms = tf.shape(atom_ener)[1]
# build energy dipole
atom_ener0 = atom_ener - tf.reshape(tf.tile(tf.reshape(energy/global_cvt_2_ener_float(natoms), [-1, 1]), [1, natoms]), [nframes, natoms])
coord = tf.reshape(coord, [nframes, natoms, 3])
atom_ener0 = tf.reshape(atom_ener0, [nframes, 1, natoms])
ener_dipole = tf.matmul(atom_ener0, coord)
ener_dipole = tf.reshape(ener_dipole, [nframes, 3])
energy_hat = label_dict['energy']
ener_dipole_hat = label_dict['energy_dipole']
find_energy = label_dict['find_energy']
find_ener_dipole = label_dict['find_energy_dipole']
l2_ener_loss = tf.reduce_mean( tf.square(energy - energy_hat), name='l2_'+suffix)
ener_dipole_reshape = tf.reshape(ener_dipole, [-1])
ener_dipole_hat_reshape = tf.reshape(ener_dipole_hat, [-1])
l2_ener_dipole_loss = tf.reduce_mean( tf.square(ener_dipole_reshape - ener_dipole_hat_reshape), name='l2_'+suffix)
# atom_norm_ener = 1./ global_cvt_2_ener_float(natoms[0])
atom_norm_ener = 1./ global_cvt_2_ener_float(natoms)
pref_e = global_cvt_2_ener_float(find_energy * (self.limit_pref_e + (self.start_pref_e - self.limit_pref_e) * learning_rate / self.starter_learning_rate) )
pref_ed = global_cvt_2_tf_float(find_ener_dipole * (self.limit_pref_ed + (self.start_pref_ed - self.limit_pref_ed) * learning_rate / self.starter_learning_rate) )
l2_loss = 0
more_loss = {}
l2_loss += atom_norm_ener * (pref_e * l2_ener_loss)
l2_loss += global_cvt_2_ener_float(pref_ed * l2_ener_dipole_loss)
more_loss['l2_ener_loss'] = l2_ener_loss
more_loss['l2_ener_dipole_loss'] = l2_ener_dipole_loss
self.l2_loss_summary = tf.summary.scalar('l2_loss', tf.sqrt(l2_loss))
self.l2_loss_ener_summary = tf.summary.scalar('l2_ener_loss', tf.sqrt(l2_ener_loss) / global_cvt_2_tf_float(natoms[0]))
self.l2_ener_dipole_loss_summary = tf.summary.scalar('l2_ener_dipole_loss', tf.sqrt(l2_ener_dipole_loss))
self.l2_l = l2_loss
self.l2_more = more_loss
return l2_loss, more_loss
def eval(self, sess, feed_dict, natoms):
run_data = [
self.l2_l,
self.l2_more['l2_ener_loss'],
self.l2_more['l2_ener_dipole_loss']
]
error, error_e, error_ed = run_sess(sess, run_data, feed_dict=feed_dict)
results = {
'natoms': natoms[0],
'rmse': np.sqrt(error),
'rmse_e': np.sqrt(error_e) / natoms[0],
'rmse_ed': np.sqrt(error_ed)
}
return results
@staticmethod
def print_header() : # depreciated
prop_fmt = ' %9s %9s'
print_str = ''
print_str += prop_fmt % ('l2_tst', 'l2_trn')
print_str += prop_fmt % ('l2_e_tst', 'l2_e_trn')
print_str += prop_fmt % ('l2_ed_tst', 'l2_ed_trn')
return print_str
def print_on_training(self,
tb_writer,
cur_batch,
sess,
natoms,
feed_dict_test,
feed_dict_batch): # depreciated
run_data = [
self.l2_l,
self.l2_more['l2_ener_loss'],
self.l2_more['l2_ener_dipole_loss']
]
# first train data
train_out = run_sess(sess, run_data, feed_dict=feed_dict_batch)
error_train, error_e_train, error_ed_train = train_out
# than test data, if tensorboard log writter is present, commpute summary
# and write tensorboard logs
if tb_writer:
summary_merged_op = tf.summary.merge([
self.l2_loss_summary,
self.l2_loss_ener_summary,
self.l2_ener_dipole_loss_summary
])
run_data.insert(0, summary_merged_op)
test_out = run_sess(sess, run_data, feed_dict=feed_dict_test)
if tb_writer:
summary = test_out.pop(0)
tb_writer.add_summary(summary, cur_batch)
error_test, error_e_test, error_ed_test = test_out
print_str = ""
prop_fmt = " %9.2e %9.2e"
print_str += prop_fmt % (np.sqrt(error_test), np.sqrt(error_train))
print_str += prop_fmt % (np.sqrt(error_e_test) / natoms[0], np.sqrt(error_e_train) / natoms[0])
print_str += prop_fmt % (np.sqrt(error_ed_test), np.sqrt(error_ed_train))
return print_str
from abc import ABCMeta, abstractmethod
from typing import Tuple, Dict
from deepmd.env import tf
class Loss(metaclass=ABCMeta):
"""The abstract class for the loss function."""
@abstractmethod
def build(self,
learning_rate: tf.Tensor,
natoms: tf.Tensor,
model_dict: Dict[str, tf.Tensor],
label_dict: Dict[str, tf.Tensor],
suffix: str) -> Tuple[tf.Tensor, Dict[str, tf.Tensor]]:
"""Build the loss function graph.
Parameters
----------
learning_rate : tf.Tensor
learning rate
natoms : tf.Tensor
number of atoms
model_dict : dict[str, tf.Tensor]
A dictionary that maps model keys to tensors
label_dict : dict[str, tf.Tensor]
A dictionary that maps label keys to tensors
suffix : str
suffix
Returns
-------
tf.Tensor
the total squared loss
dict[str, tf.Tensor]
A dictionary that maps loss keys to more loss tensors
"""
@abstractmethod
def eval(self,
sess: tf.Session,
feed_dict: Dict[tf.placeholder, tf.Tensor],
natoms: tf.Tensor) -> dict:
"""Eval the loss function.
Parameters
----------
sess : tf.Session
TensorFlow session
feed_dict : dict[tf.placeholder, tf.Tensor]
A dictionary that maps graph elements to values
natoms : tf.Tensor
number of atoms
Returns
-------
dict
A dictionary that maps keys to values. It
should contain key `natoms`
"""
import numpy as np
from deepmd.env import tf
from deepmd.common import ClassArg, add_data_requirement
from deepmd.env import global_cvt_2_tf_float
from deepmd.env import global_cvt_2_ener_float
from deepmd.utils.sess import run_sess
from .loss import Loss
class TensorLoss(Loss) :
"""
Loss function for tensorial properties.
"""
def __init__ (self, jdata, **kwarg) :
model = kwarg.get('model', None)
if model is not None:
self.type_sel = model.get_sel_type()
else:
self.type_sel = None
self.tensor_name = kwarg['tensor_name']
self.tensor_size = kwarg['tensor_size']
self.label_name = kwarg['label_name']
if jdata is not None:
self.scale = jdata.get('scale', 1.0)
else:
self.scale = 1.0
# YHT: added for global / local dipole combination
assert jdata is not None, "Please provide loss parameters!"
# YWolfeee: modify, use pref / pref_atomic, instead of pref_weight / pref_atomic_weight
self.local_weight = jdata.get('pref_atomic', None)
self.global_weight = jdata.get('pref', None)
assert (self.local_weight is not None and self.global_weight is not None), "Both `pref` and `pref_atomic` should be provided."
assert self.local_weight >= 0.0 and self.global_weight >= 0.0, "Can not assign negative weight to `pref` and `pref_atomic`"
assert (self.local_weight >0.0) or (self.global_weight>0.0), AssertionError('Can not assian zero weight both to `pref` and `pref_atomic`')
# data required
add_data_requirement("atomic_" + self.label_name,
self.tensor_size,
atomic=True,
must=False,
high_prec=False,
type_sel = self.type_sel)
add_data_requirement(self.label_name,
self.tensor_size,
atomic=False,
must=False,
high_prec=False,
type_sel = self.type_sel)
def build (self,
learning_rate,
natoms,
model_dict,
label_dict,
suffix):
polar_hat = label_dict[self.label_name]
atomic_polar_hat = label_dict["atomic_" + self.label_name]
polar = tf.reshape(model_dict[self.tensor_name], [-1])
find_global = label_dict['find_' + self.label_name]
find_atomic = label_dict['find_atomic_' + self.label_name]
# YHT: added for global / local dipole combination
l2_loss = global_cvt_2_tf_float(0.0)
more_loss = {
"local_loss":global_cvt_2_tf_float(0.0),
"global_loss":global_cvt_2_tf_float(0.0)
}
if self.local_weight > 0.0:
local_loss = global_cvt_2_tf_float(find_atomic) * tf.reduce_mean( tf.square(self.scale*(polar - atomic_polar_hat)), name='l2_'+suffix)
more_loss['local_loss'] = local_loss
l2_loss += self.local_weight * local_loss
self.l2_loss_local_summary = tf.summary.scalar('l2_local_loss',
tf.sqrt(more_loss['local_loss']))
if self.global_weight > 0.0: # Need global loss
atoms = 0
if self.type_sel is not None:
for w in self.type_sel:
atoms += natoms[2+w]
else:
atoms = natoms[0]
nframes = tf.shape(polar)[0] // self.tensor_size // atoms
# get global results
global_polar = tf.reshape(tf.reduce_sum(tf.reshape(
polar, [nframes, -1, self.tensor_size]), axis=1),[-1])
#if self.atomic: # If label is local, however
# global_polar_hat = tf.reshape(tf.reduce_sum(tf.reshape(
# polar_hat, [nframes, -1, self.tensor_size]), axis=1),[-1])
#else:
# global_polar_hat = polar_hat
global_loss = global_cvt_2_tf_float(find_global) * tf.reduce_mean( tf.square(self.scale*(global_polar - polar_hat)), name='l2_'+suffix)
more_loss['global_loss'] = global_loss
self.l2_loss_global_summary = tf.summary.scalar('l2_global_loss',
tf.sqrt(more_loss['global_loss']) / global_cvt_2_tf_float(atoms))
# YWolfeee: should only consider atoms with dipole, i.e. atoms
# atom_norm = 1./ global_cvt_2_tf_float(natoms[0])
atom_norm = 1./ global_cvt_2_tf_float(atoms)
global_loss *= atom_norm
l2_loss += self.global_weight * global_loss
self.l2_more = more_loss
self.l2_l = l2_loss
self.l2_loss_summary = tf.summary.scalar('l2_loss', tf.sqrt(l2_loss))
return l2_loss, more_loss
def eval(self, sess, feed_dict, natoms):
atoms = 0
if self.type_sel is not None:
for w in self.type_sel:
atoms += natoms[2+w]
else:
atoms = natoms[0]
run_data = [self.l2_l, self.l2_more['local_loss'], self.l2_more['global_loss']]
error, error_lc, error_gl = run_sess(sess, run_data, feed_dict=feed_dict)
results = {"natoms": atoms, "rmse": np.sqrt(error)}
if self.local_weight > 0.0:
results["rmse_lc"] = np.sqrt(error_lc)
if self.global_weight > 0.0:
results["rmse_gl"] = np.sqrt(error_gl) / atoms
return results
def print_header(self): # depreciated
prop_fmt = ' %11s %11s'
print_str = ''
print_str += prop_fmt % ('rmse_tst', 'rmse_trn')
if self.local_weight > 0.0:
print_str += prop_fmt % ('rmse_lc_tst', 'rmse_lc_trn')
if self.global_weight > 0.0:
print_str += prop_fmt % ('rmse_gl_tst', 'rmse_gl_trn')
return print_str
def print_on_training(self,
tb_writer,
cur_batch,
sess,
natoms,
feed_dict_test,
feed_dict_batch) : # depreciated
# YHT: added to calculate the atoms number
atoms = 0
if self.type_sel is not None:
for w in self.type_sel:
atoms += natoms[2+w]
else:
atoms = natoms[0]
run_data = [self.l2_l, self.l2_more['local_loss'], self.l2_more['global_loss']]
summary_list = [self.l2_loss_summary]
if self.local_weight > 0.0:
summary_list.append(self.l2_loss_local_summary)
if self.global_weight > 0.0:
summary_list.append(self.l2_loss_global_summary)
# first train data
error_train = run_sess(sess, run_data, feed_dict=feed_dict_batch)
# than test data, if tensorboard log writter is present, commpute summary
# and write tensorboard logs
if tb_writer:
#summary_merged_op = tf.summary.merge([self.l2_loss_summary])
summary_merged_op = tf.summary.merge(summary_list)
run_data.insert(0, summary_merged_op)
test_out = run_sess(sess, run_data, feed_dict=feed_dict_test)
if tb_writer:
summary = test_out.pop(0)
tb_writer.add_summary(summary, cur_batch)
error_test = test_out
print_str = ""
prop_fmt = " %11.2e %11.2e"
print_str += prop_fmt % (np.sqrt(error_test[0]), np.sqrt(error_train[0]))
if self.local_weight > 0.0:
print_str += prop_fmt % (np.sqrt(error_test[1]), np.sqrt(error_train[1]) )
if self.global_weight > 0.0:
print_str += prop_fmt % (np.sqrt(error_test[2])/atoms, np.sqrt(error_train[2])/atoms)
return print_str
from .ener import EnerModel
from .tensor import WFCModel
from .tensor import DipoleModel
from .tensor import PolarModel
from .tensor import GlobalPolarModel
import numpy as np
from typing import Tuple, List
from deepmd.env import tf
from deepmd.utils.pair_tab import PairTab
from deepmd.utils.graph import load_graph_def, get_tensor_by_name_from_graph
from deepmd.utils.errors import GraphWithoutTensorError
from deepmd.common import ClassArg
from deepmd.env import global_cvt_2_ener_float, MODEL_VERSION, GLOBAL_TF_FLOAT_PRECISION
from deepmd.env import op_module
from .model import Model
from .model_stat import make_stat_input, merge_sys_stat
class EnerModel(Model) :
"""Energy model.
Parameters
----------
descrpt
Descriptor
fitting
Fitting net
type_map
Mapping atom type to the name (str) of the type.
For example `type_map[1]` gives the name of the type 1.
data_stat_nbatch
Number of frames used for data statistic
data_stat_protect
Protect parameter for atomic energy regression
use_srtab
The table for the short-range pairwise interaction added on top of DP. The table is a text data file with (N_t + 1) * N_t / 2 + 1 columes. The first colume is the distance between atoms. The second to the last columes are energies for pairs of certain types. For example we have two atom types, 0 and 1. The columes from 2nd to 4th are for 0-0, 0-1 and 1-1 correspondingly.
smin_alpha
The short-range tabulated interaction will be swithed according to the distance of the nearest neighbor. This distance is calculated by softmin. This parameter is the decaying parameter in the softmin. It is only required when `use_srtab` is provided.
sw_rmin
The lower boundary of the interpolation between short-range tabulated interaction and DP. It is only required when `use_srtab` is provided.
sw_rmin
The upper boundary of the interpolation between short-range tabulated interaction and DP. It is only required when `use_srtab` is provided.
"""
model_type = 'ener'
def __init__ (
self,
descrpt,
fitting,
typeebd = None,
type_map : List[str] = None,
data_stat_nbatch : int = 10,
data_stat_protect : float = 1e-2,
use_srtab : str = None,
smin_alpha : float = None,
sw_rmin : float = None,
sw_rmax : float = None
) -> None:
"""
Constructor
"""
# descriptor
self.descrpt = descrpt
self.rcut = self.descrpt.get_rcut()
self.ntypes = self.descrpt.get_ntypes()
# fitting
self.fitting = fitting
self.numb_fparam = self.fitting.get_numb_fparam()
# type embedding
self.typeebd = typeebd
# other inputs
if type_map is None:
self.type_map = []
else:
self.type_map = type_map
self.data_stat_nbatch = data_stat_nbatch
self.data_stat_protect = data_stat_protect
self.srtab_name = use_srtab
if self.srtab_name is not None :
self.srtab = PairTab(self.srtab_name)
self.smin_alpha = smin_alpha
self.sw_rmin = sw_rmin
self.sw_rmax = sw_rmax
else :
self.srtab = None
def get_rcut (self) :
return self.rcut
def get_ntypes (self) :
return self.ntypes
def get_type_map (self) :
return self.type_map
def data_stat(self, data):
all_stat = make_stat_input(data, self.data_stat_nbatch, merge_sys = False)
m_all_stat = merge_sys_stat(all_stat)
self._compute_input_stat(m_all_stat, protection=self.data_stat_protect, mixed_type=data.mixed_type)
self._compute_output_stat(all_stat, mixed_type=data.mixed_type)
# self.bias_atom_e = data.compute_energy_shift(self.rcond)
def _compute_input_stat (self, all_stat, protection=1e-2, mixed_type=False):
if mixed_type:
self.descrpt.compute_input_stats(all_stat['coord'],
all_stat['box'],
all_stat['type'],
all_stat['natoms_vec'],
all_stat['default_mesh'],
all_stat,
mixed_type,
all_stat['real_natoms_vec'])
else:
self.descrpt.compute_input_stats(all_stat['coord'],
all_stat['box'],
all_stat['type'],
all_stat['natoms_vec'],
all_stat['default_mesh'],
all_stat)
self.fitting.compute_input_stats(all_stat, protection=protection)
def _compute_output_stat (self, all_stat, mixed_type=False):
if mixed_type:
self.fitting.compute_output_stats(all_stat, mixed_type=mixed_type)
else:
self.fitting.compute_output_stats(all_stat)
def build (self,
coord_,
atype_,
natoms,
box,
mesh,
input_dict,
frz_model = None,
suffix = '',
reuse = None):
if input_dict is None:
input_dict = {}
with tf.variable_scope('model_attr' + suffix, reuse = reuse) :
t_tmap = tf.constant(' '.join(self.type_map),
name = 'tmap',
dtype = tf.string)
t_mt = tf.constant(self.model_type,
name = 'model_type',
dtype = tf.string)
t_ver = tf.constant(MODEL_VERSION,
name = 'model_version',
dtype = tf.string)
if self.srtab is not None :
tab_info, tab_data = self.srtab.get()
self.tab_info = tf.get_variable('t_tab_info',
tab_info.shape,
dtype = tf.float64,
trainable = False,
initializer = tf.constant_initializer(tab_info, dtype = tf.float64))
self.tab_data = tf.get_variable('t_tab_data',
tab_data.shape,
dtype = tf.float64,
trainable = False,
initializer = tf.constant_initializer(tab_data, dtype = tf.float64))
coord = tf.reshape (coord_, [-1, natoms[1] * 3])
atype = tf.reshape (atype_, [-1, natoms[1]])
input_dict['nframes'] = tf.shape(coord)[0]
# type embedding if any
if self.typeebd is not None:
type_embedding = self.typeebd.build(
self.ntypes,
reuse = reuse,
suffix = suffix,
)
input_dict['type_embedding'] = type_embedding
input_dict['atype'] = atype_
if frz_model == None:
dout \
= self.descrpt.build(coord_,
atype_,
natoms,
box,
mesh,
input_dict,
suffix = suffix,
reuse = reuse)
dout = tf.identity(dout, name='o_descriptor')
else:
tf.constant(self.rcut,
name = 'descrpt_attr/rcut',
dtype = GLOBAL_TF_FLOAT_PRECISION)
tf.constant(self.ntypes,
name = 'descrpt_attr/ntypes',
dtype = tf.int32)
feed_dict = self.descrpt.get_feed_dict(coord_, atype_, natoms, box, mesh)
return_elements = [*self.descrpt.get_tensor_names(), 'o_descriptor:0']
imported_tensors \
= self._import_graph_def_from_frz_model(frz_model, feed_dict, return_elements)
dout = imported_tensors[-1]
self.descrpt.pass_tensors_from_frz_model(*imported_tensors[:-1])
if self.srtab is not None :
nlist, rij, sel_a, sel_r = self.descrpt.get_nlist()
nnei_a = np.cumsum(sel_a)[-1]
nnei_r = np.cumsum(sel_r)[-1]
atom_ener = self.fitting.build (dout,
natoms,
input_dict,
reuse = reuse,
suffix = suffix)
self.atom_ener = atom_ener
if self.srtab is not None :
sw_lambda, sw_deriv \
= op_module.soft_min_switch(atype,
rij,
nlist,
natoms,
sel_a = sel_a,
sel_r = sel_r,
alpha = self.smin_alpha,
rmin = self.sw_rmin,
rmax = self.sw_rmax)
inv_sw_lambda = 1.0 - sw_lambda
# NOTICE:
# atom energy is not scaled,
# force and virial are scaled
tab_atom_ener, tab_force, tab_atom_virial \
= op_module.pair_tab(self.tab_info,
self.tab_data,
atype,
rij,
nlist,
natoms,
sw_lambda,
sel_a = sel_a,
sel_r = sel_r)
energy_diff = tab_atom_ener - tf.reshape(atom_ener, [-1, natoms[0]])
tab_atom_ener = tf.reshape(sw_lambda, [-1]) * tf.reshape(tab_atom_ener, [-1])
atom_ener = tf.reshape(inv_sw_lambda, [-1]) * atom_ener
energy_raw = tab_atom_ener + atom_ener
else :
energy_raw = atom_ener
energy_raw = tf.reshape(energy_raw, [-1, natoms[0]], name = 'o_atom_energy'+suffix)
energy = tf.reduce_sum(global_cvt_2_ener_float(energy_raw), axis=1, name='o_energy'+suffix)
force, virial, atom_virial \
= self.descrpt.prod_force_virial (atom_ener, natoms)
if self.srtab is not None :
sw_force \
= op_module.soft_min_force(energy_diff,
sw_deriv,
nlist,
natoms,
n_a_sel = nnei_a,
n_r_sel = nnei_r)
force = force + sw_force + tab_force
force = tf.reshape (force, [-1, 3 * natoms[1]], name = "o_force"+suffix)
if self.srtab is not None :
sw_virial, sw_atom_virial \
= op_module.soft_min_virial (energy_diff,
sw_deriv,
rij,
nlist,
natoms,
n_a_sel = nnei_a,
n_r_sel = nnei_r)
atom_virial = atom_virial + sw_atom_virial + tab_atom_virial
virial = virial + sw_virial \
+ tf.reduce_sum(tf.reshape(tab_atom_virial, [-1, natoms[1], 9]), axis = 1)
virial = tf.reshape (virial, [-1, 9], name = "o_virial"+suffix)
atom_virial = tf.reshape (atom_virial, [-1, 9 * natoms[1]], name = "o_atom_virial"+suffix)
model_dict = {}
model_dict['energy'] = energy
model_dict['force'] = force
model_dict['virial'] = virial
model_dict['atom_ener'] = energy_raw
model_dict['atom_virial'] = atom_virial
model_dict['coord'] = coord
model_dict['atype'] = atype
return model_dict
def _import_graph_def_from_frz_model(self, frz_model, feed_dict, return_elements):
graph, graph_def = load_graph_def(frz_model)
return tf.import_graph_def(graph_def, input_map = feed_dict, return_elements = return_elements, name = "")
def init_variables(self,
graph : tf.Graph,
graph_def : tf.GraphDef,
model_type : str = "original_model",
suffix : str = "",
) -> None:
"""
Init the embedding net variables with the given frozen model
Parameters
----------
graph : tf.Graph
The input frozen model graph
graph_def : tf.GraphDef
The input frozen model graph_def
model_type : str
the type of the model
suffix : str
suffix to name scope
"""
# self.frz_model will control the self.model to import the descriptor from the given frozen model instead of building from scratch...
# initialize fitting net with the given compressed frozen model
if model_type == 'original_model':
self.descrpt.init_variables(graph, graph_def, suffix=suffix)
self.fitting.init_variables(graph, graph_def, suffix=suffix)
tf.constant("original_model", name = 'model_type', dtype = tf.string)
elif model_type == 'compressed_model':
self.fitting.init_variables(graph, graph_def, suffix=suffix)
tf.constant("compressed_model", name = 'model_type', dtype = tf.string)
else:
raise RuntimeError("Unknown model type %s" % model_type)
if self.typeebd is not None:
self.typeebd.init_variables(graph, graph_def, suffix=suffix)
from deepmd.env import tf
class Model:
def init_variables(self,
graph : tf.Graph,
graph_def : tf.GraphDef,
model_type : str = "original_model",
suffix : str = "",
) -> None:
"""
Init the embedding net variables with the given frozen model
Parameters
----------
graph : tf.Graph
The input frozen model graph
graph_def : tf.GraphDef
The input frozen model graph_def
model_type : str
the type of the model
suffix : str
suffix to name scope
"""
raise RuntimeError("The 'dp train init-frz-model' command do not support this model!")
import numpy as np
from collections import defaultdict
def _make_all_stat_ref(data, nbatches):
all_stat = defaultdict(list)
for ii in range(data.get_nsystems()) :
for jj in range(nbatches) :
stat_data = data.get_batch (sys_idx = ii)
for dd in stat_data:
if dd == "natoms_vec":
stat_data[dd] = stat_data[dd].astype(np.int32)
all_stat[dd].append(stat_data[dd])
return all_stat
def make_stat_input(data, nbatches, merge_sys = True):
"""
pack data for statistics
Parameters
----------
data:
The data
merge_sys: bool (True)
Merge system data
Returns
-------
all_stat:
A dictionary of list of list storing data for stat.
if merge_sys == False data can be accessed by
all_stat[key][sys_idx][batch_idx][frame_idx]
else merge_sys == True can be accessed by
all_stat[key][batch_idx][frame_idx]
"""
all_stat = defaultdict(list)
for ii in range(data.get_nsystems()) :
sys_stat = defaultdict(list)
for jj in range(nbatches) :
stat_data = data.get_batch (sys_idx = ii)
for dd in stat_data:
if dd == "natoms_vec":
stat_data[dd] = stat_data[dd].astype(np.int32)
sys_stat[dd].append(stat_data[dd])
for dd in sys_stat:
if merge_sys:
for bb in sys_stat[dd]:
all_stat[dd].append(bb)
else:
all_stat[dd].append(sys_stat[dd])
return all_stat
def merge_sys_stat(all_stat):
first_key = list(all_stat.keys())[0]
nsys = len(all_stat[first_key])
ret = defaultdict(list)
for ii in range(nsys):
for dd in all_stat:
for bb in all_stat[dd][ii]:
ret[dd].append(bb)
return ret
import numpy as np
from typing import Tuple, List
from deepmd.env import tf
from deepmd.common import ClassArg
from deepmd.env import global_cvt_2_ener_float, MODEL_VERSION, GLOBAL_TF_FLOAT_PRECISION
from deepmd.env import op_module
from deepmd.utils.graph import load_graph_def
from .model import Model
from .model_stat import make_stat_input, merge_sys_stat
class TensorModel(Model) :
"""Tensor model.
Parameters
----------
tensor_name
Name of the tensor.
descrpt
Descriptor
fitting
Fitting net
type_map
Mapping atom type to the name (str) of the type.
For example `type_map[1]` gives the name of the type 1.
data_stat_nbatch
Number of frames used for data statistic
data_stat_protect
Protect parameter for atomic energy regression
"""
def __init__ (
self,
tensor_name : str,
descrpt,
fitting,
type_map : List[str] = None,
data_stat_nbatch : int = 10,
data_stat_protect : float = 1e-2,
)->None:
"""
Constructor
"""
self.model_type = tensor_name
# descriptor
self.descrpt = descrpt
self.rcut = self.descrpt.get_rcut()
self.ntypes = self.descrpt.get_ntypes()
# fitting
self.fitting = fitting
# other params
if type_map is None:
self.type_map = []
else:
self.type_map = type_map
self.data_stat_nbatch = data_stat_nbatch
self.data_stat_protect = data_stat_protect
def get_rcut (self) :
return self.rcut
def get_ntypes (self) :
return self.ntypes
def get_type_map (self) :
return self.type_map
def get_sel_type(self):
return self.fitting.get_sel_type()
def get_out_size (self) :
return self.fitting.get_out_size()
def data_stat(self, data):
all_stat = make_stat_input(data, self.data_stat_nbatch, merge_sys = False)
m_all_stat = merge_sys_stat(all_stat)
self._compute_input_stat (m_all_stat, protection = self.data_stat_protect)
self._compute_output_stat(all_stat)
def _compute_input_stat(self, all_stat, protection = 1e-2) :
self.descrpt.compute_input_stats(all_stat['coord'],
all_stat['box'],
all_stat['type'],
all_stat['natoms_vec'],
all_stat['default_mesh'],
all_stat)
if hasattr(self.fitting, 'compute_input_stats'):
self.fitting.compute_input_stats(all_stat, protection = protection)
def _compute_output_stat (self, all_stat) :
if hasattr(self.fitting, 'compute_output_stats'):
self.fitting.compute_output_stats(all_stat)
def build (self,
coord_,
atype_,
natoms,
box,
mesh,
input_dict,
frz_model = None,
suffix = '',
reuse = None):
with tf.variable_scope('model_attr' + suffix, reuse = reuse) :
t_tmap = tf.constant(' '.join(self.type_map),
name = 'tmap',
dtype = tf.string)
t_st = tf.constant(self.get_sel_type(),
name = 'sel_type',
dtype = tf.int32)
t_mt = tf.constant(self.model_type,
name = 'model_type',
dtype = tf.string)
t_ver = tf.constant(MODEL_VERSION,
name = 'model_version',
dtype = tf.string)
t_od = tf.constant(self.get_out_size(),
name = 'output_dim',
dtype = tf.int32)
natomsel = sum(natoms[2+type_i] for type_i in self.get_sel_type())
nout = self.get_out_size()
if frz_model == None:
dout \
= self.descrpt.build(coord_,
atype_,
natoms,
box,
mesh,
input_dict,
suffix = suffix,
reuse = reuse)
dout = tf.identity(dout, name='o_descriptor')
else:
tf.constant(self.rcut,
name = 'descrpt_attr/rcut',
dtype = GLOBAL_TF_FLOAT_PRECISION)
tf.constant(self.ntypes,
name = 'descrpt_attr/ntypes',
dtype = tf.int32)
feed_dict = self.descrpt.get_feed_dict(coord_, atype_, natoms, box, mesh)
return_elements = [*self.descrpt.get_tensor_names(), 'o_descriptor:0']
imported_tensors \
= self._import_graph_def_from_frz_model(frz_model, feed_dict, return_elements)
dout = imported_tensors[-1]
self.descrpt.pass_tensors_from_frz_model(*imported_tensors[:-1])
rot_mat = self.descrpt.get_rot_mat()
rot_mat = tf.identity(rot_mat, name = 'o_rot_mat'+suffix)
output = self.fitting.build (dout,
rot_mat,
natoms,
reuse = reuse,
suffix = suffix)
framesize = nout if "global" in self.model_type else natomsel * nout
output = tf.reshape(output, [-1, framesize], name = 'o_' + self.model_type + suffix)
model_dict = {self.model_type: output}
if "global" not in self.model_type:
gname = "global_"+self.model_type
atom_out = tf.reshape(output, [-1, natomsel, nout])
global_out = tf.reduce_sum(atom_out, axis=1)
global_out = tf.reshape(global_out, [-1, nout], name="o_" + gname + suffix)
out_cpnts = tf.split(atom_out, nout, axis=-1)
force_cpnts = []
virial_cpnts = []
atom_virial_cpnts = []
for out_i in out_cpnts:
force_i, virial_i, atom_virial_i \
= self.descrpt.prod_force_virial(out_i, natoms)
force_cpnts.append (tf.reshape(force_i, [-1, 3*natoms[1]]))
virial_cpnts.append (tf.reshape(virial_i, [-1, 9]))
atom_virial_cpnts.append(tf.reshape(atom_virial_i, [-1, 9*natoms[1]]))
# [nframe x nout x (natom x 3)]
force = tf.concat(force_cpnts, axis=1, name="o_force" + suffix)
# [nframe x nout x 9]
virial = tf.concat(virial_cpnts, axis=1, name="o_virial" + suffix)
# [nframe x nout x (natom x 9)]
atom_virial = tf.concat(atom_virial_cpnts, axis=1, name="o_atom_virial" + suffix)
model_dict[gname] = global_out
model_dict["force"] = force
model_dict["virial"] = virial
model_dict["atom_virial"] = atom_virial
return model_dict
def _import_graph_def_from_frz_model(self, frz_model, feed_dict, return_elements):
graph, graph_def = load_graph_def(frz_model)
return tf.import_graph_def(graph_def, input_map = feed_dict, return_elements = return_elements, name = "")
def init_variables(self,
graph : tf.Graph,
graph_def : tf.GraphDef,
model_type : str = "original_model",
suffix : str = "",
) -> None:
"""
Init the embedding net variables with the given frozen model
Parameters
----------
graph : tf.Graph
The input frozen model graph
graph_def : tf.GraphDef
The input frozen model graph_def
model_type : str
the type of the model
suffix : str
suffix to name scope
"""
if model_type == 'original_model':
self.descrpt.init_variables(graph, graph_def, suffix=suffix)
self.fitting.init_variables(graph, graph_def, suffix=suffix)
tf.constant("original_model", name = 'model_type', dtype = tf.string)
elif model_type == 'compressed_model':
self.fitting.init_variables(graph, graph_def, suffix=suffix)
tf.constant("compressed_model", name = 'model_type', dtype = tf.string)
else:
raise RuntimeError("Unknown model type %s" % model_type)
class WFCModel(TensorModel):
def __init__(
self,
descrpt,
fitting,
type_map : List[str] = None,
data_stat_nbatch : int = 10,
data_stat_protect : float = 1e-2
) -> None:
TensorModel.__init__(self, 'wfc', descrpt, fitting, type_map, data_stat_nbatch, data_stat_protect)
class DipoleModel(TensorModel):
def __init__(
self,
descrpt,
fitting,
type_map : List[str] = None,
data_stat_nbatch : int = 10,
data_stat_protect : float = 1e-2
) -> None:
TensorModel.__init__(self, 'dipole', descrpt, fitting, type_map, data_stat_nbatch, data_stat_protect)
class PolarModel(TensorModel):
def __init__(
self,
descrpt,
fitting,
type_map : List[str] = None,
data_stat_nbatch : int = 10,
data_stat_protect : float = 1e-2
) -> None:
TensorModel.__init__(self, 'polar', descrpt, fitting, type_map, data_stat_nbatch, data_stat_protect)
class GlobalPolarModel(TensorModel):
def __init__(
self,
descrpt,
fitting,
type_map : List[str] = None,
data_stat_nbatch : int = 10,
data_stat_protect : float = 1e-2
) -> None:
TensorModel.__init__(self, 'global_polar', descrpt, fitting, type_map, data_stat_nbatch, data_stat_protect)
from . import data, descriptor, entrypoints, fit, utils
__all__ = [
"data",
"descriptor",
"entrypoints",
"fit",
"utils",
]
"""
nvnmd.data
==========
Provides
1. hardware configuration
2. default input script
3. title and citation
Data
----
jdata_sys
action configuration
jdata_config
hardware configuration
dscp
descriptor configuration
fitn
fitting network configuration
size
ram capacity
ctrl
control flag, such as Time Division Multiplexing (TDM)
nbit
number of bits of fixed-point number
jdata_config_16 (disable)
difference with configure fitting size as 16
jdata_config_32 (disable)
difference with configure fitting size as 32
jdata_config_64 (disable)
difference with configure fitting size as 64
jdata_config_128 (default)
difference with configure fitting size as 128
jdata_configs
all configure of jdata_config{nfit_node}
jdata_deepmd_input
default input script for nvnmd training
NVNMD_WELCOME
nvnmd title when logging
NVNMD_CITATION
citation of nvnmd
"""
jdata_sys = {
"debug": False
}
jdata_config = {
"dscp": {
"sel": [60, 60],
"rcut": 6.0,
"rcut_smth": 0.5,
"neuron": [8, 16, 32],
"resnet_dt": False,
"axis_neuron": 4,
"type_one_side": True,
"NI": 128,
"rc_lim": 0.5,
"M1": "neuron[-1]",
"M2": "axis_neuron",
"SEL": [60, 60, 0, 0],
"NNODE_FEAS": "(1, neuron)",
"nlayer_fea": "len(neuron)",
"same_net": "type_one_side",
"NIDP": "sum(sel)",
"NIX": "2^ceil(ln2(NIDP/1.5))",
"ntype": "len(sel)",
"ntypex": "same_net ? 1: ntype",
"ntypex_max": 1,
"ntype_max": 4
},
"fitn": {
"neuron": [32, 32, 32],
"resnet_dt": False,
"NNODE_FITS": "(M1*M2, neuron, 1)",
"nlayer_fit": "len(neuron)+1",
"NLAYER": "nlayer_fit"
},
"size": {
"NTYPE_MAX": 4,
"NSPU": 4096,
"MSPU": 32768,
"Na": "NSPU",
"NaX": "MSPU"
},
"ctrl": {
"NSTDM": 16,
"NSTDM_M1": 16,
"NSTDM_M2": 1,
"NSADV": "NSTDM+1",
"NSEL": "NSTDM*ntype_max",
"NSTDM_M1X": 4,
"NSTEP_DELAY": 20,
"MAX_FANOUT": 30
},
"nbit": {
"NBIT_DATA": 21,
"NBIT_DATA_FL": 13,
"NBIT_LONG_DATA": 32,
"NBIT_LONG_DATA_FL": 24,
"NBIT_DIFF_DATA": 24,
"NBIT_SPE": 2,
"NBIT_CRD": "NBIT_DATA*3",
"NBIT_LST": "ln2(NaX)",
"NBIT_SPE_MAX": 8,
"NBIT_LST_MAX": 16,
"NBIT_ATOM": "NBIT_SPE+NBIT_CRD",
"NBIT_LONG_ATOM": "NBIT_SPE+NBIT_LONG_DATA*3",
"NBIT_RIJ": "NBIT_DATA_FL+5",
"NBIT_FEA_X": 10,
"NBIT_FEA_X_FL": 4,
"NBIT_FEA_X2_FL": 6,
"NBIT_FEA": 18,
"NBIT_FEA_FL": 10,
"NBIT_SHIFT": 4,
"NBIT_DATA2": "NBIT_DATA+NBIT_DATA_FL",
"NBIT_DATA2_FL": "2*NBIT_DATA_FL",
"NBIT_DATA_FEA": "NBIT_DATA+NBIT_FEA_FL",
"NBIT_DATA_FEA_FL": "NBIT_DATA_FL+NBIT_FEA_FL",
"NBIT_FORCE": 32,
"NBIT_FORCE_FL": "2*NBIT_DATA_FL-1",
"NBIT_SUM": "NBIT_DATA_FL+8",
"NBIT_WEIGHT": 18,
"NBIT_WEIGHT_FL": 13,
"NBIT_RAM": 72,
"NBIT_ADDR": 32,
"NBTI_MODEL_HEAD": 32,
"NBIT_TH_LONG_ADD": 30,
"NBIT_ADD": 15,
"RANGE_B": [-100, 100],
"RANGE_W": [-20, 20],
"NCFG": 35,
"NNET": 4920,
"NFEA": 8192
},
"end": ""
}
jdata_config_16 = {
"dscp": {
"neuron": [8, 16, 32],
"axis_neuron": 4,
"NI": 128
},
"fitn": {
"neuron": [16, 16, 16]
},
"ctrl": {
"NSTDM": 16,
"NSTDM_M1": 16,
"NSTDM_M2": 1,
"NSTDM_M1X": 4
}
}
jdata_config_32 = {
"dscp": {
"neuron": [8, 16, 32],
"axis_neuron": 4,
"NI": 128
},
"fitn": {
"neuron": [32, 32, 32]
},
"ctrl": {
"NSTDM": 16,
"NSTDM_M1": 16,
"NSTDM_M2": 1,
"NSTDM_M1X": 4
}
}
jdata_config_64 = {
"dscp": {
"neuron": [8, 16, 32],
"axis_neuron": 4,
"NI": 128
},
"fitn": {
"neuron": [64, 64, 64]
},
"ctrl": {
"NSTDM": 32,
"NSTDM_M1": 32,
"NSTDM_M2": 1,
"NSTDM_M1X": 4
}
}
jdata_config_128 = {
"dscp": {
"neuron": [8, 16, 32],
"axis_neuron": 4,
"NI": 128
},
"fitn": {
"neuron": [128, 128, 128]
},
"ctrl": {
"NSTDM": 32,
"NSTDM_M1": 32,
"NSTDM_M2": 1,
"NSTDM_M1X": 4
}
}
jdata_configs = {
"_16": jdata_config_16,
"_32": jdata_config_32,
"_64": jdata_config_64,
"128": jdata_config_128
}
jdata_deepmd_input = {
"model": {
"descriptor": {
"seed": 1,
"type": "se_a",
"sel": [
60,
60
],
"rcut": 7.0,
"rcut_smth": 0.5,
"neuron": [
8,
16,
32
],
"type_one_side": False,
"axis_neuron": 4,
"resnet_dt": False
},
"fitting_net": {
"seed": 1,
"neuron": [
128,
128,
128
],
"resnet_dt": False
}
},
"nvnmd": {
"net_size": 128,
"config_file": "none",
"weight_file": "none",
"map_file": "none",
"enable": False,
"restore_descriptor": False,
"restore_fitting_net": False,
"quantize_descriptor": False,
"quantize_fitting_net": False
},
"learning_rate": {
"type": "exp",
"decay_steps": 5000,
"start_lr": 0.005,
"stop_lr": 8.257687192506788e-05
},
"loss": {
"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
},
"training": {
"seed": 1,
"stop_batch": 10000,
"disp_file": "lcurve.out",
"disp_freq": 100,
"numb_test": 10,
"save_freq": 1000,
"save_ckpt": "model.ckpt",
"disp_training": True,
"time_training": True,
"profiling": False,
"training_data": {
"systems": "dataset",
"set_prefix": "set",
"batch_size": 1
}
}
}
NVNMD_WELCOME = (
" _ _ __ __ _ _ __ __ ____ ",
"| \ | | \ \ / / | \ | | | \/ | | _ \ ",
"| \| | \ \ / / | \| | | |\/| | | | | |",
"| |\ | \ V / | |\ | | | | | | |_| |",
"|_| \_| \_/ |_| \_| |_| |_| |____/ ",
)
NVNMD_CITATION = (
"Please read and cite:",
"Mo et al., npj Comput Mater 8, 107 (2022)",
)
"""
nvnmd.se_a
==========
Provides
1. building descriptor with continuous embedding network
2. building descriptor with quantized embedding network
"""
import numpy as np
from deepmd.env import tf
from deepmd.env import GLOBAL_TF_FLOAT_PRECISION
from deepmd.env import GLOBAL_NP_FLOAT_PRECISION
from deepmd.env import op_module
from deepmd.utils.network import embedding_net
#
from deepmd.nvnmd.utils.config import nvnmd_cfg
from deepmd.nvnmd.utils.network import matmul3_qq
from deepmd.nvnmd.utils.weight import get_normalize, get_rng_s
def build_davg_dstd():
r"""Get the davg and dstd from the dictionary nvnmd_cfg.
The davg and dstd have been obtained by training CNN
"""
davg, dstd = get_normalize(nvnmd_cfg.weight)
return davg, dstd
def build_op_descriptor():
r"""Replace se_a.py/DescrptSeA/build
"""
if nvnmd_cfg.quantize_descriptor:
return op_module.prod_env_mat_a_nvnmd_quantize
else:
return op_module.prod_env_mat_a
def descrpt2r4(inputs, natoms):
r"""Replace :math:`r_{ji} \rightarrow r'_{ji}`
where :math:`r_{ji} = (x_{ji}, y_{ji}, z_{ji})` and
:math:`r'_{ji} = (s_{ji}, \frac{s_{ji} x_{ji}}{r_{ji}}, \frac{s_{ji} y_{ji}}{r_{ji}}, \frac{s_{ji} z_{ji}}{r_{ji}})`
"""
NBIT_DATA_FL = nvnmd_cfg.nbit['NBIT_DATA_FL']
NBIT_FEA_X_FL = nvnmd_cfg.nbit['NBIT_FEA_X_FL']
NBIT_FEA_FL = nvnmd_cfg.nbit['NBIT_FEA_FL']
prec = 1.0 / (2 ** NBIT_FEA_X_FL)
ntypes = nvnmd_cfg.dscp['ntype']
NIDP = nvnmd_cfg.dscp['NIDP']
ndescrpt = NIDP * 4
start_index = 0
# (nf, na*nd)
shape = inputs.get_shape().as_list()
# (nf*na*ni, 4)
inputs_reshape = tf.reshape(inputs, [-1, 4])
with tf.variable_scope('filter_type_all_x', reuse=True):
# u (i.e., r^2)
u = tf.reshape(tf.slice(inputs_reshape, [0, 0], [-1, 1]), [-1, 1])
with tf.variable_scope('u', reuse=True):
u = op_module.quantize_nvnmd(u, 0, -1, NBIT_DATA_FL, -1)
# print('u:', u)
u = tf.reshape(u, [-1, natoms[0] * NIDP])
# rij
rij = tf.reshape(tf.slice(inputs_reshape, [0, 1], [-1, 3]), [-1, 3])
with tf.variable_scope('rij', reuse=True):
rij = op_module.quantize_nvnmd(rij, 0, NBIT_DATA_FL, -1, -1)
# print('rij:', rij)
s = []
sr = []
for type_i in range(ntypes):
type_input = 0
postfix = f"_t{type_input}_t{type_i}"
u_i = tf.slice(
u,
[0, start_index * NIDP],
[-1, natoms[2 + type_i] * NIDP])
u_i = tf.reshape(u_i, [-1, 1])
#
keys = 's,sr'.split(',')
map_tables = [nvnmd_cfg.map[key + postfix] for key in keys]
map_tables2 = [nvnmd_cfg.map[f"d{key}_dr2" + postfix] for key in keys]
map_outs = []
for ii in range(len(keys)):
map_outs.append(op_module.map_nvnmd(
u_i,
map_tables[ii][0],
map_tables[ii][1] / prec,
map_tables2[ii][0],
map_tables2[ii][1] / prec,
prec, NBIT_FEA_FL))
s_i, sr_i = map_outs
s_i = tf.reshape(s_i, [-1, natoms[2 + type_i] * NIDP])
sr_i = tf.reshape(sr_i, [-1, natoms[2 + type_i] * NIDP])
s.append(s_i)
sr.append(sr_i)
start_index += natoms[2 + type_i]
s = tf.concat(s, axis=1)
sr = tf.concat(sr, axis=1)
with tf.variable_scope('s', reuse=True):
s = op_module.quantize_nvnmd(s, 0, NBIT_FEA_FL, NBIT_DATA_FL, -1)
with tf.variable_scope('sr', reuse=True):
sr = op_module.quantize_nvnmd(sr, 0, NBIT_FEA_FL, NBIT_DATA_FL, -1)
s = tf.reshape(s, [-1, 1])
sr = tf.reshape(sr, [-1, 1])
# R2R4
Rs = s
Rxyz = sr * rij
with tf.variable_scope('Rxyz', reuse=True):
Rxyz = op_module.quantize_nvnmd(Rxyz, 0, NBIT_DATA_FL, NBIT_DATA_FL, -1)
R4 = tf.concat([Rs, Rxyz], axis=1)
R4 = tf.reshape(R4, [-1, NIDP, 4])
inputs_reshape = R4
inputs_reshape = tf.reshape(inputs_reshape, [-1, ndescrpt])
return inputs_reshape
def filter_lower_R42GR(
type_i,
type_input,
inputs_i,
is_exclude,
activation_fn,
bavg,
stddev,
trainable,
suffix,
seed,
seed_shift,
uniform_seed,
filter_neuron,
filter_precision,
filter_resnet_dt,
embedding_net_variables):
r"""Replace se_a.py/DescrptSeA/_filter_lower
"""
shape_i = inputs_i.get_shape().as_list()
inputs_reshape = tf.reshape(inputs_i, [-1, 4])
natom = tf.shape(inputs_i)[0]
M1 = nvnmd_cfg.dscp['M1']
NBIT_DATA_FL = nvnmd_cfg.nbit['NBIT_DATA_FL']
NBIT_FEA_X_FL = nvnmd_cfg.nbit['NBIT_FEA_X_FL']
NBIT_FEA_X2_FL = nvnmd_cfg.nbit['NBIT_FEA_X2_FL']
NBIT_FEA_FL = nvnmd_cfg.nbit['NBIT_FEA_FL']
prec = 1.0 / (2 ** NBIT_FEA_X2_FL)
type_input = 0 if (type_input < 0) else type_input
postfix = f"_t{type_input}_t{type_i}"
if (nvnmd_cfg.quantize_descriptor):
s_min, smax = get_rng_s(nvnmd_cfg.weight)
s_min = -2.0
# s_min = np.floor(s_min)
s = tf.reshape(tf.slice(inputs_reshape, [0, 0], [-1, 1]), [-1, 1])
s = op_module.quantize_nvnmd(s, 0, NBIT_FEA_FL, NBIT_DATA_FL, -1)
# G
keys = 'G'.split(',')
map_tables = [nvnmd_cfg.map[key + postfix] for key in keys]
map_tables2 = [nvnmd_cfg.map[f"d{key}_ds" + postfix] for key in keys]
map_outs = []
for ii in range(len(keys)):
with tf.variable_scope(keys[ii], reuse=True):
map_outs.append(op_module.map_nvnmd(
s - s_min,
map_tables[ii][0], map_tables[ii][1] / prec,
map_tables2[ii][0], map_tables2[ii][1] / prec,
prec, NBIT_FEA_FL))
map_outs[ii] = op_module.quantize_nvnmd(map_outs[ii], 0, NBIT_FEA_FL, NBIT_DATA_FL, -1)
G = map_outs
# G
xyz_scatter = G
xyz_scatter = tf.reshape(xyz_scatter, (-1, shape_i[1] // 4, M1))
# GR
inputs_reshape = tf.reshape(inputs_reshape, [-1, shape_i[1] // 4, 4])
GR = matmul3_qq(tf.transpose(inputs_reshape, [0, 2, 1]), xyz_scatter, -1)
GR = tf.reshape(GR, [-1, 4 * M1])
return GR
else:
xyz_scatter = tf.reshape(tf.slice(inputs_reshape, [0, 0], [-1, 1]), [-1, 1])
if nvnmd_cfg.restore_descriptor:
trainable = False
embedding_net_variables = {}
for key in nvnmd_cfg.weight.keys():
if 'filter_type' in key:
key2 = key.replace('.', '/')
embedding_net_variables[key2] = nvnmd_cfg.weight[key]
if (not is_exclude):
xyz_scatter = embedding_net(
xyz_scatter,
filter_neuron,
filter_precision,
activation_fn=activation_fn,
resnet_dt=filter_resnet_dt,
name_suffix=suffix,
stddev=stddev,
bavg=bavg,
seed=seed,
trainable=trainable,
uniform_seed=uniform_seed,
initial_variables=embedding_net_variables)
if (not uniform_seed) and (seed is not None):
seed += seed_shift
else:
# we can safely return the final xyz_scatter filled with zero directly
return tf.cast(tf.fill((natom, 4, M1), 0.), GLOBAL_TF_FLOAT_PRECISION)
# natom x nei_type_i x out_size
xyz_scatter = tf.reshape(xyz_scatter, (-1, shape_i[1] // 4, M1))
# When using tf.reshape(inputs_i, [-1, shape_i[1]//4, 4]) below
# [588 24] -> [588 6 4] correct
# but if sel is zero
# [588 0] -> [147 0 4] incorrect; the correct one is [588 0 4]
# So we need to explicitly assign the shape to tf.shape(inputs_i)[0] instead of -1
return tf.matmul(tf.reshape(inputs_i, [natom, shape_i[1] // 4, 4]), xyz_scatter, transpose_a=True)
def filter_GR2D(xyz_scatter_1):
r"""Replace se_a.py/_filter
"""
NIX = nvnmd_cfg.dscp['NIX']
NBIT_DATA_FL = nvnmd_cfg.nbit['NBIT_DATA_FL']
M1 = nvnmd_cfg.dscp['M1']
M2 = nvnmd_cfg.dscp['M2']
if (nvnmd_cfg.quantize_descriptor):
xyz_scatter_1 = tf.reshape(xyz_scatter_1, [-1, 4 * M1])
# fix the number of bits of gradient
xyz_scatter_1 = op_module.quantize_nvnmd(xyz_scatter_1, 0, -1, NBIT_DATA_FL, -1)
xyz_scatter_1 = xyz_scatter_1 * (1.0 / NIX)
with tf.variable_scope('GR', reuse=True):
xyz_scatter_1 = op_module.quantize_nvnmd(xyz_scatter_1, 0, NBIT_DATA_FL, NBIT_DATA_FL, -1)
xyz_scatter_1 = tf.reshape(xyz_scatter_1, [-1, 4, M1])
# natom x 4 x outputs_size_2
xyz_scatter_2 = xyz_scatter_1
# natom x 3 x outputs_size_1
qmat = tf.slice(xyz_scatter_1, [0, 1, 0], [-1, 3, -1])
# natom x outputs_size_2 x 3
qmat = tf.transpose(qmat, perm=[0, 2, 1])
# D': natom x outputs_size x outputs_size_2
result = tf.matmul(xyz_scatter_1, xyz_scatter_2, transpose_a=True)
# D': natom x (outputs_size x outputs_size_2)
result = tf.reshape(result, [-1, M1 * M1])
#
index_subset = []
for ii in range(M1):
for jj in range(ii, ii + M2):
index_subset.append((ii * M1) + (jj % M1))
index_subset = tf.constant(np.int32(np.array(index_subset)))
result = tf.gather(result, index_subset, axis=1)
with tf.variable_scope('d', reuse=True):
result = op_module.quantize_nvnmd(result, 0, NBIT_DATA_FL, NBIT_DATA_FL, -1)
else:
# natom x 4 x outputs_size
xyz_scatter_1 = xyz_scatter_1 * (1.0 / NIX)
# natom x 4 x outputs_size_2
# xyz_scatter_2 = tf.slice(xyz_scatter_1, [0,0,0],[-1,-1,outputs_size_2])
xyz_scatter_2 = xyz_scatter_1
# natom x 3 x outputs_size_1
qmat = tf.slice(xyz_scatter_1, [0, 1, 0], [-1, 3, -1])
# natom x outputs_size_1 x 3
qmat = tf.transpose(qmat, perm=[0, 2, 1])
# natom x outputs_size x outputs_size_2
result = tf.matmul(xyz_scatter_1, xyz_scatter_2, transpose_a=True)
# natom x (outputs_size x outputs_size_2)
# result = tf.reshape(result, [-1, outputs_size_2 * outputs_size[-1]])
result = tf.reshape(result, [-1, M1 * M1])
#
index_subset = []
for ii in range(M1):
for jj in range(ii, ii + M2):
index_subset.append((ii * M1) + (jj % M1))
index_subset = tf.constant(np.int32(np.array(index_subset)))
result = tf.gather(result, index_subset, axis=1)
return result, qmat
from .freeze import save_weight
from .mapt import MapTable
from .wrap import Wrap
__all__ = [
"save_weight",
"MapTable",
"Wrap"
]
#!/usr/bin/env python3
from deepmd.env import tf
from deepmd.nvnmd.utils.fio import FioDic
def filter_tensorVariableList(tensorVariableList) -> dict:
r"""Get the name of variable for NVNMD
| :code:`descrpt_attr/t_avg:0`
| :code:`descrpt_attr/t_std:0`
| :code:`filter_type_{atom i}/matrix_{layer l}_{atomj}:0`
| :code:`filter_type_{atom i}/bias_{layer l}_{atomj}:0`
| :code:`layer_{layer l}_type_{atom i}/matrix:0`
| :code:`layer_{layer l}_type_{atom i}/bias:0`
| :code:`final_layer_type_{atom i}/matrix:0`
| :code:`final_layer_type_{atom i}/bias:0`
"""
nameList = [tv.name for tv in tensorVariableList]
nameList = [name.replace(':0', '') for name in nameList]
nameList = [name.replace('/', '.') for name in nameList]
dic_name_tv = {}
for ii in range(len(nameList)):
name = nameList[ii]
tv = tensorVariableList[ii]
p1 = name.startswith('descrpt_attr')
p1 = p1 or name.startswith('filter_type_')
p1 = p1 or name.startswith('layer_')
p1 = p1 or name.startswith('final_layer_type_')
p2 = 'Adam' not in name
p3 = 'XXX' not in name
if p1 and p2 and p3:
dic_name_tv[name] = tv
return dic_name_tv
def save_weight(sess, file_name: str = 'nvnmd/weight.npy'):
r"""Save the dictionary of weight to a npy file
"""
tvs = tf.global_variables()
dic_key_tv = filter_tensorVariableList(tvs)
dic_key_value = {}
for key in dic_key_tv.keys():
value = sess.run(dic_key_tv[key])
dic_key_value[key] = value
FioDic().save(file_name, dic_key_value)
import numpy as np
import logging
from deepmd.env import tf
from deepmd.utils.sess import run_sess
from deepmd.nvnmd.utils.fio import FioDic
from deepmd.nvnmd.utils.config import nvnmd_cfg
from deepmd.nvnmd.utils.weight import get_normalize, get_rng_s, get_filter_weight
from deepmd.nvnmd.utils.network import get_sess
from deepmd.nvnmd.data.data import jdata_deepmd_input
from typing import List, Optional
log = logging.getLogger(__name__)
class MapTable:
r"""Generate the mapping table describing the relastionship of
atomic distance, cutoff function, and embedding matrix.
three mapping table will be built:
| :math:`r^2_{ji} \rightarrow s_{ji}`
| :math:`r^2_{ji} \rightarrow sr_{ji}`
| :math:`r^2_{ji} \rightarrow \mathcal{G}_{ji}`
where :math:`s_{ji}` is cut-off function,
:math:`sr_{ji} = \frac{s(r_{ji})}{r_{ji}}`, and
:math:`\mathcal{G}_{ji}` is embedding matrix.
The mapping funciton can be define as:
| :math:`y = f(x) = y_{k} + (x - x_{k}) * dy_{k}`
| :math:`y_{k} = f(x_{k})`
| :math:`dy_{k} = \frac{f(x_{k+1}) - f(x_{k})}{dx}`
| :math:`x_{k} \leq x < x_{k+1}`
| :math:`x_{k} = k * dx`
where :math:`dx` is interpolation interval.
Parameters
----------
config_file
input file name
an .npy file containing the configuration information of NVNMD model
weight_file
input file name
an .npy file containing the weights of NVNMD model
map_file
output file name
an .npy file containing the mapping tables of NVNMD model
References
----------
DOI: 10.1038/s41524-022-00773-z
"""
def __init__(
self,
config_file: str,
weight_file: str,
map_file: str
):
self.config_file = config_file
self.weight_file = weight_file
self.map_file = map_file
jdata = jdata_deepmd_input['nvnmd']
jdata['config_file'] = config_file
jdata['weight_file'] = weight_file
jdata['enable'] = True
nvnmd_cfg.init_from_jdata(jdata)
# map_table = self.build_map()
def qqq(self, dat, NBIT_FEA_FL, NBIT_FEA_X, is_set_zero=False):
dat = dat if isinstance(dat, list) else [dat]
prec = 2 ** NBIT_FEA_FL
N = int(2 ** NBIT_FEA_X)
#
dat2 = []
for ii in range(len(dat)):
dati = dat[ii]
vi = dati[:-1] # i
vi1 = dati[1:] # i+1
# v = vi + dvi * (r - ri)
# ri = i * dt
# dvi = v(i+1) / dt
vi = np.round(vi * prec) / prec
vi1 = np.round(vi1 * prec) / prec
dvi = vi1 - vi
if is_set_zero:
dvi[0] = 0
#
v = [np.reshape(vp, [N, -1]) for vp in [vi, dvi]]
dat2.append(v)
return dat2
def build_map(self):
ntypex = nvnmd_cfg.dscp['ntypex']
ntype = nvnmd_cfg.dscp['ntype']
NBIT_FEA_FL = nvnmd_cfg.nbit['NBIT_FEA_FL']
NBIT_FEA_X = nvnmd_cfg.nbit['NBIT_FEA_X']
dic = self.run_u2s()
dic.update(self.run_s2G(dic))
# quantize s and G
prec = 2**NBIT_FEA_FL
for tt in range(ntypex):
dic['s'][tt][0] = np.round(dic['s'][tt][0] * prec) / prec
dic['sr'][tt][0] = np.round(dic['sr'][tt][0] * prec) / prec
for tt2 in range(ntype):
v = np.round(dic['G'][tt * ntype + tt2][0] * prec) / prec
dic['G'][tt * ntype + tt2][0] = v
maps = {}
keys = 's,sr,ds_dr2,dsr_dr2,G,dG_ds'.split(',')
keys2 = 'G,dG_ds'.split(',')
for key in keys:
val = self.qqq(dic[key], NBIT_FEA_FL, NBIT_FEA_X, key not in keys2)
maps[key] = val
N = int(2**NBIT_FEA_X)
maps2 = {}
maps2['r2'] = dic['r2'][0:N]
maps2['s2'] = dic['s2'][0:N]
for tt in range(ntypex):
for tt2 in range(ntype):
postfix = f'_t{tt}_t{tt2}'
for key in keys:
maps2[key + postfix] = []
maps2[key + postfix].append(maps[key][tt * ntype + tt2][0].reshape([N, -1]))
maps2[key + postfix].append(maps[key][tt * ntype + tt2][1].reshape([N, -1]))
self.map = maps2
FioDic().save(self.map_file, self.map)
log.info("NVNMD: finish building mapping table")
return self.map
# =====================================================================
# build r2s
# =====================================================================
def build_r2s(self, r2):
# limit = nvnmd_cfg.dscp['rc_lim']
rmin = nvnmd_cfg.dscp['rcut_smth']
rmax = nvnmd_cfg.dscp['rcut']
# ntypex = nvnmd_cfg.dscp['ntypex']
ntype = nvnmd_cfg.dscp['ntype']
avg, std = get_normalize(nvnmd_cfg.weight)
avg, std = np.float32(avg), np.float32(std)
r = tf.sqrt(r2)
r_ = tf.clip_by_value(r, rmin, rmax)
r__ = tf.clip_by_value(r, 0, rmax)
uu = (r_ - rmin) / (rmax - rmin)
vv = uu * uu * uu * (-6 * uu * uu + 15 * uu - 10) + 1
sl = []
srl = []
for tt in range(ntype):
s = vv / r__
sr = s / r__
s = tf.reshape(s, [-1, 1])
sr = tf.reshape(sr, [-1, 1])
s = (s - avg[tt, 0]) / std[tt, 0]
sr = sr / std[tt, 1]
sl.append(s)
srl.append(sr)
return sl, srl
def build_ds_dr(self, r2, s, sr):
# ntypex = nvnmd_cfg.dscp['ntypex']
ntype = nvnmd_cfg.dscp['ntype']
ds_drl = []
dsr_drl = []
for tt in range(ntype):
si = s[tt]
sri = sr[tt]
ds_dr = tf.gradients(si, r2)
dsr_dr = tf.gradients(sri, r2)
ds_drl.append(ds_dr[0])
dsr_drl.append(dsr_dr[0])
return ds_drl, dsr_drl
def build_r2s_r2ds(self):
dic_ph = {}
dic_ph['r2'] = tf.placeholder(tf.float32, [None, 1], 't_r2')
dic_ph['s'], dic_ph['sr'] = self.build_r2s(dic_ph['r2'])
dic_ph['ds_dr2'], dic_ph['dsr_dr2'] = self.build_ds_dr(dic_ph['r2'], dic_ph['s'], dic_ph['sr'])
return dic_ph
def run_u2s(self):
# ntypex = nvnmd_cfg.dscp['ntypex']
ntype = nvnmd_cfg.dscp['ntype']
avg, std = get_normalize(nvnmd_cfg.weight)
avg, std = np.float32(avg), np.float32(std)
NBIT_FEA_X = nvnmd_cfg.nbit['NBIT_FEA_X']
NBIT_FEA_X_FL = nvnmd_cfg.nbit['NBIT_FEA_X_FL']
dic_ph = self.build_r2s_r2ds()
sess = get_sess()
N = 2 ** NBIT_FEA_X
N2 = 2 ** NBIT_FEA_X_FL
# N+1 ranther than N for calculating defference
r2 = 1.0 * np.arange(0, N + 1) / N2
r2 = np.reshape(r2, [-1, 1])
feed_dic = {dic_ph['r2']: r2}
key = 'r2,s,sr,ds_dr2,dsr_dr2'
tlst = [dic_ph[k] for k in key.split(',')]
# res = sess.run(tlst, feed_dic)
res = run_sess(sess, tlst, feed_dict=feed_dic)
res2 = {}
key = key.split(',')
for ii in range(len(key)):
res2[key[ii]] = res[ii]
# change value
# set 0 value, when u=0
for tt in range(ntype):
res2['s'][tt][0] = -avg[tt, 0] / std[tt, 0]
res2['sr'][tt][0] = 0
res2['ds_dr2'][tt][0] = 0
res2['dsr_dr2'][tt][0] = 0
# r = np.sqrt(res2['r2'])
sess.close()
return res2
# =====================================================================
# build s2G
# =====================================================================
def build_s2G(self, s):
ntypex = nvnmd_cfg.dscp['ntypex']
ntype = nvnmd_cfg.dscp['ntype']
activation_fn = tf.tanh
outputs_size = nvnmd_cfg.dscp['NNODE_FEAS']
xyz_scatters = []
for tt in range(ntypex):
for tt2 in range(ntype):
xyz_scatter = s
for ll in range(1, len(outputs_size)):
w, b = get_filter_weight(nvnmd_cfg.weight, tt, tt2, ll)
w, b = np.float32(w), np.float32(b)
if outputs_size[ll] == outputs_size[ll - 1]:
xyz_scatter += activation_fn(tf.matmul(xyz_scatter, w) + b)
elif outputs_size[ll] == outputs_size[ll - 1] * 2:
xyz_scatter = tf.concat([xyz_scatter, xyz_scatter], 1) + activation_fn(tf.matmul(xyz_scatter, w) + b)
else:
xyz_scatter = activation_fn(tf.matmul(xyz_scatter, w) + b)
xyz_scatters.append(xyz_scatter)
return xyz_scatters
def build_dG_ds(self, G, s):
ntypex = nvnmd_cfg.dscp['ntypex']
ntype = nvnmd_cfg.dscp['ntype']
M1 = nvnmd_cfg.dscp['M1']
dG_ds = []
for tt in range(ntypex):
for tt2 in range(ntype):
Gi = G[tt * ntype + tt2]
si = s
dG_ds_i = []
for ii in range(M1):
dG_ds_ii = tf.reshape(tf.gradients(Gi[:, ii], si), [-1, 1])
dG_ds_i.append(dG_ds_ii)
dG_ds_i = tf.concat(dG_ds_i, axis=1)
dG_ds.append(dG_ds_i)
return dG_ds
def build_s2G_s2dG(self):
# ntypex = nvnmd_cfg.dscp['ntypex']
dic_ph = {}
dic_ph['s2'] = tf.placeholder(tf.float32, [None, 1], 't_s')
dic_ph['G'] = self.build_s2G(dic_ph['s2'])
dic_ph['dG_ds'] = self.build_dG_ds(dic_ph['G'], dic_ph['s2'])
return dic_ph
def run_s2G(self, dat):
NBIT_FEA_FL = nvnmd_cfg.nbit['NBIT_FEA_FL']
NBIT_FEA_X = nvnmd_cfg.nbit['NBIT_FEA_X']
NBIT_FEA_X2_FL = nvnmd_cfg.nbit['NBIT_FEA_X2_FL']
prec = 2 ** NBIT_FEA_FL
dic_ph = self.build_s2G_s2dG()
sess = get_sess()
N = 2 ** NBIT_FEA_X
N2 = 2 ** NBIT_FEA_X2_FL
s_min, s_max = get_rng_s(nvnmd_cfg.weight)
#
if (s_min < -2.0) or (s_max > 14.0):
log.warning(f"the range of s [{s_min}, {s_max}] is over the limit [-2.0, 14.0]")
s_min = -2.0
s = s_min + np.arange(0, N + 1) / N2
s = np.reshape(s, [-1, 1])
feed_dic = {dic_ph['s2']: s}
feed_dic = {dic_ph['s2']: s}
key = 's2,G,dG_ds'
tlst = [dic_ph[k] for k in key.split(',')]
# res = sess.run(tlst, feed_dic)
res = run_sess(sess, tlst, feed_dict=feed_dic)
res2 = {}
key = key.split(',')
for ii in range(len(key)):
res2[key[ii]] = res[ii]
sess.close()
return res2
def mapt(
*,
nvnmd_config: Optional[str] = 'nvnmd/config.npy',
nvnmd_weight: Optional[str] = 'nvnmd/weight.npy',
nvnmd_map: Optional[str] = 'nvnmd/map.npy',
**kwargs
):
# build mapping table
mapObj = MapTable(nvnmd_config, nvnmd_weight, nvnmd_map)
mapObj.build_map()
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