[WIP] Add collisional diagnostic buffers

Python/pywarpx/fields.py

@@ -1030,3 +1030,15 @@ def GCPPMLWrapper(level=0, include_ghosts=False):
     return _MultiFABWrapper(
         mf_name="pml_G_cp", level=level, include_ghosts=include_ghosts
     )
+
+
+def MCCEnergyWrapper(level=0, include_ghosts=False):
+    return _MultiFABWrapper(
+        mf_name="collision_energy_change", level=level, include_ghosts=include_ghosts
+    )
+
+
+def MCCIonizWrapper(level=0, include_ghosts=False):
+    return _MultiFABWrapper(
+        mf_name="collision_ionization", level=level, include_ghosts=include_ghosts
+    )

Source/Fields.H

@@ -81,7 +81,9 @@ namespace warpx::fields
         Bfield_avg_cp,
         B_old, /**< Stores the value of B at the beginning of the timestep, for the implicit solver */
         ECTRhofield,
-        Venl
+        Venl,
+        collision_energy_change, /**< Stores collisional energy transfer data for Monte Carlo Collisions */
+        collision_ionization /**< Stores ionization data for Monte Carlo Collisions */
     );
 
     /** these are vector fields */

Source/Particles/Collision/BackgroundMCC/BackgroundMCCCollision.H

@@ -45,11 +45,13 @@ public:
 
     /** Perform particle conserving MCC collisions within a tile
      *
+     * @param lev the mesh-refinement level
+     * @param species1 reference to colliding species container
      * @param pti particle iterator
      * @param t current time
      *
      */
