From 26bb807a7576b07e601bc7b52613ff16185c68cd Mon Sep 17 00:00:00 2001
From: Piyush Sharda <34922596+psharda@users.noreply.github.com>
Date: Sun, 14 Jan 2024 22:09:23 +0100
Subject: [PATCH] Save gravitational potential gpot as a derived var in
 pltfiles (#406)

Will fix #357 (when it actually works).

Testing it out on the SphericalCollapse test. The code compiles but
gives a segfault after writing chk00000

I'm sure there are bugs in what I did.

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Ben Wibking <ben@wibking.com>
Co-authored-by: Piyush Sharda <psharda@RSAA-043527.local>
---
 src/SphericalCollapse/spherical_collapse.cpp | 14 +++++++++
 src/simulation.hpp                           | 31 ++++++++++++++++++--
 tests/SphericalCollapse.in                   |  2 ++
 3 files changed, 45 insertions(+), 2 deletions(-)

diff --git a/src/SphericalCollapse/spherical_collapse.cpp b/src/SphericalCollapse/spherical_collapse.cpp
index 90c013e85..1aeaf6967 100644
--- a/src/SphericalCollapse/spherical_collapse.cpp
+++ b/src/SphericalCollapse/spherical_collapse.cpp
@@ -105,6 +105,20 @@ template <> void RadhydroSimulation<CollapseProblem>::ErrorEst(int lev, amrex::T
 	}
 }
 
+template <> void RadhydroSimulation<CollapseProblem>::ComputeDerivedVar(int lev, std::string const &dname, amrex::MultiFab &mf, const int ncomp_cc_in) const
+{
+	// compute derived variables and save in 'mf'
+	if (dname == "gpot") {
+		const int ncomp = ncomp_cc_in;
+		auto const &state = state_new_cc_[lev].const_arrays();
+		auto const &phi_data_ref = phi[lev].const_arrays();
+		auto output = mf.arrays();
+
+		amrex::ParallelFor(
+		    mf, [=] AMREX_GPU_DEVICE(int bx, int i, int j, int k) noexcept { output[bx](i, j, k, ncomp) = phi_data_ref[bx](i, j, k, ncomp); });
+	}
+}
+
 auto problem_main() -> int
 {
 	auto isNormalComp = [=](int n, int dim) {
diff --git a/src/simulation.hpp b/src/simulation.hpp
index fd14233d0..cc475de74 100644
--- a/src/simulation.hpp
+++ b/src/simulation.hpp
@@ -165,6 +165,7 @@ template <typename problem_t> class AMRSimulation : public amrex::AmrCore
 	amrex::Real reltolPoisson_ = 1.0e-5;	    // default
 	amrex::Real abstolPoisson_ = 1.0e-5;	    // default (scaled by minimum RHS value)
 	int doPoissonSolve_ = 0;		    // 1 == self-gravity enabled, 0 == disabled
+	amrex::Vector<amrex::MultiFab> phi;
 
 	amrex::Real densityFloor_ = 0.0;				// default
 	amrex::Real tempCeiling_ = std::numeric_limits<double>::max();	// default
@@ -253,6 +254,8 @@ template <typename problem_t> class AMRSimulation : public amrex::AmrCore
 	void AverageDown();
 	void AverageDownTo(int crse_lev);
 	auto timeStepWithSubcycling(int lev, amrex::Real time, bool coarseTimeBoundary, int stepsLeft) -> int;
+	void calculateGpotAllLevels();
+	void gravAccelAllLevels(amrex::Real dt);
 	void ellipticSolveAllLevels(amrex::Real dt);
 
 	void incrementFluxRegisters(amrex::MFIter &mfi, amrex::YAFluxRegister *fr_as_crse, amrex::YAFluxRegister *fr_as_fine,
@@ -576,6 +579,8 @@ template <typename problem_t> void AMRSimulation<problem_t>::setInitialCondition
 		ReadCheckpointFile();
 	}
 
+	calculateGpotAllLevels();
+
 	// abort if amrex.async_out=1, it is currently broken
 	if (amrex::AsyncOut::UseAsyncOut()) {
 		amrex::Print() << "[ERROR] [FATAL] AsyncOut is currently broken! If you want to "
@@ -893,7 +898,7 @@ template <typename problem_t> void AMRSimulation<problem_t>::evolve()
 #endif
 }
 
-template <typename problem_t> void AMRSimulation<problem_t>::ellipticSolveAllLevels(const amrex::Real dt)
+template <typename problem_t> void AMRSimulation<problem_t>::calculateGpotAllLevels()
 {
 #if AMREX_SPACEDIM == 3
 	if (doPoissonSolve_) {
@@ -911,8 +916,8 @@ template <typename problem_t> void AMRSimulation<problem_t>::ellipticSolveAllLev
 			amrex::Print() << "Doing Poisson solve...\n\n";
 		}
 
+		phi.resize(finest_level + 1);
 		// solve Poisson equation with open b.c. using the method of James (1977)
-		amrex::Vector<amrex::MultiFab> phi(finest_level + 1);
 		amrex::Vector<amrex::MultiFab> rhs(finest_level + 1);
 		const int nghost = 1;
 		const int ncomp = 1;
@@ -931,6 +936,16 @@ template <typename problem_t> void AMRSimulation<problem_t>::ellipticSolveAllLev
 		if (verbose) {
 			amrex::Print() << "\n";
 		}
+	}
+#endif
+}
+
+template <typename problem_t> void AMRSimulation<problem_t>::gravAccelAllLevels(const amrex::Real dt)
+{
+#if AMREX_SPACEDIM == 3
+	if (doPoissonSolve_) {
+
+		BL_PROFILE_REGION("GravitySolver");
 
 		// add gravitational acceleration to hydro state (using operator splitting)
 		for (int lev = 0; lev <= finest_level; ++lev) {
@@ -940,6 +955,18 @@ template <typename problem_t> void AMRSimulation<problem_t>::ellipticSolveAllLev
 #endif
 }
 
+template <typename problem_t> void AMRSimulation<problem_t>::ellipticSolveAllLevels(const amrex::Real dt)
+{
+#if AMREX_SPACEDIM == 3
+	if (doPoissonSolve_) {
+
+		calculateGpotAllLevels();
+
+		gravAccelAllLevels(dt);
+	}
+#endif
+}
+
 // N.B.: This function actually works for subcycled or not subcycled, as long as
 // nsubsteps[lev] is set correctly.
 template <typename problem_t> auto AMRSimulation<problem_t>::timeStepWithSubcycling(int lev, amrex::Real time, bool coarseTimeBoundary, int stepsLeft) -> int
diff --git a/tests/SphericalCollapse.in b/tests/SphericalCollapse.in
index db6d71804..025186a23 100644
--- a/tests/SphericalCollapse.in
+++ b/tests/SphericalCollapse.in
@@ -25,6 +25,8 @@ cfl = 0.25
 max_timesteps = 50000
 stop_time = 0.15
 
+derived_vars = gpot
+
 gravity.Gconst = 1.0    #gravitational constant
 
 do_reflux = 1