Spin lmp

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,45 +489,64 @@ extern void DP_DeepPotComputeNListf2(DP_DeepPot* dp,
                                      float* atomic_energy,
                                      float* atomic_virial);
 
-/**
- * @brief Evaluate the energy, force and virial by using a DP with the mixed
- *type. (double version)
- * @param[in] dp The DP to use.
- * @param[in] nframes The number of frames.
- * @param[in] natoms The number of atoms.
- * @param[in] coord The coordinates of atoms. The array should be of size natoms
- *x 3.
- * @param[in] atype The atom types. The array should contain nframes x natoms
- *ints.
- * @param[in] box The cell of the region. The array should be of size 9. Pass
- *NULL if pbc is not used.
- * @param[in] fparam The frame parameters. The array can be of size nframes x
- *dim_fparam.
- * @param[in] aparam The atom parameters. The array can be of size nframes x
- *dim_aparam.
- * @param[out] energy Output energy.
- * @param[out] force Output force. The array should be of size natoms x 3.
- * @param[out] virial Output virial. The array should be of size 9.
- * @param[out] atomic_energy Output atomic energy. The array should be of size
- *natoms.
- * @param[out] atomic_virial Output atomic virial. The array should be of size
- *natoms x 9.
- * @warning The output arrays should be allocated before calling this function.
- *Pass NULL if not required.
- **/
-extern void DP_DeepPotComputeMixedType(DP_DeepPot* dp,
+extern void DP_DeepPotComputeNListf2SP(DP_DeepPot* dp,
                                        const int nframes,
                                        const int natoms,
-                                       const double* coord,
+                                       const float* coord,
+                                       const float* spin,
                                        const int* atype,
-                                       const double* cell,
-                                       const double* fparam,
-                                       const double* aparam,
+                                       const float* cell,
+                                       const int nghost,
+                                       const DP_Nlist* nlist,
+                                       const int ago,
+                                       const float* fparam,
+                                       const float* aparam,
                                        double* energy,
-                                       double* force,
-                                       double* virial,
-                                       double* atomic_energy,
-                                       double* atomic_virial);
+                                       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)
+     * @param[in] dp The DP to use.
+     * @param[in] nframes The number of frames.
+     * @param[in] natoms The number of atoms.
+     * @param[in] coord The coordinates of atoms. The array should be of size
+     *natoms x 3.
+     * @param[in] atype The atom types. The array should contain nframes x
+     *natoms ints.
+     * @param[in] box The cell of the region. The array should be of size 9.
+     *Pass NULL if pbc is not used.
+     * @param[in] fparam The frame parameters. The array can be of size nframes
+     *x dim_fparam.
+     * @param[in] aparam The atom parameters. The array can be of size nframes x
+     *dim_aparam.
+     * @param[out] energy Output energy.
+     * @param[out] force Output force. The array should be of size natoms x 3.
+     * @param[out] virial Output virial. The array should be of size 9.
+     * @param[out] atomic_energy Output atomic energy. The array should be of
+     *size natoms.
+     * @param[out] atomic_virial Output atomic virial. The array should be of
+     *size natoms x 9.
+     * @warning The output arrays should be allocated before calling this
+     *function. Pass NULL if not required.
+     **/
+    extern void DP_DeepPotComputeMixedType(DP_DeepPot* dp,
+                                           const int nframes,
+                                           const int natoms,
+                                           const double* coord,
+                                           const int* atype,
+                                           const double* cell,
+                                           const double* fparam,
+                                           const double* aparam,
+                                           double* energy,
+                                           double* force,
+                                           double* virial,
+                                           double* atomic_energy,
+                                           double* atomic_virial);
 /**
  * @brief Evaluate the energy, force and virial by using a DP with the mixed
  *type. (float 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,12 +984,31 @@ void DP_DeepPotModelDeviComputeNListf2(DP_DeepPotModelDevi* dp,
                                        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.
- * @return The cutoff radius.
- */
-double DP_DeepPotModelDeviGetCutoff(DP_DeepPotModelDevi* dp);
+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.
+     * @return The cutoff radius.
+     */
+    double DP_DeepPotModelDeviGetCutoff(DP_DeepPotModelDevi* dp);
 
 /**
  * @brief Get the number of types of a DP model deviation.