Update virial and add UT.

Author: OutisLi

deepmd/tf/descriptor/se_a.py

@@ -670,6 +670,23 @@ def build(
         atype_t = tf.concat([[self.ntypes], tf.reshape(self.atype, [-1])], axis=0)
         self.nei_type_vec = tf.nn.embedding_lookup(atype_t, nlist_t)
 
+        if self.spin is not None:
+            judge = tf.equal(natoms[0], natoms[1])
+            self.diff_coord = tf.cond(
+                judge,
+                lambda: self.natoms_match(coord, natoms),
+                lambda: self.natoms_not_match(coord, natoms, atype),
+            )
+            diff_coord_reshape = tf.reshape(self.diff_coord, [-1, natoms[1], 3])
+            nlist_reshaped = tf.reshape(self.nlist, [-1, natoms[0], self.nnei])
+            rj_gathered = tf.gather(
+                diff_coord_reshape, tf.maximum(nlist_reshaped, 0), axis=1, batch_dims=1
+            )
+            ri_gathered = diff_coord_reshape[:, : natoms[0], tf.newaxis, :]
+            rij_reshaped = tf.reshape(self.rij, [-1, natoms[0], self.nnei, 3])
+            rij_update = rij_reshaped - rj_gathered + ri_gathered
+            self.rij = tf.reshape(rij_update, tf.shape(self.rij))
+
         # only used when tensorboard was set as true
         tf.summary.histogram("descrpt", self.descrpt)
         tf.summary.histogram("rij", self.rij)
@@ -1374,6 +1391,62 @@ def init_variables(
                 )
             )
 
+    def natoms_match(self, coord, natoms):
+        natoms_index = tf.concat([[0], tf.cumsum(natoms[2:])], axis=0)
+        diff_coord_loc = []
+        for i in range(self.ntypes):
+            if i + self.ntypes_spin >= self.ntypes:
+                diff_coord_loc.append(
+                    tf.slice(
+                        coord,
+                        [0, natoms_index[i] * 3],
+                        [-1, natoms[2 + i] * 3],
+                    )
+                    - tf.slice(
+                        coord,
+                        [0, natoms_index[i - len(self.spin.use_spin)] * 3],
+                        [-1, natoms[2 + i - len(self.spin.use_spin)] * 3],
+                    )
+                )
+            else:
+                diff_coord_loc.append(
+                    tf.zeros([tf.shape(coord)[0], natoms[2 + i] * 3], dtype=coord.dtype)
+                )
+        diff_coord_loc = tf.concat(diff_coord_loc, axis=1)
+        return diff_coord_loc
+
+    def natoms_not_match(self, coord, natoms, atype):
+        diff_coord_loc = self.natoms_match(coord, natoms)
+        diff_coord_ghost = []
+        aatype = atype[0, :]
+        ghost_atype = aatype[natoms[0] :]
+        _, _, ghost_natoms = tf.unique_with_counts(ghost_atype)
+        ghost_natoms_index = tf.concat([[0], tf.cumsum(ghost_natoms)], axis=0)
+        ghost_natoms_index += natoms[0]
+        for i in range(self.ntypes):
+            if i + self.ntypes_spin >= self.ntypes:
+                diff_coord_ghost.append(
+                    tf.slice(
+                        coord,
+                        [0, ghost_natoms_index[i] * 3],
+                        [-1, ghost_natoms[i] * 3],
+                    )
+                    - tf.slice(
+                        coord,
+                        [0, ghost_natoms_index[i - len(self.spin.use_spin)] * 3],
+                        [-1, ghost_natoms[i - len(self.spin.use_spin)] * 3],
+                    )
+                )
+            else:
+                diff_coord_ghost.append(
+                    tf.zeros(
+                        [tf.shape(coord)[0], ghost_natoms[i] * 3], dtype=coord.dtype
+                    )
+                )
+        diff_coord_ghost = tf.concat(diff_coord_ghost, axis=1)
+        diff_coord = tf.concat([diff_coord_loc, diff_coord_ghost], axis=1)
+        return diff_coord
+
     @property
     def explicit_ntypes(self) -> bool:
         """Explicit ntypes with type embedding."""

source/api_cc/src/DeepSpinTF.cc

@@ -809,11 +809,11 @@ void DeepSpinTF::compute(ENERGYVTYPE& dener,
   datom_energy_.resize(static_cast<size_t>(nframes) * nall);
   datom_virial_.resize(static_cast<size_t>(nframes) * nall * 9);
   for (int ii = 0; ii < nall; ++ii) {
+    int new_idx = new_idx_map[ii];
     for (int dd = 0; dd < 3; ++dd) {
-      int new_idx = new_idx_map[ii];
       dforce_[3 * ii + dd] = dforce_tmp[3 * new_idx + dd];
       datom_energy_[ii] = datom_energy_tmp[new_idx];
-      datom_virial_[ii] = datom_virial_tmp[new_idx];
+
       if (datype_[ii] < ntypes_spin && ii < nloc) {
         dforce_mag_[3 * ii + dd] = dforce_tmp[3 * (new_idx + nloc) + dd];
       } else if (datype_[ii] < ntypes_spin) {
@@ -822,6 +822,9 @@ void DeepSpinTF::compute(ENERGYVTYPE& dener,
         dforce_mag_[3 * ii + dd] = 0.0;
       }
     }
+    for (int dd = 0; dd < 9; ++dd) {
+      datom_virial_[ii * 9 + dd] = datom_virial_tmp[new_idx * 9 + dd];
+    }
   }
 }
 

source/tests/tf/test_model_spin.py

@@ -126,6 +126,7 @@ def test_model_spin(self) -> None:
         )
 
         out_ener = np.reshape(out_ener, [-1])
+        out_virial = np.reshape(out_virial, [-1])
         natoms_real = np.sum(
             test_data["natoms_vec"][2 : 2 + len(spin_param["use_spin"])]
         )
@@ -425,14 +426,36 @@ def test_model_spin(self) -> None:
             -0.0007015535942944564,
             0.004459188855221506,
         ]
+        refv = [
+            0.33691325723275595,
+            0.024301747372056412,
+            -0.06880806009046331,
+            0.026792188153995887,
+            0.3007953744219118,
+            -0.051612531097108075,
+            -0.07274496651648972,
+            -0.05156414798680478,
+            0.37692901508963417,
+            0.3205610686355494,
+            0.013102936385366228,
+            -0.04419007538301404,
+            0.014186144311082909,
+            0.31565216176483,
+            -0.058829665227551474,
+            -0.04759429793837308,
+            -0.05932221615318792,
+            0.39040431257661773,
+        ]
         refe = np.reshape(refe, [-1])
         refr = np.reshape(refr, [-1])
         refm = np.reshape(refm, [-1])
+        refv = np.reshape(refv, [-1])
 
         places = 10
         np.testing.assert_almost_equal(out_ener, refe, places)
         np.testing.assert_almost_equal(force_real, refr, places)
         np.testing.assert_almost_equal(force_mag, refm, places)
+        np.testing.assert_almost_equal(out_virial, refv, places)
 
 
 if __name__ == "__main__":