From fdb6a3d67f55755150212b849b95fc49cbce80f4 Mon Sep 17 00:00:00 2001
From: Anyang Peng <137014849+anyangml@users.noreply.github.com>
Date: Thu, 21 Mar 2024 16:17:46 +0800
Subject: [PATCH 01/30] feat: jit export in linear model

---
 .../pt/model/atomic_model/linear_atomic_model.py | 16 +++++++++++++---
 .../tests/pt/model/test_linear_atomic_model.py   | 10 ++++------
 2 files changed, 17 insertions(+), 9 deletions(-)

diff --git a/deepmd/pt/model/atomic_model/linear_atomic_model.py b/deepmd/pt/model/atomic_model/linear_atomic_model.py
index f7216f46ef..2cf1cc556c 100644
--- a/deepmd/pt/model/atomic_model/linear_atomic_model.py
+++ b/deepmd/pt/model/atomic_model/linear_atomic_model.py
@@ -92,25 +92,31 @@ def mixed_types(self) -> bool:
         """
         return True
 
+    @torch.jit.export
     def get_rcut(self) -> float:
         """Get the cut-off radius."""
         return max(self.get_model_rcuts())
-
+    
+    @torch.jit.export
     def get_type_map(self) -> List[str]:
         """Get the type map."""
         return self.type_map
-
+    
+    @torch.jit.export
     def get_model_rcuts(self) -> List[float]:
         """Get the cut-off radius for each individual models."""
         return [model.get_rcut() for model in self.models]
-
+    
+    @torch.jit.export
     def get_sel(self) -> List[int]:
         return [max([model.get_nsel() for model in self.models])]
 
+    @torch.jit.export
     def get_model_nsels(self) -> List[int]:
         """Get the processed sels for each individual models. Not distinguishing types."""
         return [model.get_nsel() for model in self.models]
 
+    @torch.jit.export
     def get_model_sels(self) -> List[List[int]]:
         """Get the sels for each individual models."""
         return [model.get_sel() for model in self.models]
@@ -289,15 +295,18 @@ def _compute_weight(
             for _ in range(nmodels)
         ]
 
+    @torch.jit.export
     def get_dim_fparam(self) -> int:
         """Get the number (dimension) of frame parameters of this atomic model."""
         # tricky...
         return max([model.get_dim_fparam() for model in self.models])
 
+    @torch.jit.export
     def get_dim_aparam(self) -> int:
         """Get the number (dimension) of atomic parameters of this atomic model."""
         return max([model.get_dim_aparam() for model in self.models])
 
+    @torch.jit.export
     def get_sel_type(self) -> List[int]:
         """Get the selected atom types of this model.
 
@@ -318,6 +327,7 @@ def get_sel_type(self) -> List[int]:
             )
         ).tolist()
 
+    @torch.jit.export
     def is_aparam_nall(self) -> bool:
         """Check whether the shape of atomic parameters is (nframes, nall, ndim).
 
diff --git a/source/tests/pt/model/test_linear_atomic_model.py b/source/tests/pt/model/test_linear_atomic_model.py
index 7f24ffdc53..adc682a41f 100644
--- a/source/tests/pt/model/test_linear_atomic_model.py
+++ b/source/tests/pt/model/test_linear_atomic_model.py
@@ -178,13 +178,11 @@ def test_self_consistency(self):
 
     def test_jit(self):
         md1 = torch.jit.script(self.md1)
-        # atomic model no more export methods
-        # self.assertEqual(md1.get_rcut(), self.rcut)
-        # self.assertEqual(md1.get_type_map(), ["foo", "bar"])
+        self.assertEqual(md1.get_rcut(), self.rcut)
+        self.assertEqual(md1.get_type_map(), ["foo", "bar"])
         md3 = torch.jit.script(self.md3)
-        # atomic model no more export methods
-        # self.assertEqual(md3.get_rcut(), self.rcut)
-        # self.assertEqual(md3.get_type_map(), ["foo", "bar"])
+        self.assertEqual(md3.get_rcut(), self.rcut)
+        self.assertEqual(md3.get_type_map(), ["foo", "bar"])
 
 
 class TestRemmapMethod(unittest.TestCase):

From a4201fe485194810eafec5d20d4a6828c1bf0579 Mon Sep 17 00:00:00 2001
From: "pre-commit-ci[bot]"
 <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Date: Thu, 21 Mar 2024 08:19:49 +0000
Subject: [PATCH 02/30] [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci
---
 deepmd/pt/model/atomic_model/linear_atomic_model.py | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/deepmd/pt/model/atomic_model/linear_atomic_model.py b/deepmd/pt/model/atomic_model/linear_atomic_model.py
index 2cf1cc556c..ae8b246233 100644
--- a/deepmd/pt/model/atomic_model/linear_atomic_model.py
+++ b/deepmd/pt/model/atomic_model/linear_atomic_model.py
@@ -96,17 +96,17 @@ def mixed_types(self) -> bool:
     def get_rcut(self) -> float:
         """Get the cut-off radius."""
         return max(self.get_model_rcuts())
-    
+
     @torch.jit.export
     def get_type_map(self) -> List[str]:
         """Get the type map."""
         return self.type_map
-    
+
     @torch.jit.export
     def get_model_rcuts(self) -> List[float]:
         """Get the cut-off radius for each individual models."""
         return [model.get_rcut() for model in self.models]
-    
+
     @torch.jit.export
     def get_sel(self) -> List[int]:
         return [max([model.get_nsel() for model in self.models])]

From c0c1dba67622a31cf122b54738e7962d9b7d4f24 Mon Sep 17 00:00:00 2001
From: Anyang Peng <137014849+anyangml@users.noreply.github.com>
Date: Fri, 22 Mar 2024 09:57:31 +0800
Subject: [PATCH 03/30] fix: revert atomic change

---
 deepmd/pt/model/atomic_model/linear_atomic_model.py | 11 +----------
 source/tests/pt/model/test_linear_atomic_model.py   | 10 ++++++++--
 2 files changed, 9 insertions(+), 12 deletions(-)

diff --git a/deepmd/pt/model/atomic_model/linear_atomic_model.py b/deepmd/pt/model/atomic_model/linear_atomic_model.py
index ae8b246233..161ac62843 100644
--- a/deepmd/pt/model/atomic_model/linear_atomic_model.py
+++ b/deepmd/pt/model/atomic_model/linear_atomic_model.py
@@ -92,12 +92,10 @@ def mixed_types(self) -> bool:
         """
         return True
 
-    @torch.jit.export
     def get_rcut(self) -> float:
         """Get the cut-off radius."""
         return max(self.get_model_rcuts())
-
-    @torch.jit.export
+    
     def get_type_map(self) -> List[str]:
         """Get the type map."""
         return self.type_map
@@ -107,16 +105,13 @@ def get_model_rcuts(self) -> List[float]:
         """Get the cut-off radius for each individual models."""
         return [model.get_rcut() for model in self.models]
 
-    @torch.jit.export
     def get_sel(self) -> List[int]:
         return [max([model.get_nsel() for model in self.models])]
 
-    @torch.jit.export
     def get_model_nsels(self) -> List[int]:
         """Get the processed sels for each individual models. Not distinguishing types."""
         return [model.get_nsel() for model in self.models]
 
-    @torch.jit.export
     def get_model_sels(self) -> List[List[int]]:
         """Get the sels for each individual models."""
         return [model.get_sel() for model in self.models]
@@ -295,18 +290,15 @@ def _compute_weight(
             for _ in range(nmodels)
         ]
 
-    @torch.jit.export
     def get_dim_fparam(self) -> int:
         """Get the number (dimension) of frame parameters of this atomic model."""
         # tricky...
         return max([model.get_dim_fparam() for model in self.models])
 
-    @torch.jit.export
     def get_dim_aparam(self) -> int:
         """Get the number (dimension) of atomic parameters of this atomic model."""
         return max([model.get_dim_aparam() for model in self.models])
 
-    @torch.jit.export
     def get_sel_type(self) -> List[int]:
         """Get the selected atom types of this model.
 
@@ -327,7 +319,6 @@ def get_sel_type(self) -> List[int]:
             )
         ).tolist()
 
-    @torch.jit.export
     def is_aparam_nall(self) -> bool:
         """Check whether the shape of atomic parameters is (nframes, nall, ndim).
 
diff --git a/source/tests/pt/model/test_linear_atomic_model.py b/source/tests/pt/model/test_linear_atomic_model.py
index adc682a41f..48f3f5ad02 100644
--- a/source/tests/pt/model/test_linear_atomic_model.py
+++ b/source/tests/pt/model/test_linear_atomic_model.py
@@ -178,11 +178,17 @@ def test_self_consistency(self):
 
     def test_jit(self):
         md1 = torch.jit.script(self.md1)
-        self.assertEqual(md1.get_rcut(), self.rcut)
-        self.assertEqual(md1.get_type_map(), ["foo", "bar"])
+        # atomic model no more export methods
+        # self.assertEqual(md1.get_rcut(), self.rcut)
+        # self.assertEqual(md1.get_type_map(), ["foo", "bar"])
+        # md3 is the model, not atomic model
         md3 = torch.jit.script(self.md3)
         self.assertEqual(md3.get_rcut(), self.rcut)
         self.assertEqual(md3.get_type_map(), ["foo", "bar"])
+        self.assertEqual(md3.get_sel(), [sum(self.sel)])
+        self.assertEqual(md3.get_dim_aparam(), 0)
+        self.assertEqual(md3.get_dim_fparam(), 0)
+        self.assertEqual(md3.is_aparam_nall(), False)
 
 
 class TestRemmapMethod(unittest.TestCase):

From 21dc19f25dd29ec2188b488d2c87325e6786cdc8 Mon Sep 17 00:00:00 2001
From: "pre-commit-ci[bot]"
 <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Date: Fri, 22 Mar 2024 02:00:46 +0000
