From 578e89b849c9b99b86ae466ee524d1e8534057f1 Mon Sep 17 00:00:00 2001 From: Han Wang Date: Fri, 26 Mar 2021 23:39:46 +0800 Subject: [PATCH 1/5] refact model --- deepmd/model/__init__.py | 5 + source/train/model.py => deepmd/model/ener.py | 244 ++++-------------- deepmd/model/model_stat.py | 60 +++++ deepmd/model/tensor.py | 151 +++++++++++ deepmd/utils/argcheck.py | 2 + setup.py | 3 +- source/train/CMakeLists.txt | 4 +- source/train/trainer.py | 40 ++- 8 files changed, 303 insertions(+), 206 deletions(-) create mode 100644 deepmd/model/__init__.py rename source/train/model.py => deepmd/model/ener.py (53%) create mode 100644 deepmd/model/model_stat.py create mode 100644 deepmd/model/tensor.py diff --git a/deepmd/model/__init__.py b/deepmd/model/__init__.py new file mode 100644 index 0000000000..e50a138aec --- /dev/null +++ b/deepmd/model/__init__.py @@ -0,0 +1,5 @@ +from .ener import EnerModel +from .tensor import WFCModel +from .tensor import DipoleModel +from .tensor import PolarModel +from .tensor import GlobalPolarModel diff --git a/source/train/model.py b/deepmd/model/ener.py similarity index 53% rename from source/train/model.py rename to deepmd/model/ener.py index 8ccee419bf..26c93ed9ba 100644 --- a/source/train/model.py +++ b/deepmd/model/ener.py @@ -1,92 +1,65 @@ import numpy as np +from typing import Tuple, List + from deepmd.env import tf -from collections import defaultdict from deepmd.utils.pair_tab import PairTab from deepmd.common import ClassArg - from deepmd.run_options import global_cvt_2_ener_float, MODEL_VERSION from deepmd.env import op_module +from .model_stat import make_stat_input, merge_sys_stat - -def _make_all_stat_ref(data, nbatches): - all_stat = defaultdict(list) - for ii in range(data.get_nsystems()) : - for jj in range(nbatches) : - stat_data = data.get_batch (sys_idx = ii) - for dd in stat_data: - if dd == "natoms_vec": - stat_data[dd] = stat_data[dd].astype(np.int32) - all_stat[dd].append(stat_data[dd]) - return all_stat - - -def make_stat_input(data, nbatches, merge_sys = True): - """ - pack data for statistics - Parameters - ---------- - data: - The data - merge_sys: bool (True) - Merge system data - Returns - ------- - all_stat: - A dictionary of list of list storing data for stat. - if merge_sys == False data can be accessed by - all_stat[key][sys_idx][batch_idx][frame_idx] - else merge_sys == True can be accessed by - all_stat[key][batch_idx][frame_idx] - """ - all_stat = defaultdict(list) - for ii in range(data.get_nsystems()) : - sys_stat = defaultdict(list) - for jj in range(nbatches) : - stat_data = data.get_batch (sys_idx = ii) - for dd in stat_data: - if dd == "natoms_vec": - stat_data[dd] = stat_data[dd].astype(np.int32) - sys_stat[dd].append(stat_data[dd]) - for dd in sys_stat: - if merge_sys: - for bb in sys_stat[dd]: - all_stat[dd].append(bb) - else: - all_stat[dd].append(sys_stat[dd]) - return all_stat - -def merge_sys_stat(all_stat): - first_key = list(all_stat.keys())[0] - nsys = len(all_stat[first_key]) - ret = defaultdict(list) - for ii in range(nsys): - for dd in all_stat: - for bb in all_stat[dd][ii]: - ret[dd].append(bb) - return ret - - -class Model() : +class EnerModel() : model_type = 'ener' - def __init__ (self, jdata, descrpt, fitting): + def __init__ ( + self, + descrpt, + fitting, + type_map : List[str], + data_stat_nbatch : int = 10, + data_stat_protect : float = 1e-2, + use_srtab : str = None, + smin_alpha : float = None, + sw_rmin : float = None, + sw_rmax : float = None + ) -> None: + """ + Constructor + + Parameters + ---------- + descrpt + Descriptor + fitting + Fitting net + type_map + Mapping atom type to the name (str) of the type. + For example `type_map[1]` gives the name of the type 1. + data_stat_nbatch + Number of frames used for data statistic + data_stat_protect + Protect parameter for atomic energy regression + use_srtab + The table for the short-range pairwise interaction added on top of DP. The table is a text data file with (N_t + 1) * N_t / 2 + 1 columes. The first colume is the distance between atoms. The second to the last columes are energies for pairs of certain types. For example we have two atom types, 0 and 1. The columes from 2nd to 4th are for 0-0, 0-1 and 1-1 correspondingly. + smin_alpha + The short-range tabulated interaction will be swithed according to the distance of the nearest neighbor. This distance is calculated by softmin. This parameter is the decaying parameter in the softmin. It is only required when `use_srtab` is provided. + sw_rmin + The lower boundary of the interpolation between short-range tabulated interaction and DP. It is only required when `use_srtab` is provided. + sw_rmin + The upper boundary of the interpolation between short-range tabulated interaction and DP. It is only required when `use_srtab` is provided. + """ + # descriptor self.descrpt = descrpt self.rcut = self.descrpt.get_rcut() self.ntypes = self.descrpt.get_ntypes() # fitting self.fitting = fitting self.numb_fparam = self.fitting.get_numb_fparam() - - args = ClassArg()\ - .add('type_map', list, default = []) \ - .add('data_stat_nbatch', int, default = 10) \ - .add('data_stat_protect',float, default = 1e-2) \ - .add('use_srtab', str) - class_data = args.parse(jdata) - self.type_map = class_data['type_map'] - self.srtab_name = class_data['use_srtab'] - self.data_stat_nbatch = class_data['data_stat_nbatch'] - self.data_stat_protect = class_data['data_stat_protect'] + # other inputs + self.type_map = type_map + self.data_stat_nbatch = data_stat_nbatch + self.data_stat_protect = data_stat_protect + self.srtab_name = use_srtab if self.srtab_name is not None : self.srtab = PairTab(self.srtab_name) args.add('smin_alpha', float, must = True)\ @@ -265,126 +238,3 @@ def build (self, return model_dict - -class TensorModel() : - def __init__ (self, jdata, descrpt, fitting, var_name): - self.model_type = var_name - self.descrpt = descrpt - self.rcut = self.descrpt.get_rcut() - self.ntypes = self.descrpt.get_ntypes() - # fitting - self.fitting = fitting - - args = ClassArg()\ - .add('type_map', list, default = []) \ - .add('data_stat_nbatch', int, default = 10) \ - .add('data_stat_protect',float, default = 1e-2) - class_data = args.parse(jdata) - self.type_map = class_data['type_map'] - self.data_stat_nbatch = class_data['data_stat_nbatch'] - self.data_stat_protect = class_data['data_stat_protect'] - - def get_rcut (self) : - return self.rcut - - def get_ntypes (self) : - return self.ntypes - - def get_type_map (self) : - return self.type_map - - def get_sel_type(self): - return self.fitting.get_sel_type() - - def get_out_size (self) : - return self.fitting.get_out_size() - - def data_stat(self, data): - all_stat = make_stat_input(data, self.data_stat_nbatch, merge_sys = False) - m_all_stat = merge_sys_stat(all_stat) - self._compute_input_stat (m_all_stat, protection = self.data_stat_protect) - self._compute_output_stat(all_stat) - - def _compute_input_stat(self, all_stat, protection = 1e-2) : - self.descrpt.compute_input_stats(all_stat['coord'], - all_stat['box'], - all_stat['type'], - all_stat['natoms_vec'], - all_stat['default_mesh'], - all_stat) - if hasattr(self.fitting, 'compute_input_stats'): - self.fitting.compute_input_stats(all_stat, protection = protection) - - def _compute_output_stat (self, all_stat) : - if hasattr(self.fitting, 'compute_output_stats'): - self.fitting.compute_output_stats(all_stat) - - def build (self, - coord_, - atype_, - natoms, - box, - mesh, - input_dict, - suffix = '', - reuse = None): - with tf.variable_scope('model_attr' + suffix, reuse = reuse) : - t_tmap = tf.constant(' '.join(self.type_map), - name = 'tmap', - dtype = tf.string) - t_st = tf.constant(self.get_sel_type(), - name = 'sel_type', - dtype = tf.int32) - t_mt = tf.constant(self.model_type, - name = 'model_type', - dtype = tf.string) - t_ver = tf.constant(MODEL_VERSION, - name = 'model_version', - dtype = tf.string) - t_od = tf.constant(self.get_out_size(), - name = 'output_dim', - dtype = tf.int32) - - dout \ - = self.descrpt.build(coord_, - atype_, - natoms, - box, - mesh, - input_dict, - suffix = suffix, - reuse = reuse) - dout = tf.identity(dout, name='o_descriptor') - rot_mat = self.descrpt.get_rot_mat() - rot_mat = tf.identity(rot_mat, name = 'o_rot_mat') - - output = self.fitting.build (dout, - rot_mat, - natoms, - reuse = reuse, - suffix = suffix) - output = tf.identity(output, name = 'o_' + self.model_type) - - return {self.model_type: output} - - -class WFCModel(TensorModel): - def __init__(self, jdata, descrpt, fitting) : - TensorModel.__init__(self, jdata, descrpt, fitting, 'wfc') - - -class DipoleModel(TensorModel): - def __init__(self, jdata, descrpt, fitting) : - TensorModel.__init__(self, jdata, descrpt, fitting, 'dipole') - - -class PolarModel(TensorModel): - def __init__(self, jdata, descrpt, fitting) : - TensorModel.__init__(self, jdata, descrpt, fitting, 'polar') - - -class GlobalPolarModel(TensorModel): - def __init__(self, jdata, descrpt, fitting) : - TensorModel.__init__(self, jdata, descrpt, fitting, 'global_polar') - - diff --git a/deepmd/model/model_stat.py b/deepmd/model/model_stat.py new file mode 100644 index 0000000000..61f151f27b --- /dev/null +++ b/deepmd/model/model_stat.py @@ -0,0 +1,60 @@ +import numpy as np +from collections import defaultdict + +def _make_all_stat_ref(data, nbatches): + all_stat = defaultdict(list) + for ii in range(data.get_nsystems()) : + for jj in range(nbatches) : + stat_data = data.get_batch (sys_idx = ii) + for dd in stat_data: + if dd == "natoms_vec": + stat_data[dd] = stat_data[dd].astype(np.int32) + all_stat[dd].append(stat_data[dd]) + return all_stat + + +def make_stat_input(data, nbatches, merge_sys = True): + """ + pack data for statistics + Parameters + ---------- + data: + The data + merge_sys: bool (True) + Merge system data + Returns + ------- + all_stat: + A dictionary of list of list storing data for stat. + if merge_sys == False data can be accessed by + all_stat[key][sys_idx][batch_idx][frame_idx] + else merge_sys == True can be accessed by + all_stat[key][batch_idx][frame_idx] + """ + all_stat = defaultdict(list) + for ii in range(data.get_nsystems()) : + sys_stat = defaultdict(list) + for jj in range(nbatches) : + stat_data = data.get_batch (sys_idx = ii) + for dd in stat_data: + if dd == "natoms_vec": + stat_data[dd] = stat_data[dd].astype(np.int32) + sys_stat[dd].append(stat_data[dd]) + for dd in sys_stat: + if merge_sys: + for bb in sys_stat[dd]: + all_stat[dd].append(bb) + else: + all_stat[dd].append(sys_stat[dd]) + return all_stat + +def merge_sys_stat(all_stat): + first_key = list(all_stat.keys())[0] + nsys = len(all_stat[first_key]) + ret = defaultdict(list) + for ii in range(nsys): + for dd in all_stat: + for bb in all_stat[dd][ii]: + ret[dd].append(bb) + return ret + diff --git a/deepmd/model/tensor.py b/deepmd/model/tensor.py new file mode 100644 index 0000000000..fa60171255 --- /dev/null +++ b/deepmd/model/tensor.py @@ -0,0 +1,151 @@ +import numpy as np +from typing import Tuple, List + +from deepmd.env import tf +from deepmd.common import ClassArg +from deepmd.run_options import global_cvt_2_ener_float, MODEL_VERSION +from deepmd.env import op_module +from .model_stat import make_stat_input, merge_sys_stat + +class TensorModel() : + def __init__ ( + self, + tensor_name : str, + descrpt, + fitting, + type_map : List[str], + data_stat_nbatch : int = 10, + data_stat_protect : float = 1e-2, + )->None: + """ + Constructor + + Parameters + ---------- + tensor_name + Name of the tensor. + descrpt + Descriptor + fitting + Fitting net + type_map + Mapping atom type to the name (str) of the type. + For example `type_map[1]` gives the name of the type 1. + data_stat_nbatch + Number of frames used for data statistic + data_stat_protect + Protect parameter for atomic energy regression + """ + self.model_type = tensor_name + # descriptor + self.descrpt = descrpt + self.rcut = self.descrpt.get_rcut() + self.ntypes = self.descrpt.get_ntypes() + # fitting + self.fitting = fitting + # other params + self.type_map = type_map + self.data_stat_nbatch = data_stat_nbatch + self.data_stat_protect = data_stat_protect + + def get_rcut (self) : + return self.rcut + + def get_ntypes (self) : + return self.ntypes + + def get_type_map (self) : + return self.type_map + + def get_sel_type(self): + return self.fitting.get_sel_type() + + def get_out_size (self) : + return self.fitting.get_out_size() + + def data_stat(self, data): + all_stat = make_stat_input(data, self.data_stat_nbatch, merge_sys = False) + m_all_stat = merge_sys_stat(all_stat) + self._compute_input_stat (m_all_stat, protection = self.data_stat_protect) + self._compute_output_stat(all_stat) + + def _compute_input_stat(self, all_stat, protection = 1e-2) : + self.descrpt.compute_input_stats(all_stat['coord'], + all_stat['box'], + all_stat['type'], + all_stat['natoms_vec'], + all_stat['default_mesh'], + all_stat) + if hasattr(self.fitting, 'compute_input_stats'): + self.fitting.compute_input_stats(all_stat, protection = protection) + + def _compute_output_stat (self, all_stat) : + if hasattr(self.fitting, 'compute_output_stats'): + self.fitting.compute_output_stats(all_stat) + + def build (self, + coord_, + atype_, + natoms, + box, + mesh, + input_dict, + suffix = '', + reuse = None): + with tf.variable_scope('model_attr' + suffix, reuse = reuse) : + t_tmap = tf.constant(' '.join(self.type_map), + name = 'tmap', + dtype = tf.string) + t_st = tf.constant(self.get_sel_type(), + name = 'sel_type', + dtype = tf.int32) + t_mt = tf.constant(self.model_type, + name = 'model_type', + dtype = tf.string) + t_ver = tf.constant(MODEL_VERSION, + name = 'model_version', + dtype = tf.string) + t_od = tf.constant(self.get_out_size(), + name = 'output_dim', + dtype = tf.int32) + + dout \ + = self.descrpt.build(coord_, + atype_, + natoms, + box, + mesh, + input_dict, + suffix = suffix, + reuse = reuse) + dout = tf.identity(dout, name='o_descriptor') + rot_mat = self.descrpt.get_rot_mat() + rot_mat = tf.identity(rot_mat, name = 'o_rot_mat') + + output = self.fitting.build (dout, + rot_mat, + natoms, + reuse = reuse, + suffix = suffix) + output = tf.identity(output, name = 'o_' + self.model_type) + + return {self.model_type: output} + + +class WFCModel(TensorModel): + def __init__(self, descrpt, fitting, tm, dsn, dsp) : + TensorModel.__init__(self, 'wfc', descrpt, fitting, tm, dsn, dsp) + +class DipoleModel(TensorModel): + def __init__(self, descrpt, fitting, tm, dsn, dsp) : + TensorModel.__init__(self, 'dipole', descrpt, fitting, tm, dsn, dsp) + +class PolarModel(TensorModel): + def __init__(self, descrpt, fitting, tm, dsn, dsp) : + TensorModel.__init__(self, 'polar', descrpt, fitting, tm, dsn, dsp) + +class GlobalPolarModel(TensorModel): + def __init__(self, descrpt, fitting, tm, dsn, dsp) : + TensorModel.__init__(self, 'global_polar', descrpt, fitting, tm, dsn, dsp) + + diff --git a/deepmd/utils/argcheck.py b/deepmd/utils/argcheck.py index a3f4d56b1e..ac993ef107 100644 --- a/deepmd/utils/argcheck.py +++ b/deepmd/utils/argcheck.py @@ -276,6 +276,7 @@ def fitting_variant_type_args(): def model_args (): doc_type_map = 'A list of strings. Give the name to each type of atoms.' doc_data_stat_nbatch = 'The model determines the normalization from the statistics of the data. This key specifies the number of `frames` in each `system` used for statistics.' + doc_data_stat_protect = 'Protect parameter for atomic energy regression.' doc_descrpt = 'The descriptor of atomic environment.' doc_fitting = 'The fitting of physical properties.' doc_use_srtab = 'The table for the short-range pairwise interaction added on top of DP. The table is a text data file with (N_t + 1) * N_t / 2 + 1 columes. The first colume is the distance between atoms. The second to the last columes are energies for pairs of certain types. For example we have two atom types, 0 and 1. The columes from 2nd to 4th are for 0-0, 0-1 and 1-1 correspondingly.' @@ -286,6 +287,7 @@ def model_args (): ca = Argument("model", dict, [Argument("type_map", list, optional = True, doc = doc_type_map), Argument("data_stat_nbatch", int, optional = True, default = 10, doc = doc_data_stat_nbatch), + Argument("data_stat_protect", float, optional = True, default = 1e-2, doc = doc_data_stat_protect), Argument("use_srtab", str, optional = True, doc = doc_use_srtab), Argument("smin_alpha", float, optional = True, doc = doc_smin_alpha), Argument("sw_rmin", float, optional = True, doc = doc_sw_rmin), diff --git a/setup.py b/setup.py index 729d62650d..a415187f35 100644 --- a/setup.py +++ b/setup.py @@ -86,7 +86,8 @@ "deepmd/loggers", "deepmd/cluster", "deepmd/entrypoints", - "deepmd/op" + "deepmd/op", + "deepmd/model", ], python_requires=">=3.6", classifiers=[ diff --git a/source/train/CMakeLists.txt b/source/train/CMakeLists.txt index 8a5317a1d0..8dd7c60d04 100644 --- a/source/train/CMakeLists.txt +++ b/source/train/CMakeLists.txt @@ -2,7 +2,7 @@ configure_file("run_config.ini" "${CMAKE_CURRENT_BINARY_DIR}/run_config.ini" @ONLY) -file(GLOB LIB_PY main.py calculator.py model.py trainer.py run_options.py ) +file(GLOB LIB_PY main.py calculator.py trainer.py run_options.py ) install( FILES ${LIB_PY} @@ -11,4 +11,4 @@ install( install( FILES ${CMAKE_CURRENT_BINARY_DIR}/run_config.ini DESTINATION deepmd/pkg_config -) \ No newline at end of file +) diff --git a/source/train/trainer.py b/source/train/trainer.py index d730a2980a..5331af84c6 100644 --- a/source/train/trainer.py +++ b/source/train/trainer.py @@ -17,7 +17,7 @@ from deepmd.descriptor import DescrptSeR from deepmd.descriptor import DescrptSeAR from deepmd.descriptor import DescrptHybrid -from deepmd.model import Model, WFCModel, DipoleModel, PolarModel, GlobalPolarModel +from deepmd.model import EnerModel, WFCModel, DipoleModel, PolarModel, GlobalPolarModel from deepmd.loss import EnerStdLoss, EnerDipoleLoss, TensorLoss from deepmd.utils.learning_rate import LearningRateExp from deepmd.utils.neighbor_stat import NeighborStat @@ -134,16 +134,44 @@ def _init_param(self, jdata): # init model # infer model type by fitting_type - if fitting_type == Model.model_type: - self.model = Model(model_param, self.descrpt, self.fitting) + if fitting_type == 'ener': + self.model = EnerModel( + self.descrpt, + self.fitting, + model_param['type_map'], + model_param['data_stat_nbatch'], + model_param['data_stat_protect'], + model_param.get('use_srtab'), + model_param.get('smin_alpha'), + model_param.get('sw_rmin'), + model_param.get('sw_rmax') + ) # elif fitting_type == 'wfc': # self.model = WFCModel(model_param, self.descrpt, self.fitting) elif fitting_type == 'dipole': - self.model = DipoleModel(model_param, self.descrpt, self.fitting) + self.model = DipoleModel( + self.descrpt, + self.fitting, + model_param['type_map'], + model_param['data_stat_nbatch'], + model_param['data_stat_protect'] + ) elif fitting_type == 'polar': - self.model = PolarModel(model_param, self.descrpt, self.fitting) + self.model = PolarModel( + self.descrpt, + self.fitting, + model_param['type_map'], + model_param['data_stat_nbatch'], + model_param['data_stat_protect'] + ) elif fitting_type == 'global_polar': - self.model = GlobalPolarModel(model_param, self.descrpt, self.fitting) + self.model = GlobalPolarModel( + self.descrpt, + self.fitting, + model_param['type_map'], + model_param['data_stat_nbatch'], + model_param['data_stat_protect'] + ) else : raise RuntimeError('get unknown fitting type when building model') From d72252622e66067923d56e70ebc2b69c4af18bd1 Mon Sep 17 00:00:00 2001 From: Han Wang Date: Sat, 27 Mar 2021 08:56:41 +0800 Subject: [PATCH 2/5] fix bug of getting inputs --- source/train/trainer.py | 24 ++++++++++++------------ 1 file changed, 12 insertions(+), 12 deletions(-) diff --git a/source/train/trainer.py b/source/train/trainer.py index 5331af84c6..fd9962ec90 100644 --- a/source/train/trainer.py +++ b/source/train/trainer.py @@ -138,9 +138,9 @@ def _init_param(self, jdata): self.model = EnerModel( self.descrpt, self.fitting, - model_param['type_map'], - model_param['data_stat_nbatch'], - model_param['data_stat_protect'], + model_param.get('type_map'), + model_param.get('data_stat_nbatch', 10), + model_param.get('data_stat_protect', 1e-2), model_param.get('use_srtab'), model_param.get('smin_alpha'), model_param.get('sw_rmin'), @@ -152,25 +152,25 @@ def _init_param(self, jdata): self.model = DipoleModel( self.descrpt, self.fitting, - model_param['type_map'], - model_param['data_stat_nbatch'], - model_param['data_stat_protect'] + model_param.get('type_map'), + model_param.get('data_stat_nbatch', 10), + model_param.get('data_stat_protect', 1e-2) ) elif fitting_type == 'polar': self.model = PolarModel( self.descrpt, self.fitting, - model_param['type_map'], - model_param['data_stat_nbatch'], - model_param['data_stat_protect'] + model_param.get('type_map'), + model_param.get('data_stat_nbatch', 10), + model_param.get('data_stat_protect', 1e-2) ) elif fitting_type == 'global_polar': self.model = GlobalPolarModel( self.descrpt, self.fitting, - model_param['type_map'], - model_param['data_stat_nbatch'], - model_param['data_stat_protect'] + model_param.get('type_map'), + model_param.get('data_stat_nbatch', 10), + model_param.get('data_stat_protect', 1e-2) ) else : raise RuntimeError('get unknown fitting type when building model') From 716a7fa0358d8765affda7035ff82636e5164785 Mon Sep 17 00:00:00 2001 From: Han Wang Date: Sat, 27 Mar 2021 09:26:43 +0800 Subject: [PATCH 3/5] fix bug of UTs and srtab --- deepmd/model/ener.py | 17 ++++++++--------- source/tests/test_gen_stat_data.py | 2 +- source/tests/test_model_loc_frame.py | 7 +++++-- source/tests/test_model_se_a.py | 4 ++-- source/tests/test_model_se_a_aparam.py | 4 ++-- source/tests/test_model_se_a_fparam.py | 4 ++-- source/tests/test_model_se_a_srtab.py | 22 ++++++++++++++++------ source/tests/test_model_se_r.py | 4 ++-- 8 files changed, 38 insertions(+), 26 deletions(-) diff --git a/deepmd/model/ener.py b/deepmd/model/ener.py index 26c93ed9ba..56f4880840 100644 --- a/deepmd/model/ener.py +++ b/deepmd/model/ener.py @@ -15,7 +15,7 @@ def __init__ ( self, descrpt, fitting, - type_map : List[str], + type_map : List[str] = None, data_stat_nbatch : int = 10, data_stat_protect : float = 1e-2, use_srtab : str = None, @@ -56,19 +56,18 @@ def __init__ ( self.fitting = fitting self.numb_fparam = self.fitting.get_numb_fparam() # other inputs - self.type_map = type_map + if type_map is None: + self.type_map = [] + else: + self.type_map = type_map self.data_stat_nbatch = data_stat_nbatch self.data_stat_protect = data_stat_protect self.srtab_name = use_srtab if self.srtab_name is not None : self.srtab = PairTab(self.srtab_name) - args.add('smin_alpha', float, must = True)\ - .add('sw_rmin', float, must = True)\ - .add('sw_rmax', float, must = True) - class_data = args.parse(jdata) - self.smin_alpha = class_data['smin_alpha'] - self.sw_rmin = class_data['sw_rmin'] - self.sw_rmax = class_data['sw_rmax'] + self.smin_alpha = smin_alpha + self.sw_rmin = sw_rmin + self.sw_rmax = sw_rmax else : self.srtab = None diff --git a/source/tests/test_gen_stat_data.py b/source/tests/test_gen_stat_data.py index f70d63c3e9..62647efba1 100644 --- a/source/tests/test_gen_stat_data.py +++ b/source/tests/test_gen_stat_data.py @@ -5,7 +5,7 @@ from deepmd.utils.data_system import DeepmdDataSystem from deepmd.fit import EnerFitting -from deepmd.model import make_stat_input, merge_sys_stat, _make_all_stat_ref +from deepmd.model.model_stat import make_stat_input, merge_sys_stat, _make_all_stat_ref def gen_sys(nframes, atom_types): natoms = len(atom_types) diff --git a/source/tests/test_model_loc_frame.py b/source/tests/test_model_loc_frame.py index b0d189e528..c94b71be01 100644 --- a/source/tests/test_model_loc_frame.py +++ b/source/tests/test_model_loc_frame.py @@ -6,7 +6,7 @@ from deepmd.utils.data_system import DataSystem from deepmd.descriptor import DescrptLocFrame from deepmd.fit import EnerFitting -from deepmd.model import Model +from deepmd.model import EnerModel from deepmd.common import j_must_have GLOBAL_ENER_FLOAT_PRECISION = tf.float64 @@ -41,7 +41,10 @@ def test_model(self): fitting = EnerFitting(descrpt, neuron = [240, 120, 60, 30, 10], seed = 1) - model = Model(jdata['model'], descrpt, fitting) + model = EnerModel( + descrpt, + fitting, + ) # model._compute_dstats([test_data['coord']], [test_data['box']], [test_data['type']], [test_data['natoms_vec']], [test_data['default_mesh']]) input_data = {'coord' : [test_data['coord']], diff --git a/source/tests/test_model_se_a.py b/source/tests/test_model_se_a.py index da6f7af3e0..837d344f90 100644 --- a/source/tests/test_model_se_a.py +++ b/source/tests/test_model_se_a.py @@ -7,7 +7,7 @@ from deepmd.utils.data_system import DataSystem from deepmd.descriptor import DescrptSeA from deepmd.fit import EnerFitting -from deepmd.model import Model +from deepmd.model import EnerModel from deepmd.common import j_must_have GLOBAL_ENER_FLOAT_PRECISION = tf.float64 @@ -40,7 +40,7 @@ def test_model(self): descrpt = DescrptSeA(**jdata['model']['descriptor']) jdata['model']['fitting_net']['descrpt'] = descrpt fitting = EnerFitting(**jdata['model']['fitting_net']) - model = Model(jdata['model'], descrpt, fitting) + model = EnerModel(descrpt, fitting) # model._compute_dstats([test_data['coord']], [test_data['box']], [test_data['type']], [test_data['natoms_vec']], [test_data['default_mesh']]) input_data = {'coord' : [test_data['coord']], diff --git a/source/tests/test_model_se_a_aparam.py b/source/tests/test_model_se_a_aparam.py index c3879604f0..7a07a30d20 100644 --- a/source/tests/test_model_se_a_aparam.py +++ b/source/tests/test_model_se_a_aparam.py @@ -6,7 +6,7 @@ from deepmd.utils.data_system import DataSystem from deepmd.descriptor import DescrptSeA from deepmd.fit import EnerFitting -from deepmd.model import Model +from deepmd.model import EnerModel from deepmd.common import j_must_have GLOBAL_ENER_FLOAT_PRECISION = tf.float64 @@ -40,7 +40,7 @@ def test_model(self): descrpt = DescrptSeA(**jdata['model']['descriptor']) jdata['model']['fitting_net']['descrpt'] = descrpt fitting = EnerFitting(**jdata['model']['fitting_net']) - model = Model(jdata['model'], descrpt, fitting) + model = EnerModel(descrpt, fitting) # model._compute_dstats([test_data['coord']], [test_data['box']], [test_data['type']], [test_data['natoms_vec']], [test_data['default_mesh']]) input_data = {'coord' : [test_data['coord']], diff --git a/source/tests/test_model_se_a_fparam.py b/source/tests/test_model_se_a_fparam.py index 36719c9a29..a045c5e520 100644 --- a/source/tests/test_model_se_a_fparam.py +++ b/source/tests/test_model_se_a_fparam.py @@ -6,7 +6,7 @@ from deepmd.utils.data_system import DataSystem from deepmd.descriptor import DescrptSeA from deepmd.fit import EnerFitting -from deepmd.model import Model +from deepmd.model import EnerModel from deepmd.common import j_must_have GLOBAL_ENER_FLOAT_PRECISION = tf.float64 @@ -41,7 +41,7 @@ def test_model(self): fitting = EnerFitting(**jdata['model']['fitting_net']) # descrpt = DescrptSeA(jdata['model']['descriptor']) # fitting = EnerFitting(jdata['model']['fitting_net'], descrpt) - model = Model(jdata['model'], descrpt, fitting) + model = EnerModel(descrpt, fitting) # model._compute_dstats([test_data['coord']], [test_data['box']], [test_data['type']], [test_data['natoms_vec']], [test_data['default_mesh']]) input_data = {'coord' : [test_data['coord']], diff --git a/source/tests/test_model_se_a_srtab.py b/source/tests/test_model_se_a_srtab.py index bd3e69cf67..f1d70a93ff 100644 --- a/source/tests/test_model_se_a_srtab.py +++ b/source/tests/test_model_se_a_srtab.py @@ -6,7 +6,7 @@ from deepmd.utils.data_system import DataSystem from deepmd.descriptor import DescrptSeA from deepmd.fit import EnerFitting -from deepmd.model import Model +from deepmd.model import EnerModel from deepmd.common import j_must_have GLOBAL_ENER_FLOAT_PRECISION = tf.float64 @@ -28,7 +28,7 @@ def setUp(self) : gen_data() def test_model(self): - jfile = 'water_se_a.json' + jfile = 'water_se_a_srtab.json' jdata = j_loader(jfile) systems = j_must_have(jdata, 'systems') @@ -51,7 +51,17 @@ def test_model(self): fitting = EnerFitting(**jdata['model']['fitting_net']) # descrpt = DescrptSeA(jdata['model']['descriptor']) # fitting = EnerFitting(jdata['model']['fitting_net'], descrpt) - model = Model(jdata['model'], descrpt, fitting) + model = EnerModel( + descrpt, + fitting, + jdata['model'].get('type_map'), + jdata['model'].get('data_stat_nbatch'), + jdata['model'].get('data_stat_protect'), + jdata['model'].get('use_srtab'), + jdata['model'].get('smin_alpha'), + jdata['model'].get('sw_rmin'), + jdata['model'].get('sw_rmax') + ) # model._compute_dstats([test_data['coord']], [test_data['box']], [test_data['type']], [test_data['natoms_vec']], [test_data['default_mesh']]) input_data = {'coord' : [test_data['coord']], @@ -111,9 +121,9 @@ def test_model(self): f = f.reshape([-1]) v = v.reshape([-1]) - refe = [6.135449167779321300e+01] - reff = [7.799691562262310585e-02,9.423098804815030483e-02,3.790560997388224204e-03,1.432522403799846578e-01,1.148392791403983204e-01,-1.321871172563671148e-02,-7.318966526325138000e-02,6.516069212737778116e-02,5.406418483320515412e-04,5.870713761026503247e-02,-1.605402669549013672e-01,-5.089516979826595386e-03,-2.554593467731766654e-01,3.092063507347833987e-02,1.510355029451411479e-02,4.869271842355533952e-02,-1.446113274345035005e-01,-1.126524434771078789e-03] - refv = [-6.076776685178300053e-01,1.103174323630009418e-01,1.984250991380156690e-02,1.103174323630009557e-01,-3.319759402259439551e-01,-6.007404107650986258e-03,1.984250991380157036e-02,-6.007404107650981921e-03,-1.200076017439753642e-03] + refe = [1.141610882066236599e+02] + reff = [-1.493121233165248043e+02,-1.831419491743885715e+02,-8.439542992300344437e+00,-1.811987095947552859e+02,-1.476380826187439084e+02,1.264271856742560018e+01,1.544377958934875323e+02,-7.816520233903435866e+00,1.287925245463442225e+00,-4.000393268449002449e+00,1.910748885843098890e+02,7.134789955349889468e+00,1.826908441979261113e+02,3.677156386479059513e+00,-1.122312112141401741e+01,-2.617413911684622008e+00,1.438445070562470391e+02,-1.402769654524568033e+00] + refv = [3.585047655925112622e+02,-7.569252978336677984e+00,-1.068382043878426124e+01,-7.569252978336677096e+00,3.618439481685132932e+02,5.448668500896081568e+00,-1.068382043878426302e+01,5.448668500896082456e+00,1.050393462151727686e+00] refe = np.reshape(refe, [-1]) reff = np.reshape(reff, [-1]) refv = np.reshape(refv, [-1]) diff --git a/source/tests/test_model_se_r.py b/source/tests/test_model_se_r.py index f7d2fa07fe..f178dbda56 100644 --- a/source/tests/test_model_se_r.py +++ b/source/tests/test_model_se_r.py @@ -6,7 +6,7 @@ from deepmd.utils.data_system import DataSystem from deepmd.descriptor import DescrptSeR from deepmd.fit import EnerFitting -from deepmd.model import Model +from deepmd.model import EnerModel from deepmd.common import j_must_have GLOBAL_ENER_FLOAT_PRECISION = tf.float64 @@ -40,7 +40,7 @@ def test_model(self): jdata['model']['fitting_net']['descrpt'] = descrpt fitting = EnerFitting(**jdata['model']['fitting_net']) # fitting = EnerFitting(jdata['model']['fitting_net'], descrpt) - model = Model(jdata['model'], descrpt, fitting) + model = EnerModel(descrpt, fitting) # model._compute_dstats([test_data['coord']], [test_data['box']], [test_data['type']], [test_data['natoms_vec']], [test_data['default_mesh']]) input_data = {'coord' : [test_data['coord']], From 5ba6456cefb6d043b6af800cf6c7214fb7a9c2be Mon Sep 17 00:00:00 2001 From: Han Wang Date: Sat, 27 Mar 2021 09:54:42 +0800 Subject: [PATCH 4/5] fix bug of gen tab.xvg --- source/tests/test_model_se_a_srtab.py | 4 ++++ source/tests/test_tab_nonsmth.py | 3 +++ source/tests/test_tab_smooth.py | 3 +++ 3 files changed, 10 insertions(+) diff --git a/source/tests/test_model_se_a_srtab.py b/source/tests/test_model_se_a_srtab.py index f1d70a93ff..7ef371ae3b 100644 --- a/source/tests/test_model_se_a_srtab.py +++ b/source/tests/test_model_se_a_srtab.py @@ -26,6 +26,10 @@ def _make_tab(ntype) : class TestModel(unittest.TestCase): def setUp(self) : gen_data() + _make_tab(2) + + def tearDown(self): + os.remove('tab.xvg') def test_model(self): jfile = 'water_se_a_srtab.json' diff --git a/source/tests/test_tab_nonsmth.py b/source/tests/test_tab_nonsmth.py index 9181ba5977..f6496bbb1d 100644 --- a/source/tests/test_tab_nonsmth.py +++ b/source/tests/test_tab_nonsmth.py @@ -51,6 +51,9 @@ def setUp (self, dtype = tf.float64, trainable = False, initializer = tf.constant_initializer(tab_data)) + + def tearDown(self): + os.remove('tab.xvg') def comp_interpl_ef (self, dcoord, diff --git a/source/tests/test_tab_smooth.py b/source/tests/test_tab_smooth.py index 2917582446..ad6ad22044 100644 --- a/source/tests/test_tab_smooth.py +++ b/source/tests/test_tab_smooth.py @@ -51,6 +51,9 @@ def setUp (self, dtype = tf.float64, trainable = False, initializer = tf.constant_initializer(tab_data)) + + def tearDown(self): + os.remove('tab.xvg') def comp_ef (self, dcoord, From aeef9b9b8c82b4e9adf1ce970183f3cb4a640bfc Mon Sep 17 00:00:00 2001 From: Han Wang Date: Sat, 27 Mar 2021 20:20:31 +0800 Subject: [PATCH 5/5] fix bugs of tensor model --- deepmd/model/tensor.py | 51 ++++++++++++++++++++++++++------- source/tests/test_polar_se_a.py | 2 +- source/tests/test_wfc.py | 2 +- 3 files changed, 43 insertions(+), 12 deletions(-) diff --git a/deepmd/model/tensor.py b/deepmd/model/tensor.py index fa60171255..8df73f4ff0 100644 --- a/deepmd/model/tensor.py +++ b/deepmd/model/tensor.py @@ -13,7 +13,7 @@ def __init__ ( tensor_name : str, descrpt, fitting, - type_map : List[str], + type_map : List[str] = None, data_stat_nbatch : int = 10, data_stat_protect : float = 1e-2, )->None: @@ -44,7 +44,10 @@ def __init__ ( # fitting self.fitting = fitting # other params - self.type_map = type_map + if type_map is None: + self.type_map = [] + else: + self.type_map = type_map self.data_stat_nbatch = data_stat_nbatch self.data_stat_protect = data_stat_protect @@ -133,19 +136,47 @@ def build (self, class WFCModel(TensorModel): - def __init__(self, descrpt, fitting, tm, dsn, dsp) : - TensorModel.__init__(self, 'wfc', descrpt, fitting, tm, dsn, dsp) + def __init__( + self, + descrpt, + fitting, + type_map : List[str] = None, + data_stat_nbatch : int = 10, + data_stat_protect : float = 1e-2 + ) -> None: + TensorModel.__init__(self, 'wfc', descrpt, fitting, type_map, data_stat_nbatch, data_stat_protect) class DipoleModel(TensorModel): - def __init__(self, descrpt, fitting, tm, dsn, dsp) : - TensorModel.__init__(self, 'dipole', descrpt, fitting, tm, dsn, dsp) + def __init__( + self, + descrpt, + fitting, + type_map : List[str] = None, + data_stat_nbatch : int = 10, + data_stat_protect : float = 1e-2 + ) -> None: + TensorModel.__init__(self, 'dipole', descrpt, fitting, type_map, data_stat_nbatch, data_stat_protect) class PolarModel(TensorModel): - def __init__(self, descrpt, fitting, tm, dsn, dsp) : - TensorModel.__init__(self, 'polar', descrpt, fitting, tm, dsn, dsp) + def __init__( + self, + descrpt, + fitting, + type_map : List[str] = None, + data_stat_nbatch : int = 10, + data_stat_protect : float = 1e-2 + ) -> None: + TensorModel.__init__(self, 'polar', descrpt, fitting, type_map, data_stat_nbatch, data_stat_protect) class GlobalPolarModel(TensorModel): - def __init__(self, descrpt, fitting, tm, dsn, dsp) : - TensorModel.__init__(self, 'global_polar', descrpt, fitting, tm, dsn, dsp) + def __init__( + self, + descrpt, + fitting, + type_map : List[str] = None, + data_stat_nbatch : int = 10, + data_stat_protect : float = 1e-2 + ) -> None: + TensorModel.__init__(self, 'global_polar', descrpt, fitting, type_map, data_stat_nbatch, data_stat_protect) diff --git a/source/tests/test_polar_se_a.py b/source/tests/test_polar_se_a.py index 1c265057a4..8ca358b585 100644 --- a/source/tests/test_polar_se_a.py +++ b/source/tests/test_polar_se_a.py @@ -40,7 +40,7 @@ def test_model(self): descrpt = DescrptSeA(**jdata['model']['descriptor']) jdata['model']['fitting_net']['descrpt'] = descrpt fitting = PolarFittingSeA(**jdata['model']['fitting_net']) - model = PolarModel(jdata['model'], descrpt, fitting) + model = PolarModel(descrpt, fitting) # model._compute_dstats([test_data['coord']], [test_data['box']], [test_data['type']], [test_data['natoms_vec']], [test_data['default_mesh']]) input_data = {'coord' : [test_data['coord']], diff --git a/source/tests/test_wfc.py b/source/tests/test_wfc.py index 9bae04ed04..9038c6d177 100644 --- a/source/tests/test_wfc.py +++ b/source/tests/test_wfc.py @@ -39,7 +39,7 @@ def test_model(self): jdata['model']['descriptor'].pop('_comment', None) descrpt = DescrptLocFrame(**jdata['model']['descriptor']) fitting = WFCFitting(jdata['model']['fitting_net'], descrpt) - model = WFCModel(jdata['model'], descrpt, fitting) + model = WFCModel(descrpt, fitting) input_data = {'coord' : [test_data['coord']], 'box': [test_data['box']],