From aa543c422101aecffbad1e78170a8342dce6d6a1 Mon Sep 17 00:00:00 2001 From: Axel Huebl Date: Thu, 11 Aug 2022 12:11:10 -0600 Subject: [PATCH] Test: Charge Deposition (#165) * Test: Charge Deposition Test the charge deposition logic & scaling of values. * Charge: `sum`->`sum_unique` Use https://github.com/AMReX-Codes/pyamrex/pull/63 * AMReX: Update Commit ... to include: https://github.com/AMReX-Codes/amrex/pull/2909 --- cmake/dependencies/ABLASTR.cmake | 2 +- examples/fodo/input_fodo.in | 2 +- examples/fodo/run_fodo.py | 2 +- src/ImpactX.H | 1 + src/initialization/InitDistribution.cpp | 4 +- src/particles/ChargeDeposition.cpp | 5 +- src/particles/ImpactXParticleContainer.cpp | 2 +- src/python/ImpactX.cpp | 24 ++++++- tests/python/test_charge_deposition.py | 74 ++++++++++++++++++++++ 9 files changed, 109 insertions(+), 7 deletions(-) create mode 100644 tests/python/test_charge_deposition.py diff --git a/cmake/dependencies/ABLASTR.cmake b/cmake/dependencies/ABLASTR.cmake index 7de956193..ad6776def 100644 --- a/cmake/dependencies/ABLASTR.cmake +++ b/cmake/dependencies/ABLASTR.cmake @@ -142,7 +142,7 @@ set(ImpactX_ablastr_branch "22.08" set(ImpactX_amrex_repo "https://github.com/AMReX-Codes/amrex.git" CACHE STRING "Repository URI to pull and build AMReX from if(ImpactX_amrex_internal)") -set(ImpactX_amrex_branch "" +set(ImpactX_amrex_branch "1bda173b489024d5f4ec79627f3f612c350e521f" CACHE STRING "Repository branch for ImpactX_amrex_repo if(ImpactX_amrex_internal)") diff --git a/examples/fodo/input_fodo.in b/examples/fodo/input_fodo.in index 9dc4f3129..7e610def1 100644 --- a/examples/fodo/input_fodo.in +++ b/examples/fodo/input_fodo.in @@ -4,7 +4,7 @@ beam.npart = 10000 beam.units = static beam.energy = 2.0e3 -beam.charge = 0.0 +beam.charge = 1.0e-9 beam.particle = electron beam.distribution = waterbag beam.sigmaX = 3.9984884770e-5 diff --git a/examples/fodo/run_fodo.py b/examples/fodo/run_fodo.py index 3da722c52..c401c73d8 100755 --- a/examples/fodo/run_fodo.py +++ b/examples/fodo/run_fodo.py @@ -23,7 +23,7 @@ # load a 2 GeV electron beam with an initial # unnormalized rms emittance of 2 nm energy_MeV = 2.0e3 # reference energy -charge_C = 0.0 # assign zero weighting to particles +charge_C = 1.0e-9 # used with space charge mass_MeV = 0.510998950 qm_qeeV = -1.0e-6/mass_MeV # charge/mass npart = 10000 # number of macro particles diff --git a/src/ImpactX.H b/src/ImpactX.H index 11658e7fe..862a310d7 100644 --- a/src/ImpactX.H +++ b/src/ImpactX.H @@ -15,6 +15,7 @@ #include "particles/ImpactXParticleContainer.H" #include +#include #include #include diff --git a/src/initialization/InitDistribution.cpp b/src/initialization/InitDistribution.cpp index a9140e09e..b85338011 100644 --- a/src/initialization/InitDistribution.cpp +++ b/src/initialization/InitDistribution.cpp @@ -46,6 +46,8 @@ namespace impactx int navg = npart / nprocs; int nleft = npart - navg * nprocs; int npart_this_proc = (myproc < nleft) ? navg+1 : navg; + auto const rel_part_this_proc = amrex::ParticleReal(npart_this_proc) / + amrex::ParticleReal(npart); std::visit([&](auto&& distribution){ x.reserve(npart_this_proc); @@ -68,7 +70,7 @@ namespace impactx int const lev = 0; m_particle_container->AddNParticles(lev, x, y, t, px, py, pt, - qm, bunch_charge); + qm, bunch_charge * rel_part_this_proc); // Resize the mesh to fit the spatial extent of the beam and then // redistribute particles, so they reside on the MPI rank that is diff --git a/src/particles/ChargeDeposition.cpp b/src/particles/ChargeDeposition.cpp index 6bb33a8fc..b31adab0a 100644 --- a/src/particles/ChargeDeposition.cpp +++ b/src/particles/ChargeDeposition.cpp @@ -66,7 +66,8 @@ namespace impactx //auto const rel_ref_ratio = ref_ratio.at(depos_lev) / ref_ratio.at(lev); amrex::ignore_unused(ref_ratio); - amrex::ParticleReal const charge = 1.0; // TODO once we implement charge + amrex::ParticleReal const q_e = 1.60217662e-19; // TODO move out + amrex::ParticleReal const charge = q_e; // cell size of the mesh to deposit to std::array const & AMREX_RESTRICT dx = {gm.CellSize(0), gm.CellSize(1), gm.CellSize(2)}; @@ -84,6 +85,8 @@ namespace impactx // start async charge communication for this level rho_at_level.SumBoundary_nowait(); + //int const comp = 0; + //rho_at_level.SumBoundary_nowait(comp, comp, rho_at_level.nGrowVect()); } // finalize communication diff --git a/src/particles/ImpactXParticleContainer.cpp b/src/particles/ImpactXParticleContainer.cpp index 04d6880e1..dde913f8f 100644 --- a/src/particles/ImpactXParticleContainer.cpp +++ b/src/particles/ImpactXParticleContainer.cpp @@ -109,7 +109,7 @@ namespace impactx pinned_tile.push_back_real(RealSoA::uy, py); pinned_tile.push_back_real(RealSoA::pt, pz); pinned_tile.push_back_real(RealSoA::m_qm, np, qm); - amrex::ParticleReal const q_e = 1.60217662e-19; + amrex::ParticleReal const q_e = 1.60217662e-19; // TODO move out pinned_tile.push_back_real(RealSoA::w, np, bchchg/q_e/np); /* Redistributes particles to their respective tiles (spatial bucket diff --git a/src/python/ImpactX.cpp b/src/python/ImpactX.cpp index b8b78cbfc..d98931d86 100644 --- a/src/python/ImpactX.cpp +++ b/src/python/ImpactX.cpp @@ -126,11 +126,33 @@ void init_ImpactX(py::module& m) py::return_value_policy::reference_internal, "Access the beam particle container." ) - //.def_readwrite("rho", &ImpactX::m_rho) + .def( + "rho", + [](ImpactX & ix, int const lev) { return &ix.m_rho.at(lev); }, + py::arg("lev"), + py::return_value_policy::reference_internal + ) .def_readwrite("lattice", &ImpactX::m_lattice, "Access the accelerator element lattice." ) + + // from AmrCore->AmrMesh + .def("Geom", + //[](ImpactX const & ix, int const lev) { return ix.Geom(lev); }, + py::overload_cast< int >(&ImpactX::Geom, py::const_), + py::arg("lev") + ) + .def("DistributionMap", + [](ImpactX const & ix, int const lev) { return ix.DistributionMap(lev); }, + //py::overload_cast< int >(&ImpactX::DistributionMap, py::const_), + py::arg("lev") + ) + .def("boxArray", + [](ImpactX const & ix, int const lev) { return ix.boxArray(lev); }, + //py::overload_cast< int >(&ImpactX::boxArray, py::const_), + py::arg("lev") + ) ; py::class_(m, "Config") diff --git a/tests/python/test_charge_deposition.py b/tests/python/test_charge_deposition.py new file mode 100644 index 000000000..4a7b4ad3a --- /dev/null +++ b/tests/python/test_charge_deposition.py @@ -0,0 +1,74 @@ +# -*- coding: utf-8 -*- + + +import math + +import amrex +import impactx +import matplotlib.pyplot as plt +import numpy as np + + +def test_charge_deposition(): + """ + Deposit charge and access/plot it + """ + sim = impactx.ImpactX() + + sim.load_inputs_file("examples/fodo/input_fodo.in") + sim.set_slice_step_diagnostics(False) + + sim.init_grids() + sim.init_beam_distribution_from_inputs() + sim.init_lattice_elements_from_inputs() + + sim.evolve() + + rho = sim.rho(lev=0) + rs = rho.sum_unique(comp=0, local=False) + + gm = sim.Geom(lev=0) + dr = gm.data().CellSize() + dV = np.prod(dr) + + beam_charge = dV*rs # in C + assert math.isclose(beam_charge, 1.0e-9) + + f = plt.figure() + ax = f.gca() + ng = rho.nGrowVect + for mfi in rho: + bx = mfi.validbox() + rbx = amrex.RealBox(bx, dr, gm.ProbLo()) + + arr = rho.array(mfi) + arr_np = np.array(arr, copy=False) + + half_z = arr_np.shape[1] // 2 # indices: comp, z, y, x + comp = 0 + mu = 1.e6 # m->mu + im = ax.imshow( + #arr_np[comp, half_z, ...] * dV, # including guard + arr_np[ + comp, + half_z, + ng[1]:-ng[1], + ng[0]:-ng[0]] * dV, # w/o guard + origin='lower', + aspect='auto', + extent=[ + rbx.lo(0) * mu, rbx.hi(0) * mu, + rbx.lo(1) * mu, rbx.hi(1) * mu + ] + ) + cb = f.colorbar(im) + cb.set_label(r"charge density [C/m$^3$]") + ax.set_xlabel(r"$x$ [$\mu$m]") + ax.set_ylabel(r"$y$ [$\mu$m]") + plt.show() + + +# implement a direct script run mode, so we can run this directly too, +# with interactive matplotlib windows, w/o pytest +if __name__ == '__main__': + test_charge_deposition()