Introduce HandleSpacecharge (#842)

* Introduce `HandleSpacecharge` Similar to our `HandleWakefield` logic for particle tracking, this simplifies the central particle tracking loop by cleanly hiding implementation details in a separate function. * Fix Doxygen String Co-authored-by: Chad Mitchell <[email protected]> --------- Co-authored-by: Chad Mitchell <[email protected]>
BLAST-ImpactX · Feb 12, 2025 · 5816acc · 5816acc
1 parent 220531e
commit 5816acc
Show file tree

Hide file tree

Showing 10 changed files with 162 additions and 81 deletions.
diff --git a/src/particles/spacecharge/CMakeLists.txt b/src/particles/spacecharge/CMakeLists.txt
@@ -2,5 +2,6 @@ target_sources(lib
   PRIVATE
     ForceFromSelfFields.cpp
     GatherAndPush.cpp
+    HandleSpacecharge.cpp
     PoissonSolve.cpp
 )
diff --git a/src/particles/spacecharge/ForceFromSelfFields.H b/src/particles/spacecharge/ForceFromSelfFields.H
@@ -19,7 +19,7 @@
 #include <unordered_map>
 
 
-namespace impactx::spacecharge
+namespace impactx::particles::spacecharge
 {
     /** Calculate the space charge force field from the electric potential
      *
@@ -36,6 +36,6 @@ namespace impactx::spacecharge
         const amrex::Vector<amrex::Geometry>& geom
     );
 
-} // namespace impactx
+}  // namespace impactx::particles::spacecharge
 
 #endif // IMPACTX_FORCEFROMSELFFIELDS_H
diff --git a/src/particles/spacecharge/ForceFromSelfFields.cpp b/src/particles/spacecharge/ForceFromSelfFields.cpp
@@ -14,7 +14,7 @@
 #include <AMReX_SPACE.H>      // for AMREX_D_DECL
 
 
-namespace impactx::spacecharge
+namespace impactx::particles::spacecharge
 {
     void ForceFromSelfFields (
         std::unordered_map<int, std::unordered_map<std::string, amrex::MultiFab> > & space_charge_field,
@@ -60,4 +60,4 @@ namespace impactx::spacecharge
             }
         }
     }
-} // namespace impactx::spacecharge
+}  // namespace impactx::particles::spacecharge
diff --git a/src/particles/spacecharge/GatherAndPush.H b/src/particles/spacecharge/GatherAndPush.H
@@ -20,7 +20,7 @@
 #include <string>
 
 
-namespace impactx::spacecharge
+namespace impactx::particles::spacecharge
 {
     /** Gather force fields and push particles in x,y,z
      *
@@ -41,6 +41,6 @@ namespace impactx::spacecharge
         amrex::ParticleReal slice_ds
     );
 
-} // namespace impactx
+}  // namespace impactx::particles::spacecharge
 
 #endif // IMPACTX_GATHER_AND_PUSH_H
diff --git a/src/particles/spacecharge/GatherAndPush.cpp b/src/particles/spacecharge/GatherAndPush.cpp
@@ -16,7 +16,7 @@
 #include <AMReX_SPACE.H>      // for AMREX_D_DECL
 
 
-namespace impactx::spacecharge
+namespace impactx::particles::spacecharge
 {
     void GatherAndPush (
         ImpactXParticleContainer & pc,
@@ -105,4 +105,4 @@ namespace impactx::spacecharge
             } // end loop over all particle boxes
         } // env mesh-refinement level loop
     }
-} // namespace impactx::spacecharge
+}  // namespace impactx::particles::spacecharge
diff --git a/src/particles/spacecharge/HandleSpacecharge.H b/src/particles/spacecharge/HandleSpacecharge.H
@@ -0,0 +1,38 @@
+/* Copyright 2022-2025 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, Chad Mitchell
+ * License: BSD-3-Clause-LBNL
+ */
+#ifndef IMPACTX_HANDLE_SPACECHARGE_H
+#define IMPACTX_HANDLE_SPACECHARGE_H
+
+#include "initialization/AmrCoreData.H"
+
+#include <AMReX_REAL.H>
+
+#include <functional>
+#include <memory>
+
+
+namespace impactx::particles::spacecharge
+{
+    /** Function to handle space charge including charge deposition, Poisson solve,
+     * field calculation, interpolation, and particle push.
+     *
+     * @param[in,out] amr_data AMR data like particle container and fields
+     * @param[in] ResizeMesh function to call to resize the mesh
+     * @param[in] slice_ds slice spacing along s
+     */
+    void HandleSpacecharge (
+        std::unique_ptr<initialization::AmrCoreData> & amr_data,
+        std::function<void()> ResizeMesh,
+        amrex::Real slice_ds
+    );
+
+} // namespace impactx::particles::spacecharge
+
+#endif // IMPACTX_HANDLE_SPACECHARGE_H
diff --git a/src/particles/spacecharge/HandleSpacecharge.cpp b/src/particles/spacecharge/HandleSpacecharge.cpp
@@ -0,0 +1,108 @@
+/* Copyright 2022-2025 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, Chad Mitchell
+ * License: BSD-3-Clause-LBNL
+ */
+#include "HandleSpacecharge.H"
+
+#include "initialization/AmrCoreData.H"
+#include "particles/ImpactXParticleContainer.H"
+#include "particles/spacecharge/ForceFromSelfFields.H"
+#include "particles/spacecharge/GatherAndPush.H"
+#include "particles/spacecharge/PoissonSolve.H"
+#include "particles/transformation/CoordinateTransformation.H"
+
+#include <AMReX_ParmParse.H>
+#include <AMReX_REAL.H>
+
+#include <memory>
+
+
+namespace impactx::particles::spacecharge
+{
+    void HandleSpacecharge (
+        std::unique_ptr<initialization::AmrCoreData> & amr_data,
+        std::function<void()> ResizeMesh,
+        amrex::Real slice_ds
+    )
+    {
+        BL_PROFILE("impactx::particles::wakefields::HandleSpacecharge")
+
+        amrex::ParmParse const pp_algo("algo");
+        bool space_charge = false;
+        pp_algo.query("space_charge", space_charge);
+
+        // turn off if disabled by user
+        if (!space_charge) { return; }
+
+        // turn off if less than 2 particles
+        if (amr_data->track_particles.m_particle_container->TotalNumberOfParticles(true, false) < 2) { return; }
+
+        // transform from x',y',t to x,y,z
+        transformation::CoordinateTransformation(
+            *amr_data->track_particles.m_particle_container,
+            CoordSystem::t
+        );
+
+        // Note: The following operation assume that
+        // the particles are in x, y, z coordinates.
+
+        // Resize the mesh, based on `amr_data->track_particles.m_particle_container` extent
+        ResizeMesh();
+
+        // Redistribute particles in the new mesh in x, y, z
+        amr_data->track_particles.m_particle_container->Redistribute();
+
+        // charge deposition
+        amr_data->track_particles.m_particle_container->DepositCharge(
+            amr_data->track_particles.m_rho,
+            amr_data->refRatio()
+        );
+
+        // poisson solve in x,y,z
+        spacecharge::PoissonSolve(
+            *amr_data->track_particles.m_particle_container,
+            amr_data->track_particles.m_rho,
+            amr_data->track_particles.m_phi,
+            amr_data->refRatio()
+        );
+
+        // calculate force in x,y,z
+        spacecharge::ForceFromSelfFields(
+            amr_data->track_particles.m_space_charge_field,
+            amr_data->track_particles.m_phi,
+            amr_data->Geom()
+        );
+
+        // gather and space-charge push in x,y,z , assuming the space-charge
+        // field is the same before/after transformation
+        // TODO: This is currently using linear order.
+        spacecharge::GatherAndPush(
+            *amr_data->track_particles.m_particle_container,
+            amr_data->track_particles.m_space_charge_field,
+            amr_data->Geom(),
+            slice_ds
+        );
+
+        // transform from x,y,z to x',y',t
+        transformation::CoordinateTransformation(
+            *amr_data->track_particles.m_particle_container,
+            CoordSystem::s
+        );
+
+        // for later: original Impact implementation as an option
+        // Redistribute particles in x',y',t
+        //   TODO: only needed if we want to gather and push space charge
+        //         in x',y',t
+        //   TODO: change geometry beforehand according to transformation
+        //m_amr_data->track_particles.m_particle_container->Redistribute();
+        //
+        // in original Impact, we gather and space-charge push in x',y',t ,
+        // assuming that the distribution did not change
+    }
+
+} // namespace impactx::particles::spacecharge
diff --git a/src/particles/spacecharge/PoissonSolve.H b/src/particles/spacecharge/PoissonSolve.H
@@ -17,7 +17,7 @@
 #include <unordered_map>
 
 
-namespace impactx::spacecharge
+namespace impactx::particles::spacecharge
 {
     /** Calculate the electric potential from charge density
      *
@@ -36,6 +36,6 @@ namespace impactx::spacecharge
         amrex::Vector<amrex::IntVect> rel_ref_ratio
     );
 
-} // namespace impactx
+} // namespace impactx::particles::spacecharge
 
 #endif // IMPACTX_POISSONSOLVE_H
diff --git a/src/particles/spacecharge/PoissonSolve.cpp b/src/particles/spacecharge/PoissonSolve.cpp
@@ -21,7 +21,7 @@
 #include <cmath>
 
 
-namespace impactx::spacecharge
+namespace impactx::particles::spacecharge
 {
     void PoissonSolve (
         ImpactXParticleContainer const & pc,
@@ -121,4 +121,4 @@ namespace impactx::spacecharge
             phi_at_level.FillBoundary(pc.GetParGDB()->Geom()[lev].periodicity());
         }
     }
-} // impactx::spacecharge
+}  // impactx::particles::spacecharge
diff --git a/src/tracking/particles.cpp b/src/tracking/particles.cpp
@@ -13,10 +13,7 @@
 #include "particles/ImpactXParticleContainer.H"
 #include "particles/Push.H"
 #include "diagnostics/DiagnosticOutput.H"
-#include "particles/spacecharge/ForceFromSelfFields.H"
-#include "particles/spacecharge/GatherAndPush.H"
-#include "particles/spacecharge/PoissonSolve.H"
-#include "particles/transformation/CoordinateTransformation.H"
+#include "particles/spacecharge/HandleSpacecharge.H"
 #include "particles/wakefields/HandleWakefield.H"
 
 #include <AMReX.H>
@@ -114,71 +111,8 @@ namespace impactx
                     // Wakefield calculation: call wakefield function to apply wake effects
                     particles::wakefields::HandleWakefield(*amr_data->track_particles.m_particle_container, element_variant, slice_ds);
 
-                    // Space-charge calculation: turn off if there is only 1 particle
-                    if (space_charge &&
-                        amr_data->track_particles.m_particle_container->TotalNumberOfParticles(true, false)) {
-
-                        // transform from x',y',t to x,y,z
-                        transformation::CoordinateTransformation(
-                                *amr_data->track_particles.m_particle_container,
-                                CoordSystem::t);
-
-                        // Note: The following operation assume that
-                        // the particles are in x, y, z coordinates.
-
-                        // Resize the mesh, based on `m_particle_container` extent
-                        ResizeMesh();
-
-                        // Redistribute particles in the new mesh in x, y, z
-                        amr_data->track_particles.m_particle_container->Redistribute();
-
-                        // charge deposition
-                        amr_data->track_particles.m_particle_container->DepositCharge(
-                            amr_data->track_particles.m_rho,
-                            amr_data->refRatio()
-                        );
-
-                        // poisson solve in x,y,z
-                        spacecharge::PoissonSolve(
-                            *amr_data->track_particles.m_particle_container,
-                            amr_data->track_particles.m_rho,
-                            amr_data->track_particles.m_phi,
-                            amr_data->refRatio()
-                        );
-
-                        // calculate force in x,y,z
-                        spacecharge::ForceFromSelfFields(
-                            amr_data->track_particles.m_space_charge_field,
-                            amr_data->track_particles.m_phi,
-                            amr_data->Geom()
-                        );
-
-                        // gather and space-charge push in x,y,z , assuming the space-charge
-                        // field is the same before/after transformation
-                        // TODO: This is currently using linear order.
-                        spacecharge::GatherAndPush(
-                            *amr_data->track_particles.m_particle_container,
-                            amr_data->track_particles.m_space_charge_field,
-                            amr_data->Geom(),
-                            slice_ds
-                        );
-
-                        // transform from x,y,z to x',y',t
-                        transformation::CoordinateTransformation(
-                            *amr_data->track_particles.m_particle_container,
-                            CoordSystem::s
-                        );
-                    }
-
-                    // for later: original Impact implementation as an option
-                    // Redistribute particles in x',y',t
-                    //   TODO: only needed if we want to gather and push space charge
-                    //         in x',y',t
-                    //   TODO: change geometry beforehand according to transformation
-                    //m_particle_container->Redistribute();
-                    //
-                    // in original Impact, we gather and space-charge push in x',y',t ,
-                    // assuming that the distribution did not change
+                    // Space-charge calculation
+                    particles::spacecharge::HandleSpacecharge(amr_data, [this](){ this->ResizeMesh(); }, slice_ds);
 
                     // push all particles with external maps
                     Push(*amr_data->track_particles.m_particle_container, element_variant, step, period);