Skip to content

Commit 7608fb3

Browse files
ax3lax3l
committed
Spin Distribution Call
1 parent b1781be commit 7608fb3

File tree

8 files changed

+101
-60
lines changed

8 files changed

+101
-60
lines changed

‎src/ImpactX.H‎

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -85,12 +85,14 @@ namespace impactx
8585
* @param bunch_charge bunch charge (C)
8686
* @param distr distribution function to draw from (object)
8787
* @param npart number of particles to draw
88+
* @param spin_distr optional spin distribution
8889
*/
8990
void
9091
add_particles (
9192
amrex::ParticleReal bunch_charge,
9293
distribution::KnownDistributions distr,
93-
amrex::Long npart
94+
amrex::Long npart,
95+
std::optional<distribution::SpinvMF> spin_distr = std::nullopt
9496
);
9597

9698
/** Validate the simulation is ready to run the particle tracking loop via @see track_particles

‎src/initialization/InitDistribution.cpp‎

Lines changed: 43 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -14,6 +14,7 @@
1414
#include "particles/CovarianceMatrix.H"
1515
#include "particles/ImpactXParticleContainer.H"
1616
#include "particles/distribution/All.H"
17+
#include "particles/distribution/SpinvMF.H"
1718
#include "particles/SplitEqually.H"
1819

1920
#include <ablastr/constant.H>
@@ -295,7 +296,8 @@ namespace impactx
295296
ImpactX::add_particles (
296297
amrex::ParticleReal bunch_charge,
297298
distribution::KnownDistributions distr,
298-
amrex::Long npart
299+
amrex::Long npart,
300+
std::optional<distribution::SpinvMF> spin_distr
299301
)
300302
{
301303
BL_PROFILE("ImpactX::add_particles");
@@ -337,13 +339,21 @@ namespace impactx
337339
// alloc data for particle attributes
338340
amrex::Gpu::DeviceVector<amrex::ParticleReal> x, y, t;
339341
amrex::Gpu::DeviceVector<amrex::ParticleReal> px, py, pt;
342+
amrex::Gpu::DeviceVector<amrex::ParticleReal> sx, sy, sz;
340343
x.resize(npart_this_proc);
341344
y.resize(npart_this_proc);
342345
t.resize(npart_this_proc);
343346
px.resize(npart_this_proc);
344347
py.resize(npart_this_proc);
345348
pt.resize(npart_this_proc);
346349

350+
bool const has_spin = spin_distr.has_value();
351+
if (has_spin) {
352+
sx.resize(npart_this_proc);
353+
sy.resize(npart_this_proc);
354+
sz.resize(npart_this_proc);
355+
}
356+
347357
std::visit([&](auto&& distribution){
348358
// initialize distributions
349359
distribution.initialize(bunch_charge, ref);
@@ -354,6 +364,9 @@ namespace impactx
354364
amrex::ParticleReal * const AMREX_RESTRICT px_ptr = px.data();
355365
amrex::ParticleReal * const AMREX_RESTRICT py_ptr = py.data();
356366
amrex::ParticleReal * const AMREX_RESTRICT pt_ptr = pt.data();
367+
amrex::ParticleReal * const AMREX_RESTRICT sx_ptr = sx.data();
368+
amrex::ParticleReal * const AMREX_RESTRICT sy_ptr = sy.data();
369+
amrex::ParticleReal * const AMREX_RESTRICT sz_ptr = sz.data();
357370

358371
using Distribution = std::decay_t<decltype(distribution)>;
359372

@@ -373,6 +386,7 @@ namespace impactx
373386
npart_this_thread = thread_chunk.size;
374387
#endif
375388

389+
// phase space init
376390
initialization::InitSingleParticleData<Distribution> const init_single_particle_data(
377391
distribution,
378392
x_ptr + my_offset,
@@ -384,6 +398,20 @@ namespace impactx
384398
);
385399

386400
amrex::ParallelForRNG(npart_this_thread, init_single_particle_data);
401+
402+
// spin init
403+
if (has_spin) {
404+
auto spinv = spin_distr.value();
405+
amrex::ParallelForRNG(npart_this_thread,
406+
[=] AMREX_GPU_DEVICE (int i, amrex::RandomEngine const& engine) noexcept {
407+
spinv(
408+
sx_ptr[i + my_offset],
409+
sy_ptr[i + my_offset],
410+
sz_ptr[i + my_offset],
411+
engine
412+
);
413+
});
414+
}
387415
}
388416

389417
// finalize distributions and deallocate temporary device global memory
@@ -548,15 +576,28 @@ namespace impactx
548576
std::string unit_type; // System of units
549577
pp_dist.get("units", unit_type);
550578

579+
// phase space distribution
551580
distribution::KnownDistributions dist = initialization::read_distribution(pp_dist);
552581
std::string distribution;
553582
pp_dist.get("distribution", distribution);
554583

584+
// spin distribution
585+
amrex::ParticleReal polarization_x = 0.0_prt,
586+
polarization_y = 0.0_prt,
587+
polarization_z = 0.0_prt;
588+
pp_dist.getWithParser("polarization_x", polarization_x);
589+
pp_dist.getWithParser("polarization_y", polarization_y);
590+
pp_dist.getWithParser("polarization_z", polarization_z);
591+
// TODO: check magnitude of x,y,z is in range [0:1]
592+
// TODO: calculate kappa from mag(x,y,z)
593+
amrex::ParticleReal kappa = 1.0; // FIXME: root finding
594+
distribution::SpinvMF spin_dist(polarization_x, polarization_y, polarization_z, kappa);
595+
555596
amrex::Long npart = 0; // Number of simulation particles
556597
if (distribution != "empty")
557598
{
558599
pp_dist.getWithParser("npart", npart);
559-
add_particles(bunch_charge, dist, npart);
600+
add_particles(bunch_charge, dist, npart, spin_dist);
560601
}
561602

562603
// print information on the initialized beam

‎src/particles/ImpactXParticleContainer.H‎

Lines changed: 7 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -179,12 +179,12 @@ namespace impactx
179179
* @param px momentum in x
180180
* @param py momentum in y
181181
* @param pt momentum in t
182-
* @param sx spin component in x
183-
* @param sy spin component in y
184-
* @param sz spin component in z
185182
* @param qm charge over mass in 1/eV
186183
* @param bunch_charge total charge within a bunch in C
187184
* @param w weight of each particle: how many real particles to represent
185+
* @param sx spin component in x
186+
* @param sy spin component in y
187+
* @param sz spin component in z
188188
*/
189189
void
190190
AddNParticles (
@@ -194,12 +194,12 @@ namespace impactx
194194
amrex::Gpu::DeviceVector<amrex::ParticleReal> const & px,
195195
amrex::Gpu::DeviceVector<amrex::ParticleReal> const & py,
196196
amrex::Gpu::DeviceVector<amrex::ParticleReal> const & pt,
197-
amrex::Gpu::DeviceVector<amrex::ParticleReal> const & sx,
198-
amrex::Gpu::DeviceVector<amrex::ParticleReal> const & sy,
199-
amrex::Gpu::DeviceVector<amrex::ParticleReal> const & sz,
200197
amrex::ParticleReal qm,
201198
std::optional<amrex::ParticleReal> bunch_charge = std::nullopt,
202-
std::optional<amrex::Gpu::DeviceVector<amrex::ParticleReal>> w = std::nullopt
199+
std::optional<amrex::Gpu::DeviceVector<amrex::ParticleReal>> w = std::nullopt,
200+
std::optional<amrex::Gpu::DeviceVector<amrex::ParticleReal>> sx = std::nullopt,
201+
std::optional<amrex::Gpu::DeviceVector<amrex::ParticleReal>> sy = std::nullopt,
202+
std::optional<amrex::Gpu::DeviceVector<amrex::ParticleReal>> sz = std::nullopt
203203
);
204204

205205
/** Register storage for lost particles

‎src/particles/ImpactXParticleContainer.cpp‎

Lines changed: 39 additions & 26 deletions
Original file line numberDiff line numberDiff line change
@@ -240,36 +240,42 @@ namespace impactx
240240
amrex::Gpu::DeviceVector<amrex::ParticleReal> const & px,
241241
amrex::Gpu::DeviceVector<amrex::ParticleReal> const & py,
242242
amrex::Gpu::DeviceVector<amrex::ParticleReal> const & pt,
243-
amrex::Gpu::DeviceVector<amrex::ParticleReal> const & sx,
244-
amrex::Gpu::DeviceVector<amrex::ParticleReal> const & sy,
245-
amrex::Gpu::DeviceVector<amrex::ParticleReal> const & sz,
246243
amrex::ParticleReal qm,
247244
std::optional<amrex::ParticleReal> bunch_charge,
248-
std::optional<amrex::Gpu::DeviceVector<amrex::ParticleReal>> w
245+
std::optional<amrex::Gpu::DeviceVector<amrex::ParticleReal>> w,
246+
std::optional<amrex::Gpu::DeviceVector<amrex::ParticleReal>> sx,
247+
std::optional<amrex::Gpu::DeviceVector<amrex::ParticleReal>> sy,
248+
std::optional<amrex::Gpu::DeviceVector<amrex::ParticleReal>> sz
249249
)
250250
{
251251
BL_PROFILE("ImpactX::AddNParticles");
252252

253253
using namespace amrex::literals; // for _rt and _prt
254254

255+
// number of particles to add
256+
std::size_t const np_s = x.size();
257+
258+
// input validation
255259
bool const has_w = w.has_value();
256260
if (!(bunch_charge.has_value() ^ has_w))
257261
{
258262
throw std::runtime_error("AddNParticles: Exactly one of bunch_charge or w must be provided!");
259263
}
260264

261-
AMREX_ALWAYS_ASSERT(x.size() == y.size());
262-
AMREX_ALWAYS_ASSERT(x.size() == t.size());
263-
AMREX_ALWAYS_ASSERT(x.size() == px.size());
264-
AMREX_ALWAYS_ASSERT(x.size() == py.size());
265-
AMREX_ALWAYS_ASSERT(x.size() == pt.size());
266-
AMREX_ALWAYS_ASSERT(x.size() == sx.size());
267-
AMREX_ALWAYS_ASSERT(x.size() == sy.size());
268-
AMREX_ALWAYS_ASSERT(x.size() == sz.size());
269-
if (has_w) { AMREX_ALWAYS_ASSERT(x.size() == w->size()); }
270-
271-
// number of particles to add
272-
amrex::Long const np = x.size();
265+
bool const has_spin = sx.has_value();
266+
267+
AMREX_ALWAYS_ASSERT(np_s == y.size());
268+
AMREX_ALWAYS_ASSERT(np_s == t.size());
269+
AMREX_ALWAYS_ASSERT(np_s == px.size());
270+
AMREX_ALWAYS_ASSERT(np_s == py.size());
271+
AMREX_ALWAYS_ASSERT(np_s == pt.size());
272+
if (has_spin) {
273+
AMREX_ALWAYS_ASSERT(sy.has_value());
274+
AMREX_ALWAYS_ASSERT(sz.has_value());
275+
AMREX_ALWAYS_ASSERT(np_s == sy.value().size());
276+
AMREX_ALWAYS_ASSERT(np_s == sz.value().size());
277+
}
278+
if (has_w) { AMREX_ALWAYS_ASSERT(np_s == w->size()); }
273279

274280
// we add particles to lev 0, grid 0
275281
int lid = 0, gid = 0;
@@ -294,7 +300,8 @@ namespace impactx
294300
DefineAndReturnParticleTile(lid, gid, ithr);
295301
}
296302

297-
amrex::Long pid = ParticleType::NextID();
303+
amrex::Long const pid = ParticleType::NextID();
304+
amrex::Long const np = np_s;
298305
ParticleType::NextID(pid + np);
299306
AMREX_ALWAYS_ASSERT_WITH_MESSAGE(
300307
pid + np < amrex::LongParticleIds::LastParticleID,
@@ -337,9 +344,9 @@ namespace impactx
337344
amrex::ParticleReal * const AMREX_RESTRICT px_arr = soa[RealSoA::px].dataPtr();
338345
amrex::ParticleReal * const AMREX_RESTRICT py_arr = soa[RealSoA::py].dataPtr();
339346
amrex::ParticleReal * const AMREX_RESTRICT pt_arr = soa[RealSoA::pt].dataPtr();
340-
amrex::ParticleReal * const AMREX_RESTRICT sx_arr = soa[RealSoA::sx].dataPtr();
341-
amrex::ParticleReal * const AMREX_RESTRICT sy_arr = soa[RealSoA::sy].dataPtr();
342-
amrex::ParticleReal * const AMREX_RESTRICT sz_arr = soa[RealSoA::sz].dataPtr();
347+
amrex::ParticleReal * const AMREX_RESTRICT sx_arr = has_spin ? soa[RealSoA::sx].dataPtr() : nullptr;
348+
amrex::ParticleReal * const AMREX_RESTRICT sy_arr = has_spin ? soa[RealSoA::sy].dataPtr() : nullptr;
349+
amrex::ParticleReal * const AMREX_RESTRICT sz_arr = has_spin ? soa[RealSoA::sz].dataPtr() : nullptr;
343350
amrex::ParticleReal * const AMREX_RESTRICT qm_arr = soa[RealSoA::qm].dataPtr();
344351
amrex::ParticleReal * const AMREX_RESTRICT w_arr = soa[RealSoA::w ].dataPtr();
345352

@@ -351,9 +358,9 @@ namespace impactx
351358
amrex::ParticleReal const * const AMREX_RESTRICT px_ptr = px.data();
352359
amrex::ParticleReal const * const AMREX_RESTRICT py_ptr = py.data();
353360
amrex::ParticleReal const * const AMREX_RESTRICT pt_ptr = pt.data();
354-
amrex::ParticleReal const * const AMREX_RESTRICT sx_ptr = sx.data();
355-
amrex::ParticleReal const * const AMREX_RESTRICT sy_ptr = sy.data();
356-
amrex::ParticleReal const * const AMREX_RESTRICT sz_ptr = sz.data();
361+
amrex::ParticleReal const * const AMREX_RESTRICT sx_ptr = has_spin ? sx.value().data() : nullptr;
362+
amrex::ParticleReal const * const AMREX_RESTRICT sy_ptr = has_spin ? sy.value().data() : nullptr;
363+
amrex::ParticleReal const * const AMREX_RESTRICT sz_ptr = has_spin ? sz.value().data() : nullptr;
357364
amrex::ParticleReal const * const AMREX_RESTRICT w_ptr = has_w ? w->data() : nullptr;
358365
amrex::ParticleReal const bunch_charge_value = has_w ? 0_prt : bunch_charge.value();
359366

@@ -370,9 +377,15 @@ namespace impactx
370377
py_arr[old_np+i] = py_ptr[my_offset+i];
371378
pt_arr[old_np+i] = pt_ptr[my_offset+i];
372379

373-
sx_arr[old_np+i] = sx_ptr[my_offset+i];
374-
sy_arr[old_np+i] = sy_ptr[my_offset+i];
375-
sz_arr[old_np+i] = sz_ptr[my_offset+i];
380+
if (has_spin) {
381+
sx_arr[old_np+i] = sx_ptr[my_offset+i];
382+
sy_arr[old_np+i] = sy_ptr[my_offset+i];
383+
sz_arr[old_np+i] = sz_ptr[my_offset+i];
384+
} else {
385+
sx_arr[old_np+i] = 0_prt;
386+
sy_arr[old_np+i] = 0_prt;
387+
sz_arr[old_np+i] = 0_prt;
388+
}
376389

377390
qm_arr[old_np+i] = qm;
378391
w_arr[old_np+i] = has_w ? w_ptr[my_offset+i] : bunch_charge_value / ablastr::constant::SI::q_e/np;

‎src/particles/distribution/All.H‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -26,6 +26,7 @@
2626

2727
namespace impactx::distribution
2828
{
29+
// Phase space distributions
2930
using KnownDistributions = std::variant<
3031
Empty, /* must be first, so KnownDistributions creates a default constructor */
3132
Gaussian,
@@ -35,8 +36,7 @@ namespace impactx::distribution
3536
Thermal,
3637
Triangle,
3738
Semigaussian,
38-
Waterbag,
39-
SpinvMF
39+
Waterbag
4040
>;
4141

