Beispiel #1
0
    def test_model(self):
        jfile = 'water.json'
        with open(jfile) as fp:
            jdata = json.load(fp)
        run_opt = RunOptions(None)
        systems = j_must_have(jdata, 'systems')
        set_pfx = j_must_have(jdata, 'set_prefix')
        batch_size = j_must_have(jdata, 'batch_size')
        test_size = j_must_have(jdata, 'numb_test')
        batch_size = 1
        test_size = 1
        stop_batch = j_must_have(jdata, 'stop_batch')
        rcut = j_must_have(jdata['model']['descriptor'], 'rcut')

        data = DataSystem(systems,
                          set_pfx,
                          batch_size,
                          test_size,
                          rcut,
                          run_opt=None)

        test_data = data.get_test()
        numb_test = 1

        descrpt = DescrptLocFrame(jdata['model']['descriptor'])
        fitting = EnerFitting(jdata['model']['fitting_net'], descrpt)
        model = Model(jdata['model'], 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']],
            'box': [test_data['box']],
            'type': [test_data['type']],
            'natoms_vec': [test_data['natoms_vec']],
            'default_mesh': [test_data['default_mesh']]
        }
        model._compute_input_stat(input_data)
        model.fitting.bias_atom_e = data.compute_energy_shift()

        t_prop_c = tf.placeholder(tf.float32, [5], name='t_prop_c')
        t_energy = tf.placeholder(global_ener_float_precision, [None],
                                  name='t_energy')
        t_force = tf.placeholder(global_tf_float_precision, [None],
                                 name='t_force')
        t_virial = tf.placeholder(global_tf_float_precision, [None],
                                  name='t_virial')
        t_atom_ener = tf.placeholder(global_tf_float_precision, [None],
                                     name='t_atom_ener')
        t_coord = tf.placeholder(global_tf_float_precision, [None],
                                 name='i_coord')
        t_type = tf.placeholder(tf.int32, [None], name='i_type')
        t_natoms = tf.placeholder(tf.int32, [model.ntypes + 2],
                                  name='i_natoms')
        t_box = tf.placeholder(global_tf_float_precision, [None, 9],
                               name='i_box')
        t_mesh = tf.placeholder(tf.int32, [None], name='i_mesh')
        is_training = tf.placeholder(tf.bool)
        t_fparam = None

        model_pred \
            = model.build (t_coord,
                           t_type,
                           t_natoms,
                           t_box,
                           t_mesh,
                           t_fparam,
                           suffix = "loc_frame",
                           reuse = False)
        energy = model_pred['energy']
        force = model_pred['force']
        virial = model_pred['virial']
        atom_ener = model_pred['atom_ener']

        feed_dict_test = {
            t_prop_c: test_data['prop_c'],
            t_energy: test_data['energy'][:numb_test],
            t_force: np.reshape(test_data['force'][:numb_test, :], [-1]),
            t_virial: np.reshape(test_data['virial'][:numb_test, :], [-1]),
            t_atom_ener: np.reshape(test_data['atom_ener'][:numb_test, :],
                                    [-1]),
            t_coord: np.reshape(test_data['coord'][:numb_test, :], [-1]),
            t_box: test_data['box'][:numb_test, :],
            t_type: np.reshape(test_data['type'][:numb_test, :], [-1]),
            t_natoms: test_data['natoms_vec'],
            t_mesh: test_data['default_mesh'],
            is_training: False
        }

        sess = tf.Session()
        sess.run(tf.global_variables_initializer())
        [e, f, v] = sess.run([energy, force, virial], feed_dict=feed_dict_test)

        e = e.reshape([-1])
        f = f.reshape([-1])
        v = v.reshape([-1])
        refe = [1.165945032784766511e+01]
        reff = [
            2.356319331246305437e-01, 1.772322096063349284e-01,
            1.455439548950788684e-02, 1.968599426000810226e-01,
            2.648214484898352983e-01, 7.595232354012236564e-02,
            -2.121321856338151401e-01, -2.463886119018566037e-03,
            -2.075636300914874069e-02, -9.360310077571798101e-03,
            -1.751965198776750943e-01, -2.046405309983102827e-02,
            -1.990194093283037535e-01, -1.828347741191920298e-02,
            -6.916374506995154325e-02, -1.197997068502068031e-02,
            -2.461097746875573200e-01, 1.987744214930105627e-02
        ]
        refv = [
            -4.998509978510510265e-01, -1.966169437179327711e-02,
            1.136130543869883977e-02, -1.966169437179334650e-02,
            -4.575353297894450555e-01, -2.668666556859019493e-03,
            1.136130543869887100e-02, -2.668666556859039876e-03,
            2.455466940358383508e-03
        ]
        refe = np.reshape(refe, [-1])
        reff = np.reshape(reff, [-1])
        refv = np.reshape(refv, [-1])

        places = 10
        for ii in range(e.size):
            self.assertAlmostEqual(e[ii], refe[ii], places=places)
        for ii in range(f.size):
            self.assertAlmostEqual(f[ii], reff[ii], places=places)
        for ii in range(v.size):
            self.assertAlmostEqual(v[ii], refv[ii], places=places)
