"""Helper class for parsing CrystFEL streams and matching them to HDF5 frames."""
from __future__ import annotations
import os
from typing import Any, Callable, Dict, List, Optional, Tuple
import h5py
[docs]
class IndexingStreamManager:
"""Encapsulates stream parsing, frame matching, and overlay caching."""
def __init__(
self,
stream_parser,
get_hdf5_path: Callable[[], Optional[str]],
get_total_frames: Callable[[], Optional[int]],
log_fn: Optional[Callable[[str], None]] = None,
):
self._stream_parser = stream_parser
self._get_hdf5_path = get_hdf5_path
self._get_total_frames = get_total_frames
self._log_fn = log_fn or (lambda _msg: None)
self.stream_path: Optional[str] = None
self._last_error: Optional[str] = None
self._parsed_stream = None
self._parsed_stream_path = None
self._chunk_sequence: List = []
self._chunks_by_event: Dict = {}
self._chunks_by_serial: Dict = {}
self._chunk_indices: Dict[int, int] = {}
self._frame_identifier_arrays: Dict[str, Optional[Any]] = {}
self._frame_lookup: Dict[str, Dict[int, int]] = {'event': {}, 'serial': {}, 'index': {}}
self._frame_ids_loaded = False
self._reflection_cache: Dict[int, List[Tuple[float, float]]] = {}
self._stream_peak_cache: Dict[int, List[Tuple[float, float]]] = {}
self._geom_min_fs: Optional[float] = None
self._geom_max_fs: Optional[float] = None
self._geom_min_ss: Optional[float] = None
self._geom_max_ss: Optional[float] = None
self._geom_panel_extents: Dict[str, Dict[str, float]] = {}
def _log(self, message: str):
try:
self._log_fn(message)
except Exception:
pass
[docs]
def has_parser(self) -> bool:
return self._stream_parser is not None
[docs]
def has_stream_file(self) -> bool:
return bool(self.stream_path and os.path.exists(self.stream_path))
[docs]
def is_stream_available(self) -> bool:
return self.has_parser() and self.has_stream_file()
[docs]
def is_loaded(self) -> bool:
return bool(
self._parsed_stream is not None
and self._parsed_stream_path == self.stream_path
)
[docs]
def set_stream_path(self, stream_path: Optional[str]):
self.stream_path = stream_path if stream_path else None
self._reset_cached_data(reset_geom=True)
if self.stream_path:
self._extract_geom_extents_from_run(self.stream_path)
[docs]
def reset_for_new_hdf5(self):
"""Invalidate caches when the loaded HDF5 changes."""
self._reset_cached_data(reset_geom=False)
def _reset_cached_data(self, reset_geom: bool):
self._parsed_stream = None
self._parsed_stream_path = None
self._chunk_sequence = []
self._chunks_by_event = {}
self._chunks_by_serial = {}
self._chunk_indices = {}
self._frame_identifier_arrays = {}
self._frame_lookup = {'event': {}, 'serial': {}, 'index': {}}
self._frame_ids_loaded = False
self._reflection_cache = {}
self._stream_peak_cache = {}
self._last_error = None
if reset_geom:
self._geom_min_fs = None
self._geom_max_fs = None
self._geom_min_ss = None
self._geom_max_ss = None
self._geom_panel_extents = {}
[docs]
def ensure_loaded(self) -> bool:
"""Parse the stream file if needed."""
if (
self._parsed_stream is not None
and self._parsed_stream_path == self.stream_path
):
return True
if not self.has_stream_file() or not self.has_parser():
return False
try:
parsed = self._stream_parser(self.stream_path)
self._log(
f"Parsed stream '{self.stream_path}' with {len(parsed.chunks)} chunks."
)
except Exception as exc:
self._last_error = str(exc)
self._log(f"Failed to parse stream '{self.stream_path}': {exc}")
return False
self._parsed_stream = parsed
self._parsed_stream_path = self.stream_path
self._chunk_sequence = parsed.chunks
self._chunks_by_event = {}
self._chunks_by_serial = {}
self._chunk_indices = {}
for idx, chunk in enumerate(parsed.chunks):
if chunk is None:
continue
self._chunk_indices[id(chunk)] = idx
try:
if chunk.event is not None:
try:
event_key = int(chunk.event)
except Exception:
event_key = chunk.event
self._chunks_by_event[event_key] = chunk
except AttributeError:
pass
try:
if chunk.serial is not None:
try:
serial_key = int(chunk.serial)
except Exception:
serial_key = chunk.serial
self._chunks_by_serial[serial_key] = chunk
except AttributeError:
pass
self._reflection_cache = {}
self._stream_peak_cache = {}
self._last_error = None
return True
[docs]
def get_parsed_stream(self):
"""Return the parsed stream object, loading it if needed."""
if not self.ensure_loaded():
return None
return self._parsed_stream
[docs]
def has_chunks(self) -> bool:
return bool(self._chunk_sequence)
[docs]
def get_last_error(self) -> Optional[str]:
return self._last_error
def _ensure_frame_identifiers_loaded(self):
if self._frame_ids_loaded:
return
h5_path = self._get_hdf5_path()
if not h5_path or not os.path.exists(h5_path):
self._frame_ids_loaded = True
return
dataset_map = {
'event': [
'event_number', 'eventNumber', 'event_numbers', 'events',
'event', 'pulse_id', 'pulseId'
],
'serial': [
'serial_number', 'serialNumber', 'serial', 'image_serial_number'
],
'index': [
'index', 'frame_index', 'frameIndex'
],
}
try:
with h5py.File(h5_path, 'r') as file:
if 'entry/data' not in file:
return
data = file['entry/data']
for kind, candidates in dataset_map.items():
if (
kind in self._frame_identifier_arrays
and self._frame_identifier_arrays[kind] is not None
):
continue
for name in candidates:
if name in data:
try:
arr = data[name][()]
self._frame_identifier_arrays[kind] = arr
self._populate_frame_lookup(kind, arr)
except Exception:
self._frame_identifier_arrays[kind] = None
break
except Exception as exc:
self._log(f"Failed to load frame identifiers: {exc}")
finally:
self._frame_ids_loaded = True
def _populate_frame_lookup(self, kind: str, arr):
if arr is None or kind not in self._frame_lookup:
return
lookup = self._frame_lookup.setdefault(kind, {})
lookup.clear()
try:
for idx, val in enumerate(arr):
try:
lookup[int(val)] = idx
except (TypeError, ValueError):
continue
except Exception:
pass
def _get_frame_identifier_value(self, kind: str, frame_index: int):
arr = self._frame_identifier_arrays.get(kind)
if arr is None:
return None
try:
if frame_index < len(arr):
return int(arr[frame_index])
except Exception:
return None
return None
[docs]
def get_chunk_for_frame(self, frame_index: int):
if not self.ensure_loaded():
return None
self._ensure_frame_identifiers_loaded()
event_val = self._get_frame_identifier_value('event', frame_index)
if event_val is not None and event_val in self._chunks_by_event:
return self._chunks_by_event[event_val]
if (
self._frame_identifier_arrays.get('event') is None
and self._chunks_by_event
):
event_key = frame_index
if event_key in self._chunks_by_event:
return self._chunks_by_event[event_key]
return None
[docs]
def get_chunk_index(self, chunk) -> Optional[int]:
if chunk is None:
return None
idx = self._chunk_indices.get(id(chunk))
if idx is not None:
return idx
try:
idx = self._chunk_sequence.index(chunk)
self._chunk_indices[id(chunk)] = idx
return idx
except (ValueError, AttributeError):
return None
[docs]
def get_frame_index_for_chunk(self, chunk) -> Optional[int]:
if chunk is None:
return None
self._ensure_frame_identifiers_loaded()
try:
event_val = getattr(chunk, 'event', None)
if event_val is not None:
lookup = self._frame_lookup.get('event') or {}
event_int = int(event_val)
if event_int in lookup:
return lookup[event_int]
if self._frame_identifier_arrays.get('event') is None:
total_frames = self._get_total_frames() or 0
if 0 <= event_int < total_frames:
return event_int
except Exception:
pass
return None
def _chunk_has_reflections(self, chunk) -> bool:
if not chunk or not getattr(chunk, 'crystals', None):
return False
for xtal in chunk.crystals:
if getattr(xtal, 'reflections', None):
return True
return False
[docs]
def find_chunk_with_reflections_from(self, start_index: int, step: int):
if not self.ensure_loaded():
return None, None
if not self._chunk_sequence:
return None, None
idx = start_index
while 0 <= idx < len(self._chunk_sequence):
chunk = self._chunk_sequence[idx]
if self._chunk_has_reflections(chunk):
return idx, chunk
idx += step
return None, None
[docs]
def get_last_chunk_with_reflections(self):
if not self.ensure_loaded():
return None, None
return self.find_chunk_with_reflections_from(len(self._chunk_sequence) - 1, -1)
[docs]
def get_reflections_for_frame(
self,
frame_index: int,
chunk,
image_shape: Tuple[Optional[int], Optional[int]] = (None, None),
):
if chunk is None:
return []
if frame_index in self._reflection_cache:
return self._reflection_cache[frame_index]
reflections: List[Tuple[float, float]] = []
total = 0
if getattr(chunk, 'crystals', None):
for crystal in chunk.crystals:
if not getattr(crystal, 'reflections', None):
continue
for refl in crystal.reflections:
panel_id = self._extract_panel_id(refl)
fs, ss = self._stream_coords_to_image_coords(
getattr(refl, 'fs_px', None),
getattr(refl, 'ss_px', None),
panel_id,
image_shape,
)
reflections.append(
(fs, ss)
)
total += 1
self._log(f"Frame {frame_index}: loaded {total} reflections from stream.")
else:
self._log(f"Frame {frame_index}: stream chunk has no crystals.")
self._reflection_cache[frame_index] = reflections
return reflections
[docs]
def get_stream_peaks_for_frame(
self,
frame_index: int,
chunk,
image_shape: Tuple[Optional[int], Optional[int]],
):
if chunk is None:
return []
if frame_index in self._stream_peak_cache:
return self._stream_peak_cache[frame_index]
peaks: List[Tuple[float, float]] = []
if getattr(chunk, 'peaks', None):
for peak in chunk.peaks:
panel_id = self._extract_panel_id(peak)
fs, ss = self._stream_coords_to_image_coords(
getattr(peak, 'fs_px', None),
getattr(peak, 'ss_px', None),
panel_id,
image_shape,
)
peaks.append((fs, ss))
self._log(f"Frame {frame_index}: loaded {len(peaks)} stream peaks.")
else:
self._log(f"Frame {frame_index}: no stream peaks available.")
self._stream_peak_cache[frame_index] = peaks
return peaks
@staticmethod
def _extract_panel_id(item) -> Optional[str]:
if item is None:
return None
for attr in ('panel', 'panel_name', 'panel_id', 'panelname'):
val = getattr(item, attr, None)
if val is None:
continue
text = str(val).strip().lower()
if text:
return text
return None
def _stream_coords_to_image_coords(
self,
fs: Optional[float],
ss: Optional[float],
panel_id: Optional[str],
image_shape: Tuple[Optional[int], Optional[int]],
):
if fs is None or ss is None:
return fs, ss
try:
fs_val = float(fs)
ss_val = float(ss)
except (TypeError, ValueError):
return None, None
panel_key = str(panel_id).strip().lower() if panel_id else None
panel = self._geom_panel_extents.get(panel_key) if panel_key else None
if panel is not None:
min_fs = panel.get('min_fs')
max_fs = panel.get('max_fs')
min_ss = panel.get('min_ss')
max_ss = panel.get('max_ss')
if None not in (min_fs, max_fs, min_ss, max_ss):
panel_w = float(max_fs - min_fs + 1.0)
panel_h = float(max_ss - min_ss + 1.0)
# CrystFEL stream coordinates can be panel-local (most common) or
# already in global image coordinates. Detect panel-local first.
if -1.0 <= fs_val <= panel_w and -1.0 <= ss_val <= panel_h:
return fs_val + float(min_fs), ss_val + float(min_ss)
if (
float(min_fs) - 1.0 <= fs_val <= float(max_fs) + 1.0
and float(min_ss) - 1.0 <= ss_val <= float(max_ss) + 1.0
):
return fs_val, ss_val
return self._apply_global_stream_offset(fs_val, ss_val, image_shape)
def _apply_global_stream_offset(
self,
fs: Optional[float],
ss: Optional[float],
image_shape: Tuple[Optional[int], Optional[int]],
):
if fs is None or ss is None:
return fs, ss
height, width = image_shape
if self._geom_min_fs is not None and self._geom_max_fs is not None:
geom_width = self._geom_max_fs - self._geom_min_fs + 1
if geom_width > 0:
fs = (fs - self._geom_min_fs)
if width:
fs *= (width / geom_width)
if self._geom_min_ss is not None and self._geom_max_ss is not None:
geom_height = self._geom_max_ss - self._geom_min_ss + 1
if geom_height > 0:
ss = (ss - self._geom_min_ss)
if height:
ss *= (height / geom_height)
return fs, ss
[docs]
def compare_peaks(self, h5_peaks, stream_peaks, tolerance: float = 0.6):
"""Return counts of unmatched peaks between HDF5 and stream within +/- tolerance."""
stream_remaining = list(stream_peaks)
unmatched_h5 = 0
for h_fs, h_ss in h5_peaks:
matched = False
for idx, (s_fs, s_ss) in enumerate(stream_remaining):
if abs(h_fs - s_fs) <= tolerance and abs(h_ss - s_ss) <= tolerance:
del stream_remaining[idx]
matched = True
break
if not matched:
unmatched_h5 += 1
unmatched_stream = len(stream_remaining)
return unmatched_h5, unmatched_stream
def _extract_geom_extents_from_run(self, stream_path: Optional[str]):
self._geom_min_fs = None
self._geom_max_fs = None
self._geom_min_ss = None
self._geom_max_ss = None
self._geom_panel_extents = {}
if not stream_path:
return
directory = os.path.dirname(stream_path)
if not directory or not os.path.isdir(directory):
return
geom_path = os.path.join(directory, 'geometry.geom')
if not os.path.exists(geom_path):
return
try:
with open(geom_path, 'r') as geom_file:
for line in geom_file:
s = line
# Strip inline comments for robust parsing.
for marker in (';', '#'):
if marker in s:
s = s.split(marker, 1)[0]
s = s.strip()
if '=' not in s:
continue
key, val = s.split('=', 1)
key = key.strip().lower()
val = val.strip()
try:
num = float(val.split()[0])
except Exception:
continue
field = key
panel_id = None
if '/' in key:
panel_id, field = key.split('/', 1)
panel_id = panel_id.strip()
field = field.strip()
if panel_id and panel_id.startswith('bad'):
continue
if field not in {'min_fs', 'max_fs', 'min_ss', 'max_ss'}:
continue
if field == 'min_fs':
self._geom_min_fs = num if self._geom_min_fs is None else min(self._geom_min_fs, num)
elif field == 'max_fs':
self._geom_max_fs = num if self._geom_max_fs is None else max(self._geom_max_fs, num)
elif field == 'min_ss':
self._geom_min_ss = num if self._geom_min_ss is None else min(self._geom_min_ss, num)
elif field == 'max_ss':
self._geom_max_ss = num if self._geom_max_ss is None else max(self._geom_max_ss, num)
if panel_id:
ext = self._geom_panel_extents.setdefault(panel_id, {})
ext[field] = num
except Exception as exc:
self._log(f"Failed to read geometry extents: {exc}")