Add Beam Spin tracking

docs/source/run/parameters.rst

@@ -1039,3 +1039,22 @@ Whether the energy loss due to classical radiation reaction of beam particles is
     Whether the beam particles undergo energy loss due to classical radiation reaction.
     The implemented radiation reaction model is based on this publication: `M. Tamburini et al., NJP 12, 123005 <https://doi.org/10.1088/1367-2630/12/12/123005>`__
     In normalized units, `hipace.background_density_SI` must be specified.
+
+Spin tracking
+-------------
+
+Track the spin of each beam particle as it is rotated by the electromagnetic fields using the
+Thomas-Bargmann-Michel-Telegdi (TBMT) model, see
+[Z. Gong et al., Matter and Radiation at Extremes 8.6 (2023), https://doi.org/10.1063/5.0152382]
+for the details of the implementation.
+This will add three extra components to each beam particle to store the spin and output
+those as part of the beam diagnostic.
+
+* ``<beam name> or beams.do_spin_tracking`` (`bool`) optional (default `0`)
+    Enable spin tracking
+
+* ``<beam name> or beams.initial_spin`` (3 `float`)
+    Initial spin ``sx sy sz`` of all particles. The length of the three components is normalized to one.
+
+* ``<beam name> or beams.spin_anom`` (`bool`) optional (default `0.00115965218128`)
+    The anomalous magnetic moment. The default value is the moment for electrons.

src/diagnostics/OpenPMDWriter.cpp

@@ -298,7 +298,8 @@ OpenPMDWriter::CopyBeams (MultiBeam& beams, const amrex::Vector< std::string > b
                     m_uint64_beam_data[ibeam][idx].get() + m_offset[ibeam]);
             }
 
-            AMREX_ALWAYS_ASSERT_WITH_MESSAGE(m_real_beam_data[ibeam].size() == soa.NumRealComps(),
+            AMREX_ALWAYS_ASSERT_WITH_MESSAGE(
+                int(m_real_beam_data[ibeam].size()) == soa.NumRealComps(),
                 "List of real names in openPMD Writer class does not match the beam");
 
             for (std::size_t idx=0; idx<m_real_beam_data[ibeam].size(); idx++) {

src/particles/beam/BeamParticleContainer.H

@@ -191,12 +191,20 @@ public:
         return m_total_num_particles;
     }
 
-    int numRealComponents () const { return BeamIdx::real_nattribs; }
+    int numRealComponents () const {
+        return BeamIdx::real_nattribs + (m_do_spin_tracking ? 3 : 0);
+    }
     int numIntComponents () const { return BeamIdx::int_nattribs; }
 
     bool communicateIdCpuComponent () const { return true; }
-    bool communicateRealComponent (int rcomp) const { return rcomp < BeamIdx::real_nattribs_in_buffer; }
-    bool communicateIntComponent (int icomp) const { return icomp < BeamIdx::int_nattribs_in_buffer; }
+    bool communicateRealComponent (int rcomp) const {
+        // communicate all compile-time and runtime real components
+        return rcomp < numRealComponents();
+    }
+    bool communicateIntComponent (int icomp) const {
+        // don't communicate nsubcycles
+        return icomp < BeamIdx::int_nattribs_in_buffer;
+    }
 
 private:
     int m_slice_permutation = 0;
@@ -228,6 +236,10 @@ public:
     amrex::Array<amrex::Parser, 6> m_external_fields_parser;
     /** If spin tracking is enabled for this beam */
     bool m_do_spin_tracking = false;
+    /** Initial spin of all particles */
+    amrex::RealVect m_initial_spin = {1, 0, 0,};
+    /** The anomalous magnetic moment */
+    amrex::Real m_spin_anom = 0.00115965218128;
 private:
     std::string m_name; /**< name of the species */
     /** injection type, fixed_width or fixed_ppc */

src/particles/beam/BeamParticleContainer.cpp

@@ -103,6 +103,15 @@ BeamParticleContainer::ReadParameters ()
         soa.GetIntData()[icomp].setArena(
             m_initialize_on_cpu ? amrex::The_Pinned_Arena() : amrex::The_Arena());
     }
+    queryWithParserAlt(pp, "do_spin_tracking", m_do_spin_tracking, pp_alt);
+    if (m_do_spin_tracking) {
+        getWithParserAlt(pp, "initial_spin", m_initial_spin, pp_alt);
+        queryWithParserAlt(pp, "spin_anom", m_spin_anom, pp_alt);
+        for (auto& beam_tile : m_slices) {
+            // Use 3 real and 0 int runtime components
+            beam_tile.define(3, 0);
+        }
+    }
 }
 
 amrex::Real
