Skip to content

Commit

Permalink
Refactor output of phi on the particles (#4839)
Browse files Browse the repository at this point in the history 
* Refactor output of phi on the particles
  • Loading branch information
@RemiLehe
RemiLehe authored Apr 9, 2024
1 parent 56a00e9 commit 5640de3
Show file tree
Hide file tree
Showing 3 changed files with 62 additions and 64 deletions.
54 changes: 54 additions & 0 deletions Source/Diagnostics/ParticleIO.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -235,3 +235,57 @@ MultiParticleContainer::WriteHeader (std::ostream& os) const
allcontainers.at(i)->WriteHeader(os);
}
}

void
storePhiOnParticles ( PinnedMemoryParticleContainer& tmp,
int electrostatic_solver_id, bool is_full_diagnostic ) {

using PinnedParIter = typename PinnedMemoryParticleContainer::ParIterType;

WARPX_ALWAYS_ASSERT_WITH_MESSAGE(
(electrostatic_solver_id == ElectrostaticSolverAlgo::LabFrame) ||
(electrostatic_solver_id == ElectrostaticSolverAlgo::LabFrameElectroMagnetostatic),
"Output of the electrostatic potential (phi) on the particles was requested, "
"but this is only available for `warpx.do_electrostatic=labframe` or `labframe-electromagnetostatic`.");
// When this is not a full diagnostic, the particles are not written at the same physical time (i.e. PIC iteration)
// that they were collected. This happens for diagnostics that use buffering (e.g. BackTransformed, BoundaryScraping).
// Here `phi` is gathered at the iteration when particles are written (not collected) and is thus mismatched.
// To avoid confusion, we raise an error in this case.
WARPX_ALWAYS_ASSERT_WITH_MESSAGE(
is_full_diagnostic,
"Output of the electrostatic potential (phi) on the particles was requested, "
"but this is only available with `diag_type = Full`.");
tmp.AddRealComp("phi");
int const phi_index = tmp.getParticleComps().at("phi");
auto& warpx = WarpX::GetInstance();
#ifdef AMREX_USE_OMP
#pragma omp parallel if (amrex::Gpu::notInLaunchRegion())
#endif
for (int lev=0; lev<=warpx.finestLevel(); lev++) {
const amrex::Geometry& geom = warpx.Geom(lev);
auto plo = geom.ProbLoArray();
auto dxi = geom.InvCellSizeArray();
amrex::MultiFab const& phi = warpx.getField( FieldType::phi_fp, lev, 0 );

for (PinnedParIter pti(tmp, lev); pti.isValid(); ++pti) {

auto phi_grid = phi[pti].array();
const auto getPosition = GetParticlePosition<PIdx>(pti);
amrex::ParticleReal* phi_particle_arr = pti.GetStructOfArrays().GetRealData(phi_index).dataPtr();

// Loop over the particles and update their position
amrex::ParallelFor( pti.numParticles(),
[=] AMREX_GPU_DEVICE (long ip) {

amrex::ParticleReal xp, yp, zp;
getPosition(ip, xp, yp, zp);
int i, j, k;
amrex::Real W[AMREX_SPACEDIM][2];
ablastr::particles::compute_weights(xp, yp, zp, plo, dxi, i, j, k, W);
amrex::Real const phi_value = ablastr::particles::interp_field_nodal(i, j, k, W, phi_grid);
phi_particle_arr[ip] = phi_value;
}
);
}
}
}
27 changes: 1 addition & 26 deletions Source/Diagnostics/WarpXOpenPMD.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -591,32 +591,7 @@ for (unsigned i = 0, n = particle_diags.size(); i < n; ++i) {

// Gather the electrostatic potential (phi) on the macroparticles
if ( particle_diags[i].m_plot_phi ) {
WARPX_ALWAYS_ASSERT_WITH_MESSAGE(
(WarpX::electrostatic_solver_id == ElectrostaticSolverAlgo::LabFrame) ||
(WarpX::electrostatic_solver_id == ElectrostaticSolverAlgo::LabFrameElectroMagnetostatic),
"Output of the electrostatic potential (phi) on the particles was requested, "
"but this is only available for `warpx.do_electrostatic=labframe` or `labframe-electromagnetostatic`.");
// Using pinned PC indicates that the particles are not written at the same physical time (i.e. PIC iteration)
// that they were collected. This happens for diagnostics that use buffering (e.g. BackTransformed, BoundaryScraping).
// Here `phi` is gathered at the iteration when particles are written (not collected) and is thus mismatched.
// To avoid confusion, we raise an error in this case.
WARPX_ALWAYS_ASSERT_WITH_MESSAGE(
use_pinned_pc == false,
"Output of the electrostatic potential (phi) on the particles was requested, "
"but this is only available with `diag_type = Full`.");
tmp.AddRealComp("phi");
int const phi_index = tmp.getParticleComps().at("phi");
auto& warpx = WarpX::GetInstance();
#ifdef AMREX_USE_OMP
#pragma omp parallel if (amrex::Gpu::notInLaunchRegion())
#endif
for (int lev=0; lev<=warpx.finestLevel(); lev++) {
const amrex::Geometry& geom = warpx.Geom(lev);
auto plo = geom.ProbLoArray();
auto dxi = geom.InvCellSizeArray();
amrex::MultiFab const& phi = warpx.getField( FieldType::phi_fp, lev, 0 );
storeFieldOnParticles( tmp, plo, dxi, phi, phi_index, lev );
}
storePhiOnParticles( tmp, WarpX::electrostatic_solver_id, !use_pinned_pc );
}

// names of amrex::Real and int particle attributes in SoA data
Expand Down
45 changes: 7 additions & 38 deletions Source/Particles/ParticleIO.H
View file
Original file line number Diff line number Diff line change
Expand Up @@ -9,6 +9,7 @@

#include "Particles/WarpXParticleContainer.H"
#include "Particles/Pusher/GetAndSetPosition.H"
#include "Particles/PinnedMemoryParticleContainer.H"

#include <ablastr/particles/NodalFieldGather.H>

Expand Down Expand Up @@ -80,47 +81,15 @@ particlesConvertUnits (ConvertDirection convert_direction, T_ParticleContainer*
}
}

