diff --git a/examples/fodo/run_fodo.py b/examples/fodo/run_fodo.py index fc48a3d2c..13c3dabb7 100755 --- a/examples/fodo/run_fodo.py +++ b/examples/fodo/run_fodo.py @@ -32,8 +32,7 @@ muxpx = -0.846574929020762, muypy = 0.846574929020762, mutpt = 0.0) -distribution.generate_add_particles( - impactX.particle_container(), qm_qeeV, charge_C, distr, npart) +impactX.add_particles(qm_qeeV, charge_C, distr, npart) # init reference particle refPart = RefPart() diff --git a/src/ImpactX.H b/src/ImpactX.H index 8c10182e4..3739cfc17 100644 --- a/src/ImpactX.H +++ b/src/ImpactX.H @@ -10,6 +10,7 @@ #ifndef IMPACT_X_H #define IMPACT_X_H +#include "particles/distribution/All.H" #include "particles/elements/All.H" #include "particles/ImpactXParticleContainer.H" @@ -63,6 +64,25 @@ namespace impactx */ void initLatticeElementsFromInputs (); + /** Generate and add n particles to the particle container + * + * Will also resize the geometry based on the updated particle + * distribution's extent and then redistribute particles in according + * AMReX grid boxes. + * + * @param qm charge/mass ratio (q_e/eV) + * @param bunch_charge bunch charge (C) + * @param distr distribution function to draw from (object) + * @param npart number of particles to draw + */ + void + add_particles ( + amrex::ParticleReal qm, + amrex::ParticleReal bunch_charge, + distribution::KnownDistributions distr, + int npart + ); + /** Run the main simulation loop for a number of steps */ void evolve (); diff --git a/src/initialization/InitDistribution.H b/src/initialization/InitDistribution.H deleted file mode 100644 index 644f38104..000000000 --- a/src/initialization/InitDistribution.H +++ /dev/null @@ -1,79 +0,0 @@ -/* Copyright 2022 The Regents of the University of California, through Lawrence - * Berkeley National Laboratory (subject to receipt of any required - * approvals from the U.S. Dept. of Energy). All rights reserved. - * - * This file is part of ImpactX. - * - * Authors: Axel Huebl - * License: BSD-3-Clause-LBNL - */ -#include "particles/ImpactXParticleContainer.H" - -#include -#include -#include -#include -#include -#include - -#include - - -namespace impactx -{ - /** Generate and add n particles to a particle container - * - * @tparam T_Distribution distribution function to draw from (type) - * @param pc a particle container to add particles to - * @param qm charge/mass ratio (q_e/eV) - * @param bunch_charge bunch charge (C) - * @param distr distribution function to draw from (object) - * @param npart number of particles to draw - */ - template - void - generate_add_particles ( - impactx::ImpactXParticleContainer & pc, - amrex::ParticleReal qm, - amrex::ParticleReal bunch_charge, - T_Distribution & distr, - int npart - ) - { - amrex::Vector x, y, t; - amrex::Vector px, py, pt; - amrex::ParticleReal ix, iy, it, ipx, ipy, ipt; - amrex::RandomEngine rng; - - // Logic: We initialize 1/Nth of particles, independent of their - // position, per MPI rank. We then measure the distribution's spatial - // extent, create a grid, resize it to fit the beam, and then - // redistribute particles so that they reside on the correct MPI rank. - int myproc = amrex::ParallelDescriptor::MyProc(); - int nprocs = amrex::ParallelDescriptor::NProcs(); - int navg = npart / nprocs; - int nleft = npart - navg * nprocs; - int npart_this_proc = (myproc < nleft) ? navg+1 : navg; - - x.reserve(npart_this_proc); - y.reserve(npart_this_proc); - t.reserve(npart_this_proc); - px.reserve(npart_this_proc); - py.reserve(npart_this_proc); - pt.reserve(npart_this_proc); - - for (amrex::Long i = 0; i < npart_this_proc; ++i) { - distr(ix, iy, it, ipx, ipy, ipt, rng); - x.push_back(ix); - y.push_back(iy); - t.push_back(it); - px.push_back(ipx); - py.push_back(ipy); - pt.push_back(ipt); - } - - int const lev = 0; - pc.AddNParticles(lev, x, y, t, px, py, pt, - qm, bunch_charge); - } -} // namespace impactx diff --git a/src/initialization/InitDistribution.cpp b/src/initialization/InitDistribution.cpp index d28e81f3f..c77a428ff 100644 --- a/src/initialization/InitDistribution.cpp +++ b/src/initialization/InitDistribution.cpp @@ -7,15 +7,9 @@ * Authors: Axel Huebl, Chad Mitchell, Ji Qiang * License: BSD-3-Clause-LBNL */ -#include "InitDistribution.H" #include "ImpactX.H" #include "particles/ImpactXParticleContainer.H" -#include "particles/distribution/Waterbag.H" -#include "particles/distribution/Kurth6D.H" -#include "particles/distribution/Gaussian.H" -#include "particles/distribution/KVdist.H" -#include "particles/distribution/Kurth4D.H" -#include "particles/distribution/Semigaussian.H" +#include "particles/distribution/All.H" #include #include @@ -28,6 +22,61 @@ namespace impactx { + void + ImpactX::add_particles ( + amrex::ParticleReal qm, + amrex::ParticleReal bunch_charge, + distribution::KnownDistributions distr, + int npart + ) + { + BL_PROFILE("ImpactX::add_particles"); + + amrex::Vector x, y, t; + amrex::Vector px, py, pt; + amrex::ParticleReal ix, iy, it, ipx, ipy, ipt; + amrex::RandomEngine rng; + + // Logic: We initialize 1/Nth of particles, independent of their + // position, per MPI rank. We then measure the distribution's spatial + // extent, create a grid, resize it to fit the beam, and then + // redistribute particles so that they reside on the correct MPI rank. + int myproc = amrex::ParallelDescriptor::MyProc(); + int nprocs = amrex::ParallelDescriptor::NProcs(); + int navg = npart / nprocs; + int nleft = npart - navg * nprocs; + int npart_this_proc = (myproc < nleft) ? navg+1 : navg; + + std::visit([&](auto&& distribution){ + x.reserve(npart_this_proc); + y.reserve(npart_this_proc); + t.reserve(npart_this_proc); + px.reserve(npart_this_proc); + py.reserve(npart_this_proc); + pt.reserve(npart_this_proc); + + for (amrex::Long i = 0; i < npart_this_proc; ++i) { + distribution(ix, iy, it, ipx, ipy, ipt, rng); + x.push_back(ix); + y.push_back(iy); + t.push_back(it); + px.push_back(ipx); + py.push_back(ipy); + pt.push_back(ipt); + } + }, distr); + + int const lev = 0; + m_particle_container->AddNParticles(lev, x, y, t, px, py, pt, + qm, bunch_charge); + + // Resize the mesh to fit the spatial extent of the beam and then + // redistribute particles, so they reside on the MPI rank that is + // responsible for the respective spatial particle position. + this->ResizeMesh(); + m_particle_container->Redistribute(); + } + void ImpactX::initBeamDistributionFromInputs () { BL_PROFILE("ImpactX::initBeamDistributionFromInputs"); @@ -78,12 +127,12 @@ namespace impactx pp_dist.query("muypy", muypy); pp_dist.query("mutpt", mutpt); - impactx::distribution::Waterbag waterbag( + distribution::KnownDistributions waterbag(distribution::Waterbag( sigx, sigy, sigt, sigpx, sigpy, sigpt, - muxpx, muypy, mutpt); + muxpx, muypy, mutpt)); - generate_add_particles(*m_particle_container, qm, bunch_charge, waterbag, npart); + add_particles(qm, bunch_charge, waterbag, npart); } else if (distribution_type == "kurth6d") { amrex::ParticleReal sigx,sigy,sigt,sigpx,sigpy,sigpt; @@ -98,12 +147,12 @@ namespace impactx pp_dist.query("muypy", muypy); pp_dist.query("mutpt", mutpt); - impactx::distribution::Kurth6D kurth6D( + distribution::KnownDistributions kurth6D(distribution::Kurth6D( sigx, sigy, sigt, sigpx, sigpy, sigpt, - muxpx, muypy, mutpt); + muxpx, muypy, mutpt)); - generate_add_particles(*m_particle_container, qm, bunch_charge, kurth6D, npart); + add_particles(qm, bunch_charge, kurth6D, npart); } else if (distribution_type == "gaussian") { amrex::ParticleReal sigx,sigy,sigt,sigpx,sigpy,sigpt; @@ -118,12 +167,12 @@ namespace impactx pp_dist.query("muypy", muypy); pp_dist.query("mutpt", mutpt); - impactx::distribution::Gaussian gaussian( + distribution::KnownDistributions gaussian(distribution::Gaussian( sigx, sigy, sigt, sigpx, sigpy, sigpt, - muxpx, muypy, mutpt); + muxpx, muypy, mutpt)); - generate_add_particles(*m_particle_container, qm, bunch_charge, gaussian, npart); + add_particles(qm, bunch_charge, gaussian, npart); } else if (distribution_type == "kvdist") { amrex::ParticleReal sigx,sigy,sigt,sigpx,sigpy,sigpt; @@ -138,12 +187,12 @@ namespace impactx pp_dist.query("muypy", muypy); pp_dist.query("mutpt", mutpt); - impactx::distribution::KVdist kvDist( + distribution::KnownDistributions kvDist(distribution::KVdist( sigx, sigy, sigt, sigpx, sigpy, sigpt, - muxpx, muypy, mutpt); + muxpx, muypy, mutpt)); - generate_add_particles(*m_particle_container, qm, bunch_charge, kvDist, npart); + add_particles(qm, bunch_charge, kvDist, npart); } else if (distribution_type == "kurth4d") { amrex::ParticleReal sigx,sigy,sigt,sigpx,sigpy,sigpt; @@ -158,12 +207,12 @@ namespace impactx pp_dist.query("muypy", muypy); pp_dist.query("mutpt", mutpt); - impactx::distribution::Kurth4D kurth4D( + distribution::KnownDistributions kurth4D(distribution::Kurth4D( sigx, sigy, sigt, sigpx, sigpy, sigpt, - muxpx, muypy, mutpt); + muxpx, muypy, mutpt)); - generate_add_particles(*m_particle_container, qm, bunch_charge, kurth4D, npart); + add_particles(qm, bunch_charge, kurth4D, npart); } else if (distribution_type == "semigaussian") { amrex::ParticleReal sigx,sigy,sigt,sigpx,sigpy,sigpt; amrex::ParticleReal muxpx = 0.0, muypy = 0.0, mutpt = 0.0; @@ -177,22 +226,16 @@ namespace impactx pp_dist.query("muypy", muypy); pp_dist.query("mutpt", mutpt); - impactx::distribution::Semigaussian semigaussian( + distribution::KnownDistributions semigaussian(distribution::Semigaussian( sigx, sigy, sigt, sigpx, sigpy, sigpt, - muxpx, muypy, mutpt); + muxpx, muypy, mutpt)); - generate_add_particles(*m_particle_container, qm, bunch_charge, semigaussian, npart); + add_particles(qm, bunch_charge, semigaussian, npart); } else { amrex::Abort("Unknown distribution: " + distribution_type); } - // Resize the mesh to fit the spatial extent of the beam and then - // redistribute particles, so they reside on the MPI rank that is - // responsible for the respective spatial particle position. - this->ResizeMesh(); - m_particle_container->Redistribute(); - // reference particle amrex::ParticleReal massE; // MeV if (particle_type == "electron") { diff --git a/src/python/ImpactX.cpp b/src/python/ImpactX.cpp index ff6609c78..435cf28af 100644 --- a/src/python/ImpactX.cpp +++ b/src/python/ImpactX.cpp @@ -6,7 +6,6 @@ #include "pyImpactX.H" #include -#include #include #include #include @@ -37,7 +36,8 @@ void init_ImpactX(py::module& m) >(m, "MultiFabPerLevel"); */ - py::class_(m, "ImpactX") + py::class_ impactx(m, "ImpactX"); + impactx .def(py::init<>()) .def("load_inputs_file", @@ -83,6 +83,10 @@ void init_ImpactX(py::module& m) .def("init_grids", &ImpactX::initGrids) .def("init_beam_distribution_from_inputs", &ImpactX::initBeamDistributionFromInputs) .def("init_lattice_elements_from_inputs", &ImpactX::initLatticeElementsFromInputs) + .def("add_particles", &ImpactX::add_particles, + py::arg("qm"), py::arg("bunch_charge"), + py::arg("distr"), py::arg("npart") + ) .def("evolve", &ImpactX::evolve) .def("particle_container", [](ImpactX & ix) -> ImpactXParticleContainer & { diff --git a/src/python/distribution.cpp b/src/python/distribution.cpp index a685880f4..95d7d6b09 100644 --- a/src/python/distribution.cpp +++ b/src/python/distribution.cpp @@ -5,7 +5,6 @@ */ #include "pyImpactX.H" -#include #include #include @@ -16,28 +15,6 @@ namespace py = pybind11; using namespace impactx; -namespace -{ - /** Register impactx::generate_add_particles for each distribution - * - * @tparam I counter through all known distributions in impactx::distribution::KnownDistributions - * @param md the module to register the function in - */ - template< std::size_t I = 0 > - void register_generate_add_particles(py::module& md) - { - using V = impactx::distribution::KnownDistributions; - using T = std::variant_alternative_t; - - md.def("generate_add_particles", &generate_add_particles, - py::arg("pc"), py::arg("qm"), py::arg("bunch_charge"), - py::arg("distr"), py::arg("npart") - ); - - if constexpr (I < std::variant_size_v - 1) - register_generate_add_particles(md); - } -} void init_distribution(py::module& m) { @@ -111,6 +88,4 @@ void init_distribution(py::module& m) py::arg("sigmaPx"), py::arg("sigmaPy"), py::arg("sigmaPt"), py::arg("muxpx")=0.0, py::arg("muypy")=0.0, py::arg("mutpt")=0.0 ); - - register_generate_add_particles(md); } diff --git a/tests/python/test_impactx.py b/tests/python/test_impactx.py index 14124eacd..c3f3a1c94 100644 --- a/tests/python/test_impactx.py +++ b/tests/python/test_impactx.py @@ -45,8 +45,7 @@ def test_impactx_nofile(): muxpx = -0.846574929020762, muypy = 0.846574929020762, mutpt = 0.0) - distribution.generate_add_particles( - impactX.particle_container(), qm_qeeV, charge_C, distr, npart) + impactX.add_particles(qm_qeeV, charge_C, distr, npart) # init reference particle refPart = RefPart()