From 9113ad7aada5c6308717f1e66bd841e43c131c3a Mon Sep 17 00:00:00 2001 From: Axel Huebl Date: Tue, 28 Jun 2022 00:34:20 +0200 Subject: [PATCH 01/10] Evolve: Check Particles Exist Make sure that particles are initialized and otherwise error out cleanly with an exception. --- src/ImpactX.cpp | 14 ++++++++++++++ 1 file changed, 14 insertions(+) diff --git a/src/ImpactX.cpp b/src/ImpactX.cpp index d67058c24..84f76b674 100644 --- a/src/ImpactX.cpp +++ b/src/ImpactX.cpp @@ -59,6 +59,20 @@ namespace impactx // before we start the evolve loop, we are in "step 0" (initial state) int global_step = 0; + // count particles - if no particles are found in our particle container, then a lot of + // AMReX routines over ParIter won't work and we have nothing to do here anyways + { + int const nLevelPC = finestLevel(); + amrex::Long nParticles = 0; + for (int lev = 0; lev <= nLevelPC; ++lev) { + nParticles += m_particle_container->NumberOfParticlesAtLevel(lev); + } + if (nParticles == 0) { + amrex::Abort("No particles found. Cannot run evolve without a beam."); + return; + } + } + amrex::ParmParse pp_diag("diag"); int file_min_digits = 6; { From ba1e98acf81af427d64d694a9a48cfae9d382bba Mon Sep 17 00:00:00 2001 From: Axel Huebl Date: Tue, 28 Jun 2022 00:35:02 +0200 Subject: [PATCH 02/10] Diagnostics: Move Copy a bit down Not needed here yet. --- src/particles/diagnostics/DiagnosticOutput.cpp | 14 +++++++------- 1 file changed, 7 insertions(+), 7 deletions(-) diff --git a/src/particles/diagnostics/DiagnosticOutput.cpp b/src/particles/diagnostics/DiagnosticOutput.cpp index 4b739d3df..40d75ce17 100644 --- a/src/particles/diagnostics/DiagnosticOutput.cpp +++ b/src/particles/diagnostics/DiagnosticOutput.cpp @@ -28,13 +28,6 @@ namespace impactx::diagnostics { using namespace amrex::literals; // for _rt and _prt - // create a host-side particle buffer - auto tmp = pc.make_alike(); - - // copy device-to-host - bool const local = true; - tmp.copyParticles(pc, local); - // write file header per MPI RANK if (!append) { if (otype == OutputType::PrintParticles) { @@ -46,6 +39,13 @@ namespace impactx::diagnostics } } + // create a host-side particle buffer + auto tmp = pc.make_alike(); + + // copy device-to-host + bool const local = true; + tmp.copyParticles(pc, local); + // loop over refinement levels int const nLevel = tmp.finestLevel(); for (int lev = 0; lev <= nLevel; ++lev) { From fb46b274aa345be2bf39ec1ae04fdeb14f84ab1e Mon Sep 17 00:00:00 2001 From: Axel Huebl Date: Tue, 28 Jun 2022 00:35:35 +0200 Subject: [PATCH 03/10] CoordinateTransformation: Formatting AMReX formatting with extra space --- src/particles/transformation/CoordinateTransformation.H | 4 ++-- src/particles/transformation/CoordinateTransformation.cpp | 4 ++-- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/src/particles/transformation/CoordinateTransformation.H b/src/particles/transformation/CoordinateTransformation.H index 13a9734fb..d6ad81597 100644 --- a/src/particles/transformation/CoordinateTransformation.H +++ b/src/particles/transformation/CoordinateTransformation.H @@ -30,8 +30,8 @@ namespace transformation * @param pc container of the particles to push * @param direction the direction (Z to T or vice versa) */ - void CoordinateTransformation(ImpactXParticleContainer &pc, - Direction const & direction); + void CoordinateTransformation (ImpactXParticleContainer &pc, + Direction const & direction); } // namespace transformation } // namespace impactx diff --git a/src/particles/transformation/CoordinateTransformation.cpp b/src/particles/transformation/CoordinateTransformation.cpp index e0b6acf7a..e488bf04a 100644 --- a/src/particles/transformation/CoordinateTransformation.cpp +++ b/src/particles/transformation/CoordinateTransformation.cpp @@ -21,8 +21,8 @@ namespace impactx { namespace transformation { - void CoordinateTransformation(ImpactXParticleContainer &pc, - Direction const &direction) { + void CoordinateTransformation (ImpactXParticleContainer &pc, + Direction const &direction) { using namespace amrex::literals; // for _rt and _prt // preparing to access reference particle data: RefPart From 2fc1600bb0686b60999207c79a689af76faf8eb1 Mon Sep 17 00:00:00 2001 From: Axel Huebl Date: Fri, 22 Jul 2022 12:13:59 -0700 Subject: [PATCH 04/10] Distribution: None Add a variant for all distributions --- src/particles/distribution/All.H | 35 +++++++++++++++++++ src/particles/distribution/None.H | 58 +++++++++++++++++++++++++++++++ 2 files changed, 93 insertions(+) create mode 100644 src/particles/distribution/All.H create mode 100644 src/particles/distribution/None.H diff --git a/src/particles/distribution/All.H b/src/particles/distribution/All.H new file mode 100644 index 000000000..cea248bf3 --- /dev/null +++ b/src/particles/distribution/All.H @@ -0,0 +1,35 @@ +/* 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 + */ +#ifndef IMPACTX_DISTRIBUTION_ALL_H +#define IMPACTX_DISTRIBUTION_ALL_H + +#include "Gaussian.H" +#include "Kurth4D.H" +#include "Kurth6D.H" +#include "KVdist.H" +#include "None.H" +#include "Semigaussian.H" +#include "Waterbag.H" + +#include + + +namespace impactx +{ +namespace distribution +{ + using KnownDistributions = std::variant< + None, /* must be first, so KnownDistributions creates a default constructor */ + Gaussian, Kurth4D, Kurth6D, KVdist, Semigaussian, Waterbag>; + +} // namespace distribution +} // namespace impactx + +#endif // IMPACTX_DISTRIBUTION_ALL_H diff --git a/src/particles/distribution/None.H b/src/particles/distribution/None.H new file mode 100644 index 000000000..c773a2a7e --- /dev/null +++ b/src/particles/distribution/None.H @@ -0,0 +1,58 @@ +/* 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: Chad Mitchell, Axel Huebl + * License: BSD-3-Clause-LBNL + */ +#ifndef IMPACTX_DISTRIBUTION_NONE +#define IMPACTX_DISTRIBUTION_NONE + +#include +#include + + +namespace impactx +{ +namespace distribution +{ + struct None + { + /** This distribution does nothing + */ + None() + { + } + + /** Return 1 6D particle coordinate + * + * Does nothing to the parameters. + * + * @param x particle position in x + * @param y particle position in y + * @param t particle position in t + * @param px particle momentum in x + * @param py particle momentum in y + * @param pt particle momentum in t + * @param engine a random number engine (with associated state) + */ + AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE + void operator() ( + [[maybe_unused]] amrex::ParticleReal & x, + [[maybe_unused]] amrex::ParticleReal & y, + [[maybe_unused]] amrex::ParticleReal & t, + [[maybe_unused]] amrex::ParticleReal & px, + [[maybe_unused]] amrex::ParticleReal & py, + [[maybe_unused]] amrex::ParticleReal & pt, + [[maybe_unused]] amrex::RandomEngine const& engine) const + { + /* nothing to do */ + } + }; + +} // namespace distribution +} // namespace impactx + +#endif // IMPACTX_DISTRIBUTION_NONE From c1ccf7e59f4977b4281669361c07ecb99cc7a256 Mon Sep 17 00:00:00 2001 From: Axel Huebl Date: Tue, 28 Jun 2022 01:17:14 +0200 Subject: [PATCH 05/10] Python: Particle Init --- src/initialization/InitDistribution.H | 79 +++++++++++++++++ src/initialization/InitDistribution.cpp | 61 +------------ src/python/CMakeLists.txt | 3 + src/python/ImpactX.cpp | 16 +++- src/python/ImpactXParticleContainer.cpp | 33 +++++++ src/python/ReferenceParticle.cpp | 28 ++++++ src/python/distribution.cpp | 113 ++++++++++++++++++++++++ src/python/pyImpactX.cpp | 7 ++ tests/python/CMakeLists.txt | 4 +- tests/python/test_impactx.py | 33 ++++++- 10 files changed, 310 insertions(+), 67 deletions(-) create mode 100644 src/initialization/InitDistribution.H create mode 100644 src/python/ImpactXParticleContainer.cpp create mode 100644 src/python/ReferenceParticle.cpp create mode 100644 src/python/distribution.cpp diff --git a/src/initialization/InitDistribution.H b/src/initialization/InitDistribution.H new file mode 100644 index 000000000..644f38104 --- /dev/null +++ b/src/initialization/InitDistribution.H @@ -0,0 +1,79 @@ +/* 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 118be783c..01332bdb2 100644 --- a/src/initialization/InitDistribution.cpp +++ b/src/initialization/InitDistribution.cpp @@ -7,6 +7,7 @@ * 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" @@ -25,65 +26,6 @@ #include -namespace -{ - /** 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 { void ImpactX::initBeamDistributionFromInputs () @@ -267,6 +209,7 @@ namespace impactx refPart.z = 0.0; refPart.px = 0.0; refPart.py = 0.0; + // make the next two lines a helper function? refPart.pt = -energy/massE - 1.0_prt; refPart.pz = sqrt(pow(refPart.pt,2) - 1.0_prt); m_particle_container->SetRefParticle(refPart); diff --git a/src/python/CMakeLists.txt b/src/python/CMakeLists.txt index 9315680de..55ee28f66 100644 --- a/src/python/CMakeLists.txt +++ b/src/python/CMakeLists.txt @@ -4,6 +4,9 @@ # target_sources(pyImpactX PRIVATE + distribution.cpp elements.cpp ImpactX.cpp + ImpactXParticleContainer.cpp + ReferenceParticle.cpp ) diff --git a/src/python/ImpactX.cpp b/src/python/ImpactX.cpp index e3619658b..ff6609c78 100644 --- a/src/python/ImpactX.cpp +++ b/src/python/ImpactX.cpp @@ -6,6 +6,14 @@ #include "pyImpactX.H" #include +#include +#include +#include +#include +#include +#include +#include + #include #include @@ -76,8 +84,12 @@ void init_ImpactX(py::module& m) .def("init_beam_distribution_from_inputs", &ImpactX::initBeamDistributionFromInputs) .def("init_lattice_elements_from_inputs", &ImpactX::initLatticeElementsFromInputs) .def("evolve", &ImpactX::evolve) - - //.def_property("particle_container", &ImpactX::m_particle_container) + .def("particle_container", + [](ImpactX & ix) -> ImpactXParticleContainer & { + return *ix.m_particle_container; + }, + py::return_value_policy::reference_internal + ) //.def_readwrite("rho", &ImpactX::m_rho) .def_readwrite("lattice", &ImpactX::m_lattice) //.def_readwrite("lattice", &ImpactX::m_lattice_test) diff --git a/src/python/ImpactXParticleContainer.cpp b/src/python/ImpactXParticleContainer.cpp new file mode 100644 index 000000000..cb087274a --- /dev/null +++ b/src/python/ImpactXParticleContainer.cpp @@ -0,0 +1,33 @@ +/* Copyright 2021-2022 The ImpactX Community + * + * Authors: Axel Huebl + * License: BSD-3-Clause-LBNL + */ +#include "pyImpactX.H" + +#include +#include +#include + +namespace py = pybind11; +using namespace impactx; + + +void init_impactxparticlecontainer(py::module& m) +{ + py::class_< + ImpactXParticleContainer, + amrex::ParticleContainer<0, 0, RealSoA::nattribs, IntSoA::nattribs> + >(m, "ImpactXParticleContainer") + //.def(py::init<>()) + .def("add_n_particles", &ImpactXParticleContainer::AddNParticles) + .def("get_ref_particle", + py::overload_cast<>(&ImpactXParticleContainer::GetRefParticle), + py::return_value_policy::reference_internal + ) + .def("set_ref_particle", &ImpactXParticleContainer::SetRefParticle) + .def("set_particle_shape", + py::overload_cast(&ImpactXParticleContainer::SetParticleShape) + ) + ; +} diff --git a/src/python/ReferenceParticle.cpp b/src/python/ReferenceParticle.cpp new file mode 100644 index 000000000..d4007409d --- /dev/null +++ b/src/python/ReferenceParticle.cpp @@ -0,0 +1,28 @@ +/* Copyright 2021-2022 The ImpactX Community + * + * Authors: Axel Huebl + * License: BSD-3-Clause-LBNL + */ +#include "pyImpactX.H" + +#include +#include + +namespace py = pybind11; +using namespace impactx; + + +void init_refparticle(py::module& m) +{ + py::class_(m, "RefPart") + .def(py::init<>()) + .def_readwrite("x", &RefPart::x, "horizontal position x, in meters") + .def_readwrite("y", &RefPart::y, "vertical position y, in meters") + .def_readwrite("z", &RefPart::z, "longitudinal position y, in meters") + .def_readwrite("t", &RefPart::t, "clock time * c in meters") + .def_readwrite("px", &RefPart::px, "momentum in x, normalized to proper velocity") + .def_readwrite("py", &RefPart::py, "momentum in y, normalized to proper velocity") + .def_readwrite("pz", &RefPart::pz, "momentum in z, normalized to proper velocity") + .def_readwrite("pt", &RefPart::pt, "energy deviation, normalized by rest energy") + ; +} diff --git a/src/python/distribution.cpp b/src/python/distribution.cpp new file mode 100644 index 000000000..48a3685d8 --- /dev/null +++ b/src/python/distribution.cpp @@ -0,0 +1,113 @@ +/* Copyright 2021-2022 The ImpactX Community + * + * Authors: Axel Huebl + * License: BSD-3-Clause-LBNL + */ +#include "pyImpactX.H" + +#include +#include + +#include +#include + +#include + +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); + + if constexpr (I < std::variant_size_v - 1) + register_generate_add_particles(md); + } +} + +void init_distribution(py::module& m) +{ + py::module_ md = m.def_submodule( + "distribution", + "Particle beam distributions in ImpactX" + ); + + py::class_(md, "Gaussian") + .def(py::init< + amrex::ParticleReal const, amrex::ParticleReal const, amrex::ParticleReal const, + amrex::ParticleReal const, amrex::ParticleReal const, amrex::ParticleReal const, + amrex::ParticleReal const, amrex::ParticleReal const, amrex::ParticleReal const + >(), + py::arg("sigmaX"), py::arg("sigmaY"), py::arg("sigmaT"), + py::arg("sigmaPx"), py::arg("sigmaPy"), py::arg("sigmaPt"), + py::arg("muxpx"), py::arg("muypy"), py::arg("mutpt") + ); + + py::class_(md, "Kurth4D") + .def(py::init< + amrex::ParticleReal const, amrex::ParticleReal const, amrex::ParticleReal const, + amrex::ParticleReal const, amrex::ParticleReal const, amrex::ParticleReal const, + amrex::ParticleReal const, amrex::ParticleReal const, amrex::ParticleReal const + >(), + py::arg("sigmaX"), py::arg("sigmaY"), py::arg("sigmaT"), + py::arg("sigmaPx"), py::arg("sigmaPy"), py::arg("sigmaPt"), + py::arg("muxpx"), py::arg("muypy"), py::arg("mutpt") + ); + + py::class_(md, "Kurth6D") + .def(py::init< + amrex::ParticleReal const, amrex::ParticleReal const, amrex::ParticleReal const, + amrex::ParticleReal const, amrex::ParticleReal const, amrex::ParticleReal const, + amrex::ParticleReal const, amrex::ParticleReal const, amrex::ParticleReal const + >(), + py::arg("sigmaX"), py::arg("sigmaY"), py::arg("sigmaT"), + py::arg("sigmaPx"), py::arg("sigmaPy"), py::arg("sigmaPt"), + py::arg("muxpx"), py::arg("muypy"), py::arg("mutpt") + ); + + py::class_(md, "KVdist") + .def(py::init< + amrex::ParticleReal const, amrex::ParticleReal const, amrex::ParticleReal const, + amrex::ParticleReal const, amrex::ParticleReal const, amrex::ParticleReal const, + amrex::ParticleReal const, amrex::ParticleReal const, amrex::ParticleReal const + >(), + py::arg("sigmaX"), py::arg("sigmaY"), py::arg("sigmaT"), + py::arg("sigmaPx"), py::arg("sigmaPy"), py::arg("sigmaPt"), + py::arg("muxpx"), py::arg("muypy"), py::arg("mutpt") + ); + + py::class_(md, "Semigaussian") + .def(py::init< + amrex::ParticleReal const, amrex::ParticleReal const, amrex::ParticleReal const, + amrex::ParticleReal const, amrex::ParticleReal const, amrex::ParticleReal const, + amrex::ParticleReal const, amrex::ParticleReal const, amrex::ParticleReal const + >(), + py::arg("sigmaX"), py::arg("sigmaY"), py::arg("sigmaT"), + py::arg("sigmaPx"), py::arg("sigmaPy"), py::arg("sigmaPt"), + py::arg("muxpx"), py::arg("muypy"), py::arg("mutpt") + ); + + py::class_(md, "Waterbag") + .def(py::init< + amrex::ParticleReal const, amrex::ParticleReal const, amrex::ParticleReal const, + amrex::ParticleReal const, amrex::ParticleReal const, amrex::ParticleReal const, + amrex::ParticleReal const, amrex::ParticleReal const, amrex::ParticleReal const + >(), + py::arg("sigmaX"), py::arg("sigmaY"), py::arg("sigmaT"), + py::arg("sigmaPx"), py::arg("sigmaPy"), py::arg("sigmaPt"), + py::arg("muxpx"), py::arg("muypy"), py::arg("mutpt") + ); + + register_generate_add_particles(md); +} diff --git a/src/python/pyImpactX.cpp b/src/python/pyImpactX.cpp index 3b0b7016b..ea8ec5061 100644 --- a/src/python/pyImpactX.cpp +++ b/src/python/pyImpactX.cpp @@ -16,8 +16,11 @@ using namespace impactx; // forward declarations of exposed classes +void init_distribution(py::module&); void init_elements(py::module&); void init_ImpactX(py::module&); +void init_impactxparticlecontainer(py::module&); +void init_refparticle(py::module&); PYBIND11_MODULE(impactx_pybind, m) { // make sure AMReX types are known @@ -31,11 +34,15 @@ PYBIND11_MODULE(impactx_pybind, m) { .. autosummary:: :toctree: _generate ImpactX + distribution elements )pbdoc"; // note: order from parent to child classes + init_distribution(m); init_elements(m); + init_refparticle(m); + init_impactxparticlecontainer(m); init_ImpactX(m); // API runtime version diff --git a/tests/python/CMakeLists.txt b/tests/python/CMakeLists.txt index bd7693084..68993df8a 100644 --- a/tests/python/CMakeLists.txt +++ b/tests/python/CMakeLists.txt @@ -2,8 +2,8 @@ file(COPY ${ImpactX_SOURCE_DIR}/examples DESTINATION ${CMAKE_PYTHON_OUTPUT_DIRECTORY}) # add pytest tests -add_test(NAME pytest - COMMAND ${Python_EXECUTABLE} -m pytest +add_test(NAME pytest.ImpactX + COMMAND ${Python_EXECUTABLE} -m pytest -s -vvvv ${ImpactX_SOURCE_DIR}/tests/python WORKING_DIRECTORY ${CMAKE_PYTHON_OUTPUT_DIRECTORY} diff --git a/tests/python/test_impactx.py b/tests/python/test_impactx.py index 7a276e468..14124eacd 100644 --- a/tests/python/test_impactx.py +++ b/tests/python/test_impactx.py @@ -1,6 +1,6 @@ # -*- coding: utf-8 -*- -from impactx import ImpactX, elements +from impactx import ImpactX, RefPart, distribution, elements def test_impactx_fodo_file(): @@ -29,7 +29,33 @@ def test_impactx_nofile(): impactX.init_grids() # init particle beam - # TODO + energy_MeV = 2.0e3 + charge_C = 0.0 + mass_MeV = 0.510998950 + qm_qeeV = -1.0/0.510998950e6 + npart = 10000 + + distr = distribution.Waterbag( + sigmaX = 3.9984884770e-5, + sigmaY = 3.9984884770e-5, + sigmaT = 1.0e-3, + sigmaPx = 2.6623538760e-5, + sigmaPy = 2.6623538760e-5, + sigmaPt = 2.0e-3, + muxpx = -0.846574929020762, + muypy = 0.846574929020762, + mutpt = 0.0) + distribution.generate_add_particles( + impactX.particle_container(), qm_qeeV, charge_C, distr, npart) + + # init reference particle + refPart = RefPart() + # make the next two lines a helper function? + refPart.pt = -energy_MeV / mass_MeV - 1.0 + refPart.pz = (refPart.pt**2 - 1.0)**0.5 + impactX.particle_container().set_ref_particle(refPart) + + assert(impactX.particle_container().TotalNumberOfParticles() == npart) # init accelerator lattice fodo = [ @@ -54,5 +80,4 @@ def test_impactx_nofile(): print(len(impactX.lattice)) assert(len(impactX.lattice) > 5) - # TODO: enable once particle beam is loaded - #impactX.evolve() + impactX.evolve() From f601017c43b1b5b93fd3f6722f4ab7b81b1856b1 Mon Sep 17 00:00:00 2001 From: Axel Huebl Date: Fri, 22 Jul 2022 16:30:13 -0700 Subject: [PATCH 06/10] Python: FODO Example --- examples/CMakeLists.txt | 75 +++++++++++++++++++++++++++++++++++---- examples/fodo/README.rst | 32 +++++++++++++++++ examples/fodo/run_fodo.py | 70 ++++++++++++++++++++++++++++++++++++ 3 files changed, 170 insertions(+), 7 deletions(-) create mode 100755 examples/fodo/run_fodo.py diff --git a/examples/CMakeLists.txt b/examples/CMakeLists.txt index 68d9c8c1e..c463bdb4f 100644 --- a/examples/CMakeLists.txt +++ b/examples/CMakeLists.txt @@ -18,6 +18,20 @@ if(ImpactX_MPI) ) endif() +function(impactx_test_set_pythonpath test_name) + if(WIN32) + string(REPLACE ";" "\\;" WIN_PYTHONPATH "$ENV{PYTHONPATH}") + string(REGEX REPLACE "/" "\\\\" WIN_PYTHON_OUTPUT_DIRECTORY ${CMAKE_PYTHON_OUTPUT_DIRECTORY}) + set_tests_properties(TEST ${test_name} + PROPERTIES ENVIRONMENT "PYTHONPATH=${WIN_PYTHON_OUTPUT_DIRECTORY}\;${WIN_PYTHONPATH}" + ) + else() + set_tests_properties(${test_name} + PROPERTIES ENVIRONMENT "PYTHONPATH=${CMAKE_PYTHON_OUTPUT_DIRECTORY}:$ENV{PYTHONPATH}" + ) + endif() +endfunction() + # FODO Cell ################################################################### # @@ -42,17 +56,17 @@ set_tests_properties(FODO.plot PROPERTIES DEPENDS "FODO.run") # if(ImpactX_MPI) add_test(NAME FODO.MPI.run - COMMAND ${MPI_TEST_EXE} - $ ${ImpactX_SOURCE_DIR}/examples/fodo/input_fodo.in - WORKING_DIRECTORY ${ImpactX_RUNTIME_OUTPUT_DIRECTORY} + COMMAND ${MPI_TEST_EXE} + $ ${ImpactX_SOURCE_DIR}/examples/fodo/input_fodo.in + WORKING_DIRECTORY ${ImpactX_RUNTIME_OUTPUT_DIRECTORY} ) add_test(NAME FODO.MPI.analysis - COMMAND ${ImpactX_SOURCE_DIR}/examples/fodo/analysis_fodo.py - WORKING_DIRECTORY ${ImpactX_RUNTIME_OUTPUT_DIRECTORY} + COMMAND ${ImpactX_SOURCE_DIR}/examples/fodo/analysis_fodo.py + WORKING_DIRECTORY ${ImpactX_RUNTIME_OUTPUT_DIRECTORY} ) add_test(NAME FODO.MPI.plot - COMMAND ${ImpactX_SOURCE_DIR}/examples/fodo/plot_fodo.py --save-png - WORKING_DIRECTORY ${ImpactX_RUNTIME_OUTPUT_DIRECTORY} + COMMAND ${ImpactX_SOURCE_DIR}/examples/fodo/plot_fodo.py --save-png + WORKING_DIRECTORY ${ImpactX_RUNTIME_OUTPUT_DIRECTORY} ) set_tests_properties(FODO.MPI.run PROPERTIES ENVIRONMENT "OMP_NUM_THREADS=1") @@ -61,6 +75,53 @@ if(ImpactX_MPI) endif() +# Python: FODO Cell ########################################################### +# +if(ImpactX_PYTHON) + add_test(NAME FODO.py.run + COMMAND ${Python_EXECUTABLE} + ${ImpactX_SOURCE_DIR}/examples/fodo/run_fodo.py + WORKING_DIRECTORY ${ImpactX_RUNTIME_OUTPUT_DIRECTORY} + ) + add_test(NAME FODO.py.analysis + COMMAND ${ImpactX_SOURCE_DIR}/examples/fodo/analysis_fodo.py + WORKING_DIRECTORY ${ImpactX_RUNTIME_OUTPUT_DIRECTORY} + ) + add_test(NAME FODO.py.plot + COMMAND ${ImpactX_SOURCE_DIR}/examples/fodo/plot_fodo.py --save-png + WORKING_DIRECTORY ${ImpactX_RUNTIME_OUTPUT_DIRECTORY} + ) + + set_tests_properties(FODO.py.analysis PROPERTIES DEPENDS "FODO.py.run") + set_tests_properties(FODO.py.plot PROPERTIES DEPENDS "FODO.py.run") + impactx_test_set_pythonpath(FODO.py.run) +endif() + + +# Python: MPI-parallel FODO Cell ############################################## +# +if(ImpactX_PYTHON AND ImpactX_MPI) + add_test(NAME FODO.py.MPI.run + COMMAND ${MPI_TEST_EXE} ${Python_EXECUTABLE} + ${ImpactX_SOURCE_DIR}/examples/fodo/run_fodo.py + WORKING_DIRECTORY ${ImpactX_RUNTIME_OUTPUT_DIRECTORY} + ) + add_test(NAME FODO.py.MPI.analysis + COMMAND ${ImpactX_SOURCE_DIR}/examples/fodo/analysis_fodo.py + WORKING_DIRECTORY ${ImpactX_RUNTIME_OUTPUT_DIRECTORY} + ) + add_test(NAME FODO.py.MPI.plot + COMMAND ${ImpactX_SOURCE_DIR}/examples/fodo/plot_fodo.py --save-png + WORKING_DIRECTORY ${ImpactX_RUNTIME_OUTPUT_DIRECTORY} + ) + + set_tests_properties(FODO.py.MPI.run PROPERTIES ENVIRONMENT "OMP_NUM_THREADS=1") + set_tests_properties(FODO.py.MPI.analysis PROPERTIES DEPENDS "FODO.py.MPI.run") + set_tests_properties(FODO.py.MPI.plot PROPERTIES DEPENDS "FODO.py.MPI.run") + impactx_test_set_pythonpath(FODO.py.MPI.run) +endif() + + # Chicane ##################################################################### # add_test(NAME chicane.run diff --git a/examples/fodo/README.rst b/examples/fodo/README.rst index 7a30f7348..7a8a46c9d 100644 --- a/examples/fodo/README.rst +++ b/examples/fodo/README.rst @@ -19,7 +19,39 @@ The second moments of the particle distribution after the FODO cell should coinc In this test, the initial and final values of :math:`\sigma_x`, :math:`\sigma_y`, :math:`\sigma_t`, :math:`\epsilon_x`, :math:`\epsilon_y`, and :math:`\epsilon_t` must agree with nominal values. +Run +--- + +This example an either be run with an inputs file (``impactx input_fodo.in``): .. literalinclude:: input_fodo.in :language: ini :caption: You can copy this file from ``examples/fodo/input_fodo.in``. + +Or as a Python script (``python3 run_fodo.py``): + +.. literalinclude:: run_fodo.py + :language: python3 + :caption: You can copy this file from ``examples/fodo/run_fodo.py``. + +Both execution modes can also be prefixed with an MPI executor, such as ``mpiexec -n 4 ...`` or ``srun -n 4 ...``, depending on the system. + + +Analyze +------- + +We run the following script to analyze correctness: + +.. literalinclude:: analysis_fodo.py + :language: python3 + :caption: You can copy this file from ``examples/fodo/analysis_fodo.py``. + + +Visualize +--------- + +You can run the following script to visualize the beam evolution over time: + +.. literalinclude:: plot_fodo.py + :language: python3 + :caption: You can copy this file from ``examples/fodo/plot_fodo.py``. diff --git a/examples/fodo/run_fodo.py b/examples/fodo/run_fodo.py new file mode 100755 index 000000000..8799933d6 --- /dev/null +++ b/examples/fodo/run_fodo.py @@ -0,0 +1,70 @@ +#!/usr/bin/env python3 +# +# Copyright 2022 ImpactX contributors +# Authors: Axel Huebl, Chad Mitchell +# License: BSD-3-Clause-LBNL +# +# -*- coding: utf-8 -*- + +import amrex +from impactx import ImpactX, RefPart, distribution, elements + + +amrex.initialize([ + # print some AMReX status messages + "amrex.verbose=1", + # abort GPU runs if out-of-memory instead of swapping to host RAM + "amrex.abort_on_out_of_gpu_memory=1" +]) + +impactX = ImpactX() + +impactX.set_particle_shape(2) +impactX.set_diags_slice_step_diagnostics(True) +impactX.init_grids() + +# init particle beam +energy_MeV = 2.0e3 +charge_C = 0.0 # assign zero weighting to particles +mass_MeV = 0.510998950 +qm_qeeV = -1.0/0.510998950e6 +npart = 10000 + +distr = distribution.Waterbag( + sigmaX = 3.9984884770e-5, + sigmaY = 3.9984884770e-5, + sigmaT = 1.0e-3, + sigmaPx = 2.6623538760e-5, + sigmaPy = 2.6623538760e-5, + sigmaPt = 2.0e-3, + muxpx = -0.846574929020762, + muypy = 0.846574929020762, + mutpt = 0.0) +distribution.generate_add_particles( + impactX.particle_container(), qm_qeeV, charge_C, distr, npart) + +# init reference particle +refPart = RefPart() +# make the next two lines a helper function? +refPart.pt = -energy_MeV / mass_MeV - 1.0 +refPart.pz = (refPart.pt**2 - 1.0)**0.5 +impactX.particle_container().set_ref_particle(refPart) + +# init accelerator lattice +ns = 25 # number of slices per ds in the element +fodo = [ + elements.Drift(ds=0.25, nslice=ns), + elements.Quad(ds=1.0, k=1.0, nslice=ns), + elements.Drift(ds=0.5, nslice=ns), + elements.Quad(ds=1.0, k=-1.0, nslice=ns), + elements.Drift(ds=0.25, nslice=ns) +] +# assign a fodo segment +impactX.lattice.extend(fodo) + +# run simulation +impactX.evolve() + +# clean shutdown +del impactX +amrex.finalize() From 2a160c478d1e8afd5f7f5c30e9b5a7fef9ff677c Mon Sep 17 00:00:00 2001 From: Axel Huebl Date: Fri, 22 Jul 2022 16:30:48 -0700 Subject: [PATCH 07/10] AMReX: Allow Delayed Initialization This makes scripting workflows less verbose. One can still manually initialize AMReX, but users can also rely on defaults now. --- examples/fodo/run_fodo.py | 9 +---- src/initialization/CMakeLists.txt | 1 + src/initialization/InitAMReX.H | 32 ++++++++++++++++ src/initialization/InitAMReX.cpp | 47 ++++++++++++++++++++++++ src/initialization/InitOneBoxPerRank.cpp | 11 ++++++ src/main.cpp | 14 +------ 6 files changed, 94 insertions(+), 20 deletions(-) create mode 100644 src/initialization/InitAMReX.H create mode 100644 src/initialization/InitAMReX.cpp diff --git a/examples/fodo/run_fodo.py b/examples/fodo/run_fodo.py index 8799933d6..82c56704e 100755 --- a/examples/fodo/run_fodo.py +++ b/examples/fodo/run_fodo.py @@ -7,15 +7,8 @@ # -*- coding: utf-8 -*- import amrex -from impactx import ImpactX, RefPart, distribution, elements - -amrex.initialize([ - # print some AMReX status messages - "amrex.verbose=1", - # abort GPU runs if out-of-memory instead of swapping to host RAM - "amrex.abort_on_out_of_gpu_memory=1" -]) +from impactx import ImpactX, RefPart, distribution, elements impactX = ImpactX() diff --git a/src/initialization/CMakeLists.txt b/src/initialization/CMakeLists.txt index f3ee4acda..27f95472c 100644 --- a/src/initialization/CMakeLists.txt +++ b/src/initialization/CMakeLists.txt @@ -1,6 +1,7 @@ target_sources(ImpactX PRIVATE AmrCoreData.cpp + InitAMReX.cpp InitDistribution.cpp InitElement.cpp InitMeshRefinement.cpp diff --git a/src/initialization/InitAMReX.H b/src/initialization/InitAMReX.H new file mode 100644 index 000000000..b8af3aa41 --- /dev/null +++ b/src/initialization/InitAMReX.H @@ -0,0 +1,32 @@ +/* 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 + */ +#ifndef IMPACT_DEFAULT_INIT_AMREX_H +#define IMPACT_DEFAULT_INIT_AMREX_H + +#include "AmrCoreData.H" + + +namespace impactx::initialization +{ + /** Initialize AMReX + * + * Initializes AMReX if not already done. + */ + void default_init_AMReX (int argc, char* argv[]); + + /** Initialize AMReX + * + * Initializes AMReX if not already done. + */ + void default_init_AMReX (); + +} // namespace impactx::initialization + +#endif // IMPACT_DEFAULT_INIT_AMREX_H diff --git a/src/initialization/InitAMReX.cpp b/src/initialization/InitAMReX.cpp new file mode 100644 index 000000000..1a8fba12c --- /dev/null +++ b/src/initialization/InitAMReX.cpp @@ -0,0 +1,47 @@ +/* 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, Chad Mitchell, Ji Qiang + * License: BSD-3-Clause-LBNL + */ +#include "InitAMReX.H" + +#include "initialization/InitParser.H" + +#include +#include + +#if defined(AMREX_USE_MPI) +# include +#else +# include +#endif + + +namespace impactx::initialization +{ + void + default_init_AMReX (int argc, char* argv[]) + { + if (!amrex::Initialized()) + { + bool const build_parm_parse = true; + amrex::Initialize( + argc, + argv, + build_parm_parse, + MPI_COMM_WORLD, + impactx::initialization::overwrite_amrex_parser_defaults + ); + } + } + + void + default_init_AMReX () + { + default_init_AMReX(0, nullptr); + } +} // namespace impactx::initialization diff --git a/src/initialization/InitOneBoxPerRank.cpp b/src/initialization/InitOneBoxPerRank.cpp index 277a6bdc0..c6692899d 100644 --- a/src/initialization/InitOneBoxPerRank.cpp +++ b/src/initialization/InitOneBoxPerRank.cpp @@ -9,6 +9,8 @@ */ #include "InitOneBoxPerRank.H" +#include "initialization/InitAMReX.H" + #include #include #include @@ -18,12 +20,21 @@ #include #include +#include + namespace impactx::initialization { AmrCoreData one_box_per_rank () { + if (!amrex::Initialized()) + { + default_init_AMReX(); + // note: due to global state, it would be too early to use amrex::Abort() here + //throw std::runtime_error("AMReX must be initialized before ImpactX simulation can be constructed."); + } + amrex::AmrInfo amr_info; const int nprocs = amrex::ParallelDescriptor::NProcs(); const amrex::IntVect high_end = amr_info.blocking_factor[0] diff --git a/src/main.cpp b/src/main.cpp index 647f14e67..32b261c74 100644 --- a/src/main.cpp +++ b/src/main.cpp @@ -8,16 +8,13 @@ * License: BSD-3-Clause-LBNL */ #include "ImpactX.H" -#include "initialization/InitParser.H" +#include "initialization/InitAMReX.H" #include #include -#include #if defined(AMREX_USE_MPI) # include -#else -# include #endif @@ -27,14 +24,7 @@ int main(int argc, char* argv[]) AMREX_ALWAYS_ASSERT(MPI_SUCCESS == MPI_Init(&argc, &argv)); #endif - bool const build_parm_parse = true; - amrex::Initialize( - argc, - argv, - build_parm_parse, - MPI_COMM_WORLD, - impactx::initialization::overwrite_amrex_parser_defaults - ); + impactx::initialization::default_init_AMReX(argc, argv); BL_PROFILE_VAR("main()", pmain); { From d847c329b32d83a99d385c7597a1cd721200d0b7 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Fri, 22 Jul 2022 23:32:56 +0000 Subject: [PATCH 08/10] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- examples/fodo/run_fodo.py | 1 - 1 file changed, 1 deletion(-) diff --git a/examples/fodo/run_fodo.py b/examples/fodo/run_fodo.py index 82c56704e..fc48a3d2c 100755 --- a/examples/fodo/run_fodo.py +++ b/examples/fodo/run_fodo.py @@ -7,7 +7,6 @@ # -*- coding: utf-8 -*- import amrex - from impactx import ImpactX, RefPart, distribution, elements impactX = ImpactX() From aa85805150dbe3783bfe93007d95986194bb8f78 Mon Sep 17 00:00:00 2001 From: Axel Huebl Date: Fri, 22 Jul 2022 17:04:27 -0700 Subject: [PATCH 09/10] Python Distributions: Default Parameters --- src/python/distribution.cpp | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/src/python/distribution.cpp b/src/python/distribution.cpp index 48a3685d8..188c8c33f 100644 --- a/src/python/distribution.cpp +++ b/src/python/distribution.cpp @@ -51,7 +51,7 @@ void init_distribution(py::module& m) >(), py::arg("sigmaX"), py::arg("sigmaY"), py::arg("sigmaT"), py::arg("sigmaPx"), py::arg("sigmaPy"), py::arg("sigmaPt"), - py::arg("muxpx"), py::arg("muypy"), py::arg("mutpt") + py::arg("muxpx")=0.0, py::arg("muypy")=0.0, py::arg("mutpt")=0.0 ); py::class_(md, "Kurth4D") @@ -62,7 +62,7 @@ void init_distribution(py::module& m) >(), py::arg("sigmaX"), py::arg("sigmaY"), py::arg("sigmaT"), py::arg("sigmaPx"), py::arg("sigmaPy"), py::arg("sigmaPt"), - py::arg("muxpx"), py::arg("muypy"), py::arg("mutpt") + py::arg("muxpx")=0.0, py::arg("muypy")=0.0, py::arg("mutpt")=0.0 ); py::class_(md, "Kurth6D") @@ -73,7 +73,7 @@ void init_distribution(py::module& m) >(), py::arg("sigmaX"), py::arg("sigmaY"), py::arg("sigmaT"), py::arg("sigmaPx"), py::arg("sigmaPy"), py::arg("sigmaPt"), - py::arg("muxpx"), py::arg("muypy"), py::arg("mutpt") + py::arg("muxpx")=0.0, py::arg("muypy")=0.0, py::arg("mutpt")=0.0 ); py::class_(md, "KVdist") @@ -84,7 +84,7 @@ void init_distribution(py::module& m) >(), py::arg("sigmaX"), py::arg("sigmaY"), py::arg("sigmaT"), py::arg("sigmaPx"), py::arg("sigmaPy"), py::arg("sigmaPt"), - py::arg("muxpx"), py::arg("muypy"), py::arg("mutpt") + py::arg("muxpx")=0.0, py::arg("muypy")=0.0, py::arg("mutpt")=0.0 ); py::class_(md, "Semigaussian") @@ -95,7 +95,7 @@ void init_distribution(py::module& m) >(), py::arg("sigmaX"), py::arg("sigmaY"), py::arg("sigmaT"), py::arg("sigmaPx"), py::arg("sigmaPy"), py::arg("sigmaPt"), - py::arg("muxpx"), py::arg("muypy"), py::arg("mutpt") + py::arg("muxpx")=0.0, py::arg("muypy")=0.0, py::arg("mutpt")=0.0 ); py::class_(md, "Waterbag") @@ -106,7 +106,7 @@ void init_distribution(py::module& m) >(), py::arg("sigmaX"), py::arg("sigmaY"), py::arg("sigmaT"), py::arg("sigmaPx"), py::arg("sigmaPy"), py::arg("sigmaPt"), - py::arg("muxpx"), py::arg("muypy"), py::arg("mutpt") + py::arg("muxpx")=0.0, py::arg("muypy")=0.0, py::arg("mutpt")=0.0 ); register_generate_add_particles(md); From 74323c292f0110e966bcefb74a664c1bee9c7549 Mon Sep 17 00:00:00 2001 From: Axel Huebl Date: Fri, 22 Jul 2022 17:28:38 -0700 Subject: [PATCH 10/10] Python: generate_add_particles args Add argument names to `distribution.generate_add_particles` --- src/python/distribution.cpp | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/src/python/distribution.cpp b/src/python/distribution.cpp index 188c8c33f..a685880f4 100644 --- a/src/python/distribution.cpp +++ b/src/python/distribution.cpp @@ -29,7 +29,10 @@ namespace using V = impactx::distribution::KnownDistributions; using T = std::variant_alternative_t; - md.def("generate_add_particles", &generate_add_particles); + 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);