From 98519132b42a33dc9bce60326ce8c2039ba554fb Mon Sep 17 00:00:00 2001
From: Anyang Peng <137014849+anyangml@users.noreply.github.com>
Date: Tue, 20 Feb 2024 15:39:58 +0800
Subject: [PATCH 01/14] feat: add dipole grad

---
 deepmd/dpmodel/fitting/dipole_fitting.py      | 19 ++++-
 .../dpmodel/fitting/polarizability_fitting.py |  4 +-
 deepmd/pt/model/model/dipole_model.py         | 77 +++++++++++++++++++
 deepmd/pt/model/task/dipole.py                | 16 +++-
 deepmd/pt/model/task/polarizability.py        |  4 +-
 source/tests/pt/model/test_dipole_fitting.py  | 50 +++++++++++-
 6 files changed, 159 insertions(+), 11 deletions(-)
 create mode 100644 deepmd/pt/model/model/dipole_model.py

diff --git a/deepmd/dpmodel/fitting/dipole_fitting.py b/deepmd/dpmodel/fitting/dipole_fitting.py
index c210945e76..69bdac1dac 100644
--- a/deepmd/dpmodel/fitting/dipole_fitting.py
+++ b/deepmd/dpmodel/fitting/dipole_fitting.py
@@ -68,9 +68,14 @@ class DipoleFitting(GeneralFitting):
     mixed_types
             If true, use a uniform fitting net for all atom types, otherwise use
             different fitting nets for different atom types.
-    exclude_types: List[int]
+    exclude_types
             Atomic contributions of the excluded atom types are set zero.
-
+    r_differentiable
+            If the variable is differentiated with respect to coordinates of atoms. 
+            Only reduciable variable are differentiable.
+    c_differentiable
+            If the variable is differentiated with respect to the cell tensor (pbc case). 
+            Only reduciable variable are differentiable.
     """
 
     def __init__(
@@ -94,6 +99,8 @@ def __init__(
         spin: Any = None,
         mixed_types: bool = False,
         exclude_types: List[int] = [],
+        r_differentiable: bool = True,
+        c_differentiable: bool = False,
         old_impl=False,
     ):
         # seed, uniform_seed are not included
@@ -109,6 +116,8 @@ def __init__(
             raise NotImplementedError("atom_ener is not implemented")
 
         self.embedding_width = embedding_width
+        self.r_differentiable = r_differentiable
+        self.c_differentiable = c_differentiable
         super().__init__(
             var_name=var_name,
             ntypes=ntypes,
@@ -139,6 +148,8 @@ def serialize(self) -> dict:
         data = super().serialize()
         data["embedding_width"] = self.embedding_width
         data["old_impl"] = self.old_impl
+        data["r_differentiable"] = self.r_differentiable
+        data["c_differentiable"] = self.c_differentiable
         return data
 
     def output_def(self):
@@ -148,8 +159,8 @@ def output_def(self):
                     self.var_name,
                     [3],
                     reduciable=True,
-                    r_differentiable=True,
-                    c_differentiable=True,
+                    r_differentiable=self.r_differentiable,
+                    c_differentiable=self.c_differentiable,
                 ),
             ]
         )
diff --git a/deepmd/dpmodel/fitting/polarizability_fitting.py b/deepmd/dpmodel/fitting/polarizability_fitting.py
index d828693fe0..c3cbe7bd1a 100644
--- a/deepmd/dpmodel/fitting/polarizability_fitting.py
+++ b/deepmd/dpmodel/fitting/polarizability_fitting.py
@@ -176,8 +176,8 @@ def output_def(self):
                     self.var_name,
                     [3, 3],
                     reduciable=True,
-                    r_differentiable=True,
-                    c_differentiable=True,
+                    r_differentiable=False,
+                    c_differentiable=False,
                 ),
             ]
         )
diff --git a/deepmd/pt/model/model/dipole_model.py b/deepmd/pt/model/model/dipole_model.py
new file mode 100644
index 0000000000..fda5940939
--- /dev/null
+++ b/deepmd/pt/model/model/dipole_model.py
@@ -0,0 +1,77 @@
+# SPDX-License-Identifier: LGPL-3.0-or-later
+from typing import (
+    Dict,
+    Optional,
+)
+
+import torch
+
+from .dp_model import (
+    DPModel,
+)
+
+
+class DipoleModel(DPModel):
+    model_type = "dipole"
+
+    def __init__(
+        self,
+        *args,
+        **kwargs,
+    ):
+        super().__init__(*args, **kwargs)
+
+    def forward(
+        self,
+        coord,
+        atype,
+        box: Optional[torch.Tensor] = None,
+        fparam: Optional[torch.Tensor] = None,
+        aparam: Optional[torch.Tensor] = None,
+        do_atomic_virial: bool = False,
+    ) -> Dict[str, torch.Tensor]:
+        model_ret = self.forward_common(
+            coord,
+            atype,
+            box,
+            fparam=fparam,
+            aparam=aparam,
+            do_atomic_virial=False,
+        )
+        if self.fitting_net is not None:
+            model_predict = {}
+            model_predict["atom_dipole"] = model_ret["dipole"]
+            model_predict["dipole"] = model_ret["dipole_redu"]
+            if self.do_grad("dipole"):
+                model_predict["dipole_grad_r"] = model_ret["dipole_derv_r"].squeeze(-2)
+        else:
+            model_predict = model_ret
+            model_predict["updated_coord"] += coord
+        return model_predict
+
+    def forward_lower(
+        self,
+        extended_coord,
+        extended_atype,
+        nlist,
+        mapping: Optional[torch.Tensor] = None,
+        fparam: Optional[torch.Tensor] = None,
+        aparam: Optional[torch.Tensor] = None,
+        do_atomic_virial: bool = False,
+    ):
+        model_ret = self.forward_common_lower(
+            extended_coord,
+            extended_atype,
+            nlist,
+            mapping,
+            fparam=fparam,
+            aparam=aparam,
+            do_atomic_virial=False,
+        )
+        if self.fitting_net is not None:
+            model_predict = {}
+            model_predict["atom_dipole"] = model_ret["dipole"]
+            model_predict["energy"] = model_ret["energy_redu"]
+        else:
+            model_predict = model_ret
+        return model_predict
diff --git a/deepmd/pt/model/task/dipole.py b/deepmd/pt/model/task/dipole.py
index 4ea66e2636..562bff8211 100644
--- a/deepmd/pt/model/task/dipole.py
+++ b/deepmd/pt/model/task/dipole.py
@@ -56,6 +56,12 @@ class DipoleFittingNet(GeneralFitting):
         The condition number for the regression of atomic energy.
     seed : int, optional
         Random seed.
+    r_differentiable
+        If the variable is differentiated with respect to coordinates of atoms. 
+        Only reduciable variable are differentiable.
+    c_differentiable
+        If the variable is differentiated with respect to the cell tensor (pbc case). 
+        Only reduciable variable are differentiable.
     """
 
     def __init__(
@@ -74,9 +80,13 @@ def __init__(
         rcond: Optional[float] = None,
         seed: Optional[int] = None,
         exclude_types: List[int] = [],
+        r_differentiable: bool = True,
+        c_differentiable: bool = False,
         **kwargs,
     ):
         self.embedding_width = embedding_width
+        self.r_differentiable = r_differentiable
+        self.c_differentiable = c_differentiable
         super().__init__(
             var_name=var_name,
             ntypes=ntypes,
@@ -103,6 +113,8 @@ def serialize(self) -> dict:
         data = super().serialize()
         data["embedding_width"] = self.embedding_width
         data["old_impl"] = self.old_impl
+        data["r_differentiable"] = self.r_differentiable
+        data["c_differentiable"] = self.c_differentiable
         return data
 
     def output_def(self) -> FittingOutputDef:
@@ -112,8 +124,8 @@ def output_def(self) -> FittingOutputDef:
                     self.var_name,
                     [3],
                     reduciable=True,
-                    r_differentiable=True,
-                    c_differentiable=True,
+                    r_differentiable=self.r_differentiable,
+                    c_differentiable=self.c_differentiable,
                 ),
             ]
         )
diff --git a/deepmd/pt/model/task/polarizability.py b/deepmd/pt/model/task/polarizability.py
index dc8d13ee84..c240567903 100644
--- a/deepmd/pt/model/task/polarizability.py
+++ b/deepmd/pt/model/task/polarizability.py
@@ -144,8 +144,8 @@ def output_def(self) -> FittingOutputDef:
                     self.var_name,
                     [3, 3],
                     reduciable=True,
-                    r_differentiable=True,
-                    c_differentiable=True,
+                    r_differentiable=False,
+                    c_differentiable=False,
                 ),
             ]
         )
