fix(api_cc): let a native-spin call name the charge state it wants served

The two artifacts this backend serves carry the condition differently, and
the charge-aware entry point treated both as the folded one: it refused any
state but the installed default and then dropped the argument, leaving
run_graph_payload to build the input from that default.

For a compressed artifact this is the contract. Its condition lives in
frozen tables that set_charge_spin rebuilds and that stand for the whole
run, and its compact lower takes no condition at all. Nothing there
changes.

An uncompressed one keeps the condition in the argument list of its
compiled forward, which was already being filled on every evaluation, just
from the default rather than from the caller. The condition now travels
with the frames it belongs to, divided by the layout the other per-frame
inputs use, so successive calls may each name their own state.

The width of the forward's argument is what tells the two apart, and the
charge-aware overload becomes the implementation the charge-unaware one
delegates to, so the condition is an ordinary input rather than a second
entry point.

Author: OutisLi

source/api_cc/include/NativeSpinPTExpt.h

@@ -347,6 +347,7 @@ class NativeSpinPTExpt : public DeepSpinBackend {
                const int& ago,
                const std::vector<VALUETYPE>& fparam,
                const std::vector<VALUETYPE>& aparam,
+               const std::vector<double>& charge_spin,
                const bool atomic);
   /**
    * @brief Evaluate frames that arrive without a neighbor list.
@@ -369,8 +370,21 @@ class NativeSpinPTExpt : public DeepSpinBackend {
                const std::vector<VALUETYPE>& box,
                const std::vector<VALUETYPE>& fparam,
                const std::vector<VALUETYPE>& aparam,
+               const std::vector<double>& charge_spin,
                const bool atomic);
 
+  /**
+   * @brief Whether a call may name the charge state it wants served.
+   *
+   * The two artifacts this class serves carry the condition differently. An
+   * uncompressed one keeps it in the argument list of its compiled forward,
+   * where a state reaches the evaluation that names it. A compressed one
+   * folds it into frozen tables, which are rebuilt by ``set_charge_spin``
+   * and stand for the whole run, so a call can only restate what they hold.
+   * The width of the forward's argument tells the two apart.
+   **/
+  bool reads_charge_spin_per_call() const { return dchgspin > 0; }
+
   /**
    * @brief Evaluate one frame: build a neighbor list, then fold ghost
    * contributions back onto their local owners.
@@ -388,6 +402,7 @@ class NativeSpinPTExpt : public DeepSpinBackend {
                      const std::vector<VALUETYPE>& box,
                      const std::vector<VALUETYPE>& fparam,
                      const std::vector<VALUETYPE>& aparam,
+                     const std::vector<double>& charge_spin,
                      const bool atomic);
 
   /**
@@ -444,7 +459,8 @@ class NativeSpinPTExpt : public DeepSpinBackend {
       const std::int64_t nloc,
       const torch::Tensor& spin,
       const std::vector<double>& fparam,
-      const std::vector<double>& aparam);
+      const std::vector<double>& aparam,
+      const std::vector<double>& charge_spin);
 
   /**
    * @brief Bind the flat artifact outputs to their metadata key names.

source/api_cc/src/NativeSpinPTExpt.cc

@@ -149,7 +149,10 @@ FrameLayout broadcast_layout(const std::vector<VALUETYPE>& values,
                              const int nframes,
                              const std::size_t stride,
                              const char* name) {
-  if (values.empty() || values.size() == stride) {
+  // A width of zero is a consumer that reads nothing. Whatever the caller
+  // supplied was answered before the frames were reached, so they divide
+  // nothing between them.
+  if (stride == 0 || values.empty() || values.size() == stride) {
     return {stride, 0};
   }
   return per_frame_layout(values, nframes, stride, name);
@@ -593,7 +596,8 @@ std::map<std::string, torch::Tensor> NativeSpinPTExpt::run_graph_payload(
     const std::int64_t nloc,
     const torch::Tensor& spin,
     const std::vector<double>& fparam,
-    const std::vector<double>& aparam) {
+    const std::vector<double>& aparam,
+    const std::vector<double>& charge_spin) {
   graph.edge_mask =
       deepmd::applyPairExclusion(graph.edge_index, graph.edge_mask, graph.atype,
                                  pair_exclude_table_, ntypes);
@@ -607,13 +611,18 @@ std::map<std::string, torch::Tensor> NativeSpinPTExpt::run_graph_payload(
     deepmd::flatten_canonical_atom_virial(output_map);
   } else {
     const torch::Device device = graph.atype.device();
+    // This lower reads the condition as an ordinary input, so a call may name
+    // the state it wants; the state the model was frozen with stands in for a
+    // call that names none.
+    const std::vector<double>& served =
+        charge_spin.empty() ? default_chg_spin_ : charge_spin;
     extract_outputs(
         output_map,
         run_model_graph(
             graph, spin,
             make_fparam_tensor(fparam, default_fparam_, dfparam, device),
             make_aparam_tensor(aparam, daparam, node_count, nloc, device),
-            make_chg_spin_tensor(default_chg_spin_, dchgspin, device)));
+            make_chg_spin_tensor(served, dchgspin, device)));
   }
   return output_map;
 }
@@ -651,6 +660,7 @@ void NativeSpinPTExpt::compute(ENERGYVTYPE& ener,
                                const int& ago,
                                const std::vector<VALUETYPE>& fparam,
                                const std::vector<VALUETYPE>& aparam,
+                               const std::vector<double>& charge_spin,
                                const bool atomic) {
   const torch::Device device = gpu_enabled ? torch::Device(torch::kCUDA, gpu_id)
                                            : torch::Device(torch::kCPU);
@@ -781,7 +791,7 @@ void NativeSpinPTExpt::compute(ENERGYVTYPE& ener,
   std::map<std::string, torch::Tensor> output_map = run_graph_payload(
       graph, node_count, nloc, spin_Tensor.slice(0, 0, node_count),
       std::vector<double>(fparam.begin(), fparam.end()),
-      std::vector<double>(aparam_real.begin(), aparam_real.end()));
+      std::vector<double>(aparam_real.begin(), aparam_real.end()), charge_spin);
 
   // The forward emits flat per-node public keys; rewrite them into the dense
   // internal-key layout the extraction below reads. The extended node set
@@ -868,6 +878,7 @@ template void NativeSpinPTExpt::compute<double, std::vector<ENERGYTYPE>>(
     const int& ago,
     const std::vector<double>& fparam,
     const std::vector<double>& aparam,
+    const std::vector<double>& charge_spin,
     const bool atomic);
 template void NativeSpinPTExpt::compute<float, std::vector<ENERGYTYPE>>(
     std::vector<ENERGYTYPE>& ener,
@@ -884,6 +895,7 @@ template void NativeSpinPTExpt::compute<float, std::vector<ENERGYTYPE>>(
     const int& ago,
     const std::vector<float>& fparam,
     const std::vector<float>& aparam,
+    const std::vector<double>& charge_spin,
     const bool atomic);
 
 // ============================================================================
@@ -903,6 +915,7 @@ void NativeSpinPTExpt::compute(ENERGYVTYPE& ener,
                                const std::vector<VALUETYPE>& box,
                                const std::vector<VALUETYPE>& fparam,
                                const std::vector<VALUETYPE>& aparam,
+                               const std::vector<double>& charge_spin,
                                const bool atomic) {
   const std::size_t nloc = atype.size();
   const int nframes = frame_count(coord, nloc);
@@ -915,6 +928,8 @@ void NativeSpinPTExpt::compute(ENERGYVTYPE& ener,
       fparam, nframes, static_cast<std::size_t>(dfparam), "fparam");
   const FrameLayout aparam_layout = broadcast_layout(
       aparam, nframes, nloc * static_cast<std::size_t>(daparam), "aparam");
+  const FrameLayout charge_spin_layout = broadcast_layout(
+      charge_spin, nframes, static_cast<std::size_t>(dchgspin), "charge_spin");
 
   ener.clear();
   force.clear();
@@ -932,7 +947,8 @@ void NativeSpinPTExpt::compute(ENERGYVTYPE& ener,
                   frame_block(spin, ff, spin_layout), atype,
                   frame_block(box, ff, box_layout),
                   frame_block(fparam, ff, fparam_layout),
-                  frame_block(aparam, ff, aparam_layout), atomic);
+                  frame_block(aparam, ff, aparam_layout),
+                  frame_block(charge_spin, ff, charge_spin_layout), atomic);
     ener.insert(ener.end(), frame_ener.begin(), frame_ener.end());
     force.insert(force.end(), frame_force.begin(), frame_force.end());
     force_mag.insert(force_mag.end(), frame_force_mag.begin(),
@@ -958,6 +974,7 @@ void NativeSpinPTExpt::compute_frame(ENERGYVTYPE& ener,
                                      const std::vector<VALUETYPE>& box,
                                      const std::vector<VALUETYPE>& fparam,
                                      const std::vector<VALUETYPE>& aparam,
+                                     const std::vector<double>& charge_spin,
                                      const bool atomic) {
   const torch::Device device = gpu_enabled ? torch::Device(torch::kCUDA, gpu_id)
                                            : torch::Device(torch::kCPU);
@@ -1036,10 +1053,10 @@ void NativeSpinPTExpt::compute_frame(ENERGYVTYPE& ener,
                        f64_options)
           .clone()
           .to(device);
-  std::map<std::string, torch::Tensor> output_map =
-      run_graph_payload(graph, nloc, nloc, spin_Tensor,
-                        std::vector<double>(fparam.begin(), fparam.end()),
-                        std::vector<double>(aparam.begin(), aparam.end()));
+  std::map<std::string, torch::Tensor> output_map = run_graph_payload(
+      graph, nloc, nloc, spin_Tensor,
+      std::vector<double>(fparam.begin(), fparam.end()),
+      std::vector<double>(aparam.begin(), aparam.end()), charge_spin);
   deepmd::remap_graph_spin_outputs_to_dense_keys(output_map, nloc, nall,
                                                  atomic);
 
@@ -1110,6 +1127,7 @@ template void NativeSpinPTExpt::compute<double, std::vector<ENERGYTYPE>>(
     const std::vector<double>& box,
     const std::vector<double>& fparam,
     const std::vector<double>& aparam,
+    const std::vector<double>& charge_spin,
     const bool atomic);
 template void NativeSpinPTExpt::compute<float, std::vector<ENERGYTYPE>>(
     std::vector<ENERGYTYPE>& ener,
@@ -1124,6 +1142,7 @@ template void NativeSpinPTExpt::compute<float, std::vector<ENERGYTYPE>>(
     const std::vector<float>& box,
     const std::vector<float>& fparam,
     const std::vector<float>& aparam,
+    const std::vector<double>& charge_spin,
     const bool atomic);
 
 // ============================================================================
@@ -1143,11 +1162,8 @@ void NativeSpinPTExpt::computew(std::vector<double>& ener,
                                 const std::vector<double>& fparam,
                                 const std::vector<double>& aparam,
                                 const bool atomic) {
-  translate_error([&] {
-    reject_unsupported_parametric_inputs(fparam, aparam, dfparam, daparam);
-    compute(ener, force, force_mag, virial, atom_energy, atom_virial, coord,
-            spin, atype, box, fparam, aparam, atomic);
-  });
+  computew(ener, force, force_mag, virial, atom_energy, atom_virial, coord,
+           spin, atype, box, fparam, aparam, std::vector<double>(), atomic);
 }
 
 void NativeSpinPTExpt::computew(std::vector<double>& ener,
@@ -1163,11 +1179,8 @@ void NativeSpinPTExpt::computew(std::vector<double>& ener,
                                 const std::vector<float>& fparam,
                                 const std::vector<float>& aparam,
                                 const bool atomic) {
-  translate_error([&] {
-    reject_unsupported_parametric_inputs(fparam, aparam, dfparam, daparam);
-    compute(ener, force, force_mag, virial, atom_energy, atom_virial, coord,
-            spin, atype, box, fparam, aparam, atomic);
-  });
+  computew(ener, force, force_mag, virial, atom_energy, atom_virial, coord,
+           spin, atype, box, fparam, aparam, std::vector<double>(), atomic);
 }
 
 void NativeSpinPTExpt::computew(std::vector<double>& ener,
@@ -1186,11 +1199,9 @@ void NativeSpinPTExpt::computew(std::vector<double>& ener,
                                 const std::vector<double>& fparam,
                                 const std::vector<double>& aparam,
                                 const bool atomic) {
-  translate_error([&] {
-    reject_unsupported_parametric_inputs(fparam, aparam, dfparam, daparam);
-    compute(ener, force, force_mag, virial, atom_energy, atom_virial, coord,
-            spin, atype, nghost, inlist, ago, fparam, aparam, atomic);
-  });
+  computew(ener, force, force_mag, virial, atom_energy, atom_virial, coord,
+           spin, atype, box, nghost, inlist, ago, fparam, aparam,
+           std::vector<double>(), atomic);
 }
 
 void NativeSpinPTExpt::computew(std::vector<double>& ener,
@@ -1209,11 +1220,9 @@ void NativeSpinPTExpt::computew(std::vector<double>& ener,
                                 const std::vector<float>& fparam,
                                 const std::vector<float>& aparam,
                                 const bool atomic) {
-  translate_error([&] {
-    reject_unsupported_parametric_inputs(fparam, aparam, dfparam, daparam);
-    compute(ener, force, force_mag, virial, atom_energy, atom_virial, coord,
-            spin, atype, nghost, inlist, ago, fparam, aparam, atomic);
-  });
+  computew(ener, force, force_mag, virial, atom_energy, atom_virial, coord,
+           spin, atype, box, nghost, inlist, ago, fparam, aparam,
+           std::vector<double>(), atomic);
 }
 
 void NativeSpinPTExpt::computew(std::vector<double>& ener,
@@ -1230,11 +1239,14 @@ void NativeSpinPTExpt::computew(std::vector<double>& ener,
                                 const std::vector<double>& aparam,
                                 const std::vector<double>& charge_spin,
                                 const bool atomic) {
-  check_call_charge_spin(
-      charge_spin, frame_count(coord, atype.size()), settable_chgspin,
-      /*applied_per_call=*/false, default_chg_spin_, chg_spin_table_ranges_);
-  computew(ener, force, force_mag, virial, atom_energy, atom_virial, coord,
-           spin, atype, box, fparam, aparam, atomic);
+  translate_error([&] {
+    reject_unsupported_parametric_inputs(fparam, aparam, dfparam, daparam);
+    check_call_charge_spin(charge_spin, frame_count(coord, atype.size()),
+                           settable_chgspin, reads_charge_spin_per_call(),
+                           default_chg_spin_, chg_spin_table_ranges_);
+    compute(ener, force, force_mag, virial, atom_energy, atom_virial, coord,
+            spin, atype, box, fparam, aparam, charge_spin, atomic);
+  });
 }
 
 void NativeSpinPTExpt::computew(std::vector<double>& ener,
@@ -1251,11 +1263,14 @@ void NativeSpinPTExpt::computew(std::vector<double>& ener,
                                 const std::vector<float>& aparam,
                                 const std::vector<double>& charge_spin,
                                 const bool atomic) {
-  check_call_charge_spin(
-      charge_spin, frame_count(coord, atype.size()), settable_chgspin,
-      /*applied_per_call=*/false, default_chg_spin_, chg_spin_table_ranges_);
-  computew(ener, force, force_mag, virial, atom_energy, atom_virial, coord,
-           spin, atype, box, fparam, aparam, atomic);
+  translate_error([&] {
+    reject_unsupported_parametric_inputs(fparam, aparam, dfparam, daparam);
+    check_call_charge_spin(charge_spin, frame_count(coord, atype.size()),
+                           settable_chgspin, reads_charge_spin_per_call(),
+                           default_chg_spin_, chg_spin_table_ranges_);
+    compute(ener, force, force_mag, virial, atom_energy, atom_virial, coord,
+            spin, atype, box, fparam, aparam, charge_spin, atomic);
+  });
 }
 
 void NativeSpinPTExpt::computew(std::vector<double>& ener,
@@ -1275,11 +1290,15 @@ void NativeSpinPTExpt::computew(std::vector<double>& ener,
                                 const std::vector<double>& aparam,
                                 const std::vector<double>& charge_spin,
                                 const bool atomic) {
-  check_call_charge_spin(charge_spin, 1, settable_chgspin,
-                         /*applied_per_call=*/false, default_chg_spin_,
-                         chg_spin_table_ranges_);
-  computew(ener, force, force_mag, virial, atom_energy, atom_virial, coord,
-           spin, atype, box, nghost, inlist, ago, fparam, aparam, atomic);
+  translate_error([&] {
+    reject_unsupported_parametric_inputs(fparam, aparam, dfparam, daparam);
+    check_call_charge_spin(charge_spin, 1, settable_chgspin,
+                           reads_charge_spin_per_call(), default_chg_spin_,
+                           chg_spin_table_ranges_);
+    compute(ener, force, force_mag, virial, atom_energy, atom_virial, coord,
+            spin, atype, nghost, inlist, ago, fparam, aparam, charge_spin,
+            atomic);
+  });
 }
 
 void NativeSpinPTExpt::computew(std::vector<double>& ener,
@@ -1299,11 +1318,15 @@ void NativeSpinPTExpt::computew(std::vector<double>& ener,
                                 const std::vector<float>& aparam,
                                 const std::vector<double>& charge_spin,
                                 const bool atomic) {
-  check_call_charge_spin(charge_spin, 1, settable_chgspin,
-                         /*applied_per_call=*/false, default_chg_spin_,
-                         chg_spin_table_ranges_);
-  computew(ener, force, force_mag, virial, atom_energy, atom_virial, coord,
-           spin, atype, box, nghost, inlist, ago, fparam, aparam, atomic);
+  translate_error([&] {
+    reject_unsupported_parametric_inputs(fparam, aparam, dfparam, daparam);
+    check_call_charge_spin(charge_spin, 1, settable_chgspin,
+                           reads_charge_spin_per_call(), default_chg_spin_,
+                           chg_spin_table_ranges_);
+    compute(ener, force, force_mag, virial, atom_energy, atom_virial, coord,
+            spin, atype, nghost, inlist, ago, fparam, aparam, charge_spin,
+            atomic);
+  });
 }
 
 void NativeSpinPTExpt::compute_canonical_graph_gpu(