Space charge 2D envelope push

src/particles/spacecharge/CMakeLists.txt

@@ -3,4 +3,5 @@ target_sources(lib
     ForceFromSelfFields.cpp
     GatherAndPush.cpp
     PoissonSolve.cpp
+    EnvelopeSpaceChargePush.cpp
 )

src/particles/spacecharge/EnvelopeSpaceChargePush.H

@@ -0,0 +1,41 @@
+/* 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: Chad Mitchell, Axel Huebl
+ * License: BSD-3-Clause-LBNL
+ */
+#ifndef IMPACTX_ENVELOPESPACECHARGEPUSH_H
+#define IMPACTX_ENVELOPESPACECHARGEPUSH_H
+
+#include "particles/ImpactXParticleContainer.H"
+#include "particles/CovarianceMatrix.H"
+
+#include <AMReX_REAL.H>
+
+
+namespace impactx::spacecharge
+{
+     /** This function pushes the 6x6 beam covariance matrix for a slice
+      *  of length ds, using the linear space charge fields in an rms
+      *  equivalent 2D ellipse, as determined from the beam covariance matrix.
+      *  Note: This is a reduced model of 2D space charge.
+      *
+      * @param[in] refpart reference particle
+      * @param[inout] cm covariance matrix
+      * @param[in] current beam current [A]
+      * @param[in] ds step size [m]
+      */
+     void
+     envelope_space_charge2D_push (
+         RefPart const & refpart,
+         Map6x6 & cm,
+         amrex::ParticleReal & current,
+         amrex::ParticleReal & ds
+    );
+
+} // namespace impactx
+
+#endif // IMPACTX_ENVELOPESPACECHARGEPUSH_H

src/particles/spacecharge/EnvelopeSpaceChargePush.cpp

@@ -0,0 +1,68 @@
+/* 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: Marco Garten, Axel Huebl
+ * License: BSD-3-Clause-LBNL
+ */
+#include "EnvelopeSpaceChargePush.H"
+
+#include <AMReX_BLProfiler.H>
+#include <AMReX_REAL.H>       // for Real
+#include <AMReX_SmallMatrix.H>
+
+#include <cmath>
+
+namespace impactx::spacecharge
+{
+        /** This function returns the linear transport map associated with a
+         *  reduced 2D space charge model, based on the beam covariance matrix.
+         *
+         * @param[in] refpart reference particle
+         * @param[in,out] cm covariance matrix
+         * @param[in] current beam current [A]
+         * @param[in] ds step size [m]
+         * @returns 6x6 transport matrix
+         */
+        AMREX_GPU_HOST AMREX_FORCE_INLINE
+        void
+        envelope_space_charge2D_push (
+        RefPart const & refpart,
+        Map6x6 & cm,
+        amrex::ParticleReal & current,
+        amrex::ParticleReal & ds
+    )
+    {
+        using namespace amrex::literals;
+
+        // initialize the linear transport map
+        Map6x6 R = Map6x6::Identity();
+
+        // access reference particle values to find beta*gamma^2
+        amrex::ParticleReal const pt_ref = refpart.pt;
+        amrex::ParticleReal const betgam2 = std::pow(pt_ref, 2) - 1.0_prt;
+
+        // evaluate the beam space charge perveance from current
+        // TODO
+        amrex::ParticleReal const Kpv=0.0_prt;
+
+        // evaluate the linear transfer map
+        amrex::ParticleReal const sigma2 = cm(1,1)*cm(3,3)-cm(1,3)*cm(1,3);
+        amrex::ParticleReal const sigma = std::sqrt(sigma2);
+        amrex::ParticleReal const D = (sigma + cm(1,1)) * (sigma + cm(3,3)) - cm(1,3)*cm(1,3);
+        amrex::ParticleReal const coeff = Kpv*ds/(2.0_prt*D);
+
+        R(2,1) = coeff * (sigma + cm(3,3));
+        R(2,3) = coeff * (-cm(1,3));
+        R(4,1) = R(2,3);
+        R(4,3) = coeff * (sigma + cm(1,1));
+
+        // update the beam covariance matrix
+        cm = R * cm * R.transpose();
+
+    }
+
+
+} // namespace impactx::spacecharge

src/tracking/envelope.cpp

@@ -11,6 +11,7 @@
 #include "initialization/InitAmrCore.H"
 #include "particles/ImpactXParticleContainer.H"
 #include "particles/Push.H"
+#include "particles/spacecharge/EnvelopeSpaceChargePush.H"
 #include "diagnostics/DiagnosticOutput.H"
 
 #include <AMReX.H>
@@ -77,7 +78,7 @@ namespace impactx
         amrex::ParmParse const pp_algo("algo");
         bool space_charge = false;
         pp_algo.query("space_charge", space_charge);
-        AMREX_ALWAYS_ASSERT_WITH_MESSAGE(!space_charge, "Space charge not yet implemented for envelope tracking.");
+        //AMREX_ALWAYS_ASSERT_WITH_MESSAGE("2D space charge only is implemented for envelope tracking.");
         if (verbose > 0)
         {
             amrex::Print() << " Space Charge effects: " << space_charge << "\n";
@@ -123,15 +124,22 @@ namespace impactx
                                        << " slice_step=" << slice_step << "\n";
                     }
 
-                    std::visit([&ref, &cm](auto&& element)
+                    std::visit([&ref, &cm, &slice_ds, &space_charge](auto&& element)
                     {
                         // push reference particle in global coordinates
                         {
                             BL_PROFILE("impactx::Push::RefPart");
                             element(ref);
                         }
 
-                        // push Covariance Matrix
+                        if (space_charge)
+                        {
+                            // push Covariance Matrix in 2D space charge fields
+                            amrex::ParticleReal current=0.0;  //TODO: This must be set.
+                            spacecharge::envelope_space_charge2D_push(ref,cm,current,slice_ds);
+                        }
+
+                        // push Covariance Matrix in external fields
                         element(cm, ref);
 
                     }, element_variant);