add particle drift

quokka-astro · Jan 18, 2024 · 8dcc8f7 · 8dcc8f7
1 parent bb32493
commit 8dcc8f7
Show file tree

Hide file tree

Showing 2 changed files with 21 additions and 6 deletions.
diff --git a/src/CICParticles.hpp b/src/CICParticles.hpp
@@ -15,6 +15,7 @@
 #include "AMReX_MultiFabUtil.H"
 #include "AMReX_ParticleInterpolators.H"
 #include "AMReX_Particles.H"
+#include "AMReX_ParIter.H"
 namespace quokka
 {
 
@@ -23,7 +24,9 @@ constexpr int CICParticleInts = 0;
 
 enum ParticleDataIdx { ParticleMassIdx = 0, ParticleVxIdx, ParticleVyIdx, ParticleVzIdx };
 
-using CICParticleContainer = amrex::AmrParticleContainer<0, 0, CICParticleReals, CICParticleInts>;
+// TODO(bwibking): switch to using struct-of-arrays (SOA) particles
+using CICParticleContainer = amrex::AmrParticleContainer<CICParticleReals, CICParticleInts>;
+using CICParticleIterator = amrex::ParIter<CICParticleReals, CICParticleInts>;
 
 struct CICDeposition {
 	amrex::Real Gconst{};

diff --git a/src/simulation.hpp b/src/simulation.hpp
@@ -1080,7 +1080,9 @@ template <typename problem_t> void AMRSimulation<problem_t>::kickParticlesAllLev
 
 	if (do_cic_particles != 0) {
 		for (int lev = 0; lev <= finest_level; ++lev) {
-			const int nc = AMREX_SPACEDIM;
+			const int nstart_mesh = 0;
+			const int nstart_particle = quokka::ParticleVxIdx;
+			const int ncomps = AMREX_SPACEDIM;
 			const auto plo = geom[lev].ProbLoArray();
 			const auto dx_inv = geom[lev].InvCellSizeArray();
 
@@ -1091,7 +1093,7 @@ template <typename problem_t> void AMRSimulation<problem_t>::kickParticlesAllLev
 				    amrex::ParticleInterpolator::Linear interp(p, plo, dx_inv);
 
 				    interp.MeshToParticle(
-					p, acc, 0, 1 + AMREX_SPACEDIM, nc,
+					p, acc, nstart_mesh, nstart_particle, ncomps,
 					[=] AMREX_GPU_DEVICE(amrex::Array4<const amrex::Real> const &arr, int i, int j, int k, int comp) {
 						// compute cell-centered acceleration -grad(phi)
 						if (comp == 0) {
@@ -1105,9 +1107,9 @@ template <typename problem_t> void AMRSimulation<problem_t>::kickParticlesAllLev
 						}
 					},
 
-					[=] AMREX_GPU_DEVICE(quokka::CICParticleContainer::ParticleType & part, int comp, amrex::Real val) {
+					[=] AMREX_GPU_DEVICE(quokka::CICParticleContainer::ParticleType & p, int const comp, amrex::Real const acc_comp) {
 						// kick particle 'part' by updating its velocity
-						part.rdata(comp) += 0.5 * dt * amrex::ParticleReal(val);
+						p.rdata(comp) += 0.5 * dt * amrex::ParticleReal(acc_comp);
 					});
 			    });
 		}
@@ -1119,7 +1121,17 @@ template <typename problem_t> void AMRSimulation<problem_t>::driftParticlesAllLe
 	// drift all particles (do: pos[i] += dt * vel[i])
 
 	if (do_cic_particles != 0) {
-		// TODO(bwibking): implement
+		for (int lev = 0; lev <= finest_level; ++lev) {
+			for (quokka::CICParticleIterator pIter(*CICParticles, lev); pIter.isValid(); ++pIter) {
+				auto &particles = pIter.GetArrayOfStructs();
+				for (auto &particle : particles) {
+					// update particle position
+					for (int i = 0; i < AMREX_SPACEDIM; ++i) {
+						particle.pos(i) += dt * particle.rdata(quokka::ParticleVxIdx + i);
+					}
+				}
+			}
+		}
 	}
 }