Solenoid lattice element

Source/AcceleratorLattice/AcceleratorLattice.H

@@ -11,6 +11,8 @@
 #include "LatticeElements/Drift.H"
 #include "LatticeElements/HardEdgedQuadrupole.H"
 #include "LatticeElements/HardEdgedPlasmaLens.H"
+#include "LatticeElements/Solenoid.H"
+#include "LatticeElements/SolenoidRF.H"
 
 #include <memory>
 #include <string>
@@ -73,6 +75,8 @@ public:
     Drift h_drift;
     HardEdgedQuadrupole h_quad;
     HardEdgedPlasmaLens h_plasmalens;
+    Solenoid h_solenoid;
+    SolenoidRF h_solenoidrf;
 
 };
 

Source/AcceleratorLattice/AcceleratorLattice.cpp

@@ -11,6 +11,8 @@
 #include "LatticeElements/Drift.H"
 #include "LatticeElements/HardEdgedQuadrupole.H"
 #include "LatticeElements/HardEdgedPlasmaLens.H"
+#include "LatticeElements/Solenoid.H"
+#include "LatticeElements/SolenoidRF.H"
 
 #include <AMReX_REAL.H>
 
@@ -28,6 +30,8 @@ AcceleratorLattice::AcceleratorLattice ()
 
     h_quad.WriteToDevice();
     h_plasmalens.WriteToDevice();
+    h_solenoid.WriteToDevice();
+    h_solenoidrf.WriteToDevice();
 }
 
 void
@@ -65,6 +69,12 @@ AcceleratorLattice::ReadLattice (std::string const & root_name, amrex::ParticleR
         else if (element_type == "plasmalens") {
             h_plasmalens.AddElement(pp_element, z_location);
         }
+        else if (element_type == "solenoid") {
+            h_solenoid.AddElement(pp_element, z_location);
+        }
+        else if (element_type == "solenoidrf") {
+            h_solenoidrf.AddElement(pp_element, z_location);
+        }
         else if (element_type == "line") {
             ReadLattice(element_name, z_location);
         }

Source/AcceleratorLattice/LatticeElementFinder.H

@@ -9,6 +9,8 @@
 
 #include "LatticeElements/HardEdgedQuadrupole.H"
 #include "LatticeElements/HardEdgedPlasmaLens.H"
+#include "LatticeElements/Solenoid.H"
+#include "LatticeElements/SolenoidRF.H"
 #include "Particles/Pusher/GetAndSetPosition.H"
 #include "Particles/WarpXParticleContainer.H"
 
@@ -78,6 +80,8 @@ struct LatticeElementFinder
     /* The index lookup tables for each lattice element type */
     amrex::Gpu::DeviceVector<int> d_quad_indices;
     amrex::Gpu::DeviceVector<int> d_plasmalens_indices;
+    amrex::Gpu::DeviceVector<int> d_solenoid_indices;
+    amrex::Gpu::DeviceVector<int> d_solenoidrf_indices;
 
     /**
      * \brief Fill in the index lookup tables
@@ -130,10 +134,14 @@ struct LatticeElementFinderDevice
     /* Device level instances for each lattice element type */
     HardEdgedQuadrupoleDevice d_quad;
     HardEdgedPlasmaLensDevice d_plasmalens;
+    SolenoidDevice d_solenoid;
+    SolenoidRFDevice d_solenoidrf;
 
     /* Device level index lookup tables for each element type */
     int const* d_quad_indices_arr = nullptr;
     int const* d_plasmalens_indices_arr = nullptr;
+    int const* d_solenoid_indices_arr = nullptr;
+    int const* d_solenoidrf_indices_arr = nullptr;
 
     /**
      * \brief Gather the field for the particle from the lattice elements
@@ -174,10 +182,10 @@ struct LatticeElementFinderDevice
 
         amrex::ParticleReal Ex_sum = 0._prt;
         amrex::ParticleReal Ey_sum = 0._prt;
-        const amrex::ParticleReal Ez_sum = 0._prt;
+        amrex::ParticleReal Ez_sum = 0._prt;
         amrex::ParticleReal Bx_sum = 0._prt;
         amrex::ParticleReal By_sum = 0._prt;
-        const amrex::ParticleReal Bz_sum = 0._prt;
+        amrex::ParticleReal Bz_sum = 0._prt;
 
         if (d_quad.nelements > 0) {
             if (d_quad_indices_arr[iz] > -1) {
@@ -203,6 +211,34 @@ struct LatticeElementFinderDevice
             }
         }
 
+        if (d_solenoid.nelements > 0) {
+            if (d_solenoid_indices_arr[iz] > -1) {
+                const auto ielement = d_solenoid_indices_arr[iz];
+                amrex::ParticleReal Ex, Ey, Ez, Bx, By, Bz;
+                d_solenoid.get_field(ielement, x, y, z, Ex, Ey, Ez, Bx, By, Bz);
+                Ex_sum += Ex;
+                Ey_sum += Ey;
+                Ez_sum += Ez;
+                Bx_sum += Bx;
+                By_sum += By;
+                Bz_sum += Bz;
+            }
+        }
+
+        if (d_solenoidrf.nelements > 0) {
+            if (d_solenoidrf_indices_arr[iz] > -1) {
+                const auto ielement = d_solenoidrf_indices_arr[iz];
+                amrex::ParticleReal Ex, Ey, Ez, Bx, By, Bz;
+                d_solenoidrf.get_field(ielement, x, y, z, m_time, Ex, Ey, Ez, Bx, By, Bz);
+                Ex_sum += Ex;
+                Ey_sum += Ey;
+                Ez_sum += Ez;
+                Bx_sum += Bx;
+                By_sum += By;
+                Bz_sum += Bz;
+            }
+        }
+
         if (m_gamma_boost > 1._prt) {
             // The fields returned from get_field is in the lab frame
             // Transform the fields to the boosted frame

Source/AcceleratorLattice/LatticeElementFinder.cpp

@@ -8,6 +8,8 @@
 #include "LatticeElementFinder.H"
 #include "LatticeElements/HardEdgedQuadrupole.H"
 #include "LatticeElements/HardEdgedPlasmaLens.H"
+#include "LatticeElements/Solenoid.H"
+#include "LatticeElements/SolenoidRF.H"
 
 #include <AMReX_ParmParse.H>
 #include <AMReX_REAL.H>
@@ -48,6 +50,14 @@ LatticeElementFinder::AllocateIndices (AcceleratorLattice const& accelerator_lat
     if (accelerator_lattice.h_plasmalens.nelements > 0) {
         d_plasmalens_indices.resize(m_nz);
     }
+
+    if (accelerator_lattice.h_solenoid.nelements > 0) {
+        d_solenoid_indices.resize(m_nz);
+    }
+
+    if (accelerator_lattice.h_solenoidrf.nelements > 0) {
+        d_solenoidrf_indices.resize(m_nz);
+    }
 }
 
 void
@@ -74,6 +84,18 @@ LatticeElementFinder::UpdateIndices (int const lev, amrex::MFIter const& a_mfi,
                               accelerator_lattice.h_plasmalens.d_ze,
                               d_plasmalens_indices);
     }
+
+    if (accelerator_lattice.h_solenoid.nelements > 0) {
+        setup_lattice_indices(accelerator_lattice.h_solenoid.d_zs,
+                              accelerator_lattice.h_solenoid.d_ze,
+                              d_solenoid_indices);
+    }
+
+    if (accelerator_lattice.h_solenoidrf.nelements > 0) {
+        setup_lattice_indices(accelerator_lattice.h_solenoidrf.d_zs,
+                              accelerator_lattice.h_solenoidrf.d_ze,
+                              d_solenoidrf_indices);
+    }
 }
 
 LatticeElementFinderDevice
@@ -119,6 +141,16 @@ LatticeElementFinderDevice::InitLatticeElementFinderDevice (WarpXParIter const&
         d_plasmalens_indices_arr = h_finder.d_plasmalens_indices.data();
     }
 
+    if (accelerator_lattice.h_solenoid.nelements > 0) {
+        d_solenoid = accelerator_lattice.h_solenoid.GetDeviceInstance();
+        d_solenoid_indices_arr = h_finder.d_solenoid_indices.data();
+    }
+
+    if (accelerator_lattice.h_solenoidrf.nelements > 0) {
+        d_solenoidrf = accelerator_lattice.h_solenoidrf.GetDeviceInstance();
+        d_solenoidrf_indices_arr = h_finder.d_solenoidrf_indices.data();
+    }
+
 }
 
 void

Source/AcceleratorLattice/LatticeElements/CMakeLists.txt

@@ -6,5 +6,7 @@ foreach(D IN LISTS WarpX_DIMS)
         Drift.cpp
         HardEdgedQuadrupole.cpp
         HardEdgedPlasmaLens.cpp
+        Solenoid.cpp
+        SolenoidRF.cpp
     )
 endforeach()

Source/AcceleratorLattice/LatticeElements/Make.package

@@ -2,5 +2,7 @@ CEXE_sources += LatticeElementBase.cpp
 CEXE_sources += Drift.cpp
 CEXE_sources += HardEdgedQuadrupole.cpp
 CEXE_sources += HardEdgedPlasmaLens.cpp
+CEXE_sources += Solenoid.cpp
+CEXE_sources += SolenoidRF.cpp
 
 VPATH_LOCATIONS   += $(WARPX_HOME)/Source/AcceleratorLattice/LatticeElements

Source/AcceleratorLattice/LatticeElements/Solenoid.H

@@ -0,0 +1,145 @@
+/* Copyright 2023 David Grote and David Bizzozero
+ *
+ * This file is part of WarpX.
+ *
+ * License: BSD-3-Clause-LBNL
+ */
+#ifndef WARPX_ACCELERATORLATTICE_LATTICEELEMENTS_SOLENOID_H_
+#define WARPX_ACCELERATORLATTICE_LATTICEELEMENTS_SOLENOID_H_
+
+#include "LatticeElementBase.H"
+#include "HardEdged_K.H"
+#include "Utils/WarpXConst.H"
+
+#include <AMReX_REAL.H>
+#include <AMReX_ParmParse.H>
+#include <AMReX_GpuContainers.H>
+
+#include <string>
+#include <cmath>
+
+// Specifies a field that models the field generated by a solenoidal field,
+// an axisymmetric field expansion
+
+struct SolenoidDevice;
+
+struct Solenoid
+    : LatticeElementBase
+{
+
+    Solenoid ();
+
+    ~Solenoid () = default;
+
+    /**
+     * \brief Read in an element and add it to the lists
+     *
+     * @param[in] pp_element The ParmParse instance to read in the data
+     * @param[inout] z_location The current z location in the lattice
+     */
+    void
+    AddElement (amrex::ParmParse & pp_element, amrex::ParticleReal & z_location);
+
+    /**
+     * \brief Write the element information to the device
+     */
+    void
+    WriteToDevice ();
+
+    // Host variables read from input file
+    std::vector<amrex::ParticleReal> h_scale;
+    std::vector<std::vector<amrex::ParticleReal>> h_b_coef;
+
+    // Device variable read from input files
+    amrex::Gpu::DeviceVector<amrex::ParticleReal> d_scale;
+    amrex::Gpu::DeviceVector<amrex::Vector<amrex::ParticleReal> *> d_b_coef;
+
+    /**
+     * \brief Returns the device level instance with the lattice information
+     */
+    SolenoidDevice GetDeviceInstance () const;
+
+
+};
+
+// Instance that is trivially copyable to the device.
+
+struct SolenoidDevice
+{
+
+    /**
+     * \brief Initializes the data and pointer needed to reference the lattice element info
+     *
+     * @param[in] h_solenoid host level instance that this is associated with
+     */
+    void InitSolenoidDevice (Solenoid const& h_solenoid);
+
+    int nelements = 0;
+
+    const amrex::ParticleReal* AMREX_RESTRICT d_zs_arr;
+    const amrex::ParticleReal* AMREX_RESTRICT d_ze_arr;
+    const amrex::ParticleReal* AMREX_RESTRICT d_scale_arr;
+    unsigned long* AMREX_RESTRICT d_num_b_coef;
+    const amrex::ParticleReal** AMREX_RESTRICT d_b_coef_arr;
+
+    /**
+     * \brief Fetch the field of the specified element at the given location
+     *
+     * @param[in] ielement the element number
+     * @param[in] x, y, z the particle position in the lab frame
+     * @param[in] m_time the current time in the lab frame
+     * @param[out] Ex, Ey, Ez, Bx, By, Bz the fetched field in the lab frame
+     */
+    AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE
+    void get_field (const int ielement,
+                    const amrex::ParticleReal x,
+                    const amrex::ParticleReal y,
+                    const amrex::ParticleReal z,
+                    amrex::ParticleReal& Ex,
+                    amrex::ParticleReal& Ey,
+                    amrex::ParticleReal& Ez,
+                    amrex::ParticleReal& Bx,
+                    amrex::ParticleReal& By,
+                    amrex::ParticleReal& Bz) const
+    {
+
+        const amrex::ParticleReal zs = d_zs_arr[ielement];
+        const amrex::ParticleReal ze = d_ze_arr[ielement];
+        const amrex::ParticleReal scale = d_scale_arr[ielement];
+        const amrex::ParticleReal* b_coef = d_b_coef_arr[ielement];
+
+        constexpr amrex::ParticleReal pi = MathConst::pi;   // Everyone's favorite circle ratio
+
+        const int n_coef = (d_num_b_coef[ielement]-1) / 2;  // Number of coef. is half the length of the list: zero and odd inds are cosine, even inds are sine
+        const amrex::ParticleReal zlen = ze - zs;           // Length of element
+
+        amrex::ParticleReal zz = z - ze + zlen/2.0;         // Shifted coordinate to midpoint of element (zz = 0 is midpoint of element)
+        amrex::ParticleReal bz = b_coef[0]/2.0;             // A_0 cosine fourier coef. is halved (sine coef. is trivially zero)
+        amrex::ParticleReal bzp = 0.0;                      // Variable to store Bz'(zz) (z-derivative of on-axis longitudinal field at z = zz)
+        amrex::ParticleReal bzpp = 0.0;                     // Variable to store Bz''(zz)
+        amrex::ParticleReal bzppp = 0.0;                    // Variable to store Bz'''(zz)
+        for (int k = 1; k < n_coef; ++k){
+            bz = bz + b_coef[2*k-1]*cos(2*pi*k*zz/zlen) + b_coef[2*k]*sin(2.0*pi*k*zz/zlen);
+            bzp = bzp - (2.0*pi*k*zz/zlen)*b_coef[2*k-1]*sin(2.0*pi*k*zz/zlen) + (2.0*pi*k*zz/zlen)*b_coef[2*k]*cos(2.0*pi*k*zz/zlen);
+            bzpp = bzpp - pow(2.0*pi*k*zz/zlen,2)*b_coef[2*k-1]*cos(2.0*pi*k*zz/zlen) - pow(2.0*pi*k*zz/zlen,2)*b_coef[2*k]*sin(2.0*pi*k*zz/zlen);
+            bzppp = bzppp + pow(2.0*pi*k*zz/zlen,3)*b_coef[2*k-1]*sin(2.0*pi*k*zz/zlen) - pow(2.0*pi*k*zz/zlen,3)*b_coef[2*k]*cos(2.0*pi*k*zz/zlen);
+        }
+        bz = bz*scale;          // Scale fields and derivatives by user-provided scaling factor (from input file)
+        bzp = bzp*scale;
+        bzpp = bzpp*scale;
+        bzppp = bzppp*scale;
+
+        amrex::ParticleReal r2 = x*x + y*y;
+
+        Ex = 0.0;      // Update E and B fields from Fourier coefficients and scaling factors
+        Ey = 0.0;
+        Ez = 0.0;
+        Bx = -(bzp*x/2) + (bzppp*x*r2/16);
+        By = -(bzp*y/2) + (bzppp*y*r2/16);
+        Bz = bz - (bzpp*r2/4);
+
+    }
+
+};
+
+#endif // WARPX_ACCELERATORLATTICE_LATTICEELEMENTS_SOLENOID_H_