Gauss 3D Space Charge Pusher

src/initialization/Algorithms.H

@@ -21,6 +21,7 @@ namespace impactx
     AMREX_ENUM(SpaceChargeAlgo,
         False,    /**< Disabled: no space charge calculation */
         True_3D,  /**< 3D beam distribution */
+        Gauss3D,  /**< 3D Gaussian beam distribution */
         True_2D   /**< averaged 2D transverse beam distribution with a current along s */
     );
 

src/initialization/Algorithms.cpp

@@ -62,6 +62,10 @@ namespace impactx
         {
             return SpaceChargeAlgo::True_3D;
         }
+        else if (space_charge == "Gauss3D")
+        {
+            return SpaceChargeAlgo::Gauss3D;
+        }
         else if (space_charge == "2D")
         {
             return SpaceChargeAlgo::True_2D;

src/particles/spacecharge/CMakeLists.txt

@@ -1,6 +1,7 @@
 target_sources(lib
   PRIVATE
     ForceFromSelfFields.cpp
+    Gauss3dPush.cpp
     GatherAndPush.cpp
     HandleSpacecharge.cpp
     PoissonSolve.cpp

src/particles/spacecharge/Gauss3dPush.H

@@ -0,0 +1,40 @@
+/* Copyright 2022-2023 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: Ji Qiang
+ * License: BSD-3-Clause-LBNL
+ */
+#ifndef IMPACT_GAUSS3D_PUSH_H
+#define IMPACT_GAUSS3D_PUSH_H
+
+#include "particles/ImpactXParticleContainer.H"
+
+#include <unordered_map>
+#include <string>
+
+
+namespace impactx::particles::spacecharge
+{
+    /**------------------------
+     *
+     * This calculates the space charge fields from a 3D Gaussian distribution
+     * with respect to particle position from the following paper.
+     * "Two-and-a-half dimensional symplectic space-charge slover"
+     * 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] slice_ds segment length in meters
+     */
+    void Gauss3dPush (
+        ImpactXParticleContainer & pc,
+        amrex::ParticleReal slice_ds
+    );
+
+}  // namespace impactx::particles::spacecharge
+
+#endif // IMPACT_GAUSS3D_PUSH_H

src/particles/spacecharge/Gauss3dPush.cpp

@@ -0,0 +1,187 @@
+/* Copyright 2022-2023 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: Ji Qiang
+ * License: BSD-3-Clause-LBNL
+ */
+#include "Gauss3dPush.H"
+
+#include <AMReX_BLProfiler.H>
+#include <AMReX_REAL.H>       // for Real
+#include <AMReX_SPACE.H>      // for AMREX_D_DECL
+#include "diagnostics/ReducedBeamCharacteristics.H"
+
+
+namespace impactx::particles::spacecharge
+{
+// compute integrals Eqs 45-47 used in the space-charge fields from a 3D Gaussian distribution
+// Input particle locations (x,y,z) and RMS sizes (sigx,sigy,sigz) and return the integrals for SC fields.
+    void efldgauss (int nint, amrex::ParticleReal const xin, amrex::ParticleReal const yin, amrex::ParticleReal const zin,
+               amrex::ParticleReal const sigx, amrex::ParticleReal const sigy, amrex::ParticleReal const sigz, amrex::ParticleReal const gamma,
+               amrex::ParticleReal &pintex, amrex::ParticleReal &pintey, amrex::ParticleReal &pintez)
+    {
+      // enforce an odd grid points
+      if (nint % 2 == 0) nint += 1;
+
+      amrex::ParticleReal const xmin = 0.0, xmax = 1.0;
+      amrex::ParticleReal const h = (xmax - xmin) / (nint - 1);
+      amrex::ParticleReal const xp = xin / sigx;
+      amrex::ParticleReal const yp = yin / sigy;
+      amrex::ParticleReal const zp = zin / sigz;
+      amrex::ParticleReal const asp = sigx / sigy;
+      amrex::ParticleReal const basp = sigx / (sigz * gamma);
+
+      std::vector<amrex::ParticleReal> xvec(nint), t2vec(nint), t2sqrt(nint), bt2vec(nint), bt2sqrt(nint);
+      std::vector<amrex::ParticleReal> f2xvec(nint), f2yvec(nint), f2zvec(nint), f1vec(nint), fxvec(nint);
+
+      for (int i = 0; i < nint; ++i)
+        xvec[i] = xmin + i * h;
+
+      for (int i = 0; i < nint; ++i)
+        t2vec[i] = asp * asp + (1.0 - asp * asp) * xvec[i] * xvec[i];
+
+      for (int i = 0; i < nint; ++i)
+        t2sqrt[i] = std::sqrt(t2vec[i]);
+
+      for (int i = 0; i < nint; ++i)
+        bt2vec[i] = basp * basp + (1.0 - basp * basp) * xvec[i] * xvec[i];
+
+      for (int i = 0; i < nint; ++i)
+        bt2sqrt[i] = std::sqrt(bt2vec[i]);
+
+      for (int i = 0; i < nint; ++i)
+        f2xvec[i] = t2sqrt[i] * bt2sqrt[i];
+
+      for (int i = 0; i < nint; ++i)
+        f2yvec[i] = t2vec[i] * t2sqrt[i] * bt2sqrt[i];
+
+      for (int i = 0; i < nint; ++i)
+        f2zvec[i] = bt2vec[i] * t2sqrt[i] * bt2sqrt[i];
+
+      amrex::ParticleReal xp2 = xp * xp, yp2 = yp * yp, zp2 = zp * zp;
+      for (int i = 0; i < nint; ++i)
+        f1vec[i] = xvec[i] * xvec[i] * std::exp(-xvec[i] * xvec[i] * (xp2 + yp2 / t2vec[i] + zp2 / bt2vec[i]) / 2.0);
+
+      // Simpson's rule integration
+      amrex::ParticleReal sum0ex = 0.0, sum0ey = 0.0, sum0ez = 0.0;
+
+      // Ex
+      for (int i = 0; i < nint; ++i)
+        fxvec[i] = f1vec[i] / f2xvec[i];
+      for (int i = 0; i < nint; i += 2)
+        sum0ex += 2.0 * fxvec[i] * h;
+      for (int i = 1; i < nint; i += 2)
+        sum0ex += 4.0 * fxvec[i] * h;
+      pintex = sum0ex / 3.0;
+
+      // Ey
+      for (int i = 0; i < nint; ++i)
+        fxvec[i] = f1vec[i] / f2yvec[i];
+      for (int i = 0; i < nint; i += 2)
+        sum0ey += 2.0 * fxvec[i] * h;
+      for (int i = 1; i < nint; i += 2)
+        sum0ey += 4.0 * fxvec[i] * h;
+      pintey = sum0ey / 3.0;
+
+      // Ez
+      for (int i = 0; i < nint; ++i)
+        fxvec[i] = f1vec[i] / f2zvec[i];
+      for (int i = 0; i < nint; i += 2)
+        sum0ez += 2.0 * fxvec[i] * h;
+      for (int i = 1; i < nint; i += 2)
+        sum0ez += 4.0 * fxvec[i] * h;
+      pintez = sum0ez / 3.0;
+    }
+
+    void Gauss3dPush (
+        ImpactXParticleContainer & pc,
+        amrex::ParticleReal const slice_ds
+    )
+    {
+        BL_PROFILE("impactx::spacecharge::GatherAndPush");
+
+        using namespace amrex::literals;
+
+        amrex::ParticleReal const charge = pc.GetRefParticle().charge;
+
+        std::unordered_map<std::string, amrex::ParticleReal> const rbc = diagnostics::reduced_beam_characteristics(pc);
+        // get RMS sigma of beam
+        // Standard deviation of the particle positions (speed of light times time delay for t) (unit: meter)
+        amrex::ParticleReal const sigx = rbc.at("sig_x");
+        amrex::ParticleReal const sigy = rbc.at("sig_y");
+        amrex::ParticleReal const sigz = rbc.at("sig_t");
+        amrex::ParticleReal const bchchg = rbc.at("charge_C");
+
+        // physical constants and reference quantities
+        amrex::ParticleReal const c0_SI = 2.99792458e8;  // TODO move out
+        amrex::ParticleReal const mc_SI = pc.GetRefParticle().mass * c0_SI;
+        amrex::ParticleReal const pz_ref_SI = pc.GetRefParticle().beta_gamma() * mc_SI;
+        amrex::ParticleReal const gamma = pc.GetRefParticle().gamma();
+        amrex::ParticleReal const inv_gamma2 = 1.0_prt / (gamma * gamma);
+        amrex::ParticleReal const rfpiepslon = 1.0_prt /(c0_SI*c0_SI*1.0e-7);
+
+        amrex::ParticleReal const dt = slice_ds / pc.GetRefParticle().beta() / c0_SI;
+
+        // group together constants for the momentum push
+        amrex::ParticleReal const push_consts = dt * charge * inv_gamma2 / pz_ref_SI;
+
+        // loop over refinement levels
+        int const nLevel = pc.finestLevel();
+        for (int lev = 0; lev <= nLevel; ++lev)
+        {
+
+            // loop over all particle boxes
+            using ParIt = ImpactXParticleContainer::iterator;
+#ifdef AMREX_USE_OMP
+#pragma omp parallel if (amrex::Gpu::notInLaunchRegion())
+#endif
+            for (ParIt pti(pc, lev); pti.isValid(); ++pti) {
+                const int np = pti.numParticles();
+
+                // preparing access to particle data: SoA of Reals
+                auto& soa_real = pti.GetStructOfArrays().GetRealData();
+                amrex::ParticleReal* const AMREX_RESTRICT part_x = soa_real[RealSoA::x].dataPtr();
+                amrex::ParticleReal* const AMREX_RESTRICT part_y = soa_real[RealSoA::y].dataPtr();
+                amrex::ParticleReal* const AMREX_RESTRICT part_z = soa_real[RealSoA::z].dataPtr(); // note: currently for a fixed t
+                amrex::ParticleReal* const AMREX_RESTRICT part_px = soa_real[RealSoA::px].dataPtr();
+                amrex::ParticleReal* const AMREX_RESTRICT part_py = soa_real[RealSoA::py].dataPtr();
+                amrex::ParticleReal* const AMREX_RESTRICT part_pz = soa_real[RealSoA::pz].dataPtr(); // note: currently for a fixed t
+
+                // group together constants for the momentum push
+                amrex::ParticleReal const push_consts = dt * charge * inv_gamma2 / pz_ref_SI;
+
+                // gather to each particle and push momentum
+                amrex::ParallelFor(np, [=] AMREX_GPU_DEVICE (int i) {
+                    // access SoA Real data
+                    amrex::ParticleReal & AMREX_RESTRICT x = part_x[i];
+                    amrex::ParticleReal & AMREX_RESTRICT y = part_y[i];
+                    amrex::ParticleReal & AMREX_RESTRICT z = part_z[i];
+                    amrex::ParticleReal & AMREX_RESTRICT px = part_px[i];
+                    amrex::ParticleReal & AMREX_RESTRICT py = part_py[i];
+                    amrex::ParticleReal & AMREX_RESTRICT pz = part_pz[i];
+
+                    //field integrals from a 3D Gaussian bunch
+                    int const nint = 401;
+                    amrex::ParticleReal eintx, einty, eintz;
+                    efldgauss(nint,x,y,z,sigx,sigy,sigz,gamma,eintx,einty,eintz);
+
+                    amrex::ParticleReal const asp = sigx/sigy;
+
+                    // push momentum
+                    using ablastr::constant::math::pi;
+                    px += x*eintx*2*asp/sigx/sigx/sigz/sqrt(2*pi)*bchchg*rfpiepslon * push_consts;
+                    py += y*einty*2*asp/sigy/sigy/sigz/sqrt(2*pi)*bchchg*rfpiepslon * push_consts;
+                    pz += z*eintz*2*asp/sigz/sigz/sigz/sqrt(2*pi)*bchchg*rfpiepslon * push_consts;
+
+                    // push position is done in the lattice elements
+                });
+
+
+            } // end loop over all particle boxes
+        } // env mesh-refinement level loop
+    }
+
+}  // namespace impactx::particles::spacecharge

src/particles/spacecharge/HandleSpacecharge.cpp

@@ -14,6 +14,7 @@
 #include "particles/ImpactXParticleContainer.H"
 #include "particles/spacecharge/ForceFromSelfFields.H"
 #include "particles/spacecharge/GatherAndPush.H"
+#include "particles/spacecharge/Gauss3dPush.H"
 #include "particles/spacecharge/PoissonSolve.H"
 #include "particles/transformation/CoordinateTransformation.H"
 
@@ -47,45 +48,56 @@ namespace impactx::particles::spacecharge
             CoordSystem::t
         );
 
-        // Note: The following operation assume that
-        // the particles are in x, y, z coordinates.
+        if (space_charge == SpaceChargeAlgo::Gauss3D) {
 
-        // Resize the mesh, based on `amr_data->track_particles.m_particle_container` extent
-        ResizeMesh();
+            // update momentum assming a 3D Gaussian bunch
+            Gauss3dPush(*amr_data->track_particles.m_particle_container, slice_ds);
 
-        // Redistribute particles in the new mesh in x, y, z
-        amr_data->track_particles.m_particle_container->Redistribute();
+        }
+        else
+        {
 
-        // charge deposition
-        amr_data->track_particles.m_particle_container->DepositCharge(
+          // 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(
+          // 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(
+          // 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(
+          // 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
-        );
+            slice_ds);
+
+        }
+
 
         // transform from x,y,z to x',y',t
         transformation::CoordinateTransformation(