From 0cee03d24dffb3843f25f6669f2d4d0ff2130f7e Mon Sep 17 00:00:00 2001
From: Seth R Johnson <johnsonsr@ornl.gov>
Date: Sat, 30 Nov 2024 13:46:02 -0500
Subject: [PATCH] Trim muppet data and fix vector allocation

---
 app/celer-export-geant.cc             |  2 +-
 src/celeritas/io/ImportDataTrimmer.cc | 27 +++++++++++++++++++++++----
 src/celeritas/io/ImportDataTrimmer.hh |  4 +++-
 3 files changed, 27 insertions(+), 6 deletions(-)

diff --git a/app/celer-export-geant.cc b/app/celer-export-geant.cc
index e09be54045..e63055b950 100644
--- a/app/celer-export-geant.cc
+++ b/app/celer-export-geant.cc
@@ -103,7 +103,7 @@ void run(std::string const& gdml_filename,
     if (gen_test)
     {
         ImportDataTrimmer::Input options;
-        options.max_size = 3;
+        options.mupp = true;
         ImportDataTrimmer trim(options);
         trim(&imported);
     }
diff --git a/src/celeritas/io/ImportDataTrimmer.cc b/src/celeritas/io/ImportDataTrimmer.cc
index 918a519380..e0002d59ab 100644
--- a/src/celeritas/io/ImportDataTrimmer.cc
+++ b/src/celeritas/io/ImportDataTrimmer.cc
@@ -41,6 +41,8 @@ void ImportDataTrimmer::operator()(ImportData* data)
         (*this)(&data->geo_materials);
         (*this)(&data->phys_materials);
         (*this)(&data->optical_materials);
+        (*this)(&data->neutron_elastic_data);
+        (*this)(&data->atomic_relaxation_data);
     }
 
     if (options_.physics)
@@ -52,6 +54,11 @@ void ImportDataTrimmer::operator()(ImportData* data)
         (*this)(&data->volumes);
     }
 
+    if (options_.mupp)
+    {
+        (*this)(&data->mu_pair_production_data);
+    }
+
     for (auto& e : data->elements)
     {
         (*this)(&e);
@@ -79,9 +86,15 @@ void ImportDataTrimmer::operator()(ImportData* data)
 
     (*this)(&data->sb_data);
     (*this)(&data->livermore_pe_data);
-    (*this)(&data->neutron_elastic_data);
-    (*this)(&data->atomic_relaxation_data);
-    (*this)(&data->mu_pair_production_data);
+
+    if (this->options_.physics)
+    {
+        for (auto&& kv : data->neutron_elastic_data)
+        {
+            (*this)(&kv.second);
+        }
+        // TODO: trim relaxation shells
+    }
 
     for (auto& m : data->optical_models)
     {
@@ -305,8 +318,14 @@ void ImportDataTrimmer::operator()(std::vector<T>* vec)
 {
     CELER_EXPECT(vec);
 
+    if (options_.max_size == numeric_limits<size_type>::max())
+    {
+        // Don't trim
+        return;
+    }
+
     std::vector<T> result;
-    result.reserve(options_.max_size);
+    result.reserve(std::min(options_.max_size + 1, vec->size()));
 
     auto filter = this->make_filterer(vec->size());
     for (auto i : range(vec->size()))
diff --git a/src/celeritas/io/ImportDataTrimmer.hh b/src/celeritas/io/ImportDataTrimmer.hh
index 0c7687a12f..b54593700a 100644
--- a/src/celeritas/io/ImportDataTrimmer.hh
+++ b/src/celeritas/io/ImportDataTrimmer.hh
@@ -27,8 +27,10 @@ class ImportDataTrimmer
         bool materials{false};
         //! Reduce the number of physics models and processes
         bool physics{false};
+        //! Reduce the MuPPET table fidelity
+        bool mupp{false};
         //! Maximum number of elements in a vector (approximate)
-        size_type max_size{numeric_limits<size_type>::max()};
+        std::size_t max_size{numeric_limits<std::size_t>::max()};
     };
 
   public: