Skip to content

Commit

Permalink
ParticleInitData for PureSoA
Browse files Browse the repository at this point in the history
  • Loading branch information
ax3l committed Feb 2, 2024
1 parent de5215b commit 58be049
Show file tree
Hide file tree
Showing 4 changed files with 13 additions and 16 deletions.
15 changes: 8 additions & 7 deletions src/Particle/ParticleContainer.H
Original file line number Diff line number Diff line change
Expand Up @@ -380,13 +380,17 @@ void make_ParticleContainer_and_Iterators (py::module &m, std::string allocstr)
// }
;

py_pc
.def("InitRandom", py::overload_cast<Long, ULong, const ParticleInitData&, bool, RealBox>(&ParticleContainerType::InitRandom))
;

// TODO for pure SoA
// depends on https://github.com/AMReX-Codes/amrex/pull/3280
if constexpr (!T_ParticleType::is_soa_particle) {
py_pc
.def("InitRandom", py::overload_cast<Long, ULong, const ParticleInitData&, bool, RealBox>(&ParticleContainerType::InitRandom)) // TODO pure SoA
.def("InitRandomPerBox", py::overload_cast<Long, ULong, const ParticleInitData&>(&ParticleContainerType::InitRandomPerBox)) // TODO pure SoA
.def("InitOnePerCell", &ParticleContainerType::InitOnePerCell);
.def("InitRandomPerBox", py::overload_cast<Long, ULong, const ParticleInitData&>(&ParticleContainerType::InitRandomPerBox))
.def("InitOnePerCell", &ParticleContainerType::InitOnePerCell)
;
}

using iterator = amrex::ParIter_impl<ParticleType, T_NArrayReal, T_NArrayInt, Allocator>;
Expand All @@ -408,10 +412,7 @@ void make_ParticleContainer_and_Iterators (py::module &m, std::string allocstr)
template <typename T_ParticleType, int T_NArrayReal=0, int T_NArrayInt=0>
void make_ParticleContainer_and_Iterators (py::module &m)
{
// TODO for pure SoA
// depends on https://github.com/AMReX-Codes/amrex/pull/3280
if constexpr (!T_ParticleType::is_soa_particle)
make_ParticleInitData<T_ParticleType, T_NArrayReal, T_NArrayInt>(m);
make_ParticleInitData<T_ParticleType, T_NArrayReal, T_NArrayInt>(m);

// first, because used as copy target in methods in containers with other allocators
make_ParticleContainer_and_Iterators<T_ParticleType, T_NArrayReal, T_NArrayInt,
Expand Down
1 change: 0 additions & 1 deletion src/Particle/ParticleContainer_ImpactX.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -14,6 +14,5 @@ void init_ParticleContainer_ImpactX(py::module& m) {

// TODO: we might need to move all or most of the defines in here into a
// test/example submodule, so they do not collide with downstream projects
make_ParticleContainer_and_Iterators<Particle<0, 0>, 5, 0>(m); // ImpactX 22.07 - 24.02
make_ParticleContainer_and_Iterators<SoAParticle<8, 0>, 8, 0>(m); // ImpactX 24.03+
}
3 changes: 0 additions & 3 deletions src/Particle/ParticleContainer_WarpX.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -13,9 +13,6 @@ void init_ParticleContainer_WarpX(py::module& m) {

// TODO: we might need to move all or most of the defines in here into a
// test/example submodule, so they do not collide with downstream projects
make_ParticleContainer_and_Iterators<Particle<0, 0>, 4, 0>(m); // WarpX 22.07 - 24.02 1D-3D
//make_ParticleContainer_and_Iterators<Particle<0, 0>, 5, 0> (m); // WarpX 22.07 - 24.02 RZ

#if AMREX_SPACEDIM == 1
make_ParticleContainer_and_Iterators<SoAParticle<5, 0>, 5, 0>(m); // WarpX 24.03+ 1D
#elif AMREX_SPACEDIM == 2
Expand Down
10 changes: 5 additions & 5 deletions tests/test_particleContainer.py
Original file line number Diff line number Diff line change
Expand Up @@ -300,19 +300,19 @@ class Config:
# pc = sim.get_particles()
# Config = sim.extension.Config

# iterate over every mesh-refinement levels (no MR: lev=0)
# iterate over mesh-refinement level (no MR: lev=0)
for lvl in range(pc.finest_level + 1):
# get every local chunk of particles
for pti in pc.iterator(pc, level=lvl):
# additional compile-time and runtime attributes in SoA format
# compile-time and runtime attributes in SoA format
soa = pti.soa().to_cupy() if Config.have_gpu else pti.soa().to_numpy()

# notes:
# Only the next lines are the "HOT LOOP" of the computation.
# For efficiency, use numpy array operation for speed.
# For speed, use array operations.

# write to all particles in the chunk
# note: careful, if you change particle positions, you need to
# note: careful, if you change particle positions, you might need to
# redistribute particles before continuing the simulation step
print(soa.real)
soa.real[0][()] = 0.30 # x
Expand Down Expand Up @@ -341,7 +341,7 @@ class Config:
# pc = sim.get_particles()
# Config = sim.extension.Config

# iterate over every mesh-refinement levels (no MR: lev=0)
# iterate over mesh-refinement level (no MR: lev=0)
for lvl in range(pc.finest_level + 1):
# get every local chunk of particles
for pti in pc.iterator(pc, level=lvl):
Expand Down

0 comments on commit 58be049

Please sign in to comment.