diff --git a/include/bout/interpolation_xz.hxx b/include/bout/interpolation_xz.hxx index 1a8be9ef6b..5c741a7d3a 100644 --- a/include/bout/interpolation_xz.hxx +++ b/include/bout/interpolation_xz.hxx @@ -29,6 +29,7 @@ #include #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 + 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 { diff --git a/include/bout/petsc_operators.hxx b/include/bout/petsc_operators.hxx index 5a7df5b088..91a983efc7 100644 --- a/include/bout/petsc_operators.hxx +++ b/include/bout/petsc_operators.hxx @@ -24,6 +24,7 @@ #include #include +#include #include #include #include @@ -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 cache_forward; + mutable std::optional cache_backward; /// @brief Read a Field3D from the mesh, throwing if the field is absent. /// /// @param mesh Non-owning mesh pointer. diff --git a/src/mesh/interpolation/hermite_spline_xz.cxx b/src/mesh/interpolation/hermite_spline_xz.cxx index 706fea9c7a..e948c815b8 100644 --- a/src/mesh/interpolation/hermite_spline_xz.cxx +++ b/src/mesh/interpolation/hermite_spline_xz.cxx @@ -160,7 +160,8 @@ XZHermiteSplineBase::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::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::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::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::interpolate( } ASSERT2(std::isfinite(cptr[int(i)])); } - VecRestoreArrayRead(result, &cptr); + BOUT_DO_PETSC(VecRestoreArrayRead(result, &cptr)); #endif } @@ -531,6 +533,36 @@ Field3D XZHermiteSplineBase::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 Field3D XZHermiteSplineBase::interpolate(const Field3D& f, const Field3D& delta_x, diff --git a/src/mesh/interpolation_xz.cxx b/src/mesh/interpolation_xz.cxx index 04d20769e4..bbc859361e 100644 --- a/src/mesh/interpolation_xz.cxx +++ b/src/mesh/interpolation_xz.cxx @@ -101,4 +101,5 @@ RegisterXZInterpolation registerinterpmonotonichermitesplinel{"monotonichermitesplinelegacy"}; RegisterXZInterpolation registerinterplagrange4pt{"lagrange4pt"}; RegisterXZInterpolation registerinterpbilinear{"bilinear"}; +RegisterXZInterpolation registerinterpprecomp{"precomputed"}; } // namespace diff --git a/src/mesh/petsc_operators.cxx b/src/mesh/petsc_operators.cxx index b05f013023..4ab9842b3a 100644 --- a/src/mesh/petsc_operators.cxx +++ b/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