From aa51b3fc411b34e9db7fa221dada16e6731eab65 Mon Sep 17 00:00:00 2001
From: Alexander Steppke <alexander.steppke@psi.ch>
Date: Wed, 12 Nov 2025 16:53:59 +0100
Subject: [PATCH] finalized JF_stacks and unfinished scan in scan_info

---
 src/cristallina/analysis.py | 28 ++++++++++++++++++----------
 src/cristallina/utils.py    |  9 ++++++++-
 2 files changed, 26 insertions(+), 11 deletions(-)

diff --git a/src/cristallina/analysis.py b/src/cristallina/analysis.py
index 037337d..7566d1d 100644
--- a/src/cristallina/analysis.py
+++ b/src/cristallina/analysis.py
@@ -339,7 +339,17 @@ def calculate_JF_stacks(scan : SFScanInfo,
                         exclude_steps: list[int] | None = None,
                         recompute : bool=False,
                         detector : str = "JF16T03V02"):
-    """Calculate image stacks for JF detectors in a scan."""
+    """Calculate image stacks for JF detectors in a scan.
+    Args:
+        scan (SFScanInfo): The scan object containing scan steps.
+        lower_cutoff_threshold: Threshold to apply to images before summation in keV.
+        exclude_steps: List of step indices to exclude from processing. Defaults to None.
+        recompute: If True, forces recomputation even if cached results are available. Defaults to False.
+        detector: The detector channel to process. Defaults to "JF16T03V02" (JF 1.5M).
+    Returns:
+        stacks: List of summed image stacks for each step.
+        I0_norms: List of I0 normalizations for each step.
+    """
     
     stacks = []
     I0_norms = []
@@ -349,18 +359,17 @@ def calculate_JF_stacks(scan : SFScanInfo,
             logger.info(f"skipping step {i}")
             continue
 
-        #subset = step["JF17T16V01", "JF20T01V01"]
         JF_channels = [channels.JF, channels.JF_8M, channels.JF_I0, channels.GASMONITOR]
         available_channels = [ch.name for ch in step.channels]
         selected_channels = [ch for ch in JF_channels if ch in available_channels]
 
-        subset = step[*selected_channels] # only I0 for test cases
-        pids_before = len(subset.pids)
+        subset = step[*selected_channels]
+        pids_before = subset.pids.copy()
         subset.drop_missing()
-        pids_after = len(subset.pids)
+        pids_after = subset.pids.copy()
 
-        if pids_before != pids_after:
-            logger.warning(f"Step {i}: dropped {pids_before - pids_after} pulse IDs due to missing data.")
+        if set(pids_before) != set(pids_after):
+            logger.warning(f"Step {i}: dropped {set(pids_before) - set(pids_after)} pulse IDs due to missing data.")
         
         # we only consider the JF files here
         files = [f.fname for f in step.files if "JF" in f.fname]
@@ -371,11 +380,10 @@ def calculate_JF_stacks(scan : SFScanInfo,
             
         # TODO: define roi for I0 detector
         stack_I0 = perform_image_roi_crop(files, channel=channels.JF_I0, lower_cutoff=2,)
-        I0_norm = np.sum(stack_I0, axis=(1, 2))
+        I0_norm = np.sum(stack_I0, axis=(0, 1, 2))
         I0_norms.append(I0_norm)
 
-    return stacks, I0_norms
-
+    return np.array(stacks), np.array(I0_norms)
 
 
 @memory.cache(ignore=["batch_size"])  # we ignore batch_size for caching purposes
diff --git a/src/cristallina/utils.py b/src/cristallina/utils.py
index c7fd491..4616624 100644
--- a/src/cristallina/utils.py
+++ b/src/cristallina/utils.py
@@ -29,6 +29,7 @@ import matplotlib.pyplot as plt
 from scipy.optimize import curve_fit
 
 from . import channels
+from .uscan import UnfinishedScanInfo
 
 
 def collect_runs_metadata(pgroup_data_dir: str | os.PathLike):
@@ -103,13 +104,19 @@ def collect_runs_metadata(pgroup_data_dir: str | os.PathLike):
     return df
 
 
-def scan_info(run_number, base_path=None, small_data=True, pgroup=None):
+def scan_info(run_number, base_path=None, small_data=True, pgroup=None, use_uscan=False):
     """Returns SFScanInfo object for a given run number.
     If there is are small data channels, they will be added (small_data=False to suppress their loading).
+
+    use_uscan=True will use UnfinishedScanInfo to load scans that are still running.
     """
     if base_path == None:
         base_path = heuristic_extract_base_path(pgroup)
 
+    if use_uscan:
+        scan = UnfinishedScanInfo(f"{base_path}run{run_number:04}/meta/scan.json")
+        return scan
+    
     scan = SFScanInfo(f"{base_path}run{run_number:04}/meta/scan.json")
 
     try: