Skip to content

Commit

Permalink
Clean up of ParticleExtrema
Browse files Browse the repository at this point in the history
  • Loading branch information
dpgrote committed Jun 5, 2024
1 parent 1d22597 commit 641ef52
Showing 1 changed file with 64 additions and 173 deletions.
237 changes: 64 additions & 173 deletions Source/Diagnostics/ReducedDiags/ParticleExtrema.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -51,15 +51,15 @@
#include <map>
#include <vector>

using namespace amrex;
using namespace amrex::literals;
using namespace warpx::fields;

// constructor
ParticleExtrema::ParticleExtrema (const std::string& rd_name)
: ReducedDiags{rd_name}
{
// read species name
const ParmParse pp_rd_name(rd_name);
const amrex::ParmParse pp_rd_name(rd_name);
pp_rd_name.get("species",m_species_name);

// get WarpX class object
Expand Down Expand Up @@ -122,7 +122,7 @@ ParticleExtrema::ParticleExtrema (const std::string& rd_name)

m_data.resize(all_diag_names.size());

if (ParallelDescriptor::IOProcessor())
if (amrex::ParallelDescriptor::IOProcessor())
{
if ( m_write_header )
{
Expand Down Expand Up @@ -165,16 +165,7 @@ void ParticleExtrema::ComputeDiags (int step)
const auto species_names = mypc.GetSpeciesNames();

// inverse of speed of light squared
Real constexpr inv_c2 = 1.0_rt / (PhysConst::c * PhysConst::c);

// If 2D-XZ, p.pos(1) is z, rather than p.pos(2).
#if (defined WARPX_DIM_3D)
int const index_z = 2;
#elif (defined WARPX_DIM_XZ || defined WARPX_DIM_RZ)
int const index_z = 1;
#elif (defined WARPX_DIM_1D_Z)
int const index_z = 0;
#endif
amrex::Real constexpr inv_c2 = 1.0_rt / (PhysConst::c * PhysConst::c);

// loop over species
for (int i_s = 0; i_s < nSpecies; ++i_s)
Expand All @@ -193,172 +184,72 @@ void ParticleExtrema::ComputeDiags (int step)
}

using PType = typename WarpXParticleContainer::SuperParticleType;

// xmin
#if (defined WARPX_DIM_RZ)
Real xmin = ReduceMin( myspc,
[=] AMREX_GPU_HOST_DEVICE (const PType& p)
{ return p.pos(0)*std::cos(p.rdata(PIdx::theta)); });
#elif (defined WARPX_DIM_1D_Z)
Real xmin = 0.0_rt;
#else
Real xmin = ReduceMin( myspc,
[=] AMREX_GPU_HOST_DEVICE (const PType& p)
{ return p.pos(0); });
#endif

// xmax
#if (defined WARPX_DIM_RZ)
Real xmax = ReduceMax( myspc,
[=] AMREX_GPU_HOST_DEVICE (const PType& p)
{ return p.pos(0)*std::cos(p.rdata(PIdx::theta)); });
#elif (defined WARPX_DIM_1D_Z)
Real xmax = 0.0_rt;
#else
Real xmax = ReduceMax( myspc,
[=] AMREX_GPU_HOST_DEVICE (const PType& p)
{ return p.pos(0); });
#endif

// ymin
#if (defined WARPX_DIM_RZ)
Real ymin = ReduceMin( myspc,
[=] AMREX_GPU_HOST_DEVICE (const PType& p)
{ return p.pos(0)*std::sin(p.rdata(PIdx::theta)); });
#elif (defined WARPX_DIM_XZ || WARPX_DIM_1D_Z)
Real ymin = 0.0_rt;
#else
Real ymin = ReduceMin( myspc,
[=] AMREX_GPU_HOST_DEVICE (const PType& p)
{ return p.pos(1); });
#endif

// ymax
#if (defined WARPX_DIM_RZ)
Real ymax = ReduceMax( myspc,
[=] AMREX_GPU_HOST_DEVICE (const PType& p)
{ return p.pos(0)*std::sin(p.rdata(PIdx::theta)); });
#elif (defined WARPX_DIM_XZ || WARPX_DIM_1D_Z)
Real ymax = 0.0_rt;
#else
Real ymax = ReduceMax( myspc,
[=] AMREX_GPU_HOST_DEVICE (const PType& p)
{ return p.pos(1); });
#endif

// zmin
Real zmin = ReduceMin( myspc,
[=] AMREX_GPU_HOST_DEVICE (const PType& p)
{ return p.pos(index_z); });

// zmax
Real zmax = ReduceMax( myspc,
[=] AMREX_GPU_HOST_DEVICE (const PType& p)
{ return p.pos(index_z); });

// uxmin
Real uxmin = ReduceMin( myspc,
[=] AMREX_GPU_HOST_DEVICE (const PType& p)
{ return p.rdata(PIdx::ux); });

// uxmax
Real uxmax = ReduceMax( myspc,
[=] AMREX_GPU_HOST_DEVICE (const PType& p)
{ return p.rdata(PIdx::ux); });

// uymin
Real uymin = ReduceMin( myspc,
[=] AMREX_GPU_HOST_DEVICE (const PType& p)
{ return p.rdata(PIdx::uy); });

// uymax
Real uymax = ReduceMax( myspc,
[=] AMREX_GPU_HOST_DEVICE (const PType& p)
{ return p.rdata(PIdx::uy); });

// uzmin
Real uzmin = ReduceMin( myspc,
[=] AMREX_GPU_HOST_DEVICE (const PType& p)
{ return p.rdata(PIdx::uz); });

// uzmax
Real uzmax = ReduceMax( myspc,
[=] AMREX_GPU_HOST_DEVICE (const PType& p)
{ return p.rdata(PIdx::uz); });

// gmin
Real gmin = 0.0_rt;
if ( is_photon ) {
gmin = ReduceMin( myspc,
[=] AMREX_GPU_HOST_DEVICE (const PType& p)
using OpMin = amrex::ReduceOpMin;
using OpMax = amrex::ReduceOpMax;

amrex::ReduceOps<OpMin, OpMin, OpMin, OpMin,
OpMax, OpMax, OpMax, OpMax> reduce_ops;
auto posminmax = amrex::ParticleReduce<amrex::ReduceData<amrex::Real, amrex::Real, amrex::Real, amrex::Real,
amrex::Real, amrex::Real, amrex::Real, amrex::Real>>(
myspc,
[=] AMREX_GPU_DEVICE(const PType& p) noexcept -> amrex::GpuTuple<amrex::Real, amrex::Real, amrex::Real, amrex::Real,
amrex::Real, amrex::Real, amrex::Real, amrex::Real>
{
const Real ux = p.rdata(PIdx::ux);
const Real uy = p.rdata(PIdx::uy);
const Real uz = p.rdata(PIdx::uz);
const Real us = ux*ux + uy*uy + uz*uz;
return std::sqrt(us*inv_c2);
});
} else {
gmin = ReduceMin( myspc,
[=] AMREX_GPU_HOST_DEVICE (const PType& p)
amrex::ParticleReal x, y, z;
get_particle_position(p, x, y, z);
amrex::Real const w = p.rdata(PIdx::w);
return {w, x, y, z, w, x, y, z};
},
reduce_ops);

amrex::Real wmin = amrex::get<0>(posminmax);
amrex::Real xmin = amrex::get<1>(posminmax);
amrex::Real ymin = amrex::get<2>(posminmax);
amrex::Real zmin = amrex::get<3>(posminmax);
amrex::Real wmax = amrex::get<4>(posminmax);
amrex::Real xmax = amrex::get<5>(posminmax);
amrex::Real ymax = amrex::get<6>(posminmax);
amrex::Real zmax = amrex::get<7>(posminmax);

amrex::Real gfactor = (is_photon ? 0._rt : 1._rt);
auto uminmax = amrex::ParticleReduce<amrex::ReduceData<amrex::Real, amrex::Real, amrex::Real, amrex::Real,
amrex::Real, amrex::Real, amrex::Real, amrex::Real>>(
myspc,
[=] AMREX_GPU_DEVICE(const PType& p) noexcept -> amrex::GpuTuple<amrex::Real, amrex::Real, amrex::Real, amrex::Real,
amrex::Real, amrex::Real, amrex::Real, amrex::Real>
{
const Real ux = p.rdata(PIdx::ux);
const Real uy = p.rdata(PIdx::uy);
const Real uz = p.rdata(PIdx::uz);
const Real us = ux*ux + uy*uy + uz*uz;
return std::sqrt(1.0_rt + us*inv_c2);
});
}

// gmax
Real gmax = 0.0_rt;
if ( is_photon ) {
gmax = ReduceMax( myspc,
[=] AMREX_GPU_HOST_DEVICE (const PType& p)
{
const Real ux = p.rdata(PIdx::ux);
const Real uy = p.rdata(PIdx::uy);
const Real uz = p.rdata(PIdx::uz);
const Real us = ux*ux + uy*uy + uz*uz;
return std::sqrt(us*inv_c2);
});
} else {
gmax = ReduceMax( myspc,
[=] AMREX_GPU_HOST_DEVICE (const PType& p)
{
const Real ux = p.rdata(PIdx::ux);
const Real uy = p.rdata(PIdx::uy);
const Real uz = p.rdata(PIdx::uz);
const Real us = ux*ux + uy*uy + uz*uz;
return std::sqrt(1.0_rt + us*inv_c2);
});
}

// wmin
Real wmin = ReduceMin( myspc,
[=] AMREX_GPU_HOST_DEVICE (const PType& p)
{ return p.rdata(PIdx::w); });

// wmax
Real wmax = ReduceMax( myspc,
[=] AMREX_GPU_HOST_DEVICE (const PType& p)
{ return p.rdata(PIdx::w); });

ParallelDescriptor::ReduceRealMin({xmin,ymin,zmin,uxmin,uymin,uzmin,gmin,wmin});
ParallelDescriptor::ReduceRealMax({xmax,ymax,zmax,uxmax,uymax,uzmax,gmax,wmax});
amrex::Real const ux = p.rdata(PIdx::ux);
amrex::Real const uy = p.rdata(PIdx::uy);
amrex::Real const uz = p.rdata(PIdx::uz);
amrex::Real const g = std::sqrt(gfactor + (ux*ux + uy*uy + uz*uz)*inv_c2);
return {g, ux, uy, uz, g, ux, uy, uz};
},
reduce_ops);

