[Draft] 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

@@ -49,6 +49,7 @@ namespace impactx
                     auto & AMREX_RESTRICT soa_real = pti.GetStructOfArrays().GetRealData();
                     // after https://github.com/ECP-WarpX/WarpX/pull/2838 add const:
                     auto const wp = soa_real[RealSoA::w];
+                    amrex::Print() << "wp[0]=" << wp[0] << "\n";
                     int const * const AMREX_RESTRICT ion_lev = nullptr;
 
                     // physical lower corner of the current box
@@ -66,7 +67,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)};

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

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

tests/python/test_charge_deposition.py

@@ -0,0 +1,54 @@
+# -*- 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(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)
+    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())  # TODO
+        print(f"bx={bx}")
+        #print(f"rbx={rbx}")
+
+        arr = rho.array(mfi)
+        arr_np = np.array(arr, copy=False)
+
+        half_z = arr_np.shape[1] // 2
+        im = ax.imshow(
+            arr_np[0, half_z, ...]
+            #extent=rbx
+        )
+        f.colorbar(im)
+        plt.show()