Initial commit for 3D space charge envelope routines.

src/envelope/spacecharge/Elliptic_Integral.H

@@ -0,0 +1,162 @@
+/* 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 ELLIPTIC_INTEGRAL_H
+#define ELLIPTIC_INTEGRAL_H
+
+#include <ablastr/warn_manager/WarnManager.H>
+
+#include <AMReX_Extension.H>
+#include <AMReX_REAL.H>
+#include <AMReX_Print.H>
+
+#include <cmath>
+
+namespace impactx::envelope::spacecharge
+{
+
+    /** Function to return the Carlson elliptic integral denoted
+     *  by RD(x,y,z).  This is a slight modification of
+     *  the C++ function written by John Burkardt (GNU LGPL license):
+     *
+     *  https://people.math.sc.edu/Burkardt/cpp_src
+     *        /elliptic_integral/elliptic_integral.html
+     *
+     *  which in turn is based on:
+     *
+     * Bille Carlson,
+     * Computing Elliptic Integrals by Duplication,
+     * Numerische Mathematik,
+     * Volume 33, 1979, pages 1-16.
+     *    
+     * Bille Carlson, Elaine Notis,
+     * Algorithm 577, Algorithms for Incomplete Elliptic Integrals,
+     * ACM Transactions on Mathematical Software,
+     * Volume 7, Number 3, pages 398-403, September 1981.
+     * 
+     * @param[in] x first argument
+     * @param[in] y second argument
+     * @param[in] z third argument
+     * @param[in] errtol error tolerance
+     * @returns the elliptic integral RD(x,y,z)
+     */
+    amrex::ParticleReal
+    Elliptic_RD (
+        amrex::ParticleReal x,   
+        amrex::ParticleReal y,
+        amrex::ParticleReal z,
+        amrex::ParticleReal errtol
+    )
+    {
+        using namespace amrex::literals;
+
+        amrex::ParticleReal c1;
+        amrex::ParticleReal c2;
+        amrex::ParticleReal c3;
+        amrex::ParticleReal c4;
+        amrex::ParticleReal ea;
+        amrex::ParticleReal eb;
+        amrex::ParticleReal ec;
+        amrex::ParticleReal ed;
+        amrex::ParticleReal ef;
+        amrex::ParticleReal epslon;
+        amrex::ParticleReal lamda;
+        amrex::ParticleReal const lolim = 6.0E-51;
+        amrex::ParticleReal mu;
+        amrex::ParticleReal power4;
+        amrex::ParticleReal sigma;
+        amrex::ParticleReal s1;
+        amrex::ParticleReal s2;
+        amrex::ParticleReal const uplim = 1.0E+48;
+        amrex::ParticleReal value;
+        amrex::ParticleReal xn;
+        amrex::ParticleReal xndev;
+        amrex::ParticleReal xnroot;
+        amrex::ParticleReal yn;
+        amrex::ParticleReal yndev;
+        amrex::ParticleReal ynroot;
+        amrex::ParticleReal zn;
+        amrex::ParticleReal zndev;
+        amrex::ParticleReal znroot;
+      //
+      //   LOLIM and UPLIM determine the range of valid arguments.
+      //   LOLIM IS NOT LESS THAN 2 / (MACHINE MAXIMUM) ^ (2/3).
+      //   UPLIM IS NOT GREATER THAN (0.1 * ERRTOL / MACHINE
+      //   MINIMUM) ^ (2/3), WHERE ERRTOL IS DESCRIBED BELOW.
+      //   IN THE FOLLOWING TABLE IT IS ASSUMED THAT ERRTOL WILL
+      //   NEVER BE CHOSEN SMALLER THAN 1.D-5.
+      //
+        if ( 
+          x < 0.0 || 
+          y < 0.0 || 
+          x + y < lolim || 
+          z < lolim || 
+          uplim < x || 
+          uplim < y || 
+          uplim < z )
+        {
+          amrex::Print() << "\n";
+          amrex::Print() << "RD - Error!\n";
+          amrex::Print() << "  Invalid input arguments.\n";
+          amrex::Print() << "  X = " << x << "\n";
+          amrex::Print() << "  Y = " << y << "\n";
+          amrex::Print() << "  Z = " << z << "\n";
+          amrex::Print() << "\n";
+          value = 0.0_prt;
+          return value;
+        }
+
+        xn = x;
+        yn = y;
+        zn = z;
+        sigma = 0.0_prt;
+        power4 = 1.0_prt;
+
+        while ( true )
+        {
+          mu = ( xn + yn + 3.0_prt * zn ) * 0.2_prt;
+          xndev = ( mu - xn ) / mu;
+          yndev = ( mu - yn ) / mu;
+          zndev = ( mu - zn ) / mu;
+          epslon = std::max ( std::abs ( xndev ), 
+            std::max ( std::abs ( yndev ), std::abs ( zndev ) ) );
+
+          if ( epslon < errtol )
+          {
+            c1 = 3.0_prt / 14.0_prt;
+            c2 = 1.0_prt / 6.0_prt;
+            c3 = 9.0_prt / 22.0_prt;
+            c4 = 3.0_prt / 26.0_prt;
+            ea = xndev * yndev;
+            eb = zndev * zndev;
+            ec = ea - eb;
+            ed = ea - 6.0_prt * eb;
+            ef = ed + ec + ec;
+            s1 = ed * ( - c1 + 0.25_prt * c3 * ed - 1.5_prt * c4 * zndev * ef );
+            s2 = zndev  * ( c2 * ef + zndev * ( - c3 * ec + zndev * c4 * ea ) );
+            value = 3.0_prt * sigma  + power4 * ( 1.0_prt + s1 + s2 ) / ( mu * std::sqrt ( mu ) );
+            return value;
+          }
+
+          xnroot = std::sqrt ( xn );
+          ynroot = std::sqrt ( yn );
+          znroot = std::sqrt ( zn );
+          lamda = xnroot * ( ynroot + znroot ) + ynroot * znroot;
+          sigma = sigma + power4 / ( znroot * ( zn + lamda ) );
+          power4 = power4 * 0.25_prt;
+          xn = ( xn + lamda ) * 0.25_prt;
+          yn = ( yn + lamda ) * 0.25_prt;
+          zn = ( zn + lamda ) * 0.25_prt;
+        }
+
+      }
+
+} // namespace impactx::envelope::spacecharge
+
+#endif // ELLIPTIC_INTEGRAL_H

