ideal_mhd_model: share the magnetic-pressure kernel

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

@@ -24,6 +24,7 @@
 #include "vmecpp/vmec/ideal_mhd_model/bcontra_kernel.h"
 #include "vmecpp/vmec/ideal_mhd_model/jacobian_kernel.h"
 #include "vmecpp/vmec/ideal_mhd_model/metric_kernel.h"
+#include "vmecpp/vmec/ideal_mhd_model/pressure_kernel.h"
 #include "vmecpp/vmec/radial_partitioning/radial_partitioning.h"
 #include "vmecpp/vmec/radial_profiles/radial_profiles.h"
 #include "vmecpp/vmec/vmec_constants/vmec_algorithm_constants.h"
@@ -1492,7 +1493,9 @@ void IdealMhdModel::pressureAndEnergies() {
   // Compute as a vectorized operation over all half-grid points
   // temporarily re-use `totalPressure` to store only magnetic pressure; kinetic
   // pressure presH will be added below
-  totalPressure = 0.5 * (bsupu.cwiseProduct(bsubu) + bsupv.cwiseProduct(bsubv));
+  ComputeMagneticPressure(bsupu.data(), bsubu.data(), bsupv.data(),
+                          bsubv.data(), static_cast<int>(bsupu.size()),
+                          totalPressure.data());
 
   // Accumulate magnetic energy and add kinetic pressure per surface
   double localMagneticEnergy = 0.0;

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

@@ -0,0 +1,25 @@
+// SPDX-FileCopyrightText: 2024-present Proxima Fusion GmbH
+// <info@proximafusion.com>
+//
+// SPDX-License-Identifier: MIT
+#ifndef VMECPP_VMEC_IDEAL_MHD_MODEL_PRESSURE_KERNEL_H_
+#define VMECPP_VMEC_IDEAL_MHD_MODEL_PRESSURE_KERNEL_H_
+
+namespace vmecpp {
+
+// Magnetic pressure |B|^2/2 = 0.5 (B^u B_u + B^v B_v) at each half-grid point.
+// This is the field-dependent (nonlinear) part of pressureAndEnergies; the
+// kinetic pressure profile and the energy volume integrals stay in the solver.
+// Shared, allocation-free over flat buffers (n = number of half-grid points),
+// between IdealMhdModel::pressureAndEnergies and the Enzyme autodiff path.
+inline void ComputeMagneticPressure(const double* bsupu, const double* bsubu,
+                                    const double* bsupv, const double* bsubv,
+                                    int n, double* total_pressure) {
+  for (int i = 0; i < n; ++i) {
+    total_pressure[i] = 0.5 * (bsupu[i] * bsubu[i] + bsupv[i] * bsubv[i]);
+  }
+}
+
+}  // namespace vmecpp
+
+#endif  // VMECPP_VMEC_IDEAL_MHD_MODEL_PRESSURE_KERNEL_H_