exact profile evaluation

Author: jurasic-pf

README.md

@@ -318,9 +318,6 @@ VMEC++:
    * `cubic_spline_ip`
    * `pedestal`
    * `rational`
-   * `line_segment`
-   * `line_segment_i`
-   * `line_segment_ip`
    * `nice_quadratic`
    * `sum_cossq_s`
    * `sum_cossq_sqrts`

src/vmecpp/__init__.py

@@ -47,11 +47,11 @@
 
 AuxFType = typing.Annotated[
     _ArrayType,
-    pydantic.BeforeValidator(lambda x: _util.pad_to_target(x, ndfmax, 0.0)),
+    pydantic.BeforeValidator(lambda x: _util.right_pad(x, ndfmax, 0.0)),
 ]
 AuxSType = typing.Annotated[
     _ArrayType,
-    pydantic.BeforeValidator(lambda x: _util.pad_to_target(x, ndfmax, -1.0)),
+    pydantic.BeforeValidator(lambda x: _util.right_pad(x, ndfmax, -1.0)),
 ]
 
 MgridModeType: typing.TypeAlias = typing.Annotated[
@@ -770,7 +770,8 @@ class VmecWOut(BaseModelWithNumpy):
     """Radial derivative of enclosed toroidal magnetic flux ``phi'`` on the full-
     grid."""
 
-    chipf: jt.Float[np.ndarray, "n_surfaces"]
+    # Defaulted for backwards compatibility with old wout files
+    chipf: jt.Float[np.ndarray, "n_surfaces"] = np.array([])
     """Radial derivative of enclosed poloidal magnetic flux ``chi'`` on the full-
     grid."""
 
@@ -1095,11 +1096,13 @@ class VmecWOut(BaseModelWithNumpy):
     """Volume-averaged magnetic field strength."""
 
     # In the C++ WOutFileContents this is called safety_factor.
-    q_factor: jt.Float[np.ndarray, "n_surfaces"]
+    # Defaulted for backwards compatibility with old wout files.
+    q_factor: jt.Float[np.ndarray, "n_surfaces"] = np.array([])
     r"""Safety factor :math:`q = 1/\iota` on the full-grid."""
 
     # In the C++ WOutFileContents this is called poloidal_flux.
-    chi: jt.Float[np.ndarray, "n_surfaces"]
+    # Defaulted for backwards compatibility with old wout files.
+    chi: jt.Float[np.ndarray, "n_surfaces"] = np.array([])
     r"""Enclosed poloidal magnetic flux :math:`\chi` on the full-grid."""
 
     # In the C++ WOutFileContents this is called spectral_width.
@@ -1417,45 +1420,15 @@ def _from_cpp_wout(cpp_wout: _vmecpp.VmecppWOut) -> VmecWOut:
         attrs["gmnc"] = _pad_and_transpose(cpp_wout.gmnc, attrs["mnmax_nyq"])
 
         # These attributes have zero-padding at the end up to a fixed length
-        attrs["am"] = np.pad(cpp_wout.am, (0, preset - len(cpp_wout.am)))
-        attrs["ac"] = np.pad(cpp_wout.ac, (0, preset - len(cpp_wout.ac)))
-        attrs["ai"] = np.pad(cpp_wout.ai, (0, preset - len(cpp_wout.ai)))
-        attrs["am_aux_s"] = np.pad(
-            cpp_wout.am_aux_s,
-            (0, ndfmax - len(cpp_wout.am_aux_s)),
-            mode="constant",
-            constant_values=-1.0,
-        )
-        attrs["am_aux_f"] = np.pad(
-            cpp_wout.am_aux_f,
-            (0, ndfmax - len(cpp_wout.am_aux_f)),
-            mode="constant",
-            constant_values=0.0,
-        )
-        attrs["ac_aux_s"] = np.pad(
-            cpp_wout.ac_aux_s,
-            (0, ndfmax - len(cpp_wout.ac_aux_s)),
-            mode="constant",
-            constant_values=-1.0,
-        )
-        attrs["ac_aux_f"] = np.pad(
-            cpp_wout.ac_aux_f,
-            (0, ndfmax - len(cpp_wout.ac_aux_f)),
-            mode="constant",
-            constant_values=0.0,
-        )
-        attrs["ai_aux_s"] = np.pad(
-            cpp_wout.ai_aux_s,
-            (0, ndfmax - len(cpp_wout.ai_aux_s)),
-            mode="constant",
-            constant_values=-1.0,
-        )
-        attrs["ai_aux_f"] = np.pad(
-            cpp_wout.ai_aux_f,
-            (0, ndfmax - len(cpp_wout.ai_aux_f)),
-            mode="constant",
-            constant_values=0.0,
-        )
+        attrs["am"] = _util.right_pad(cpp_wout.am, preset)
+        attrs["ac"] = _util.right_pad(cpp_wout.ac, preset)
+        attrs["ai"] = _util.right_pad(cpp_wout.ai, preset)
+        attrs["am_aux_s"] = _util.right_pad(cpp_wout.am_aux_s, ndfmax, -1.0)
+        attrs["am_aux_f"] = _util.right_pad(cpp_wout.am_aux_f, ndfmax)
+        attrs["ac_aux_s"] = _util.right_pad(cpp_wout.ac_aux_s, ndfmax, -1.0)
+        attrs["ac_aux_f"] = _util.right_pad(cpp_wout.ac_aux_f, ndfmax)
+        attrs["ai_aux_s"] = _util.right_pad(cpp_wout.ai_aux_s, ndfmax, -1.0)
+        attrs["ai_aux_f"] = _util.right_pad(cpp_wout.ai_aux_f, ndfmax)
 
         attrs["restart_reason_timetrace"] = cpp_wout.restart_reasons
 
@@ -2020,6 +1993,46 @@ def _pad_and_transpose(
     return stacked
 
 
+def populate_raw_profile(
+    vmec_input: VmecInput,
+    field: typing.Literal["pressure", "iota", "current"],
+    f: typing.Callable[[np.ndarray], np.ndarray],
+) -> None:
+    """Populate a line segment profile using callable ``f``.
+
+    The callable is evaluated on all unique ``s`` values required for the
+    multi-grid steps (full and half grids). The resulting knots and values are
+    stored in the auxiliary arrays for the chosen profile.
+    """
+    s_values: set[float] = set()
+    for ns in vmec_input.ns_array:
+        full_grid = np.linspace(0.0, 1.0, ns)
+        half_grid = full_grid - 0.5 * (full_grid[1] - full_grid[0])
+        s_values.update(full_grid)
+        s_values.update(half_grid)
+    knots = np.array(np.sort(np.array(list(s_values))))
+    values = np.array(f(knots))
+
+    if field == "pressure":
+        vmec_input.pmass_type = "line_segment"
+        vmec_input.am_aux_s = knots
+        vmec_input.am_aux_f = values
+        vmec_input.am = np.array([])
+    elif field == "iota":
+        vmec_input.piota_type = "line_segment"
+        vmec_input.ai_aux_s = knots
+        vmec_input.ai_aux_f = values
+        vmec_input.ai = np.array([])
+    elif field == "current":
+        vmec_input.pcurr_type = "line_segment_i"
+        vmec_input.ac_aux_s = knots
+        vmec_input.ac_aux_f = values
+        vmec_input.ac = np.array([])
+    else:
+        msg = "field must be one of 'pressure', 'iota', 'current'"
+        raise ValueError(msg)
+
+
 # Ordered this way to ensure run, VmecInput, and VmecOutput are the first three
 # items in the generated documentation.
 __all__ = [
@@ -2032,4 +2045,5 @@ def _pad_and_transpose(
     "Threed1Volumetrics",
     "MakegridParameters",
     "MagneticFieldResponseTable",
+    "populate_raw_profile",
 ]

src/vmecpp/_util.py

@@ -414,7 +414,9 @@ def sparse_to_dense_coefficients_implicit(
     return sparse_to_dense_coefficients(maybe_sparse_list, mpol, ntor)
 
 
-def pad_to_target(value, target_length: int, default_value: float):
+def pad_to_target(
+    value: np.ndarray, target_length: int, default_value: float
+) -> np.ndarray:
     if len(value) <= target_length:
         return np.pad(
             value,
@@ -427,3 +429,13 @@ def pad_to_target(value, target_length: int, default_value: float):
         f"length {target_length} allowed for serialization"
     )
     raise ValueError(msg)
+
+
+def right_pad(
+    arr: np.ndarray, target_length: int, default_value: float = 0.0
+) -> np.ndarray:
+    """Right-pad an array with zeros to a given length.
+
+    If the array is longer than the target length, leave it unchanged.
+    """
+    return pad_to_target(arr, max(len(arr), target_length), default_value)

src/vmecpp/cpp/vmecpp/vmec/radial_profiles/radial_profiles.cc

@@ -653,24 +653,74 @@ double RadialProfiles::evalRational(const std::vector<double>& coeffs,
   return 0.0;
 }
 
+// Linear interpolation between closest points and associated knots.
+// Clamp the profile if x is outside the range of knots.
 double RadialProfiles::evalLineSegment(const std::vector<double>& splineKnots,
                                        const std::vector<double>& splineValues,
                                        double x) {
-  // TODO(jons): implement `line_segment`
-  (void)splineKnots;
-  (void)splineValues;
-  (void)x;
-  return 0.0;
+  const int n = static_cast<int>(splineKnots.size());
+  if (n < 2 || n != static_cast<int>(splineValues.size())) {
+    return 0.0;
+  }
+
+  auto it = std::lower_bound(splineKnots.begin(), splineKnots.end(), x);
+  if (it == splineKnots.end()) {
+    return splineKnots.back();  // x is out of bounds
+  }
+  if (it == splineKnots.begin()) {
+    return splineValues[0];  // x is below the first knot
+  }
+  const double x0 = *it;
+  const double x1 = *(it + 1);
+  int ilow = static_cast<int>(std::distance(splineKnots.begin(), it));
+  const double y0 = splineValues[ilow];
+  const double y1 = splineValues[ilow + 1];
+  const double t = (x - x0) / (x1 - x0);
+  return (1.0 - t) * y0 + t * y1;
 }
 
 double RadialProfiles::evalLineSegmentIntegrated(
     const std::vector<double>& splineKnots,
     const std::vector<double>& splineValues, double x) {
-  // TODO(jons): implement `line_segment_i`
-  (void)splineKnots;
-  (void)splineValues;
-  (void)x;
-  return 0.0;
+  const int n = static_cast<int>(splineKnots.size());
+  if (n < 2 || n != static_cast<int>(splineValues.size())) {
+    return 0.0;
+  }
+
+  auto integrate_segment = [](double x0, double x1, double y0, double y1) {
+    const double m = (y1 - y0) / (x1 - x0);
+    const double b = y0 - m * x0;
+    return m * 0.5 * (x1 * x1 - x0 * x0) + b * (x1 - x0);
+  };
+
+  double xi = x;
+  double result = 0.0;
+
+  if (xi <= splineKnots.front()) {
+    result += integrate_segment(0.0, xi, splineValues[0],
+                                evalLineSegment(splineKnots, splineValues, xi));
+    return result;
+  }
+
+  int idx = 0;
+  while (idx < n - 1 && xi > splineKnots[idx + 1]) {
+    result += integrate_segment(splineKnots[idx], splineKnots[idx + 1],
+                                splineValues[idx], splineValues[idx + 1]);
+    ++idx;
+  }
+
+  const double x0 = splineKnots[idx];
+  const double x1 = std::min(xi, splineKnots[idx + 1]);
+  const double y0 = splineValues[idx];
+  const double y1 = splineValues[idx + 1];
+  result += integrate_segment(x0, x1, y0, y1);
+
+  if (xi > splineKnots.back()) {
+    result += integrate_segment(splineKnots.back(), xi, splineValues[n - 1],
+                                evalLineSegment(splineKnots, splineValues, xi));
+  }
+
+  return result;
 }
 
 double RadialProfiles::evalNiceQuadratic(const std::vector<double>& coeffs,

tests/test_init.py

@@ -582,6 +582,28 @@ def assert_aux_defaults(wout: vmecpp.VmecWOut):
     np.testing.assert_almost_equal(wout.am_aux_f[:2], np.array([2.0, 3.0]))
 
 
+def test_populate_raw_profile_knots():
+    vmec_input = vmecpp.VmecInput.default()
+    vmec_input.ns_array = np.array([5, 9])
+
+    def f(s):
+        return s**2
+
+    vmecpp.populate_raw_profile(vmec_input, "pressure", f)
+
+    s_values = set()
+    for ns in vmec_input.ns_array:
+        delta = 1.0 / float(ns - 1)
+        s_values.update(i * delta for i in range(ns))
+        s_values.update((i - 0.5) * delta for i in range(ns))
+    expected_knots = np.array(sorted(s_values), dtype=float)
+
+    n = len(expected_knots)
+    np.testing.assert_allclose(vmec_input.am_aux_s[:n], expected_knots)
+    np.testing.assert_allclose(vmec_input.am_aux_f[:n], expected_knots**2)
+    assert vmec_input.pmass_type == "line_segment"
+
+
 def test_default_preset():
     # Default construction doesn't throw an exception
     default_preset = vmecpp.VmecInput.default()