Beam read_from_file also reads spin (#1172)

Co-authored-by: Alexander Sinn <[email protected]>

Author: MaxThevenet

src/particles/beam/BeamParticleContainer.cpp

@@ -92,26 +92,29 @@ BeamParticleContainer::ReadParameters ()
         "Tilted beams and correlated energy spreads are only implemented for fixed weight beams");
     }
     queryWithParserAlt(pp, "initialize_on_cpu", m_initialize_on_cpu, pp_alt);
+    queryWithParserAlt(pp, "do_spin_tracking", m_do_spin_tracking, pp_alt);
+    if (m_do_spin_tracking) {
+        if (m_injection_type != "from_file") {
+            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);
+        }
+        getBeamInitSlice().define(3, 0);
+    }
     auto& soa = getBeamInitSlice().GetStructOfArrays();
     soa.GetIdCPUData().setArena(
         m_initialize_on_cpu ? amrex::The_Pinned_Arena() : amrex::The_Arena());
     for (int rcomp = 0; rcomp < soa.NumRealComps(); ++rcomp) {
-        soa.GetRealData()[rcomp].setArena(
+        soa.GetRealData(rcomp).setArena(
             m_initialize_on_cpu ? amrex::The_Pinned_Arena() : amrex::The_Arena());
     }
     for (int icomp = 0; icomp < soa.NumIntComps(); ++icomp) {
-        soa.GetIntData()[icomp].setArena(
+        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
@@ -369,6 +372,7 @@ BeamParticleContainer::initializeSlice (int slice, int which_slice) {
 
         const int slice_offset = m_init_sorter.m_box_offsets_cpu[slice];
         const auto permutations = m_init_sorter.m_box_permutations.dataPtr();
+        const bool do_spin_tracking = m_do_spin_tracking;
 
         amrex::ParallelFor(num_particles,
             [=] AMREX_GPU_DEVICE (const int ip) {
@@ -380,20 +384,23 @@ BeamParticleContainer::initializeSlice (int slice, int which_slice) {
                 ptd.rdata(BeamIdx::ux)[ip] = ptd_init.rdata(BeamIdx::ux)[idx_src];
                 ptd.rdata(BeamIdx::uy)[ip] = ptd_init.rdata(BeamIdx::uy)[idx_src];
                 ptd.rdata(BeamIdx::uz)[ip] = ptd_init.rdata(BeamIdx::uz)[idx_src];
-
+                if (do_spin_tracking) {
+                    ptd.m_runtime_rdata[0][ip] = ptd_init.m_runtime_rdata[0][idx_src];
+                    ptd.m_runtime_rdata[1][ip] = ptd_init.m_runtime_rdata[1][idx_src];
+                    ptd.m_runtime_rdata[2][ip] = ptd_init.m_runtime_rdata[2][idx_src];
+                }
                 ptd.idcpu(ip) = ptd_init.idcpu(idx_src);
                 ptd.idata(BeamIdx::nsubcycles)[ip] = 0;
                 ptd.idata(BeamIdx::mr_level)[ip] = 0;
             }
         );
     }
 
-    if (m_do_spin_tracking) {
+    if (m_do_spin_tracking && m_injection_type != "from_file") {
         HIPACE_PROFILE("BeamParticleContainer::initializeSpin()");
         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];

src/particles/beam/BeamParticleContainerInit.cpp

@@ -42,10 +42,12 @@ namespace
      */
     AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE
     void AddOneBeamParticle (
-        const BeamTileInit::ParticleTileDataType& ptd, const amrex::Real& x,
-        const amrex::Real& y, const amrex::Real& z, const amrex::Real& ux, const amrex::Real& uy,
-        const amrex::Real& uz, const amrex::Real& weight, const amrex::Long pid,
-        const amrex::Long ip, const amrex::Real& speed_of_light, const EnforceBC& enforceBC) noexcept
+        const BeamTileInit::ParticleTileDataType& ptd,
+        const amrex::Real& x, const amrex::Real& y, const amrex::Real& z,
+        const amrex::Real& ux, const amrex::Real& uy, const amrex::Real& uz,
+        const amrex::Real& sx, const amrex::Real& sy, const amrex::Real& sz,
+        const amrex::Real& weight, const amrex::Long pid, const amrex::Long ip,
+        const amrex::Real& speed_of_light, const EnforceBC& enforceBC, bool do_spin) noexcept
     {
         amrex::Real xp = x;
         amrex::Real yp = y;
@@ -59,6 +61,11 @@ namespace
         ptd.rdata(BeamIdx::ux )[ip] = uxp;
         ptd.rdata(BeamIdx::uy )[ip] = uyp;
         ptd.rdata(BeamIdx::uz )[ip] = uz * speed_of_light;
+        if (do_spin) {
+            ptd.m_runtime_rdata[0][ip] = sx;
+            ptd.m_runtime_rdata[1][ip] = sy;
+            ptd.m_runtime_rdata[2][ip] = sz;
+        }
         ptd.rdata(BeamIdx::w  )[ip] = std::abs(weight);
 
         ptd.idcpu(ip) = pid + ip;
@@ -790,6 +797,7 @@ InitBeamFromFile (const std::string input_file,
     std::string name_particle ="";
     std::string name_r ="", name_rx ="", name_ry ="", name_rz ="";
     std::string name_u ="", name_ux ="", name_uy ="", name_uz ="";
+    std::string name_s ="", name_sx ="", name_sy ="", name_sz ="";
     std::string name_m ="", name_mm ="";
     std::string name_q ="", name_qq ="";
     std::string name_g ="", name_gg ="";
@@ -873,6 +881,24 @@ InitBeamFromFile (const std::string input_file,
                         }
                     }
                 }
+                else if(units == std::array<double,7> {2., 1., -1., 0., 0., 0., 0.} &&
+                        m_do_spin_tracking) {
+                    // spin
+                    if(physical_quantity.first == "spin") {
+                        name_s = physical_quantity.first;
+                        for( auto const& axes_direction : physical_quantity.second ) {
+                            if(axes_direction.first == "x" || axes_direction.first == "X") {
+                                name_sx = axes_direction.first;
+                            }
+                            if(axes_direction.first == "y" || axes_direction.first == "Y") {
+                                name_sy = axes_direction.first;
+                            }
+                            if(axes_direction.first == "z" || axes_direction.first == "Z") {
+                                name_sz = axes_direction.first;
+                            }
+                        }
+                    }
+                }
             }
         }
     }
@@ -953,15 +979,28 @@ InitBeamFromFile (const std::string input_file,
         }
     }
 
+    if (m_do_spin_tracking) {
+        for(std::string name_s_c : {name_sx, name_sy, name_sz}) {
+            if(!series.iterations[num_iteration].particles[name_particle][name_s].contains(name_s_c)) {
+                amrex::Abort("Beam input file does not contain " + name_s_c + " coordinate in " +
+                             name_s + " (spin). An attempt to read these was done because " +
+                             "do_spin_tracking is on for at least one beam.\n");
+            }
+        }
+    }
     // print the names of the arrays in the file that are used for the simulation
     if(Hipace::m_verbose >= 3){
        amrex::Print() << "Beam Input File '" << input_file << "' in Iteration '" << num_iteration
           << "' and Paricle '" << name_particle
-          << "' imported with:\nPositon '" << name_r << "' (coordinates '" << name_rx << "', '"
+          << "' imported with:\nPosition '" << name_r << "' (coordinates '" << name_rx << "', '"
           << name_ry << "', '" << name_rz << "')\n"
           << momentum_type << " '" << name_u << "' (coordinates '" << name_ux
           << "', '" << name_uy << "', '" << name_uz << "')\n"
           << weighting_type << " '" << name_w << "' (in '" << name_ww << "')\n";
+       if (m_do_spin_tracking) {
+           amrex::Print() << name_u << "' (coordinates '" << name_ux
+                          << "', '" << name_uy << "', '" << name_uz << "')\n";
+       }
     }
 
     auto electrons = series.iterations[num_iteration].particles[name_particle];
@@ -987,6 +1026,15 @@ InitBeamFromFile (const std::string input_file,
         amrex::The_Pinned_Arena()->alloc(sizeof(input_type)*num_to_add) ), del};
     std::shared_ptr<input_type> u_z_data{ reinterpret_cast<input_type*>(
         amrex::The_Pinned_Arena()->alloc(sizeof(input_type)*num_to_add) ), del};
+    std::shared_ptr<input_type> s_x_data{ reinterpret_cast<input_type*>(
+        amrex::The_Pinned_Arena()->alloc(m_do_spin_tracking ?
+                                         sizeof(input_type)*num_to_add : 0) ), del};
+    std::shared_ptr<input_type> s_y_data{ reinterpret_cast<input_type*>(
+        amrex::The_Pinned_Arena()->alloc(m_do_spin_tracking ?
+                                         sizeof(input_type)*num_to_add : 0) ), del};
+    std::shared_ptr<input_type> s_z_data{ reinterpret_cast<input_type*>(
+        amrex::The_Pinned_Arena()->alloc(m_do_spin_tracking ?
+                                         sizeof(input_type)*num_to_add: 0 ) ), del};
     std::shared_ptr<input_type> w_w_data{ reinterpret_cast<input_type*>(
         amrex::The_Pinned_Arena()->alloc(sizeof(input_type)*num_to_add) ), del};
 
@@ -997,6 +1045,11 @@ InitBeamFromFile (const std::string input_file,
     electrons[name_u][name_uy].loadChunk<input_type>(u_y_data, {0u}, {num_to_add});
     electrons[name_u][name_uz].loadChunk<input_type>(u_z_data, {0u}, {num_to_add});
     electrons[name_w][name_ww].loadChunk<input_type>(w_w_data, {0u}, {num_to_add});
+    if (m_do_spin_tracking) {
+        electrons[name_s][name_sx].loadChunk<input_type>(s_x_data, {0u}, {num_to_add});
+        electrons[name_s][name_sy].loadChunk<input_type>(s_y_data, {0u}, {num_to_add});
+        electrons[name_s][name_sz].loadChunk<input_type>(s_z_data, {0u}, {num_to_add});
+    }
 
     series.flush();
 
@@ -1052,6 +1105,7 @@ InitBeamFromFile (const std::string input_file,
     auto old_size = particle_tile.size();
     auto new_size = old_size + num_to_add;
     particle_tile.resize(new_size);
+
     const auto ptd = particle_tile.getParticleTileData();
     const auto enforceBC = EnforceBC();
 
@@ -1064,7 +1118,11 @@ InitBeamFromFile (const std::string input_file,
     const input_type * const u_x_ptr = u_x_data.get();
     const input_type * const u_y_ptr = u_y_data.get();
     const input_type * const u_z_ptr = u_z_data.get();
+    const input_type * const s_x_ptr = m_do_spin_tracking ? s_x_data.get() : nullptr;
+    const input_type * const s_y_ptr = m_do_spin_tracking ? s_y_data.get() : nullptr;
+    const input_type * const s_z_ptr = m_do_spin_tracking ? s_z_data.get() : nullptr;
     const input_type * const w_w_ptr = w_w_data.get();
+    const bool do_spin_tracking = m_do_spin_tracking;
 
     amrex::ParallelFor(amrex::Long(num_to_add),
         [=] AMREX_GPU_DEVICE (const amrex::Long i) {
@@ -1075,8 +1133,11 @@ InitBeamFromFile (const std::string input_file,
                 static_cast<amrex::Real>(u_x_ptr[i] * unit_ux), // = gamma * beta
                 static_cast<amrex::Real>(u_y_ptr[i] * unit_uy),
                 static_cast<amrex::Real>(u_z_ptr[i] * unit_uz),
+                do_spin_tracking ? static_cast<amrex::Real>(s_x_ptr[i]) : 0.,
+                do_spin_tracking ? static_cast<amrex::Real>(s_y_ptr[i]) : 0.,
+                do_spin_tracking ? static_cast<amrex::Real>(s_z_ptr[i]) : 0.,
                 static_cast<amrex::Real>(w_w_ptr[i] * unit_ww),
-                pid, i, phys_const.c, enforceBC);
+                pid, i, phys_const.c, enforceBC, do_spin_tracking);
         });
 
     amrex::Gpu::streamSynchronize();