From f9e0a5a8d34fbb0ac3f18e5157957a2c5abb8262 Mon Sep 17 00:00:00 2001
From: Jinzhe Zeng <jinzhe.zeng@rutgers.edu>
Date: Tue, 4 Jun 2024 17:25:40 -0400
Subject: [PATCH] fix(tf): fix modifier_type in DeepEval

A downgrade in #3213.

Signed-off-by: Jinzhe Zeng <jinzhe.zeng@rutgers.edu>
---
 deepmd/tf/infer/deep_eval.py |  7 ++++-
 source/tests/tf/test_dplr.py | 60 ++++++++++++++++++++++++++++++++++++
 2 files changed, 66 insertions(+), 1 deletion(-)
 create mode 100644 source/tests/tf/test_dplr.py

diff --git a/deepmd/tf/infer/deep_eval.py b/deepmd/tf/infer/deep_eval.py
index 825ac6704a..313b22842a 100644
--- a/deepmd/tf/infer/deep_eval.py
+++ b/deepmd/tf/infer/deep_eval.py
@@ -139,7 +139,6 @@ def __init__(
         self.has_fparam = self.tensors["fparam"] is not None
         self.has_aparam = self.tensors["aparam"] is not None
         self.has_spin = self.ntypes_spin > 0
-        self.modifier_type = None
 
         # looks ugly...
         if self.modifier_type == "dipole_charge":
@@ -201,6 +200,8 @@ def _init_tensors(self):
             "ntypes_spin": "spin_attr/ntypes_spin:0",
             # descriptor
             "descriptor": "o_descriptor:0",
+            # modifier
+            "modifier_type": "modifier_attr/type:0",
         }
         # output tensors
         output_tensor_names = {}
@@ -260,6 +261,10 @@ def _init_attr(self):
         else:
             self.numb_dos = 0
         self.tmap = tmap.decode("utf-8").split()
+        if self.tensors["modifier_type"] is not None:
+            self.modifier_type = run_sess(self.sess, [self.tensors["modifier_type"]])[0]
+        else:
+            self.modifier_type = None
 
     @property
     @lru_cache(maxsize=None)
diff --git a/source/tests/tf/test_dplr.py b/source/tests/tf/test_dplr.py
new file mode 100644
index 0000000000..6788a22547
--- /dev/null
+++ b/source/tests/tf/test_dplr.py
@@ -0,0 +1,60 @@
+# SPDX-License-Identifier: LGPL-3.0-or-later
+import unittest
+from pathlib import (
+    Path,
+)
+
+import numpy as np
+
+from deepmd.infer.deep_pot import (
+    DeepPot,
+)
+from deepmd.tf.utils.convert import (
+    convert_pbtxt_to_pb,
+)
+
+
+class TestDPLR(unittest.TestCase):
+    def setUp(self):
+        # a bit strange path, need to move to the correct directory
+        pbtxt_file = (
+            Path(__file__).parent.parent.parent / "lmp" / "tests" / "lrmodel.pbtxt"
+        )
+        convert_pbtxt_to_pb(pbtxt_file, "lrmodel.pb")
+
+        self.expected_e_lr_efield_variable = -40.56538550
+        self.expected_f_lr_efield_variable = np.array(
+            [
+                [0.35019748, 0.27802691, -0.38443156],
+                [-0.42115581, -0.20474826, -0.02701100],
+                [-0.56357653, 0.34154004, 0.78389512],
+                [0.21023870, -0.60684635, -0.39875165],
+                [0.78732106, 0.00610023, 0.17197636],
+                [-0.36302488, 0.18592742, -0.14567727],
+            ]
+        )
+
+        self.box = np.eye(3).reshape(1, 9) * 20.0
+        self.coord = np.array(
+            [
+                [1.25545000, 1.27562200, 0.98873000],
+                [0.96101000, 3.25750000, 1.33494000],
+                [0.66417000, 1.31153700, 1.74354000],
+                [1.29187000, 0.33436000, 0.73085000],
+                [1.88885000, 3.51130000, 1.42444000],
+                [0.51617000, 4.04330000, 0.90904000],
+                [1.25545000, 1.27562200, 0.98873000],
+                [0.96101000, 3.25750000, 1.33494000],
+            ]
+        ).reshape(1, 8, 3)
+        self.atype = np.array([0, 0, 1, 1, 1, 1, 2, 2])
+
+    def test_eval(self):
+        dp = DeepPot("lrmodel.pb")
+        e, f, v, ae, av = dp.eval(
+            self.coord[:, :6], self.box, self.atype[:6], atomic=True
+        )
+        np.testing.assert_allclose(e, self.expected_e_lr_efield_variable, atol=1e-6)
+        np.testing.assert_allclose(
+            f.ravel(), self.expected_f_lr_efield_variable.ravel(), atol=1e-6
+        )