From 3fa255d833aa2fdd04755911d9d61466d0b90fe4 Mon Sep 17 00:00:00 2001
From: Ben Wibking <ben@wibking.com>
Date: Sun, 7 Apr 2024 22:09:59 -0400
Subject: [PATCH] add computeVolumeIntegral function (#597)

### Description
Adds a function that can be used to a compute a volume integral of a
user-defined function that takes an Array4 for a box of the
`state_new_cc_` MultiFab as input.

### Related issues
N/A

### Checklist
_Before this pull request can be reviewed, all of these tasks should be
completed. Denote completed tasks with an `x` inside the square brackets
`[ ]` in the Markdown source below:_
- [x] I have added a description (see above).
- [x] I have added a link to any related issues see (see above).
- [x] I have read the [Contributing
Guide](https://github.com/quokka-astro/quokka/blob/development/CONTRIBUTING.md).
- [x] I have added tests for any new physics that this PR adds to the
code.
- [x] I have tested this PR on my local computer and all tests pass.
- [x] I have manually triggered the GPU tests with the magic comment
`/azp run`.
- [ ] I have requested a reviewer for this PR.
---
 src/simulation.hpp | 29 +++++++++++++++++++++++++++++
 1 file changed, 29 insertions(+)

diff --git a/src/simulation.hpp b/src/simulation.hpp
index 10dfd1445..200865f8c 100644
--- a/src/simulation.hpp
+++ b/src/simulation.hpp
@@ -307,6 +307,9 @@ template <typename problem_t> class AMRSimulation : public amrex::AmrCore
 
 	template <typename ReduceOp, typename F> auto computePlaneProjection(F const &user_f, int dir) const -> amrex::BaseFab<amrex::Real>;
 
+	// compute volume integrals
+	template <typename F> auto computeVolumeIntegral(F const &user_f) -> amrex::Real;
+
 	// I/O functions
 	[[nodiscard]] auto PlotFileName(int lev) const -> std::string;
 	[[nodiscard]] auto CustomPlotFileName(const char *base, int lev) const -> std::string;
@@ -1958,6 +1961,32 @@ template <typename problem_t> void AMRSimulation<problem_t>::AverageDownTo(int c
 	}
 }
 
+template <typename problem_t> template <typename F> auto AMRSimulation<problem_t>::computeVolumeIntegral(F const &user_f) -> amrex::Real
+{
+	// compute integral of user_f(i, j, k, state) along the given axis.
+	const BL_PROFILE("AMRSimulation::computeVolumeIntegral()");
+
+	// allocate temporary multifabs
+	amrex::Vector<amrex::MultiFab> q;
+	q.resize(finest_level + 1);
+	for (int lev = 0; lev <= finest_level; ++lev) {
+		q[lev].define(boxArray(lev), DistributionMap(lev), 1, 0);
+	}
+
+	// evaluate user_f on all levels
+	// (note: it is not necessary to average down)
+	for (int lev = 0; lev <= finest_level; ++lev) {
+		auto const &state = state_new_cc_[lev].const_arrays();
+		auto const &result = q[lev].arrays();
+		amrex::ParallelFor(q[lev], [=] AMREX_GPU_DEVICE(int bx, int i, int j, int k) { result[bx](i, j, k) = user_f(i, j, k, state[bx]); });
+	}
+	amrex::Gpu::streamSynchronize();
+
+	// call amrex::volumeWeightedSum
+	const amrex::Real result = amrex::volumeWeightedSum(amrex::GetVecOfConstPtrs(q), 0, geom, ref_ratio);
+	return result;
+}
+
 #ifdef AMREX_PARTICLES
 template <typename problem_t> void AMRSimulation<problem_t>::InitParticles()
 {