import numpy as np import logging from deepmd.nvnmd.data.data import jdata_config, jdata_configs, jdata_deepmd_input from deepmd.nvnmd.data.data import NVNMD_WELCOME, NVNMD_CITATION from deepmd.nvnmd.utils.fio import FioDic log = logging.getLogger(__name__) 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) 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() def init_value(self): r"""Initial member with dict """ self.dscp = self.config['dscp'] self.fitn = self.config['fitn'] self.size = self.config['size'] self.ctrl = self.config['ctrl'] self.nbit = self.config['nbit'] 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 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['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() 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") def init_from_deepmd_input(self, jdata): r"""Initial members with input script of deepmd """ self.config['dscp'] = FioDic().update(jdata['descriptor'], self.config['dscp']) self.config['fitn'] = FioDic().update(jdata['fitting_net'], self.config['fitn']) self.config['dscp'] = self.init_dscp(self.config['dscp'], self.config) self.config['fitn'] = self.init_fitn(self.config['fitn'], self.config) # self.init_net_size() self.init_value() def init_dscp(self, jdata: dict, jdata_parent: dict = {}) -> dict: r"""Initial members about descriptor """ jdata['M1'] = jdata['neuron'][-1] jdata['M2'] = jdata['axis_neuron'] 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) jdata['NIDP'] = int(np.sum(jdata['sel'])) jdata['NIX'] = 2 ** int(np.ceil(np.log2(jdata['NIDP'] / 1.5))) jdata['SEL'] = (jdata['sel'] + [0, 0, 0, 0])[0:4] jdata['ntype'] = len(jdata['sel']) jdata['ntypex'] = 1 if(jdata['same_net']) else jdata['ntype'] return jdata 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 def init_size(self, jdata: dict, jdata_parent: dict = {}) -> dict: r"""Initial members about ram capacity """ jdata['Na'] = jdata['NSPU'] jdata['NaX'] = jdata['MSPU'] return jdata 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 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'] jdata['NBIT_CRD'] = jdata['NBIT_DATA'] * 3 jdata['NBIT_LST'] = int(np.ceil(np.log2(NaX))) jdata['NBIT_ATOM'] = jdata['NBIT_SPE'] + jdata['NBIT_CRD'] jdata['NBIT_LONG_ATOM'] = jdata['NBIT_SPE'] + jdata['NBIT_LONG_DATA'] * 3 jdata['NBIT_RIJ'] = jdata['NBIT_DATA_FL'] + 5 jdata['NBIT_SUM'] = jdata['NBIT_DATA_FL'] + 8 jdata['NBIT_DATA2'] = jdata['NBIT_DATA'] + jdata['NBIT_DATA_FL'] jdata['NBIT_DATA2_FL'] = 2 * jdata['NBIT_DATA_FL'] jdata['NBIT_DATA_FEA'] = jdata['NBIT_DATA'] + jdata['NBIT_FEA_FL'] jdata['NBIT_DATA_FEA_FL'] = jdata['NBIT_DATA_FL'] + jdata['NBIT_FEA_FL'] jdata['NBIT_FORCE_FL'] = 2 * jdata['NBIT_DATA_FL'] - 1 return jdata 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 FioDic().save(file_name, self.config) 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 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 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() return jdata 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 def get_learning_rate_jdata(self): r"""Generate `learning_rate` in input script """ return jdata_deepmd_input['learning_rate'] def get_loss_jdata(self): r"""Generate `loss` in input script """ return jdata_deepmd_input['loss'] def get_training_jdata(self): r"""Generate `training` in input script """ return jdata_deepmd_input['training'] 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 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'])