Test: Charge Deposition

Test the charge deposition logic & scaling of values.

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

src/ImpactX.H

@@ -15,6 +15,7 @@
 #include "particles/ImpactXParticleContainer.H"
 
 #include <AMReX_AmrCore.H>
+#include <AMReX_MultiFab.H>
 
 #include <list>
 #include <memory>

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

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<amrex::Real, 3> 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

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

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_<Config>(m, "Config")

tests/python/test_charge_deposition.py

@@ -0,0 +1,74 @@
+# -*- coding: utf-8 -*-
+
+
+import amrex
+import impactx
+import math
+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 non-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()