177 lines
6.4 KiB
Python
177 lines
6.4 KiB
Python
"""
|
|
Pipeline processing for SATOP31-ATT01 arrival time tool.
|
|
|
|
Produces per-shot AND N-shot averaged outputs:
|
|
- Single-shot: {device}:signal, edge_pos, xcorr, xcorr_ampl, arrival_time, ...
|
|
- Averaged: {device}:avg_signal, avg_edge_pos, avg_xcorr, avg_xcorr_ampl, avg_arrival_time
|
|
|
|
The averaged channels are always present once the first valid FEL-on shot
|
|
arrives (average of 1 = the shot itself), so output types and sizes never
|
|
change mid-stream. The averaging depth grows until it reaches
|
|
``params["avg_length"]`` (default 50), after which it is a sliding window.
|
|
"""
|
|
|
|
from collections import deque
|
|
from logging import getLogger
|
|
import numpy as np
|
|
from scipy import signal
|
|
|
|
_logger = getLogger(__name__)
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Module-level state (persists across calls within one pipeline instance)
|
|
# ---------------------------------------------------------------------------
|
|
buffer_dark = deque()
|
|
buffer_savgol = deque()
|
|
buffer_signal = None # deque for N-shot averaged signal
|
|
initialized = False
|
|
|
|
|
|
def initialize(params):
|
|
global buffer_dark, buffer_savgol, buffer_signal, initialized
|
|
|
|
buf_len = params["buffer_length"]
|
|
avg_len = params.get("avg_length", 50)
|
|
|
|
buffer_dark = deque(maxlen=buf_len)
|
|
buffer_savgol = deque(maxlen=buf_len)
|
|
buffer_signal = deque(maxlen=avg_len)
|
|
|
|
initialized = True
|
|
_logger.info(
|
|
f"ATT01 pipeline initialised: buffer_length={buf_len}, avg_length={avg_len}"
|
|
)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Edge-finding (unchanged logic, extracted for reuse)
|
|
# ---------------------------------------------------------------------------
|
|
def find_edge(data, step_length=50, edge_type="falling", roi=None):
|
|
"""
|
|
Find an edge via cross-correlation with a step waveform.
|
|
|
|
If *roi* is ``[start, end]`` the search is restricted to that slice
|
|
and the returned edge position is offset back to full-frame coordinates.
|
|
"""
|
|
if roi is not None:
|
|
data_roi = data[roi[0]:roi[1]]
|
|
else:
|
|
data_roi = data
|
|
|
|
step_waveform = np.ones(step_length)
|
|
if edge_type == "rising":
|
|
step_waveform[: step_length // 2] = -1
|
|
elif edge_type == "falling":
|
|
step_waveform[step_length // 2 :] = -1
|
|
|
|
xcorr = signal.correlate(data_roi, step_waveform, mode="valid")
|
|
edge_position = np.argmax(xcorr) + np.floor(step_length / 2)
|
|
|
|
if roi is not None:
|
|
edge_position += roi[0]
|
|
|
|
return {
|
|
"edge_pos": edge_position,
|
|
"xcorr": xcorr,
|
|
"xcorr_ampl": float(np.amax(xcorr)),
|
|
"signal": data,
|
|
}
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Main entry point
|
|
# ---------------------------------------------------------------------------
|
|
def process(data, pulse_id, timestamp, params):
|
|
global initialized
|
|
|
|
device = params["device"]
|
|
step_length = params["step_length"]
|
|
edge_type = params["edge_type"]
|
|
dark_event = params["dark_event"]
|
|
fel_on_event = params["fel_on_event"]
|
|
calib = params["calib"]
|
|
filter_window = params["filter_window"]
|
|
prof_sig = data[params["prof_sig"]]
|
|
events = data[params["events"]]
|
|
|
|
if not initialized:
|
|
initialize(params)
|
|
|
|
# ------------------------------------------------------------------
|
|
# Pre-processing (unchanged)
|
|
# ------------------------------------------------------------------
|
|
prof_sig_savgol = signal.savgol_filter(prof_sig, filter_window, 3)
|
|
|
|
if events[dark_event]:
|
|
buffer_dark.append(prof_sig)
|
|
buffer_savgol.append(prof_sig_savgol)
|
|
|
|
if buffer_savgol:
|
|
prof_sig_norm = prof_sig_savgol / np.mean(np.array(buffer_savgol), axis=0)
|
|
else:
|
|
prof_sig_norm = prof_sig_savgol
|
|
|
|
# ------------------------------------------------------------------
|
|
# Single-shot edge finding
|
|
# ------------------------------------------------------------------
|
|
if events[fel_on_event] and not events[dark_event]:
|
|
edge_results = find_edge(
|
|
prof_sig_norm, step_length, edge_type, roi=params.get("roi")
|
|
)
|
|
edge_results["arrival_time"] = np.polyval(calib, edge_results["edge_pos"])
|
|
else:
|
|
edge_results = {
|
|
"edge_pos": None,
|
|
"xcorr": None,
|
|
"xcorr_ampl": None,
|
|
"signal": prof_sig_norm,
|
|
"arrival_time": None,
|
|
}
|
|
|
|
# ------------------------------------------------------------------
|
|
# Build single-shot output (unchanged keys)
|
|
# ------------------------------------------------------------------
|
|
output = {}
|
|
for key, value in edge_results.items():
|
|
output[f"{device}:{key}"] = value
|
|
|
|
output[f"{device}:raw_wf"] = prof_sig
|
|
output[f"{device}:raw_wf_savgol"] = prof_sig_savgol
|
|
|
|
if buffer_dark:
|
|
output[f"{device}:avg_dark_wf"] = np.mean(buffer_dark, axis=0)
|
|
else:
|
|
output[f"{device}:avg_dark_wf"] = None
|
|
|
|
# ------------------------------------------------------------------
|
|
# N-shot signal averaging + averaged edge finding
|
|
# ------------------------------------------------------------------
|
|
# Only accumulate valid FEL-on normalised signals
|
|
if events[fel_on_event] and not events[dark_event]:
|
|
buffer_signal.append(prof_sig_norm.copy())
|
|
|
|
if len(buffer_signal) > 0:
|
|
# Compute average over available shots (1..avg_length)
|
|
avg_signal = np.mean(np.array(buffer_signal), axis=0)
|
|
|
|
# Run edge finding on the averaged signal
|
|
avg_edge_results = find_edge(
|
|
avg_signal, step_length, edge_type, roi=params.get("roi")
|
|
)
|
|
avg_edge_results["arrival_time"] = np.polyval(
|
|
calib, avg_edge_results["edge_pos"]
|
|
)
|
|
|
|
output[f"{device}:avg_signal"] = avg_signal
|
|
output[f"{device}:avg_edge_pos"] = avg_edge_results["edge_pos"]
|
|
output[f"{device}:avg_xcorr"] = avg_edge_results["xcorr"]
|
|
output[f"{device}:avg_xcorr_ampl"] = avg_edge_results["xcorr_ampl"]
|
|
output[f"{device}:avg_arrival_time"] = avg_edge_results["arrival_time"]
|
|
output[f"{device}:avg_nshots"] = len(buffer_signal)
|
|
else:
|
|
# No valid shots yet — don't include avg keys at all.
|
|
# (Once the first FEL-on shot arrives, these keys appear and
|
|
# stay with consistent types/sizes for the rest of the run.)
|
|
pass
|
|
|
|
return output |