From 925f351a3157d32e05778156b7a0d8d490d71cf0 Mon Sep 17 00:00:00 2001 From: nahso Date: Mon, 25 Sep 2023 13:59:37 +0800 Subject: [PATCH] ut for se_e2_a_ebd_v2 model --- source/tests/test_model_se_a_ebd_v2.py | 159 +++++++++++++++++++++++++ 1 file changed, 159 insertions(+) create mode 100644 source/tests/test_model_se_a_ebd_v2.py diff --git a/source/tests/test_model_se_a_ebd_v2.py b/source/tests/test_model_se_a_ebd_v2.py new file mode 100644 index 0000000000..e91a0dbf1f --- /dev/null +++ b/source/tests/test_model_se_a_ebd_v2.py @@ -0,0 +1,159 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +import numpy as np +from common import ( + DataSystem, + gen_data, + j_loader, +) + +from deepmd.common import ( + j_must_have, +) +from deepmd.descriptor.se_a_ebd_v2 import ( + DescrptSeAEbdV2, +) +from deepmd.env import ( + tf, +) +from deepmd.fit import ( + EnerFitting, +) +from deepmd.model import ( + EnerModel, +) +from deepmd.utils.type_embed import ( + TypeEmbedNet, +) + +GLOBAL_ENER_FLOAT_PRECISION = tf.float64 +GLOBAL_TF_FLOAT_PRECISION = tf.float64 +GLOBAL_NP_FLOAT_PRECISION = np.float64 + + +class TestModel(tf.test.TestCase): + def setUp(self): + gen_data() + + def test_model(self): + jfile = "water_se_a_ebd.json" + jdata = j_loader(jfile) + + 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 + + jdata["model"]["descriptor"].pop("type", None) + jdata["model"]["type_embedding"] = {} + jdata["model"]["type_embedding"]["neuron"] = [1] + jdata["model"]["type_embedding"]["resnet_dt"] = False + jdata["model"]["type_embedding"]["seed"] = 1 + typeebd_param = jdata["model"]["type_embedding"] + typeebd = TypeEmbedNet( + neuron=typeebd_param["neuron"], + activation_function=None, + resnet_dt=typeebd_param["resnet_dt"], + seed=typeebd_param["seed"], + uniform_seed=True, + padding=True, + ) + descrpt = DescrptSeAEbdV2( + **jdata["model"]["descriptor"], + ) + jdata["model"]["fitting_net"]["descrpt"] = descrpt + fitting = EnerFitting( + **jdata["model"]["fitting_net"], + ) + # fitting = EnerFitting(jdata['model']['fitting_net'], descrpt) + model = EnerModel(descrpt, fitting, typeebd) + + # 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_a_ebd_v2", + 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 = self.cached_session().__enter__() + 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 = [5.435394596262052014e-01 ] + reff = [ + 6.583728125594628944e-02,7.228993116083935744e-02,1.971543579114074483e-03,6.567474563776359853e-02,7.809421727465599983e-02,-4.866958849094786890e-03,-8.670511901715304004e-02,3.525374157021862048e-02,1.415748959800727487e-03,6.375813001810648473e-02,-1.139053242798149790e-01,-4.178593754384440744e-03,-1.471737787218250215e-01,4.189712704724830872e-02,7.011731363309440038e-03,3.860874082716164030e-02,-1.136296927731473005e-01,-1.353471298745012206e-03 + ] + refv = [ + -4.243979601186427253e-01,1.097173849143971286e-01,1.227299373463585502e-02,1.097173849143970314e-01,-2.462891443164323124e-01,-5.711664180530139426e-03,1.227299373463585502e-02,-5.711664180530143763e-03,-6.217348853341628408e-04 + ] + 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)