Measure the discretization by the method of manufactured solutions

pyproject.toml

@@ -46,6 +46,9 @@ vmecpp = "vmecpp.__main__:main"
 [project.optional-dependencies]
 test = [
   "pytest",
+  # tests/test_manufactured_solution.py differentiates the ideal-MHD energy
+  # density symbolically; the solver itself does not use it
+  "sympy",
 ]
 benchmark = [
   "pytest",

src/vmecpp/__init__.py

@@ -449,6 +449,13 @@ class VmecInput(BaseModelWithNumpy):
     lforbal: bool = False
     """Hack: directly compute innermost flux surface geometry from radial force balance"""
 
+    enable_force_source: bool = False
+    """Permit an additive spectral force source (``VmecModel.set_force_source``).
+
+    A run that carries one solves a modified problem rather than ideal MHD, so
+    installing a source is refused unless this is set.
+    """
+
     return_outputs_even_if_not_converged: bool = False
     """If true, return a wout even if VMEC++ did not converge, instead of raising a
     RuntimeError.

src/vmecpp/cpp/vmecpp/common/vmec_indata/vmec_indata.cc

@@ -228,6 +228,7 @@ VmecINDATA::VmecINDATA() {
   delt = 1.0;
   tcon0 = 1.0;
   lforbal = false;
+  enable_force_source = false;
   iteration_style = IterationStyle::VMEC_8_52;
   return_outputs_even_if_not_converged = false;
 
@@ -354,6 +355,7 @@ absl::Status VmecINDATA::WriteTo(H5::H5File& file) const {
   WriteH5Dataset(delt, "/indata/delt", file);
   WriteH5Dataset(tcon0, "/indata/tcon0", file);
   WriteH5Dataset(lforbal, "/indata/lforbal", file);
+  WriteH5Dataset(enable_force_source, "/indata/enable_force_source", file);
   WriteH5Dataset(return_outputs_even_if_not_converged,
                  "/indata/return_outputs_even_if_not_converged", file);
 
@@ -459,6 +461,10 @@ absl::Status VmecINDATA::LoadInto(VmecINDATA& m_indata, H5::H5File& from_file) {
   ReadH5Dataset(m_indata.delt, "/indata/delt", from_file);
   ReadH5Dataset(m_indata.tcon0, "/indata/tcon0", from_file);
   ReadH5Dataset(m_indata.lforbal, "/indata/lforbal", from_file);
+  if (from_file.nameExists("/indata/enable_force_source")) {
+    ReadH5Dataset(m_indata.enable_force_source, "/indata/enable_force_source",
+                  from_file);
+  }
 
   // Legacy way of checking for dataset existence
   if (H5Lexists(from_file.getId(),
@@ -943,6 +949,14 @@ absl::StatusOr<VmecINDATA> VmecINDATA::FromJson(
     vmec_indata.lforbal = maybe_lforbal->value();
   }
 
+  auto maybe_enable_force_source = JsonReadBool(j, "enable_force_source");
+  if (!maybe_enable_force_source.ok()) {
+    return maybe_enable_force_source.status();
+  }
+  if (maybe_enable_force_source->has_value()) {
+    vmec_indata.enable_force_source = maybe_enable_force_source->value();
+  }
+
   auto maybe_iteration_style = JsonReadString(j, "iteration_style");
   if (!maybe_iteration_style.ok()) {
     return maybe_iteration_style.status();
@@ -1267,6 +1281,7 @@ absl::StatusOr<std::string> VmecINDATA::ToJson() const {
   output["delt"] = delt;
   output["tcon0"] = tcon0;
   output["lforbal"] = lforbal;
+  output["enable_force_source"] = enable_force_source;
   output["iteration_style"] = ToString(iteration_style);
   output["return_outputs_even_if_not_converged"] =
       return_outputs_even_if_not_converged;

src/vmecpp/cpp/vmecpp/common/vmec_indata/vmec_indata.h

@@ -219,6 +219,11 @@ class VmecINDATA {
   // balance
   bool lforbal;
 
+  // Permit an additive spectral force source (VmecModel.set_force_source). A
+  // run that carries one solves a modified problem rather than ideal MHD, so
+  // installing one is refused unless the input asks for it.
+  bool enable_force_source;
+
   // allows to switch between VMEC 8.52 and PARVMEC iteration style
   // default: VMEC 8.52 (Golden Reference for V&V, and what educational_VMEC is
   // based on)

src/vmecpp/cpp/vmecpp/vmec/fourier_forces/fourier_forces.cc

@@ -158,4 +158,38 @@ void FourierForces::residuals(Eigen::Vector3d& fRes,
   fRes[2] = local_fResL;
 }
 
+std::vector<std::span<double>> FourierForces::ActiveSpans() {
+  std::vector<std::span<double>> out = {frcc};
+  if (s_.lthreed) {
+    out.push_back(frss);
+  }
+  if (s_.lasym) {
+    out.push_back(frsc);
+  }
+  if (s_.lasym && s_.lthreed) {
+    out.push_back(frcs);
+  }
+  out.push_back(fzsc);
+  if (s_.lthreed) {
+    out.push_back(fzcs);
+  }
+  if (s_.lasym) {
+    out.push_back(fzcc);
+  }
+  if (s_.lasym && s_.lthreed) {
+    out.push_back(fzss);
+  }
+  out.push_back(flsc);
+  if (s_.lthreed) {
+    out.push_back(flcs);
+  }
+  if (s_.lasym) {
+    out.push_back(flcc);
+  }
+  if (s_.lasym && s_.lthreed) {
+    out.push_back(flss);
+  }
+  return out;
+}
+
 }  // namespace vmecpp

src/vmecpp/cpp/vmecpp/vmec/fourier_forces/fourier_forces.h

@@ -7,6 +7,7 @@
 
 #include <Eigen/Dense>
 #include <span>
+#include <vector>
 
 #include "vmecpp/vmec/fourier_coefficients/fourier_coefficients.h"
 
@@ -23,6 +24,12 @@ class FourierForces : public FourierCoeffs {
   void zeroZForceForM1();
   void residuals(Eigen::Vector3d& fRes, bool includeEdgeRZ) const;
 
+  // The parities this run actually carries, in a fixed order: R, then Z, then
+  // lambda, each cc/ss/sc/cs as the symmetry and dimensionality allow. Code
+  // that treats the whole force as one flat object, such as
+  // IdealMhdModel::SetForceSource, uses this so the layout has one definition.
+  std::vector<std::span<double>> ActiveSpans();
+
   // appropriately-named variables for the data in FourierCoeffs
   std::span<double> frcc;
   std::span<double> frss;

src/vmecpp/cpp/vmecpp/vmec/ideal_mhd_model/ideal_mhd_model.cc

@@ -927,6 +927,8 @@ absl::StatusOr<bool> IdealMhdModel::update(
 
   m_physical_f.decomposeInto(m_decomposed_f, m_p_.scalxc);
 
+  addForceSource(m_decomposed_f);
+
   // ----- start of residue
 
   // re-establish m=1 constraint
@@ -2067,6 +2069,48 @@ void IdealMhdModel::computeForceNorms(const FourierGeometry& decomposed_x) {
 #endif  // _OPENMP
 }
 
+absl::Status IdealMhdModel::SetForceSource(const Eigen::VectorXd& source) {
+  if (source.size() == 0) {
+    force_source_.resize(0);
+    return absl::OkStatus();
+  }
+  const int mnsize = s_.mpol * (s_.ntor + 1);
+  const int num_parities = s_.num_basis * 3;
+  const Eigen::Index expected =
+      static_cast<Eigen::Index>(num_parities) * m_fc_.ns * mnsize;
+  if (source.size() != expected) {
+    return absl::InvalidArgumentError(absl::StrFormat(
+        "force source has %d entries, but %d are required for %d parities on "
+        "ns=%d with mpol=%d, ntor=%d",
+        source.size(), expected, num_parities, m_fc_.ns, s_.mpol, s_.ntor));
+  }
+  force_source_ = source;
+  return absl::OkStatus();
+}
+
+void IdealMhdModel::addForceSource(FourierForces& m_decomposed_f) const {
+  if (force_source_.size() == 0) {
+    return;
+  }
+  const int mnsize = s_.mpol * (s_.ntor + 1);
+  const int block = m_fc_.ns * mnsize;
+  const std::vector<std::span<double>> targets = m_decomposed_f.ActiveSpans();
+  for (size_t parity = 0; parity < targets.size(); ++parity) {
+    const double* src =
+        force_source_.data() + static_cast<Eigen::Index>(parity) * block;
+    std::span<double> dst = targets[parity];
+    // FourierCoeffs allocates through the LCFS when this thread owns it, which
+    // is one row beyond nsMax(); that row carries the lambda force.
+    const int rows = static_cast<int>(dst.size()) / mnsize;
+    for (int row = 0; row < rows; ++row) {
+      const int global = (m_decomposed_f.nsMin() + row) * mnsize;
+      for (int i = 0; i < mnsize; ++i) {
+        dst[row * mnsize + i] += src[global + i];
+      }
+    }
+  }
+}
+
 void IdealMhdModel::computeMHDForces() {
   int jMaxRZ = std::min(r_.nsMaxF, m_fc_.ns - 1);
   if (m_fc_.lfreeb) {

src/vmecpp/cpp/vmecpp/vmec/ideal_mhd_model/ideal_mhd_model.h

@@ -89,6 +89,24 @@ class IdealMhdModel {
   std::int64_t forceEvaluationCount() const { return force_evaluation_count_; }
   void resetForceEvaluationCount() { force_evaluation_count_ = 0; }
 
+  // Add a fixed spectral force to every force evaluation.
+  //
+  // The source is laid out over the whole radial grid as
+  // ((jF * mpol + m) * (ntor + 1) + n), one block of ns * mpol * (ntor + 1)
+  // entries per active Fourier parity, in the order
+  //   frcc [frss] [frsc] [frcs] fzsc [fzcs] [fzcc] [fzss] flsc [flcs] ...
+  // matching FourierForces. It is added to the decomposed force before the
+  // m = 1 gauge rotation, so it passes through the same chain as the force it
+  // augments and a state whose total force vanishes is a fixed point of the
+  // iteration. An empty vector clears the source.
+  //
+  // This makes the discretization testable by the method of manufactured
+  // solutions: with the source set to the negative of the continuum ideal-MHD
+  // force of a chosen analytic mapping, that mapping is the exact solution of
+  // the modified problem and the distance to the converged discrete state
+  // measures the discretization error.
+  absl::Status SetForceSource(const Eigen::VectorXd& source);
+
   // Coordinates which inverse-DFT routine to call for computing
   // the flux surface geometry and lambda on it from the provided Fourier
   // coefficients. Also computes the net dR/dTheta and dZ/dTheta, without the
@@ -590,6 +608,12 @@ class IdealMhdModel {
   VacuumPressureState& m_vacuum_pressure_state_;
   std::int64_t force_evaluation_count_ = 0;
 
+  // Optional additive spectral force source; see SetForceSource.
+  Eigen::VectorXd force_source_;
+
+  // Add the radial slice of force_source_ owned by this thread.
+  void addForceSource(FourierForces& m_decomposed_f) const;
+
 #ifdef VMECPP_USE_FFTX
   // Pre-computed FFTX kernels for the toroidal (zeta) Fourier transforms.
   // Created once at construction and reused across iterations. Execution is

src/vmecpp/cpp/vmecpp/vmec/pybind11/pybind_vmec.cc

@@ -12,7 +12,9 @@
 #include <Eigen/Dense>
 #include <filesystem>
 #include <optional>
+#include <span>
 #include <string>
+#include <tuple>
 #include <type_traits>  // std::is_same_v
 #include <utility>      // std::move
 
@@ -152,6 +154,17 @@ void UnflattenActive(FourierObject &m_x, const vmecpp::Sizes &s,
   }
 }
 
+inline void CheckFourierShape(const vmecpp::RowMatrixXd &m, const char *name,
+                              int mnmax, int ns) {
+  if (m.rows() != mnmax || m.cols() != ns) {
+    throw std::runtime_error(std::string("VmecModel.set_state_from_fourier: ") +
+                             name + " has shape (" + std::to_string(m.rows()) +
+                             ", " + std::to_string(m.cols()) + "), expected (" +
+                             std::to_string(mnmax) + ", " + std::to_string(ns) +
+                             ")");
+  }
+}
+
 // Single-resolution, single-threaded VMEC++ iteration model.
 //
 // Exposes the VMEC++ forward model (flux-surface geometry -> MHD forces) and
@@ -405,13 +418,53 @@ class VmecModel {
   }
 
   // Flat decision vector (decomposed, i.e. preconditioner-scaled coefficients).
+  void SetForceSource(const Eigen::VectorXd &source) {
+    const absl::Status s = vmec_->SetForceSource(source);
+    if (!s.ok()) {
+      throw std::runtime_error(std::string(s.message()));
+    }
+  }
   Eigen::VectorXd GetState() const {
     return FlattenActive(*vmec_->decomposed_x_[0], vmec_->s_);
   }
   void SetState(const Eigen::VectorXd &flat) const {
     UnflattenActive(*vmec_->decomposed_x_[0], vmec_->s_, flat);
     exact_primal_valid_ = false;  // primal geometry cache is stale
   }
+
+  // Set the state from Fourier coefficients in the combined basis the wout
+  // file uses, R = sum rmnc cos(m u - n v) [+ rmns sin(m u - n v)] and likewise
+  // for Z and lambda, each an [mnmax, ns] array in the standard mode ordering.
+  // The conversion is FourierGeometry::InitFromState, the routine a hot restart
+  // already uses, so the basis normalization, the m = 1 poloidal-origin gauge
+  // and lambda's phip / lamscale scaling have a single implementation. The
+  // asymmetric arrays are ignored for a stellarator-symmetric run and required
+  // for a lasym one.
+  void SetStateFromFourier(const vmecpp::RowMatrixXd &rmnc,
+                           const vmecpp::RowMatrixXd &zmns,
+                           const vmecpp::RowMatrixXd &lmns,
+                           const vmecpp::RowMatrixXd &rmns,
+                           const vmecpp::RowMatrixXd &zmnc,
+                           const vmecpp::RowMatrixXd &lmnc) const {
+    const vmecpp::Sizes &s = vmec_->s_;
+    const int ns = vmec_->fc_.ns;
+    CheckFourierShape(rmnc, "rmnc", s.mnmax, ns);
+    CheckFourierShape(zmns, "zmns", s.mnmax, ns);
+    CheckFourierShape(lmns, "lmns", s.mnmax, ns);
+    if (s.lasym) {
+      CheckFourierShape(rmns, "rmns", s.mnmax, ns);
+      CheckFourierShape(zmnc, "zmnc", s.mnmax, ns);
+      CheckFourierShape(lmnc, "lmnc", s.mnmax, ns);
+    }
+    // A null Boundaries pointer makes InitFromState take the last surface from
+    // the given state rather than from the input boundary, which is what a
+    // free-boundary run needs and what a fixed-boundary one already agrees
+    // with.
+    vmec_->decomposed_x_[0]->InitFromState(
+        vmec_->t_, rmnc, zmns, lmns, rmns, zmnc, lmnc, *vmec_->p_[0],
+        vmec_->constants_, vmec_->indata_.signgs, nullptr);
+  }
+
   // Flat force vector (decomposed/preconditioned), valid after Evaluate().
   Eigen::VectorXd GetForces() const {
     return FlattenActive(*vmec_->decomposed_f_[0], vmec_->s_);
@@ -833,6 +886,7 @@ PYBIND11_MODULE(_vmecpp, m) {
   pyindata.def_readwrite("delt", &VmecINDATA::delt)
       .def_readwrite("tcon0", &VmecINDATA::tcon0)
       .def_readwrite("lforbal", &VmecINDATA::lforbal)
+      .def_readwrite("enable_force_source", &VmecINDATA::enable_force_source)
       .def_readwrite("iteration_style", &VmecINDATA::iteration_style)
       .def_readwrite("return_outputs_even_if_not_converged",
                      &VmecINDATA::return_outputs_even_if_not_converged)
@@ -1551,7 +1605,13 @@ PYBIND11_MODULE(_vmecpp, m) {
       .def("solve", &VmecModel::Solve)
       .def("get_state", &VmecModel::GetState)
       .def("set_state", &VmecModel::SetState, py::arg("state"))
+      .def("set_state_from_fourier", &VmecModel::SetStateFromFourier,
+           py::arg("rmnc"), py::arg("zmns"), py::arg("lmns"),
+           py::arg("rmns") = vmecpp::RowMatrixXd(),
+           py::arg("zmnc") = vmecpp::RowMatrixXd(),
+           py::arg("lmnc") = vmecpp::RowMatrixXd())
       .def("get_forces", &VmecModel::GetForces)
+      .def("set_force_source", &VmecModel::SetForceSource, py::arg("source"))
       .def("get_geometry", &VmecModel::GetGeometry)
       .def("geometry_state_vjp", &VmecModel::GeometryStateVjp,
            py::arg("coefficient_bar"))

src/vmecpp/cpp/vmecpp/vmec/vmec/BUILD.bazel

@@ -57,6 +57,20 @@ cc_test(
     size = "large",
 )
 
+cc_test(
+    name = "force_source_test",
+    srcs = ["force_source_test.cc"],
+    data = [
+        "//vmecpp/test_data:cth_like_fixed_bdy",
+    ],
+    deps = [
+        "//util/file_io:file_io",
+        "//vmecpp/vmec/vmec",
+        "@googletest//:gtest_main",
+    ],
+    size = "medium",
+)
+
 cc_test(
     name = "vmec_allocation_test",
     srcs = ["vmec_allocation_test.cc"],