FCI with precomputed weights

include/bout/interpolation_xz.hxx

@@ -29,6 +29,7 @@
 #include <array>
 
 #if BOUT_HAS_PETSC
+#include "bout/petsc_operators.hxx"
 #include "bout/petsclib.hxx"
 #endif
 
@@ -314,6 +315,51 @@ public:
                       const BoutMask& mask,
                       const std::string& region = "RGN_NOBNDRY") override;
 };
+class PrecomputedInterpolation : public XZInterpolation {
+#if BOUT_HAS_PETSC
+  PetscOperators op;
+#endif
+public:
+  PrecomputedInterpolation(Mesh* mesh = nullptr, Options* options = nullptr)
+      : PrecomputedInterpolation(0, mesh, options) {}
+  PrecomputedInterpolation(int y_offset = 0, Mesh* mesh = nullptr,
+                           Options* options = nullptr);
+  PrecomputedInterpolation([[maybe_unused]] const BoutMask& mask, int y_offset = 0,
+                           Mesh* mesh = nullptr, Options* options = nullptr)
+      : PrecomputedInterpolation(y_offset, mesh, options) {}
+
+  void calcWeights([[maybe_unused]] const Field3D& delta_x,
+                   [[maybe_unused]] const Field3D& delta_z,
+                   [[maybe_unused]] const std::string& region = "RGN_NOBNDRY") override {}
+  void calcWeights([[maybe_unused]] const Field3D& delta_x,
+                   [[maybe_unused]] const Field3D& delta_z,
+                   [[maybe_unused]] const BoutMask& mask,
+                   [[maybe_unused]] const std::string& region = "RGN_NOBNDRY") override {}
+
+  // Use precalculated weights
+  Field3D interpolate(const Field3D& f,
+                      const std::string& region = "RGN_NOBNDRY") const override;
+  // Calculate weights and interpolate
+  Field3D
+  interpolate([[maybe_unused]] const Field3D& f, [[maybe_unused]] const Field3D& delta_x,
+              [[maybe_unused]] const Field3D& delta_z,
+              [[maybe_unused]] const std::string& region = "RGN_NOBNDRY") override {
+    throw BoutException("Not Implemented!");
+  }
+  Field3D
+  interpolate([[maybe_unused]] const Field3D& f, [[maybe_unused]] const Field3D& delta_x,
+              [[maybe_unused]] const Field3D& delta_z,
+              [[maybe_unused]] const BoutMask& mask,
+              [[maybe_unused]] const std::string& region = "RGN_NOBNDRY") override {
+    throw BoutException("Not Implemented!");
+  }
+  std::vector<ParallelTransform::PositionsAndWeights>
+  getWeightsForYApproximation([[maybe_unused]] int i, [[maybe_unused]] int j,
+                              [[maybe_unused]] int k,
+                              [[maybe_unused]] int yoffset) override {
+    throw BoutException("NotImplemented");
+  }
+};
 
 class XZInterpolationFactory
     : public Factory<XZInterpolation, XZInterpolationFactory, Mesh*, Options*> {

include/bout/petsc_operators.hxx

@@ -24,6 +24,7 @@
 
 #include <functional>
 #include <memory>
+#include <optional>
 #include <string>
 #include <type_traits>
 #include <unordered_map>
@@ -899,6 +900,22 @@ public:
   /// @returns A diagonal PetscBackwardLegOperator.
   PetscBackwardLegOperator diagonal(const PetscBackwardLegVector& v) const;
 
+  /// @brief Compute the forward parallel slice using the SOM weights.
+  ///
+  /// @param f Field whose yup() slice should be (re)computed. Must have
+  ///          parallel slices allocated.
+  /// @returns A Field3D containing the forward-interpolated values, indexed
+  ///          the same way as f.yup() (i.e. at i.yp()).
+  Field3D computeYup(const Field3D& f) const;
+
+  /// @brief Compute the backward parallel slice using the SOM weights.
+  ///
+  /// @param f Field whose ydown() slice should be (re)computed. Must have
+  ///          parallel slices allocated.
+  /// @returns A Field3D containing the forward-interpolated values, indexed
+  ///          the same way as f.yup() (i.e. at i.yp()).
+  Field3D computeYdown(const Field3D& f) const;
+
   /// @brief Collection of parallel differential operators assembled for the
   ///        Support Operator Method (SOM).
   ///
@@ -976,6 +993,8 @@ public:
   Parallel getParallel() const;
 
 private:
+  mutable std::optional<PetscForwardOperator> cache_forward;
+  mutable std::optional<PetscBackwardOperator> cache_backward;
   /// @brief Read a Field3D from the mesh, throwing if the field is absent.
   ///
   /// @param mesh Non-owning mesh pointer.

src/mesh/interpolation/hermite_spline_xz.cxx

@@ -160,7 +160,8 @@ XZHermiteSplineBase<monotonic, imp_type>::XZHermiteSplineBase(int y_offset, Mesh
         &Options::root()["mesh:paralleltransform:xzinterpolation:hermitespline"]);
     const int m = localmesh->LocalNx * localmesh->LocalNy * localmesh->LocalNz;
     const int M = m * localmesh->getNXPE() * localmesh->getNYPE() * localmesh->getNZPE();
-    MatCreateAIJ(BoutComm::get(), m, m, M, M, 16, nullptr, 16, nullptr, &petscWeights);
+    BOUT_DO_PETSC(MatCreateAIJ(BoutComm::get(), m, m, M, M, 16, nullptr, 16, nullptr,
+                               &petscWeights));
 #endif
   }
   if constexpr (monotonic) {
@@ -333,7 +334,8 @@ void XZHermiteSplineBase<monotonic, imp_type>::calcWeights(
                                             k_corner(x, y, z) - 1 + k);
             vals[k] = newWeights[(j * 4) + k][i];
           }
