Merge branch 'BenWibking/tracer-particles' of github.com:quokka-astro…

…/quokka into BenWibking/tracer-particles
quokka-astro · Jan 15, 2024 · 3b60b53 · 3b60b53
2 parents 00da54e + 7f6c84d
commit 3b60b53
Show file tree

Hide file tree

Showing 3 changed files with 45 additions and 2 deletions.
diff --git 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
@@ -171,6 +171,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
@@ -284,6 +285,8 @@ template <typename problem_t> class AMRSimulation : public amrex::AmrCore
 	void AverageDown();
 	void AverageDownTo(int crse_lev);
 	void timeStepWithSubcycling(int lev, amrex::Real time, int iteration);
+	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,
@@ -629,6 +632,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 "
@@ -945,7 +950,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_) {
@@ -963,8 +968,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;
@@ -983,6 +988,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) {
@@ -992,6 +1007,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> void AMRSimulation<problem_t>::timeStepWithSubcycling(int lev, amrex::Real time, int iteration)

diff --git 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