fix(tf): handle multiframe DeepSpin extension (#5851)

Closes #5660.

## Summary

- build TensorFlow DeepSpin physical and virtual coordinate buffers for
every input frame in both direct and external-neighbor-list paths
- separate extended coordinate frame strides from original-atom `aparam`
strides during NULL-atom selection
- copy physical force, magnetic force, atomic energy, and atomic virial
outputs with explicit source/destination frame offsets
- derive the atomic-energy source width from the TensorFlow tensor,
since spin models emit energies only for physical atoms while C++
preserves extended virtual slots
- add distinct two-frame direct and InputNlist atomic regressions for
both `double` and `float`

## Why existing tests missed this

All TensorFlow DeepSpin C++ tests used scalar-energy, single-frame
calls. The extension helpers therefore only needed frame-zero buffers,
and frame-zero copyback filled every asserted output. No test exercised
the documented vector-energy overload with `nframes > 1`, so undersized
extension buffers, incorrect inferred frame counts, and missing output
frame offsets remained invisible.

The new tests perturb both coordinates and spins in frame 2, evaluate
each frame separately, and compare every batched slice against its
one-frame reference. They also assert the frame references differ,
preventing duplicated-frame behavior from passing.

## Validation

- `ruff format .`
- `ruff check .`
- `clang-format --dry-run --Werror` on all changed C++ files
- built `deepmd_op`, `deepmd_backend_tf`, and `runUnitTests_cc`
- all 16 `TestInferDeepSpin` and `TestInferDeepSpinNopbc` typed tests
passed
- focused four-case two-frame matrix passed: direct/InputNlist ×
double/float

Coding agent: Codex
Codex version: codex-cli 0.144.4
Model: gpt-5.6-sol
Reasoning effort: xhigh


<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->

## Summary by CodeRabbit

* **New Features**
* Added support for processing multiple frames in spin-model
calculations.
* Extended coordinate and neighbor-list handling to support frame-aware
inputs.
* Preserved correct real-atom outputs across force, magnetic-force,
energy, and virial results.

* **Bug Fixes**
  * Improved validation of coordinate, spin, and output dimensions.
* Corrected handling of virtual-atom slots and per-frame output layouts.
* Fixed multi-frame behavior for standard and LMP neighbor-list
calculations.

<!-- end of auto-generated comment: release notes by coderabbit.ai -->

---------

Co-authored-by: njzjz-bot <njzjz.bot@gmail.com>
Co-authored-by: njzjz-bot <njzjz-bot@users.noreply.github.com>

Author: njzjz-bot

source/api_cc/include/DeepSpinTF.h

@@ -259,14 +259,16 @@ class DeepSpinTF : public DeepSpinBackend {
               const int nghost,
               const std::vector<VALUETYPE>& spin,
               const int numb_types,
-              const int numb_types_spin);
+              const int numb_types_spin,
+              const int nframes);
 
   template <typename VALUETYPE>
   void extend_nlist(std::vector<VALUETYPE>& extend_dcoord,
                     std::vector<int>& extend_atype,
                     const std::vector<VALUETYPE>& dcoord_,
                     const std::vector<VALUETYPE>& dspin_,
-                    const std::vector<int>& datype_);
+                    const std::vector<int>& datype_,
+                    const int nframes);
 
   void cum_sum(std::map<int, int>&, std::map<int, int>&);
 

source/api_cc/src/DeepSpinTF.cc