-          MatSetValues(petscWeights, 1, idxn, 4, idxm, vals, INSERT_VALUES);
+          BOUT_DO_PETSC(
+              MatSetValues(petscWeights, 1, idxn, 4, idxm, vals, INSERT_VALUES));
         }
 #endif
       }
@@ -350,10 +352,10 @@ void XZHermiteSplineBase<monotonic, imp_type>::calcWeights(
   }
   if constexpr (imp_type == bout::details::implementation_type::petsc) {
 #if BOUT_HAS_PETSC
-    MatAssemblyBegin(petscWeights, MAT_FINAL_ASSEMBLY);
-    MatAssemblyEnd(petscWeights, MAT_FINAL_ASSEMBLY);
+    BOUT_DO_PETSC(MatAssemblyBegin(petscWeights, MAT_FINAL_ASSEMBLY));
+    BOUT_DO_PETSC(MatAssemblyEnd(petscWeights, MAT_FINAL_ASSEMBLY));
     if (!isInit) {
-      MatCreateVecs(petscWeights, &rhs, &result);
+      BOUT_DO_PETSC(MatCreateVecs(petscWeights, &rhs, &result));
     }
     isInit = true;
 #endif
@@ -423,11 +425,11 @@ Field3D XZHermiteSplineBase<monotonic, imp_type>::interpolate(
 #if BOUT_HAS_PETSC
     BoutReal* ptr = nullptr;
     const BoutReal* cptr = nullptr;
-    VecGetArray(rhs, &ptr);
+    BOUT_DO_PETSC(VecGetArray(rhs, &ptr));
     BOUT_FOR(i, f.getRegion("RGN_NOY")) { ptr[int(i)] = f[i]; }
-    VecRestoreArray(rhs, &ptr);
-    MatMult(petscWeights, rhs, result);
-    VecGetArrayRead(result, &cptr);
+    BOUT_DO_PETSC(VecRestoreArray(rhs, &ptr));
+    BOUT_DO_PETSC(MatMult(petscWeights, rhs, result));
+    BOUT_DO_PETSC(VecGetArrayRead(result, &cptr));
     BOUT_FOR(i, f.getRegion(region2)) {
       f_interp[i] = cptr[int(i)];
       if constexpr (monotonic) {
@@ -443,7 +445,7 @@ Field3D XZHermiteSplineBase<monotonic, imp_type>::interpolate(
       }
       ASSERT2(std::isfinite(cptr[int(i)]));
     }
-    VecRestoreArrayRead(result, &cptr);
+    BOUT_DO_PETSC(VecRestoreArrayRead(result, &cptr));
 #endif
   }
 
@@ -531,6 +533,36 @@ Field3D XZHermiteSplineBase<monotonic, imp_type>::interpolate(
   return f_interp;
 }
 
+#if BOUT_HAS_PETSC
+PrecomputedInterpolation::PrecomputedInterpolation(int y_offset, Mesh* mesh,
+                                                   [[maybe_unused]] Options* options)
+    : XZInterpolation(y_offset, mesh), op(mesh) {}
+#else
+PrecomputedInterpolation::PrecomputedInterpolation([[maybe_unused]] int y_offset,
+                                                   [[maybe_unused]] Mesh* mesh,
+                                                   [[maybe_unused]] Options* options) {
+  throw BoutException("PrecomputedInterpolation requires PETSc");
+}
+#endif
+
+Field3D
+PrecomputedInterpolation::interpolate(const Field3D& fin,
+                                      [[maybe_unused]] const std::string& region) const {
+#if BOUT_HAS_PETSC
+  Field3D f = fin;
+  f.splitParallelSlices();
+  f.yup().allocate();
+  f.ydown().allocate();
+  if (y_offset == 1) {
+    return op.computeYup(f);
+  }
+  if (y_offset == -1) {
+    return op.computeYdown(f);
+  }
+#endif
+  throw BoutException("Offset {} not supported", y_offset);
+}
+
 template <bool monotonic, bout::details::implementation_type imp_type>
 Field3D XZHermiteSplineBase<monotonic, imp_type>::interpolate(const Field3D& f,
                                                               const Field3D& delta_x,

src/mesh/interpolation_xz.cxx

@@ -101,4 +101,5 @@ RegisterXZInterpolation<XZMonotonicHermiteSplineLegacy>
     registerinterpmonotonichermitesplinel{"monotonichermitesplinelegacy"};
 RegisterXZInterpolation<XZLagrange4pt> registerinterplagrange4pt{"lagrange4pt"};
 RegisterXZInterpolation<XZBilinear> registerinterpbilinear{"bilinear"};
+RegisterXZInterpolation<PrecomputedInterpolation> registerinterpprecomp{"precomputed"};
 } // namespace

src/mesh/petsc_operators.cxx

@@ -559,4 +559,73 @@ PetscCellOperator makeNeumannOperator(const PetscCellMappingPtr& mapping,
   return PetscCellOperator(mapping, mapping, std::move(out));
 }
 
+Field3D PetscOperators::computeYup(const Field3D& f) const {
+  ASSERT1(f.hasParallelSlices());
+  ASSERT1(f.yup().isAllocated());
+  ASSERT1(f.ydown().isAllocated());
+
+  if (!cache_forward.has_value()) {
+    cache_forward = this->forward();
+  }
+  auto leg = cache_forward.value()(f);
+
+  const PetscScalar* leg_x{nullptr};
+  BOUT_DO_PETSC(VecGetArrayRead(leg.raw(), &leg_x));
+
+  Field3D result{emptyFrom(f)};
+
+  const PetscInt leg_row_start = forward_leg_mapping->rowStart();
+
+  cell_mapping->mapOwnedInteriorCells([&](PetscInt /*cell_row*/, const Ind3D& i, int) {
+    const int interior = ROUND(forward_leg_interior_number[i]);
+
+    BoutReal value = 0.0;
+    if (interior >= 0) {
+      const PetscInt row = forward_leg_mapping->storedToPetsc(interior) - leg_row_start;
+      value = leg_x[row];
+    } else {
+      value = BoutNaN;
+    }
+    result[i.yp()] = value;
+  });
+
+  BOUT_DO_PETSC(VecRestoreArrayRead(leg.raw(), &leg_x));
+
+  return result;
+}
+
+Field3D PetscOperators::computeYdown(const Field3D& f) const {
+  ASSERT1(f.hasParallelSlices());
+  ASSERT1(f.yup().isAllocated());
+  ASSERT1(f.ydown().isAllocated());
+
+  if (!cache_backward.has_value()) {
+    cache_backward = this->backward();
+  }
+  auto leg = cache_backward.value()(f);
+
+  const PetscScalar* leg_x{nullptr};
+  BOUT_DO_PETSC(VecGetArrayRead(leg.raw(), &leg_x));
+
+  Field3D result{emptyFrom(f)};
+
+  const PetscInt leg_row_start = backward_leg_mapping->rowStart();
+
+  cell_mapping->mapOwnedInteriorCells([&](PetscInt /*cell_row*/, const Ind3D& i, int) {
+    const int interior = ROUND(backward_leg_interior_number[i]);
+
+    BoutReal value = 0.0;
+    if (interior >= 0) {
+      const PetscInt row = backward_leg_mapping->storedToPetsc(interior) - leg_row_start;
+      value = leg_x[row];
+    } else {
+      value = BoutNaN;
+    }
+    result[i.ym()] = value;
+  });
+
+  BOUT_DO_PETSC(VecRestoreArrayRead(leg.raw(), &leg_x));
+
+  return result;
+}
 #endif