feat(tf/pt): add/refact lammps support for spin model

deepmd/pt/model/descriptor/repformers.py

@@ -28,6 +28,9 @@
 from deepmd.pt.utils.exclude_mask import (
     PairExcludeMask,
 )
+from deepmd.pt.utils.spin import (
+    concat_switch_virtual,
+)
 from deepmd.pt.utils.utils import (
     ActivationFn,
 )
@@ -422,6 +425,7 @@ def forward(
             atype_embd = extended_atype_embd
         assert isinstance(atype_embd, torch.Tensor)  # for jit
         g1 = self.act(atype_embd)
+        ng1 = g1.shape[-1]
         # nb x nloc x nnei x 1,  nb x nloc x nnei x 3
         if not self.direct_dist:
             g2, h2 = torch.split(dmatrix, [1, 3], dim=-1)
@@ -448,10 +452,27 @@ def forward(
                 assert mapping is not None
                 g1_ext = torch.gather(g1, 1, mapping)
             else:
-                n_padding = nall - nloc
-                g1 = torch.nn.functional.pad(
-                    g1.squeeze(0), (0, 0, 0, n_padding), value=0.0
-                )
+                has_spin = "has_spin" in comm_dict
+                if not has_spin:
+                    n_padding = nall - nloc
+                    g1 = torch.nn.functional.pad(
+                        g1.squeeze(0), (0, 0, 0, n_padding), value=0.0
+                    )
+                    real_nloc = nloc
+                    real_nall = nall
+                else:
+                    # for spin
+                    real_nloc = nloc // 2
+                    real_nall = nall // 2
+                    real_n_padding = real_nall - real_nloc
+                    g1_real, g1_virtual = torch.split(g1, [real_nloc, real_nloc], dim=1)
+                    # mix_g1: nb x real_nloc x (ng1 * 2)
+                    mix_g1 = torch.cat([g1_real, g1_virtual], dim=2)
+                    # nb x real_nall x (ng1 * 2)
+                    g1 = torch.nn.functional.pad(
+                        mix_g1.squeeze(0), (0, 0, 0, real_n_padding), value=0.0
+                    )
+
                 assert "send_list" in comm_dict
                 assert "send_proc" in comm_dict
                 assert "recv_proc" in comm_dict
@@ -466,10 +487,15 @@ def forward(
                     comm_dict["recv_num"],
                     g1,
                     comm_dict["communicator"],
-                    torch.tensor(nloc),  # pylint: disable=no-explicit-dtype,no-explicit-device
-                    torch.tensor(nall - nloc),  # pylint: disable=no-explicit-dtype,no-explicit-device
+                    torch.tensor(real_nloc),  # pylint: disable=no-explicit-dtype,no-explicit-device
+                    torch.tensor(real_nall - real_nloc),  # pylint: disable=no-explicit-dtype,no-explicit-device
                 )
                 g1_ext = ret[0].unsqueeze(0)
+                if has_spin:
+                    g1_real_ext, g1_virtual_ext = torch.split(g1_ext, [ng1, ng1], dim=2)
+                    g1_ext = concat_switch_virtual(
+                        g1_real_ext, g1_virtual_ext, real_nloc
+                    )
             g1, g2, h2 = ll.forward(
                 g1_ext,
                 g2,

deepmd/pt/model/model/spin_model.py

@@ -15,6 +15,9 @@
 from deepmd.pt.model.atomic_model import (
     DPAtomicModel,
 )
+from deepmd.pt.utils.spin import (
+    concat_switch_virtual,
+)
 from deepmd.pt.utils.utils import (
     to_torch_tensor,
 )
@@ -79,15 +82,15 @@ def process_spin_input_lower(
             self.virtual_scale_mask.to(extended_atype.device)
         )[extended_atype].reshape([nframes, nall, 1])
         virtual_extended_atype = extended_atype + self.ntypes_real
-        extended_coord_updated = self.concat_switch_virtual(
+        extended_coord_updated = concat_switch_virtual(
             extended_coord, virtual_extended_coord, nloc
         )
-        extended_atype_updated = self.concat_switch_virtual(
+        extended_atype_updated = concat_switch_virtual(
             extended_atype, virtual_extended_atype, nloc
         )
         if mapping is not None:
             virtual_mapping = mapping + nloc
-            mapping_updated = self.concat_switch_virtual(mapping, virtual_mapping, nloc)
+            mapping_updated = concat_switch_virtual(mapping, virtual_mapping, nloc)
         else:
             mapping_updated = None
         # extend the nlist
@@ -203,33 +206,6 @@ def extend_nlist(extended_atype, nlist):
         extended_nlist[second_part_index] -= nall - nloc
         return extended_nlist
 
-    @staticmethod
-    def concat_switch_virtual(extended_tensor, extended_tensor_virtual, nloc: int):
-        """
-        Concat real and virtual extended tensors, and switch all the local ones to the first nloc * 2 atoms.
-        - [:, :nloc]: original nloc real atoms.
-        - [:, nloc: nloc + nloc]: virtual atoms corresponding to nloc real atoms.
-        - [:, nloc + nloc: nloc + nall]: ghost real atoms.
-        - [:, nloc + nall: nall + nall]: virtual atoms corresponding to ghost real atoms.
-        """
-        nframes, nall = extended_tensor.shape[:2]
-        out_shape = list(extended_tensor.shape)
-        out_shape[1] *= 2
-        extended_tensor_updated = torch.zeros(
-            out_shape,
-            dtype=extended_tensor.dtype,
-            device=extended_tensor.device,
-        )
-        extended_tensor_updated[:, :nloc] = extended_tensor[:, :nloc]
-        extended_tensor_updated[:, nloc : nloc + nloc] = extended_tensor_virtual[
-            :, :nloc
-        ]
-        extended_tensor_updated[:, nloc + nloc : nloc + nall] = extended_tensor[
-            :, nloc:
-        ]
-        extended_tensor_updated[:, nloc + nall :] = extended_tensor_virtual[:, nloc:]
-        return extended_tensor_updated.view(out_shape)
-
     @staticmethod
     def expand_aparam(aparam, nloc: int):
         """Expand the atom parameters for virtual atoms if necessary."""
@@ -469,6 +445,7 @@ def forward_common_lower(
         fparam: Optional[torch.Tensor] = None,
         aparam: Optional[torch.Tensor] = None,
         do_atomic_virial: bool = False,
+        comm_dict: Optional[dict[str, torch.Tensor]] = None,
         extra_nlist_sort: bool = False,
     ):
         nframes, nloc = nlist.shape[:2]
@@ -490,6 +467,7 @@ def forward_common_lower(
             fparam=fparam,
             aparam=aparam,
             do_atomic_virial=do_atomic_virial,
+            comm_dict=comm_dict,
             extra_nlist_sort=extra_nlist_sort,
         )
         model_output_type = self.backbone_model.model_output_type()
@@ -605,6 +583,7 @@ def forward_lower(
         fparam: Optional[torch.Tensor] = None,
         aparam: Optional[torch.Tensor] = None,
         do_atomic_virial: bool = False,
+        comm_dict: Optional[dict[str, torch.Tensor]] = None,
     ):
         model_ret = self.forward_common_lower(
             extended_coord,
@@ -615,6 +594,7 @@ def forward_lower(
             fparam=fparam,
             aparam=aparam,
             do_atomic_virial=do_atomic_virial,
+            comm_dict=comm_dict,
             extra_nlist_sort=self.backbone_model.need_sorted_nlist_for_lower(),
         )
         model_predict = {}

deepmd/pt/utils/spin.py

@@ -0,0 +1,30 @@
+# SPDX-License-Identifier: LGPL-3.0-or-later
+
+import torch
+
+
+def concat_switch_virtual(
+    extended_tensor,
+    extended_tensor_virtual,
+    nloc: int,
+):
+    """
+    Concat real and virtual extended tensors, and switch all the local ones to the first nloc * 2 atoms.
+    - [:, :nloc]: original nloc real atoms.
+    - [:, nloc: nloc + nloc]: virtual atoms corresponding to nloc real atoms.
+    - [:, nloc + nloc: nloc + nall]: ghost real atoms.
+    - [:, nloc + nall: nall + nall]: virtual atoms corresponding to ghost real atoms.
+    """
+    nframes, nall = extended_tensor.shape[:2]
+    out_shape = list(extended_tensor.shape)
+    out_shape[1] *= 2
+    extended_tensor_updated = torch.zeros(
+        out_shape,
+        dtype=extended_tensor.dtype,
+        device=extended_tensor.device,
+    )
+    extended_tensor_updated[:, :nloc] = extended_tensor[:, :nloc]
+    extended_tensor_updated[:, nloc : nloc + nloc] = extended_tensor_virtual[:, :nloc]
+    extended_tensor_updated[:, nloc + nloc : nloc + nall] = extended_tensor[:, nloc:]
+    extended_tensor_updated[:, nloc + nall :] = extended_tensor_virtual[:, nloc:]
+    return extended_tensor_updated.view(out_shape)

deepmd/tf/entrypoints/freeze.py

@@ -123,6 +123,8 @@ def _make_node_names(
             "o_atom_energy",
             "o_atom_virial",
             "spin_attr/ntypes_spin",
+            "spin_attr/virtual_len",
+            "spin_attr/spin_norm",
             "fitting_attr/dfparam",
             "fitting_attr/daparam",
             "fitting_attr/aparam_nall",
@@ -258,6 +260,8 @@ def freeze_graph(
         "train_attr/min_nbor_dist",
         "fitting_attr/aparam_nall",
         "spin_attr/ntypes_spin",
+        "spin_attr/virtual_len",
+        "spin_attr/spin_norm",
     ]
     different_set = set(output_node) - set(input_node)
     if different_set: