Include nearly all seed info post-synthesis.

Author: tymorrow

riid/data/synthetic/base.py

@@ -14,6 +14,9 @@
                                  _get_utc_timestamp)
 
 
+IGNORED_INFO_COLS = ["snr", "live_time", "real_time", "total_counts", "timestamp"]
+
+
 class Synthesizer():
     """Base class for synthesizers."""
 
@@ -159,12 +162,9 @@ def _get_minimal_ss(spectra, sources, live_times, snrs, total_counts=None) -> Sa
     ss.spectra_state = SpectraState.Counts
     ss.spectra = pd.DataFrame(spectra)
     ss.sources = sources
-    ss.info.description = np.full(n_samples, "")  # Ensures the length of info equal n_samples
     ss.info.snr = snrs
     ss.info.total_counts = total_counts if total_counts is not None else spectra.sum(axis=1)
     ss.info.live_time = live_times
-    ss.info.occupancy_flag = 0
-    ss.info.tag = " "  # TODO: test if this can be an empty string
 
     return ss
 
@@ -259,3 +259,17 @@ def get_samples_per_seed(columns: pd.MultiIndex, min_samples_per_seed: int, bala
     total_samples_expected = sum([x * y for x, y in zip(occurences, samples_per_level_value)])
 
     return lv_to_samples_per_seed, total_samples_expected
+
+
+def set_ss_info_from_seed_info(ss: SampleSet, info: pd.Series, timestamp: str,
+                               ignored_cols=IGNORED_INFO_COLS):
+    if ss is None:
+        return
+
+    info_indices_to_set = info.index.difference(ignored_cols)
+    info_to_set = info[info_indices_to_set]
+    for c, v in info_to_set.items():
+        ss.info[c] = v
+
+    ss.info.real_time = ss.info.live_time / (1 - info.dead_time_prop)
+    ss.info.timestamp = timestamp

riid/data/synthetic/passby.py

@@ -12,7 +12,8 @@
 from numpy.random import Generator
 
 from riid import SampleSet
-from riid.data.synthetic.base import Synthesizer, get_distribution_values
+from riid.data.synthetic.base import (Synthesizer, get_distribution_values,
+                                      set_ss_info_from_seed_info)
 
 
 class PassbySynthesizer(Synthesizer):
@@ -282,45 +283,16 @@ def generate(self, fg_seeds_ss: SampleSet, bg_seeds_ss: SampleSet,
                                fg_sources, bg_sources)
                     args.append(pb_args)
 
-        # TODO: follow prevents periodic progress reports
+        # TODO: the following prevents periodic progress reports
         passbys = []
         for a in args:
             f, fwhm, snr, dwell_time, fg_pmf, bg_pmf, fg_sources, bg_sources = a
             fg_passby_ss, gross_passby_ss = self._generate_single_passby(
                 fwhm, snr, dwell_time, fg_pmf, bg_pmf, fg_sources, bg_sources
             )
-            live_times = None
-            if fg_passby_ss is not None:
-                live_times = fg_passby_ss.info.live_time
-            elif gross_passby_ss is not None:
-                live_times = gross_passby_ss.info.live_time
-            else:
-                live_times = 1.0
-
-            fg_seed_ecal = fg_seeds_ss.ecal[f]
             fg_seed_info = fg_seeds_ss.info.iloc[f]
-            batch_rt_targets = live_times * (1 - fg_seed_info.dead_time_prop)
-            fg_seed_distance_cm = fg_seed_info.distance_cm
-            fg_seed_dead_time_prop = fg_seed_info.dead_time_prop
-            fg_seed_ad = fg_seed_info.areal_density
-            fg_seed_an = fg_seed_info.atomic_number
-            fg_seed_neutron_counts = fg_seed_info.neutron_counts
-
-            def _set_remaining_info(ss):
-                if ss is None:
-                    return
-                ss: SampleSet = ss
-                ss.ecal = fg_seed_ecal
-                ss.info.real_time = batch_rt_targets
-                ss.info.distance_cm = fg_seed_distance_cm
-                ss.info.dead_time_prop = fg_seed_dead_time_prop
-                ss.info.areal_density = fg_seed_ad
-                ss.info.atomic_number = fg_seed_an
-                ss.info.neutron_counts = fg_seed_neutron_counts
-                ss.info.timestamp = self._synthesis_start_dt
-
-            _set_remaining_info(fg_passby_ss)
-            _set_remaining_info(gross_passby_ss)
+            set_ss_info_from_seed_info(fg_passby_ss, fg_seed_info, self._synthesis_start_dt)
+            set_ss_info_from_seed_info(gross_passby_ss, fg_seed_info, self._synthesis_start_dt)
 
             passbys.append((fg_passby_ss, gross_passby_ss))
 

riid/data/synthetic/static.py

@@ -11,7 +11,8 @@
 from numpy.random import Generator
 
 from riid import SampleSet, SpectraState, SpectraType
-from riid.data.synthetic.base import Synthesizer, get_distribution_values
+from riid.data.synthetic.base import (Synthesizer, get_distribution_values,
+                                      set_ss_info_from_seed_info)
 
 
 class StaticSynthesizer(Synthesizer):
@@ -190,30 +191,9 @@ def _get_synthetic_samples(self, fg_seeds_ss: SampleSet, bg_seeds_ss: SampleSet,
                 batch_snr_targets
             )
 
-            fg_seed_ecal = fg_seeds_ss.ecal[f]
             fg_seed_info = fg_seeds_ss.info.iloc[f]
-            batch_rt_targets = batch_lt_targets / (1 - fg_seed_info.dead_time_prop)
-            fg_seed_distance_cm = fg_seed_info.distance_cm
-            fg_seed_dead_time_prop = fg_seed_info.dead_time_prop
-            fg_seed_ad = fg_seed_info.areal_density
-            fg_seed_an = fg_seed_info.atomic_number
-            fg_seed_neutron_counts = fg_seed_info.neutron_counts
-
-            def _set_remaining_info(ss):
-                if ss is None:
-                    return
-                ss: SampleSet = ss
-                ss.ecal = fg_seed_ecal
-                ss.info.real_time = batch_rt_targets
-                ss.info.distance_cm = fg_seed_distance_cm
-                ss.info.dead_time_prop = fg_seed_dead_time_prop
-                ss.info.areal_density = fg_seed_ad
-                ss.info.atomic_number = fg_seed_an
-                ss.info.neutron_counts = fg_seed_neutron_counts
-                ss.info.timestamp = self._synthesis_start_dt
-
-            _set_remaining_info(fg_batch_ss)
-            _set_remaining_info(gross_batch_ss)
+            set_ss_info_from_seed_info(fg_batch_ss, fg_seed_info, self._synthesis_start_dt)
+            set_ss_info_from_seed_info(gross_batch_ss, fg_seed_info, self._synthesis_start_dt)
 
             fg_ss_batches.append(fg_batch_ss)
             gross_ss_batches.append(gross_batch_ss)