1 files changed, 29 insertions, 4 deletions
diff --git a/utils/sddm/plot_energy.py b/utils/sddm/plot_energy.py
index 391fa66..f81df67 100644
--- a/utils/sddm/plot_energy.py
+++ b/utils/sddm/plot_energy.py
@@ -19,11 +19,11 @@ import numpy as np
 from scipy.stats import poisson, dirichlet, multinomial, beta
 from scipy.special import spence, gammaln
 from .dc import DC_MUON, DC_JUNK, DC_CRATE_ISOTROPY, DC_QVNHIT, DC_NECK, DC_FLASHER, DC_ESUM, DC_OWL, DC_OWL_TRIGGER, DC_FTS, DC_ITC, DC_BREAKDOWN
-from . import grouper, print_warning, AV_RADIUS, PSUP_RADIUS, read_hdf
+from . import grouper, print_warning, AV_RADIUS, PSUP_RADIUS, read_hdf, write_hdf
 import pandas as pd
 from scipy.special import logsumexp
-from os.path import exists, join
-from os import environ
+from os.path import exists, join, split, getmtime
+from os import environ, mkdir, remove
 import hashlib
 
 EPSILON = 1e-10
@@ -482,7 +482,21 @@ def burst_cut(ev):
     return ev.groupby(burst.cumsum()).filter(lambda ev: len(ev[ev.prompt_50]) <= BURST_MAX_EVENTS).reset_index(drop=True)
 
 def get_events(filenames, merge_fits=False, nhit_thresh=None, mc=False):
-    ev = pd.concat([read_hdf(filename, "ev").assign(filename=filename) for filename in filenames],ignore_index=True)
+    h = hashlib.sha1()
+
+    for filename in filenames:
+        head, tail = split(filename)
+        mtime = getmtime(filename)
+        h.update(tail.encode("utf-8"))
+        h.update(str(mtime).encode("utf-8"))
+
+    hash_filename = join(environ['HOME'],'.cached_hdf5','%s.hdf5' % h.hexdigest())
+
+    if exists(hash_filename) and merge_fits:
+        print("loading result from '%s'" % hash_filename)
+        return read_hdf(hash_filename,"ev")
+
+    ev = pd.concat([read_hdf(filename, "ev") for filename in filenames],ignore_index=True)
     fits = pd.concat([read_hdf(filename, "fits") for filename in filenames],ignore_index=True)
     rhdr = pd.concat([read_hdf(filename, "rhdr") for filename in filenames],ignore_index=True)
     if mc:
@@ -727,10 +741,21 @@ def get_events(filenames, merge_fits=False, nhit_thresh=None, mc=False):
                       np.cos(ev_single_particle.theta1)*ev_single_particle.z
         ev['udotr'] /= ev_single_particle.r
 
+        ev = ev.reset_index()
+
         # Yes, this is super hacky, but I don't want to change what's returned from
         # this function.
         ev.attrs = {'time_pulse_gt': time_pulse_gt, 'time_10_mhz': time_10_mhz}
 
+        head, tail = split(hash_filename)
+        if not exists(head):
+            mkdir(head)
+
+        try:
+            write_hdf(hash_filename,ev,"ev")
+        except KeyboardInterrupt as e:
+            remove(hash_filename)
+
         return ev
 
     # Yes, this is super hacky, but I don't want to change what's returned from