diff --git a/source/tests/pt/model/test_dipole_fitting.py b/source/tests/pt/model/test_dipole_fitting.py
index 3f67043767..c816e80c0c 100644
--- a/source/tests/pt/model/test_dipole_fitting.py
+++ b/source/tests/pt/model/test_dipole_fitting.py
@@ -28,9 +28,24 @@
 from .test_env_mat import (
     TestCaseSingleFrameWithNlist,
 )
+from deepmd.pt.model.model.dipole_model import (
+    DipoleModel
+)
 
 dtype = env.GLOBAL_PT_FLOAT_PRECISION
 
+def finite_difference(f, x, a, delta=1e-6):
+    in_shape = x.shape
+    y0 = f(x, a)
+    out_shape = y0.shape
+    res = np.empty(out_shape + in_shape)
+    for idx in np.ndindex(*in_shape):
+        diff = np.zeros(in_shape)
+        diff[idx] += delta
+        y1p = f(x + diff, a)
+        y1n = f(x - diff, a)
+        res[(Ellipsis, *idx)] = (y1p - y1n) / (2 * delta)
+    return res
 
 class TestDipoleFitting(unittest.TestCase, TestCaseSingleFrameWithNlist):
     def setUp(self):
@@ -120,7 +135,6 @@ def test_jit(
             ).to(env.DEVICE)
             torch.jit.script(ft0)
 
-
 class TestEquivalence(unittest.TestCase):
     def setUp(self) -> None:
         self.natoms = 5
@@ -269,5 +283,39 @@ def test_trans(self):
         np.testing.assert_allclose(to_numpy_array(res[0]), to_numpy_array(res[1]))
 
 
+class TestDipoleModel(unittest.TestCase):
+    def setUp(self):
+        self.natoms = 5
+        self.rcut = 4
+        self.nt = 3
+        self.rcut_smth = 0.5
+        self.sel = [46, 92, 4]
+        self.nf = 1
+        self.coord = 2 * torch.rand([self.natoms, 3], dtype=dtype).to(env.DEVICE)
+        self.atype = torch.IntTensor([0, 0, 0, 1, 1]).to(env.DEVICE)
+        self.dd0 = DescrptSeA(self.rcut, self.rcut_smth, self.sel).to(env.DEVICE)
+        self.ft0 = DipoleFittingNet(
+                "dipole",
+                self.nt,
+                self.dd0.dim_out,
+                embedding_width=self.dd0.get_dim_emb(),
+                numb_fparam=0,
+                numb_aparam=0,
+                mixed_types=True,
+            ).to(env.DEVICE)
+        self.type_mapping = ["O", "H", "B"]
+        self.model = DipoleModel(self.dd0, self.ft0, self.type_mapping)
+
+    def test_auto_diff(self):
+        places = 5
+        delta = 1e-5
+        atype = self.atype.view(self.nf, self.natoms)
+        def ff(coord, atype):
+            return self.model(coord, atype)["dipole"].detach().numpy()
+        fdf = -finite_difference(ff, self.coord, atype, delta=delta)
+        rff = self.model(self.coord, atype)["dipole_grad_r"].detach().numpy()
+
+        np.testing.assert_almost_equal(fdf, rff.transpose(0, 2, 1, 3), decimal=places)
+
 if __name__ == "__main__":
     unittest.main()

From 2a1c7eaa515e4bad7f2ccbaa5f580b4697f0d6f0 Mon Sep 17 00:00:00 2001
From: "pre-commit-ci[bot]"
 <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Date: Tue, 20 Feb 2024 07:47:19 +0000
Subject: [PATCH 02/14] [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci
---
 deepmd/dpmodel/fitting/dipole_fitting.py     |  4 +--
 deepmd/pt/model/task/dipole.py               |  4 +--
 source/tests/pt/model/test_dipole_fitting.py | 28 ++++++++++++--------
 3 files changed, 21 insertions(+), 15 deletions(-)

diff --git a/deepmd/dpmodel/fitting/dipole_fitting.py b/deepmd/dpmodel/fitting/dipole_fitting.py
index 69bdac1dac..be28fb75ce 100644
--- a/deepmd/dpmodel/fitting/dipole_fitting.py
+++ b/deepmd/dpmodel/fitting/dipole_fitting.py
@@ -71,10 +71,10 @@ class DipoleFitting(GeneralFitting):
     exclude_types
             Atomic contributions of the excluded atom types are set zero.
     r_differentiable
-            If the variable is differentiated with respect to coordinates of atoms. 
+            If the variable is differentiated with respect to coordinates of atoms.
             Only reduciable variable are differentiable.
     c_differentiable
-            If the variable is differentiated with respect to the cell tensor (pbc case). 
+            If the variable is differentiated with respect to the cell tensor (pbc case).
             Only reduciable variable are differentiable.
     """
 
diff --git a/deepmd/pt/model/task/dipole.py b/deepmd/pt/model/task/dipole.py
index 562bff8211..c86dbba35f 100644
--- a/deepmd/pt/model/task/dipole.py
+++ b/deepmd/pt/model/task/dipole.py
@@ -57,10 +57,10 @@ class DipoleFittingNet(GeneralFitting):
     seed : int, optional
         Random seed.
     r_differentiable
-        If the variable is differentiated with respect to coordinates of atoms. 
+        If the variable is differentiated with respect to coordinates of atoms.
         Only reduciable variable are differentiable.
     c_differentiable
-        If the variable is differentiated with respect to the cell tensor (pbc case). 
+        If the variable is differentiated with respect to the cell tensor (pbc case).
         Only reduciable variable are differentiable.
     """
 
diff --git a/source/tests/pt/model/test_dipole_fitting.py b/source/tests/pt/model/test_dipole_fitting.py
index c816e80c0c..9fe8da09ef 100644
--- a/source/tests/pt/model/test_dipole_fitting.py
+++ b/source/tests/pt/model/test_dipole_fitting.py
@@ -12,6 +12,9 @@
 from deepmd.pt.model.descriptor.se_a import (
     DescrptSeA,
 )
+from deepmd.pt.model.model.dipole_model import (
+    DipoleModel,
+)
 from deepmd.pt.model.task.dipole import (
     DipoleFittingNet,
 )
@@ -28,12 +31,10 @@
 from .test_env_mat import (
     TestCaseSingleFrameWithNlist,
 )
-from deepmd.pt.model.model.dipole_model import (
-    DipoleModel
-)
 
 dtype = env.GLOBAL_PT_FLOAT_PRECISION
 
+
 def finite_difference(f, x, a, delta=1e-6):
     in_shape = x.shape
     y0 = f(x, a)
@@ -47,6 +48,7 @@ def finite_difference(f, x, a, delta=1e-6):
         res[(Ellipsis, *idx)] = (y1p - y1n) / (2 * delta)
     return res
 
+
 class TestDipoleFitting(unittest.TestCase, TestCaseSingleFrameWithNlist):
     def setUp(self):
         TestCaseSingleFrameWithNlist.setUp(self)
@@ -135,6 +137,7 @@ def test_jit(
             ).to(env.DEVICE)
             torch.jit.script(ft0)
 
+
 class TestEquivalence(unittest.TestCase):
     def setUp(self) -> None:
         self.natoms = 5
@@ -295,14 +298,14 @@ def setUp(self):
         self.atype = torch.IntTensor([0, 0, 0, 1, 1]).to(env.DEVICE)
         self.dd0 = DescrptSeA(self.rcut, self.rcut_smth, self.sel).to(env.DEVICE)
         self.ft0 = DipoleFittingNet(
-                "dipole",
-                self.nt,
-                self.dd0.dim_out,
-                embedding_width=self.dd0.get_dim_emb(),
-                numb_fparam=0,
-                numb_aparam=0,
-                mixed_types=True,
-            ).to(env.DEVICE)
+            "dipole",
+            self.nt,
+            self.dd0.dim_out,
+            embedding_width=self.dd0.get_dim_emb(),
+            numb_fparam=0,
+            numb_aparam=0,
+            mixed_types=True,
+        ).to(env.DEVICE)
         self.type_mapping = ["O", "H", "B"]
         self.model = DipoleModel(self.dd0, self.ft0, self.type_mapping)
 
@@ -310,12 +313,15 @@ def test_auto_diff(self):
         places = 5
         delta = 1e-5
         atype = self.atype.view(self.nf, self.natoms)
+
         def ff(coord, atype):
             return self.model(coord, atype)["dipole"].detach().numpy()
