ut for se_e2_a_ebd_v2 model

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)