/** Gathers the field from a MultiFab to the macroparticles
* and stores it into a runtime component of the particle container
/** Gathers phi (electrostatic potential) from a MultiFab to the macroparticles.
* Adds a runtime component of the particle container to store it.
*
* @tparam T_ParticleContainer a WarpX particle container or AmrParticleContainer
* @param tmp the particle container on which to store the gathered field
* @param plo the coordinates of the low bound of the box, along each dimension
* @param dxi the inverse of the cell size, along each dimension
* @param phi a multifab that contains the field to be gathered
* @param phi_index the index of the runtime component where to store the gathered field
* @param lev the MR level
* @param electrostatic_solver_id the type of electrostatic solver used
* @param is_full_diagnostic whether this diagnostic is a full diagnostic
*/
template< typename T_ParticleContainer >
void
storeFieldOnParticles ( T_ParticleContainer& tmp,
amrex::GpuArray<amrex::Real,AMREX_SPACEDIM> const& plo,
amrex::GpuArray<amrex::Real,AMREX_SPACEDIM> const& dxi,
amrex::MultiFab const& phi,
int const phi_index, int const lev ) {

using PinnedParIter = typename T_ParticleContainer::ParIterType;

for (PinnedParIter pti(tmp, lev); pti.isValid(); ++pti) {

auto phi_grid = phi[pti].array();
const auto getPosition = GetParticlePosition<PIdx>(pti);
amrex::ParticleReal* phi_particle_arr = pti.GetStructOfArrays().GetRealData(phi_index).dataPtr();

// Loop over the particles and update their position
amrex::ParallelFor( pti.numParticles(),
[=] AMREX_GPU_DEVICE (long ip) {

amrex::ParticleReal xp, yp, zp;
getPosition(ip, xp, yp, zp);
int i, j, k;
amrex::Real W[AMREX_SPACEDIM][2];
ablastr::particles::compute_weights(xp, yp, zp, plo, dxi, i, j, k, W);
amrex::Real const phi_value = ablastr::particles::interp_field_nodal(i, j, k, W, phi_grid);
phi_particle_arr[ip] = phi_value;
}
);
}
}
storePhiOnParticles ( PinnedMemoryParticleContainer& tmp,
int electrostatic_solver_id, bool is_full_diagnostic );

#endif /* PARTICLEIO_H_ */

0 comments on commit 5640de3

Please sign in to comment.