From f78d30f10ad3e4187dc081f7dc2e139d1042c46f Mon Sep 17 00:00:00 2001
From: Alexander Steppke <alexander.steppke@psi.ch>
Date: Tue, 30 Aug 2022 16:49:33 +0200
Subject: [PATCH] added heuristic pgroup for cachedir

---
 src/cristallina/analysis.py | 65 +++++++++++++++++++++++++++++++++++--
 src/cristallina/plot.py     |  8 +++--
 2 files changed, 68 insertions(+), 5 deletions(-)

diff --git a/src/cristallina/analysis.py b/src/cristallina/analysis.py
index 2956554..62ede27 100644
--- a/src/cristallina/analysis.py
+++ b/src/cristallina/analysis.py
@@ -9,11 +9,19 @@ import joblib
 from joblib import Parallel, delayed, Memory
 
 # TODO: generalize this for all analysis, i.e. find appropriate p-group
-location = "/das/work/units/cristallina/p19739/cachedir"
-memory = Memory(location, verbose=0, compress=2)
 
+from . import utils
 from .utils import ROI
 
+try:
+    pgroup = utils.heuristic_extract_pgroup()
+    location = f"/das/work/units/cristallina/p{pgroup}/cachedir"
+except KeyError as e:
+    print(e)
+    location = "/das/work/units/cristallina/p19739/cachedir"
+
+memory = Memory(location, verbose=0, compress=2)
+
 
 @memory.cache(ignore=["batch_size"])  # we ignore batch_size for caching purposes
 def sum_images(
@@ -65,3 +73,56 @@ def sum_images(
                 break
 
     return res
+
+@memory.cache(ignore=["batch_size"])  # we ignore batch_size for caching purposes
+def get_contrast_images(
+    fileset,
+    channel="JF16T03V01",
+    alignment_channels=None,
+    batch_size=10,
+    roi: ROI = None,
+    preview=False,
+):
+    """
+    Sums a given region of interest (roi) for an image channel from a
+    given fileset (e.g. "run0352/data/acq0001.*.h5").
+
+    Allows alignment, i.e. reducing only to a common subset with other channels.
+
+    Summation is performed in batches, preview only sums and returns the first batch.
+    """
+
+    with SFDataFiles(*fileset) as data:
+        if alignment_channels is not None:
+            channels = [channel] + [ch for ch in alignment_channels]
+        else:
+            channels = [channel]
+
+        subset = data[channels]
+
+        subset.drop_missing()
+
+        Images = subset[channel]
+
+        res = defaultdict(list)
+        res["roi"] = repr(roi)
+
+        for image_slice in Images.in_batches(batch_size):
+
+            index_slice, im = image_slice
+
+            if roi is None:
+                im_ROI = im[:]
+            else:
+                im_ROI = im[:, roi.rows, roi.cols]
+
+            res[f"{channel}_mean"].extend(np.mean(im_ROI, axis=(1, 2)))
+            res[f"{channel}_std"].extend(np.std(im_ROI, axis=(1, 2)))
+
+            res["pids"].extend(Images.pids[index_slice])
+
+            # only return first batch
+            if preview:
+                break
+
+    return res
diff --git a/src/cristallina/plot.py b/src/cristallina/plot.py
index 526c08b..85d2b13 100644
--- a/src/cristallina/plot.py
+++ b/src/cristallina/plot.py
@@ -17,8 +17,10 @@ from pathlib import Path
 from sfdata import SFDataFiles, sfdatafile, SFScanInfo
 import jungfrau_utils as ju
 
+from . import utils
 from .utils import ROI
 
+
 def ju_patch_less_verbose(ju_module):
     """Quick monkey patch to suppress verbose messages from gain & pedestal file searcher."""
     ju_module.swissfel_helpers._locate_gain_file = ju_module.swissfel_helpers.locate_gain_file
@@ -188,7 +190,7 @@ def plot_image_channel(data, channel_name, pulse=0, ax=None, rois=None, norms=No
     axis_styling(ax, channel_name, description)
     plt.legend(loc=4)    
 
-def plot_spectrum_channel(data, channel_name_x, channel_name_y, average=True, frame=0, ax=None):
+def plot_spectrum_channel(data, channel_name_x, channel_name_y, average=True, pulse=0, ax=None):
     """
     Plots channel data for two channels where the first is taken as the (constant) x-axis
     and the second as the y-axis (here we take by default the mean over the individual pulses).
@@ -203,7 +205,7 @@ def plot_spectrum_channel(data, channel_name_x, channel_name_y, average=True, fr
     if average:
         y_data = mean_over_frames
     else:
-        y_data = data[channel_name_y].data[frame]
+        y_data = data[channel_name_y].data[pulse]
         
 
     if ax is None:
@@ -212,4 +214,4 @@ def plot_spectrum_channel(data, channel_name_x, channel_name_y, average=True, fr
     ax.plot(data[channel_name_x].data[0], y_data)
     description = None # f"mean: {mean:.2e},\nstd: {std:.2e}"
     ax.set_xlabel(channel_name_x)
-    axis_styling(ax, channel_name_y, description)
\ No newline at end of file
+    axis_styling(ax, channel_name_y, description)