-    void doBackgroundCollisionsWithinTile ( WarpXParIter& pti, amrex::Real t);
+    void doBackgroundCollisionsWithinTile ( int lev, WarpXParticleContainer& species1, WarpXParIter& pti, amrex::Real t);
 
     /** Perform MCC ionization interactions
      *

Source/Particles/Collision/BackgroundMCC/BackgroundMCCCollision.cpp

@@ -296,7 +296,7 @@ BackgroundMCCCollision::doCollisions (amrex::Real cur_time, amrex::Real dt, Mult
             }
             auto wt = static_cast<amrex::Real>(amrex::second());
 
-            doBackgroundCollisionsWithinTile(pti, cur_time);
+            doBackgroundCollisionsWithinTile(lev, species1, pti, cur_time);
 
             if (cost && WarpX::load_balance_costs_update_algo == LoadBalanceCostsUpdateAlgo::Timers)
             {
@@ -315,13 +315,17 @@ BackgroundMCCCollision::doCollisions (amrex::Real cur_time, amrex::Real dt, Mult
 
 
 void BackgroundMCCCollision::doBackgroundCollisionsWithinTile
-( WarpXParIter& pti, amrex::Real t )
+( int lev, WarpXParticleContainer& species1, WarpXParIter& pti, amrex::Real t )
 {
     using namespace amrex::literals;
+    using warpx::fields::FieldType;
 
     // So that CUDA code gets its intrinsic, not the host-only C++ library version
     using std::sqrt;
 
+    // get pointer to warpx instance, for multifab retrieval
+    auto & warpx = WarpX::GetInstance();
+
     // get particle count
     const long np = pti.numParticles();
 
@@ -353,6 +357,17 @@ void BackgroundMCCCollision::doBackgroundCollisionsWithinTile
     amrex::ParticleReal* const AMREX_RESTRICT ux = attribs[PIdx::ux].dataPtr();
     amrex::ParticleReal* const AMREX_RESTRICT uy = attribs[PIdx::uy].dataPtr();
     amrex::ParticleReal* const AMREX_RESTRICT uz = attribs[PIdx::uz].dataPtr();
+    amrex::ParticleReal* const AMREX_RESTRICT w  = attribs[PIdx::w].dataPtr();
+
+    // get multi-fab for collisional energy transfer (does this need to be moved out of the pti loop?)
+    const amrex::MultiFab &coll_E_change = *warpx.m_fields.get(FieldType::collision_energy_change, lev);
+    const auto& coll_E_change_arr = coll_E_change[pti].array();
+
+    // get parameters for getParticleCell (method used here is similar to TemperatureFunctor.cpp)
+    auto& tile = pti.GetParticleTile();
+    auto ptd = tile.getParticleTileData();
+    const auto plo = species1.Geom(lev).ProbLoArray();
+    const auto dxi = species1.Geom(lev).InvCellSizeArray();
 
     amrex::ParallelForRNG(np,
                           [=] AMREX_GPU_HOST_DEVICE (long ip, amrex::RandomEngine const& engine)
@@ -366,8 +381,8 @@ void BackgroundMCCCollision::doBackgroundCollisionsWithinTile
                               const amrex::ParticleReal n_a = n_a_func(x, y, z, t);
                               const amrex::ParticleReal T_a = T_a_func(x, y, z, t);
 
-                              amrex::ParticleReal v_coll, v_coll2, sigma_E, nu_i = 0;
-                              double gamma, E_coll;
+                              amrex::ParticleReal v_coll, v_coll2, sigma_E, nu_i = 0, v_lost2;
+                              double gamma, E_coll, E_lost;
                               amrex::ParticleReal ua_x, ua_y, ua_z, vx, vy, vz;
                               amrex::ParticleReal uCOM_x, uCOM_y, uCOM_z;
                               const amrex::ParticleReal col_select = amrex::Random(engine);
@@ -406,6 +421,11 @@ void BackgroundMCCCollision::doBackgroundCollisionsWithinTile
                                   // check if this collision should be performed
                                   if (col_select > nu_i) { continue; }
 
+                                  // get the initial particle energy for tracking
+                                  if (WarpX::do_MCC_energy_tracking) {
+                                      v_lost2 = ux[ip]*ux[ip] + uy[ip]*uy[ip] + uz[ip]*uz[ip];
+                                  }
+
                                   // charge exchange is implemented as a simple swap of the projectile
                                   // and target velocities which doesn't require any of the Lorentz
                                   // transformations below; note that if the projectile and target
@@ -414,6 +434,20 @@ void BackgroundMCCCollision::doBackgroundCollisionsWithinTile
                                       ux[ip] = ua_x;
                                       uy[ip] = ua_y;
                                       uz[ip] = ua_z;
+
+                                      if (WarpX::do_MCC_energy_tracking) {
+                                          // Calculate the energy lost for tracking
+                                          v_lost2 -= ux[ip]*ux[ip] + uy[ip]*uy[ip] + uz[ip]*uz[ip];
+                                          ParticleUtils::getEnergy(v_lost2, m, E_lost);
+
+                                          // store the energy (add E_lost * w[ip] to the buffer of energy sent to the background)
+                                          // get the particle's cell index (method used here is similar to TemperatureFunctor.cpp)
+                                          const auto p = WarpXParticleContainer::ParticleType(ptd, ip);
+                                          const auto [ii, jj, kk] = amrex::getParticleCell(p , plo, dxi).dim3();
+
+                                          // store energy lost in the buffer
+                                          coll_E_change_arr(ii, jj, kk, 0) += E_lost * w[ip];
+                                      }
                                       break;
                                   }
 
@@ -459,6 +493,20 @@ void BackgroundMCCCollision::doBackgroundCollisionsWithinTile
                                   ux[ip] = vx + ua_x;
                                   uy[ip] = vy + ua_y;
                                   uz[ip] = vz + ua_z;
+
+                                  if (WarpX::do_MCC_energy_tracking) {
+                                      // Calculate difference in energy for tracking
+                                      v_lost2 -= ux[ip]*ux[ip] + uy[ip]*uy[ip] + uz[ip]*uz[ip];
+                                      ParticleUtils::getEnergy(v_lost2, m, E_lost);
+
+                                      // store the energy (add E_lost * w[ip] to the buffer of energy sent to the background)
+                                      // get the particle's cell index (method used here is similar to TemperatureFunctor.cpp)
+                                      const auto p = WarpXParticleContainer::ParticleType(ptd, ip);
+                                      const auto [ii, jj, kk] = amrex::getParticleCell(p , plo, dxi).dim3();
+
+                                      // store energy lost in the buffer
+                                      coll_E_change_arr(ii, jj, kk, 0) += E_lost * w[ip];
+                                  }
                                   break;
                               }
                           }
@@ -470,8 +518,13 @@ void BackgroundMCCCollision::doBackgroundIonization
 ( int lev, amrex::LayoutData<amrex::Real>* cost,
   WarpXParticleContainer& species1, WarpXParticleContainer& species2, amrex::Real t)
 {
+    using warpx::fields::FieldType;
+
     WARPX_PROFILE("BackgroundMCCCollision::doBackgroundIonization()");
 
+    // get pointer to warpx instance, for multifab retrieval
+    auto & warpx = WarpX::GetInstance();
+
     const SmartCopyFactory copy_factory_elec(species1, species1);
     const SmartCopyFactory copy_factory_ion(species1, species2);
     const auto CopyElec = copy_factory_elec.getSmartCopy();
@@ -485,6 +538,13 @@ void BackgroundMCCCollision::doBackgroundIonization
 
     const amrex::ParticleReal sqrt_kb_m = std::sqrt(PhysConst::kb / m_background_mass);
 
+    // get multi-fab for ionization creation tracking
+    const amrex::MultiFab &coll_ionization = *warpx.m_fields.get(FieldType::collision_ionization, lev);
+
+    // get parameters for getParticleCell
+    const auto plo = species1.Geom(lev).ProbLoArray();
+    const auto dxi = species1.Geom(lev).InvCellSizeArray();
+
 #ifdef AMREX_USE_OMP
 #pragma omp parallel if (amrex::Gpu::notInLaunchRegion())
 #endif
@@ -515,6 +575,36 @@ void BackgroundMCCCollision::doBackgroundIonization
         setNewParticleIDs(elec_tile, np_elec, num_added);
         setNewParticleIDs(ion_tile, np_ion, num_added);
 
+        if (WarpX::do_MCC_ionization_tracking)
+        {
+            // get parameters for getParticleCell (method used here is similar to TemperatureFunctor.cpp)
+            // should these lines be removed and we just do ptd = elec_tile.getParticleTileData()?
+            auto& tile = pti.GetParticleTile();
+            auto ptd = tile.getParticleTileData();
+
+            // get multi-fab array for ionization creation tracking
+            const auto& coll_ionization_arr = coll_ionization[pti].array();
+
+            // get Struct-Of-Array particle data
+            // as long as w is set after the copy, will it include newly created particles (be of length np_elec + num_added)?
+            auto& attribs = pti.GetAttribs();
+            amrex::ParticleReal* const AMREX_RESTRICT w  = attribs[PIdx::w].dataPtr();
+
+            // add each created particle to tracking buffer
+            amrex::ParallelFor(num_added,
+                               [=] AMREX_GPU_HOST_DEVICE (long ip)
+                               {
+                               // store the energy (add E_lost * w[ip] to the buffer of energy sent to the background)
+                               // get the particle's cell index (method used here is similar to TemperatureFunctor.cpp)
+                               const auto p = WarpXParticleContainer::ParticleType(ptd, ip + np_elec);
+                               const auto [ii, jj, kk] = amrex::getParticleCell(p , plo, dxi).dim3();
+
+                               // store energy lost in the buffer
+                               coll_ionization_arr(ii, jj, kk, 0) += w[ip + np_elec];
+                               }
+                               );
+        }
+
         if (cost && WarpX::load_balance_costs_update_algo == LoadBalanceCostsUpdateAlgo::Timers)
         {
             amrex::Gpu::synchronize();

Source/WarpX.H

@@ -868,6 +868,10 @@ public:
     static amrex::Real moving_window_v;
     static bool fft_do_time_averaging;
 
+    // Control whether Monte Carlo Collision statistics are collected
+    static bool do_MCC_energy_tracking;
+    static bool do_MCC_ionization_tracking;
+
     // these should be private, but can't due to Cuda limitations
     static void ComputeDivB (amrex::MultiFab& divB, int dcomp,
                              ablastr::fields::VectorField const & B,

Source/WarpX.cpp

@@ -104,6 +104,9 @@ Real WarpX::moving_window_v = std::numeric_limits<amrex::Real>::max();
 
 bool WarpX::fft_do_time_averaging = false;
 
+bool WarpX::do_MCC_energy_tracking = false;
+bool WarpX::do_MCC_ionization_tracking = false;
+
 amrex::IntVect WarpX::m_fill_guards_fields  = amrex::IntVect(0);
 amrex::IntVect WarpX::m_fill_guards_current = amrex::IntVect(0);
 
@@ -1744,6 +1747,53 @@ WarpX::ReadParameters ()
             }
         }
     }
+
+    // set flags for MCC collision tracking
+    // TODO: set an option for only tracking certain collisions
+    {
+        const ParmParse pp_collisions("collisions");
+
+        // Check if MCC tracking is turned off
+        int MCC_tracking = 1;
+        pp_collisions.query("mcc_tracking", MCC_tracking);
+
+        if (MCC_tracking) {
+            std::vector<std::string> collision_names;
+            pp_collisions.queryarr("collision_names", collision_names);
+
+            // Look through collision types for MCC
+            auto const ncollisions = collision_names.size();
+            for (int i = 0; i < static_cast<int>(ncollisions); ++i) {
+                const ParmParse pp_collision_name(collision_names[i]);
+
+                // For legacy, pairwisecoulomb is the default
+                std::string type = "pairwisecoulomb";
+
+                pp_collision_name.query("type", type);
+
+                // Check any MCC scattering processes
+                if (type == "background_mcc") {
+                    std::vector<std::string> scattering_processes;
+                    pp_collision_name.queryarr("scattering_processes", scattering_processes);
+
+                    // Look through scattering processes for ionization
+                    auto const nprocesses = scattering_processes.size();
+                    for (int j = 0; j < static_cast<int>(nprocesses); ++j) {
+                        // Set ionization tracking to true
+                        if (scattering_processes[j] == "ionization") {
+                            do_MCC_ionization_tracking = true;
+                        }
+                    }
+
+                    // If there are other processes set energy tracking to true
+                    if ((do_MCC_ionization_tracking && nprocesses > 1) ||
+                        (!do_MCC_ionization_tracking && nprocesses > 0)) {
+                        do_MCC_energy_tracking = true;
+                    }
+                }
+            }
+        }
+    }
 }
 
 void
@@ -2225,6 +2275,25 @@ WarpX::AllocLevelMFs (int lev, const BoxArray& ba, const DistributionMapping& dm
     m_fields.alloc_init(FieldType::current_fp, Direction{1}, lev, amrex::convert(ba, jy_nodal_flag), dm, ncomps, ngJ, 0.0_rt);
     m_fields.alloc_init(FieldType::current_fp, Direction{2}, lev, amrex::convert(ba, jz_nodal_flag), dm, ncomps, ngJ, 0.0_rt);
 
+    // Allocate extra multifabs needed for MCC collision tracking
+    if (do_MCC_ionization_tracking || do_MCC_energy_tracking) {
+        // 1 component per cell (either energy change or ionization)
+        int MCC_ncomps = 1;
+
+        // Zero guard cells
+        IntVect ngMCC = IntVect::TheZeroVector();
+
+        // Prepare to initialize a cell-centered MultiFab
+        IntVect MCC_nodal_flag = IntVect::TheCellVector();
+
+        if (do_MCC_energy_tracking) {
+            m_fields.alloc_init(FieldType::collision_energy_change, lev, amrex::convert(ba, MCC_nodal_flag), dm, MCC_ncomps, ngMCC, 0.0_rt);
+        }
+        if (do_MCC_ionization_tracking) {
+            m_fields.alloc_init(FieldType::collision_ionization, lev, amrex::convert(ba, MCC_nodal_flag), dm, MCC_ncomps, ngMCC, 0.0_rt);
+        }
+    }
+
     if (do_current_centering)
     {
         amrex::BoxArray const& nodal_ba = amrex::convert(ba, amrex::IntVect::TheNodeVector());