[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/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,7 @@ 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 charge = 1.0; // TODO: once we implement charge
 
                     // 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/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::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,66 @@
+# -*- coding: utf-8 -*-
+
+import amrex
+from impactx import ImpactX
+import matplotlib.pyplot as plt
+import numpy as np
+
+
+def test_charge_deposition():
+    """
+    Deposit charge and access/plot it
+    """
+    impactX = ImpactX()
+
+    impactX.load_inputs_file("examples/fodo/input_fodo.in")
+
+    impactX.init_grids()
+    impactX.init_beam_distribution_from_inputs()
+    impactX.init_lattice_elements_from_inputs()
+
+    impactX.evolve()
+
+    f = plt.figure()
+    ax = f.gca()
+    lev = 0
+    rho = impactX.rho(lev)
+    print(f"rho={rho}")
+    print(f"rho.sum={rho.sum(0)}")
+
+    gm = impactX.Geom(lev)
+    print(f"gm={gm}")
+
+    for mfi in rho:
+        bx = mfi.validbox()
+        #rbx = amrex.RealBox(bx, gm.CellSize(), gm.ProbLo())
+        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()
+
+
+if __name__ == '__main__':
+    amrex.initialize([
+        # print AMReX status messages
+        "amrex.verbose=2",
+        # throw exceptions and create core dumps instead of
+        # AMReX backtrace files: allows to attach to
+        # debuggers
+        "amrex.throw_exception=1",
+        "amrex.signal_handling=0",
+        # abort GPU runs if out-of-memory instead of swapping to host RAM
+        "amrex.abort_on_out_of_gpu_memory=1"
+    ])
+
+    test_charge_deposition()
+
+    amrex.finalize()