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/src/ImpactX.H b/src/ImpactX.H index 3739cfc17..1c36c79fe 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 864fa7ae1..e6d51f8c2 100644 --- a/src/python/ImpactX.cpp +++ b/src/python/ImpactX.cpp @@ -112,9 +112,30 @@ void init_ImpactX(py::module& m) }, py::return_value_policy::reference_internal ) - //.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) - //.def_readwrite("lattice", &ImpactX::m_lattice_test) + + // 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..8d606ef8e --- /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) + # TODO: for cell-centered data, this double counts non-owned cells with MPI + rs = rho.sum(comp=0, local=False) + + gm = sim.Geom(lev=0) + dr = gm.data().CellSize() + dV = np.prod(dr) + + beam_charge = dV*rs # in C + # TODO: does not yet pass with MPI runs (too large value, see above) + #assert math.isclose(beam_charge, 1.0e-9) + + f = plt.figure() + ax = f.gca() + 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 + comp=0 + mu = 1.e6 # m->mu + im = ax.imshow( + #arr_np[comp, half_z, ...], # including guard + arr_np[ + comp, + half_z, + bx.lo_vect[1]:bx.hi_vect[1], + bx.lo_vect[0]:bx.hi_vect[0]], # w/o guard + origin='lower', + aspect='auto', + extent=[ + rbx.lo(1) * mu, rbx.hi(1) * mu, + rbx.lo(0) * mu, rbx.hi(0) * mu + ] + ) + cb = f.colorbar(im) + cb.set_label(r"charge density [C/m$^3$]") + ax.set_xlabel(r"$y$ [$\mu$m]") + ax.set_ylabel(r"$x$ [$\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()