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

deepmd/pt/model/model/spin_model.py

@@ -471,6 +471,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]
@@ -492,6 +493,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()
@@ -607,6 +609,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,
@@ -617,6 +620,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/tf/entrypoints/freeze.py

@@ -124,6 +124,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",
@@ -259,6 +261,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:

source/api_c/include/c_api.h

@@ -230,6 +230,22 @@ extern void DP_DeepPotComputeNList(DP_DeepPot* dp,
                                    double* atomic_energy,
                                    double* atomic_virial);
 
+extern void DP_DeepPotComputeNListSP(DP_DeepPot* dp,
+                                     const int natoms,
+                                     const double* coord,
+                                     const double* spin,
+                                     const int* atype,
+                                     const double* cell,
+                                     const int nghost,
+                                     const DP_Nlist* nlist,
+                                     const int ago,
+                                     double* energy,
+                                     double* force,
+                                     double* force_mag,
+                                     double* virial,
+                                     double* atomic_energy,
+                                     double* atomic_virial);
+
 /**
  * @brief Evaluate the energy, force and virial by using a DP with the neighbor
  *list. (float version)
@@ -268,6 +284,22 @@ extern void DP_DeepPotComputeNListf(DP_DeepPot* dp,
                                     float* atomic_energy,
                                     float* atomic_virial);
 
+extern void DP_DeepPotComputeNListfSP(DP_DeepPot* dp,
+                                      const int natoms,
+                                      const float* coord,
+                                      const float* spin,
+                                      const int* atype,
+                                      const float* cell,
+                                      const int nghost,
+                                      const DP_Nlist* nlist,
+                                      const int ago,
+                                      double* energy,
+                                      float* force,
+                                      float* force_mag,
+                                      float* virial,
+                                      float* atomic_energy,
+                                      float* atomic_virial);
+
 /**
  * @brief Evaluate the energy, force and virial by using a DP. (double version)
  * @version 2
@@ -392,6 +424,25 @@ extern void DP_DeepPotComputeNList2(DP_DeepPot* dp,
                                     double* atomic_energy,
                                     double* atomic_virial);
 
+extern void DP_DeepPotComputeNList2SP(DP_DeepPot* dp,
+                                      const int nframes,
+                                      const int natoms,
+                                      const double* coord,
+                                      const double* spin,
+                                      const int* atype,
+                                      const double* cell,
+                                      const int nghost,
+                                      const DP_Nlist* nlist,
+                                      const int ago,
+                                      const double* fparam,
+                                      const double* aparam,
+                                      double* energy,
+                                      double* force,
+                                      double* force_mag,
+                                      double* virial,
+                                      double* atomic_energy,
+                                      double* atomic_virial);
+
 /**
  * @brief Evaluate the energy, force and virial by using a DP with the neighbor
  *list. (float version)
@@ -438,6 +489,25 @@ extern void DP_DeepPotComputeNListf2(DP_DeepPot* dp,
                                      float* atomic_energy,
                                      float* atomic_virial);
 
+extern void DP_DeepPotComputeNListf2SP(DP_DeepPot* dp,
+                                       const int nframes,
+                                       const int natoms,
+                                       const float* coord,
+                                       const float* spin,
+                                       const int* atype,
+                                       const float* cell,
+                                       const int nghost,
+                                       const DP_Nlist* nlist,
+                                       const int ago,
+                                       const float* fparam,
+                                       const float* aparam,
+                                       double* energy,
+                                       float* force,
+                                       float* force_mag,
+                                       float* virial,
+                                       float* atomic_energy,
+                                       float* atomic_virial);
+
 /**
  * @brief Evaluate the energy, force and virial by using a DP with the mixed
  *type. (double version)
@@ -734,6 +804,22 @@ extern void DP_DeepPotModelDeviComputeNList(DP_DeepPotModelDevi* dp,
                                             double* atomic_energy,
                                             double* atomic_virial);
 
+extern void DP_DeepPotModelDeviComputeNListSP(DP_DeepPotModelDevi* dp,
+                                              const int natoms,
+                                              const double* coord,
+                                              const double* spin,
+                                              const int* atype,
+                                              const double* cell,
+                                              const int nghost,
+                                              const DP_Nlist* nlist,
+                                              const int ago,
+                                              double* energy,
+                                              double* force,
+                                              double* force_mag,
+                                              double* virial,
+                                              double* atomic_energy,
+                                              double* atomic_virial);
+
 /**
  * @brief Evaluate the energy, force and virial by using a DP model deviation
  *with neighbor list. (float version)
@@ -771,6 +857,22 @@ extern void DP_DeepPotModelDeviComputeNListf(DP_DeepPotModelDevi* dp,
                                              float* atomic_energy,
                                              float* atomic_virial);
 
+extern void DP_DeepPotModelDeviComputeNListfSP(DP_DeepPotModelDevi* dp,
+                                               const int natoms,
+                                               const float* coord,
+                                               const float* spin,
+                                               const int* atype,
+                                               const float* cell,
+                                               const int nghost,
+                                               const DP_Nlist* nlist,
+                                               const int ago,
+                                               double* energy,
+                                               float* force,
+                                               float* force_mag,
+                                               float* virial,
+                                               float* atomic_energy,
+                                               float* atomic_virial);
+
 /**
  * @brief Evaluate the energy, force and virial by using a DP model deviation
  *with neighbor list. (double version)
@@ -816,6 +918,26 @@ void DP_DeepPotModelDeviComputeNList2(DP_DeepPotModelDevi* dp,
                                       double* virial,
                                       double* atomic_energy,
                                       double* atomic_virial);
+
+void DP_DeepPotModelDeviComputeNList2SP(DP_DeepPotModelDevi* dp,
+                                        const int nframes,
+                                        const int natoms,
+                                        const double* coord,
+                                        const double* spin,
+                                        const int* atype,
+                                        const double* cell,
+                                        const int nghost,
+                                        const DP_Nlist* nlist,
+                                        const int ago,
+                                        const double* fparam,
+                                        const double* aparam,
+                                        double* energy,
+                                        double* force,
+                                        double* force_mag,
+                                        double* virial,
+                                        double* atomic_energy,
+                                        double* atomic_virial);
+
 /**
  * @brief Evaluate the energy, force and virial by using a DP model deviation
  *with neighbor list. (float version)
@@ -862,6 +984,25 @@ void DP_DeepPotModelDeviComputeNListf2(DP_DeepPotModelDevi* dp,
                                        float* atomic_energy,
                                        float* atomic_virial);
 
+void DP_DeepPotModelDeviComputeNListf2SP(DP_DeepPotModelDevi* dp,
+                                         const int nframes,
+                                         const int natoms,
+                                         const float* coord,
+                                         const float* spin,
+                                         const int* atype,
+                                         const float* cell,
+                                         const int nghost,
+                                         const DP_Nlist* nlist,
+                                         const int ago,
+                                         const float* fparam,
+                                         const float* aparam,
+                                         double* energy,
+                                         float* force,
+                                         float* force_mag,
+                                         float* virial,
+                                         float* atomic_energy,
+                                         float* atomic_virial);
+
 /**
  * @brief Get the type map of a DP model deviation.
  * @param[in] dp The DP model deviation to use.