Merge branch 'master' into 0304_add_chg_spin

Author: iProzd

.github/workflows/build_wheel.yml

@@ -68,7 +68,7 @@ jobs:
           CUDA_VERSION: ${{ matrix.cuda_version }}
           DP_PKG_NAME: ${{ matrix.dp_pkg_name }}
           CIBW_BUILD_FRONTEND: "build[uv]"
-      - uses: actions/upload-artifact@v6
+      - uses: actions/upload-artifact@v7
         with:
           name: cibw-cp${{ matrix.python }}-${{ matrix.platform_id }}-cu${{ matrix.cuda_version }}-${{ strategy.job-index }}
           path: ./wheelhouse/*.whl
@@ -82,7 +82,7 @@ jobs:
       - name: Build sdist
         run: pipx run uv tool run --with build[uv] --from build python -m build --installer uv --sdist
 
-      - uses: actions/upload-artifact@v6
+      - uses: actions/upload-artifact@v7
         with:
           name: cibw-sdist
           path: dist/*.tar.gz
@@ -95,7 +95,7 @@ jobs:
       id-token: write
     if: github.event_name == 'push' && startsWith(github.event.ref, 'refs/tags/v')
     steps:
-      - uses: actions/download-artifact@v7
+      - uses: actions/download-artifact@v8
         with:
           pattern: cibw-*
           path: dist
@@ -124,13 +124,13 @@ jobs:
           swap-storage: true
           docker-images: true
       - uses: actions/checkout@v6
-      - uses: actions/download-artifact@v7
+      - uses: actions/download-artifact@v8
         with:
           path: source/install/docker/dist
           pattern: cibw-*-manylinux_x86_64-cu${{ matrix.cuda_version }}*
           merge-multiple: true
       - name: Log in to the Container registry
-        uses: docker/login-action@v3
+        uses: docker/login-action@v4
         with:
           registry: ghcr.io
           username: ${{ github.actor }}
@@ -157,7 +157,7 @@ jobs:
     needs: [build_wheels, build_sdist]
     runs-on: ubuntu-latest
     steps:
-      - uses: actions/download-artifact@v7
+      - uses: actions/download-artifact@v8
         with:
           path: dist/packages
           pattern: cibw-*

.github/workflows/package_c.yml

@@ -47,7 +47,7 @@ jobs:
         if: matrix.filename != 'libdeepmd_c.tar.gz'
       # for download and debug
       - name: Upload artifact
-        uses: actions/upload-artifact@v6
+        uses: actions/upload-artifact@v7
         with:
           name: libdeepmd_c-${{ strategy.job-index }}-${{ matrix.filename }}
           path: ${{ matrix.filename }}
@@ -65,7 +65,7 @@ jobs:
     steps:
       - uses: actions/checkout@v6
       - name: Download artifact
-        uses: actions/download-artifact@v7
+        uses: actions/download-artifact@v8
         with:
           pattern: libdeepmd_c-*
           merge-multiple: true

.github/workflows/test_python.yml

@@ -71,7 +71,7 @@ jobs:
         if: matrix.group == 1
       - run: mv .test_durations .test_durations_${{ matrix.group }}
       - name: Upload partial durations
-        uses: actions/upload-artifact@v6
+        uses: actions/upload-artifact@v7
         with:
           name: split-${{ matrix.python }}-${{ matrix.group }}
           path: .test_durations_${{ matrix.group }}
@@ -100,7 +100,7 @@ jobs:
           key: test2-durations-combined-${{ matrix.python }}-${{ github.sha }}
           restore-keys: test2-durations-combined-${{ matrix.python }}
       - name: Download artifacts
-        uses: actions/download-artifact@v7
+        uses: actions/download-artifact@v8
         with:
           pattern: split-${{ matrix.python }}-*
           merge-multiple: true

deepmd/dpmodel/model/make_hessian_model.py

@@ -0,0 +1,59 @@
+# SPDX-License-Identifier: LGPL-3.0-or-later
+import copy
+from typing import (
+    Any,
+)
+
+from deepmd.dpmodel.output_def import (
+    FittingOutputDef,
+)
+
+
+def make_hessian_model(T_Model: type) -> type:
+    """Make a model that can compute Hessian.
+
+    With the JAX-mirrored approach, hessian is computed in
+    ``forward_common_atomic`` (in make_model.py) on extended coordinates.
+    This wrapper only needs to override ``atomic_output_def()`` to set
+    ``r_hessian=True``, and ``communicate_extended_output`` in dpmodel
+    naturally maps it from nall to nloc.
+
+    Parameters
+    ----------
+    T_Model
+        The model. Should provide the ``atomic_output_def`` method.
+
+    Returns
+    -------
+    The model that computes hessian.
+
+    """
+
+    class CM(T_Model):
+        def __init__(
+            self,
+            *args: Any,
+            **kwargs: Any,
+        ) -> None:
+            super().__init__(
+                *args,
+                **kwargs,
+            )
+            self.hess_fitting_def = copy.deepcopy(super().atomic_output_def())
+
+        def requires_hessian(
+            self,
+            keys: str | list[str],
+        ) -> None:
+            """Set which output variable(s) requires hessian."""
+            if isinstance(keys, str):
+                keys = [keys]
+            for kk in self.hess_fitting_def.keys():
+                if kk in keys:
+                    self.hess_fitting_def[kk].r_hessian = True
+
+        def atomic_output_def(self) -> FittingOutputDef:
+            """Get the fitting output def."""
+            return self.hess_fitting_def
+
+    return CM

deepmd/pt_expt/model/__init__.py

@@ -1,4 +1,8 @@
 # SPDX-License-Identifier: LGPL-3.0-or-later
+from deepmd.dpmodel.model.make_hessian_model import (
+    make_hessian_model,
+)
+
 from .dipole_model import (
     DipoleModel,
 )
@@ -33,4 +37,5 @@
     "PolarModel",
     "PropertyModel",
     "get_model",
+    "make_hessian_model",
 ]

deepmd/pt_expt/model/ener_model.py

@@ -1,4 +1,5 @@
 # SPDX-License-Identifier: LGPL-3.0-or-later
+import copy
 from typing import (
     Any,
 )
@@ -14,6 +15,9 @@
 from deepmd.dpmodel.model.dp_model import (
     DPModelCommon,
 )
+from deepmd.dpmodel.model.make_hessian_model import (
+    make_hessian_model,
+)
 
 from .make_model import (
     make_model,
@@ -34,6 +38,17 @@ def __init__(
     ) -> None:
         DPModelCommon.__init__(self)
         DPEnergyModel_.__init__(self, *args, **kwargs)
+        self._hessian_enabled = False
+
+    def enable_hessian(self) -> None:
+        if self._hessian_enabled:
+            return
+        self.__class__ = make_hessian_model(type(self))
+        self.hess_fitting_def = copy.deepcopy(
+            super(type(self), self).atomic_output_def()
+        )
+        self.requires_hessian("energy")
+        self._hessian_enabled = True
 
     def forward(
         self,
@@ -63,6 +78,8 @@ def forward(
                 model_predict["atom_virial"] = model_ret["energy_derv_c"].squeeze(-2)
         if "mask" in model_ret:
             model_predict["mask"] = model_ret["mask"]
+        if self.atomic_output_def()["energy"].r_hessian:
+            model_predict["hessian"] = model_ret["energy_derv_r_derv_r"].squeeze(-3)
         return model_predict
 
     def forward_lower(
@@ -115,6 +132,8 @@ def translated_output_def(self) -> dict[str, Any]:
             output_def["atom_virial"].squeeze(-2)
         if "mask" in out_def_data:
             output_def["mask"] = out_def_data["mask"]
+        if self.atomic_output_def()["energy"].r_hessian:
+            output_def["hessian"] = out_def_data["energy_derv_r_derv_r"]
         return output_def
 
     def forward_lower_exportable(