+
         fdf = -finite_difference(ff, self.coord, atype, delta=delta)
         rff = self.model(self.coord, atype)["dipole_grad_r"].detach().numpy()
 
         np.testing.assert_almost_equal(fdf, rff.transpose(0, 2, 1, 3), decimal=places)
 
+
 if __name__ == "__main__":
     unittest.main()

From 3987283a8ec34ef6660ca8ee07450769153e44b7 Mon Sep 17 00:00:00 2001
From: Anyang Peng <137014849+anyangml@users.noreply.github.com>
Date: Wed, 21 Feb 2024 10:25:19 +0800
Subject: [PATCH 03/14] feat: add UTs

---
 deepmd/dpmodel/fitting/dipole_fitting.py     |  2 +-
 deepmd/infer/deep_tensor.py                  | 18 +++++++----
 deepmd/pt/model/model/dipole_model.py        | 24 ++++++++++----
 deepmd/pt/model/task/dipole.py               |  2 +-
 source/tests/pt/model/test_dipole_fitting.py | 34 ++++++++++++++++----
 5 files changed, 59 insertions(+), 21 deletions(-)

diff --git a/deepmd/dpmodel/fitting/dipole_fitting.py b/deepmd/dpmodel/fitting/dipole_fitting.py
index be28fb75ce..f5acabf7b1 100644
--- a/deepmd/dpmodel/fitting/dipole_fitting.py
+++ b/deepmd/dpmodel/fitting/dipole_fitting.py
@@ -100,7 +100,7 @@ def __init__(
         mixed_types: bool = False,
         exclude_types: List[int] = [],
         r_differentiable: bool = True,
-        c_differentiable: bool = False,
+        c_differentiable: bool = True,
         old_impl=False,
     ):
         # seed, uniform_seed are not included
diff --git a/deepmd/infer/deep_tensor.py b/deepmd/infer/deep_tensor.py
index a6cefa63c1..36671861a2 100644
--- a/deepmd/infer/deep_tensor.py
+++ b/deepmd/infer/deep_tensor.py
@@ -105,6 +105,8 @@ def eval(
             **kwargs,
         )
         sel_natoms = self._get_sel_natoms(atom_types[0])
+        if sel_natoms == 0:
+            sel_natoms = atom_types.shape[-1] # set to natoms
         if atomic:
             return results[self.output_tensor_name].reshape(nframes, sel_natoms, -1)
         else:
@@ -184,7 +186,10 @@ def eval_full(
             aparam=aparam,
             **kwargs,
         )
+        
         sel_natoms = self._get_sel_natoms(atom_types[0])
