test: better structure for the dpmodel uts (#3232)

- mv all dpmodel related UTs to a separate directory
- split the large test_model_format_utils.py as tests for different
modules.

---------

Co-authored-by: Han Wang <[email protected]>

source/tests/common/dpmodel/README

@@ -0,0 +1 @@
+test deepmd-kit/source/deepmd/dpmodel

source/tests/common/dpmodel/case_single_frame_with_nlist.py

@@ -0,0 +1,32 @@
+# SPDX-License-Identifier: LGPL-3.0-or-later
+import numpy as np
+
+
+class TestCaseSingleFrameWithNlist:
+    def setUp(self):
+        # nloc == 3, nall == 4
+        self.nloc = 3
+        self.nall = 4
+        self.nf, self.nt = 1, 2
+        self.coord_ext = np.array(
+            [
+                [0, 0, 0],
+                [0, 1, 0],
+                [0, 0, 1],
+                [0, -2, 0],
+            ],
+            dtype=np.float64,
+        ).reshape([1, self.nall * 3])
+        self.atype_ext = np.array([0, 0, 1, 0], dtype=int).reshape([1, self.nall])
+        # sel = [5, 2]
+        self.sel = [5, 2]
+        self.nlist = np.array(
+            [
+                [1, 3, -1, -1, -1, 2, -1],
+                [0, -1, -1, -1, -1, 2, -1],
+                [0, 1, -1, -1, -1, -1, -1],
+            ],
+            dtype=int,
+        ).reshape([1, self.nloc, sum(self.sel)])
+        self.rcut = 0.4
+        self.rcut_smth = 2.2

source/tests/common/dpmodel/test_descriptor_se_e2_a.py

@@ -0,0 +1,35 @@
+# SPDX-License-Identifier: LGPL-3.0-or-later
+import unittest
+
+import numpy as np
+
+from deepmd.dpmodel.descriptor import (
+    DescrptSeA,
+)
+
+from .case_single_frame_with_nlist import (
+    TestCaseSingleFrameWithNlist,
+)
+
+
+class TestDescrptSeA(unittest.TestCase, TestCaseSingleFrameWithNlist):
+    def setUp(self):
+        TestCaseSingleFrameWithNlist.setUp(self)
+
+    def test_self_consistency(
+        self,
+    ):
+        rng = np.random.default_rng()
+        nf, nloc, nnei = self.nlist.shape
+        davg = rng.normal(size=(self.nt, nnei, 4))
+        dstd = rng.normal(size=(self.nt, nnei, 4))
+        dstd = 0.1 + np.abs(dstd)
+
+        em0 = DescrptSeA(self.rcut, self.rcut_smth, self.sel)
+        em0.davg = davg
+        em0.dstd = dstd
+        em1 = DescrptSeA.deserialize(em0.serialize())
+        mm0 = em0.call(self.coord_ext, self.atype_ext, self.nlist)
+        mm1 = em1.call(self.coord_ext, self.atype_ext, self.nlist)
+        for ii in [0, 1, 4]:
+            np.testing.assert_allclose(mm0[ii], mm1[ii])

source/tests/common/dpmodel/test_dp_atomic_model.py

@@ -0,0 +1,47 @@
+# SPDX-License-Identifier: LGPL-3.0-or-later
+import unittest
+
+import numpy as np
+
+from deepmd.dpmodel.descriptor import (
+    DescrptSeA,
+)
+from deepmd.dpmodel.fitting import (
+    InvarFitting,
+)
+from deepmd.dpmodel.model import (
+    DPAtomicModel,
+)
+
+from .case_single_frame_with_nlist import (
+    TestCaseSingleFrameWithNlist,
+)
+
+
+class TestDPAtomicModel(unittest.TestCase, TestCaseSingleFrameWithNlist):
+    def setUp(self):
+        TestCaseSingleFrameWithNlist.setUp(self)
+
+    def test_self_consistency(
+        self,
+    ):
+        ds = DescrptSeA(
+            self.rcut,
+            self.rcut_smth,
+            self.sel,
+        )
+        ft = InvarFitting(
+            "energy",
+            self.nt,
+            ds.get_dim_out(),
+            1,
+            distinguish_types=ds.distinguish_types(),
+        )
+        type_map = ["foo", "bar"]
+        md0 = DPAtomicModel(ds, ft, type_map=type_map)
+        md1 = DPAtomicModel.deserialize(md0.serialize())
+
+        ret0 = md0.forward_atomic(self.coord_ext, self.atype_ext, self.nlist)
+        ret1 = md1.forward_atomic(self.coord_ext, self.atype_ext, self.nlist)
+
+        np.testing.assert_allclose(ret0["energy"], ret1["energy"])

source/tests/common/dpmodel/test_dp_model.py

@@ -0,0 +1,49 @@
+# SPDX-License-Identifier: LGPL-3.0-or-later
+import unittest
+
+import numpy as np
+
+from deepmd.dpmodel.descriptor import (
+    DescrptSeA,
+)
+from deepmd.dpmodel.fitting import (
+    InvarFitting,
+)
+from deepmd.dpmodel.model import (
+    DPModel,
+)
+
+from .case_single_frame_with_nlist import (
+    TestCaseSingleFrameWithNlist,
+)
+
+
+class TestDPModel(unittest.TestCase, TestCaseSingleFrameWithNlist):
+    def setUp(self):
+        TestCaseSingleFrameWithNlist.setUp(self)
+
+    def test_self_consistency(
+        self,
+    ):
+        nf, nloc, nnei = self.nlist.shape
+        ds = DescrptSeA(
+            self.rcut,
+            self.rcut_smth,
+            self.sel,
+        )
+        ft = InvarFitting(
+            "energy",
+            self.nt,
+            ds.get_dim_out(),
+            1,
+            distinguish_types=ds.distinguish_types(),
+        )
+        type_map = ["foo", "bar"]
+        md0 = DPModel(ds, ft, type_map=type_map)
+        md1 = DPModel.deserialize(md0.serialize())
+
+        ret0 = md0.call_lower(self.coord_ext, self.atype_ext, self.nlist)
+        ret1 = md1.call_lower(self.coord_ext, self.atype_ext, self.nlist)
+
+        np.testing.assert_allclose(ret0["energy"], ret1["energy"])
+        np.testing.assert_allclose(ret0["energy_redu"], ret1["energy_redu"])

source/tests/common/dpmodel/test_env_mat.py

@@ -0,0 +1,32 @@
+# SPDX-License-Identifier: LGPL-3.0-or-later
+import unittest
+
+import numpy as np
+
+from deepmd.dpmodel.utils import (
+    EnvMat,
+)
+
+from .case_single_frame_with_nlist import (
+    TestCaseSingleFrameWithNlist,
+)
+
+
+class TestEnvMat(unittest.TestCase, TestCaseSingleFrameWithNlist):
+    def setUp(self):
+        TestCaseSingleFrameWithNlist.setUp(self)
+
+    def test_self_consistency(
+        self,
+    ):
+        rng = np.random.default_rng()
+        nf, nloc, nnei = self.nlist.shape
+        davg = rng.normal(size=(self.nt, nnei, 4))
+        dstd = rng.normal(size=(self.nt, nnei, 4))
+        dstd = 0.1 + np.abs(dstd)
+        em0 = EnvMat(self.rcut, self.rcut_smth)
+        em1 = EnvMat.deserialize(em0.serialize())
+        mm0, ww0 = em0.call(self.coord_ext, self.atype_ext, self.nlist, davg, dstd)
+        mm1, ww1 = em1.call(self.coord_ext, self.atype_ext, self.nlist, davg, dstd)
+        np.testing.assert_allclose(mm0, mm1)
+        np.testing.assert_allclose(ww0, ww1)

source/tests/common/dpmodel/test_fitting_invar_fitting.py

@@ -0,0 +1,136 @@
+# SPDX-License-Identifier: LGPL-3.0-or-later
+import itertools
+import unittest
+
+import numpy as np
+
+from deepmd.dpmodel.descriptor import (
+    DescrptSeA,
+)
+from deepmd.dpmodel.fitting import (
+    InvarFitting,
+)
+
+from .case_single_frame_with_nlist import (
+    TestCaseSingleFrameWithNlist,
+)
+
+
+class TestInvarFitting(unittest.TestCase, TestCaseSingleFrameWithNlist):
+    def setUp(self):
+        TestCaseSingleFrameWithNlist.setUp(self)
+
+    def test_self_consistency(
+        self,
+    ):
+        rng = np.random.default_rng()
+        nf, nloc, nnei = self.nlist.shape
+        ds = DescrptSeA(self.rcut, self.rcut_smth, self.sel)
+        dd = ds.call(self.coord_ext, self.atype_ext, self.nlist)
+        atype = self.atype_ext[:, :nloc]
+
+        for (
+            distinguish_types,
+            od,
+            nfp,
+            nap,
+        ) in itertools.product(
+            [True, False],
+            [1, 2],
+            [0, 3],
+            [0, 4],
+        ):
+            ifn0 = InvarFitting(
+                "energy",
+                self.nt,
+                ds.dim_out,
+                od,
+                numb_fparam=nfp,
+                numb_aparam=nap,
+                distinguish_types=distinguish_types,
+            )
+            ifn1 = InvarFitting.deserialize(ifn0.serialize())
+            if nfp > 0:
+                ifp = rng.normal(size=(self.nf, nfp))
+            else:
+                ifp = None
+            if nap > 0:
+                iap = rng.normal(size=(self.nf, self.nloc, nap))
+            else:
+                iap = None
+            ret0 = ifn0(dd[0], atype, fparam=ifp, aparam=iap)
+            ret1 = ifn1(dd[0], atype, fparam=ifp, aparam=iap)
+            np.testing.assert_allclose(ret0["energy"], ret1["energy"])
+
+    def test_self_exception(
+        self,
+    ):
+        rng = np.random.default_rng()
+        nf, nloc, nnei = self.nlist.shape
+        ds = DescrptSeA(self.rcut, self.rcut_smth, self.sel)
+        dd = ds.call(self.coord_ext, self.atype_ext, self.nlist)
+        atype = self.atype_ext[:, :nloc]
+
+        for (
+            distinguish_types,
+            od,
+            nfp,
+            nap,
+        ) in itertools.product(
+            [True, False],
+            [1, 2],
+            [0, 3],
+            [0, 4],
+        ):
+            ifn0 = InvarFitting(
+                "energy",
+                self.nt,
+                ds.dim_out,
+                od,
+                numb_fparam=nfp,
+                numb_aparam=nap,
+                distinguish_types=distinguish_types,
+            )
+
+            if nfp > 0:
+                ifp = rng.normal(size=(self.nf, nfp))
+            else:
+                ifp = None
+            if nap > 0:
+                iap = rng.normal(size=(self.nf, self.nloc, nap))
+            else:
+                iap = None
+            with self.assertRaises(ValueError) as context:
+                ret0 = ifn0(dd[0][:, :, :-2], atype, fparam=ifp, aparam=iap)
+                self.assertIn("input descriptor", context.exception)
+
+            if nfp > 0:
+                ifp = rng.normal(size=(self.nf, nfp - 1))
+                with self.assertRaises(ValueError) as context:
+                    ret0 = ifn0(dd[0], atype, fparam=ifp, aparam=iap)
+                    self.assertIn("input fparam", context.exception)
+
+            if nap > 0:
+                iap = rng.normal(size=(self.nf, self.nloc, nap - 1))
+                with self.assertRaises(ValueError) as context:
+                    ifn0(dd[0], atype, fparam=ifp, aparam=iap)
+                    self.assertIn("input aparam", context.exception)
+
+    def test_get_set(self):
+        ifn0 = InvarFitting(
+            "energy",
+            self.nt,
+            3,
+            1,
+        )
+        rng = np.random.default_rng()
+        foo = rng.normal([3, 4])
+        for ii in [
+            "bias_atom_e",
+            "fparam_avg",
+            "fparam_inv_std",
+            "aparam_avg",
+            "aparam_inv_std",
+        ]:
+            ifn0[ii] = foo
+            np.testing.assert_allclose(foo, ifn0[ii])