import logging
import numpy as np
from deepmd.nvnmd.data.data import (
NVNMD_CITATION,
NVNMD_WELCOME,
jdata_config,
jdata_configs,
jdata_deepmd_input,
)
from deepmd.nvnmd.utils.fio import (
FioDic,
)
from deepmd.nvnmd.utils.op import (
r2s,
)
log = logging.getLogger(__name__)
[docs]class NvnmdConfig:
r"""Configuration for NVNMD
record the message of model such as size, using nvnmd or not.
Parameters
----------
jdata
a dictionary of input script
References
----------
DOI: 10.1038/s41524-022-00773-z
"""
def __init__(self, jdata: dict):
self.map = {}
self.config = jdata_config
self.save_path = "nvnmd/config.npy"
self.weight = {}
self.init_from_jdata(jdata)
[docs] def init_from_jdata(self, jdata: dict = {}):
r"""Initial this class with `jdata` loaded from input script."""
if jdata == {}:
return None
self.net_size = jdata["net_size"]
self.map_file = jdata["map_file"]
self.config_file = jdata["config_file"]
self.enable = jdata["enable"]
self.weight_file = jdata["weight_file"]
self.restore_descriptor = jdata["restore_descriptor"]
self.restore_fitting_net = jdata["restore_fitting_net"]
self.quantize_descriptor = jdata["quantize_descriptor"]
self.quantize_fitting_net = jdata["quantize_fitting_net"]
# load data
if self.enable:
self.map = FioDic().load(self.map_file, {})
self.weight = FioDic().load(self.weight_file, {})
jdata_config_ = jdata_config.copy()
jdata_config_["fitn"]["neuron"][0] = self.net_size
load_config = FioDic().load(self.config_file, jdata_config_)
self.init_from_config(load_config)
# if load the file, set net_size
self.init_net_size()
[docs] def init_value(self):
r"""Initial member with dict."""
self.dscp = self.config["dscp"]
self.fitn = self.config["fitn"]
self.dpin = self.config["dpin"]
self.size = self.config["size"]
self.ctrl = self.config["ctrl"]
self.nbit = self.config["nbit"]
[docs] def update_config(self):
r"""Update config from dict."""
self.config["dscp"] = self.dscp
self.config["fitn"] = self.fitn
self.config["dpin"] = self.dpin
self.config["size"] = self.size
self.config["ctrl"] = self.ctrl
self.config["nbit"] = self.nbit
[docs] def init_train_mode(self, mod="cnn"):
r"""Configure for taining cnn or qnn."""
if mod == "cnn":
self.restore_descriptor = False
self.restore_fitting_net = False
self.quantize_descriptor = False
self.quantize_fitting_net = False
elif mod == "qnn":
self.restore_descriptor = True
self.restore_fitting_net = True
self.quantize_descriptor = True
self.quantize_fitting_net = True
[docs] def init_from_config(self, jdata):
r"""Initial member element one by one."""
self.config = FioDic().update(jdata, self.config)
self.config["dscp"] = self.init_dscp(self.config["dscp"], self.config)
self.config["fitn"] = self.init_fitn(self.config["fitn"], self.config)
self.config["dpin"] = self.init_dpin(self.config["dpin"], self.config)
self.config["size"] = self.init_size(self.config["size"], self.config)
self.config["ctrl"] = self.init_ctrl(self.config["ctrl"], self.config)
self.config["nbit"] = self.init_nbit(self.config["nbit"], self.config)
self.init_value()
[docs] def init_net_size(self):
r"""Initial net_size."""
# self.net_size = self.fitn['neuron'][0]
self.net_size = self.config["fitn"]["neuron"][0]
if self.enable:
key = str(self.net_size)
if key in jdata_configs.keys():
# log.info(f"NVNMD: configure the net_size is {key}")
self.init_from_config(jdata_configs[key])
else:
log.error("NVNMD: don't have the configure of net_size")
[docs] def init_dscp(self, jdata: dict, jdata_parent: dict = {}) -> dict:
r"""Initial members about descriptor."""
# embedding
jdata["M1"] = jdata["neuron"][-1]
jdata["M2"] = jdata["axis_neuron"]
jdata["SEL"] = (jdata["sel"] + [0, 0, 0, 0])[0:4]
jdata["NNODE_FEAS"] = [1] + jdata["neuron"]
jdata["nlayer_fea"] = len(jdata["neuron"])
jdata["same_net"] = int(1) if jdata["type_one_side"] else int(0)
# neighbor
jdata["NIDP"] = int(np.sum(jdata["sel"]))
jdata["NIX"] = 2 ** int(np.ceil(np.log2(jdata["NIDP"] / 1.5)))
# type
jdata["ntype"] = len(jdata["sel"])
jdata["ntypex"] = 1 if (jdata["same_net"]) else jdata["ntype"]
return jdata
[docs] def init_fitn(self, jdata: dict, jdata_parent: dict = {}) -> dict:
r"""Initial members about fitting network."""
M1 = jdata_parent["dscp"]["M1"]
M2 = jdata_parent["dscp"]["M2"]
jdata["NNODE_FITS"] = [int(M1 * M2)] + jdata["neuron"] + [1]
jdata["nlayer_fit"] = len(jdata["neuron"]) + 1
jdata["NLAYER"] = jdata["nlayer_fit"]
return jdata
[docs] def init_dpin(self, jdata: dict, jdata_parent: dict = {}) -> dict:
r"""Initial members about other deepmd input."""
return jdata
[docs] def init_size(self, jdata: dict, jdata_parent: dict = {}) -> dict:
r"""Initial members about ram capacity."""
jdata["NAEXT"] = jdata["Na"]
jdata["NTYPE"] = jdata_parent["dscp"]["ntype_max"]
jdata["NTYPEX"] = jdata_parent["dscp"]["ntypex_max"]
return jdata
[docs] def init_ctrl(self, jdata: dict, jdata_parent: dict = {}) -> dict:
r"""Initial members about control signal."""
ntype_max = jdata_parent["dscp"]["ntype_max"]
jdata["NSADV"] = jdata["NSTDM"] + 1
jdata["NSEL"] = jdata["NSTDM"] * ntype_max
if 32 % jdata["NSTDM_M1X"] > 0:
log.warning(
"NVNMD: NSTDM_M1X must be divisor of 32 for the right runing in data_merge module"
)
return jdata
[docs] def init_nbit(self, jdata: dict, jdata_parent: dict = {}) -> dict:
r"""Initial members about quantification precision."""
Na = jdata_parent["size"]["Na"]
NaX = jdata_parent["size"]["NaX"]
ntype_max = jdata_parent["dscp"]["ntype_max"]
NSEL = jdata_parent["ctrl"]["NSEL"]
# general
jdata["NBIT_FLTM"] = 1 + jdata["NBIT_FLTF"]
jdata["NBIT_FLTH"] = 1 + jdata["NBIT_FLTM"]
# atom
jdata["NBIT_ENE_FL"] = jdata["NBIT_FIT_DATA_FL"]
jdata["NBIT_SPE"] = int(np.ceil(np.log2(ntype_max)))
jdata["NBIT_LST"] = int(np.ceil(np.log2(NaX)))
jdata["NBIT_CRD3"] = jdata["NBIT_CRD"] * 3
jdata["NBIT_ATOM"] = jdata["NBIT_SPE"] + jdata["NBIT_CRD3"]
# middle result
jdata["NBIT_SEL"] = int(np.ceil(np.log2(NSEL)))
return jdata
[docs] def save(self, file_name=None):
r"""Save all configuration to file."""
if file_name is None:
file_name = self.save_path
else:
self.save_path = file_name
self.update_config()
FioDic().save(file_name, self.config)
[docs] def get_s_range(self, davg, dstd):
rmin = nvnmd_cfg.dscp["rcut_smth"]
rmax = nvnmd_cfg.dscp["rcut"]
ntype = self.dscp["ntype"]
dmin = self.dscp["dmin"]
#
s0 = r2s(dmin, rmin, rmax)
smin_ = -davg[:ntype, 0] / dstd[:ntype, 0]
smax_ = (s0 - davg[:ntype, 0]) / dstd[:ntype, 0]
smin = np.min(smin_)
smax = np.max(smax_)
self.dscp["smin"] = smin
self.dscp["smax"] = smax
nvnmd_cfg.save()
# check
log.info(f"the range of s is [{smin}, {smax}]")
if smax - smin > 16.0:
log.warning("the range of s is over the limit (smax - smin) > 16.0")
log.warning(
"Please reset the rcut_smth as a bigger value to fix this warning"
)
[docs] def get_dscp_jdata(self):
r"""Generate `model/descriptor` in input script."""
dscp = self.dscp
jdata = jdata_deepmd_input["model"]["descriptor"]
jdata["sel"] = dscp["sel"]
jdata["rcut"] = dscp["rcut"]
jdata["rcut_smth"] = dscp["rcut_smth"]
jdata["neuron"] = dscp["neuron"]
jdata["type_one_side"] = dscp["type_one_side"]
jdata["axis_neuron"] = dscp["axis_neuron"]
return jdata
[docs] def get_fitn_jdata(self):
r"""Generate `model/fitting_net` in input script."""
fitn = self.fitn
jdata = jdata_deepmd_input["model"]["fitting_net"]
jdata["neuron"] = fitn["neuron"]
return jdata
[docs] def get_model_jdata(self):
r"""Generate `model` in input script."""
jdata = jdata_deepmd_input["model"]
jdata["descriptor"] = self.get_dscp_jdata()
jdata["fitting_net"] = self.get_fitn_jdata()
if len(self.dpin["type_map"]) > 0:
jdata["type_map"] = self.dpin["type_map"]
return jdata
[docs] def get_nvnmd_jdata(self):
r"""Generate `nvnmd` in input script."""
jdata = jdata_deepmd_input["nvnmd"]
jdata["net_size"] = self.net_size
jdata["config_file"] = self.config_file
jdata["weight_file"] = self.weight_file
jdata["map_file"] = self.map_file
jdata["enable"] = self.enable
jdata["restore_descriptor"] = self.restore_descriptor
jdata["restore_fitting_net"] = self.restore_fitting_net
jdata["quantize_descriptor"] = self.quantize_descriptor
jdata["quantize_fitting_net"] = self.quantize_fitting_net
return jdata
[docs] def get_learning_rate_jdata(self):
r"""Generate `learning_rate` in input script."""
return jdata_deepmd_input["learning_rate"]
[docs] def get_loss_jdata(self):
r"""Generate `loss` in input script."""
return jdata_deepmd_input["loss"]
[docs] def get_training_jdata(self):
r"""Generate `training` in input script."""
return jdata_deepmd_input["training"]
[docs] def get_deepmd_jdata(self):
r"""Generate input script with member element one by one."""
jdata = jdata_deepmd_input.copy()
jdata["model"] = self.get_model_jdata()
jdata["nvnmd"] = self.get_nvnmd_jdata()
jdata["learning_rate"] = self.get_learning_rate_jdata()
jdata["loss"] = self.get_loss_jdata()
jdata["training"] = self.get_training_jdata()
return jdata
[docs] def disp_message(self):
r"""Display the log of NVNMD."""
NVNMD_CONFIG = (
f"enable: {self.enable}",
f"net_size: {self.net_size}",
f"map_file: {self.map_file}",
f"config_file: {self.config_file}",
f"weight_file: {self.weight_file}",
f"restore_descriptor: {self.restore_descriptor}",
f"restore_fitting_net: {self.restore_fitting_net}",
f"quantize_descriptor: {self.quantize_descriptor}",
f"quantize_fitting_net: {self.quantize_fitting_net}",
)
for message in NVNMD_WELCOME + NVNMD_CITATION + NVNMD_CONFIG:
log.info(message)
# global configuration for nvnmd
nvnmd_cfg = NvnmdConfig(jdata_deepmd_input["nvnmd"])