Beispiel #2
0
    def test_model(self):
        jfile = 'water_se_a.json'
        with open(jfile) as fp:
            jdata = json.load(fp)
        run_opt = RunOptions(None)
        systems = j_must_have(jdata, 'systems')
        set_pfx = j_must_have(jdata, 'set_prefix')
        batch_size = j_must_have(jdata, 'batch_size')
        test_size = j_must_have(jdata, 'numb_test')
        batch_size = 1
        test_size = 1
        stop_batch = j_must_have(jdata, 'stop_batch')
        rcut = j_must_have(jdata['model']['descriptor'], 'rcut')

        data = DataSystem(systems,
                          set_pfx,
                          batch_size,
                          test_size,
                          rcut,
                          run_opt=None)

        test_data = data.get_test()
        numb_test = 1

        descrpt = DescrptSeA(jdata['model']['descriptor'])
        fitting = EnerFitting(jdata['model']['fitting_net'], descrpt)
        model = Model(jdata['model'], 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']],
            'box': [test_data['box']],
            'type': [test_data['type']],
            'natoms_vec': [test_data['natoms_vec']],
            'default_mesh': [test_data['default_mesh']]
        }
        model._compute_dstats(input_data)
        model.bias_atom_e = data.compute_energy_shift()

        t_prop_c = tf.placeholder(tf.float32, [5], name='t_prop_c')
        t_energy = tf.placeholder(global_ener_float_precision, [None],
                                  name='t_energy')
        t_force = tf.placeholder(global_tf_float_precision, [None],
                                 name='t_force')
        t_virial = tf.placeholder(global_tf_float_precision, [None],
                                  name='t_virial')
        t_atom_ener = tf.placeholder(global_tf_float_precision, [None],
                                     name='t_atom_ener')
        t_coord = tf.placeholder(global_tf_float_precision, [None],
                                 name='i_coord')
        t_type = tf.placeholder(tf.int32, [None], name='i_type')
        t_natoms = tf.placeholder(tf.int32, [model.ntypes + 2],
                                  name='i_natoms')
        t_box = tf.placeholder(global_tf_float_precision, [None, 9],
                               name='i_box')
        t_mesh = tf.placeholder(tf.int32, [None], name='i_mesh')
        is_training = tf.placeholder(tf.bool)
        t_fparam = None

        model_pred \
            = model.build (t_coord,
                           t_type,
                           t_natoms,
                           t_box,
                           t_mesh,
                           t_fparam,
                           suffix = "se_a",
                           reuse = False)
        energy = model_pred['energy']
        force = model_pred['force']
        virial = model_pred['virial']
        atom_ener = model_pred['atom_ener']

        feed_dict_test = {
            t_prop_c: test_data['prop_c'],
            t_energy: test_data['energy'][:numb_test],
            t_force: np.reshape(test_data['force'][:numb_test, :], [-1]),
            t_virial: np.reshape(test_data['virial'][:numb_test, :], [-1]),
            t_atom_ener: np.reshape(test_data['atom_ener'][:numb_test, :],
                                    [-1]),
            t_coord: np.reshape(test_data['coord'][:numb_test, :], [-1]),
            t_box: test_data['box'][:numb_test, :],
            t_type: np.reshape(test_data['type'][:numb_test, :], [-1]),
            t_natoms: test_data['natoms_vec'],
            t_mesh: test_data['default_mesh'],
            is_training: False
        }

        sess = tf.Session()
        sess.run(tf.global_variables_initializer())
        [e, f, v] = sess.run([energy, force, virial], feed_dict=feed_dict_test)

        e = e.reshape([-1])
        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 = np.reshape(refe, [-1])
        reff = np.reshape(reff, [-1])
        refv = np.reshape(refv, [-1])

        places = 10
        for ii in range(e.size):
            self.assertAlmostEqual(e[ii], refe[ii], places=places)
        for ii in range(f.size):
            self.assertAlmostEqual(f[ii], reff[ii], places=places)
        for ii in range(v.size):
            self.assertAlmostEqual(v[ii], refv[ii], places=places)
