1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
|
#include "ParticleReductionFunctor.H"
#include "Diagnostics/ComputeDiagFunctors/ComputeDiagFunctor.H"
#include "Particles/MultiParticleContainer.H"
#include "Particles/WarpXParticleContainer.H"
#include "Utils/CoarsenIO.H"
#include "WarpX.H"
#include <AMReX_Array.H>
#include <AMReX_BLassert.H>
#include <AMReX_IntVect.H>
#include <AMReX_MultiFab.H>
#include <AMReX_REAL.H>
using namespace amrex::literals;
ParticleReductionFunctor::ParticleReductionFunctor (const amrex::MultiFab* mf_src, const int lev,
const amrex::IntVect crse_ratio, const std::string fn_str,
const int ispec, const bool do_average,
const bool do_filter, const std::string filter_str, const int ncomp)
: ComputeDiagFunctor(ncomp, crse_ratio), m_lev(lev), m_ispec(ispec), m_do_average(do_average), m_do_filter(do_filter)
{
// mf_src will not be used, let's make sure it's null.
AMREX_ALWAYS_ASSERT(mf_src == nullptr);
// Write only in one output component.
AMREX_ALWAYS_ASSERT(ncomp == 1);
// Allocate and compile a parser based on the input string fn_str
m_map_fn_parser = std::make_unique<amrex::Parser>(makeParser(
fn_str, {"x", "y", "z", "ux", "uy", "uz"}));
m_map_fn = m_map_fn_parser->compile<6>();
// Do the same for filter function, if it exists
if (m_do_filter) {
m_filter_fn_parser = std::make_unique<amrex::Parser>(makeParser(
filter_str, {"x", "y", "z", "ux", "uy", "uz"}));
m_filter_fn = m_filter_fn_parser->compile<6>();
}
}
void
ParticleReductionFunctor::operator() (amrex::MultiFab& mf_dst, const int dcomp, const int /*i_buffer*/) const
{
auto& warpx = WarpX::GetInstance();
// Guard cell is set to 1 for generality. However, for a cell-centered
// output Multifab, mf_dst, the guard-cell data is not needed especially considering
// the operations performend in the CoarsenAndInterpolate function.
constexpr int ng = 1;
// Temporary cell-centered, multi-component MultiFab for storing particles per cell.
amrex::MultiFab red_mf(warpx.boxArray(m_lev), warpx.DistributionMap(m_lev), 1, ng);
// Set value to 0, and increment the value in each cell with ppc.
red_mf.setVal(0._rt);
auto& pc = warpx.GetPartContainer().GetParticleContainer(m_ispec);
// Copy over member variables so they can be captured in the lambda
auto map_fn = m_map_fn;
auto filter_fn = m_filter_fn;
const bool do_filter = m_do_filter;
ParticleToMesh(pc, red_mf, m_lev,
[=] AMREX_GPU_DEVICE (const WarpXParticleContainer::SuperParticleType& p,
amrex::Array4<amrex::Real> const& out_array,
amrex::GpuArray<amrex::Real,AMREX_SPACEDIM> const& plo,
amrex::GpuArray<amrex::Real,AMREX_SPACEDIM> const& dxi)
{
// Get position in WarpX convention to use in parser. Will be different from
// p.pos() for 1D and 2D simulations.
amrex::ParticleReal xw = 0._rt, yw = 0._rt, zw = 0._rt;
get_particle_position(p, xw, yw, zw);
// Get position in AMReX convention to calculate corresponding index.
// Ideally this will be replaced with the AMReX NGP interpolator
// Always do x direction. No RZ case because it's not implemented, and code
// will have aborted
int ii = 0, jj = 0, kk = 0;
amrex::ParticleReal x = p.pos(0);
amrex::Real lx = (x - plo[0]) * dxi[0];
ii = static_cast<int>(amrex::Math::floor(lx));
#if defined(WARPX_DIM_XZ) || defined(WARPX_DIM_3D)
amrex::ParticleReal y = p.pos(1);
amrex::Real ly = (y - plo[1]) * dxi[1];
jj = static_cast<int>(amrex::Math::floor(ly));
#endif
#if defined(WARPX_DIM_3D)
amrex::ParticleReal z = p.pos(2);
amrex::Real lz = (z - plo[2]) * dxi[2];
kk = static_cast<int>(amrex::Math::floor(lz));
#endif
// Fix dimensions since parser assumes u = gamma * v / c
amrex::ParticleReal ux = p.rdata(PIdx::ux) / PhysConst::c;
amrex::ParticleReal uy = p.rdata(PIdx::uy) / PhysConst::c;
amrex::ParticleReal uz = p.rdata(PIdx::uz) / PhysConst::c;
amrex::Real value;
if ((do_filter) && (!filter_fn(xw, yw, zw, ux, uy, uz))) value = 0._rt;
else value = map_fn(xw, yw, zw, ux, uy, uz);
amrex::Gpu::Atomic::AddNoRet(&out_array(ii, jj, kk, 0), p.rdata(PIdx::w) * value);
});
if (m_do_average) {
amrex::MultiFab ppc_mf(warpx.boxArray(m_lev), warpx.DistributionMap(m_lev), 1, ng);
ppc_mf.setVal(0._rt);
// Add the weight for each particle -- total number of particles of this species
ParticleToMesh(pc, ppc_mf, m_lev,
[=] AMREX_GPU_DEVICE (const WarpXParticleContainer::SuperParticleType& p,
amrex::Array4<amrex::Real> const& out_array,
amrex::GpuArray<amrex::Real,AMREX_SPACEDIM> const& plo,
amrex::GpuArray<amrex::Real,AMREX_SPACEDIM> const& dxi)
{
// Get position in WarpX convention to use in parser. Will be different from
// p.pos() for 1D and 2D simulations.
amrex::ParticleReal xw = 0._rt, yw = 0._rt, zw = 0._rt;
get_particle_position(p, xw, yw, zw);
// Get position in AMReX convention to calculate corresponding index.
// Ideally this will be replaced with the AMReX NGP interpolator
// Always do x direction. No RZ case because it's not implemented, and code
// will have aborted
int ii = 0, jj = 0, kk = 0;
amrex::ParticleReal x = p.pos(0);
amrex::Real lx = (x - plo[0]) * dxi[0];
ii = static_cast<int>(amrex::Math::floor(lx));
#if defined(WARPX_DIM_XZ) || defined(WARPX_DIM_3D)
amrex::ParticleReal y = p.pos(1);
amrex::Real ly = (y - plo[1]) * dxi[1];
jj = static_cast<int>(amrex::Math::floor(ly));
#endif
#if defined(WARPX_DIM_3D)
amrex::ParticleReal z = p.pos(2);
amrex::Real lz = (z - plo[2]) * dxi[2];
kk = static_cast<int>(amrex::Math::floor(lz));
#endif
// Fix dimensions since parser assumes u = gamma * v / c
amrex::ParticleReal ux = p.rdata(PIdx::ux) / PhysConst::c;
amrex::ParticleReal uy = p.rdata(PIdx::uy) / PhysConst::c;
amrex::ParticleReal uz = p.rdata(PIdx::uz) / PhysConst::c;
amrex::Real filter;
if ((do_filter) && (!filter_fn(xw, yw, zw, ux, uy, uz))) filter = 0._rt;
else filter = 1._rt;
amrex::Gpu::Atomic::AddNoRet(&out_array(ii, jj, kk, 0), p.rdata(PIdx::w) * filter);
});
// Divide value by number of particles for average. Set average to zero if there are no particles
for (amrex::MFIter mfi(red_mf, amrex::TilingIfNotGPU()); mfi.isValid(); ++mfi)
{
const amrex::Box& box = mfi.tilebox();
amrex::Array4<amrex::Real> const& a_red = red_mf.array(mfi);
amrex::Array4<amrex::Real const> const& a_ppc = ppc_mf.const_array(mfi);
amrex::ParallelFor(box,
[=] AMREX_GPU_DEVICE (int i, int j, int k) {
if (a_ppc(i,j,k,0) == 0) a_red(i,j,k,0) = 0;
else a_red(i,j,k,0) = a_red(i,j,k,0) / a_ppc(i,j,k,0);
});
}
}
// Coarsen and interpolate from ppc_mf to the output diagnostic MultiFab, mf_dst.
CoarsenIO::Coarsen(mf_dst, red_mf, dcomp, 0, nComp(), 0, m_crse_ratio);
}
|
