Test: Charge Deposition (#165)

* Test: Charge Deposition Test the charge deposition logic & scaling of values. * Charge: `sum`->`sum_unique` Use AMReX-Codes/pyamrex#63 * AMReX: Update Commit ... to include: AMReX-Codes/amrex#2909
ECP-WarpX · Aug 11, 2022 · aa543c4 · aa543c4
1 parent a16ffa9
commit aa543c4
Show file tree

Hide file tree

Showing 9 changed files with 109 additions and 7 deletions.
diff --git 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
@@ -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
@@ -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
@@ -15,6 +15,7 @@
 #include "particles/ImpactXParticleContainer.H"
 
 #include <AMReX_AmrCore.H>
+#include <AMReX_MultiFab.H>
 
 #include <list>
 #include <memory>

diff --git 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
@@ -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

diff --git 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
@@ -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")

diff --git a/tests/python/test_charge_deposition.py 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()