Beispiel #3
0
    def test_model(self):
        jfile = 'water_se_r.json'
        with open(jfile) as fp:
            jdata = json.load (fp)
        run_opt = RunOptions(None) 
        systems = j_must_have(jdata, 'systems')
        set_pfx = j_must_have(jdata, 'set_prefix')
        batch_size = j_must_have(jdata, 'batch_size')
        test_size = j_must_have(jdata, 'numb_test')
        batch_size = 1
        test_size = 1
        stop_batch = j_must_have(jdata, 'stop_batch')
        rcut = j_must_have (jdata['model']['descriptor'], 'rcut')
        
        data = DataSystem(systems, set_pfx, batch_size, test_size, rcut, run_opt = None)
        
        test_data = data.get_test ()
        numb_test = 1
        
        descrpt = DescrptSeR(jdata['model']['descriptor'])
        fitting = EnerFitting(jdata['model']['fitting_net'], descrpt)
        model = Model(jdata['model'], 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']], 
                      'box': [test_data['box']], 
                      'type': [test_data['type']],
                      'natoms_vec' : [test_data['natoms_vec']],
                      'default_mesh' : [test_data['default_mesh']]
        }
        model._compute_input_stat(input_data)
        model.descrpt.bias_atom_e = data.compute_energy_shift()

        t_prop_c           = tf.placeholder(tf.float32, [5],    name='t_prop_c')
        t_energy           = tf.placeholder(global_ener_float_precision, [None], name='t_energy')
        t_force            = tf.placeholder(global_tf_float_precision, [None], name='t_force')
        t_virial           = tf.placeholder(global_tf_float_precision, [None], name='t_virial')
        t_atom_ener        = tf.placeholder(global_tf_float_precision, [None], name='t_atom_ener')
        t_coord            = tf.placeholder(global_tf_float_precision, [None], name='i_coord')
        t_type             = tf.placeholder(tf.int32,   [None], name='i_type')
        t_natoms           = tf.placeholder(tf.int32,   [model.ntypes+2], name='i_natoms')
        t_box              = tf.placeholder(global_tf_float_precision, [None, 9], name='i_box')
        t_mesh             = tf.placeholder(tf.int32,   [None], name='i_mesh')
        is_training        = tf.placeholder(tf.bool)
        t_fparam = None

        model_pred\
            = model.build (t_coord, 
                           t_type, 
                           t_natoms, 
                           t_box, 
                           t_mesh,
                           t_fparam,
                           suffix = "se_r", 
                           reuse = False)
        energy = model_pred['energy']
        force  = model_pred['force']
        virial = model_pred['virial']
        atom_ener =  model_pred['atom_ener']

        feed_dict_test = {t_prop_c:        test_data['prop_c'],
                          t_energy:        test_data['energy']              [:numb_test],
                          t_force:         np.reshape(test_data['force']    [:numb_test, :], [-1]),
                          t_virial:        np.reshape(test_data['virial']   [:numb_test, :], [-1]),
                          t_atom_ener:     np.reshape(test_data['atom_ener'][:numb_test, :], [-1]),
                          t_coord:         np.reshape(test_data['coord']    [:numb_test, :], [-1]),
                          t_box:           test_data['box']                 [:numb_test, :],
                          t_type:          np.reshape(test_data['type']     [:numb_test, :], [-1]),
                          t_natoms:        test_data['natoms_vec'],
                          t_mesh:          test_data['default_mesh'],
                          is_training:     False}

        sess = tf.Session()
        sess.run(tf.global_variables_initializer())
        [e, f, v] = sess.run([energy, force, virial], 
                             feed_dict = feed_dict_test)

        e = e.reshape([-1])
        f = f.reshape([-1])
        v = v.reshape([-1])
        refe = [6.152085988309423925e+01]
        reff = [-1.714443151616400110e-04,-1.315836609370952051e-04,-5.584120460897444674e-06,-7.197863450669731334e-05,-1.384609799994930676e-04,8.856091902774708468e-06,1.120578238869146797e-04,-7.428703645877488470e-05,9.370560731488587317e-07,-1.048347129617610465e-04,1.977876923815685781e-04,7.522050342771599598e-06,2.361772659657814205e-04,-5.774651813388292487e-05,-1.233143271630744828e-05,2.257277740226381951e-08,2.042905031476775584e-04,6.003548585097267914e-07]
        refv = [1.035180911513190792e-03,-1.118982949050497126e-04,-2.383287813436022850e-05,-1.118982949050497126e-04,4.362023915782403281e-04,8.119543218224559240e-06,-2.383287813436022850e-05,8.119543218224559240e-06,1.201142938802945237e-06]
        refe = np.reshape(refe, [-1])
        reff = np.reshape(reff, [-1])
        refv = np.reshape(refv, [-1])

        places = 6
        for ii in range(e.size) :
            self.assertAlmostEqual(e[ii], refe[ii], places = places)
        for ii in range(f.size) :
            self.assertAlmostEqual(f[ii], reff[ii], places = places)
        for ii in range(v.size) :
            self.assertAlmostEqual(v[ii], refv[ii], places = places)
