From 4287df850fabfa397a4135a0d08b5652349b07b6 Mon Sep 17 00:00:00 2001
From: Axel Huebl <axel.huebl@plasma.ninja>
Date: Tue, 14 Jan 2025 15:45:05 -0800
Subject: [PATCH] Fix Lost Particle w/ Runtime Attr (#795)

* Fix Lost Particle w/ Runtime Attr

For some processes and/or diagnostics, we add extra runtime
attributes to our beam (particle container). The logic that
collected "lost" particles in the beamline, i.e., as marked as
lost in apertures, did not account for extra runtime attributes
and thus got lost in bookkeeping.

This fixes the collection logic to be more robust and also copy
any extra runtime attributes over to the "lost" particle recording.

This was first seen with an input that used the nonlinear lens (NLL),
NLL-invariant diagnostics and an aperture at the same time.

* More Robust Implementation
---
 src/ImpactX.cpp               |  7 -------
 src/particles/CollectLost.cpp | 32 +++++++++++++++++++++++++++++++-
 2 files changed, 31 insertions(+), 8 deletions(-)

diff --git a/src/ImpactX.cpp b/src/ImpactX.cpp
index 2b356cb72..bc71e0de1 100644
--- a/src/ImpactX.cpp
+++ b/src/ImpactX.cpp
@@ -90,13 +90,6 @@ namespace impactx {
         amr_data->InitFromScratch(0.0);
 
         // alloc particle containers
-        //   the lost particles have an extra runtime attribute: s when it was lost
-        if (!amr_data->m_particles_lost->HasRealComp("s_lost"))
-        {
-            bool comm = true;
-            amr_data->m_particles_lost->AddRealComp("s_lost", comm);
-        }
-
         //   have to resize here, not in the constructor because grids have not
         //   been built when constructor was called.
         amr_data->m_particle_container->reserveData();
diff --git a/src/particles/CollectLost.cpp b/src/particles/CollectLost.cpp
index d2390c268..24bfd575b 100644
--- a/src/particles/CollectLost.cpp
+++ b/src/particles/CollectLost.cpp
@@ -55,8 +55,38 @@ namespace impactx
         BL_PROFILE("impactX::collect_lost_particles");
 
         using SrcData = ImpactXParticleContainer::ParticleTileType::ConstParticleTileDataType;
-
         ImpactXParticleContainer& dest = *source.GetLostParticleContainer();
+
+        // Check destination has the same attributes as source + "s_lost"
+        for (auto & name : source.RealSoA_names())
+        {
+            amrex::Print() << "name: " << name << std::endl;
+            if (!dest.HasRealComp(name)) {
+                amrex::Print() << "adding " << name << std::endl;
+                dest.AddRealComp(name);
+            }
+            AMREX_ALWAYS_ASSERT_WITH_MESSAGE(source.GetRealCompIndex(name) == dest.GetRealCompIndex(name),
+                                             "Source and destination Real attributes misaligned!");
+        }
+        for (auto & name : source.intSoA_names())
+        {
+            if (!dest.HasIntComp(name)) {
+                dest.AddIntComp(name);
+            }
+            AMREX_ALWAYS_ASSERT_WITH_MESSAGE(source.GetIntCompIndex(name) == dest.GetIntCompIndex(name),
+                                             "Source and destination Int attributes misaligned!");
+        }
+        // the lost particles have an extra runtime attribute: s when it was lost
+        if (!dest.HasRealComp("s_lost"))
+        {
+            bool comm = true;
+            dest.AddRealComp("s_lost", comm);
+        }
+        AMREX_ALWAYS_ASSERT_WITH_MESSAGE(source.RealSoA_names().size() + 1 == dest.RealSoA_names().size(),
+                                         "Source and destination have different Real attributes!");
+        AMREX_ALWAYS_ASSERT_WITH_MESSAGE(source.intSoA_names().size() == dest.intSoA_names().size(),
+                                         "Source and destination have different Int attributes!");
+
         const int s_runtime_index = dest.GetRealCompIndex("s_lost") - dest.NArrayReal;
 
         RefPart const ref_part = source.GetRefParticle();