1 files changed, 38 insertions, 11 deletions

diff --git a/Source/Particles/MultiParticleContainer.cpp b/Source/Particles/MultiParticleContainer.cpp
index 983530569..9d39ec2f9 100644
--- a/Source/Particles/MultiParticleContainer.cpp
+++ b/Source/Particles/MultiParticleContainer.cpp
@@ -29,16 +29,31 @@ MultiParticleContainer::MultiParticleContainer (AmrCore* amr_core)
 
     pc_tmp.reset(new PhysicalParticleContainer(amr_core));
 
-    if (WarpX::do_boosted_frame_diagnostic && WarpX::do_boosted_frame_particles)
+    // Compute the number of species for which lab-frame data is dumped
+    // nspecies_lab_frame_diags, and map their ID to MultiParticleContainer
+    // particle IDs in map_species_lab_diags.
+    map_species_boosted_frame_diags.resize(nspecies);
+    nspecies_boosted_frame_diags = 0;
+    for (int i=0; i<nspecies; i++){
+        auto& pc = allcontainers[i];
+        if (pc->do_boosted_frame_diags){
+            map_species_boosted_frame_diags[nspecies_boosted_frame_diags] = i;
+            do_boosted_frame_diags = 1;
+            nspecies_boosted_frame_diags += 1;
+        }
+    }
+
+    if (WarpX::do_boosted_frame_diagnostic && do_boosted_frame_diags)
     {
-        for (int i = 0; i < nspecies + nlasers; ++i)
+        for (int i = 0; i < nspecies_boosted_frame_diags; ++i)
         {
-            allcontainers[i]->AddRealComp("xold");
-            allcontainers[i]->AddRealComp("yold");
-            allcontainers[i]->AddRealComp("zold");
-            allcontainers[i]->AddRealComp("uxold");
-            allcontainers[i]->AddRealComp("uyold");
-            allcontainers[i]->AddRealComp("uzold");
+            int is = map_species_boosted_frame_diags[i];
+            allcontainers[is]->AddRealComp("xold");
+            allcontainers[is]->AddRealComp("yold");
+            allcontainers[is]->AddRealComp("zold");
+            allcontainers[is]->AddRealComp("uxold");
+            allcontainers[is]->AddRealComp("uyold");
+            allcontainers[is]->AddRealComp("uzold");
         }
         pc_tmp->AddRealComp("xold");
         pc_tmp->AddRealComp("yold");
@@ -374,13 +389,25 @@ MultiParticleContainer
 
     BL_PROFILE("MultiParticleContainer::GetLabFrameData");
     
-    for (int i = 0; i < nspecies; ++i){
-        WarpXParticleContainer* pc = allcontainers[i].get();
+    // Loop over particle species
+    for (int i = 0; i < nspecies_boosted_frame_diags; ++i){
+        int isp = map_species_boosted_frame_diags[i];
+        WarpXParticleContainer* pc = allcontainers[isp].get();
         WarpXParticleContainer::DiagnosticParticles diagnostic_particles;
         pc->GetParticleSlice(direction, z_old, z_new, t_boost, t_lab, dt, diagnostic_particles);
-
+        // Here, diagnostic_particles[lev][index] is a WarpXParticleContainer::DiagnosticParticleData
+        // where "lev" is the AMR level and "index" is a [grid index][tile index] pair.
+        
+        // Loop over AMR levels
         for (int lev = 0; lev <= pc->finestLevel(); ++lev){
+            // Loop over [grid index][tile index] pairs 
+            // and Fills parts[species number i] with particle data from all grids and 
+            // tiles in diagnostic_particles. parts contains particles from all 
+            // AMR levels indistinctly.
             for (auto it = diagnostic_particles[lev].begin(); it != diagnostic_particles[lev].end(); ++it){
+                // it->first is the [grid index][tile index] key
+                // it->second is the corresponding 
+                // WarpXParticleContainer::DiagnosticParticleData value
                 parts[i].GetRealData(DiagIdx::w).insert(  parts[i].GetRealData(DiagIdx::w  ).end(),
                                                           it->second.GetRealData(DiagIdx::w  ).begin(),
                                                           it->second.GetRealData(DiagIdx::w  ).end());