1 files changed, 10 insertions, 7 deletions
diff --git a/utils/sddm/plot_energy.py b/utils/sddm/plot_energy.py
index 1fbe67d..84a7094 100755
--- a/utils/sddm/plot_energy.py
+++ b/utils/sddm/plot_energy.py
@@ -109,13 +109,13 @@ def tag_michels(ev):
     """
     Looks for Michel electrons after muons.
     """
-    time_since_last_prompt_plus_muon = ev.gtr - ev.gtr.where(ev.prompt | ((ev.dc & DC_MUON) != 0)).ffill()
-    ev['muon_gtid'] = ev.gtid.where(ev.prompt | ((ev.dc & DC_MUON) != 0)).ffill()
-    ev['muon_nhit'] = ev.nhit_cal.where(ev.prompt | ((ev.dc & DC_MUON) != 0)).ffill()
+    time_since_last_prompt = ev.gtr - ev.gtr.where(ev.prompt).ffill()
+    ev['muon_gtid'] = ev.gtid.where(ev.prompt).ffill()
+    ev['muon_nhit'] = ev.nhit_cal.where(ev.prompt).ffill()
     ev['michel'] = ~ev.prompt
     ev['michel'] &= ev.dc & (DC_JUNK | DC_CRATE_ISOTROPY | DC_QVNHIT | DC_FLASHER | DC_NECK | DC_ESUM | DC_OWL | DC_OWL_TRIGGER | DC_FTS) == 0
     ev['michel'] &= ev.nhit >= 100
-    ev['michel'] &= (time_since_last_prompt_plus_muon > 800) & (time_since_last_prompt_plus_muon < 200e3)
+    ev['michel'] &= (time_since_last_prompt > 800) & (time_since_last_prompt < 20e3)
     ev['muon_gtid'] = ev.muon_gtid.where(ev.michel,other=-1)
     ev['muon_nhit'] = ev.muon_nhit.where(ev.michel,other=-1)
     ev['stopping_muon'] = np.ones(len(ev),dtype=np.bool)
@@ -126,14 +126,17 @@ def atmospheric_events(ev):
     """
     Tags atmospheric events which have a neutron follower.
     """
+    time_since_last_prompt = ev.gtr - ev.gtr.where(ev.prompt).ffill()
+    ev['atm_gtid'] = ev.gtid.where(ev.prompt).ffill()
     ev['neutron'] = ~ev.prompt
     ev['neutron'] &= ev.dc & (DC_JUNK | DC_CRATE_ISOTROPY | DC_QVNHIT | DC_FLASHER | DC_NECK | DC_ESUM | DC_OWL | DC_OWL_TRIGGER | DC_FTS) == 0
     ev['neutron'] &= ~np.isnan(ev.ftp_x) & ~np.isnan(ev.rsp_energy)
     ev['neutron'] &= ev.ftp_r < AV_RADIUS
     ev['neutron'] &= ev.rsp_energy > 4.0
-    time_till_next_neutron = ev.gtr.where(ev.neutron).bfill() - ev.gtr
-    ev['atm'] = ev.prompt & (time_till_next_neutron > 20e3) & (time_till_next_neutron < 250e6)
-
+    ev['neutron'] &= (time_since_last_prompt > 20e3) & (time_since_last_prompt < 250e6)
+    ev['atm_gtid'] = ev.atm_gtid.where(ev.neutron,other=-1)
+    ev['atm'] = np.ones(len(ev),dtype=np.bool)
+    ev['atm'] = ev.atm.where(ev.gtid.isin(ev.atm_gtid[ev.atm_gtid > 0].values),other=False)
     return ev
 
 def gtid_sort(ev, first_gtid):