@@ -383,6 +392,20 @@ BeamParticleContainer::intializeSlice (int slice, int which_slice) {
             }
         );
     }
+
+    if (m_do_spin_tracking) {
+        auto ptd = getBeamSlice(which_slice).getParticleTileData();
+
+        const amrex::RealVect initial_spin_norm = m_initial_spin / m_initial_spin.vectorLength();
+
+        amrex::ParallelFor(getNumParticles(which_slice),
+            [=] AMREX_GPU_DEVICE (const int ip) {
+                ptd.m_runtime_rdata[0][ip] = initial_spin_norm[0];
+                ptd.m_runtime_rdata[1][ip] = initial_spin_norm[1];
+                ptd.m_runtime_rdata[2][ip] = initial_spin_norm[2];
+            }
+        );
+    }
 }
 
 void

src/particles/pusher/BeamParticleAdvance.cpp

@@ -32,6 +32,8 @@ AdvanceBeamParticlesSlice (
     const amrex::Real dt = Hipace::GetInstance().m_dt / n_subcycles;
     const amrex::Real background_density_SI = Hipace::m_background_density_SI;
     const bool normalized_units = Hipace::m_normalized_units;
+    const bool spin_tracking = beam.m_do_spin_tracking;
+    const amrex::Real spin_anom = beam.m_spin_anom;
 
     if (normalized_units && radiation_reaction) {
         AMREX_ALWAYS_ASSERT_WITH_MESSAGE(background_density_SI!=0,
@@ -140,6 +142,13 @@ AdvanceBeamParticlesSlice (
 
             int i = ptd.idata(BeamIdx::nsubcycles)[ip];
 
+            amrex::RealVect spin {0._rt, 0._rt, 0._rt};
+            if (spin_tracking) {
+                spin[0] = ptd.m_runtime_rdata[0][ip];
+                spin[1] = ptd.m_runtime_rdata[1][ip];
+                spin[2] = ptd.m_runtime_rdata[2][ip];
+            }
+
             for (; i < n_subcycles; i++) {
 
                 if (zp < min_z) {
@@ -214,6 +223,28 @@ AdvanceBeamParticlesSlice (
                        + uy_intermediate*uy_intermediate
                        + uz_intermediate*uz_intermediate )*inv_c2 );
 
+                if (spin_tracking) {
+                    const amrex::RealVect E {ExmByp + clight*Byp, EypBxp - clight*Bxp, Ezp};
+                    const amrex::RealVect B {Bxp, Byp, Bzp};
+                    const amrex::RealVect u {ux_intermediate*inv_clight, uy_intermediate*inv_clight,
+                                             uz_intermediate*inv_clight};
+                    const amrex::RealVect beta = u*gamma_intermediate_inv;
+                    const amrex::Real gamma_inv_p1 =
+                        gamma_intermediate_inv / (1._rt + gamma_intermediate_inv);
+
+                    const amrex::RealVect omega = std::abs(charge_mass_ratio) * (
+                        B * gamma_intermediate_inv - beta.crossProduct(E) * inv_clight * gamma_inv_p1
+                        + spin_anom * (
+                            B - gamma_inv_p1 * u * beta.dotProduct(B) - beta.crossProduct(E) * inv_clight
+                        )
+                    );
+
+                    const amrex::RealVect h = omega * dt * 0.5_rt;
+                    const amrex::RealVect s_prime = spin + h.crossProduct(spin);
+                    const amrex::Real o = 1._rt / (1._rt + h.dotProduct(h));
+                    spin = o * (s_prime + (h.dotProduct(s_prime) * h + h.crossProduct(s_prime)));
+                }
+
                 amrex::ParticleReal uz_next = uz + dt * charge_mass_ratio
                     * ( Ezp + ( ux_intermediate * Byp - uy_intermediate * Bxp )
                     * gamma_intermediate_inv );
@@ -300,5 +331,11 @@ AdvanceBeamParticlesSlice (
             ptd.rdata(BeamIdx::ux)[ip] = ux;
             ptd.rdata(BeamIdx::uy)[ip] = uy;
             ptd.rdata(BeamIdx::uz)[ip] = uz;
+
+            if (spin_tracking) {
+                ptd.m_runtime_rdata[0][ip] = spin[0];
+                ptd.m_runtime_rdata[1][ip] = spin[1];
+                ptd.m_runtime_rdata[2][ip] = spin[2];
+            }
         });
 }