Beispiel #4
0
    def _init_param(self, jdata):
        # model config
        model_param = j_must_have(jdata, 'model')
        descrpt_param = j_must_have(model_param, 'descriptor')
        fitting_param = j_must_have(model_param, 'fitting_net')

        # descriptor
        descrpt_type = j_must_have(descrpt_param, 'type')
        if descrpt_type == 'loc_frame':
            self.descrpt = DescrptLocFrame(descrpt_param)
        elif descrpt_type == 'se_a':
            self.descrpt = DescrptSeA(descrpt_param)
        elif descrpt_type == 'se_r':
            self.descrpt = DescrptSeR(descrpt_param)
        elif descrpt_type == 'se_ar':
            self.descrpt = DescrptSeAR(descrpt_param)
        else:
            raise RuntimeError('unknow model type ' + descrpt_type)

        # fitting net
        try:
            fitting_type = fitting_param['type']
        except:
            fitting_type = 'ener'
        if fitting_type == 'ener':
            self.fitting = EnerFitting(fitting_param, self.descrpt)
        elif fitting_type == 'wfc':
            self.fitting = WFCFitting(fitting_param, self.descrpt)
        elif fitting_type == 'dipole':
            if descrpt_type == 'se_a':
                self.fitting = DipoleFittingSeA(fitting_param, self.descrpt)
            else:
                raise RuntimeError(
                    'fitting dipole only supports descrptors: se_a')
        elif fitting_type == 'polar':
            if descrpt_type == 'loc_frame':
                self.fitting = PolarFittingLocFrame(fitting_param,
                                                    self.descrpt)
            elif descrpt_type == 'se_a':
                self.fitting = PolarFittingSeA(fitting_param, self.descrpt)
            else:
                raise RuntimeError(
                    'fitting polar only supports descrptors: loc_frame and se_a'
                )
        elif fitting_type == 'global_polar':
            if descrpt_type == 'se_a':
                self.fitting = GlobalPolarFittingSeA(fitting_param,
                                                     self.descrpt)
            else:
                raise RuntimeError(
                    'fitting global_polar only supports descrptors: loc_frame and se_a'
                )
        else:
            raise RuntimeError('unknow fitting type ' + fitting_type)

        # init model
        # infer model type by fitting_type
        if fitting_type == Model.model_type:
            self.model = Model(model_param, self.descrpt, self.fitting)
        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)
        elif fitting_type == 'polar':
            self.model = PolarModel(model_param, self.descrpt, self.fitting)
        elif fitting_type == 'global_polar':
            self.model = GlobalPolarModel(model_param, self.descrpt,
                                          self.fitting)
        else:
            raise RuntimeError('get unknown fitting type when building model')

        # learning rate
        lr_param = j_must_have(jdata, 'learning_rate')
        try:
            lr_type = lr_param['type']
        except:
            lr_type = 'exp'
        if lr_type == 'exp':
            self.lr = LearningRateExp(lr_param)
        else:
            raise RuntimeError('unknown learning_rate type ' + lr_type)

        # loss
        # infer loss type by fitting_type
        try:
            loss_param = jdata['loss']
            loss_type = loss_param.get('type', 'std')
        except:
            loss_param = None
            loss_type = 'std'

        if fitting_type == 'ener':
            if loss_type == 'std':
                self.loss = EnerStdLoss(
                    loss_param, starter_learning_rate=self.lr.start_lr())
            elif loss_type == 'ener_dipole':
                self.loss = EnerDipoleLoss(
                    loss_param, starter_learning_rate=self.lr.start_lr())
            else:
                raise RuntimeError('unknow loss type')
        elif fitting_type == 'wfc':
            self.loss = TensorLoss(loss_param,
                                   model=self.model,
                                   tensor_name='wfc',
                                   tensor_size=self.model.get_out_size(),
                                   label_name='wfc')
        elif fitting_type == 'dipole':
            self.loss = TensorLoss(loss_param,
                                   model=self.model,
                                   tensor_name='dipole',
                                   tensor_size=3,
                                   label_name='dipole')
        elif fitting_type == 'polar':
            self.loss = TensorLoss(loss_param,
                                   model=self.model,
                                   tensor_name='polar',
                                   tensor_size=9,
                                   label_name='polarizability')
        elif fitting_type == 'global_polar':
            self.loss = TensorLoss(loss_param,
                                   model=self.model,
                                   tensor_name='global_polar',
                                   tensor_size=9,
                                   atomic=False,
                                   label_name='polarizability')
        else:
            raise RuntimeError(
                'get unknown fitting type when building loss function')

        # training
        training_param = j_must_have(jdata, 'training')

        tr_args = ClassArg()\
                  .add('numb_test',     int,    default = 1)\
                  .add('disp_file',     str,    default = 'lcurve.out')\
                  .add('disp_freq',     int,    default = 100)\
                  .add('save_freq',     int,    default = 1000)\
                  .add('save_ckpt',     str,    default = 'model.ckpt')\
                  .add('display_in_training', bool, default = True)\
                  .add('timing_in_training',  bool, default = True)\
                  .add('profiling',     bool,   default = False)\
                  .add('profiling_file',str,    default = 'timeline.json')\
                  .add('sys_probs',   list    )\
                  .add('auto_prob_style', str, default = "prob_sys_size")
        tr_data = tr_args.parse(training_param)
        self.numb_test = tr_data['numb_test']
        self.disp_file = tr_data['disp_file']
        self.disp_freq = tr_data['disp_freq']
        self.save_freq = tr_data['save_freq']
        self.save_ckpt = tr_data['save_ckpt']
        self.display_in_training = tr_data['display_in_training']
        self.timing_in_training = tr_data['timing_in_training']
        self.profiling = tr_data['profiling']
        self.profiling_file = tr_data['profiling_file']
        self.sys_probs = tr_data['sys_probs']
        self.auto_prob_style = tr_data['auto_prob_style']
        self.useBN = False
        if fitting_type == 'ener' and self.fitting.get_numb_fparam() > 0:
            self.numb_fparam = self.fitting.get_numb_fparam()
        else:
            self.numb_fparam = 0
    def test_model(self):
        jfile = 'water_se_a_aparam.json'
        with open(jfile) as fp:
            jdata = json.load(fp)
        run_opt = RunOptions(None)
        systems = j_must_have(jdata, 'systems')
        set_pfx = j_must_have(jdata, 'set_prefix')
        batch_size = j_must_have(jdata, 'batch_size')
        test_size = j_must_have(jdata, 'numb_test')
        batch_size = 1
        test_size = 1
        stop_batch = j_must_have(jdata, 'stop_batch')
        rcut = j_must_have(jdata['model']['descriptor'], 'rcut')

        data = DataSystem(systems,
                          set_pfx,
                          batch_size,
                          test_size,
                          rcut,
                          run_opt=None)

        test_data = data.get_test()
        # manually set aparam
        test_data['aparam'] = np.load('system/set.000/aparam.npy')
        numb_test = 1

        descrpt = DescrptSeA(jdata['model']['descriptor'])
        fitting = EnerFitting(jdata['model']['fitting_net'], descrpt)
        model = Model(jdata['model'], 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']],
            'box': [test_data['box']],
            'type': [test_data['type']],
            'natoms_vec': [test_data['natoms_vec']],
            'default_mesh': [test_data['default_mesh']],
            'aparam': [test_data['aparam']],
        }
        model._compute_dstats(input_data)
        model.bias_atom_e = data.compute_energy_shift()

        t_prop_c = tf.placeholder(tf.float32, [5], name='t_prop_c')
        t_energy = tf.placeholder(global_ener_float_precision, [None],
                                  name='t_energy')
        t_force = tf.placeholder(global_tf_float_precision, [None],
                                 name='t_force')
        t_virial = tf.placeholder(global_tf_float_precision, [None],
                                  name='t_virial')
        t_atom_ener = tf.placeholder(global_tf_float_precision, [None],
                                     name='t_atom_ener')
        t_coord = tf.placeholder(global_tf_float_precision, [None],
                                 name='i_coord')
        t_type = tf.placeholder(tf.int32, [None], name='i_type')
        t_natoms = tf.placeholder(tf.int32, [model.ntypes + 2],
                                  name='i_natoms')
        t_box = tf.placeholder(global_tf_float_precision, [None, 9],
                               name='i_box')
        t_mesh = tf.placeholder(tf.int32, [None], name='i_mesh')
        t_aparam = tf.placeholder(global_tf_float_precision, [None],
                                  name='i_aparam')
        is_training = tf.placeholder(tf.bool)
        input_dict = {}
        input_dict['aparam'] = t_aparam

        model_pred\
            = model.build (t_coord,
                           t_type,
                           t_natoms,
                           t_box,
                           t_mesh,
                           input_dict,
                           suffix = "se_a_aparam",
                           reuse = False)
        energy = model_pred['energy']
        force = model_pred['force']
        virial = model_pred['virial']
        atom_ener = model_pred['atom_ener']

        feed_dict_test = {
            t_prop_c: test_data['prop_c'],
            t_energy: test_data['energy'][:numb_test],
            t_force: np.reshape(test_data['force'][:numb_test, :], [-1]),
            t_virial: np.reshape(test_data['virial'][:numb_test, :], [-1]),
            t_atom_ener: np.reshape(test_data['atom_ener'][:numb_test, :],
                                    [-1]),
            t_coord: np.reshape(test_data['coord'][:numb_test, :], [-1]),
            t_box: test_data['box'][:numb_test, :],
            t_type: np.reshape(test_data['type'][:numb_test, :], [-1]),
            t_natoms: test_data['natoms_vec'],
            t_mesh: test_data['default_mesh'],
            t_aparam: np.reshape(test_data['aparam'][:numb_test, :], [-1]),
            is_training: False
        }

        sess = tf.Session()
        sess.run(tf.global_variables_initializer())
        [e, f, v] = sess.run([energy, force, virial], feed_dict=feed_dict_test)

        e = e.reshape([-1])
        f = f.reshape([-1])
        v = v.reshape([-1])
        refe = [61.35473702079649]
        reff = [
            7.789591210641927388e-02, 9.411176646369459609e-02,
            3.785806413688173194e-03, 1.430830954178063386e-01,
            1.146964190520970150e-01, -1.320340288927138173e-02,
            -7.308720494747594776e-02, 6.508269338140809657e-02,
            5.398739145542804643e-04, 5.863268336973800898e-02,
            -1.603409523950408699e-01, -5.083084610994957619e-03,
            -2.551569799443983988e-01, 3.087934885732580501e-02,
            1.508590526622844222e-02, 4.863249399791078065e-02,
            -1.444292753594846324e-01, -1.125098094204559241e-03
        ]
        refv = [
            -6.069498397488943819e-01, 1.101778888191114192e-01,
            1.981907430646132409e-02, 1.101778888191114608e-01,
            -3.315612988100872793e-01, -5.999739184898976799e-03,
            1.981907430646132756e-02, -5.999739184898974197e-03,
            -1.198656608172396325e-03
        ]
        refe = np.reshape(refe, [-1])
        reff = np.reshape(reff, [-1])
        refv = np.reshape(refv, [-1])

        places = 10
        for ii in range(e.size):
            self.assertAlmostEqual(e[ii], refe[ii], places=places)
        for ii in range(f.size):
            self.assertAlmostEqual(f[ii], reff[ii], places=places)
        for ii in range(v.size):
            self.assertAlmostEqual(v[ii], refv[ii], places=places)
    def test_model(self):
        jfile = 'water_se_a_fparam.json'
        with open(jfile) as fp:
            jdata = json.load(fp)
        run_opt = RunOptions(None)
        systems = j_must_have(jdata, 'systems')
        set_pfx = j_must_have(jdata, 'set_prefix')
        batch_size = j_must_have(jdata, 'batch_size')
        test_size = j_must_have(jdata, 'numb_test')
        batch_size = 1
        test_size = 1
        stop_batch = j_must_have(jdata, 'stop_batch')
        rcut = j_must_have(jdata['model']['descriptor'], 'rcut')

        data = DataSystem(systems,
                          set_pfx,
                          batch_size,
                          test_size,
                          rcut,
                          run_opt=None)

        test_data = data.get_test()
        numb_test = 1

        descrpt = DescrptSeA(jdata['model']['descriptor'])
        fitting = EnerFitting(jdata['model']['fitting_net'], descrpt)
        model = Model(jdata['model'], descrpt, fitting)

        model._compute_dstats([test_data['coord']], [test_data['box']],
                              [test_data['type']], [test_data['natoms_vec']],
                              [test_data['default_mesh']])
        model.bias_atom_e = data.compute_energy_shift()

        t_prop_c = tf.placeholder(tf.float32, [5], name='t_prop_c')
        t_energy = tf.placeholder(global_ener_float_precision, [None],
                                  name='t_energy')
        t_force = tf.placeholder(global_tf_float_precision, [None],
                                 name='t_force')
        t_virial = tf.placeholder(global_tf_float_precision, [None],
                                  name='t_virial')
        t_atom_ener = tf.placeholder(global_tf_float_precision, [None],
                                     name='t_atom_ener')
        t_coord = tf.placeholder(global_tf_float_precision, [None],
                                 name='i_coord')
        t_type = tf.placeholder(tf.int32, [None], name='i_type')
        t_natoms = tf.placeholder(tf.int32, [model.ntypes + 2],
                                  name='i_natoms')
        t_box = tf.placeholder(global_tf_float_precision, [None, 9],
                               name='i_box')
        t_mesh = tf.placeholder(tf.int32, [None], name='i_mesh')
        t_fparam = tf.placeholder(global_tf_float_precision, [None],
                                  name='i_fparam')
        is_training = tf.placeholder(tf.bool)
        input_dict = {}
        input_dict['fparam'] = t_fparam

        model_pred\
            = model.build (t_coord,
                           t_type,
                           t_natoms,
                           t_box,
                           t_mesh,
                           input_dict,
                           suffix = "se_a_fparam",
                           reuse = False)
        energy = model_pred['energy']
        force = model_pred['force']
        virial = model_pred['virial']
        atom_ener = model_pred['atom_ener']

        feed_dict_test = {
            t_prop_c: test_data['prop_c'],
            t_energy: test_data['energy'][:numb_test],
            t_force: np.reshape(test_data['force'][:numb_test, :], [-1]),
            t_virial: np.reshape(test_data['virial'][:numb_test, :], [-1]),
            t_atom_ener: np.reshape(test_data['atom_ener'][:numb_test, :],
                                    [-1]),
            t_coord: np.reshape(test_data['coord'][:numb_test, :], [-1]),
            t_box: test_data['box'][:numb_test, :],
            t_type: np.reshape(test_data['type'][:numb_test, :], [-1]),
            t_natoms: test_data['natoms_vec'],
            t_mesh: test_data['default_mesh'],
            t_fparam: np.reshape(test_data['fparam'][:numb_test, :], [-1]),
            is_training: False
        }

        sess = tf.Session()
        sess.run(tf.global_variables_initializer())
        [e, f, v] = sess.run([energy, force, virial], feed_dict=feed_dict_test)

        e = e.reshape([-1])
        f = f.reshape([-1])
        v = v.reshape([-1])
        refe = [6.135136929183754972e+01]
        reff = [
            7.761477777656561328e-02, 9.383013575207051205e-02,
            3.776776376267230399e-03, 1.428268971463224069e-01,
            1.143858253900619654e-01, -1.318441687719179231e-02,
            -7.271897092708884403e-02, 6.494907553857684479e-02,
            5.355599592111062821e-04, 5.840910251709752199e-02,
            -1.599042555763417750e-01, -5.067165555590445389e-03,
            -2.546246315216804113e-01, 3.073296814647456451e-02,
            1.505994759166155023e-02, 4.849282500878367153e-02,
            -1.439937492508420736e-01, -1.120701494357654411e-03
        ]
        refv = [
            -6.054303146013112480e-01, 1.097859194719944115e-01,
            1.977605183964963390e-02, 1.097859194719943976e-01,
            -3.306167096812382966e-01, -5.978855662865613894e-03,
            1.977605183964964083e-02, -5.978855662865616497e-03,
            -1.196331922996723236e-03
        ]
        refe = np.reshape(refe, [-1])
        reff = np.reshape(reff, [-1])
        refv = np.reshape(refv, [-1])

        places = 10
        for ii in range(e.size):
            self.assertAlmostEqual(e[ii], refe[ii], places=places)
        for ii in range(f.size):
            self.assertAlmostEqual(f[ii], reff[ii], places=places)
        for ii in range(v.size):
            self.assertAlmostEqual(v[ii], refv[ii], places=places)