fix TestDeepPotAPBCExcludeTypes (#2680)

I found that #956 broke this unit test, making it only testing a
standard model.

Signed-off-by: Jinzhe Zeng <jinzhe.zeng@rutgers.edu>

source/tests/test_model_compression_se_a.py

@@ -64,7 +64,29 @@ def _init_models():
     return INPUT, frozen_model, compressed_model
 
 
+def _init_models_exclude_types():
+    data_file = str(tests_path / os.path.join("model_compression", "data"))
+    frozen_model = str(tests_path / "dp-original-exclude-types.pb")
+    compressed_model = str(tests_path / "dp-compressed-exclude-types.pb")
+    INPUT = str(tests_path / "input.json")
+    jdata = j_loader(str(tests_path / os.path.join("model_compression", "input.json")))
+    jdata["model"]["descriptor"]["exclude_types"] = [[0, 1]]
+    jdata["training"]["training_data"]["systems"] = data_file
+    jdata["training"]["validation_data"]["systems"] = data_file
+    with open(INPUT, "w") as fp:
+        json.dump(jdata, fp, indent=4)
+
+    ret = run_dp("dp train " + INPUT)
+    np.testing.assert_equal(ret, 0, "DP train failed!")
+    ret = run_dp("dp freeze -o " + frozen_model)
+    np.testing.assert_equal(ret, 0, "DP freeze failed!")
+    ret = run_dp("dp compress " + " -i " + frozen_model + " -o " + compressed_model)
+    np.testing.assert_equal(ret, 0, "DP model compression failed!")
+    return INPUT, frozen_model, compressed_model
+
+
 INPUT, FROZEN_MODEL, COMPRESSED_MODEL = _init_models()
+INPUT_ET, FROZEN_MODEL_ET, COMPRESSED_MODEL_ET = _init_models_exclude_types()
 
 
 class TestDeepPotAPBC(unittest.TestCase):
@@ -410,8 +432,8 @@ def test_ase(self):
 class TestDeepPotAPBCExcludeTypes(unittest.TestCase):
     @classmethod
     def setUpClass(self):
-        self.dp_original = DeepPot(FROZEN_MODEL)
-        self.dp_compressed = DeepPot(COMPRESSED_MODEL)
+        self.dp_original = DeepPot(FROZEN_MODEL_ET)
+        self.dp_compressed = DeepPot(COMPRESSED_MODEL_ET)
         self.coords = np.array(
             [
                 12.83,
@@ -439,9 +461,9 @@ def setUpClass(self):
 
     @classmethod
     def tearDownClass(self):
-        _file_delete(INPUT)
-        _file_delete(FROZEN_MODEL)
-        _file_delete(COMPRESSED_MODEL)
+        _file_delete(INPUT_ET)
+        _file_delete(FROZEN_MODEL_ET)
+        _file_delete(COMPRESSED_MODEL_ET)
         _file_delete("out.json")
         _file_delete("compress.json")
         _file_delete("checkpoint")