amrex::Real gmin = amrex::get<0>(uminmax);
amrex::Real uxmin = amrex::get<1>(uminmax);
amrex::Real uymin = amrex::get<2>(uminmax);
amrex::Real uzmin = amrex::get<3>(uminmax);
amrex::Real gmax = amrex::get<4>(uminmax);
amrex::Real uxmax = amrex::get<5>(uminmax);
amrex::Real uymax = amrex::get<6>(uminmax);
amrex::Real uzmax = amrex::get<7>(uminmax);

amrex::ParallelDescriptor::ReduceRealMin({xmin,ymin,zmin,uxmin,uymin,uzmin,gmin,wmin});
amrex::ParallelDescriptor::ReduceRealMax({xmax,ymax,zmax,uxmax,uymax,uzmax,gmax,wmax});

#if (defined WARPX_QED)
// get number of level (int)
const auto level_number = WarpX::GetInstance().finestLevel();

// compute chimin and chimax
Real chimin_f = 0.0_rt;
Real chimax_f = 0.0_rt;
amrex::Real chimin_f = 0.0_rt;
amrex::Real chimax_f = 0.0_rt;

if (myspc.DoQED())
{
// declare chi arrays
std::vector<Real> chimin, chimax;
std::vector<amrex::Real> chimin, chimax;
chimin.resize(level_number+1,0.0_rt);
chimax.resize(level_number+1,0.0_rt);

Expand All @@ -383,8 +274,8 @@ void ParticleExtrema::ComputeDiags (int step)
const MultiFab & Bz = warpx.getField(FieldType::Bfield_aux, lev,2);

// declare reduce_op
ReduceOps<ReduceOpMin, ReduceOpMax> reduce_op;
ReduceData<Real, Real> reduce_data(reduce_op);
amrex::ReduceOps<amrex::ReduceOpMin, amrex::ReduceOpMax> reduce_op;
amrex::ReduceData<amrex::Real, amrex::Real> reduce_data(reduce_op);
using ReduceTuple = typename decltype(reduce_data)::Type;

// Loop over boxes
Expand Down Expand Up @@ -429,7 +320,7 @@ void ParticleExtrema::ComputeDiags (int step)
[=] AMREX_GPU_DEVICE (int i) -> ReduceTuple
{
// get external fields
ParticleReal xp, yp, zp;
amrex::ParticleReal xp, yp, zp;
GetPosition(i, xp, yp, zp);
amrex::ParticleReal ex = Ex_external_particle;
amrex::ParticleReal ey = Ey_external_particle;
Expand All @@ -449,7 +340,7 @@ void ParticleExtrema::ComputeDiags (int step)
dx_arr, xyzmin_arr, lo,
n_rz_azimuthal_modes, nox, galerkin_interpolation);
// compute chi
Real chi = 0.0_rt;
amrex::Real chi = 0.0_rt;
if ( is_photon ) {
chi = QedUtils::chi_photon(ux[i]*m, uy[i]*m, uz[i]*m,
ex, ey, ez, bx, by, bz);
Expand All @@ -466,8 +357,8 @@ void ParticleExtrema::ComputeDiags (int step)
}
chimin_f = *std::min_element(chimin.begin(), chimin.end());
chimax_f = *std::max_element(chimax.begin(), chimax.end());
ParallelDescriptor::ReduceRealMin(chimin_f, ParallelDescriptor::IOProcessorNumber());
ParallelDescriptor::ReduceRealMax(chimax_f, ParallelDescriptor::IOProcessorNumber());
amrex::ParallelDescriptor::ReduceRealMin(chimin_f, amrex::ParallelDescriptor::IOProcessorNumber());
amrex::ParallelDescriptor::ReduceRealMax(chimax_f, amrex::ParallelDescriptor::IOProcessorNumber());
}
#endif

Expand Down

0 comments on commit 641ef52

Please sign in to comment.