@@ -176,9 +176,23 @@ static void run_model(
   for (size_t ii = 0; ii < static_cast<size_t>(nframes) * nall * 3; ++ii) {
     dforce[ii] = of(ii);
   }
+  if (output_ae.NumElements() % nframes != 0) {
+    throw deepmd::deepmd_exception(
+        "TensorFlow atomic-energy output is not divisible by nframes.");
+  }
+  const size_t nloc_energy = output_ae.NumElements() / nframes;
+  if (nloc_energy > nall) {
+    throw deepmd::deepmd_exception(
+        "TensorFlow atomic-energy output has more atoms than the extended "
+        "DeepSpin system.");
+  }
+  // Spin models emit atomic energies for physical atoms only. The extended
+  // virtual atoms sort after all physical types, so keep their slots zero and
+  // use the actual output width as the per-frame source stride.
   for (int ii = 0; ii < nframes; ++ii) {
-    for (int jj = 0; jj < nloc; ++jj) {
-      datom_energy[ii * nall + jj] = oae(ii * nloc + jj);
+    for (size_t jj = 0; jj < nloc_energy; ++jj) {
+      datom_energy[static_cast<size_t>(ii) * nall + jj] =
+          oae(static_cast<size_t>(ii) * nloc_energy + jj);
     }
   }
   for (size_t ii = 0; ii < static_cast<size_t>(nframes) * nall * 9; ++ii) {
@@ -614,7 +628,7 @@ void DeepSpinTF::compute(ENERGYVTYPE& dener,
 
   std::vector<VALUETYPE> extend_dcoord;
   std::vector<int> extend_atype;
-  extend_nlist(extend_dcoord, extend_atype, dcoord_, dspin_, datype_);
+  extend_nlist(extend_dcoord, extend_atype, dcoord_, dspin_, datype_, nframes);
 
   atommap = deepmd::AtomMap(extend_atype.begin(), extend_atype.end());
 
@@ -645,16 +659,43 @@ void DeepSpinTF::compute(ENERGYVTYPE& dener,
                        atommap, nframes);
     }
   }
-  // backward force and mag.
+  // Atomic outputs from the TensorFlow graph include the appended virtual spin
+  // atoms.  Keep a temporary copy so the public API can return the documented
+  // real-atom layout, matching the neighbor-list overload below.
+  std::vector<VALUETYPE> datom_energy_tmp, datom_virial_tmp;
+  if (atomic) {
+    datom_energy_tmp.swap(datom_energy_);
+    datom_virial_tmp.swap(datom_virial_);
+    datom_energy_.resize(static_cast<size_t>(nframes) * nloc);
+    datom_virial_.resize(static_cast<size_t>(nframes) * nloc * 9);
+  }
+
+  // Backward force, magnetic force, and optional atomic outputs.
   dforce_.resize(static_cast<size_t>(nframes) * nloc * 3);
   dforce_mag_.resize(static_cast<size_t>(nframes) * nloc * 3);
-  for (int ii = 0; ii < nloc; ++ii) {
-    for (int dd = 0; dd < 3; ++dd) {
-      dforce_[3 * ii + dd] = dforce_tmp[3 * ii + dd];
-      if (datype_[ii] < ntypes_spin) {
-        dforce_mag_[3 * ii + dd] = dforce_tmp[3 * (ii + nloc) + dd];
-      } else {
-        dforce_mag_[3 * ii + dd] = 0.0;
+  const size_t extend_nall = extend_atype.size();
+  for (int ff = 0; ff < nframes; ++ff) {
+    for (int ii = 0; ii < nloc; ++ii) {
+      const size_t output_atom = static_cast<size_t>(ff) * nloc + ii;
+      const size_t extended_atom = static_cast<size_t>(ff) * extend_nall + ii;
+      if (atomic) {
+        datom_energy_[output_atom] = datom_energy_tmp[extended_atom];
+      }
+      for (int dd = 0; dd < 3; ++dd) {
+        dforce_[output_atom * 3 + dd] = dforce_tmp[extended_atom * 3 + dd];
+        if (datype_[ii] < ntypes_spin) {
+          const size_t virtual_atom =
+              static_cast<size_t>(ff) * extend_nall + ii + nloc;
+          dforce_mag_[output_atom * 3 + dd] = dforce_tmp[virtual_atom * 3 + dd];
+        } else {
+          dforce_mag_[output_atom * 3 + dd] = 0.0;
+        }
+      }
+      if (atomic) {
+        for (int dd = 0; dd < 9; ++dd) {
+          datom_virial_[output_atom * 9 + dd] =
+              datom_virial_tmp[extended_atom * 9 + dd];
+        }
       }
     }
   }
@@ -747,7 +788,7 @@ void DeepSpinTF::compute(ENERGYVTYPE& dener,
   extend(extend_inum, extend_ilist, extend_numneigh, extend_neigh,
          extend_firstneigh, extend_dcoord, extend_dtype, extend_nghost,
          new_idx_map, old_idx_map, lmp_list, dcoord_, datype_, nghost, dspin_,
-         ntypes, ntypes_spin);
+         ntypes, ntypes_spin, nframes);
   InputNlist extend_lmp_list(extend_inum, &extend_ilist[0], &extend_numneigh[0],
                              &extend_firstneigh[0]);
   extend_lmp_list.set_mask(lmp_list.mask);
@@ -820,22 +861,30 @@ void DeepSpinTF::compute(ENERGYVTYPE& dener,
   dforce_mag_.resize(static_cast<size_t>(nframes) * nall * 3);
   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) {
-      dforce_[3 * ii + dd] = dforce_tmp[3 * new_idx + dd];
-      datom_energy_[ii] = datom_energy_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) {
-        dforce_mag_[3 * ii + dd] = dforce_tmp[3 * (new_idx + nghost) + dd];
-      } else {
-        dforce_mag_[3 * ii + dd] = 0.0;
+  const size_t extended_nall = fwd_map.size();
+  for (int ff = 0; ff < nframes; ++ff) {
+    for (int ii = 0; ii < nall; ++ii) {
+      const int new_idx = new_idx_map[ii];
+      const size_t output_atom = static_cast<size_t>(ff) * nall + ii;
+      const size_t extended_atom =
+          static_cast<size_t>(ff) * extended_nall + new_idx;
+      datom_energy_[output_atom] = datom_energy_tmp[extended_atom];
+      for (int dd = 0; dd < 3; ++dd) {
+        dforce_[output_atom * 3 + dd] = dforce_tmp[extended_atom * 3 + dd];
+
+        if (datype_[ii] < ntypes_spin) {
+          const int virtual_idx = new_idx + (ii < nloc ? nloc : nghost);
+          const size_t virtual_atom =
+              static_cast<size_t>(ff) * extended_nall + virtual_idx;
+          dforce_mag_[output_atom * 3 + dd] = dforce_tmp[virtual_atom * 3 + dd];
+        } else {
+          dforce_mag_[output_atom * 3 + dd] = 0.0;
+        }
+      }
+      for (int dd = 0; dd < 9; ++dd) {
+        datom_virial_[output_atom * 9 + dd] =
+            datom_virial_tmp[extended_atom * 9 + dd];
       }
-    }
-    for (int dd = 0; dd < 9; ++dd) {
-      datom_virial_[ii * 9 + dd] = datom_virial_tmp[new_idx * 9 + dd];
     }
   }
 }
@@ -1016,7 +1065,8 @@ void DeepSpinTF::extend(int& extend_inum,
                         const int nghost,
                         const std::vector<VALUETYPE>& spin,
                         const int numb_types,
-                        const int numb_types_spin) {
+                        const int numb_types_spin,
+                        const int nframes) {
   extend_ilist.clear();
   extend_numneigh.clear();
   extend_neigh.clear();
@@ -1033,8 +1083,10 @@ void DeepSpinTF::extend(int& extend_inum,
     get_vector<float>(spin_norm, "spin_attr/spin_norm");
   }
 
-  int nall = dcoord.size() / 3;
+  int nall = atype.size();
   int nloc = nall - nghost;
+  assert(static_cast<size_t>(nframes) * nall * 3 == dcoord.size());
+  assert(dcoord.size() == spin.size());
   assert(nloc == lmp_list.inum);
 
   // record numb_types_real and nloc_virt
@@ -1140,26 +1192,36 @@ void DeepSpinTF::extend(int& extend_inum,
   }
 
   // extend coord
-  extend_dcoord.resize(static_cast<size_t>(extend_nall) * 3);
-  for (int ii = 0; ii < nloc; ii++) {
-    for (int jj = 0; jj < 3; jj++) {
-      extend_dcoord[new_idx_map[ii] * 3 + jj] = dcoord[ii * 3 + jj];
-      if (atype[ii] < numb_types_spin) {
-        double temp_dcoord = dcoord[ii * 3 + jj] + spin[ii * 3 + jj] /
-                                                       spin_norm[atype[ii]] *
-                                                       virtual_len[atype[ii]];
-        extend_dcoord[(new_idx_map[ii] + nloc) * 3 + jj] = temp_dcoord;
+  extend_dcoord.resize(static_cast<size_t>(nframes) * extend_nall * 3);
+  for (int ff = 0; ff < nframes; ++ff) {
+    const size_t input_offset = static_cast<size_t>(ff) * nall * 3;
+    const size_t output_offset = static_cast<size_t>(ff) * extend_nall * 3;
+    for (int ii = 0; ii < nloc; ii++) {
+      for (int jj = 0; jj < 3; jj++) {
+        extend_dcoord[output_offset + new_idx_map[ii] * 3 + jj] =
+            dcoord[input_offset + ii * 3 + jj];
+        if (atype[ii] < numb_types_spin) {
+          const VALUETYPE temp_dcoord = dcoord[input_offset + ii * 3 + jj] +
+                                        spin[input_offset + ii * 3 + jj] /
+                                            spin_norm[atype[ii]] *
+                                            virtual_len[atype[ii]];
+          extend_dcoord[output_offset + (new_idx_map[ii] + nloc) * 3 + jj] =
+              temp_dcoord;
+        }
       }
     }
-  }
-  for (int ii = nloc; ii < nall; ii++) {
-    for (int jj = 0; jj < 3; jj++) {
-      extend_dcoord[new_idx_map[ii] * 3 + jj] = dcoord[ii * 3 + jj];
-      if (atype[ii] < numb_types_spin) {
-        double temp_dcoord = dcoord[ii * 3 + jj] + spin[ii * 3 + jj] /
-                                                       spin_norm[atype[ii]] *
-                                                       virtual_len[atype[ii]];
-        extend_dcoord[(new_idx_map[ii] + nghost) * 3 + jj] = temp_dcoord;
+    for (int ii = nloc; ii < nall; ii++) {
+      for (int jj = 0; jj < 3; jj++) {
+        extend_dcoord[output_offset + new_idx_map[ii] * 3 + jj] =
+            dcoord[input_offset + ii * 3 + jj];
+        if (atype[ii] < numb_types_spin) {
+          const VALUETYPE temp_dcoord = dcoord[input_offset + ii * 3 + jj] +
+                                        spin[input_offset + ii * 3 + jj] /
+                                            spin_norm[atype[ii]] *
+                                            virtual_len[atype[ii]];
+          extend_dcoord[output_offset + (new_idx_map[ii] + nghost) * 3 + jj] =
+              temp_dcoord;
+        }
       }
     }
   }
@@ -1195,7 +1257,8 @@ template void DeepSpinTF::extend<double>(
     const int nghost,
     const std::vector<double>& spin,
     const int numb_types,
-    const int numb_types_spin);
+    const int numb_types_spin,
+    const int nframes);
 
 template void DeepSpinTF::extend<float>(
     int& extend_inum,
@@ -1214,14 +1277,16 @@ template void DeepSpinTF::extend<float>(
     const int nghost,
     const std::vector<float>& spin,
     const int numb_types,
-    const int numb_types_spin);
+    const int numb_types_spin,
+    const int nframes);
 
 template <typename VALUETYPE>
 void DeepSpinTF::extend_nlist(std::vector<VALUETYPE>& extend_dcoord,
                               std::vector<int>& extend_atype,
                               const std::vector<VALUETYPE>& dcoord_,
                               const std::vector<VALUETYPE>& dspin_,
-                              const std::vector<int>& datype_) {
+                              const std::vector<int>& datype_,
+                              const int nframes) {
   if (dtype == tensorflow::DT_DOUBLE) {
     get_vector<double>(virtual_len, "spin_attr/virtual_len");
     get_vector<double>(spin_norm, "spin_attr/spin_norm");
@@ -1240,20 +1305,27 @@ void DeepSpinTF::extend_nlist(std::vector<VALUETYPE>& extend_dcoord,
     }
   }
   int extend_nall = nloc + nloc_spin;
-  extend_dcoord.resize(static_cast<size_t>(extend_nall) * 3);
+  assert(static_cast<size_t>(nframes) * nloc * 3 == dcoord_.size());
+  assert(dcoord_.size() == dspin_.size());
+  extend_dcoord.resize(static_cast<size_t>(nframes) * extend_nall * 3);
   extend_atype.resize(extend_nall);
   for (int ii = 0; ii < nloc; ii++) {
     extend_atype[ii] = datype_[ii];
     if (datype_[ii] < ntypes_spin) {
       extend_atype[ii + nloc] = datype_[ii] + ntypes - ntypes_spin;
     }
-    for (int jj = 0; jj < 3; jj++) {
-      extend_dcoord[ii * 3 + jj] = dcoord_[ii * 3 + jj];
-      if (datype_[ii] < ntypes_spin) {
-        extend_dcoord[(ii + nloc) * 3 + jj] =
-            dcoord_[ii * 3 + jj] + dspin_[ii * 3 + jj] /
-                                       spin_norm[datype_[ii]] *
-                                       virtual_len[datype_[ii]];
+    for (int ff = 0; ff < nframes; ++ff) {
+      const size_t input_offset = static_cast<size_t>(ff) * nloc * 3;
+      const size_t output_offset = static_cast<size_t>(ff) * extend_nall * 3;
+      for (int jj = 0; jj < 3; jj++) {
+        extend_dcoord[output_offset + ii * 3 + jj] =
+            dcoord_[input_offset + ii * 3 + jj];
+        if (datype_[ii] < ntypes_spin) {
+          extend_dcoord[output_offset + (ii + nloc) * 3 + jj] =
+              dcoord_[input_offset + ii * 3 + jj] +
+              dspin_[input_offset + ii * 3 + jj] / spin_norm[datype_[ii]] *
+                  virtual_len[datype_[ii]];
+        }
       }
     }
   }
@@ -1264,11 +1336,13 @@ template void DeepSpinTF::extend_nlist<double>(
     std::vector<int>& extend_atype,
     const std::vector<double>& dcoord_,
     const std::vector<double>& dspin_,
-    const std::vector<int>& datype_);
+    const std::vector<int>& datype_,
+    const int nframes);
 
 template void DeepSpinTF::extend_nlist<float>(std::vector<float>& extend_dcoord,
                                               std::vector<int>& extend_atype,
                                               const std::vector<float>& dcoord_,
                                               const std::vector<float>& dspin_,
-                                              const std::vector<int>& datype_);
+                                              const std::vector<int>& datype_,
+                                              const int nframes);
 #endif

source/api_cc/src/common.cc

@@ -184,8 +184,14 @@ void deepmd::select_real_atoms_coord(std::vector<VALUETYPE>& dcoord,
   nloc_real = nall_real - nghost_real;
   dcoord.resize(static_cast<size_t>(nframes) * nall_real * 3);
   datype.resize(nall_real);
-  // fwd map
-  select_map<VALUETYPE>(dcoord, dcoord_, fwd_map, 3, nframes, nall_real, nall);
+  // Coordinate buffers can contain an extended atom set while aparam keeps
+  // the caller's original atom stride (for example DeepSpin virtual atoms).
+  // Infer the coordinate stride from its own frame-major buffer instead of
+  // reusing the aparam atom count supplied through ``nall``.
+  const int coord_nall = dcoord_.size() / static_cast<size_t>(nframes) / 3;
+  assert(static_cast<size_t>(nframes) * coord_nall * 3 == dcoord_.size());
+  select_map<VALUETYPE>(dcoord, dcoord_, fwd_map, 3, nframes, nall_real,
+                        coord_nall);
   select_map<int>(datype, datype_, fwd_map, 1);
   // aparam
   if (daparam > 0) {