Fix ComputeSoundSpeed for particles with MHD (quokka-astro#1481)

### Description
Modifies ComputeSoundSpeed in the particle creation machinery to
correctly subtract magnetic energy.

### Related issues
Fixes quokka-astro#1467.

### 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 (if applicable; see
above).
- [x] I have read the [Contributing
Guide](https://github.com/quokka-astro/quokka/blob/development/CONTRIBUTING.md).
- [ ] I have added tests for any new physics that this PR adds to the
code.
- [ ] *(For quokka-astro org members)* I have manually triggered the GPU
tests with the magic comment `/azp run`.

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: ChongChong He <chongchonghe99@gmail.com>

Author: 3 people

inputs/ParticleCreationAMR.in

@@ -15,8 +15,8 @@ amr.v              = 1       # verbosity in Amr
 # *****************************************************************
 amr.n_cell          = 32 32 32
 amr.max_level       = 1     # number of levels = max_level + 1
-amr.blocking_factor = 8     # grid size must be divisible by this
-amr.max_grid_size   = 8    # at least 128 for GPUs
+amr.blocking_factor = 16    # needs to be >= ceil(nghost/ref_ratio)*ref_ratio for proper nesting
+amr.max_grid_size   = 16    # at least 128 for GPUs
 amr.n_error_buf     = 2     # minimum 2 cell buffer around tagged cells
 amr.grid_eff        = 1.0
 

src/hydro/EOS.hpp

@@ -321,6 +321,14 @@ AMREX_FORCE_INLINE AMREX_GPU_HOST_DEVICE auto EOS<problem_t>::ComputePressure(am
 {
 	// return pressure for an ideal gas
 	amrex::Real P = NAN;
+
+	if constexpr (gamma_ == 1.0) {
+		static_assert(EOS_Traits<problem_t>::cs_isothermal > 0.0, "EOS_Traits<problem_t>::cs_isothermal must be set when gamma=1.");
+		amrex::ignore_unused(Eint);
+		amrex::ignore_unused(massScalars);
+		P = rho * EOS_Traits<problem_t>::cs_isothermal * EOS_Traits<problem_t>::cs_isothermal;
+		return P;
+	}
 #ifdef CHEMISTRY
 	eos_t chemstate;
 	chemstate.rho = rho;
@@ -342,19 +350,17 @@ AMREX_FORCE_INLINE AMREX_GPU_HOST_DEVICE auto EOS<problem_t>::ComputePressure(am
 #else
 	amrex::ignore_unused(massScalars);
 
-	if constexpr (gamma_ != 1.0) {
-		chem_eos_t estate;
-		estate.rho = rho;
-		// if rho is 0, pass 0 to state.e
-		if (rho == 0.0) {
-			estate.e = 0;
-		} else {
-			estate.e = Eint / rho;
-		}
-		estate.mu = mean_molecular_weight_ / C::m_u;
-		eos(eos_input_re, estate);
-		P = estate.p;
+	chem_eos_t estate;
+	estate.rho = rho;
+	// if rho is 0, pass 0 to state.e
+	if (rho == 0.0) {
+		estate.e = 0;
+	} else {
+		estate.e = Eint / rho;
 	}
+	estate.mu = mean_molecular_weight_ / C::m_u;
+	eos(eos_input_re, estate);
+	P = estate.p;
 #endif
 	return P;
 }
@@ -367,6 +373,15 @@ AMREX_FORCE_INLINE AMREX_GPU_HOST_DEVICE auto EOS<problem_t>::ComputeSoundSpeed(
 	// return sound speed for an ideal gas
 	amrex::Real cs = NAN;
 
+	if constexpr (gamma_ == 1.0) {
+		static_assert(EOS_Traits<problem_t>::cs_isothermal > 0.0, "EOS_Traits<problem_t>::cs_isothermal must be set when gamma=1.");
+		amrex::ignore_unused(rho);
+		amrex::ignore_unused(Pressure);
+		amrex::ignore_unused(massScalars);
+		cs = EOS_Traits<problem_t>::cs_isothermal;
+		return cs;
+	}
+
 #ifdef CHEMISTRY
 	eos_t chemstate;
 	chemstate.rho = rho;
@@ -388,14 +403,12 @@ AMREX_FORCE_INLINE AMREX_GPU_HOST_DEVICE auto EOS<problem_t>::ComputeSoundSpeed(
 #else
 	amrex::ignore_unused(massScalars);
 
-	if constexpr (gamma_ != 1.0) {
-		chem_eos_t estate;
-		estate.rho = rho;
-		estate.p = Pressure;
-		estate.mu = mean_molecular_weight_ / C::m_u;
-		eos(eos_input_rp, estate);
-		cs = estate.cs;
-	}
+	chem_eos_t estate;
+	estate.rho = rho;
+	estate.p = Pressure;
+	estate.mu = mean_molecular_weight_ / C::m_u;
+	eos(eos_input_rp, estate);
+	cs = estate.cs;
 #endif
 	return cs;
 }

src/particles/PhysicsParticles.hpp

@@ -1,9 +1,11 @@
 #ifndef PHYSICS_PARTICLES_HPP_
 #define PHYSICS_PARTICLES_HPP_
 
+#include <array>
 #include <cstdint>
 #include <map>
 #include <memory>
+#include <ranges>
 #include <string>
 
 #include <fmt/format.h>
@@ -136,18 +138,27 @@ class PhysicsParticleDescriptorBase
 
 	//----- Methods that are implemented for some but not all particle types, so they cannot be pure virtual -----
 
-	virtual auto depositSN(amrex::MultiFab & /*state*/, int /*lev*/, amrex::Real /*time*/, amrex::Real /*dt*/) -> amrex::Real { return 0.0_rt; }
+	virtual auto depositSN(amrex::MultiFab & /*state*/, std::array<amrex::MultiFab, AMREX_SPACEDIM> const * /*state_fc*/, int /*lev*/, amrex::Real /*time*/,
+			       amrex::Real /*dt*/) -> amrex::Real
+	{
+		return 0.0_rt;
+	}
 
-	virtual void computeSinkAccretion(amrex::MultiFab &state, amrex::MultiFab &state_accretion_rate, int lev, amrex::Real time, amrex::Real dt)
-	{ /* Default empty implementation */ }
+	virtual void computeSinkAccretion(amrex::MultiFab &state, amrex::MultiFab &state_accretion_rate,
+					  std::array<amrex::MultiFab, AMREX_SPACEDIM> const *state_fc, int lev, amrex::Real time, amrex::Real dt)
+	{ /* Default empty implementation */
+	}
 
 	// Create particles from hydro state at the finest level
 	// Note: particles are not allowed to spawn outside of real cells. If they do, we will need a redistribution immediately after this call in order to
 	// make particle-mesh interaction work.
-	virtual void createParticlesFromState(amrex::MultiFab &state, amrex::MultiFab &accretion_rate, int lev, amrex::Real current_time, amrex::Real dt)
-	{ /* Default empty implementation */ }
+	virtual void createParticlesFromState(amrex::MultiFab &state, amrex::MultiFab &accretion_rate, int lev, amrex::Real current_time, amrex::Real dt,
+					      std::array<amrex::MultiFab, AMREX_SPACEDIM> const *state_fc)
+	{ /* Default empty implementation */
+	}
 
-	virtual void applySinkAccretion(amrex::MultiFab &state, amrex::MultiFab &state_accretion_rate, const amrex::Geometry &geom, int lev, amrex::Real time,
+	virtual void applySinkAccretion(amrex::MultiFab &state, amrex::MultiFab &state_accretion_rate,
+					std::array<amrex::MultiFab, AMREX_SPACEDIM> const *state_fc, const amrex::Geometry &geom, int lev, amrex::Real time,
 					amrex::Real dt)
 	{ /* Default empty implementation */
 	}
@@ -562,7 +573,8 @@ template <typename ContainerType, typename problem_t, ParticleType particleType>
 	}
 
 	// Implementation of supernova energy and momentum deposition from particles to grid
-	auto depositSN(amrex::MultiFab &state, int lev, amrex::Real time, amrex::Real dt) -> amrex::Real override
+	auto depositSN(amrex::MultiFab &state, std::array<amrex::MultiFab, AMREX_SPACEDIM> const *state_fc, int lev, amrex::Real time, amrex::Real dt)
+	    -> amrex::Real override
 	{
 		amrex::Real max_velocity = 0.0;
 
@@ -573,7 +585,7 @@ template <typename ContainerType, typename problem_t, ParticleType particleType>
 								 "UnitSystem must be CGS for particleMeshInteraction");
 
 				// Deposit supernova energy and momentum from all particles. This also updates the evolution stage of the particles.
-				max_velocity = SNDeposition<ContainerType, problem_t>(this->container_, state, lev, time, dt, this->getMassIndex(),
+				max_velocity = SNDeposition<ContainerType, problem_t>(this->container_, state, state_fc, lev, time, dt, this->getMassIndex(),
 										      this->getEvolutionStageIndex(), this->getBirthTimeIndex());
 			} else {
 				// Only update evolution stage but not deposit energy/momentum
@@ -586,26 +598,28 @@ template <typename ContainerType, typename problem_t, ParticleType particleType>
 	}
 
 	// compute accretion rate
-	void computeSinkAccretion(amrex::MultiFab &state, amrex::MultiFab &state_accretion_rate, int lev, amrex::Real time, amrex::Real dt) override
+	void computeSinkAccretion(amrex::MultiFab &state, amrex::MultiFab &state_accretion_rate, std::array<amrex::MultiFab, AMREX_SPACEDIM> const *state_fc,
+				  int lev, amrex::Real time, amrex::Real dt) override
 	{
-		SinkAccretionUtils::computeAccretion<ContainerType, problem_t>(this->container_, state, state_accretion_rate, lev, time, dt,
+		SinkAccretionUtils::computeAccretion<ContainerType, problem_t>(this->container_, state, state_accretion_rate, state_fc, lev, time, dt,
 									       this->getMassIndex());
 	}
 
 	// apply accretion
-	void applySinkAccretion(amrex::MultiFab &state, amrex::MultiFab &state_accretion_rate, const amrex::Geometry &geom, int lev, amrex::Real time,
-				amrex::Real dt) override
+	void applySinkAccretion(amrex::MultiFab &state, amrex::MultiFab &state_accretion_rate, std::array<amrex::MultiFab, AMREX_SPACEDIM> const *state_fc,
+				const amrex::Geometry &geom, int lev, amrex::Real time, amrex::Real dt) override
 	{
-		SinkAccretionUtils::applyAccretion<ContainerType, problem_t>(this->container_, state, state_accretion_rate, geom, lev, time, dt,
+		SinkAccretionUtils::applyAccretion<ContainerType, problem_t>(this->container_, state, state_accretion_rate, state_fc, geom, lev, time, dt,
 									     this->getMassIndex());
 	}
 
-	void createParticlesFromState(amrex::MultiFab &state, amrex::MultiFab &accretion_rate, int lev, amrex::Real current_time, amrex::Real dt) override
+	void createParticlesFromState(amrex::MultiFab &state, amrex::MultiFab &accretion_rate, int lev, amrex::Real current_time, amrex::Real dt,
+				      std::array<amrex::MultiFab, AMREX_SPACEDIM> const *state_fc) override
 	{
 		// Use the traits class to implement the specialized behavior
 		ParticleCreationTraits<particleType>::template createParticles<problem_t, ContainerType>(
 		    this->container_, this->getMassIndex(), state, accretion_rate, lev, current_time, dt, this->getEvolutionStageIndex(),
-		    this->getBirthTimeIndex());
+		    this->getBirthTimeIndex(), state_fc);
 	}
 #endif // AMREX_SPACEDIM == 3
 };
@@ -756,37 +770,39 @@ template <typename problem_t> class PhysicsParticleRegister
 	}
 
 	// Deposit supernova energy and momentum from all particles
-	auto depositSN(amrex::MultiFab &state, int lev, amrex::Real time, amrex::Real dt) -> amrex::Real
+	auto depositSN(amrex::MultiFab &state, std::array<amrex::MultiFab, AMREX_SPACEDIM> const *state_fc, int lev, amrex::Real time, amrex::Real dt)
+	    -> amrex::Real
 	{
 		const BL_PROFILE("PhysicsParticleRegister::depositSN()");
 		amrex::Real max_velocity = 0.0;
 		// Each particle type handles its own buffer creation and roundoff independently
 		for (const auto &[type, descriptor] : particleRegistry_) {
-			const amrex::Real max_velocity_ = descriptor->depositSN(state, lev, time, dt);
+			const amrex::Real max_velocity_ = descriptor->depositSN(state, state_fc, lev, time, dt);
 			max_velocity = std::max(max_velocity, max_velocity_);
 		}
 		return max_velocity;
 	}
 
 	// Implementation of computeSinkAccretion
-	void computeSinkAccretion(amrex::MultiFab &state, amrex::MultiFab &state_accretion_rate, int lev, amrex::Real time, amrex::Real dt)
+	void computeSinkAccretion(amrex::MultiFab &state, amrex::MultiFab &state_accretion_rate, std::array<amrex::MultiFab, AMREX_SPACEDIM> const *state_fc,
+				  int lev, amrex::Real time, amrex::Real dt)
 	{
 		const BL_PROFILE("PhysicsParticleRegister::computeSinkAccretion()");
 		for (const auto &[type, descriptor] : particleRegistry_) {
 			if (descriptor->getAllowsAccretion()) {
-				descriptor->computeSinkAccretion(state, state_accretion_rate, lev, time, dt);
+				descriptor->computeSinkAccretion(state, state_accretion_rate, state_fc, lev, time, dt);
 			}
 		}
 	}
 
 	// Implementation of applySinkAccretion
-	void applySinkAccretion(amrex::MultiFab &state, amrex::MultiFab &state_accretion_rate, const amrex::Geometry &geom, int lev, amrex::Real time,
-				amrex::Real dt)
+	void applySinkAccretion(amrex::MultiFab &state, amrex::MultiFab &state_accretion_rate, std::array<amrex::MultiFab, AMREX_SPACEDIM> const *state_fc,
+				const amrex::Geometry &geom, int lev, amrex::Real time, amrex::Real dt)
 	{
 		const BL_PROFILE("PhysicsParticleRegister::applySinkAccretion()");
 		for (const auto &[type, descriptor] : particleRegistry_) {
 			if (descriptor->getAllowsAccretion()) {
-				descriptor->applySinkAccretion(state, state_accretion_rate, geom, lev, time, dt);
+				descriptor->applySinkAccretion(state, state_accretion_rate, state_fc, geom, lev, time, dt);
 			}
 		}
 	}
@@ -830,7 +846,7 @@ template <typename problem_t> class PhysicsParticleRegister
 			const std::string typeName = getParticleTypeName(type);
 
 			// Check if this particle type is in the requested list
-			if (std::find(particleTypeNames.begin(), particleTypeNames.end(), typeName) != particleTypeNames.end()) {
+			if (std::ranges::find(particleTypeNames, typeName) != particleTypeNames.end()) {
 				descriptor->writePlotFile(plotfilename, typeName);
 				descriptor->writeUnitsFile(plotfilename, typeName);
 			}
@@ -887,14 +903,15 @@ template <typename problem_t> class PhysicsParticleRegister
 	}
 
 	// Create particles based on particle type
-	void createParticlesFromState(amrex::MultiFab &state, amrex::MultiFab &accretion_rate, int lev, amrex::Real current_time, amrex::Real dt)
+	void createParticlesFromState(amrex::MultiFab &state, amrex::MultiFab &accretion_rate, int lev, amrex::Real current_time, amrex::Real dt,
+				      std::array<amrex::MultiFab, AMREX_SPACEDIM> const *state_fc = nullptr)
 	{
 		const BL_PROFILE("PhysicsParticleRegister::createParticlesFromState()");
 		for (const auto &[type, descriptor] : particleRegistry_) {
 			// Only create particles if the descriptor allows creation
 			if (descriptor->getAllowsCreation()) {
 				// Call the appropriate particle creation method based on the particle type
-				descriptor->createParticlesFromState(state, accretion_rate, lev, current_time, dt);
+				descriptor->createParticlesFromState(state, accretion_rate, lev, current_time, dt, state_fc);
 
 				// redistribute particles
 				// descriptor->redistribute(lev);