Subject: [PATCH 04/30] [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci
---
 deepmd/pt/model/atomic_model/linear_atomic_model.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/deepmd/pt/model/atomic_model/linear_atomic_model.py b/deepmd/pt/model/atomic_model/linear_atomic_model.py
index 161ac62843..c623b4f7f0 100644
--- a/deepmd/pt/model/atomic_model/linear_atomic_model.py
+++ b/deepmd/pt/model/atomic_model/linear_atomic_model.py
@@ -95,7 +95,7 @@ def mixed_types(self) -> bool:
     def get_rcut(self) -> float:
         """Get the cut-off radius."""
         return max(self.get_model_rcuts())
-    
+
     def get_type_map(self) -> List[str]:
         """Get the type map."""
         return self.type_map

From cba5823e11560bfcf09f01fe59e0691f6eae874c Mon Sep 17 00:00:00 2001
From: Anyang Peng <137014849+anyangml@users.noreply.github.com>
Date: Fri, 22 Mar 2024 19:51:19 +0800
Subject: [PATCH 05/30] fix: revert change

---
 deepmd/pt/model/atomic_model/linear_atomic_model.py | 1 -
 1 file changed, 1 deletion(-)

diff --git a/deepmd/pt/model/atomic_model/linear_atomic_model.py b/deepmd/pt/model/atomic_model/linear_atomic_model.py
index c623b4f7f0..f7216f46ef 100644
--- a/deepmd/pt/model/atomic_model/linear_atomic_model.py
+++ b/deepmd/pt/model/atomic_model/linear_atomic_model.py
@@ -100,7 +100,6 @@ def get_type_map(self) -> List[str]:
         """Get the type map."""
         return self.type_map
 
-    @torch.jit.export
     def get_model_rcuts(self) -> List[float]:
         """Get the cut-off radius for each individual models."""
         return [model.get_rcut() for model in self.models]

From 08e60537b85f3e29f618ea84492354e19ae93710 Mon Sep 17 00:00:00 2001
From: Anyang Peng <137014849+anyangml@users.noreply.github.com>
Date: Sun, 24 Mar 2024 19:08:45 +0800
Subject: [PATCH 06/30] fix: linear bias

---
 .../atomic_model/linear_atomic_model.py       | 20 +++++++++++--------
 .../model/atomic_model/linear_atomic_model.py | 18 +++++++----------
 2 files changed, 19 insertions(+), 19 deletions(-)

diff --git a/deepmd/dpmodel/atomic_model/linear_atomic_model.py b/deepmd/dpmodel/atomic_model/linear_atomic_model.py
index 93a885f3ab..7895ac7820 100644
--- a/deepmd/dpmodel/atomic_model/linear_atomic_model.py
+++ b/deepmd/dpmodel/atomic_model/linear_atomic_model.py
@@ -67,6 +67,7 @@ def __init__(
         assert len(err_msg) == 0, "\n".join(err_msg)
         self.mixed_types_list = [model.mixed_types() for model in self.models]
         super().__init__(**kwargs)
+        self.atomic_bias = None
 
     def mixed_types(self) -> bool:
         """If true, the model
@@ -162,7 +163,6 @@ def forward_atomic(
             )
         ]
         ener_list = []
-
         for i, model in enumerate(self.models):
             mapping = self.mapping_list[i]
             ener_list.append(
@@ -176,11 +176,16 @@ def forward_atomic(
                 )["energy"]
             )
         self.weights = self._compute_weight(extended_coord, extended_atype, nlists_)
-        self.atomic_bias = None
-        if self.atomic_bias is not None:
-            raise NotImplementedError("Need to add bias in a future PR.")
-        else:
-            fit_ret = {
+        atype = extended_atype[:, :nloc]
+        bias_list = []
+        for idx, model in enumerate(self.models):
+            bias_atom_e = model.get_out_bias()
+            
+            ener_list[idx] += bias_atom_e[atype]
+            bias_list[idx] = bias_atom_e[atype]
+
+        self.atomic_bias = np.sum(np.stack(bias_list) * np.stack(self.weights), axis=0)
+        fit_ret = {
                 "energy": np.sum(np.stack(ener_list) * np.stack(self.weights), axis=0),
             }  # (nframes, nloc, 1)
         return fit_ret
@@ -293,8 +298,7 @@ def set_out_bias(self, out_bias: np.ndarray, add=False) -> None:
 
     def get_out_bias(self) -> np.ndarray:
         """Return the weighted output bias of the linear atomic model."""
-        # TODO add get_out_bias for linear atomic model
-        raise NotImplementedError
+        return self.atomic_bias
 
     def is_aparam_nall(self) -> bool:
         """Check whether the shape of atomic parameters is (nframes, nall, ndim).
diff --git a/deepmd/pt/model/atomic_model/linear_atomic_model.py b/deepmd/pt/model/atomic_model/linear_atomic_model.py
index f599399e66..3d922ec6b8 100644
--- a/deepmd/pt/model/atomic_model/linear_atomic_model.py
+++ b/deepmd/pt/model/atomic_model/linear_atomic_model.py
@@ -202,17 +202,14 @@ def forward_atomic(
         weights = self._compute_weight(extended_coord, extended_atype, nlists_)
 
         atype = extended_atype[:, :nloc]
+        bias_list = []
         for idx, model in enumerate(self.models):
-            # TODO: provide interfaces for atomic models to access bias_atom_e
-            if isinstance(model, DPAtomicModel):
-                bias_atom_e = model.fitting_net.bias_atom_e
-            elif isinstance(model, PairTabAtomicModel):
-                bias_atom_e = model.bias_atom_e
-            else:
-                bias_atom_e = None
-            if bias_atom_e is not None:
-                ener_list[idx] += bias_atom_e[atype]
+            bias_atom_e = model.get_out_bias()
+            
+            ener_list[idx] += bias_atom_e[atype]
+            bias_list[idx] = bias_atom_e[atype]
 
+        self.atomic_bias = torch.sum(torch.stack(bias_list) * torch.stack(weights), dim=0)
         fit_ret = {
             "energy": torch.sum(torch.stack(ener_list) * torch.stack(weights), dim=0),
         }  # (nframes, nloc, 1)
@@ -307,8 +304,7 @@ def set_out_bias(self, out_bias: torch.Tensor, add=False) -> None:
 
     def get_out_bias(self) -> torch.Tensor:
         """Return the weighted output bias of the linear atomic model."""
-        # TODO add get_out_bias for linear atomic model
-        raise NotImplementedError
+        return self.atomic_bias
 
     def get_dim_fparam(self) -> int:
         """Get the number (dimension) of frame parameters of this atomic model."""

From 83e8a5b1a0c000ff51310e72c2fbaa55d8407213 Mon Sep 17 00:00:00 2001
From: "pre-commit-ci[bot]"
 <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Date: Sun, 24 Mar 2024 11:09:21 +0000
Subject: [PATCH 07/30] [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci
---
 deepmd/dpmodel/atomic_model/linear_atomic_model.py  | 6 +++---
 deepmd/pt/model/atomic_model/linear_atomic_model.py | 6 ++++--
 2 files changed, 7 insertions(+), 5 deletions(-)

diff --git a/deepmd/dpmodel/atomic_model/linear_atomic_model.py b/deepmd/dpmodel/atomic_model/linear_atomic_model.py
index 7895ac7820..660898fb7d 100644
--- a/deepmd/dpmodel/atomic_model/linear_atomic_model.py
+++ b/deepmd/dpmodel/atomic_model/linear_atomic_model.py
@@ -180,14 +180,14 @@ def forward_atomic(
         bias_list = []
         for idx, model in enumerate(self.models):
             bias_atom_e = model.get_out_bias()
-            
+
             ener_list[idx] += bias_atom_e[atype]
             bias_list[idx] = bias_atom_e[atype]
 
         self.atomic_bias = np.sum(np.stack(bias_list) * np.stack(self.weights), axis=0)
         fit_ret = {
-                "energy": np.sum(np.stack(ener_list) * np.stack(self.weights), axis=0),
-            }  # (nframes, nloc, 1)
+            "energy": np.sum(np.stack(ener_list) * np.stack(self.weights), axis=0),
+        }  # (nframes, nloc, 1)
         return fit_ret
 
     @staticmethod
diff --git a/deepmd/pt/model/atomic_model/linear_atomic_model.py b/deepmd/pt/model/atomic_model/linear_atomic_model.py
index 3d922ec6b8..65556cf658 100644
--- a/deepmd/pt/model/atomic_model/linear_atomic_model.py
+++ b/deepmd/pt/model/atomic_model/linear_atomic_model.py
@@ -205,11 +205,13 @@ def forward_atomic(
         bias_list = []
         for idx, model in enumerate(self.models):
             bias_atom_e = model.get_out_bias()
-            
+
             ener_list[idx] += bias_atom_e[atype]
             bias_list[idx] = bias_atom_e[atype]
 
-        self.atomic_bias = torch.sum(torch.stack(bias_list) * torch.stack(weights), dim=0)
+        self.atomic_bias = torch.sum(
+            torch.stack(bias_list) * torch.stack(weights), dim=0
+        )
         fit_ret = {
             "energy": torch.sum(torch.stack(ener_list) * torch.stack(weights), dim=0),
         }  # (nframes, nloc, 1)

From e57ac56a6e9216a77778f103cf4862bd968f961f Mon Sep 17 00:00:00 2001
From: Anyang Peng <137014849+anyangml@users.noreply.github.com>
Date: Sun, 24 Mar 2024 22:29:18 +0800
Subject: [PATCH 08/30] fix: index

---
 deepmd/dpmodel/atomic_model/linear_atomic_model.py  | 2 +-
 deepmd/pt/model/atomic_model/linear_atomic_model.py | 2 +-
 2 files changed, 2 insertions(+), 2 deletions(-)

diff --git a/deepmd/dpmodel/atomic_model/linear_atomic_model.py b/deepmd/dpmodel/atomic_model/linear_atomic_model.py
index 660898fb7d..c6c0fc2058 100644
--- a/deepmd/dpmodel/atomic_model/linear_atomic_model.py
+++ b/deepmd/dpmodel/atomic_model/linear_atomic_model.py
@@ -182,7 +182,7 @@ def forward_atomic(
             bias_atom_e = model.get_out_bias()
 
             ener_list[idx] += bias_atom_e[atype]
-            bias_list[idx] = bias_atom_e[atype]
+            bias_list.append(bias_atom_e[atype])
 
         self.atomic_bias = np.sum(np.stack(bias_list) * np.stack(self.weights), axis=0)
         fit_ret = {
diff --git a/deepmd/pt/model/atomic_model/linear_atomic_model.py b/deepmd/pt/model/atomic_model/linear_atomic_model.py
index 65556cf658..0dd0ebe454 100644
--- a/deepmd/pt/model/atomic_model/linear_atomic_model.py
+++ b/deepmd/pt/model/atomic_model/linear_atomic_model.py
@@ -207,7 +207,7 @@ def forward_atomic(
             bias_atom_e = model.get_out_bias()
 
             ener_list[idx] += bias_atom_e[atype]
-            bias_list[idx] = bias_atom_e[atype]
+            bias_list.append(bias_atom_e[atype])
 
         self.atomic_bias = torch.sum(
             torch.stack(bias_list) * torch.stack(weights), dim=0

From abb2455f4c24685c255c74eee9a622115e54f1b2 Mon Sep 17 00:00:00 2001
From: anyangml <anyangpeng.ca@gmail.com>
Date: Mon, 25 Mar 2024 01:39:23 +0000
Subject: [PATCH 09/30] fix: UTs

---
 deepmd/dpmodel/atomic_model/pairtab_atomic_model.py | 3 +++
 1 file changed, 3 insertions(+)

diff --git a/deepmd/dpmodel/atomic_model/pairtab_atomic_model.py b/deepmd/dpmodel/atomic_model/pairtab_atomic_model.py
index 30ab58928b..470db4acb5 100644
--- a/deepmd/dpmodel/atomic_model/pairtab_atomic_model.py
+++ b/deepmd/dpmodel/atomic_model/pairtab_atomic_model.py
@@ -81,6 +81,9 @@ def __init__(
         else:
             self.tab_info, self.tab_data = None, None
 
+        self.bias_atom_e = np.zeros(
+            (self.ntypes, 1), dtype=np.float64
+        )
         if isinstance(sel, int):
             self.sel = sel
         elif isinstance(sel, list):

From 7d8afcbcddc7a51550f29d8cee41994f98dbbf32 Mon Sep 17 00:00:00 2001
From: "pre-commit-ci[bot]"
 <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Date: Mon, 25 Mar 2024 01:41:08 +0000
Subject: [PATCH 10/30] [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci
---
 deepmd/dpmodel/atomic_model/pairtab_atomic_model.py | 4 +---
 1 file changed, 1 insertion(+), 3 deletions(-)

diff --git a/deepmd/dpmodel/atomic_model/pairtab_atomic_model.py b/deepmd/dpmodel/atomic_model/pairtab_atomic_model.py
index 470db4acb5..17977aea27 100644
--- a/deepmd/dpmodel/atomic_model/pairtab_atomic_model.py
+++ b/deepmd/dpmodel/atomic_model/pairtab_atomic_model.py
@@ -81,9 +81,7 @@ def __init__(
         else:
             self.tab_info, self.tab_data = None, None
 
-        self.bias_atom_e = np.zeros(
-            (self.ntypes, 1), dtype=np.float64
-        )
+        self.bias_atom_e = np.zeros((self.ntypes, 1), dtype=np.float64)
         if isinstance(sel, int):
             self.sel = sel
         elif isinstance(sel, list):

From 6307ad5fe654dd473320050d01df80ffa3554a60 Mon Sep 17 00:00:00 2001
From: Anyang Peng <137014849+anyangml@users.noreply.github.com>
Date: Mon, 25 Mar 2024 09:44:28 +0800
Subject: [PATCH 11/30] chore: revert changes

---
 source/tests/pt/model/test_linear_atomic_model.py | 11 ++++-------
 1 file changed, 4 insertions(+), 7 deletions(-)

diff --git a/source/tests/pt/model/test_linear_atomic_model.py b/source/tests/pt/model/test_linear_atomic_model.py
index 48f3f5ad02..e0f859acae 100644
--- a/source/tests/pt/model/test_linear_atomic_model.py
+++ b/source/tests/pt/model/test_linear_atomic_model.py
@@ -181,14 +181,11 @@ def test_jit(self):
         # atomic model no more export methods
         # self.assertEqual(md1.get_rcut(), self.rcut)
         # self.assertEqual(md1.get_type_map(), ["foo", "bar"])
-        # md3 is the model, not atomic model
         md3 = torch.jit.script(self.md3)
-        self.assertEqual(md3.get_rcut(), self.rcut)
-        self.assertEqual(md3.get_type_map(), ["foo", "bar"])
-        self.assertEqual(md3.get_sel(), [sum(self.sel)])
-        self.assertEqual(md3.get_dim_aparam(), 0)
-        self.assertEqual(md3.get_dim_fparam(), 0)
-        self.assertEqual(md3.is_aparam_nall(), False)
+        # atomic model no more export methods
+        # self.assertEqual(md3.get_rcut(), self.rcut)
+        # self.assertEqual(md3.get_type_map(), ["foo", "bar"])
+
 
 
 class TestRemmapMethod(unittest.TestCase):

From 82f1c39650f4f5acda24eb2a1663d8443e040ace Mon Sep 17 00:00:00 2001
From: Anyang Peng <137014849+anyangml@users.noreply.github.com>
Date: Mon, 25 Mar 2024 10:02:27 +0800
Subject: [PATCH 12/30] fix: placeholder

---
 deepmd/dpmodel/atomic_model/linear_atomic_model.py  | 5 ++++-
 deepmd/pt/model/atomic_model/linear_atomic_model.py | 4 +++-
 2 files changed, 7 insertions(+), 2 deletions(-)

diff --git a/deepmd/dpmodel/atomic_model/linear_atomic_model.py b/deepmd/dpmodel/atomic_model/linear_atomic_model.py
index c6c0fc2058..cf47fca406 100644
--- a/deepmd/dpmodel/atomic_model/linear_atomic_model.py
+++ b/deepmd/dpmodel/atomic_model/linear_atomic_model.py
@@ -66,8 +66,11 @@ def __init__(
             self.mapping_list.append(self.remap_atype(tpmp, self.type_map))
         assert len(err_msg) == 0, "\n".join(err_msg)
         self.mixed_types_list = [model.mixed_types() for model in self.models]
+        ntypes = len(self.type_map)
+        # a placeholder for storing weighted bias.
+        self.atomic_bias = np.zeros((ntypes, 1),dtype=np.float64)
         super().__init__(**kwargs)
-        self.atomic_bias = None
+
 
     def mixed_types(self) -> bool:
         """If true, the model
diff --git a/deepmd/pt/model/atomic_model/linear_atomic_model.py b/deepmd/pt/model/atomic_model/linear_atomic_model.py
index 0dd0ebe454..4ffb1789da 100644
--- a/deepmd/pt/model/atomic_model/linear_atomic_model.py
+++ b/deepmd/pt/model/atomic_model/linear_atomic_model.py
@@ -72,7 +72,9 @@ def __init__(
             self.mapping_list.append(self.remap_atype(tpmp, self.type_map))
         assert len(err_msg) == 0, "\n".join(err_msg)
 
-        self.atomic_bias = None
+        ntypes = len(self.type_map)
+        # a placeholder for storing weighted bias.
+        self.atomic_bias = torch.zeros(ntypes, 1, dtype=env.GLOBAL_PT_ENER_FLOAT_PRECISION, device=env.DEVICE)
         self.mixed_types_list = [model.mixed_types() for model in self.models]
         self.rcuts = torch.tensor(
             self.get_model_rcuts(), dtype=torch.float64, device=env.DEVICE

From 798cb3c1b20a0975b09bd5cae30530ae39736321 Mon Sep 17 00:00:00 2001
From: "pre-commit-ci[bot]"
 <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Date: Mon, 25 Mar 2024 02:16:45 +0000
Subject: [PATCH 13/30] [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci
---
 deepmd/dpmodel/atomic_model/linear_atomic_model.py  | 3 +--
 deepmd/pt/model/atomic_model/linear_atomic_model.py | 4 +++-
 source/tests/pt/model/test_linear_atomic_model.py   | 1 -
 3 files changed, 4 insertions(+), 4 deletions(-)

diff --git a/deepmd/dpmodel/atomic_model/linear_atomic_model.py b/deepmd/dpmodel/atomic_model/linear_atomic_model.py
index cf47fca406..ac9814c874 100644
--- a/deepmd/dpmodel/atomic_model/linear_atomic_model.py
+++ b/deepmd/dpmodel/atomic_model/linear_atomic_model.py
@@ -68,10 +68,9 @@ def __init__(
         self.mixed_types_list = [model.mixed_types() for model in self.models]
         ntypes = len(self.type_map)
         # a placeholder for storing weighted bias.
-        self.atomic_bias = np.zeros((ntypes, 1),dtype=np.float64)
+        self.atomic_bias = np.zeros((ntypes, 1), dtype=np.float64)
         super().__init__(**kwargs)
 
-
     def mixed_types(self) -> bool:
         """If true, the model
         1. assumes total number of atoms aligned across frames;
diff --git a/deepmd/pt/model/atomic_model/linear_atomic_model.py b/deepmd/pt/model/atomic_model/linear_atomic_model.py
index 4ffb1789da..a5a2ce1fcb 100644
--- a/deepmd/pt/model/atomic_model/linear_atomic_model.py
+++ b/deepmd/pt/model/atomic_model/linear_atomic_model.py
@@ -74,7 +74,9 @@ def __init__(
 
         ntypes = len(self.type_map)
         # a placeholder for storing weighted bias.
-        self.atomic_bias = torch.zeros(ntypes, 1, dtype=env.GLOBAL_PT_ENER_FLOAT_PRECISION, device=env.DEVICE)
+        self.atomic_bias = torch.zeros(
+            ntypes, 1, dtype=env.GLOBAL_PT_ENER_FLOAT_PRECISION, device=env.DEVICE
+        )
         self.mixed_types_list = [model.mixed_types() for model in self.models]
         self.rcuts = torch.tensor(
             self.get_model_rcuts(), dtype=torch.float64, device=env.DEVICE
diff --git a/source/tests/pt/model/test_linear_atomic_model.py b/source/tests/pt/model/test_linear_atomic_model.py
index e0f859acae..7f24ffdc53 100644
--- a/source/tests/pt/model/test_linear_atomic_model.py
+++ b/source/tests/pt/model/test_linear_atomic_model.py
@@ -187,7 +187,6 @@ def test_jit(self):
         # self.assertEqual(md3.get_type_map(), ["foo", "bar"])
 
 
-
 class TestRemmapMethod(unittest.TestCase):
     def test_valid(self):
         atype = torch.randint(0, 3, (4, 20), device=env.DEVICE)

From 0f6f0f2ea69a5f987becfd75cf5fe25c1704e52b Mon Sep 17 00:00:00 2001
From: anyangml <anyangpeng.ca@gmail.com>
Date: Mon, 25 Mar 2024 02:23:02 +0000
Subject: [PATCH 14/30] fix: Jit

---
 deepmd/pt/model/atomic_model/dp_atomic_model.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/deepmd/pt/model/atomic_model/dp_atomic_model.py b/deepmd/pt/model/atomic_model/dp_atomic_model.py
index 13b8f09a79..7b711d1f6d 100644
--- a/deepmd/pt/model/atomic_model/dp_atomic_model.py
+++ b/deepmd/pt/model/atomic_model/dp_atomic_model.py
@@ -242,7 +242,7 @@ def set_out_bias(self, out_bias: torch.Tensor, add=False) -> None:
 
     def get_out_bias(self) -> torch.Tensor:
         """Return the output bias of the atomic model."""
-        return self.fitting_net["bias_atom_e"]
+        return self.fitting_net.bias_atom_e
 
     def get_dim_fparam(self) -> int:
         """Get the number (dimension) of frame parameters of this atomic model."""

From 6fda3d6d313de4d07ad46532449f4dde7bfc854f Mon Sep 17 00:00:00 2001
From: Anyang Peng <137014849+anyangml@users.noreply.github.com>
Date: Sun, 7 Apr 2024 11:58:17 +0800
Subject: [PATCH 15/30] fix: remove get/set_out_bias

---
 .../dpmodel/atomic_model/dp_atomic_model.py   | 21 -------
 .../atomic_model/linear_atomic_model.py       | 28 ---------
 .../atomic_model/make_base_atomic_model.py    | 19 ------
 .../atomic_model/pairtab_atomic_model.py      | 19 ------
 .../model/atomic_model/base_atomic_model.py   |  2 -
 .../pt/model/atomic_model/dp_atomic_model.py  | 21 -------
 .../model/atomic_model/linear_atomic_model.py | 59 ++++++++-----------
 .../atomic_model/pairtab_atomic_model.py      | 19 ------
 deepmd/pt/model/model/make_model.py           |  2 -
 9 files changed, 25 insertions(+), 165 deletions(-)

diff --git a/deepmd/dpmodel/atomic_model/dp_atomic_model.py b/deepmd/dpmodel/atomic_model/dp_atomic_model.py
index cca46d3710..d39e236d07 100644
--- a/deepmd/dpmodel/atomic_model/dp_atomic_model.py
+++ b/deepmd/dpmodel/atomic_model/dp_atomic_model.py
@@ -79,27 +79,6 @@ def mixed_types(self) -> bool:
         """
         return self.descriptor.mixed_types()
 
-    def set_out_bias(self, out_bias: np.ndarray, add=False) -> None:
-        """
-        Modify the output bias for the atomic model.
-
-        Parameters
-        ----------
-        out_bias : np.ndarray
-            The new bias to be applied.
-        add : bool, optional
-            Whether to add the new bias to the existing one.
-            If False, the output bias will be directly replaced by the new bias.
-            If True, the new bias will be added to the existing one.
-        """
-        self.fitting["bias_atom_e"] = (
-            out_bias + self.fitting["bias_atom_e"] if add else out_bias
-        )
-
-    def get_out_bias(self) -> np.ndarray:
-        """Return the output bias of the atomic model."""
-        return self.fitting["bias_atom_e"]
-
     def forward_atomic(
         self,
         extended_coord: np.ndarray,
diff --git a/deepmd/dpmodel/atomic_model/linear_atomic_model.py b/deepmd/dpmodel/atomic_model/linear_atomic_model.py
index d01bd67826..e6296316a5 100644
--- a/deepmd/dpmodel/atomic_model/linear_atomic_model.py
+++ b/deepmd/dpmodel/atomic_model/linear_atomic_model.py
@@ -175,15 +175,7 @@ def forward_atomic(
                 )["energy"]
             )
         self.weights = self._compute_weight(extended_coord, extended_atype, nlists_)
-        atype = extended_atype[:, :nloc]
-        bias_list = []
-        for idx, model in enumerate(self.models):
-            bias_atom_e = model.get_out_bias()
 
-            ener_list[idx] += bias_atom_e[atype]
-            bias_list.append(bias_atom_e[atype])
-
-        self.atomic_bias = np.sum(np.stack(bias_list) * np.stack(self.weights), axis=0)
         fit_ret = {
             "energy": np.sum(np.stack(ener_list) * np.stack(self.weights), axis=0),
         }  # (nframes, nloc, 1)
@@ -279,26 +271,6 @@ def get_sel_type(self) -> List[int]:
         # join all the selected types
         return list(set().union(*[model.get_sel_type() for model in self.models]))
 
-    def set_out_bias(self, out_bias: np.ndarray, add=False) -> None:
-        """
-        Modify the output bias for all the models in the linear atomic model.
-
-        Parameters
-        ----------
-        out_bias : torch.Tensor
-            The new bias to be applied.
-        add : bool, optional
-            Whether to add the new bias to the existing one.
-            If False, the output bias will be directly replaced by the new bias.
-            If True, the new bias will be added to the existing one.
-        """
-        for model in self.models:
-            model.set_out_bias(out_bias, add=add)
-
-    def get_out_bias(self) -> np.ndarray:
-        """Return the weighted output bias of the linear atomic model."""
-        return self.atomic_bias
-
     def is_aparam_nall(self) -> bool:
         """Check whether the shape of atomic parameters is (nframes, nall, ndim).
 
diff --git a/deepmd/dpmodel/atomic_model/make_base_atomic_model.py b/deepmd/dpmodel/atomic_model/make_base_atomic_model.py
index 3e02a5d076..936c2b0943 100644
--- a/deepmd/dpmodel/atomic_model/make_base_atomic_model.py
+++ b/deepmd/dpmodel/atomic_model/make_base_atomic_model.py
@@ -95,25 +95,6 @@ def get_sel_type(self) -> List[int]:
             If returning an empty list, all atom types are selected.
             """
 
-        @abstractmethod
-        def set_out_bias(self, out_bias: t_tensor, add=False) -> None:
-            """
-            Modify the output bias for the atomic model.
-
-            Parameters
-            ----------
-            out_bias : t_tensor
-                The new bias to be applied.
-            add : bool, optional
-                Whether to add the new bias to the existing one.
-                If False, the output bias will be directly replaced by the new bias.
-                If True, the new bias will be added to the existing one.
-            """
-
-        @abstractmethod
-        def get_out_bias(self) -> t_tensor:
-            """Return the output bias of the atomic model."""
-
         @abstractmethod
         def is_aparam_nall(self) -> bool:
             """Check whether the shape of atomic parameters is (nframes, nall, ndim).
diff --git a/deepmd/dpmodel/atomic_model/pairtab_atomic_model.py b/deepmd/dpmodel/atomic_model/pairtab_atomic_model.py
index 9c1f355f0d..51e46901e9 100644
--- a/deepmd/dpmodel/atomic_model/pairtab_atomic_model.py
+++ b/deepmd/dpmodel/atomic_model/pairtab_atomic_model.py
@@ -131,25 +131,6 @@ def mixed_types(self) -> bool:
         # to match DPA1 and DPA2.
         return True
 
-    def set_out_bias(self, out_bias: np.ndarray, add=False) -> None:
-        """
-        Modify the output bias for the atomic model.
-
-        Parameters
-        ----------
-        out_bias : torch.Tensor
-            The new bias to be applied.
-        add : bool, optional
-            Whether to add the new bias to the existing one.
-            If False, the output bias will be directly replaced by the new bias.
-            If True, the new bias will be added to the existing one.
-        """
-        self.bias_atom_e = out_bias + self.bias_atom_e if add else out_bias
-
-    def get_out_bias(self) -> np.ndarray:
-        """Return the output bias of the atomic model."""
-        return self.bias_atom_e
-
     def serialize(self) -> dict:
         dd = BaseAtomicModel.serialize(self)
         dd.update(
diff --git a/deepmd/pt/model/atomic_model/base_atomic_model.py b/deepmd/pt/model/atomic_model/base_atomic_model.py
index 57ca21a826..e750b6a54e 100644
--- a/deepmd/pt/model/atomic_model/base_atomic_model.py
+++ b/deepmd/pt/model/atomic_model/base_atomic_model.py
@@ -366,7 +366,6 @@ def change_out_bias(
                 rcond=self.rcond,
                 preset_bias=self.preset_out_bias,
             )
-            # self.set_out_bias(delta_bias, add=True)
             self._store_out_stat(delta_bias, out_std, add=True)
         elif bias_adjust_mode == "set-by-statistic":
             bias_out, std_out = compute_output_stats(
@@ -377,7 +376,6 @@ def change_out_bias(
                 rcond=self.rcond,
                 preset_bias=self.preset_out_bias,
             )
-            # self.set_out_bias(bias_out)
             self._store_out_stat(bias_out, std_out)
         else:
             raise RuntimeError("Unknown bias_adjust_mode mode: " + bias_adjust_mode)
diff --git a/deepmd/pt/model/atomic_model/dp_atomic_model.py b/deepmd/pt/model/atomic_model/dp_atomic_model.py
index 82758672a1..14975732bc 100644
--- a/deepmd/pt/model/atomic_model/dp_atomic_model.py
+++ b/deepmd/pt/model/atomic_model/dp_atomic_model.py
@@ -217,27 +217,6 @@ def wrapped_sampler():
         self.descriptor.compute_input_stats(wrapped_sampler, stat_file_path)
         self.compute_or_load_out_stat(wrapped_sampler, stat_file_path)
 
-    def set_out_bias(self, out_bias: torch.Tensor, add=False) -> None:
-        """
-        Modify the output bias for the atomic model.
-
-        Parameters
-        ----------
-        out_bias : torch.Tensor
-            The new bias to be applied.
-        add : bool, optional
-            Whether to add the new bias to the existing one.
-            If False, the output bias will be directly replaced by the new bias.
-            If True, the new bias will be added to the existing one.
-        """
-        self.fitting_net["bias_atom_e"] = (
-            out_bias + self.fitting_net["bias_atom_e"] if add else out_bias
-        )
-
-    def get_out_bias(self) -> torch.Tensor:
-        """Return the output bias of the atomic model."""
-        return self.fitting_net.bias_atom_e
-
     def get_dim_fparam(self) -> int:
         """Get the number (dimension) of frame parameters of this atomic model."""
         return self.fitting_net.get_dim_fparam()
diff --git a/deepmd/pt/model/atomic_model/linear_atomic_model.py b/deepmd/pt/model/atomic_model/linear_atomic_model.py
index 709b59cb45..58d81a5037 100644
--- a/deepmd/pt/model/atomic_model/linear_atomic_model.py
+++ b/deepmd/pt/model/atomic_model/linear_atomic_model.py
@@ -192,8 +192,7 @@ def forward_atomic(
 
         for i, model in enumerate(self.models):
             mapping = self.mapping_list[i]
-            ener_list.append(
-                model.forward_atomic(
+            raw_ret = model.forward_atomic(
                     extended_coord,
                     mapping[extended_atype],
                     nlists_[i],
@@ -201,25 +200,37 @@ def forward_atomic(
                     fparam,
                     aparam,
                 )["energy"]
+            # apply bias to each individual model
+            ener_list.append(
+                model.apply_out_stat(
+                    raw_ret, mapping[extended_atype]
+                )    
             )
-
         weights = self._compute_weight(extended_coord, extended_atype, nlists_)
 
-        atype = extended_atype[:, :nloc]
-        bias_list = []
-        for idx, model in enumerate(self.models):
-            bias_atom_e = model.get_out_bias()
-
-            ener_list[idx] += bias_atom_e[atype]
-            bias_list.append(bias_atom_e[atype])
-
-        self.atomic_bias = torch.sum(
-            torch.stack(bias_list) * torch.stack(weights), dim=0
-        )
         fit_ret = {
             "energy": torch.sum(torch.stack(ener_list) * torch.stack(weights), dim=0),
         }  # (nframes, nloc, 1)
         return fit_ret
+    
+    def apply_out_stat(
+        self,
+        ret: Dict[str, torch.Tensor],
+        atype: torch.Tensor,
+    ):
+        """Apply the stat to each atomic output.
+        The developer may override the method to define how the bias is applied
+        to the atomic output of the model. 
+
+        Parameters
+        ----------
+        ret
+            The returned dict by the forward_atomic method
+        atype
+            The atom types. nf x nloc
+
+        """
+        return ret
 
     @staticmethod
     def remap_atype(ori_map: List[str], new_map: List[str]) -> torch.Tensor:
@@ -292,26 +303,6 @@ def _compute_weight(
             for _ in range(nmodels)
         ]
 
-    def set_out_bias(self, out_bias: torch.Tensor, add=False) -> None:
-        """
-        Modify the output bias for all the models in the linear atomic model.
-
-        Parameters
-        ----------
-        out_bias : torch.Tensor
-            The new bias to be applied.
-        add : bool, optional
-            Whether to add the new bias to the existing one.
-            If False, the output bias will be directly replaced by the new bias.
-            If True, the new bias will be added to the existing one.
-        """
-        for model in self.models:
-            model.set_out_bias(out_bias, add=add)
-
-    def get_out_bias(self) -> torch.Tensor:
-        """Return the weighted output bias of the linear atomic model."""
-        return self.atomic_bias
-
     def get_dim_fparam(self) -> int:
         """Get the number (dimension) of frame parameters of this atomic model."""
         # tricky...
diff --git a/deepmd/pt/model/atomic_model/pairtab_atomic_model.py b/deepmd/pt/model/atomic_model/pairtab_atomic_model.py
index 627dffd620..579e3efda4 100644
--- a/deepmd/pt/model/atomic_model/pairtab_atomic_model.py
+++ b/deepmd/pt/model/atomic_model/pairtab_atomic_model.py
@@ -226,25 +226,6 @@ def compute_or_load_stat(
         """
         self.compute_or_load_out_stat(merged, stat_file_path)
 
-    def set_out_bias(self, out_bias: torch.Tensor, add=False) -> None:
-        """
-        Modify the output bias for the atomic model.
-
-        Parameters
-        ----------
-        out_bias : torch.Tensor
-            The new bias to be applied.
-        add : bool, optional
-            Whether to add the new bias to the existing one.
-            If False, the output bias will be directly replaced by the new bias.
-            If True, the new bias will be added to the existing one.
-        """
-        self.bias_atom_e = out_bias + self.bias_atom_e if add else out_bias
-
-    def get_out_bias(self) -> torch.Tensor:
-        """Return the output bias of the atomic model."""
-        return self.bias_atom_e
-
     def forward_atomic(
         self,
         extended_coord: torch.Tensor,
diff --git a/deepmd/pt/model/model/make_model.py b/deepmd/pt/model/model/make_model.py
index 25a8ec9201..0c36b9783c 100644
--- a/deepmd/pt/model/model/make_model.py
+++ b/deepmd/pt/model/model/make_model.py
@@ -172,8 +172,6 @@ def forward_common(
             model_predict = self.output_type_cast(model_predict, input_prec)
             return model_predict
 
-        def get_out_bias(self) -> torch.Tensor:
-            return self.atomic_model.get_out_bias()
 
         def change_out_bias(
             self,

From a5c6f2db22e56bd1a94065078350077e3a4bb8a3 Mon Sep 17 00:00:00 2001
From: "pre-commit-ci[bot]"
 <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Date: Sun, 7 Apr 2024 03:58:49 +0000
Subject: [PATCH 16/30] [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci
---
 .../model/atomic_model/linear_atomic_model.py | 24 ++++++++-----------
 deepmd/pt/model/model/make_model.py           |  1 -
 2 files changed, 10 insertions(+), 15 deletions(-)

diff --git a/deepmd/pt/model/atomic_model/linear_atomic_model.py b/deepmd/pt/model/atomic_model/linear_atomic_model.py
index 58d81a5037..5c1db5a550 100644
--- a/deepmd/pt/model/atomic_model/linear_atomic_model.py
+++ b/deepmd/pt/model/atomic_model/linear_atomic_model.py
@@ -193,26 +193,22 @@ def forward_atomic(
         for i, model in enumerate(self.models):
             mapping = self.mapping_list[i]
             raw_ret = model.forward_atomic(
-                    extended_coord,
-                    mapping[extended_atype],
-                    nlists_[i],
-                    mapping,
-                    fparam,
-                    aparam,
-                )["energy"]
+                extended_coord,
+                mapping[extended_atype],
+                nlists_[i],
+                mapping,
+                fparam,
+                aparam,
+            )["energy"]
             # apply bias to each individual model
-            ener_list.append(
-                model.apply_out_stat(
-                    raw_ret, mapping[extended_atype]
-                )    
-            )
+            ener_list.append(model.apply_out_stat(raw_ret, mapping[extended_atype]))
         weights = self._compute_weight(extended_coord, extended_atype, nlists_)
 
         fit_ret = {
             "energy": torch.sum(torch.stack(ener_list) * torch.stack(weights), dim=0),
         }  # (nframes, nloc, 1)
         return fit_ret
-    
+
     def apply_out_stat(
         self,
         ret: Dict[str, torch.Tensor],
@@ -220,7 +216,7 @@ def apply_out_stat(
     ):
         """Apply the stat to each atomic output.
         The developer may override the method to define how the bias is applied
-        to the atomic output of the model. 
+        to the atomic output of the model.
 
         Parameters
         ----------
diff --git a/deepmd/pt/model/model/make_model.py b/deepmd/pt/model/model/make_model.py
index 0c36b9783c..f1a0334335 100644
--- a/deepmd/pt/model/model/make_model.py
+++ b/deepmd/pt/model/model/make_model.py
@@ -172,7 +172,6 @@ def forward_common(
             model_predict = self.output_type_cast(model_predict, input_prec)
             return model_predict
 
-
         def change_out_bias(
             self,
             merged,

From cbffc2b9e71d0588ba794026666f2bfe51b4a407 Mon Sep 17 00:00:00 2001
From: Anyang Peng <137014849+anyangml@users.noreply.github.com>
Date: Sun, 7 Apr 2024 12:00:36 +0800
Subject: [PATCH 17/30] fix: log

---
 deepmd/pt/train/training.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/deepmd/pt/train/training.py b/deepmd/pt/train/training.py
index 73404b0c83..e198ee3d1a 100644
--- a/deepmd/pt/train/training.py
+++ b/deepmd/pt/train/training.py
@@ -1152,12 +1152,12 @@ def _model_change_out_bias(
     _sample_func,
     _model_params,
 ):
-    old_bias = _model.get_out_bias()
+    old_bias = _model.out_bias
     _model.change_out_bias(
         _sample_func,
         bias_adjust_mode=_model_params.get("bias_adjust_mode", "change-by-statistic"),
     )
-    new_bias = _model.get_out_bias()
+    new_bias = _model.out_bias
 
     model_type_map = _model.get_type_map()
     sorter = np.argsort(model_type_map)

From 33b33be2cca1a355b5e3ee60c84baec95e17d268 Mon Sep 17 00:00:00 2001
From: anyangml <anyangpeng.ca@gmail.com>
Date: Sun, 7 Apr 2024 06:42:22 +0000
Subject: [PATCH 18/30] fix:UTs

---
 deepmd/pt/model/atomic_model/linear_atomic_model.py | 4 ++--
 deepmd/pt/model/model/make_model.py                 | 4 ++++
 2 files changed, 6 insertions(+), 2 deletions(-)

diff --git a/deepmd/pt/model/atomic_model/linear_atomic_model.py b/deepmd/pt/model/atomic_model/linear_atomic_model.py
index 5c1db5a550..6cede6a265 100644
--- a/deepmd/pt/model/atomic_model/linear_atomic_model.py
+++ b/deepmd/pt/model/atomic_model/linear_atomic_model.py
@@ -199,9 +199,9 @@ def forward_atomic(
                 mapping,
                 fparam,
                 aparam,
-            )["energy"]
+            )
             # apply bias to each individual model
-            ener_list.append(model.apply_out_stat(raw_ret, mapping[extended_atype]))
+            ener_list.append(model.apply_out_stat(raw_ret, mapping[extended_atype][:,:nloc])["energy"])
         weights = self._compute_weight(extended_coord, extended_atype, nlists_)
 
         fit_ret = {
diff --git a/deepmd/pt/model/model/make_model.py b/deepmd/pt/model/model/make_model.py
index f1a0334335..ab3f3d0a8f 100644
--- a/deepmd/pt/model/model/make_model.py
+++ b/deepmd/pt/model/model/make_model.py
@@ -172,6 +172,10 @@ def forward_common(
             model_predict = self.output_type_cast(model_predict, input_prec)
             return model_predict
 
+        @property
+        def out_bias(self)-> torch.Tensor:
+            return self.atomic_model.out_bias
+
         def change_out_bias(
             self,
             merged,

From 3b6f65c9a23f07059de33e71ffbe5316de16b9b4 Mon Sep 17 00:00:00 2001
From: "pre-commit-ci[bot]"
 <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Date: Sun, 7 Apr 2024 06:42:53 +0000
Subject: [PATCH 19/30] [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci
---
 deepmd/pt/model/atomic_model/linear_atomic_model.py | 6 +++++-
 deepmd/pt/model/model/make_model.py                 | 2 +-
 2 files changed, 6 insertions(+), 2 deletions(-)

diff --git a/deepmd/pt/model/atomic_model/linear_atomic_model.py b/deepmd/pt/model/atomic_model/linear_atomic_model.py
index 6cede6a265..e6faabf5c5 100644
--- a/deepmd/pt/model/atomic_model/linear_atomic_model.py
+++ b/deepmd/pt/model/atomic_model/linear_atomic_model.py
@@ -201,7 +201,11 @@ def forward_atomic(
                 aparam,
             )
             # apply bias to each individual model
-            ener_list.append(model.apply_out_stat(raw_ret, mapping[extended_atype][:,:nloc])["energy"])
+            ener_list.append(
+                model.apply_out_stat(raw_ret, mapping[extended_atype][:, :nloc])[
+                    "energy"
+                ]
+            )
         weights = self._compute_weight(extended_coord, extended_atype, nlists_)
 
         fit_ret = {
diff --git a/deepmd/pt/model/model/make_model.py b/deepmd/pt/model/model/make_model.py
index ab3f3d0a8f..5a1edb7c1f 100644
--- a/deepmd/pt/model/model/make_model.py
+++ b/deepmd/pt/model/model/make_model.py
@@ -173,7 +173,7 @@ def forward_common(
             return model_predict
 
         @property
-        def out_bias(self)-> torch.Tensor:
+        def out_bias(self) -> torch.Tensor:
             return self.atomic_model.out_bias
 
         def change_out_bias(

From 9f0979dea9067e19a0507e3dfda8f65e07b58853 Mon Sep 17 00:00:00 2001
From: anyangml <anyangpeng.ca@gmail.com>
Date: Sun, 7 Apr 2024 07:34:10 +0000
Subject: [PATCH 20/30] fix: UTs

---
 deepmd/pt/model/atomic_model/dp_atomic_model.py      | 3 +++
 deepmd/pt/model/atomic_model/linear_atomic_model.py  | 3 +++
 deepmd/pt/model/atomic_model/pairtab_atomic_model.py | 3 +++
 deepmd/pt/model/model/make_model.py                  | 5 ++---
 deepmd/pt/train/training.py                          | 8 ++++----
 5 files changed, 15 insertions(+), 7 deletions(-)

diff --git a/deepmd/pt/model/atomic_model/dp_atomic_model.py b/deepmd/pt/model/atomic_model/dp_atomic_model.py
index 14975732bc..5090cb2c4a 100644
--- a/deepmd/pt/model/atomic_model/dp_atomic_model.py
+++ b/deepmd/pt/model/atomic_model/dp_atomic_model.py
@@ -176,6 +176,9 @@ def forward_atomic(
         )
         return fit_ret
 
+    def get_out_bias(self) -> torch.Tensor:
+        return self.out_bias
+        
     def compute_or_load_stat(
         self,
         sampled_func,
diff --git a/deepmd/pt/model/atomic_model/linear_atomic_model.py b/deepmd/pt/model/atomic_model/linear_atomic_model.py
index e6faabf5c5..8fdb7c8a93 100644
--- a/deepmd/pt/model/atomic_model/linear_atomic_model.py
+++ b/deepmd/pt/model/atomic_model/linear_atomic_model.py
@@ -96,6 +96,9 @@ def mixed_types(self) -> bool:
         """
         return True
 
+    def get_out_bias(self) -> torch.Tensor:
+        return self.out_bias
+        
     def get_rcut(self) -> float:
         """Get the cut-off radius."""
         return max(self.get_model_rcuts())
diff --git a/deepmd/pt/model/atomic_model/pairtab_atomic_model.py b/deepmd/pt/model/atomic_model/pairtab_atomic_model.py
index 579e3efda4..81203348c2 100644
--- a/deepmd/pt/model/atomic_model/pairtab_atomic_model.py
+++ b/deepmd/pt/model/atomic_model/pairtab_atomic_model.py
@@ -136,6 +136,9 @@ def fitting_output_def(self) -> FittingOutputDef:
             ]
         )
 
+    def get_out_bias(self) -> torch.Tensor:
+        return self.out_bias
+        
     def get_rcut(self) -> float:
         return self.rcut
 
diff --git a/deepmd/pt/model/model/make_model.py b/deepmd/pt/model/model/make_model.py
index 5a1edb7c1f..fe6084c965 100644
--- a/deepmd/pt/model/model/make_model.py
+++ b/deepmd/pt/model/model/make_model.py
@@ -172,9 +172,8 @@ def forward_common(
             model_predict = self.output_type_cast(model_predict, input_prec)
             return model_predict
 
-        @property
-        def out_bias(self) -> torch.Tensor:
-            return self.atomic_model.out_bias
+        def get_out_bias(self)-> torch.Tensor:
+            return self.atomic_model.get_out_bias()
 
         def change_out_bias(
             self,
diff --git a/deepmd/pt/train/training.py b/deepmd/pt/train/training.py
index e198ee3d1a..f6fbfedf15 100644
--- a/deepmd/pt/train/training.py
+++ b/deepmd/pt/train/training.py
@@ -1152,12 +1152,12 @@ def _model_change_out_bias(
     _sample_func,
     _model_params,
 ):
-    old_bias = _model.out_bias
+    old_bias = _model.get_out_bias()
     _model.change_out_bias(
         _sample_func,
         bias_adjust_mode=_model_params.get("bias_adjust_mode", "change-by-statistic"),
     )
-    new_bias = _model.out_bias
+    new_bias = _model.get_out_bias()
 
     model_type_map = _model.get_type_map()
     sorter = np.argsort(model_type_map)
@@ -1168,7 +1168,7 @@ def _model_change_out_bias(
     idx_type_map = sorter[np.searchsorted(model_type_map, new_type_map, sorter=sorter)]
     log.info(
         f"Change output bias of {new_type_map!s} "
-        f"from {to_numpy_array(old_bias[idx_type_map]).reshape(-1)!s} "
-        f"to {to_numpy_array(new_bias[idx_type_map]).reshape(-1)!s}."
+        f"from {to_numpy_array(old_bias[:,idx_type_map]).reshape(-1)!s} "
+        f"to {to_numpy_array(new_bias[:,idx_type_map]).reshape(-1)!s}."
     )
     return _model

From 919dad16c5121cda2574f9086b189c92ed7f2771 Mon Sep 17 00:00:00 2001
From: "pre-commit-ci[bot]"
 <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Date: Sun, 7 Apr 2024 07:35:28 +0000
Subject: [PATCH 21/30] [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci
---
 deepmd/pt/model/atomic_model/dp_atomic_model.py      | 2 +-
 deepmd/pt/model/atomic_model/linear_atomic_model.py  | 2 +-
 deepmd/pt/model/atomic_model/pairtab_atomic_model.py | 2 +-
 deepmd/pt/model/model/make_model.py                  | 2 +-
 4 files changed, 4 insertions(+), 4 deletions(-)

diff --git a/deepmd/pt/model/atomic_model/dp_atomic_model.py b/deepmd/pt/model/atomic_model/dp_atomic_model.py
index 5090cb2c4a..6cbbf5ed7f 100644
--- a/deepmd/pt/model/atomic_model/dp_atomic_model.py
+++ b/deepmd/pt/model/atomic_model/dp_atomic_model.py
@@ -178,7 +178,7 @@ def forward_atomic(
 
     def get_out_bias(self) -> torch.Tensor:
         return self.out_bias
-        
+
     def compute_or_load_stat(
         self,
         sampled_func,
diff --git a/deepmd/pt/model/atomic_model/linear_atomic_model.py b/deepmd/pt/model/atomic_model/linear_atomic_model.py
index 8fdb7c8a93..9070cc81d1 100644
--- a/deepmd/pt/model/atomic_model/linear_atomic_model.py
+++ b/deepmd/pt/model/atomic_model/linear_atomic_model.py
@@ -98,7 +98,7 @@ def mixed_types(self) -> bool:
 
     def get_out_bias(self) -> torch.Tensor:
         return self.out_bias
-        
+
     def get_rcut(self) -> float:
         """Get the cut-off radius."""
         return max(self.get_model_rcuts())
diff --git a/deepmd/pt/model/atomic_model/pairtab_atomic_model.py b/deepmd/pt/model/atomic_model/pairtab_atomic_model.py
index 81203348c2..b4639fcbb4 100644
--- a/deepmd/pt/model/atomic_model/pairtab_atomic_model.py
+++ b/deepmd/pt/model/atomic_model/pairtab_atomic_model.py
@@ -138,7 +138,7 @@ def fitting_output_def(self) -> FittingOutputDef:
 
     def get_out_bias(self) -> torch.Tensor:
         return self.out_bias
-        
+
     def get_rcut(self) -> float:
         return self.rcut
 
diff --git a/deepmd/pt/model/model/make_model.py b/deepmd/pt/model/model/make_model.py
index fe6084c965..25a8ec9201 100644
--- a/deepmd/pt/model/model/make_model.py
+++ b/deepmd/pt/model/model/make_model.py
@@ -172,7 +172,7 @@ def forward_common(
             model_predict = self.output_type_cast(model_predict, input_prec)
             return model_predict
 
-        def get_out_bias(self)-> torch.Tensor:
+        def get_out_bias(self) -> torch.Tensor:
             return self.atomic_model.get_out_bias()
 
         def change_out_bias(

From 89fe0429e22a76e840b022996611cee865f37b47 Mon Sep 17 00:00:00 2001
From: anyangml <anyangpeng.ca@gmail.com>
Date: Sun, 7 Apr 2024 08:00:46 +0000
Subject: [PATCH 22/30] chore: remove bias_atom_e

---
 deepmd/dpmodel/atomic_model/pairtab_atomic_model.py | 1 -
 deepmd/pt/model/atomic_model/linear_atomic_model.py | 5 -----
 2 files changed, 6 deletions(-)

diff --git a/deepmd/dpmodel/atomic_model/pairtab_atomic_model.py b/deepmd/dpmodel/atomic_model/pairtab_atomic_model.py
index 51e46901e9..2d3bccb258 100644
--- a/deepmd/dpmodel/atomic_model/pairtab_atomic_model.py
+++ b/deepmd/dpmodel/atomic_model/pairtab_atomic_model.py
@@ -85,7 +85,6 @@ def __init__(
         else:
             self.tab_info, self.tab_data = None, None
 
-        self.bias_atom_e = np.zeros((self.ntypes, 1), dtype=np.float64)
         if isinstance(sel, int):
             self.sel = sel
         elif isinstance(sel, list):
diff --git a/deepmd/pt/model/atomic_model/linear_atomic_model.py b/deepmd/pt/model/atomic_model/linear_atomic_model.py
index 9070cc81d1..5740533810 100644
--- a/deepmd/pt/model/atomic_model/linear_atomic_model.py
+++ b/deepmd/pt/model/atomic_model/linear_atomic_model.py
@@ -73,11 +73,6 @@ def __init__(
             self.mapping_list.append(self.remap_atype(tpmp, self.type_map))
         assert len(err_msg) == 0, "\n".join(err_msg)
 
-        ntypes = len(self.type_map)
-        # a placeholder for storing weighted bias.
-        self.atomic_bias = torch.zeros(
-            ntypes, 1, dtype=env.GLOBAL_PT_ENER_FLOAT_PRECISION, device=env.DEVICE
-        )
         self.mixed_types_list = [model.mixed_types() for model in self.models]
         self.rcuts = torch.tensor(
             self.get_model_rcuts(), dtype=torch.float64, device=env.DEVICE

From 90b8d565faeb4daa7c3d7692ad9de0f144570987 Mon Sep 17 00:00:00 2001
From: Anyang Peng <137014849+anyangml@users.noreply.github.com>
Date: Sun, 7 Apr 2024 16:31:40 +0800
Subject: [PATCH 23/30] chore:use forward_common_atomic

---
 deepmd/pt/model/atomic_model/linear_atomic_model.py | 11 ++++-------
 1 file changed, 4 insertions(+), 7 deletions(-)

diff --git a/deepmd/pt/model/atomic_model/linear_atomic_model.py b/deepmd/pt/model/atomic_model/linear_atomic_model.py
index 5740533810..7fe23523a7 100644
--- a/deepmd/pt/model/atomic_model/linear_atomic_model.py
+++ b/deepmd/pt/model/atomic_model/linear_atomic_model.py
@@ -190,19 +190,16 @@ def forward_atomic(
 
         for i, model in enumerate(self.models):
             mapping = self.mapping_list[i]
-            raw_ret = model.forward_atomic(
+            # apply bias to each individual model
+            ener_list.append(
+                model.forward_common_atomic(
                 extended_coord,
                 mapping[extended_atype],
                 nlists_[i],
                 mapping,
                 fparam,
                 aparam,
-            )
-            # apply bias to each individual model
-            ener_list.append(
-                model.apply_out_stat(raw_ret, mapping[extended_atype][:, :nloc])[
-                    "energy"
-                ]
+            )["energy"]
             )
         weights = self._compute_weight(extended_coord, extended_atype, nlists_)
 

From 1ef891e96847449a5a35c785e27ad7aa2b3b7b57 Mon Sep 17 00:00:00 2001
From: "pre-commit-ci[bot]"
 <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Date: Sun, 7 Apr 2024 08:34:21 +0000
Subject: [PATCH 24/30] [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci
---
 .../pt/model/atomic_model/linear_atomic_model.py   | 14 +++++++-------
 1 file changed, 7 insertions(+), 7 deletions(-)

diff --git a/deepmd/pt/model/atomic_model/linear_atomic_model.py b/deepmd/pt/model/atomic_model/linear_atomic_model.py
index 7fe23523a7..87b58388f2 100644
--- a/deepmd/pt/model/atomic_model/linear_atomic_model.py
+++ b/deepmd/pt/model/atomic_model/linear_atomic_model.py
@@ -193,13 +193,13 @@ def forward_atomic(
             # apply bias to each individual model
             ener_list.append(
                 model.forward_common_atomic(
-                extended_coord,
-                mapping[extended_atype],
-                nlists_[i],
-                mapping,
-                fparam,
-                aparam,
-            )["energy"]
+                    extended_coord,
+                    mapping[extended_atype],
+                    nlists_[i],
+                    mapping,
+                    fparam,
+                    aparam,
+                )["energy"]
             )
         weights = self._compute_weight(extended_coord, extended_atype, nlists_)
 

From f32614fd30100e44544a7f2ae73962a4ec564634 Mon Sep 17 00:00:00 2001
From: Anyang Peng <137014849+anyangml@users.noreply.github.com>
Date: Sun, 7 Apr 2024 20:29:42 +0800
Subject: [PATCH 25/30] feat: add UTs

---
 .../atomic_model/linear_atomic_model.py       |   3 +-
 .../model/atomic_model/linear_atomic_model.py |  75 ++++--
 .../pt/model/test_linear_atomic_model_stat.py | 229 ++++++++++++++++++
 3 files changed, 282 insertions(+), 25 deletions(-)
 create mode 100644 source/tests/pt/model/test_linear_atomic_model_stat.py

diff --git a/deepmd/dpmodel/atomic_model/linear_atomic_model.py b/deepmd/dpmodel/atomic_model/linear_atomic_model.py
index e6296316a5..e4a85d7bc2 100644
--- a/deepmd/dpmodel/atomic_model/linear_atomic_model.py
+++ b/deepmd/dpmodel/atomic_model/linear_atomic_model.py
@@ -248,7 +248,8 @@ def _compute_weight(
     ) -> List[np.ndarray]:
         """This should be a list of user defined weights that matches the number of models to be combined."""
         nmodels = len(self.models)
-        return [np.ones(1) / nmodels for _ in range(nmodels)]
+        nframes, nloc, _ = nlists_[0].shape
+        return [np.ones((nframes, nloc, 1)) / nmodels for _ in range(nmodels)]
 
     def get_dim_fparam(self) -> int:
         """Get the number (dimension) of frame parameters of this atomic model."""
diff --git a/deepmd/pt/model/atomic_model/linear_atomic_model.py b/deepmd/pt/model/atomic_model/linear_atomic_model.py
index 87b58388f2..9bf033953e 100644
--- a/deepmd/pt/model/atomic_model/linear_atomic_model.py
+++ b/deepmd/pt/model/atomic_model/linear_atomic_model.py
@@ -5,6 +5,8 @@
     List,
     Optional,
     Tuple,
+    Union,
+    Callable,
 )
 
 import torch
@@ -293,8 +295,9 @@ def _compute_weight(
     ) -> List[torch.Tensor]:
         """This should be a list of user defined weights that matches the number of models to be combined."""
         nmodels = len(self.models)
+        nframes, nloc, _ = nlists_[0].shape
         return [
-            torch.ones(1, dtype=torch.float64, device=env.DEVICE) / nmodels
+            torch.ones((nframes, nloc, 1), dtype=torch.float64, device=env.DEVICE) / nmodels
             for _ in range(nmodels)
         ]
 
@@ -333,6 +336,53 @@ def is_aparam_nall(self) -> bool:
         If False, the shape is (nframes, nloc, ndim).
         """
         return False
+    
+    def compute_or_load_out_stat(
+        self,
+        merged: Union[Callable[[], List[dict]], List[dict]],
+        stat_file_path: Optional[DPPath] = None,
+    ):
+        """
+        Compute the output statistics (e.g. energy bias) for the fitting net from packed data.
+
+        Parameters
+        ----------
+        merged : Union[Callable[[], List[dict]], List[dict]]
+            - List[dict]: A list of data samples from various data systems.
+                Each element, `merged[i]`, is a data dictionary containing `keys`: `torch.Tensor`
+                originating from the `i`-th data system.
+            - Callable[[], List[dict]]: A lazy function that returns data samples in the above format
+                only when needed. Since the sampling process can be slow and memory-intensive,
+                the lazy function helps by only sampling once.
+        stat_file_path : Optional[DPPath]
+            The path to the stat file.
+
+        """
+        for md in self.models:
+            md.compute_or_load_out_stat(merged, stat_file_path)
+    
+    def compute_or_load_stat(
+        self,
+        sampled_func,
+        stat_file_path: Optional[DPPath] = None,
+    ):
+        """
+        Compute or load the statistics parameters of the model,
+        such as mean and standard deviation of descriptors or the energy bias of the fitting net.
+        When `sampled` is provided, all the statistics parameters will be calculated (or re-calculated for update),
+        and saved in the `stat_file_path`(s).
+        When `sampled` is not provided, it will check the existence of `stat_file_path`(s)
+        and load the calculated statistics parameters.
+
+        Parameters
+        ----------
+        sampled_func
+            The lazy sampled function to get data frames from different data systems.
+        stat_file_path
+            The dictionary of paths to the statistics files.
+        """
+        for md in self.models:
+            md.compute_or_load_stat(sampled_func, stat_file_path)
 
 
 class DPZBLLinearEnergyAtomicModel(LinearEnergyAtomicModel):
@@ -376,29 +426,6 @@ def __init__(
         # this is a placeholder being updated in _compute_weight, to handle Jit attribute init error.
         self.zbl_weight = torch.empty(0, dtype=torch.float64, device=env.DEVICE)
 
-    def compute_or_load_stat(
-        self,
-        sampled_func,
-        stat_file_path: Optional[DPPath] = None,
-    ):
-        """
-        Compute or load the statistics parameters of the model,
-        such as mean and standard deviation of descriptors or the energy bias of the fitting net.
-        When `sampled` is provided, all the statistics parameters will be calculated (or re-calculated for update),
-        and saved in the `stat_file_path`(s).
-        When `sampled` is not provided, it will check the existence of `stat_file_path`(s)
-        and load the calculated statistics parameters.
-
-        Parameters
-        ----------
-        sampled_func
-            The lazy sampled function to get data frames from different data systems.
-        stat_file_path
-            The dictionary of paths to the statistics files.
-        """
-        self.models[0].compute_or_load_stat(sampled_func, stat_file_path)
-        self.models[1].compute_or_load_stat(sampled_func, stat_file_path)
-
     def serialize(self) -> dict:
         dd = BaseAtomicModel.serialize(self)
         dd.update(
diff --git a/source/tests/pt/model/test_linear_atomic_model_stat.py b/source/tests/pt/model/test_linear_atomic_model_stat.py
new file mode 100644
index 0000000000..ae1ca84419
--- /dev/null
+++ b/source/tests/pt/model/test_linear_atomic_model_stat.py
@@ -0,0 +1,229 @@
+# SPDX-License-Identifier: LGPL-3.0-or-later
+import tempfile
+import unittest
+from pathlib import (
+    Path,
+)
+from typing import (
+    Optional,
+)
+
+import h5py
+import numpy as np
+import torch
+
+from deepmd.dpmodel.output_def import (
+    FittingOutputDef,
+    OutputVariableDef,
+)
+from deepmd.pt.model.atomic_model import (
+    BaseAtomicModel,
+    DPAtomicModel,
+    LinearEnergyAtomicModel,
+)
+from deepmd.pt.model.descriptor.dpa1 import (
+    DescrptDPA1,
+)
+from deepmd.pt.model.task.base_fitting import (
+    BaseFitting,
+)
+from deepmd.pt.utils import (
+    env,
+)
+from deepmd.pt.utils.utils import (
+    to_numpy_array,
+    to_torch_tensor,
+)
+from deepmd.utils.path import (
+    DPPath,
+)
+
+from .test_env_mat import (
+    TestCaseSingleFrameWithNlist,
+)
+
+dtype = env.GLOBAL_PT_FLOAT_PRECISION
+
+
+class FooFittingA(torch.nn.Module, BaseFitting):
+    def output_def(self):
+        return FittingOutputDef(
+            [
+                OutputVariableDef(
+                    "energy",
+                    [1],
+                    reduciable=True,
+                    r_differentiable=True,
+                    c_differentiable=True,
+                ),
+            ]
+        )
+
+    def serialize(self) -> dict:
+        raise NotImplementedError
+
+    def forward(
+        self,
+        descriptor: torch.Tensor,
+        atype: torch.Tensor,
+        gr: Optional[torch.Tensor] = None,
+        g2: Optional[torch.Tensor] = None,
+        h2: Optional[torch.Tensor] = None,
+        fparam: Optional[torch.Tensor] = None,
+        aparam: Optional[torch.Tensor] = None,
+    ):
+        nf, nloc, _ = descriptor.shape
+        ret = {}
+        ret["energy"] = (
+            torch.Tensor(
+                [
+                    [1.0, 2.0, 3.0],
+                    [4.0, 5.0, 6.0],
+                ]
+            )
+            .view([nf, nloc] + self.output_def()["energy"].shape)
+            .to(env.GLOBAL_PT_FLOAT_PRECISION)
+            .to(env.DEVICE)
+        )
+        
+        return ret
+
+class FooFittingB(torch.nn.Module, BaseFitting):
+    def output_def(self):
+        return FittingOutputDef(
+            [
+                OutputVariableDef(
+                    "energy",
+                    [1],
+                    reduciable=True,
+                    r_differentiable=True,
+                    c_differentiable=True,
+                ),
+            ]
+        )
+
+    def serialize(self) -> dict:
+        raise NotImplementedError
+
+    def forward(
+        self,
+        descriptor: torch.Tensor,
+        atype: torch.Tensor,
+        gr: Optional[torch.Tensor] = None,
+        g2: Optional[torch.Tensor] = None,
+        h2: Optional[torch.Tensor] = None,
+        fparam: Optional[torch.Tensor] = None,
+        aparam: Optional[torch.Tensor] = None,
+    ):
+        nf, nloc, _ = descriptor.shape
+        ret = {}
+        ret["energy"] = (
+            torch.Tensor(
+                [
+                    [7.0, 8.0, 9.0],
+                    [10.0, 11.0, 12.0],
+                ]
+            )
+            .view([nf, nloc] + self.output_def()["energy"].shape)
+            .to(env.GLOBAL_PT_FLOAT_PRECISION)
+            .to(env.DEVICE)
+        )
+        
+        return ret
+
+class TestAtomicModelStat(unittest.TestCase, TestCaseSingleFrameWithNlist):
+    def tearDown(self):
+        self.tempdir.cleanup()
+
+    def setUp(self):
+        TestCaseSingleFrameWithNlist.setUp(self)
+        nf, nloc, nnei = self.nlist.shape
+        self.merged_output_stat = [
+            {
+                "coord": to_torch_tensor(np.zeros([2, 3, 3])),
+                "atype": to_torch_tensor(
+                    np.array([[0, 0, 1], [0, 1, 1]], dtype=np.int32)
+                ),
+                "atype_ext": to_torch_tensor(
+                    np.array([[0, 0, 1, 0], [0, 1, 1, 0]], dtype=np.int32)
+                ),
+                "box": to_torch_tensor(np.zeros([2, 3, 3])),
+                "natoms": to_torch_tensor(
+                    np.array([[3, 3, 2, 1], [3, 3, 1, 2]], dtype=np.int32)
+                ),
+                # bias of foo: 1, 3
+                "energy": to_torch_tensor(np.array([5.0, 7.0]).reshape(2, 1)),
+
+            }
+        ]
+        self.tempdir = tempfile.TemporaryDirectory()
+        h5file = str((Path(self.tempdir.name) / "testcase.h5").resolve())
+        with h5py.File(h5file, "w") as f:
+            pass
+        self.stat_file_path = DPPath(h5file, "a")
+
+    def test_linear_atomic_model_stat_with_bias(self):
+        nf, nloc, nnei = self.nlist.shape
+        ds = DescrptDPA1(
+            self.rcut,
+            self.rcut_smth,
+            sum(self.sel),
+            self.nt,
+        ).to(env.DEVICE)
+        ft_a = FooFittingA().to(env.DEVICE)
+        ft_b = FooFittingB().to(env.DEVICE)
+        type_map = ["foo", "bar"]
+        md0 = DPAtomicModel(
+            ds,
+            ft_a,
+            type_map=type_map,
+        ).to(env.DEVICE)
+        md1 = DPAtomicModel(
+            ds,
+            ft_b,
+            type_map=type_map,
+        ).to(env.DEVICE)
+        linear_model = LinearEnergyAtomicModel(
+            [md0,md1],type_map=type_map
+        ).to(env.DEVICE)
+
+        args = [
+            to_torch_tensor(ii) for ii in [self.coord_ext, self.atype_ext, self.nlist]
+        ]
+        # nf x nloc
+        at = self.atype_ext[:, :nloc]
+
+        # 1. test run without bias
+        # nf x na x odim
+        ret0 = linear_model.forward_common_atomic(*args)
+        
+        ret0 = to_numpy_array(ret0["energy"])
+        ret_no_bias = []
+        for md in linear_model.models:
+            ret_no_bias.append(md.forward_common_atomic(*args)["energy"])
+        expected_ret0 = np.array(
+            [
+                [4.0, 5.0, 6.0],
+                [7.0, 8.0, 9.0],
+            ]
+        ).reshape([nf, nloc] + linear_model.fitting_output_def()["energy"].shape)
+        
+        np.testing.assert_almost_equal(ret0, expected_ret0)
+
+        # 2. test bias is applied
+        linear_model.compute_or_load_out_stat(
+            self.merged_output_stat, stat_file_path=self.stat_file_path
+        )
+        # bias applied to sub atomic models.
+        ener_bias = np.array([1.0, 3.0]).reshape(2, 1)
+        linear_ret = []
+        for idx, md in enumerate(linear_model.models):
+            ret = md.forward_common_atomic(*args)
+            ret = to_numpy_array(ret["energy"])
+            linear_ret.append(ret_no_bias[idx] + ener_bias[at])
+            np.testing.assert_almost_equal((ret_no_bias[idx] + ener_bias[at]), ret)
+        
+        # linear model not adding bias again
+        ret1 = linear_model.forward_common_atomic(*args)
+        ret1 = to_numpy_array(ret1["energy"])
+        np.testing.assert_almost_equal(torch.mean(torch.stack(linear_ret),dim=0), ret1)

From 1b233e737908925a46d941920cbe651f960d125b Mon Sep 17 00:00:00 2001
From: "pre-commit-ci[bot]"
 <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Date: Sun, 7 Apr 2024 12:30:33 +0000
Subject: [PATCH 26/30] [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci
---
 .../model/atomic_model/linear_atomic_model.py |  9 ++++----
 .../pt/model/test_linear_atomic_model_stat.py | 22 +++++++++----------
 2 files changed, 16 insertions(+), 15 deletions(-)

diff --git a/deepmd/pt/model/atomic_model/linear_atomic_model.py b/deepmd/pt/model/atomic_model/linear_atomic_model.py
index 9bf033953e..c5abc4575c 100644
--- a/deepmd/pt/model/atomic_model/linear_atomic_model.py
+++ b/deepmd/pt/model/atomic_model/linear_atomic_model.py
@@ -1,12 +1,12 @@
 # SPDX-License-Identifier: LGPL-3.0-or-later
 import copy
 from typing import (
+    Callable,
     Dict,
     List,
     Optional,
     Tuple,
     Union,
-    Callable,
 )
 
 import torch
@@ -297,7 +297,8 @@ def _compute_weight(
         nmodels = len(self.models)
         nframes, nloc, _ = nlists_[0].shape
         return [
-            torch.ones((nframes, nloc, 1), dtype=torch.float64, device=env.DEVICE) / nmodels
+            torch.ones((nframes, nloc, 1), dtype=torch.float64, device=env.DEVICE)
+            / nmodels
             for _ in range(nmodels)
         ]
 
@@ -336,7 +337,7 @@ def is_aparam_nall(self) -> bool:
         If False, the shape is (nframes, nloc, ndim).
         """
         return False
-    
+
     def compute_or_load_out_stat(
         self,
         merged: Union[Callable[[], List[dict]], List[dict]],
@@ -360,7 +361,7 @@ def compute_or_load_out_stat(
         """
         for md in self.models:
             md.compute_or_load_out_stat(merged, stat_file_path)
-    
+
     def compute_or_load_stat(
         self,
         sampled_func,
diff --git a/source/tests/pt/model/test_linear_atomic_model_stat.py b/source/tests/pt/model/test_linear_atomic_model_stat.py
index ae1ca84419..d8137c17c0 100644
--- a/source/tests/pt/model/test_linear_atomic_model_stat.py
+++ b/source/tests/pt/model/test_linear_atomic_model_stat.py
@@ -17,7 +17,6 @@
     OutputVariableDef,
 )
 from deepmd.pt.model.atomic_model import (
-    BaseAtomicModel,
     DPAtomicModel,
     LinearEnergyAtomicModel,
 )
@@ -85,9 +84,10 @@ def forward(
             .to(env.GLOBAL_PT_FLOAT_PRECISION)
             .to(env.DEVICE)
         )
-        
+
         return ret
 
+
 class FooFittingB(torch.nn.Module, BaseFitting):
     def output_def(self):
         return FittingOutputDef(
@@ -128,9 +128,10 @@ def forward(
             .to(env.GLOBAL_PT_FLOAT_PRECISION)
             .to(env.DEVICE)
         )
-        
+
         return ret
 
+
 class TestAtomicModelStat(unittest.TestCase, TestCaseSingleFrameWithNlist):
     def tearDown(self):
         self.tempdir.cleanup()
@@ -153,7 +154,6 @@ def setUp(self):
                 ),
                 # bias of foo: 1, 3
                 "energy": to_torch_tensor(np.array([5.0, 7.0]).reshape(2, 1)),
-
             }
         ]
         self.tempdir = tempfile.TemporaryDirectory()
@@ -183,9 +183,9 @@ def test_linear_atomic_model_stat_with_bias(self):
             ft_b,
             type_map=type_map,
         ).to(env.DEVICE)
-        linear_model = LinearEnergyAtomicModel(
-            [md0,md1],type_map=type_map
-        ).to(env.DEVICE)
+        linear_model = LinearEnergyAtomicModel([md0, md1], type_map=type_map).to(
+            env.DEVICE
+        )
 
         args = [
             to_torch_tensor(ii) for ii in [self.coord_ext, self.atype_ext, self.nlist]
@@ -196,7 +196,7 @@ def test_linear_atomic_model_stat_with_bias(self):
         # 1. test run without bias
         # nf x na x odim
         ret0 = linear_model.forward_common_atomic(*args)
-        
+
         ret0 = to_numpy_array(ret0["energy"])
         ret_no_bias = []
         for md in linear_model.models:
@@ -207,7 +207,7 @@ def test_linear_atomic_model_stat_with_bias(self):
                 [7.0, 8.0, 9.0],
             ]
         ).reshape([nf, nloc] + linear_model.fitting_output_def()["energy"].shape)
-        
+
         np.testing.assert_almost_equal(ret0, expected_ret0)
 
         # 2. test bias is applied
@@ -222,8 +222,8 @@ def test_linear_atomic_model_stat_with_bias(self):
             ret = to_numpy_array(ret["energy"])
             linear_ret.append(ret_no_bias[idx] + ener_bias[at])
             np.testing.assert_almost_equal((ret_no_bias[idx] + ener_bias[at]), ret)
-        
+
         # linear model not adding bias again
         ret1 = linear_model.forward_common_atomic(*args)
         ret1 = to_numpy_array(ret1["energy"])
-        np.testing.assert_almost_equal(torch.mean(torch.stack(linear_ret),dim=0), ret1)
+        np.testing.assert_almost_equal(torch.mean(torch.stack(linear_ret), dim=0), ret1)

From 97bb808d3631a04b39b581a921e9c9ea18a097ba Mon Sep 17 00:00:00 2001
From: Anyang Peng <137014849+anyangml@users.noreply.github.com>
Date: Sun, 7 Apr 2024 20:42:09 +0800
Subject: [PATCH 27/30] fix: precommit

---
 source/tests/pt/model/test_linear_atomic_model_stat.py | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

diff --git a/source/tests/pt/model/test_linear_atomic_model_stat.py b/source/tests/pt/model/test_linear_atomic_model_stat.py
index d8137c17c0..f3fe4f63d4 100644
--- a/source/tests/pt/model/test_linear_atomic_model_stat.py
+++ b/source/tests/pt/model/test_linear_atomic_model_stat.py
@@ -80,7 +80,7 @@ def forward(
                     [4.0, 5.0, 6.0],
                 ]
             )
-            .view([nf, nloc] + self.output_def()["energy"].shape)
+            .view([nf, nloc, *self.output_def()["energy"].shape])
             .to(env.GLOBAL_PT_FLOAT_PRECISION)
             .to(env.DEVICE)
         )
@@ -124,7 +124,7 @@ def forward(
                     [10.0, 11.0, 12.0],
                 ]
             )
-            .view([nf, nloc] + self.output_def()["energy"].shape)
+            .view([nf, nloc, *self.output_def()["energy"].shape])
             .to(env.GLOBAL_PT_FLOAT_PRECISION)
             .to(env.DEVICE)
         )
@@ -206,8 +206,8 @@ def test_linear_atomic_model_stat_with_bias(self):
                 [4.0, 5.0, 6.0],
                 [7.0, 8.0, 9.0],
             ]
-        ).reshape([nf, nloc] + linear_model.fitting_output_def()["energy"].shape)
-
+        ).reshape(nf, nloc,  *linear_model.fitting_output_def()["energy"].shape)
+        
         np.testing.assert_almost_equal(ret0, expected_ret0)
 
         # 2. test bias is applied

From 2dec4cde008c5084cb055977ece5be937371e0db Mon Sep 17 00:00:00 2001
From: "pre-commit-ci[bot]"
 <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Date: Sun, 7 Apr 2024 12:43:34 +0000
Subject: [PATCH 28/30] [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci
---
 source/tests/pt/model/test_linear_atomic_model_stat.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/source/tests/pt/model/test_linear_atomic_model_stat.py b/source/tests/pt/model/test_linear_atomic_model_stat.py
index f3fe4f63d4..490d777cb1 100644
--- a/source/tests/pt/model/test_linear_atomic_model_stat.py
+++ b/source/tests/pt/model/test_linear_atomic_model_stat.py
@@ -206,8 +206,8 @@ def test_linear_atomic_model_stat_with_bias(self):
                 [4.0, 5.0, 6.0],
                 [7.0, 8.0, 9.0],
             ]
-        ).reshape(nf, nloc,  *linear_model.fitting_output_def()["energy"].shape)
-        
+        ).reshape(nf, nloc, *linear_model.fitting_output_def()["energy"].shape)
+
         np.testing.assert_almost_equal(ret0, expected_ret0)
 
         # 2. test bias is applied

From 666993981c83b22665810a6324425b7d68a202e2 Mon Sep 17 00:00:00 2001
From: anyangml <anyangpeng.ca@gmail.com>
Date: Mon, 8 Apr 2024 02:17:17 +0000
Subject: [PATCH 29/30] fix: cuda

---
 source/tests/pt/model/test_linear_atomic_model_stat.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/source/tests/pt/model/test_linear_atomic_model_stat.py b/source/tests/pt/model/test_linear_atomic_model_stat.py
index 490d777cb1..95632f373a 100644
--- a/source/tests/pt/model/test_linear_atomic_model_stat.py
+++ b/source/tests/pt/model/test_linear_atomic_model_stat.py
@@ -200,7 +200,7 @@ def test_linear_atomic_model_stat_with_bias(self):
         ret0 = to_numpy_array(ret0["energy"])
         ret_no_bias = []
         for md in linear_model.models:
-            ret_no_bias.append(md.forward_common_atomic(*args)["energy"])
+            ret_no_bias.append(to_numpy_array(md.forward_common_atomic(*args)["energy"]))
         expected_ret0 = np.array(
             [
                 [4.0, 5.0, 6.0],
@@ -226,4 +226,4 @@ def test_linear_atomic_model_stat_with_bias(self):
         # linear model not adding bias again
         ret1 = linear_model.forward_common_atomic(*args)
         ret1 = to_numpy_array(ret1["energy"])
-        np.testing.assert_almost_equal(torch.mean(torch.stack(linear_ret), dim=0), ret1)
+        np.testing.assert_almost_equal(np.mean(np.stack(linear_ret), axis=0), ret1)

From c8010bc2dddc423e8ec95b1d0dddb3e2f146b254 Mon Sep 17 00:00:00 2001
From: "pre-commit-ci[bot]"
 <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Date: Mon, 8 Apr 2024 02:17:49 +0000
Subject: [PATCH 30/30] [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci
---
 source/tests/pt/model/test_linear_atomic_model_stat.py | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/source/tests/pt/model/test_linear_atomic_model_stat.py b/source/tests/pt/model/test_linear_atomic_model_stat.py
index 95632f373a..010cecf9f8 100644
--- a/source/tests/pt/model/test_linear_atomic_model_stat.py
+++ b/source/tests/pt/model/test_linear_atomic_model_stat.py
@@ -200,7 +200,9 @@ def test_linear_atomic_model_stat_with_bias(self):
         ret0 = to_numpy_array(ret0["energy"])
         ret_no_bias = []
         for md in linear_model.models:
-            ret_no_bias.append(to_numpy_array(md.forward_common_atomic(*args)["energy"]))
+            ret_no_bias.append(
+                to_numpy_array(md.forward_common_atomic(*args)["energy"])
+            )
         expected_ret0 = np.array(
             [
                 [4.0, 5.0, 6.0],