From fc09e77edc7f4e68fc3e24d626ab85a44cc05be9 Mon Sep 17 00:00:00 2001
From: Han Wang <92130845+wanghan-iapcm@users.noreply.github.com>
Date: Fri, 13 Oct 2023 12:36:39 +0800
Subject: [PATCH] unittest for the compression of smooth se_atten descriptor
 (#2916)

Co-authored-by: Han Wang <wang_han@iapcm.ac.cn>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
---
 .../tests/test_model_compression_se_atten.py  | 44 ++++++++++++++++---
 source/tests/test_model_se_atten.py           |  2 +
 2 files changed, 39 insertions(+), 7 deletions(-)

diff --git a/source/tests/test_model_compression_se_atten.py b/source/tests/test_model_compression_se_atten.py
index 30d78669ce..6bab1a3881 100644
--- a/source/tests/test_model_compression_se_atten.py
+++ b/source/tests/test_model_compression_se_atten.py
@@ -46,10 +46,31 @@ def _subprocess_run(command):
 # - type embedding FP32, se_atten FP64
 # - type embedding FP32, se_atten FP32
 tests = [
-    {"se_atten precision": "float64", "type embedding precision": "float64"},
-    {"se_atten precision": "float64", "type embedding precision": "float32"},
-    {"se_atten precision": "float32", "type embedding precision": "float64"},
-    {"se_atten precision": "float32", "type embedding precision": "float32"},
+    {
+        "se_atten precision": "float64",
+        "type embedding precision": "float64",
+        "smooth_type_embdding": True,
+    },
+    {
+        "se_atten precision": "float64",
+        "type embedding precision": "float64",
+        "smooth_type_embdding": False,
+    },
+    {
+        "se_atten precision": "float64",
+        "type embedding precision": "float32",
+        "smooth_type_embdding": True,
+    },
+    {
+        "se_atten precision": "float32",
+        "type embedding precision": "float64",
+        "smooth_type_embdding": True,
+    },
+    {
+        "se_atten precision": "float32",
+        "type embedding precision": "float32",
+        "smooth_type_embdding": True,
+    },
 ]
 
 
@@ -73,6 +94,9 @@ def _init_models():
         jdata["model"]["descriptor"]["stripped_type_embedding"] = True
         jdata["model"]["descriptor"]["sel"] = 120
         jdata["model"]["descriptor"]["attn_layer"] = 0
+        jdata["model"]["descriptor"]["smooth_type_embdding"] = tests[i][
+            "smooth_type_embdding"
+        ]
         jdata["model"]["type_embedding"] = {}
         jdata["model"]["type_embedding"]["precision"] = tests[i][
             "type embedding precision"
@@ -479,9 +503,15 @@ def test_1frame(self):
             self.assertEqual(ff1.shape, (nframes, natoms, 3))
             self.assertEqual(vv1.shape, (nframes, 9))
             # check values
-            np.testing.assert_almost_equal(ff0, ff1, default_places)
-            np.testing.assert_almost_equal(ee0, ee1, default_places)
-            np.testing.assert_almost_equal(vv0, vv1, default_places)
+            np.testing.assert_almost_equal(
+                ff0, ff1, default_places, err_msg=str(tests[i])
+            )
+            np.testing.assert_almost_equal(
+                ee0, ee1, default_places, err_msg=str(tests[i])
+            )
+            np.testing.assert_almost_equal(
+                vv0, vv1, default_places, err_msg=str(tests[i])
+            )
 
     def test_1frame_atm(self):
         for i in range(len(tests)):
diff --git a/source/tests/test_model_se_atten.py b/source/tests/test_model_se_atten.py
index ad037e2931..5417201a9f 100644
--- a/source/tests/test_model_se_atten.py
+++ b/source/tests/test_model_se_atten.py
@@ -753,6 +753,8 @@ def test_smoothness_of_stripped_type_embedding_smooth_model(self):
         jdata["model"]["descriptor"]["stripped_type_embedding"] = True
         jdata["model"]["descriptor"]["smooth_type_embdding"] = True
         jdata["model"]["descriptor"]["attn_layer"] = 1
+        jdata["model"]["descriptor"]["rcut"] = 6.0
+        jdata["model"]["descriptor"]["rcut_smth"] = 4.0
         descrpt = DescrptSeAtten(**jdata["model"]["descriptor"], uniform_seed=True)
         jdata["model"]["fitting_net"]["descrpt"] = descrpt
         fitting = EnerFitting(**jdata["model"]["fitting_net"], uniform_seed=True)