feat(pt_expt): implement DeepSpin model in pt_expt backend

deepmd/dpmodel/model/make_model.py

@@ -68,6 +68,7 @@ def model_call_from_call_lower(
     fparam: Array | None = None,
     aparam: Array | None = None,
     do_atomic_virial: bool = False,
+    coord_corr_for_virial: Array | None = None,
 ) -> dict[str, Array]:
     """Return model prediction from lower interface.
 
@@ -119,6 +120,20 @@ def model_call_from_call_lower(
         distinguish_types=False,
     )
     extended_coord = extended_coord.reshape(nframes, -1, 3)
+    if coord_corr_for_virial is not None:
+        xp = array_api_compat.array_namespace(coord_corr_for_virial)
+        # mapping: nf x nall -> nf x nall x 1, then tile to nf x nall x 3
+        mapping_idx = xp.tile(
+            xp.reshape(mapping, (nframes, -1, 1)),
+            (1, 1, 3),
+        )
+        extended_coord_corr = xp.take_along_axis(
+            coord_corr_for_virial,
+            mapping_idx,
+            axis=1,
+        )
+    else:
+        extended_coord_corr = None
     model_predict_lower = call_lower(
         extended_coord,
         extended_atype,
@@ -127,6 +142,7 @@ def model_call_from_call_lower(
         fparam=fp,
         aparam=ap,
         do_atomic_virial=do_atomic_virial,
+        extended_coord_corr=extended_coord_corr,
     )
     model_predict = communicate_extended_output(
         model_predict_lower,
@@ -237,6 +253,7 @@ def call_common(
             fparam: Array | None = None,
             aparam: Array | None = None,
             do_atomic_virial: bool = False,
+            coord_corr_for_virial: Array | None = None,
         ) -> dict[str, Array]:
             """Return model prediction.
 
@@ -255,6 +272,9 @@ def call_common(
                 atomic parameter. nf x nloc x nda
             do_atomic_virial
                 If calculate the atomic virial.
+            coord_corr_for_virial
+                The coordinates correction for virial.
+                shape: nf x (nloc x 3)
 
             Returns
             -------
@@ -279,6 +299,7 @@ def call_common(
                 fparam=fp,
                 aparam=ap,
                 do_atomic_virial=do_atomic_virial,
+                coord_corr_for_virial=coord_corr_for_virial,
             )
             model_predict = self._output_type_cast(model_predict, input_prec)
             return model_predict
@@ -292,6 +313,7 @@ def call_common_lower(
             fparam: Array | None = None,
             aparam: Array | None = None,
             do_atomic_virial: bool = False,
+            extended_coord_corr: Array | None = None,
         ) -> dict[str, Array]:
             """Return model prediction. Lower interface that takes
             extended atomic coordinates and types, nlist, and mapping
@@ -314,6 +336,9 @@ def call_common_lower(
                 atomic parameter. nf x nloc x nda
             do_atomic_virial
                 whether calculate atomic virial
+            extended_coord_corr
+                coordinates correction for virial in extended region.
+                nf x (nall x 3)
 
             Returns
             -------
@@ -341,6 +366,7 @@ def call_common_lower(
                 fparam=fp,
                 aparam=ap,
                 do_atomic_virial=do_atomic_virial,
+                extended_coord_corr=extended_coord_corr,
             )
             model_predict = self._output_type_cast(model_predict, input_prec)
             return model_predict
@@ -354,6 +380,7 @@ def forward_common_atomic(
             fparam: Array | None = None,
             aparam: Array | None = None,
             do_atomic_virial: bool = False,
+            extended_coord_corr: Array | None = None,
         ) -> dict[str, Array]:
             atomic_ret = self.atomic_model.forward_common_atomic(
                 extended_coord,

deepmd/dpmodel/model/model.py

@@ -164,6 +164,7 @@ def get_spin_model(data: dict) -> SpinModel:
     data : dict
         The data to construct the model.
     """
+    data = copy.deepcopy(data)
     # include virtual spin and placeholder types
     data["type_map"] += [item + "_spin" for item in data["type_map"]]
     spin = Spin(