Space Charge Solver (#162)

* New Fields: Phi & Force

Add new fields (MultiFabs) for the scalar potential and the space
charge force.

* Poisson Solver

* Fix Space Charge Force Constructor

* Add ForceFromSelfFields

Add function ForceFromSelfFields to calculate space charge self
fields and their force.

* Phi & Force: Expose to Python

* rename: space_charge_field

* Implement: Nodal Gather and Push Momentum

* Force Calc: Fix Stray Include

* Cleaning: Docs, Formatting, Constants, Constness

* Test: Analyze Expanding Beam

Co-authored-by: Marco Garten <mgarten@lbl.gov>

examples/CMakeLists.txt

@@ -286,7 +286,7 @@ add_impactx_test(expanding_beam
     examples/expanding_beam/input_expanding.in
       OFF  # ImpactX MPI-parallel
       OFF  # ImpactX Python interface
-    OFF  # examples/expanding_beam/analysis_expanding.py
+    examples/expanding_beam/analysis_expanding.py
     OFF  # no plot script yet
 )
 
@@ -296,7 +296,7 @@ add_impactx_test(expanding_beam.py
     examples/expanding_beam/run_expanding.py
       OFF  # ImpactX MPI-parallel
       ON   # ImpactX Python interface
-    OFF  # examples/expanding_beam/analysis_expanding.py
+    examples/expanding_beam/analysis_expanding.py
     OFF  # no plot script yet
 )
 

src/ImpactX.H

@@ -124,8 +124,12 @@ namespace impactx
         /** these are the physical/beam particles of the simulation */
         std::unique_ptr<ImpactXParticleContainer> m_particle_container;
 
-        /** charge per level */
+        /** charge density per level */
         std::unordered_map<int, amrex::MultiFab> m_rho;
+        /** scalar potential per level */
+        std::unordered_map<int, amrex::MultiFab> m_phi;
+        /** space charge field (vector) per level */
+        std::unordered_map<int, std::unordered_map<std::string, amrex::MultiFab> > m_space_charge_field;
 
         /** these are elements defining the accelerator lattice */
         std::list<KnownElements> m_lattice;

src/ImpactX.cpp

@@ -11,8 +11,11 @@
 #include "initialization/InitAmrCore.H"
 #include "particles/ImpactXParticleContainer.H"
 #include "particles/Push.H"
-#include "particles/transformation/CoordinateTransformation.H"
 #include "particles/diagnostics/DiagnosticOutput.H"
+#include "particles/spacecharge/ForceFromSelfFields.H"
+#include "particles/spacecharge/GatherAndPush.H"
+#include "particles/spacecharge/PoissonSolve.H"
+#include "particles/transformation/CoordinateTransformation.H"
 
 #include <AMReX.H>
 #include <AMReX_AmrParGDB.H>
@@ -101,7 +104,11 @@ namespace impactx
         {
             // number of slices used for the application of space charge
             int nslice = 1;
-            std::visit([&nslice](auto&& element){ nslice = element.nslice(); }, element_variant);
+            amrex::ParticleReal slice_ds; // in meters
+            std::visit([&nslice, &slice_ds](auto&& element){
+                nslice = element.nslice();
+                slice_ds = element.ds() / nslice;
+            }, element_variant);
 
             // sub-steps for space charge within the element
             for (int slice_step = 0; slice_step < nslice; ++slice_step)
@@ -117,8 +124,9 @@ namespace impactx
                 {
 
                     // transform from x',y',t to x,y,z
-                    transformation::CoordinateTransformation(*m_particle_container,
-                                                             transformation::Direction::to_fixed_t);
+                    transformation::CoordinateTransformation(
+                        *m_particle_container,
+                        transformation::Direction::to_fixed_t);
 
                     // Note: The following operation assume that
                     // the particles are in x, y, z coordinates.
@@ -133,11 +141,20 @@ namespace impactx
                     m_particle_container->DepositCharge(m_rho, this->refRatio());
 
                     // poisson solve in x,y,z
-                    //   TODO
+                    spacecharge::PoissonSolve(*m_particle_container, m_rho, m_phi);
+
+                    // calculate force in x,y,z
+                    spacecharge::ForceFromSelfFields(m_space_charge_field,
+                                                     m_phi,
+                                                     this->geom);
 
                     // gather and space-charge push in x,y,z , assuming the space-charge
                     // field is the same before/after transformation
-                    //   TODO
+                    // TODO: This is currently using linear order.
+                    spacecharge::GatherAndPush(*m_particle_container,
+                                               m_space_charge_field,
+                                               this->geom,
+                                               slice_ds);
 
                     // transform from x,y,z to x',y',t
                     transformation::CoordinateTransformation(*m_particle_container,

src/initialization/InitMeshRefinement.cpp

@@ -68,8 +68,35 @@ namespace impactx
         else  // odd shape orders
             num_guards_rho = (particle_shape + 1) / 2;
 
-        m_rho.emplace(lev,
-                      amrex::MultiFab{amrex::convert(cba, rho_nodal_flag), dm, num_components_rho, num_guards_rho, tag("rho")});
+        m_rho.emplace(
+            lev,
+            amrex::MultiFab{amrex::convert(cba, rho_nodal_flag), dm, num_components_rho, num_guards_rho, tag("rho")});
+
+        // scalar potential
+        auto const phi_nodal_flag = rho_nodal_flag;
+        int const num_components_phi = 1;
+        int const num_guards_phi = num_guards_rho + 1; // todo: I think this just depends on max(MLMG, force calc)
+        m_phi.emplace(
+            lev,
+            amrex::MultiFab{amrex::convert(cba, phi_nodal_flag), dm, num_components_phi, num_guards_phi, tag("phi")});
+
+        // space charge force
+        std::unordered_map<std::string, amrex::MultiFab> f_comp;
+        for (std::string const comp : {"x", "y", "z"})
+        {
+            std::string const str_tag = "space_charge_field_" + comp;
+            f_comp.emplace(
+                comp,
+                amrex::MultiFab{
+                    amrex::convert(cba, rho_nodal_flag),
+                    dm,
+                    num_components_rho,
+                    num_guards_rho,
+                    tag(str_tag)
+                }
+            );
+        }
+        m_space_charge_field.emplace(lev, std::move(f_comp));
     }
 
     /** Make a new level using provided BoxArray and DistributionMapping and fill
@@ -101,6 +128,8 @@ namespace impactx
     void ImpactX::ClearLevel (int lev)
     {
         m_rho.erase(lev);
+        m_phi.erase(lev);
+        m_space_charge_field.erase(lev);
     }
 
     void ImpactX::ResizeMesh ()

src/particles/CMakeLists.txt

@@ -5,5 +5,6 @@ target_sources(ImpactX
     Push.cpp
 )
 
-add_subdirectory(transformation)
 add_subdirectory(diagnostics)
+add_subdirectory(spacecharge)
+add_subdirectory(transformation)

src/particles/spacecharge/CMakeLists.txt

@@ -0,0 +1,6 @@
+target_sources(ImpactX
+  PRIVATE
+    ForceFromSelfFields.cpp
+    GatherAndPush.cpp
+    PoissonSolve.cpp
+)

src/particles/spacecharge/ForceFromSelfFields.H

@@ -0,0 +1,41 @@
+/* Copyright 2022 The Regents of the University of California, through Lawrence
+ *           Berkeley National Laboratory (subject to receipt of any required
+ *           approvals from the U.S. Dept. of Energy). All rights reserved.
+ *
+ * This file is part of ImpactX.
+ *
+ * Authors: Marco Garten, Axel Huebl
+ * License: BSD-3-Clause-LBNL
+ */
+#ifndef IMPACTX_FORCEFROMSELFFIELDS_H
+#define IMPACTX_FORCEFROMSELFFIELDS_H
+
+#include "particles/ImpactXParticleContainer.H"
+
+#include <AMReX_Geometry.H>
+#include <AMReX_MultiFab.H>
+#include <AMReX_Vector.H>
+
+#include <unordered_map>
+
+
+namespace impactx::spacecharge
+{
+    /** Calculate the space charge force field from the electric potential
+     *
+     * This resets the values in scf_<component> to zero and then calculates the space
+     * charge force field.
+     *
+     * @param[inout] space_charge_field space charge force component in x,y,z per level
+     * @param[in] phi scalar potential per level
+     * @param[in] geom geometry object
+     */
+    void ForceFromSelfFields (
+        std::unordered_map<int, std::unordered_map<std::string, amrex::MultiFab> > & space_charge_field,
+        std::unordered_map<int, amrex::MultiFab> const & phi,
+        const amrex::Vector<amrex::Geometry>& geom
+    );
+
+} // namespace impactx
+
+#endif // IMPACTX_FORCEFROMSELFFIELDS_H

src/particles/spacecharge/ForceFromSelfFields.cpp

@@ -0,0 +1,60 @@
+/* Copyright 2022 The Regents of the University of California, through Lawrence
+ *           Berkeley National Laboratory (subject to receipt of any required
+ *           approvals from the U.S. Dept. of Energy). All rights reserved.
+ *
+ * This file is part of ImpactX.
+ *
+ * Authors: Marco Garten, Axel Huebl
+ * License: BSD-3-Clause-LBNL
+ */
+#include "ForceFromSelfFields.H"
+
+#include <AMReX_BLProfiler.H>
+#include <AMReX_REAL.H>       // for Real
+#include <AMReX_SPACE.H>      // for AMREX_D_DECL
+
+
+namespace impactx::spacecharge
+{
+    void ForceFromSelfFields (
+        std::unordered_map<int, std::unordered_map<std::string, amrex::MultiFab> > & space_charge_field,
+        std::unordered_map<int, amrex::MultiFab> const & phi,
+        amrex::Vector<amrex::Geometry> const & geom
+    )
+    {
+        BL_PROFILE("impactx::spacecharge::ForceFromSelfFields");
+
+        using namespace amrex::literals;
+
+        // loop over refinement levels
+        int const finest_level = phi.size() - 1u;
+        for (int lev = 0; lev <= finest_level; ++lev) {
+
+            // stencil coefficients: 0.5 * inverse cell size
+            auto const &gm = geom[lev];
+            auto const dr = gm.CellSizeArray();
+            amrex::GpuArray<amrex::Real, 3> const inv2dr{AMREX_D_DECL(0.5_rt/dr[0], 0.5_rt/dr[1], 0.5_rt/dr[2])};
+
+            // reset the values in space_charge_field to zero
+            space_charge_field.at(lev).at("x").setVal(0.);
+            space_charge_field.at(lev).at("y").setVal(0.);
+            space_charge_field.at(lev).at("z").setVal(0.);
+
+            for (amrex::MFIter mfi(phi.at(lev)); mfi.isValid(); ++mfi) {
+
+                amrex::Box bx = mfi.validbox();
+                auto const phi_arr = (phi.at(lev))[mfi].const_array();
+
+                auto scf_arr_x = space_charge_field[lev]["x"][mfi].array();
+                auto scf_arr_y = space_charge_field[lev]["y"][mfi].array();
+                auto scf_arr_z = space_charge_field[lev]["z"][mfi].array();
+
+                amrex::ParallelFor(bx, [=] AMREX_GPU_DEVICE (int i, int j, int k) noexcept {
+                    scf_arr_x(i, j, k) = inv2dr[0] * (phi_arr(i-1, j, k) - phi_arr(i+1, j, k));
+                    scf_arr_y(i, j, k) = inv2dr[1] * (phi_arr(i, j-1, k) - phi_arr(i, j+1, k));
+                    scf_arr_z(i, j, k) = inv2dr[2] * (phi_arr(i, j, k-1) - phi_arr(i, j, k+1));
+                });
+            }
+        }
+    }
+} // namespace impactx::spacecharge

src/particles/spacecharge/GatherAndPush.H

@@ -0,0 +1,46 @@
+/* Copyright 2022 The Regents of the University of California, through Lawrence
+ *           Berkeley National Laboratory (subject to receipt of any required
+ *           approvals from the U.S. Dept. of Energy). All rights reserved.
+ *
+ * This file is part of ImpactX.
+ *
+ * Authors: Axel Huebl, Remi Lehe
+ * License: BSD-3-Clause-LBNL
+ */
+#ifndef IMPACTX_GATHER_AND_PUSH_H
+#define IMPACTX_GATHER_AND_PUSH_H
+
+#include "particles/ImpactXParticleContainer.H"
+
+#include <AMReX_Geometry.H>
+#include <AMReX_MultiFab.H>
+#include <AMReX_Vector.H>
+
+#include <unordered_map>
+#include <string>
+
+
+namespace impactx::spacecharge
+{
+    /** Gather force fields and push particles in x,y,z
+     *
+     * This gathers the space charge field with respect to particle position
+     * and shape. The momentum of all particles is then pushed using a common
+     * time step given by the reference particle speed and ds slice. The
+     * position push is done in the lattice elements and not here.
+     *
+     * @param[inout] pc container of the particles that deposited rho
+     * @param[in] space_charge_field space charge force component in x,y,z per level
+     * @param[in] geom geometry object
+     * @param[in] slice_ds segment length in meters
+     */
+    void GatherAndPush (
+        ImpactXParticleContainer & pc,
+        std::unordered_map<int, std::unordered_map<std::string, amrex::MultiFab> > const & space_charge_field,
+        const amrex::Vector<amrex::Geometry>& geom,
+        amrex::ParticleReal const slice_ds
+    );
+
+} // namespace impactx
+
+#endif // IMPACTX_GATHER_AND_PUSH_H