4242
} // namespace impactx::distribution

‎src/particles/distribution/SpinvMF.H‎

Lines changed: 2 additions & 22 deletions
Original file line numberDiff line numberDiff line change
@@ -51,26 +51,6 @@ namespace impactx::distribution
5151
{
5252
}
5353

54-
/** Initialize the distribution.
55-
*
56-
* Nothing to do here.
57-
*
58-
* @param bunch_charge charge of the beam in C
59-
* @param ref the reference particle
60-
*/
61-
void initialize ([[maybe_unused]] amrex::ParticleReal bunch_charge, [[maybe_unused]] RefPart const & ref)
62-
{
63-
}
64-
65-
/** Close and deallocate all data and handles.
66-
*
67-
* Nothing to do here.
68-
*/
69-
void
70-
finalize ()
71-
{
72-
}
73-
7454
/** Return 1 3D spin (unit) vector
7555
*
7656
* @param sx particle spin component in x
@@ -113,13 +93,13 @@ namespace impactx::distribution
11393
// D. Frisch and U. D. Hanebeck, "Deterministic Von Mises-Fisher Sampling on the Sphere Using
11494
// Fibonacci Lattices", IEEE, 2023, DOI: 10.1109/SDF-MFI59545.2023.10361396
11595
if(m_kappa > 0) {
116-
t = 1_prt + std::log1p(x2 * std::expm1(-2*m_kappa))/m_kappa;
96+
t = 1_prt + std::log1p(x2 * std::expm1(-2_prt * m_kappa)) / m_kappa;
11797
} else {
11898
t = 1_prt - 2_prt * x2;
11999
}
120100

121101
// Evaluation of the spin components
122-
amrex::ParticleReal u=std::sqrt(1_prt-powi<2>(t));
102+
amrex::ParticleReal u = std::sqrt(1_prt-powi<2>(t));
123103
sx = u * (b*c + a*muz*s) + t*mux;
124104
sy = u * (-a*c + b*muz*s) + t*muy;
125105
sz = u * (-muperp*s) + t*muz;

‎src/python/ImpactX.cpp‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -607,6 +607,7 @@ void init_ImpactX (py::module& m)
607607
.def("add_particles", &ImpactX::add_particles,
608608
py::arg("bunch_charge"),
609609
py::arg("distr"), py::arg("npart"),
610+
py::arg("spin_distr") = py::none(), // TODO: bind the distribution::SpinvMF object to Python
610611
"Particle tracking mode:"
611612
"Generate and add n particles to the particle container.\n\n"
612613
"Will also resize the geometry based on the updated particle\n"

‎src/python/ImpactXParticleContainer.cpp‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -84,6 +84,7 @@ void init_impactxparticlecontainer(py::module& m)
8484
py::arg("x"), py::arg("y"), py::arg("t"),
8585
py::arg("px"), py::arg("py"), py::arg("pt"),
8686
py::arg("qm"), py::arg("bunch_charge")=py::none(), py::arg("w")=py::none(),
87+
py::arg("sx")=py::none(), py::arg("sy")=py::none(), py::arg("sz")=py::none(),
8788
"Add new particles to the container for fixed s.\n\n"
8889
"Either the total charge (bunch_charge) or the weight of each\n"
8990
"particle (w) must be provided.\n\n"
@@ -99,6 +100,9 @@ void init_impactxparticlecontainer(py::module& m)
99100
":param qm: charge over mass in 1/eV\n"
100101
":param bunch_charge: total charge within a bunch in C"
101102
":param w: weight of each particle: how many real particles to represent"
103+
":param sx: spin component in x\n"
104+
":param sy: spin component in y\n"
105+
":param sz: spin component in z\n"
102106
)
103107
.def("ref_particle",
104108
py::overload_cast<>(&ImpactXParticleContainer::GetRefParticle),

0 commit comments

Comments
 (0)