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 51% rename from source/train/model.py rename to deepmd/model/ener.py index 8ccee419bf..56f4880840 100644 --- a/source/train/model.py +++ b/deepmd/model/ener.py @@ -1,101 +1,73 @@ 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] = None, + data_stat_nbatch : int = 10, + data_stat_protect : float = 1e-2, + use_srtab : str = None, + smin_alpha : float = None, + sw_rmin : float = None, + sw_rmax : float = None + ) -> None: + """ + Constructor + + 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 + 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 @@ -265,126 +237,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..8df73f4ff0 --- /dev/null +++ b/deepmd/model/tensor.py @@ -0,0 +1,182 @@ +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] = None, + 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 + if type_map is None: + self.type_map = [] + else: + self.type_map = type_map + self.data_stat_nbatch = data_stat_nbatch + self.data_stat_protect = data_stat_protect + + def get_rcut (self) : + return self.rcut + + def get_ntypes (self) : + return self.ntypes + + def get_type_map (self) : + return self.type_map + + def get_sel_type(self): + return self.fitting.get_sel_type() + + def get_out_size (self) : + return self.fitting.get_out_size() + + def data_stat(self, data): + all_stat = make_stat_input(data, self.data_stat_nbatch, merge_sys = False) + m_all_stat = merge_sys_stat(all_stat) + self._compute_input_stat (m_all_stat, protection = self.data_stat_protect) + self._compute_output_stat(all_stat) + + def _compute_input_stat(self, all_stat, protection = 1e-2) : + self.descrpt.compute_input_stats(all_stat['coord'], + all_stat['box'], + all_stat['type'], + all_stat['natoms_vec'], + all_stat['default_mesh'], + all_stat) + if hasattr(self.fitting, 'compute_input_stats'): + self.fitting.compute_input_stats(all_stat, protection = protection) + + def _compute_output_stat (self, all_stat) : + if hasattr(self.fitting, 'compute_output_stats'): + self.fitting.compute_output_stats(all_stat) + + def build (self, + coord_, + atype_, + natoms, + box, + mesh, + input_dict, + 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, + type_map : List[str] = None, + data_stat_nbatch : int = 10, + data_stat_protect : float = 1e-2 + ) -> None: + TensorModel.__init__(self, 'wfc', descrpt, fitting, type_map, data_stat_nbatch, data_stat_protect) + +class DipoleModel(TensorModel): + def __init__( + self, + descrpt, + fitting, + type_map : List[str] = None, + data_stat_nbatch : int = 10, + data_stat_protect : float = 1e-2 + ) -> None: + TensorModel.__init__(self, 'dipole', descrpt, fitting, type_map, data_stat_nbatch, data_stat_protect) + +class PolarModel(TensorModel): + def __init__( + self, + descrpt, + fitting, + type_map : List[str] = None, + data_stat_nbatch : int = 10, + data_stat_protect : float = 1e-2 + ) -> None: + TensorModel.__init__(self, 'polar', descrpt, fitting, type_map, data_stat_nbatch, data_stat_protect) + +class GlobalPolarModel(TensorModel): + def __init__( + self, + descrpt, + fitting, + type_map : List[str] = None, + data_stat_nbatch : int = 10, + data_stat_protect : float = 1e-2 + ) -> None: + TensorModel.__init__(self, 'global_polar', descrpt, fitting, type_map, data_stat_nbatch, data_stat_protect) + + 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/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..7ef371ae3b 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 @@ -26,9 +26,13 @@ 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.json' + jfile = 'water_se_a_srtab.json' jdata = j_loader(jfile) systems = j_must_have(jdata, 'systems') @@ -51,7 +55,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 +125,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']], 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_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, 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']], 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..fd9962ec90 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.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'), + 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.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(model_param, self.descrpt, self.fitting) + self.model = PolarModel( + self.descrpt, + self.fitting, + 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(model_param, self.descrpt, self.fitting) + self.model = GlobalPolarModel( + self.descrpt, + self.fitting, + 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')