src/envelope/spacecharge/EnvelopeSpaceChargePush.H

@@ -37,6 +37,24 @@ namespace impactx::envelope::spacecharge
          amrex::ParticleReal ds
     );
 
+     /** This function pushes the 6x6 beam covariance matrix for a slice
+      *  of length ds, using the linear space charge fields in an rms
+      *  equivalent 3D ellipsoid, as determined from the beam covariance matrix.
+      *  Note: This is a reduced model of 3D space charge.
+      *
+      * @param[in] refpart reference particle
+      * @param[inout] cm covariance matrix
+      * @param[in] bunch_charge bunch charge [C]
+      * @param[in] ds step size [m]
+      */
+     void
+     space_charge3D_push (
+         RefPart const & AMREX_RESTRICT refpart,
+         Map6x6 & AMREX_RESTRICT cm,
+         amrex::ParticleReal current,
+         amrex::ParticleReal ds
+    );
+
 } // namespace impactx
 
 #endif // IMPACTX_ENVELOPESPACECHARGEPUSH_H

src/envelope/spacecharge/EnvelopeSpaceChargePush.cpp

@@ -8,6 +8,7 @@
  * License: BSD-3-Clause-LBNL
  */
 #include "EnvelopeSpaceChargePush.H"
+#include "Elliptic_Integral.H"
 
 #include <ablastr/warn_manager/WarnManager.H>
 
@@ -67,5 +68,59 @@ namespace impactx::envelope::spacecharge
 
     }
 
+    void
+    space_charge3D_push (
+        RefPart const & AMREX_RESTRICT refpart,
+        Map6x6 & AMREX_RESTRICT cm,
+        amrex::ParticleReal bunch_charge,
+        amrex::ParticleReal ds
+    )
+    {
+        using namespace amrex::literals;
+
+        // skip calculations for trivial case
+        if (bunch_charge == 0_prt) { return; }
+
+        // initialize the linear transport map
+        Map6x6 R = Map6x6::Identity();
+
+        // physical constants and reference quantities
+        using ablastr::constant::SI::c;
+        using ablastr::constant::SI::ep0;
+        using ablastr::constant::math::pi;
+        amrex::ParticleReal const mass = refpart.mass;
+        amrex::ParticleReal const charge = refpart.charge;
+        amrex::ParticleReal const pt_ref = refpart.pt;
+        amrex::ParticleReal const betgam2 = std::pow(pt_ref, 2) - 1_prt;
+        amrex::ParticleReal const betgam = std::sqrt(betgam2);
+
+        // evaluate the 3D space charge intensity parameter from bunch charge
+        amrex::ParticleReal const rcN = charge * bunch_charge / (4_prt * pi * ep0 * mass * std::pow(c,2));
+        amrex::ParticleReal const coeff = ds * rcN / betgam2 * (1_prt/(5_prt * std::sqrt(5_prt)));
+
+        // set parameters for elliptic integrals
+        amrex::ParticleReal const errtol = 1.0e-3;
+        amrex::ParticleReal const x = cm(1,1);
+        amrex::ParticleReal const y = cm(3,3);
+        amrex::ParticleReal const z = betgam * cm(5,5);
+
+        // evaluate the off-identity elements of the linear transfer map
+        R(2,1) = coeff * Elliptic_RD(y,z,x,errtol);  
+        R(4,3) = coeff * Elliptic_RD(z,x,y,errtol);
+        R(6,5) = coeff * Elliptic_RD(x,y,z,errtol);
+
+        // update the beam covariance matrix
+        cm = R * cm * R.transpose();
+
+        // test of elliptic integral evaluation only   
+        //amrex::ParticleReal const x = 1.0;
+        //amrex::ParticleReal const y = 5.0;
+        //amrex::ParticleReal const z = 0.02;
+        //amrex::ParticleReal const errtol = 1.0e-3;
+        //amrex::ParticleReal rd = Elliptic_RD(x,y,z,errtol);
+        //amrex::Print() << "RD( " << x << "," << y << "," << z << " ) = " << rd << "\n";
+
+    }
+
 
 } // namespace impactx::spacecharge