+        if sel_natoms == 0:
+            sel_natoms = atom_types.shape[-1] # set to natoms
         energy = results[f"{self.output_tensor_name}_redu"].reshape(nframes, -1)
         force = results[f"{self.output_tensor_name}_derv_r"].reshape(
             nframes, -1, natoms, 3
@@ -192,14 +197,13 @@ def eval_full(
         virial = results[f"{self.output_tensor_name}_derv_c_redu"].reshape(
             nframes, -1, 9
         )
-        atomic_energy = results[self.output_tensor_name].reshape(
-            nframes, sel_natoms, -1
-        )
-        atomic_virial = results[f"{self.output_tensor_name}_derv_c"].reshape(
-            nframes, -1, natoms, 9
-        )
-
         if atomic:
+            atomic_energy = results[self.output_tensor_name].reshape(
+                nframes, sel_natoms, -1
+            )
+            atomic_virial = results[f"{self.output_tensor_name}_derv_c"].reshape(
+                nframes, -1, natoms, 9
+            )
             return (
                 energy,
                 force,
diff --git a/deepmd/pt/model/model/dipole_model.py b/deepmd/pt/model/model/dipole_model.py
index fda5940939..10fc55607a 100644
--- a/deepmd/pt/model/model/dipole_model.py
+++ b/deepmd/pt/model/model/dipole_model.py
@@ -36,14 +36,19 @@ def forward(
             box,
             fparam=fparam,
             aparam=aparam,
-            do_atomic_virial=False,
+            do_atomic_virial=do_atomic_virial,
         )
         if self.fitting_net is not None:
             model_predict = {}
-            model_predict["atom_dipole"] = model_ret["dipole"]
-            model_predict["dipole"] = model_ret["dipole_redu"]
+            model_predict["dipole"] = model_ret["dipole"]
+            model_predict["global_dipole"] = model_ret["dipole_redu"]
             if self.do_grad("dipole"):
-                model_predict["dipole_grad_r"] = model_ret["dipole_derv_r"].squeeze(-2)
+                model_predict["force"] = model_ret["dipole_derv_r"].squeeze(-2)
+                if do_atomic_virial:
+                    model_predict["atom_virial"] = model_ret["dipole_derv_c"].squeeze(
+                        -3
+                    )
+                model_predict["virial"] = model_ret["dipole_derv_c_redu"].squeeze(-2)
         else:
             model_predict = model_ret
             model_predict["updated_coord"] += coord
@@ -70,8 +75,15 @@ def forward_lower(
         )
         if self.fitting_net is not None:
             model_predict = {}
-            model_predict["atom_dipole"] = model_ret["dipole"]
-            model_predict["energy"] = model_ret["energy_redu"]
+            model_predict["dipole"] = model_ret["dipole"]
+            model_predict["global_dipole"] = model_ret["dipole_redu"]
+            if self.do_grad("dipole"):
+                model_predict["force"] = model_ret["dipole_derv_r"].squeeze(-2)
+                if do_atomic_virial:
+                    model_predict["atom_virial"] = model_ret["dipole_derv_c"].squeeze(
+                        -3
+                    )
+                model_predict["virial"] = model_ret["dipole_derv_c_redu"].squeeze(-2)
         else:
             model_predict = model_ret
         return model_predict
diff --git a/deepmd/pt/model/task/dipole.py b/deepmd/pt/model/task/dipole.py
index c86dbba35f..88391b1922 100644
--- a/deepmd/pt/model/task/dipole.py
+++ b/deepmd/pt/model/task/dipole.py
@@ -81,7 +81,7 @@ def __init__(
         seed: Optional[int] = None,
         exclude_types: List[int] = [],
         r_differentiable: bool = True,
-        c_differentiable: bool = False,
+        c_differentiable: bool = True,
         **kwargs,
     ):
         self.embedding_width = embedding_width
diff --git a/source/tests/pt/model/test_dipole_fitting.py b/source/tests/pt/model/test_dipole_fitting.py
index 9fe8da09ef..43fac8691d 100644
--- a/source/tests/pt/model/test_dipole_fitting.py
+++ b/source/tests/pt/model/test_dipole_fitting.py
@@ -2,6 +2,7 @@
 import itertools
 import unittest
 
+import os
 import numpy as np
 import torch
 from scipy.stats import (
@@ -31,6 +32,9 @@
 from .test_env_mat import (
     TestCaseSingleFrameWithNlist,
 )
+from deepmd.infer.deep_dipole import (
+    DeepDipole,
+)
 
 dtype = env.GLOBAL_PT_FLOAT_PRECISION
 
@@ -289,12 +293,14 @@ def test_trans(self):
 class TestDipoleModel(unittest.TestCase):
     def setUp(self):
         self.natoms = 5
-        self.rcut = 4
+        self.rcut = 4.
         self.nt = 3
         self.rcut_smth = 0.5
         self.sel = [46, 92, 4]
         self.nf = 1
         self.coord = 2 * torch.rand([self.natoms, 3], dtype=dtype).to(env.DEVICE)
+        cell = torch.rand([3, 3], dtype=dtype).to(env.DEVICE)
+        self.cell = (cell + cell.T) + 5.0 * torch.eye(3).to(env.DEVICE)
         self.atype = torch.IntTensor([0, 0, 0, 1, 1]).to(env.DEVICE)
         self.dd0 = DescrptSeA(self.rcut, self.rcut_smth, self.sel).to(env.DEVICE)
         self.ft0 = DipoleFittingNet(
@@ -308,20 +314,36 @@ def setUp(self):
         ).to(env.DEVICE)
         self.type_mapping = ["O", "H", "B"]
         self.model = DipoleModel(self.dd0, self.ft0, self.type_mapping)
-
+        self.file_path = "model_output.pth"
+        
     def test_auto_diff(self):
         places = 5
         delta = 1e-5
         atype = self.atype.view(self.nf, self.natoms)
 
         def ff(coord, atype):
-            return self.model(coord, atype)["dipole"].detach().numpy()
+            return self.model(coord, atype)["global_dipole"].detach().numpy()
 
         fdf = -finite_difference(ff, self.coord, atype, delta=delta)
-        rff = self.model(self.coord, atype)["dipole_grad_r"].detach().numpy()
+        rff = self.model(self.coord, atype)["force"].detach().numpy()
 
         np.testing.assert_almost_equal(fdf, rff.transpose(0, 2, 1, 3), decimal=places)
-
-
+    
+    def test_deepdipole_infer(self):
+        atype = self.atype.view(self.nf, self.natoms)
+        coord = self.coord.reshape(1,5,3)
+        cell = self.cell.reshape(1,9)
+        jit_md = torch.jit.script(self.model)
+        torch.jit.save(jit_md, self.file_path)
+        load_md = DeepDipole(self.file_path)
+        load_md.eval(coords=coord, atom_types=atype, cells=cell, atomic=True)
+        load_md.eval(coords=coord, atom_types=atype, cells=cell, atomic=False)
+        load_md.eval_full(coords=coord, atom_types=atype, cells=cell, atomic=True)
+        load_md.eval_full(coords=coord, atom_types=atype, cells=cell, atomic=False)
+
+    def tearDown(self) -> None:
+        if os.path.exists(self.file_path):
+            os.remove(self.file_path)
+            
 if __name__ == "__main__":
     unittest.main()

From 278b17ba1568bd5f76bcd41743759aa47775b017 Mon Sep 17 00:00:00 2001
From: "pre-commit-ci[bot]"
 <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Date: Wed, 21 Feb 2024 02:25:56 +0000
Subject: [PATCH 04/14] [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci
---
 deepmd/infer/deep_tensor.py                  |  6 +++---
 source/tests/pt/model/test_dipole_fitting.py | 21 ++++++++++----------
 2 files changed, 14 insertions(+), 13 deletions(-)

diff --git a/deepmd/infer/deep_tensor.py b/deepmd/infer/deep_tensor.py
index 36671861a2..1bdc459920 100644
--- a/deepmd/infer/deep_tensor.py
+++ b/deepmd/infer/deep_tensor.py
@@ -106,7 +106,7 @@ def eval(
         )
         sel_natoms = self._get_sel_natoms(atom_types[0])
         if sel_natoms == 0:
-            sel_natoms = atom_types.shape[-1] # set to natoms
+            sel_natoms = atom_types.shape[-1]  # set to natoms
         if atomic:
             return results[self.output_tensor_name].reshape(nframes, sel_natoms, -1)
         else:
@@ -186,10 +186,10 @@ def eval_full(
             aparam=aparam,
             **kwargs,
         )
-        
+
         sel_natoms = self._get_sel_natoms(atom_types[0])
         if sel_natoms == 0:
-            sel_natoms = atom_types.shape[-1] # set to natoms
+            sel_natoms = atom_types.shape[-1]  # set to natoms
         energy = results[f"{self.output_tensor_name}_redu"].reshape(nframes, -1)
         force = results[f"{self.output_tensor_name}_derv_r"].reshape(
             nframes, -1, natoms, 3
diff --git a/source/tests/pt/model/test_dipole_fitting.py b/source/tests/pt/model/test_dipole_fitting.py
index 43fac8691d..5285c4d1b3 100644
--- a/source/tests/pt/model/test_dipole_fitting.py
+++ b/source/tests/pt/model/test_dipole_fitting.py
@@ -1,8 +1,8 @@
 # SPDX-License-Identifier: LGPL-3.0-or-later
 import itertools
+import os
 import unittest
 
-import os
 import numpy as np
 import torch
 from scipy.stats import (
@@ -10,6 +10,9 @@
 )
 
 from deepmd.dpmodel.fitting import DipoleFitting as DPDipoleFitting
+from deepmd.infer.deep_dipole import (
+    DeepDipole,
+)
 from deepmd.pt.model.descriptor.se_a import (
     DescrptSeA,
 )
@@ -32,9 +35,6 @@
 from .test_env_mat import (
     TestCaseSingleFrameWithNlist,
 )
-from deepmd.infer.deep_dipole import (
-    DeepDipole,
-)
 
 dtype = env.GLOBAL_PT_FLOAT_PRECISION
 
@@ -293,7 +293,7 @@ def test_trans(self):
 class TestDipoleModel(unittest.TestCase):
     def setUp(self):
         self.natoms = 5
-        self.rcut = 4.
+        self.rcut = 4.0
         self.nt = 3
         self.rcut_smth = 0.5
         self.sel = [46, 92, 4]
@@ -315,7 +315,7 @@ def setUp(self):
         self.type_mapping = ["O", "H", "B"]
         self.model = DipoleModel(self.dd0, self.ft0, self.type_mapping)
         self.file_path = "model_output.pth"
-        
+
     def test_auto_diff(self):
         places = 5
         delta = 1e-5
@@ -328,11 +328,11 @@ def ff(coord, atype):
         rff = self.model(self.coord, atype)["force"].detach().numpy()
 
         np.testing.assert_almost_equal(fdf, rff.transpose(0, 2, 1, 3), decimal=places)
-    
+
     def test_deepdipole_infer(self):
         atype = self.atype.view(self.nf, self.natoms)
-        coord = self.coord.reshape(1,5,3)
-        cell = self.cell.reshape(1,9)
+        coord = self.coord.reshape(1, 5, 3)
+        cell = self.cell.reshape(1, 9)
         jit_md = torch.jit.script(self.model)
         torch.jit.save(jit_md, self.file_path)
         load_md = DeepDipole(self.file_path)
@@ -344,6 +344,7 @@ def test_deepdipole_infer(self):
     def tearDown(self) -> None:
         if os.path.exists(self.file_path):
             os.remove(self.file_path)
-            
+
+
 if __name__ == "__main__":
     unittest.main()

From 38af4b17b91d805de6cda8c19c1d2da73b76f2f9 Mon Sep 17 00:00:00 2001
From: Anyang Peng <137014849+anyangml@users.noreply.github.com>
Date: Wed, 21 Feb 2024 10:35:00 +0800
Subject: [PATCH 05/14] chore: refactor

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

diff --git a/source/tests/pt/model/test_dipole_fitting.py b/source/tests/pt/model/test_dipole_fitting.py
index 5285c4d1b3..986f8e800c 100644
--- a/source/tests/pt/model/test_dipole_fitting.py
+++ b/source/tests/pt/model/test_dipole_fitting.py
@@ -298,10 +298,10 @@ def setUp(self):
         self.rcut_smth = 0.5
         self.sel = [46, 92, 4]
         self.nf = 1
-        self.coord = 2 * torch.rand([self.natoms, 3], dtype=dtype).to(env.DEVICE)
-        cell = torch.rand([3, 3], dtype=dtype).to(env.DEVICE)
-        self.cell = (cell + cell.T) + 5.0 * torch.eye(3).to(env.DEVICE)
-        self.atype = torch.IntTensor([0, 0, 0, 1, 1]).to(env.DEVICE)
+        self.coord = 2 * torch.rand([self.natoms, 3], dtype=dtype, device=env.DEVICE)
+        cell = torch.rand([3, 3], dtype=dtype, device=env.DEVICE)
+        self.cell = (cell + cell.T) + 5.0 * torch.eye(3, device=env.DEVICE)
+        self.atype = torch.IntTensor([0, 0, 0, 1, 1], device=env.DEVICE)
         self.dd0 = DescrptSeA(self.rcut, self.rcut_smth, self.sel).to(env.DEVICE)
         self.ft0 = DipoleFittingNet(
             "dipole",

From 2eef3c0a0850685ebb09c135fea8b66f9b52d709 Mon Sep 17 00:00:00 2001
From: Anyang Peng <137014849+anyangml@users.noreply.github.com>
Date: Wed, 21 Feb 2024 10:56:21 +0800
Subject: [PATCH 06/14] feat: add deeppolar

---
 deepmd/pt/model/model/dipole_model.py         |  2 +-
 deepmd/pt/model/model/polar_model.py          | 75 +++++++++++++++++++
 .../pt/model/test_polarizability_fitting.py   | 46 ++++++++++++
 3 files changed, 122 insertions(+), 1 deletion(-)
 create mode 100644 deepmd/pt/model/model/polar_model.py

diff --git a/deepmd/pt/model/model/dipole_model.py b/deepmd/pt/model/model/dipole_model.py
index 10fc55607a..1e5bde8217 100644
--- a/deepmd/pt/model/model/dipole_model.py
+++ b/deepmd/pt/model/model/dipole_model.py
@@ -71,7 +71,7 @@ def forward_lower(
             mapping,
             fparam=fparam,
             aparam=aparam,
-            do_atomic_virial=False,
+            do_atomic_virial=do_atomic_virial,
         )
         if self.fitting_net is not None:
             model_predict = {}
diff --git a/deepmd/pt/model/model/polar_model.py b/deepmd/pt/model/model/polar_model.py
new file mode 100644
index 0000000000..85aeebc0f5
--- /dev/null
+++ b/deepmd/pt/model/model/polar_model.py
@@ -0,0 +1,75 @@
+# SPDX-License-Identifier: LGPL-3.0-or-later
+from typing import (
+    Dict,
+    Optional,
+)
+
+import torch
+
+from .dp_model import (
+    DPModel,
+)
+
+
+class PolarModel(DPModel):
+    model_type = "polar"
+
+    def __init__(
+        self,
+        *args,
+        **kwargs,
+    ):
+        super().__init__(*args, **kwargs)
+
+    def forward(
+        self,
+        coord,
+        atype,
+        box: Optional[torch.Tensor] = None,
+        fparam: Optional[torch.Tensor] = None,
+        aparam: Optional[torch.Tensor] = None,
+        do_atomic_virial: bool = False,
+    ) -> Dict[str, torch.Tensor]:
+        model_ret = self.forward_common(
+            coord,
+            atype,
+            box,
+            fparam=fparam,
+            aparam=aparam,
+            do_atomic_virial=do_atomic_virial,
+        )
+        if self.fitting_net is not None:
+            model_predict = {}
+            model_predict["polar"] = model_ret["polar"]
+            model_predict["global_polar"] = model_ret["polar_redu"]
+        else:
+            model_predict = model_ret
+            model_predict["updated_coord"] += coord
+        return model_predict
+
+    def forward_lower(
+        self,
+        extended_coord,
+        extended_atype,
+        nlist,
+        mapping: Optional[torch.Tensor] = None,
+        fparam: Optional[torch.Tensor] = None,
+        aparam: Optional[torch.Tensor] = None,
+        do_atomic_virial: bool = False,
+    ):
+        model_ret = self.forward_common_lower(
+            extended_coord,
+            extended_atype,
+            nlist,
+            mapping,
+            fparam=fparam,
+            aparam=aparam,
+            do_atomic_virial=do_atomic_virial,
+        )
+        if self.fitting_net is not None:
+            model_predict = {}
+            model_predict["polar"] = model_ret["polar"]
+            model_predict["global_polar"] = model_ret["polar_redu"]
+        else:
+            model_predict = model_ret
+        return model_predict
diff --git a/source/tests/pt/model/test_polarizability_fitting.py b/source/tests/pt/model/test_polarizability_fitting.py
index de43c57b8b..b7a04244fd 100644
--- a/source/tests/pt/model/test_polarizability_fitting.py
+++ b/source/tests/pt/model/test_polarizability_fitting.py
@@ -2,6 +2,7 @@
 import itertools
 import unittest
 
+import os
 import numpy as np
 import torch
 from scipy.stats import (
@@ -15,9 +16,15 @@
 from deepmd.pt.model.task.polarizability import (
     PolarFittingNet,
 )
+from deepmd.infer.deep_polar import (
+    DeepPolar,
+)
 from deepmd.pt.utils import (
     env,
 )
+from deepmd.pt.model.model.polar_model import (
+    PolarModel,
+)
 from deepmd.pt.utils.nlist import (
     extend_input_and_build_neighbor_list,
 )
@@ -307,6 +314,45 @@ def test_trans(self):
 
             np.testing.assert_allclose(to_numpy_array(res[0]), to_numpy_array(res[1]))
 
+class TestDipoleModel(unittest.TestCase):
+    def setUp(self):
+        self.natoms = 5
+        self.rcut = 4.0
+        self.nt = 3
+        self.rcut_smth = 0.5
+        self.sel = [46, 92, 4]
+        self.nf = 1
+        self.coord = 2 * torch.rand([self.natoms, 3], dtype=dtype, device=env.DEVICE)
+        cell = torch.rand([3, 3], dtype=dtype, device=env.DEVICE)
+        self.cell = (cell + cell.T) + 5.0 * torch.eye(3, device=env.DEVICE)
+        self.atype = torch.IntTensor([0, 0, 0, 1, 1], device=env.DEVICE)
+        self.dd0 = DescrptSeA(self.rcut, self.rcut_smth, self.sel).to(env.DEVICE)
+        self.ft0 = PolarFittingNet(
+            "polar",
+            self.nt,
+            self.dd0.dim_out,
+            embedding_width=self.dd0.get_dim_emb(),
+            numb_fparam=0,
+            numb_aparam=0,
+            mixed_types=True,
+        ).to(env.DEVICE)
+        self.type_mapping = ["O", "H", "B"]
+        self.model = PolarModel(self.dd0, self.ft0, self.type_mapping)
+        self.file_path = "model_output.pth"
+    
+    def test_deepdipole_infer(self):
+        atype = self.atype.view(self.nf, self.natoms)
+        coord = self.coord.reshape(1, 5, 3)
+        cell = self.cell.reshape(1, 9)
+        jit_md = torch.jit.script(self.model)
+        torch.jit.save(jit_md, self.file_path)
+        load_md = DeepPolar(self.file_path)
+        load_md.eval(coords=coord, atom_types=atype, cells=cell, atomic=True)
+        load_md.eval(coords=coord, atom_types=atype, cells=cell, atomic=False)
+
+    def tearDown(self) -> None:
+        if os.path.exists(self.file_path):
+            os.remove(self.file_path)
 
 if __name__ == "__main__":
     unittest.main()

From 442d27d94814ab3f13e77816f401e6ca6efe8d9f Mon Sep 17 00:00:00 2001
From: "pre-commit-ci[bot]"
 <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Date: Wed, 21 Feb 2024 02:56:58 +0000
Subject: [PATCH 07/14] [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci
---
 .../pt/model/test_polarizability_fitting.py    | 18 ++++++++++--------
 1 file changed, 10 insertions(+), 8 deletions(-)

diff --git a/source/tests/pt/model/test_polarizability_fitting.py b/source/tests/pt/model/test_polarizability_fitting.py
index b7a04244fd..bc474758c5 100644
--- a/source/tests/pt/model/test_polarizability_fitting.py
+++ b/source/tests/pt/model/test_polarizability_fitting.py
@@ -1,8 +1,8 @@
 # SPDX-License-Identifier: LGPL-3.0-or-later
 import itertools
+import os
 import unittest
 
-import os
 import numpy as np
 import torch
 from scipy.stats import (
@@ -10,21 +10,21 @@
 )
 
 from deepmd.dpmodel.fitting import PolarFitting as DPPolarFitting
+from deepmd.infer.deep_polar import (
+    DeepPolar,
+)
 from deepmd.pt.model.descriptor.se_a import (
     DescrptSeA,
 )
+from deepmd.pt.model.model.polar_model import (
+    PolarModel,
+)
 from deepmd.pt.model.task.polarizability import (
     PolarFittingNet,
 )
-from deepmd.infer.deep_polar import (
-    DeepPolar,
-)
 from deepmd.pt.utils import (
     env,
 )
-from deepmd.pt.model.model.polar_model import (
-    PolarModel,
-)
 from deepmd.pt.utils.nlist import (
     extend_input_and_build_neighbor_list,
 )
@@ -314,6 +314,7 @@ def test_trans(self):
 
             np.testing.assert_allclose(to_numpy_array(res[0]), to_numpy_array(res[1]))
 
+
 class TestDipoleModel(unittest.TestCase):
     def setUp(self):
         self.natoms = 5
@@ -339,7 +340,7 @@ def setUp(self):
         self.type_mapping = ["O", "H", "B"]
         self.model = PolarModel(self.dd0, self.ft0, self.type_mapping)
         self.file_path = "model_output.pth"
-    
+
     def test_deepdipole_infer(self):
         atype = self.atype.view(self.nf, self.natoms)
         coord = self.coord.reshape(1, 5, 3)
@@ -354,5 +355,6 @@ def tearDown(self) -> None:
         if os.path.exists(self.file_path):
             os.remove(self.file_path)
 
+
 if __name__ == "__main__":
     unittest.main()

From 658dd92c21c1c0b5c313aff4579a5e2f80832da3 Mon Sep 17 00:00:00 2001
From: Anyang Peng <137014849+anyangml@users.noreply.github.com>
Date: Wed, 21 Feb 2024 12:08:05 +0800
Subject: [PATCH 08/14] feat: split do_grad r, c

---
 .../atomic_model/make_base_atomic_model.py    | 36 ++++++++++++++-----
 .../pt/model/atomic_model/dp_atomic_model.py  |  2 +-
 .../model/atomic_model/linear_atomic_model.py |  2 +-
 .../atomic_model/pairtab_atomic_model.py      |  2 +-
 deepmd/pt/model/model/dipole_model.py         | 10 +++---
 deepmd/pt/model/model/dp_zbl_model.py         | 16 ++++-----
 deepmd/pt/model/model/ener_model.py           | 18 +++++-----
 7 files changed, 52 insertions(+), 34 deletions(-)

diff --git a/deepmd/dpmodel/atomic_model/make_base_atomic_model.py b/deepmd/dpmodel/atomic_model/make_base_atomic_model.py
index e262404762..fb7f76de18 100644
--- a/deepmd/dpmodel/atomic_model/make_base_atomic_model.py
+++ b/deepmd/dpmodel/atomic_model/make_base_atomic_model.py
@@ -118,33 +118,51 @@ def serialize(self) -> dict:
         def deserialize(cls):
             pass
 
-        def do_grad(
+        def do_grad_r(
             self,
             var_name: Optional[str] = None,
         ) -> bool:
-            """Tell if the output variable `var_name` is differentiable.
-            if var_name is None, returns if any of the variable is differentiable.
+            """Tell if the output variable `var_name` is r_differentiable.
+            if var_name is None, returns if any of the variable is r_differentiable.
 
             """
             odef = self.fitting_output_def()
             if var_name is None:
                 require: List[bool] = []
                 for vv in odef.keys():
-                    require.append(self.do_grad_(vv))
+                    require.append(self.do_grad_(vv, "r"))
                 return any(require)
             else:
-                return self.do_grad_(var_name)
+                return self.do_grad_(var_name, "r")
+        
+        def do_grad_c(
+            self,
+            var_name: Optional[str] = None,
+        ) -> bool:
+            """Tell if the output variable `var_name` is c_differentiable.
+            if var_name is None, returns if any of the variable is c_differentiable.
+
+            """
+            odef = self.fitting_output_def()
+            if var_name is None:
+                require: List[bool] = []
+                for vv in odef.keys():
+                    require.append(self.do_grad_(vv, "c"))
+                return any(require)
+            else:
+                return self.do_grad_(var_name, "c")
 
         def do_grad_(
             self,
             var_name: str,
+            base: str
         ) -> bool:
             """Tell if the output variable `var_name` is differentiable."""
             assert var_name is not None
-            return (
-                self.fitting_output_def()[var_name].r_differentiable
-                or self.fitting_output_def()[var_name].c_differentiable
-            )
+            assert base in ["c", "r"]
+            if base == "c":
+                return self.fitting_output_def()[var_name].c_differentiable
+            return self.fitting_output_def()[var_name].r_differentiable
 
     setattr(BAM, fwd_method_name, BAM.fwd)
     delattr(BAM, "fwd")
diff --git a/deepmd/pt/model/atomic_model/dp_atomic_model.py b/deepmd/pt/model/atomic_model/dp_atomic_model.py
index e6dc395500..ecac50737b 100644
--- a/deepmd/pt/model/atomic_model/dp_atomic_model.py
+++ b/deepmd/pt/model/atomic_model/dp_atomic_model.py
@@ -146,7 +146,7 @@ def forward_atomic(
         """
         nframes, nloc, nnei = nlist.shape
         atype = extended_atype[:, :nloc]
-        if self.do_grad():
+        if self.do_grad_r() or self.do_grad_c():
             extended_coord.requires_grad_(True)
         descriptor, rot_mat, g2, h2, sw = self.descriptor(
             extended_coord,
diff --git a/deepmd/pt/model/atomic_model/linear_atomic_model.py b/deepmd/pt/model/atomic_model/linear_atomic_model.py
index b688362b73..d1bcd8e25e 100644
--- a/deepmd/pt/model/atomic_model/linear_atomic_model.py
+++ b/deepmd/pt/model/atomic_model/linear_atomic_model.py
@@ -134,7 +134,7 @@ def forward_atomic(
             the result dict, defined by the fitting net output def.
         """
         nframes, nloc, nnei = nlist.shape
-        if self.do_grad():
+        if self.do_grad_r or self.do_grad_c():
             extended_coord.requires_grad_(True)
         extended_coord = extended_coord.view(nframes, -1, 3)
         sorted_rcuts, sorted_sels = self._sort_rcuts_sels()
diff --git a/deepmd/pt/model/atomic_model/pairtab_atomic_model.py b/deepmd/pt/model/atomic_model/pairtab_atomic_model.py
index 4019fda423..8408e4d507 100644
--- a/deepmd/pt/model/atomic_model/pairtab_atomic_model.py
+++ b/deepmd/pt/model/atomic_model/pairtab_atomic_model.py
@@ -145,7 +145,7 @@ def forward_atomic(
     ) -> Dict[str, torch.Tensor]:
         nframes, nloc, nnei = nlist.shape
         extended_coord = extended_coord.view(nframes, -1, 3)
-        if self.do_grad():
+        if self.do_grad_r or self.do_grad_c():
             extended_coord.requires_grad_(True)
 
         # this will mask all -1 in the nlist
diff --git a/deepmd/pt/model/model/dipole_model.py b/deepmd/pt/model/model/dipole_model.py
index 1e5bde8217..220fdbb273 100644
--- a/deepmd/pt/model/model/dipole_model.py
+++ b/deepmd/pt/model/model/dipole_model.py
@@ -42,13 +42,14 @@ def forward(
             model_predict = {}
             model_predict["dipole"] = model_ret["dipole"]
             model_predict["global_dipole"] = model_ret["dipole_redu"]
-            if self.do_grad("dipole"):
+            if self.do_grad_r("dipole"):
                 model_predict["force"] = model_ret["dipole_derv_r"].squeeze(-2)
+            if self.do_grad_c("dipole"):
+                model_predict["virial"] = model_ret["dipole_derv_c_redu"].squeeze(-2)
                 if do_atomic_virial:
                     model_predict["atom_virial"] = model_ret["dipole_derv_c"].squeeze(
                         -3
                     )
-                model_predict["virial"] = model_ret["dipole_derv_c_redu"].squeeze(-2)
         else:
             model_predict = model_ret
             model_predict["updated_coord"] += coord
@@ -77,13 +78,14 @@ def forward_lower(
             model_predict = {}
             model_predict["dipole"] = model_ret["dipole"]
             model_predict["global_dipole"] = model_ret["dipole_redu"]
-            if self.do_grad("dipole"):
+            if self.do_grad_r("dipole"):
                 model_predict["force"] = model_ret["dipole_derv_r"].squeeze(-2)
+            if self.do_grad_c("dipole"):
+                model_predict["virial"] = model_ret["dipole_derv_c_redu"].squeeze(-2)
                 if do_atomic_virial:
                     model_predict["atom_virial"] = model_ret["dipole_derv_c"].squeeze(
                         -3
                     )
-                model_predict["virial"] = model_ret["dipole_derv_c_redu"].squeeze(-2)
         else:
             model_predict = model_ret
         return model_predict
diff --git a/deepmd/pt/model/model/dp_zbl_model.py b/deepmd/pt/model/model/dp_zbl_model.py
index 8d71157b60..448d408b3f 100644
--- a/deepmd/pt/model/model/dp_zbl_model.py
+++ b/deepmd/pt/model/model/dp_zbl_model.py
@@ -48,11 +48,12 @@ def forward(
         model_predict = {}
         model_predict["atom_energy"] = model_ret["energy"]
         model_predict["energy"] = model_ret["energy_redu"]
-        if self.do_grad("energy"):
+        if self.do_grad_r("energy"):
             model_predict["force"] = model_ret["energy_derv_r"].squeeze(-2)
-            if do_atomic_virial:
-                model_predict["atom_virial"] = model_ret["energy_derv_c"].squeeze(-3)
+        if self.do_grad_c("energy"):
             model_predict["virial"] = model_ret["energy_derv_c_redu"].squeeze(-2)
+            if do_atomic_virial:
+                model_predict["atom_virial"] = model_ret["energy_derv_c"].squeeze(-3)   
         else:
             model_predict["force"] = model_ret["dforce"]
         return model_predict
@@ -80,13 +81,12 @@ def forward_lower(
         model_predict = {}
         model_predict["atom_energy"] = model_ret["energy"]
         model_predict["energy"] = model_ret["energy_redu"]
-        if self.do_grad("energy"):
-            model_predict["extended_force"] = model_ret["energy_derv_r"].squeeze(-2)
+        if self.do_grad_r("energy"):
+            model_predict["force"] = model_ret["energy_derv_r"].squeeze(-2)
+        if self.do_grad_c("energy"):
             model_predict["virial"] = model_ret["energy_derv_c_redu"].squeeze(-2)
             if do_atomic_virial:
-                model_predict["extended_virial"] = model_ret["energy_derv_c"].squeeze(
-                    -2
-                )
+                model_predict["atom_virial"] = model_ret["energy_derv_c"].squeeze(-3)
         else:
             assert model_ret["dforce"] is not None
             model_predict["dforce"] = model_ret["dforce"]
diff --git a/deepmd/pt/model/model/ener_model.py b/deepmd/pt/model/model/ener_model.py
index 2afeb2762b..c747a84e28 100644
--- a/deepmd/pt/model/model/ener_model.py
+++ b/deepmd/pt/model/model/ener_model.py
@@ -42,13 +42,12 @@ def forward(
             model_predict = {}
             model_predict["atom_energy"] = model_ret["energy"]
             model_predict["energy"] = model_ret["energy_redu"]
-            if self.do_grad("energy"):
+            if self.do_grad_r("energy"):
                 model_predict["force"] = model_ret["energy_derv_r"].squeeze(-2)
-                if do_atomic_virial:
-                    model_predict["atom_virial"] = model_ret["energy_derv_c"].squeeze(
-                        -3
-                    )
+            if self.do_grad_c("energy"):
                 model_predict["virial"] = model_ret["energy_derv_c_redu"].squeeze(-2)
+                if do_atomic_virial:
+                    model_predict["atom_virial"] = model_ret["energy_derv_c"].squeeze(-3)
             else:
                 model_predict["force"] = model_ret["dforce"]
         else:
@@ -79,13 +78,12 @@ def forward_lower(
             model_predict = {}
             model_predict["atom_energy"] = model_ret["energy"]
             model_predict["energy"] = model_ret["energy_redu"]
-            if self.do_grad("energy"):
-                model_predict["extended_force"] = model_ret["energy_derv_r"].squeeze(-2)
+            if self.do_grad_r("energy"):
+                model_predict["force"] = model_ret["energy_derv_r"].squeeze(-2)
+            if self.do_grad_c("energy"):
                 model_predict["virial"] = model_ret["energy_derv_c_redu"].squeeze(-2)
                 if do_atomic_virial:
-                    model_predict["extended_virial"] = model_ret[
-                        "energy_derv_c"
-                    ].squeeze(-2)
+                    model_predict["atom_virial"] = model_ret["energy_derv_c"].squeeze(-3)
             else:
                 assert model_ret["dforce"] is not None
                 model_predict["dforce"] = model_ret["dforce"]

From 9eafbd4bb8391c665eea80bb31910906fae127da Mon Sep 17 00:00:00 2001
From: "pre-commit-ci[bot]"
 <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Date: Wed, 21 Feb 2024 04:10:10 +0000
Subject: [PATCH 09/14] [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci
---
 deepmd/dpmodel/atomic_model/make_base_atomic_model.py | 8 ++------
 deepmd/pt/model/model/dp_zbl_model.py                 | 2 +-
 deepmd/pt/model/model/ener_model.py                   | 8 ++++++--
 3 files changed, 9 insertions(+), 9 deletions(-)

diff --git a/deepmd/dpmodel/atomic_model/make_base_atomic_model.py b/deepmd/dpmodel/atomic_model/make_base_atomic_model.py
index fb7f76de18..b6c6b8460f 100644
--- a/deepmd/dpmodel/atomic_model/make_base_atomic_model.py
+++ b/deepmd/dpmodel/atomic_model/make_base_atomic_model.py
@@ -134,7 +134,7 @@ def do_grad_r(
                 return any(require)
             else:
                 return self.do_grad_(var_name, "r")
-        
+
         def do_grad_c(
             self,
             var_name: Optional[str] = None,
@@ -152,11 +152,7 @@ def do_grad_c(
             else:
                 return self.do_grad_(var_name, "c")
 
-        def do_grad_(
-            self,
-            var_name: str,
-            base: str
-        ) -> bool:
+        def do_grad_(self, var_name: str, base: str) -> bool:
             """Tell if the output variable `var_name` is differentiable."""
             assert var_name is not None
             assert base in ["c", "r"]
diff --git a/deepmd/pt/model/model/dp_zbl_model.py b/deepmd/pt/model/model/dp_zbl_model.py
index 448d408b3f..4683f62466 100644
--- a/deepmd/pt/model/model/dp_zbl_model.py
+++ b/deepmd/pt/model/model/dp_zbl_model.py
@@ -53,7 +53,7 @@ def forward(
         if self.do_grad_c("energy"):
             model_predict["virial"] = model_ret["energy_derv_c_redu"].squeeze(-2)
             if do_atomic_virial:
-                model_predict["atom_virial"] = model_ret["energy_derv_c"].squeeze(-3)   
+                model_predict["atom_virial"] = model_ret["energy_derv_c"].squeeze(-3)
         else:
             model_predict["force"] = model_ret["dforce"]
         return model_predict
diff --git a/deepmd/pt/model/model/ener_model.py b/deepmd/pt/model/model/ener_model.py
index c747a84e28..b22f73275f 100644
--- a/deepmd/pt/model/model/ener_model.py
+++ b/deepmd/pt/model/model/ener_model.py
@@ -47,7 +47,9 @@ def forward(
             if self.do_grad_c("energy"):
                 model_predict["virial"] = model_ret["energy_derv_c_redu"].squeeze(-2)
                 if do_atomic_virial:
-                    model_predict["atom_virial"] = model_ret["energy_derv_c"].squeeze(-3)
+                    model_predict["atom_virial"] = model_ret["energy_derv_c"].squeeze(
+                        -3
+                    )
             else:
                 model_predict["force"] = model_ret["dforce"]
         else:
@@ -83,7 +85,9 @@ def forward_lower(
             if self.do_grad_c("energy"):
                 model_predict["virial"] = model_ret["energy_derv_c_redu"].squeeze(-2)
                 if do_atomic_virial:
-                    model_predict["atom_virial"] = model_ret["energy_derv_c"].squeeze(-3)
+                    model_predict["atom_virial"] = model_ret["energy_derv_c"].squeeze(
+                        -3
+                    )
             else:
                 assert model_ret["dforce"] is not None
                 model_predict["dforce"] = model_ret["dforce"]

From 5fad8b1ab4d3fc2b66754dc073f057c0190eda47 Mon Sep 17 00:00:00 2001
From: Anyang Peng <137014849+anyangml@users.noreply.github.com>
Date: Wed, 21 Feb 2024 12:33:48 +0800
Subject: [PATCH 10/14] fix: UTs

---
 deepmd/pt/infer/deep_eval.py | 3 +++
 1 file changed, 3 insertions(+)

diff --git a/deepmd/pt/infer/deep_eval.py b/deepmd/pt/infer/deep_eval.py
index 601bd6f755..3cb0b96270 100644
--- a/deepmd/pt/infer/deep_eval.py
+++ b/deepmd/pt/infer/deep_eval.py
@@ -373,6 +373,9 @@ def _eval_model(
                 shape = self._get_output_shape(odef, nframes, natoms)
                 out = batch_output[pt_name].reshape(shape).detach().cpu().numpy()
                 results.append(out)
+            else:
+                shape = self._get_output_shape(odef, nframes, natoms)
+                results.append(np.full(np.abs(shape), np.nan)) # this is kinda hacky
         return tuple(results)
 
     def _get_output_shape(self, odef, nframes, natoms):

From 5113288ddcf0074952aae331420d867ecb693fb5 Mon Sep 17 00:00:00 2001
From: "pre-commit-ci[bot]"
 <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Date: Wed, 21 Feb 2024 04:34:38 +0000
Subject: [PATCH 11/14] [pre-commit.ci] auto fixes from pre-commit.com hooks

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

diff --git a/deepmd/pt/infer/deep_eval.py b/deepmd/pt/infer/deep_eval.py
index 3cb0b96270..f642d34d61 100644
--- a/deepmd/pt/infer/deep_eval.py
+++ b/deepmd/pt/infer/deep_eval.py
@@ -375,7 +375,7 @@ def _eval_model(
                 results.append(out)
             else:
                 shape = self._get_output_shape(odef, nframes, natoms)
-                results.append(np.full(np.abs(shape), np.nan)) # this is kinda hacky
+                results.append(np.full(np.abs(shape), np.nan))  # this is kinda hacky
         return tuple(results)
 
     def _get_output_shape(self, odef, nframes, natoms):

From 800833a732336dff03a21814efba6fca254350f8 Mon Sep 17 00:00:00 2001
From: Anyang Peng <137014849+anyangml@users.noreply.github.com>
Date: Wed, 21 Feb 2024 13:26:11 +0800
Subject: [PATCH 12/14] fix: UTs

---
 deepmd/pt/model/atomic_model/linear_atomic_model.py  | 2 +-
 deepmd/pt/model/atomic_model/pairtab_atomic_model.py | 2 +-
 deepmd/pt/model/model/ener_model.py                  | 4 ++--
 source/tests/pt/model/test_dipole_fitting.py         | 8 ++++----
 4 files changed, 8 insertions(+), 8 deletions(-)

diff --git a/deepmd/pt/model/atomic_model/linear_atomic_model.py b/deepmd/pt/model/atomic_model/linear_atomic_model.py
index 7016629fb6..16b06b2211 100644
--- a/deepmd/pt/model/atomic_model/linear_atomic_model.py
+++ b/deepmd/pt/model/atomic_model/linear_atomic_model.py
@@ -145,7 +145,7 @@ def forward_atomic(
             the result dict, defined by the fitting net output def.
         """
         nframes, nloc, nnei = nlist.shape
-        if self.do_grad_r or self.do_grad_c():
+        if self.do_grad_r() or self.do_grad_c():
             extended_coord.requires_grad_(True)
         extended_coord = extended_coord.view(nframes, -1, 3)
         sorted_rcuts, sorted_sels = self._sort_rcuts_sels()
diff --git a/deepmd/pt/model/atomic_model/pairtab_atomic_model.py b/deepmd/pt/model/atomic_model/pairtab_atomic_model.py
index 162549fcd0..d8b830d1eb 100644
--- a/deepmd/pt/model/atomic_model/pairtab_atomic_model.py
+++ b/deepmd/pt/model/atomic_model/pairtab_atomic_model.py
@@ -148,7 +148,7 @@ def forward_atomic(
     ) -> Dict[str, torch.Tensor]:
         nframes, nloc, nnei = nlist.shape
         extended_coord = extended_coord.view(nframes, -1, 3)
-        if self.do_grad_r or self.do_grad_c():
+        if self.do_grad_r() or self.do_grad_c():
             extended_coord.requires_grad_(True)
 
         # this will mask all -1 in the nlist
diff --git a/deepmd/pt/model/model/ener_model.py b/deepmd/pt/model/model/ener_model.py
index b22f73275f..64ecfc8d10 100644
--- a/deepmd/pt/model/model/ener_model.py
+++ b/deepmd/pt/model/model/ener_model.py
@@ -81,11 +81,11 @@ def forward_lower(
             model_predict["atom_energy"] = model_ret["energy"]
             model_predict["energy"] = model_ret["energy_redu"]
             if self.do_grad_r("energy"):
-                model_predict["force"] = model_ret["energy_derv_r"].squeeze(-2)
+                model_predict["extended_force"] = model_ret["energy_derv_r"].squeeze(-2)
             if self.do_grad_c("energy"):
                 model_predict["virial"] = model_ret["energy_derv_c_redu"].squeeze(-2)
                 if do_atomic_virial:
-                    model_predict["atom_virial"] = model_ret["energy_derv_c"].squeeze(
+                    model_predict["extended_virial"] = model_ret["energy_derv_c"].squeeze(
                         -3
                     )
             else:
diff --git a/source/tests/pt/model/test_dipole_fitting.py b/source/tests/pt/model/test_dipole_fitting.py
index f64201d7d7..e6ac98f451 100644
--- a/source/tests/pt/model/test_dipole_fitting.py
+++ b/source/tests/pt/model/test_dipole_fitting.py
@@ -298,10 +298,10 @@ def setUp(self):
         self.rcut_smth = 0.5
         self.sel = [46, 92, 4]
         self.nf = 1
-        self.coord = 2 * torch.rand([self.natoms, 3], dtype=dtype, device=env.DEVICE)
-        cell = torch.rand([3, 3], dtype=dtype, device=env.DEVICE)
-        self.cell = (cell + cell.T) + 5.0 * torch.eye(3, device=env.DEVICE)
-        self.atype = torch.IntTensor([0, 0, 0, 1, 1], device=env.DEVICE)
+        self.coord = 2 * torch.rand([self.natoms, 3], dtype=dtype, device="cpu")
+        cell = torch.rand([3, 3], dtype=dtype, device="cpu")
+        self.cell = (cell + cell.T) + 5.0 * torch.eye(3, device="cpu")
+        self.atype = torch.IntTensor([0, 0, 0, 1, 1], device="cpu")
         self.dd0 = DescrptSeA(self.rcut, self.rcut_smth, self.sel).to(env.DEVICE)
         self.ft0 = DipoleFittingNet(
             "dipole",

From fdbc8fcae812a8e8879fca0aef685def3d2fe70f Mon Sep 17 00:00:00 2001
From: "pre-commit-ci[bot]"
 <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Date: Wed, 21 Feb 2024 05:28:36 +0000
Subject: [PATCH 13/14] [pre-commit.ci] auto fixes from pre-commit.com hooks

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

diff --git a/deepmd/pt/model/model/ener_model.py b/deepmd/pt/model/model/ener_model.py
index 64ecfc8d10..946cfd20f8 100644
--- a/deepmd/pt/model/model/ener_model.py
+++ b/deepmd/pt/model/model/ener_model.py
@@ -85,9 +85,9 @@ def forward_lower(
             if self.do_grad_c("energy"):
                 model_predict["virial"] = model_ret["energy_derv_c_redu"].squeeze(-2)
                 if do_atomic_virial:
-                    model_predict["extended_virial"] = model_ret["energy_derv_c"].squeeze(
-                        -3
-                    )
+                    model_predict["extended_virial"] = model_ret[
+                        "energy_derv_c"
+                    ].squeeze(-3)
             else:
                 assert model_ret["dforce"] is not None
                 model_predict["dforce"] = model_ret["dforce"]

From 5f1aec97e5f279c5f32e3d7fb396e12e8fdd9be8 Mon Sep 17 00:00:00 2001
From: Anyang Peng <137014849+anyangml@users.noreply.github.com>
Date: Wed, 21 Feb 2024 14:15:43 +0800
Subject: [PATCH 14/14] fix: CUDA

---
 source/tests/pt/model/test_dipole_fitting.py         | 4 ++--
 source/tests/pt/model/test_polarizability_fitting.py | 8 ++++----
 2 files changed, 6 insertions(+), 6 deletions(-)

diff --git a/source/tests/pt/model/test_dipole_fitting.py b/source/tests/pt/model/test_dipole_fitting.py
index e6ac98f451..fb04e49484 100644
--- a/source/tests/pt/model/test_dipole_fitting.py
+++ b/source/tests/pt/model/test_dipole_fitting.py
@@ -322,10 +322,10 @@ def test_auto_diff(self):
         atype = self.atype.view(self.nf, self.natoms)
 
         def ff(coord, atype):
-            return self.model(coord, atype)["global_dipole"].detach().numpy()
+            return self.model(coord, atype)["global_dipole"].detach().cpu().numpy()
 
         fdf = -finite_difference(ff, self.coord, atype, delta=delta)
-        rff = self.model(self.coord, atype)["force"].detach().numpy()
+        rff = self.model(self.coord, atype)["force"].detach().cpu().numpy()
 
         np.testing.assert_almost_equal(fdf, rff.transpose(0, 2, 1, 3), decimal=places)
 
diff --git a/source/tests/pt/model/test_polarizability_fitting.py b/source/tests/pt/model/test_polarizability_fitting.py
index ebe666e331..3f154383b2 100644
--- a/source/tests/pt/model/test_polarizability_fitting.py
+++ b/source/tests/pt/model/test_polarizability_fitting.py
@@ -323,10 +323,10 @@ def setUp(self):
         self.rcut_smth = 0.5
         self.sel = [46, 92, 4]
         self.nf = 1
-        self.coord = 2 * torch.rand([self.natoms, 3], dtype=dtype, device=env.DEVICE)
-        cell = torch.rand([3, 3], dtype=dtype, device=env.DEVICE)
-        self.cell = (cell + cell.T) + 5.0 * torch.eye(3, device=env.DEVICE)
-        self.atype = torch.IntTensor([0, 0, 0, 1, 1], device=env.DEVICE)
+        self.coord = 2 * torch.rand([self.natoms, 3], dtype=dtype, device="cpu")
+        cell = torch.rand([3, 3], dtype=dtype, device="cpu")
+        self.cell = (cell + cell.T) + 5.0 * torch.eye(3, device="cpu")
+        self.atype = torch.IntTensor([0, 0, 0, 1, 1], device="cpu")
         self.dd0 = DescrptSeA(self.rcut, self.rcut_smth, self.sel).to(env.DEVICE)
         self.ft0 = PolarFittingNet(
             "polar",