Skip to content

Commit

Permalink
Local Move of Lost Particles
Browse files Browse the repository at this point in the history
  • Loading branch information
ax3l committed Nov 22, 2023
1 parent 5c9041e commit 2d1f96d
Show file tree
Hide file tree
Showing 3 changed files with 140 additions and 36 deletions.
15 changes: 14 additions & 1 deletion src/ImpactX.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -34,7 +34,8 @@ namespace impactx
{
ImpactX::ImpactX ()
: AmrCore(initialization::init_amr_core()),
m_particle_container(std::make_unique<ImpactXParticleContainer>(this))
m_particle_container(std::make_unique<ImpactXParticleContainer>(this)),
m_particles_lost(std::make_unique<ImpactXParticleContainer>(this))
{
// todo: if amr.n_cells is provided, overwrite/redefine AmrCore object

Expand Down Expand Up @@ -82,6 +83,18 @@ namespace impactx
// init blocks / grids & MultiFabs
AmrCore::InitFromScratch(0.0);
amrex::Print() << "boxArray(0) " << boxArray(0) << std::endl;

// alloc particle containers
// the lost particles have an extra runtime attribute: s when it was lost
bool comm = true;
m_particles_lost->AddRealComp(comm);

// have to resize here, not in the constructor because grids have not
// been built when constructor was called.
m_particle_container->reserveData();
m_particle_container->resizeData();
m_particles_lost->reserveData();
m_particles_lost->resizeData();
}

void ImpactX::evolve ()
Expand Down
156 changes: 126 additions & 30 deletions src/particles/CollectLost.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -11,48 +11,144 @@

#include <AMReX_GpuLaunch.H>
#include <AMReX_GpuQualifiers.H>
#include <AMReX_Math.H>
#include <AMReX_ParticleTransformation.H>
#include <AMReX_RandomEngine.H>


namespace impactx
{
void collect_lost_particles (ImpactXParticleContainer& source, ImpactXParticleContainer& dest)
{
BL_PROFILE("impactX::collect_lost_particles");

using SrcData = ImpactXParticleContainer::ParticleTileType::ConstParticleTileDataType;
using DstData = ImpactXParticleContainer::ParticleTileType::ParticleTileDataType;

// copy all particles marked with a negative ID from source to destination
bool const local = true;
dest.copyParticles(
source,
[=] AMREX_GPU_HOST_DEVICE(const SrcData &src, int ip) {
return src.id(ip) < 0;
},
local
);
RefPart const ref_part = source.GetRefParticle();
auto const s_lost = ref_part.s;

// flip IDs back to positive in destination
int const nLevel = dest.finestLevel();
// have to resize here, not in the constructor because grids have not
// been built when constructor was called.
dest.reserveData();
dest.resizeData();

// copy all particles marked with a negative ID from source to destination
int const nLevel = source.finestLevel();
for (int lev = 0; lev <= nLevel; ++lev) {
// loop over all particle boxes
using ParIt = ImpactXParticleContainer::iterator;
auto& plevel = source.GetParticles(lev);

// TODO: is it safe to add OpenMP parallelism here?
for (ParIt pti(source, lev); pti.isValid(); ++pti) {
auto index = std::make_pair(pti.index(), pti.LocalTileIndex());
if (plevel.find(index) == plevel.end()) continue;

auto& ptile_source = plevel.at(index);
auto const np = ptile_source.numParticles();
if (np == 0) continue; // no particles in source tile

// we will copy particles that were marked as lost, with a negative id
auto predicate = [] AMREX_GPU_HOST_DEVICE (const SrcData& src, int ip, const amrex::RandomEngine& /*engine*/) noexcept
{
return src.id(ip) < 0;
};

auto& ptile_dest = dest.DefineAndReturnParticleTile(
lev, pti.index(), pti.LocalTileIndex());

// count how many particles we will copy
amrex::ReduceOps<amrex::ReduceOpSum> reduce_op;
amrex::ReduceData<int> reduce_data(reduce_op);
{
const auto src_data = ptile_source.getConstParticleTileData();

const amrex::RandomEngine rng{}; // unused
reduce_op.eval(np, reduce_data, [=] AMREX_GPU_HOST_DEVICE (int ip) {
return predicate(src_data, ip, rng) ? 1 : 0;
});
}
int const np_to_move = amrex::get<0>(reduce_data.value());
if (np_to_move == 0) continue; // no particles to move from source tile

// allocate memory in destination
int const dst_index = ptile_dest.numParticles();
ptile_dest.resize(dst_index + np_to_move);

// copy particles
// skipped in loop below: integer compile-time or runtime attributes
AMREX_ALWAYS_ASSERT(SrcData::NAI == 0);
AMREX_ALWAYS_ASSERT(ptile_source.NumRuntimeIntComps() == 0);

// first runtime attribute in destination is s position when particle got lost
int const s_index = dest.NumRuntimeRealComps() - 1;
auto copy_and_mark_negative = [&s_index, &s_lost](DstData& dst, const SrcData& src, int src_ip, int dst_ip) noexcept
{
dst.m_aos[dst_ip] = src.m_aos[src_ip];

for (int j = 0; j < SrcData::NAR; ++j)
dst.m_rdata[j][dst_ip] = src.m_rdata[j][src_ip];
for (int j = 0; j < src.m_num_runtime_real; ++j)
dst.m_runtime_rdata[j][dst_ip] = src.m_runtime_rdata[j][src_ip];

// unused: integer compile-time or runtime attributes
//for (int j = 0; j < SrcData::NAI; ++j)
// dst.m_idata[j][dst_ip] = src.m_idata[j][src_ip];
//for (int j = 0; j < src.m_num_runtime_int; ++j)
// dst.m_runtime_idata[j][dst_ip] = src.m_runtime_idata[j][src_ip];

// flip id to positive in destination
dst.id(dst_ip) = amrex::Math::abs(dst.id(dst_ip));

// remember the current s of the ref particle when lost
dst.m_runtime_rdata[s_index][dst_ip] = s_lost;
};

amrex::filterAndTransformParticles(
ptile_dest,
ptile_source,
predicate,
copy_and_mark_negative,
0,
dst_index
);

// remove particles with negative ids in source
{
int n_removed = 0;
auto ptile_src_data = ptile_source.getParticleTileData();
for (int ip = 0; ip < np; ++ip)
{
if (ptile_source.id(ip) < 0)
n_removed++;
else
{
if (n_removed > 0)
{
// move down
int const new_index = ip - n_removed;

ptile_src_data.m_aos[new_index] = ptile_src_data.m_aos[ip];

for (int j = 0; j < SrcData::NAR; ++j)
ptile_src_data.m_rdata[j][new_index] = ptile_src_data.m_rdata[j][ip];
for (int j = 0; j < ptile_src_data.m_num_runtime_real; ++j)
ptile_src_data.m_runtime_rdata[j][new_index] = ptile_src_data.m_runtime_rdata[j][ip];

// unused: integer compile-time or runtime attributes
//for (int j = 0; j < SrcData::NAI; ++j)
// dst.m_idata[j][new_index] = src.m_idata[j][ip];
//for (int j = 0; j < ptile_src_data.m_num_runtime_int; ++j)
// dst.m_runtime_idata[j][new_index] = src.m_runtime_idata[j][ip];
}
}
}
AMREX_ALWAYS_ASSERT(np_to_move == n_removed);
ptile_source.resize(np - n_removed);
}

#ifdef AMREX_USE_OMP
#pragma omp parallel if (amrex::Gpu::notInLaunchRegion())
#endif
for (ParIt pti(dest, lev); pti.isValid(); ++pti) {
const int np = pti.numParticles();

// preparing access to particle data: AoS
using PType = ImpactXParticleContainer::ParticleType;
auto &aos = pti.GetArrayOfStructs();
PType *AMREX_RESTRICT aos_ptr = aos().dataPtr();

amrex::ParallelFor(np, [=] AMREX_GPU_DEVICE(long i) {
PType &p = aos_ptr[i];
p.id() = -p.id();
});
}
}

// TODO remove particles with negative ids in source
} // particle tile loop
} // lev
}
} // namespace impactx
5 changes: 0 additions & 5 deletions src/particles/ImpactXParticleContainer.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -104,11 +104,6 @@ namespace impactx
// number of particles to add
int const np = x.size();

// have to resize here, not in the constructor because grids have not
// been built when constructor was called.
reserveData();
resizeData();

auto& particle_tile = DefineAndReturnParticleTile(0, 0, 0);

/* Create a temporary tile to obtain data from simulation. This data
Expand Down

0 comments on commit 2d1f96d

Please sign in to comment.