Skip to content

GitLab

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

BladeDISC DeePMD code update

parent 6b33aeb8
Pipeline #180 failed with stages
in 0 seconds
Hide whitespace changes
Inline Side-by-side
Showing with 0 additions and 5052 deletions
_skbuild/linux-x86_64-3.6/cmake-install/deepmd/entrypoints/neighbor_stat.py deleted 100644 → 0
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/cmake-install/deepmd/entrypoints/test.py deleted 100644 → 0
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/cmake-install/deepmd/entrypoints/train.py deleted 100644 → 0
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/cmake-install/deepmd/entrypoints/transfer.py deleted 100644 → 0
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/cmake-install/deepmd/env.py deleted 100644 → 0
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/cmake-install/deepmd/fit/__init__.py deleted 100644 → 0
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/cmake-install/deepmd/fit/dipole.py deleted 100644 → 0
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/cmake-install/deepmd/fit/ener.py deleted 100644 → 0
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/cmake-install/deepmd/fit/fitting.py deleted 100644 → 0
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/cmake-install/deepmd/fit/polar.py deleted 100644 → 0
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/cmake-install/deepmd/fit/wfc.py deleted 100644 → 0
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/cmake-install/deepmd/infer/__init__.py deleted 100644 → 0
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/cmake-install/deepmd/infer/data_modifier.py deleted 100644 → 0
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/cmake-install/deepmd/infer/deep_dipole.py deleted 100644 → 0
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")
_skbuild/linux-x86_64-3.6/cmake-install/deepmd/infer/deep_eval.py deleted 100644 → 0
View file @ 6b33aeb8
import os
from typing import List, Optional, TYPE_CHECKING, Union
from functools import lru_cache
import numpy as np
from deepmd.common import make_default_mesh
from deepmd.env import default_tf_session_config, tf, MODEL_VERSION
from deepmd.utils.sess import run_sess
from deepmd.utils.batch_size import AutoBatchSize
if TYPE_CHECKING:
from pathlib import Path
class DeepEval:
"""Common methods for DeepPot, DeepWFC, DeepPolar, ...
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
auto_batch_size : bool or int or AutomaticBatchSize, default: False
If True, automatic batch size will be used. If int, it will be used
as the initial batch size.
"""
load_prefix: str # set by subclass
def __init__(
self,
model_file: "Path",
load_prefix: str = "load",
default_tf_graph: bool = False,
auto_batch_size: Union[bool, int, AutoBatchSize] = False,
):
self.graph = self._load_graph(
model_file, prefix=load_prefix, default_tf_graph=default_tf_graph
)
self.load_prefix = load_prefix
# graph_compatable should be called after graph and prefix are set
if not self._graph_compatable():
raise RuntimeError(
f"model in graph (version {self.model_version}) is incompatible"
f"with the model (version {MODEL_VERSION}) supported by the current code."
)
# set default to False, as subclasses may not support
if isinstance(auto_batch_size, bool):
if auto_batch_size:
self.auto_batch_size = AutoBatchSize()
else:
self.auto_batch_size = None
elif isinstance(auto_batch_size, int):
self.auto_batch_size = AutoBatchSize(auto_batch_size)
elif isinstance(auto_batch_size, AutoBatchSize):
self.auto_batch_size = auto_batch_size
else:
raise TypeError("auto_batch_size should be bool, int, or AutoBatchSize")
@property
@lru_cache(maxsize=None)
def model_type(self) -> str:
"""Get type of model.
:type:str
"""
t_mt = self._get_tensor("model_attr/model_type:0")
[mt] = run_sess(self.sess, [t_mt], feed_dict={})
return mt.decode("utf-8")
@property
@lru_cache(maxsize=None)
def model_version(self) -> str:
"""Get version of model.
Returns
-------
str
version of model
"""
try:
t_mt = self._get_tensor("model_attr/model_version:0")
except KeyError:
# For deepmd-kit version 0.x - 1.x, set model version to 0.0
return "0.0"
else:
[mt] = run_sess(self.sess, [t_mt], feed_dict={})
return mt.decode("utf-8")
@property
@lru_cache(maxsize=None)
def sess(self) -> tf.Session:
"""Get TF session."""
# start a tf session associated to the graph
return tf.Session(graph=self.graph, config=default_tf_session_config)
def _graph_compatable(
self
) -> bool :
""" Check the model compatability
Returns
-------
bool
If the model stored in the graph file is compatable with the current code
"""
model_version_major = int(self.model_version.split('.')[0])
model_version_minor = int(self.model_version.split('.')[1])
MODEL_VERSION_MAJOR = int(MODEL_VERSION.split('.')[0])
MODEL_VERSION_MINOR = int(MODEL_VERSION.split('.')[1])
if (model_version_major != MODEL_VERSION_MAJOR) or \
(model_version_minor > MODEL_VERSION_MINOR) :
return False
else:
return True
def _get_tensor(
self, tensor_name: str, attr_name: Optional[str] = None
) -> tf.Tensor:
"""Get TF graph tensor and assign it to class namespace.
Parameters
----------
tensor_name : str
name of tensor to get
attr_name : Optional[str], optional
if specified, class attribute with this name will be created and tensor will
be assigned to it, by default None
Returns
-------
tf.Tensor
loaded tensor
"""
tensor_path = os.path.join(self.load_prefix, tensor_name)
tensor = self.graph.get_tensor_by_name(tensor_path)
if attr_name:
setattr(self, attr_name, tensor)
return tensor
else:
return tensor
@staticmethod
def _load_graph(
frozen_graph_filename: "Path", prefix: str = "load", default_tf_graph: bool = False
):
# We load the protobuf file from the disk and parse it to retrieve the
# unserialized graph_def
with tf.gfile.GFile(str(frozen_graph_filename), "rb") as f:
graph_def = tf.GraphDef()
graph_def.ParseFromString(f.read())
if default_tf_graph:
tf.import_graph_def(
graph_def,
input_map=None,
return_elements=None,
name=prefix,
producer_op_list=None
)
graph = tf.get_default_graph()
else :
# Then, we can use again a convenient built-in function to import
# a graph_def into the current default Graph
with tf.Graph().as_default() as graph:
tf.import_graph_def(
graph_def,
input_map=None,
return_elements=None,
name=prefix,
producer_op_list=None
)
return graph
@staticmethod
def sort_input(
coord : np.ndarray, atom_type : np.ndarray, sel_atoms : List[int] = None
):
"""
Sort atoms in the system according their types.
Parameters
----------
coord
The coordinates of atoms.
Should be of shape [nframes, natoms, 3]
atom_type
The type of atoms
Should be of shape [natoms]
sel_atom
The selected atoms by type
Returns
-------
coord_out
The coordinates after sorting
atom_type_out
The atom types after sorting
idx_map
The index mapping from the input to the output.
For example coord_out = coord[:,idx_map,:]
sel_atom_type
Only output if sel_atoms is not None
The sorted selected atom types
sel_idx_map
Only output if sel_atoms is not None
The index mapping from the selected atoms to sorted selected atoms.
"""
if sel_atoms is not None:
selection = [False] * np.size(atom_type)
for ii in sel_atoms:
selection += (atom_type == ii)
sel_atom_type = atom_type[selection]
natoms = atom_type.size
idx = np.arange (natoms)
idx_map = np.lexsort ((idx, atom_type))
nframes = coord.shape[0]
coord = coord.reshape([nframes, -1, 3])
coord = np.reshape(coord[:,idx_map,:], [nframes, -1])
atom_type = atom_type[idx_map]
if sel_atoms is not None:
sel_natoms = np.size(sel_atom_type)
sel_idx = np.arange(sel_natoms)
sel_idx_map = np.lexsort((sel_idx, sel_atom_type))
sel_atom_type = sel_atom_type[sel_idx_map]
return coord, atom_type, idx_map, sel_atom_type, sel_idx_map
else:
return coord, atom_type, idx_map
@staticmethod
def reverse_map(vec : np.ndarray, imap : List[int]) -> np.ndarray:
"""Reverse mapping of a vector according to the index map
Parameters
----------
vec
Input vector. Be of shape [nframes, natoms, -1]
imap
Index map. Be of shape [natoms]
Returns
-------
vec_out
Reverse mapped vector.
"""
ret = np.zeros(vec.shape)
# for idx,ii in enumerate(imap) :
# ret[:,ii,:] = vec[:,idx,:]
ret[:, imap, :] = vec
return ret
def make_natoms_vec(self, atom_types : np.ndarray) -> np.ndarray :
"""Make the natom vector used by deepmd-kit.
Parameters
----------
atom_types
The type of atoms
Returns
-------
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
"""
natoms_vec = np.zeros (self.ntypes+2).astype(int)
natoms = atom_types.size
natoms_vec[0] = natoms
natoms_vec[1] = natoms
for ii in range (self.ntypes) :
natoms_vec[ii+2] = np.count_nonzero(atom_types == ii)
return natoms_vec
_skbuild/linux-x86_64-3.6/cmake-install/deepmd/infer/deep_polar.py deleted 100644 → 0
View file @ 6b33aeb8
from deepmd.infer.deep_tensor import DeepTensor
import numpy as np
from typing import TYPE_CHECKING, List, Optional
if TYPE_CHECKING:
from pathlib import Path
class DeepPolar(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_polar: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")
class DeepGlobalPolar(DeepTensor):
"""Constructor.
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
"""
def __init__(
self, model_file: str, load_prefix: str = "load", default_tf_graph: bool = False
) -> None:
self.tensors.update(
{
"t_sel_type": "model_attr/sel_type:0",
# output tensor
"t_tensor": "o_global_polar:0",
}
)
DeepTensor.__init__(
self,
model_file,
load_prefix=load_prefix,
default_tf_graph=default_tf_graph,
)
def eval(
self,
coords: np.ndarray,
cells: np.ndarray,
atom_types: List[int],
atomic: bool = False,
fparam: Optional[np.ndarray] = None,
aparam: Optional[np.ndarray] = None,
efield: Optional[np.ndarray] = None,
) -> np.ndarray:
"""Evaluate the model.
Parameters
----------
coords
The coordinates of atoms.
The array should be of size nframes x natoms x 3
cells
The cell of the region.
If None then non-PBC is assumed, otherwise using PBC.
The array should be of size nframes x 9
atom_types
The atom types
The list should contain natoms ints
atomic
Not used in this model
fparam
Not used in this model
aparam
Not used in this model
efield
Not used in this model
Returns
-------
tensor
The returned tensor
If atomic == False then of size nframes x variable_dof
else of size nframes x natoms x variable_dof
"""
return DeepTensor.eval(self, coords, cells, atom_types, atomic=False)
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")
_skbuild/linux-x86_64-3.6/cmake-install/deepmd/infer/deep_pot.py deleted 100644 → 0
View file @ 6b33aeb8
import logging
from typing import TYPE_CHECKING, List, Optional, Tuple, Union, Callable
import numpy as np
from deepmd.common import make_default_mesh
from deepmd.env import default_tf_session_config, tf
from deepmd.infer.data_modifier import DipoleChargeModifier
from deepmd.infer.deep_eval import DeepEval
from deepmd.utils.sess import run_sess
from deepmd.utils.batch_size import AutoBatchSize
if TYPE_CHECKING:
from pathlib import Path
log = logging.getLogger(__name__)
class DeepPot(DeepEval):
"""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
auto_batch_size : bool or int or AutomaticBatchSize, default: True
If True, automatic batch size will be used. If int, it will be used
as the initial batch size.
Examples
--------
>>> from deepmd.infer import DeepPot
>>> import numpy as np
>>> dp = DeepPot('graph.pb')
>>> 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]
>>> e, f, v = dp.eval(coord, cell, atype)
where `e`, `f` and `v` are predicted energy, force and virial of the system, respectively.
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,
auto_batch_size: Union[bool, int, AutoBatchSize] = True,
) -> None:
# add these tensors on top of what is defined by DeepTensor Class
# use this in favor of dict update to move attribute from class to
# instance namespace
self.tensors = dict(
{
# descrpt attrs
"t_ntypes": "descrpt_attr/ntypes:0",
"t_rcut": "descrpt_attr/rcut:0",
# fitting attrs
"t_dfparam": "fitting_attr/dfparam:0",
"t_daparam": "fitting_attr/daparam:0",
# model attrs
"t_tmap": "model_attr/tmap:0",
# inputs
"t_coord": "t_coord:0",
"t_type": "t_type:0",
"t_natoms": "t_natoms:0",
"t_box": "t_box:0",
"t_mesh": "t_mesh:0",
# add output tensors
"t_energy": "o_energy:0",
"t_force": "o_force:0",
"t_virial": "o_virial:0",
"t_ae": "o_atom_energy:0",
"t_av": "o_atom_virial:0",
"t_descriptor": "o_descriptor:0",
},
)
DeepEval.__init__(
self,
model_file,
load_prefix=load_prefix,
default_tf_graph=default_tf_graph,
auto_batch_size=auto_batch_size,
)
# load optional tensors
operations = [op.name for op in self.graph.get_operations()]
# check if the graph has these operations:
# if yes add them
if 't_efield' in operations:
self._get_tensor("t_efield:0", "t_efield")
self.has_efield = True
else:
log.debug(f"Could not get tensor 't_efield:0'")
self.t_efield = None
self.has_efield = False
if 'load/t_fparam' in operations:
self.tensors.update({"t_fparam": "t_fparam:0"})
self.has_fparam = True
else:
log.debug(f"Could not get tensor 't_fparam:0'")
self.t_fparam = None
self.has_fparam = False
if 'load/t_aparam' in operations:
self.tensors.update({"t_aparam": "t_aparam:0"})
self.has_aparam = True
else:
log.debug(f"Could not get tensor 't_aparam:0'")
self.t_aparam = None
self.has_aparam = False
# now load tensors to object attributes
for attr_name, tensor_name in self.tensors.items():
self._get_tensor(tensor_name, attr_name)
self._run_default_sess()
self.tmap = self.tmap.decode('UTF-8').split()
# setup modifier
try:
t_modifier_type = self._get_tensor("modifier_attr/type:0")
self.modifier_type = run_sess(self.sess, t_modifier_type).decode("UTF-8")
except (ValueError, KeyError):
self.modifier_type = None
if self.modifier_type == "dipole_charge":
t_mdl_name = self._get_tensor("modifier_attr/mdl_name:0")
t_mdl_charge_map = self._get_tensor("modifier_attr/mdl_charge_map:0")
t_sys_charge_map = self._get_tensor("modifier_attr/sys_charge_map:0")
t_ewald_h = self._get_tensor("modifier_attr/ewald_h:0")
t_ewald_beta = self._get_tensor("modifier_attr/ewald_beta:0")
[mdl_name, mdl_charge_map, sys_charge_map, ewald_h, ewald_beta] = run_sess(self.sess, [t_mdl_name, t_mdl_charge_map, t_sys_charge_map, t_ewald_h, t_ewald_beta])
mdl_name = mdl_name.decode("UTF-8")
mdl_charge_map = [int(ii) for ii in mdl_charge_map.decode("UTF-8").split()]
sys_charge_map = [int(ii) for ii in sys_charge_map.decode("UTF-8").split()]
self.dm = DipoleChargeModifier(mdl_name, mdl_charge_map, sys_charge_map, ewald_h = ewald_h, ewald_beta = ewald_beta)
def _run_default_sess(self):
[self.ntypes, self.rcut, self.dfparam, self.daparam, self.tmap] = run_sess(self.sess,
[self.t_ntypes, self.t_rcut, self.t_dfparam, self.t_daparam, self.t_tmap]
)
def get_ntypes(self) -> int:
"""Get the number of atom types of this model."""
return self.ntypes
def get_rcut(self) -> float:
"""Get the cut-off radius of this model."""
return self.rcut
def get_type_map(self) -> List[int]:
"""Get the type map (element name of the atom types) of this model."""
return self.tmap
def get_sel_type(self) -> List[int]:
"""Unsupported in this model."""
raise NotImplementedError("This model type does not support this attribute")
def get_dim_fparam(self) -> int:
"""Get the number (dimension) of frame parameters of this DP."""
return self.dfparam
def get_dim_aparam(self) -> int:
"""Get the number (dimension) of atomic parameters of this DP."""
return self.daparam
def _eval_func(self, inner_func: Callable, numb_test: int, natoms: int) -> Callable:
"""Wrapper method with auto batch size.
Parameters
----------
inner_func : Callable
the method to be wrapped
numb_test: int
number of tests
natoms : int
number of atoms
Returns
-------
Callable
the wrapper
"""
if self.auto_batch_size is not None:
def eval_func(*args, **kwargs):
return self.auto_batch_size.execute_all(inner_func, numb_test, natoms, *args, **kwargs)
else:
eval_func = inner_func
return eval_func
def _get_natoms_and_nframes(self, coords: np.ndarray, atom_types: List[int]) -> Tuple[int, int]:
natoms = len(atom_types)
coords = np.reshape(np.array(coords), [-1, natoms * 3])
nframes = coords.shape[0]
return natoms, nframes
def eval(
self,
coords: np.ndarray,
cells: np.ndarray,
atom_types: List[int],
atomic: bool = False,
fparam: Optional[np.ndarray] = None,
aparam: Optional[np.ndarray] = None,
efield: Optional[np.ndarray] = None,
) -> Tuple[np.ndarray, ...]:
"""Evaluate the energy, force and virial by using this DP.
Parameters
----------
coords
The coordinates of atoms.
The array should be of size nframes x natoms x 3
cells
The cell of the region.
If None then non-PBC is assumed, otherwise using PBC.
The array should be of size nframes x 9
atom_types
The atom types
The list should contain natoms ints
atomic
Calculate the atomic energy and virial
fparam
The frame parameter.
The array can be of size :
- nframes x dim_fparam.
- dim_fparam. Then all frames are assumed to be provided with the same fparam.
aparam
The atomic parameter
The array can be of size :
- nframes x natoms x dim_aparam.
- natoms x dim_aparam. Then all frames are assumed to be provided with the same aparam.
- dim_aparam. Then all frames and atoms are provided with the same aparam.
efield
The external field on atoms.
The array should be of size nframes x natoms x 3
Returns
-------
energy
The system energy.
force
The force on each atom
virial
The virial
atom_energy
The atomic energy. Only returned when atomic == True
atom_virial
The atomic virial. Only returned when atomic == True
"""
# reshape coords before getting shape
natoms, numb_test = self._get_natoms_and_nframes(coords, atom_types)
output = self._eval_func(self._eval_inner, numb_test, natoms)(coords, cells, atom_types, fparam = fparam, aparam = aparam, atomic = atomic, efield = efield)
if self.modifier_type is not None:
if atomic:
raise RuntimeError('modifier does not support atomic modification')
me, mf, mv = self.dm.eval(coords, cells, atom_types)
output = list(output) # tuple to list
e, f, v = output[:3]
output[0] += me.reshape(e.shape)
output[1] += mf.reshape(f.shape)
output[2] += mv.reshape(v.shape)
output = tuple(output)
return output
def _prepare_feed_dict(
self,
coords,
cells,
atom_types,
fparam=None,
aparam=None,
atomic=False,
efield=None
):
# standarize the shape of inputs
natoms, nframes = self._get_natoms_and_nframes(coords, atom_types)
atom_types = np.array(atom_types, dtype = int).reshape([-1])
coords = np.reshape(np.array(coords), [-1, natoms * 3])
if cells is None:
pbc = False
# make cells to work around the requirement of pbc
cells = np.tile(np.eye(3), [nframes, 1]).reshape([nframes, 9])
else:
pbc = True
cells = np.array(cells).reshape([nframes, 9])
if self.has_fparam :
assert(fparam is not None)
fparam = np.array(fparam)
if self.has_aparam :
assert(aparam is not None)
aparam = np.array(aparam)
if self.has_efield :
assert(efield is not None), "you are using a model with external field, parameter efield should be provided"
efield = np.array(efield)
# reshape the inputs
if self.has_fparam :
fdim = self.get_dim_fparam()
if fparam.size == nframes * fdim :
fparam = np.reshape(fparam, [nframes, fdim])
elif fparam.size == fdim :
fparam = np.tile(fparam.reshape([-1]), [nframes, 1])
else :
raise RuntimeError('got wrong size of frame param, should be either %d x %d or %d' % (nframes, fdim, fdim))
if self.has_aparam :
fdim = self.get_dim_aparam()
if aparam.size == nframes * natoms * fdim:
aparam = np.reshape(aparam, [nframes, natoms * fdim])
elif aparam.size == natoms * fdim :
aparam = np.tile(aparam.reshape([-1]), [nframes, 1])
elif aparam.size == fdim :
aparam = np.tile(aparam.reshape([-1]), [nframes, natoms])
else :
raise RuntimeError('got wrong size of frame param, should be either %d x %d x %d or %d x %d or %d' % (nframes, natoms, fdim, natoms, fdim, fdim))
# sort inputs
coords, atom_types, imap = self.sort_input(coords, atom_types)
if self.has_efield:
efield = np.reshape(efield, [nframes, natoms, 3])
efield = efield[:,imap,:]
efield = np.reshape(efield, [nframes, natoms*3])
# make natoms_vec and default_mesh
natoms_vec = self.make_natoms_vec(atom_types)
assert(natoms_vec[0] == natoms)
# evaluate
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] = np.reshape(coords, [-1])
feed_dict_test[self.t_box ] = np.reshape(cells , [-1])
if self.has_efield:
feed_dict_test[self.t_efield]= np.reshape(efield, [-1])
if pbc:
feed_dict_test[self.t_mesh ] = make_default_mesh(cells)
else:
feed_dict_test[self.t_mesh ] = np.array([], dtype = np.int32)
if self.has_fparam:
feed_dict_test[self.t_fparam] = np.reshape(fparam, [-1])
if self.has_aparam:
feed_dict_test[self.t_aparam] = np.reshape(aparam, [-1])
return feed_dict_test, imap
def _eval_inner(
self,
coords,
cells,
atom_types,
fparam=None,
aparam=None,
atomic=False,
efield=None
):
natoms, nframes = self._get_natoms_and_nframes(coords, atom_types)
feed_dict_test, imap = self._prepare_feed_dict(coords, cells, atom_types, fparam, aparam, efield)
t_out = [self.t_energy,
self.t_force,
self.t_virial]
if atomic :
t_out += [self.t_ae,
self.t_av]
v_out = run_sess(self.sess, t_out, feed_dict = feed_dict_test)
energy = v_out[0]
force = v_out[1]
virial = v_out[2]
if atomic:
ae = v_out[3]
av = v_out[4]
# reverse map of the outputs
force = self.reverse_map(np.reshape(force, [nframes,-1,3]), imap)
if atomic :
ae = self.reverse_map(np.reshape(ae, [nframes,-1,1]), imap)
av = self.reverse_map(np.reshape(av, [nframes,-1,9]), imap)
energy = np.reshape(energy, [nframes, 1])
force = np.reshape(force, [nframes, natoms, 3])
virial = np.reshape(virial, [nframes, 9])
if atomic:
ae = np.reshape(ae, [nframes, natoms, 1])
av = np.reshape(av, [nframes, natoms, 9])
return energy, force, virial, ae, av
else :
return energy, force, virial
def eval_descriptor(self,
coords: np.ndarray,
cells: np.ndarray,
atom_types: List[int],
fparam: Optional[np.ndarray] = None,
aparam: Optional[np.ndarray] = None,
efield: Optional[np.ndarray] = None,
) -> np.array:
"""Evaluate descriptors by using this DP.
Parameters
----------
coords
The coordinates of atoms.
The array should be of size nframes x natoms x 3
cells
The cell of the region.
If None then non-PBC is assumed, otherwise using PBC.
The array should be of size nframes x 9
atom_types
The atom types
The list should contain natoms ints
fparam
The frame parameter.
The array can be of size :
- nframes x dim_fparam.
- dim_fparam. Then all frames are assumed to be provided with the same fparam.
aparam
The atomic parameter
The array can be of size :
- nframes x natoms x dim_aparam.
- natoms x dim_aparam. Then all frames are assumed to be provided with the same aparam.
- dim_aparam. Then all frames and atoms are provided with the same aparam.
efield
The external field on atoms.
The array should be of size nframes x natoms x 3
Returns
-------
descriptor
Descriptors.
"""
natoms, numb_test = self._get_natoms_and_nframes(coords, atom_types)
descriptor = self._eval_func(self._eval_descriptor_inner, numb_test, natoms)(coords, cells, atom_types, fparam = fparam, aparam = aparam, efield = efield)
return descriptor
def _eval_descriptor_inner(self,
coords: np.ndarray,
cells: np.ndarray,
atom_types: List[int],
fparam: Optional[np.ndarray] = None,
aparam: Optional[np.ndarray] = None,
efield: Optional[np.ndarray] = None,
) -> np.array:
natoms, nframes = self._get_natoms_and_nframes(coords, atom_types)
feed_dict_test, imap = self._prepare_feed_dict(coords, cells, atom_types, fparam, aparam, efield)
descriptor, = run_sess(self.sess, [self.t_descriptor], feed_dict = feed_dict_test)
return self.reverse_map(np.reshape(descriptor, [nframes, natoms, -1]), imap)
_skbuild/linux-x86_64-3.6/cmake-install/deepmd/infer/deep_tensor.py deleted 100644 → 0
View file @ 6b33aeb8
import os
from typing import List, Optional, TYPE_CHECKING, Tuple
import numpy as np
from deepmd.common import make_default_mesh
from deepmd.env import default_tf_session_config, tf
from deepmd.infer.deep_eval import DeepEval
from deepmd.utils.sess import run_sess
if TYPE_CHECKING:
from pathlib import Path
class DeepTensor(DeepEval):
"""Evaluates a tensor model.
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
"""
tensors = {
# descriptor attrs
"t_ntypes": "descrpt_attr/ntypes:0",
"t_rcut": "descrpt_attr/rcut:0",
# model attrs
"t_tmap": "model_attr/tmap:0",
"t_sel_type": "model_attr/sel_type:0",
"t_ouput_dim": "model_attr/output_dim:0",
# inputs
"t_coord": "t_coord:0",
"t_type": "t_type:0",
"t_natoms": "t_natoms:0",
"t_box": "t_box:0",
"t_mesh": "t_mesh:0",
}
def __init__(
self,
model_file: "Path",
load_prefix: str = 'load',
default_tf_graph: bool = False
) -> None:
"""Constructor"""
DeepEval.__init__(
self,
model_file,
load_prefix=load_prefix,
default_tf_graph=default_tf_graph
)
# check model type
model_type = self.tensors["t_tensor"][2:-2]
assert self.model_type == model_type, \
f"expect {model_type} model but got {self.model_type}"
# now load tensors to object attributes
for attr_name, tensor_name in self.tensors.items():
self._get_tensor(tensor_name, attr_name)
# load optional tensors if possible
optional_tensors = {
"t_global_tensor": f"o_global_{model_type}:0",
"t_force": "o_force:0",
"t_virial": "o_virial:0",
"t_atom_virial": "o_atom_virial:0"
}
try:
# first make sure these tensor all exists (but do not modify self attr)
for attr_name, tensor_name in optional_tensors.items():
self._get_tensor(tensor_name)
# then put those into self.attrs
for attr_name, tensor_name in optional_tensors.items():
self._get_tensor(tensor_name, attr_name)
except KeyError:
self._support_gfv = False
else:
self.tensors.update(optional_tensors)
self._support_gfv = True
self._run_default_sess()
self.tmap = self.tmap.decode('UTF-8').split()
def _run_default_sess(self):
[self.ntypes, self.rcut, self.tmap, self.tselt, self.output_dim] \
= run_sess(self.sess,
[self.t_ntypes, self.t_rcut, self.t_tmap, self.t_sel_type, self.t_ouput_dim]
)
def get_ntypes(self) -> int:
"""Get the number of atom types of this model."""
return self.ntypes
def get_rcut(self) -> float:
"""Get the cut-off radius of this model."""
return self.rcut
def get_type_map(self) -> List[int]:
"""Get the type map (element name of the atom types) of this model."""
return self.tmap
def get_sel_type(self) -> List[int]:
"""Get the selected atom types of this model."""
return self.tselt
def get_dim_fparam(self) -> int:
"""Get the number (dimension) of frame parameters of this DP."""
return self.dfparam
def get_dim_aparam(self) -> int:
"""Get the number (dimension) of atomic parameters of this DP."""
return self.daparam
def eval(
self,
coords: np.ndarray,
cells: np.ndarray,
atom_types: List[int],
atomic: bool = True,
fparam: Optional[np.ndarray] = None,
aparam: Optional[np.ndarray] = None,
efield: Optional[np.ndarray] = None
) -> np.ndarray:
"""Evaluate the model.
Parameters
----------
coords
The coordinates of atoms.
The array should be of size nframes x natoms x 3
cells
The cell of the region.
If None then non-PBC is assumed, otherwise using PBC.
The array should be of size nframes x 9
atom_types
The atom types
The list should contain natoms ints
atomic
If True (default), return the atomic tensor
Otherwise return the global tensor
fparam
Not used in this model
aparam
Not used in this model
efield
Not used in this model
Returns
-------
tensor
The returned tensor
If atomic == False then of size nframes x output_dim
else of size nframes x natoms x output_dim
"""
# standarize the shape of inputs
atom_types = np.array(atom_types, dtype = int).reshape([-1])
natoms = atom_types.size
coords = np.reshape(np.array(coords), [-1, natoms * 3])
nframes = coords.shape[0]
if cells is None:
pbc = False
cells = np.tile(np.eye(3), [nframes, 1]).reshape([nframes, 9])
else:
pbc = True
cells = np.array(cells).reshape([nframes, 9])
# 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)
# evaluate
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] = np.reshape(coords, [-1])
feed_dict_test[self.t_box ] = np.reshape(cells , [-1])
if pbc:
feed_dict_test[self.t_mesh ] = make_default_mesh(cells)
else:
feed_dict_test[self.t_mesh ] = np.array([], dtype = np.int32)
if atomic:
assert "global" not in self.model_type, \
f"cannot do atomic evaluation with model type {self.model_type}"
t_out = [self.t_tensor]
else:
assert self._support_gfv or "global" in self.model_type, \
f"do not support global tensor evaluation with old {self.model_type} model"
t_out = [self.t_global_tensor if self._support_gfv else self.t_tensor]
v_out = self.sess.run (t_out, feed_dict = feed_dict_test)
tensor = v_out[0]
# reverse map of the outputs
if atomic:
tensor = np.array(tensor)
tensor = self.reverse_map(np.reshape(tensor, [nframes,-1,self.output_dim]), sel_imap)
tensor = np.reshape(tensor, [nframes, len(sel_at), self.output_dim])
else:
tensor = np.reshape(tensor, [nframes, self.output_dim])
return tensor
def eval_full(
self,
coords: np.ndarray,
cells: np.ndarray,
atom_types: List[int],
atomic: bool = False,
fparam: Optional[np.array] = None,
aparam: Optional[np.array] = None,
efield: Optional[np.array] = None
) -> Tuple[np.ndarray, ...]:
"""Evaluate the model with interface similar to the energy model.
Will return global tensor, component-wise force and virial
and optionally atomic tensor and atomic virial.
Parameters
----------
coords
The coordinates of atoms.
The array should be of size nframes x natoms x 3
cells
The cell of the region.
If None then non-PBC is assumed, otherwise using PBC.
The array should be of size nframes x 9
atom_types
The atom types
The list should contain natoms ints
atomic
Whether to calculate atomic tensor and virial
fparam
Not used in this model
aparam
Not used in this model
efield
Not used in this model
Returns
-------
tensor
The global tensor.
shape: [nframes x nout]
force
The component-wise force (negative derivative) on each atom.
shape: [nframes x nout x natoms x 3]
virial
The component-wise virial of the tensor.
shape: [nframes x nout x 9]
atom_tensor
The atomic tensor. Only returned when atomic == True
shape: [nframes x natoms x nout]
atom_virial
The atomic virial. Only returned when atomic == True
shape: [nframes x nout x natoms x 9]
"""
assert self._support_gfv, \
f"do not support eval_full with old tensor model"
# standarize the shape of inputs
atom_types = np.array(atom_types, dtype = int).reshape([-1])
natoms = atom_types.size
coords = np.reshape(np.array(coords), [-1, natoms * 3])
nframes = coords.shape[0]
if cells is None:
pbc = False
cells = np.tile(np.eye(3), [nframes, 1]).reshape([nframes, 9])
else:
pbc = True
cells = np.array(cells).reshape([nframes, 9])
nout = self.output_dim
# 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)
# evaluate
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] = np.reshape(coords, [-1])
feed_dict_test[self.t_box ] = np.reshape(cells , [-1])
if pbc:
feed_dict_test[self.t_mesh ] = make_default_mesh(cells)
else:
feed_dict_test[self.t_mesh ] = np.array([], dtype = np.int32)
t_out = [self.t_global_tensor,
self.t_force,
self.t_virial]
if atomic :
t_out += [self.t_tensor,
self.t_atom_virial]
v_out = self.sess.run (t_out, feed_dict = feed_dict_test)
gt = v_out[0] # global tensor
force = v_out[1]
virial = v_out[2]
if atomic:
at = v_out[3] # atom tensor
av = v_out[4] # atom virial
# please note here the shape are wrong!
force = self.reverse_map(np.reshape(force, [nframes*nout, natoms ,3]), imap)
if atomic:
at = self.reverse_map(np.reshape(at, [nframes, len(sel_at), nout]), sel_imap)
av = self.reverse_map(np.reshape(av, [nframes*nout, natoms, 9]), imap)
# make sure the shapes are correct here
gt = np.reshape(gt, [nframes, nout])
force = np.reshape(force, [nframes, nout, natoms, 3])
virial = np.reshape(virial, [nframes, nout, 9])
if atomic:
at = np.reshape(at, [nframes, len(sel_at), self.output_dim])
av = np.reshape(av, [nframes, nout, natoms, 9])
return gt, force, virial, at, av
else:
return gt, force, virial
\ No newline at end of file
_skbuild/linux-x86_64-3.6/cmake-install/deepmd/infer/deep_wfc.py deleted 100644 → 0
View file @ 6b33aeb8
from deepmd.infer.deep_tensor import DeepTensor
from typing import TYPE_CHECKING
if TYPE_CHECKING:
from pathlib import Path
class DeepWFC(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_wfc: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")
_skbuild/linux-x86_64-3.6/cmake-install/deepmd/infer/ewald_recp.py deleted 100644 → 0
View file @ 6b33aeb8
import numpy as np
from typing import Tuple, List
from deepmd.env import tf
from deepmd.common import ClassArg
from deepmd.env import 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.env import default_tf_session_config
from deepmd.utils.sess import run_sess
class EwaldRecp () :
"""
Evaluate the reciprocal part of the Ewald sum
"""
def __init__(self,
hh,
beta):
"""
Constructor
Parameters
----------
hh
Grid spacing of the reciprocal part of Ewald sum. Unit: A
beta
Splitting parameter of the Ewald sum. Unit: A^{-1}
"""
self.hh = hh
self.beta = beta
with tf.Graph().as_default() as graph:
# place holders
self.t_nloc = tf.placeholder(tf.int32, [1], name = "t_nloc")
self.t_coord = tf.placeholder(GLOBAL_TF_FLOAT_PRECISION, [None], name='t_coord')
self.t_charge = tf.placeholder(GLOBAL_TF_FLOAT_PRECISION, [None], name='t_charge')
self.t_box = tf.placeholder(GLOBAL_TF_FLOAT_PRECISION, [None], name='t_box')
# output
self.t_energy, self.t_force, self.t_virial \
= op_module.ewald_recp(self.t_coord, self.t_charge, self.t_nloc, self.t_box,
ewald_h = self.hh,
ewald_beta = self.beta)
self.sess = tf.Session(graph=graph, config=default_tf_session_config)
def eval(self,
coord : np.ndarray,
charge : np.ndarray,
box : np.ndarray
) -> Tuple[np.ndarray, np.ndarray, np.ndarray] :
"""
Evaluate
Parameters
----------
coord
The coordinates of atoms
charge
The atomic charge
box
The simulation region. PBC is assumed
Returns
-------
e
The energy
f
The force
v
The virial
"""
coord = np.array(coord)
charge = np.array(charge)
box = np.array(box)
nframes = charge.shape[0]
natoms = charge.shape[1]
coord = np.reshape(coord, [nframes * 3 * natoms])
charge = np.reshape(charge, [nframes * natoms])
box = np.reshape(box, [nframes * 9])
[energy, force, virial] \
= run_sess(self.sess, [self.t_energy, self.t_force, self.t_virial],
feed_dict = {
self.t_coord: coord,
self.t_charge: charge,
self.t_box: box,
self.t_nloc: [natoms],
})
return energy, force, virial
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