"""SRAS analysis: scope calibration, image reducers, and the SAW matched-filter pipeline. Qt-free; operates on the per-angle arrays returned by core.sras_format.SrasFile.load_angle(). Channel semantics (fixed by the acquisition app): CH1 — RF Acoustic Packet: FFT → peak frequency CH3 — Bias A (DC): waveform mean CH4 — Bias B (DC): waveform mean — also the RF valid-pixel mask source """ from __future__ import annotations import os import re from concurrent.futures import ThreadPoolExecutor from dataclasses import dataclass import numpy as np from scipy.signal import butter, hilbert, sosfiltfilt # Channel indices into the on-disk channel axis (order fixed by SCAN_CHANNELS) CH1_IDX, CH3_IDX, CH4_IDX = 0, 1, 2 # Fallback scope calibration for preambles missing YMULT/YOFF/YZERO: # 50 mV/div, 8 div full-scale, int8 ADC, position = -2.72 div FALLBACK_YMULT_MV = 1.5625 # mV per ADC count FALLBACK_YOFF_ADC = -87.04 # ADC count that represents 0 V def parse_preamble(preamble: str) -> dict[str, float]: """Extract YMULT, YOFF, YZERO from a Tektronix WFMOutpre string.""" result = {} for key in ("YMULT", "YOFF", "YZERO"): m = re.search(rf'\b{key}\s+([-+]?\d*\.?\d+(?:[Ee][+-]?\d+)?)', preamble) if m: result[key] = float(m.group(1)) return result @dataclass class ChannelCalibration: """Per-channel ADC↔mV conversion, parsed from the file's preambles.""" ymult_mv: list[float] yoff_adc: list[float] yzero_mv: list[float] @classmethod def from_preambles(cls, preambles: list[str]) -> "ChannelCalibration": cal = cls([], [], []) for p in preambles: vals = parse_preamble(p) # YMULT/YZERO from the scope are in V; stored here in mV cal.ymult_mv.append(vals.get("YMULT", FALLBACK_YMULT_MV / 1000) * 1000) cal.yoff_adc.append(vals.get("YOFF", FALLBACK_YOFF_ADC)) cal.yzero_mv.append(vals.get("YZERO", 0.0) * 1000) return cal def adc_to_mv(self, adc, ch: int): return (adc - self.yoff_adc[ch]) * self.ymult_mv[ch] + self.yzero_mv[ch] def mv_to_adc(self, mv, ch: int): return (mv - self.yzero_mv[ch]) / self.ymult_mv[ch] + self.yoff_adc[ch] def power_spectrum(waveform: np.ndarray) -> np.ndarray: """FFT power with the DC bin suppressed.""" power = np.abs(np.fft.rfft(waveform)) ** 2 power[..., 0] = 0.0 return power # --------------------------------------------------------------------------- # Image reducers — all take the (n_rows, n_ch, n_frames, spf) angle view # --------------------------------------------------------------------------- def compute_dc_image(angle_view: np.ndarray, ch_idx: int) -> np.ndarray: """Mean of each waveform → (n_rows, n_frames) float32. Computed directly on the int8 view — no float32 copy of the block. """ return angle_view[:, ch_idx].mean(axis=-1, dtype=np.float32) def _valid_ch1_waveforms(angle_view: np.ndarray, calib: ChannelCalibration, dc_threshold_mv: float, background: np.ndarray | None, ) -> tuple[np.ndarray, np.ndarray]: """CH4-DC mask + float32 CH1 waveforms for only the pixels that pass. Materializes float32 for the valid pixels alone (fancy-index on the int8 view first), so a mostly-masked angle costs almost nothing. """ dc4_mv = calib.adc_to_mv(compute_dc_image(angle_view, CH4_IDX), CH4_IDX) valid = dc4_mv >= dc_threshold_mv if not valid.any(): return valid, np.empty((0, angle_view.shape[-1]), dtype=np.float32) waves = angle_view[:, CH1_IDX][valid].astype(np.float32) # (n_valid, spf) if background is not None: waves -= background return valid, waves def compute_rf_image(angle_view: np.ndarray, calib: ChannelCalibration, freq_axis_mhz: np.ndarray, dc_threshold_mv: float, background: np.ndarray | None = None, gate_start_ns: float | None = None, gate_end_ns: float | None = None, time_axis_ns: np.ndarray | None = None) -> np.ndarray: """FFT of each CH1 waveform; pixel = peak frequency in MHz. Pixels whose CH4 DC mean (in mV) is below dc_threshold_mv are 0 and the FFT is skipped for them. Optional time gate zeroes samples outside [gate_start_ns, gate_end_ns] before the FFT. """ valid, waves = _valid_ch1_waveforms(angle_view, calib, dc_threshold_mv, background) img = np.zeros(valid.shape, dtype=np.float32) if len(waves): if (gate_start_ns is not None or gate_end_ns is not None) and time_axis_ns is not None: keep = np.ones(len(time_axis_ns), dtype=bool) if gate_start_ns is not None: keep &= time_axis_ns >= gate_start_ns if gate_end_ns is not None: keep &= time_axis_ns <= gate_end_ns waves[:, ~keep] = 0.0 peak_bins = np.argmax(power_spectrum(waves), axis=-1) img[valid] = freq_axis_mhz[peak_bins] return img def compute_saw_image(angle_view: np.ndarray, calib: ChannelCalibration, dc_threshold_mv: float, pipeline: "SawPipeline", mode: str, background: np.ndarray | None = None) -> np.ndarray: """Matched-filter pipeline over every valid pixel. mode : "amplitude" → MF envelope peak in the SAW window "tof" → arrival time (ns) of that peak Only the requested scalar is kept per pixel — the per-shot intermediate arrays are dropped inside the worker instead of being accumulated. """ valid, waves = _valid_ch1_waveforms(angle_view, calib, dc_threshold_mv, background) img = np.zeros(valid.shape, dtype=np.float32) if len(waves): key = "peak_amplitude" if mode == "amplitude" else "peak_time_ns" def _scalar(w): return pipeline.process_shot_metrics(w)[key] n_workers = min(os.cpu_count() or 4, len(waves)) with ThreadPoolExecutor(max_workers=n_workers) as executor: img[valid] = np.fromiter(executor.map(_scalar, waves), dtype=np.float32, count=len(waves)) return img # --------------------------------------------------------------------------- # SAW signal processing pipeline # --------------------------------------------------------------------------- class SawPipeline: """EMI-cleaning and SAW extraction pipeline. Stages (each independently bypassable): 1. EMI gate — cosine-taper the first `emi_gate_ns` ns to suppress the laser-firing burst at t≈0; leaves the SAW packet alone. 2. Bandpass — 6th-order Butterworth zero-phase (sosfiltfilt). 3. Matched filter — FFT cross-correlation with a Hann-windowed template built from the average of N clean shots. 4. Analytic — Hilbert transform of MF output → amplitude envelope. """ def __init__(self, sample_rate_hz: float, emi_gate_ns: float = 50.0, bp_lo_mhz: float = 85.0, bp_hi_mhz: float = 200.0, saw_window_ns: tuple[float, float] = (80.0, 350.0)): self.sample_rate_hz = float(sample_rate_hz) self.emi_gate_ns = float(emi_gate_ns) self.bp_lo_mhz = float(bp_lo_mhz) self.bp_hi_mhz = float(bp_hi_mhz) self.saw_window_ns = (float(saw_window_ns[0]), float(saw_window_ns[1])) self.template: np.ndarray | None = None self._template_fft: dict[int, np.ndarray] = {} # nfft → rfft(template) self._emi_gate_samples = max(1, int(round( self.emi_gate_ns * 1e-9 * self.sample_rate_hz))) nyq = self.sample_rate_hz / 2.0 lo = np.clip(self.bp_lo_mhz * 1e6 / nyq, 1e-6, 0.999) hi = np.clip(self.bp_hi_mhz * 1e6 / nyq, lo + 1e-6, 0.9999) # 6th-order Butterworth → 12th-order bandpass; ~120 dB/decade rolloff self._sos = butter(6, [lo, hi], btype='bandpass', output='sos') def gate_emi(self, signal: np.ndarray) -> np.ndarray: """Cosine-taper (raised cosine 0→1) the first emi_gate samples. The taper rolls up smoothly from zero so the abrupt EMI burst is suppressed without introducing a step discontinuity at the gate edge. """ n = min(self._emi_gate_samples, len(signal)) out = signal.copy() out[:n] *= 0.5 * (1.0 - np.cos(np.pi * np.arange(n) / n)) return out def bandpass(self, signal: np.ndarray) -> np.ndarray: """Zero-phase IIR Butterworth bandpass (sosfiltfilt), float32 in/out.""" return sosfiltfilt(self._sos, signal).astype(np.float32, copy=False) def build_template(self, waveforms: np.ndarray) -> None: """Average N shots (EMI-gated + bandpassed), Hann-windowed to the declared SAW window, to form the matched-filter template.""" processed = np.stack([ self.bandpass(self.gate_emi(np.asarray(w, dtype=np.float32))) for w in waveforms ]) avg = processed.mean(axis=0) n = len(avg) t_ns = np.arange(n) / self.sample_rate_hz * 1e9 i0 = max(0, int(np.searchsorted(t_ns, self.saw_window_ns[0]))) i1 = min(n, int(np.searchsorted(t_ns, self.saw_window_ns[1]))) windowed = np.zeros(n, dtype=np.float32) if i1 > i0: windowed[i0:i1] = avg[i0:i1] * np.hanning(i1 - i0) self.template = windowed self._template_fft.clear() def matched_filter(self, signal: np.ndarray) -> tuple[np.ndarray, np.ndarray]: """FFT cross-correlation with the template → (mf_output, envelope).""" if self.template is None: raise RuntimeError("No template — call build_template() first") n = len(signal) nfft = 1 << (n + len(self.template) - 1).bit_length() T = self._template_fft.get(nfft) if T is None: T = np.conj(np.fft.rfft(self.template, nfft)) self._template_fft[nfft] = T S = np.fft.rfft(signal, nfft) mf = np.fft.irfft(S * T, nfft)[:n] env = np.abs(hilbert(mf)) return mf.astype(np.float32, copy=False), env.astype(np.float32, copy=False) def process_shot(self, signal: np.ndarray) -> dict: """EMI gate → bandpass → matched filter on one shot; returns every stage plus metrics (for diagnostics displays).""" raw = np.asarray(signal, dtype=np.float32) gated = self.gate_emi(raw) filtered = self.bandpass(gated) if self.template is not None: mf_out, env = self.matched_filter(filtered) else: mf_out = filtered.copy() env = np.abs(hilbert(filtered)).astype(np.float32) metrics = self._envelope_metrics(env, filtered) return { "raw": raw, "gated": gated, "filtered": filtered, "mf_output": mf_out, "envelope": env, "sample_rate_hz": self.sample_rate_hz, **metrics, } def process_shot_metrics(self, signal: np.ndarray) -> dict: """Like process_shot but returns only the scalar metrics — used for whole-image sweeps where retaining per-shot arrays would multiply memory by the pixel count.""" filtered = self.bandpass(self.gate_emi(np.asarray(signal, dtype=np.float32))) if self.template is not None: _, env = self.matched_filter(filtered) else: env = np.abs(hilbert(filtered)) return self._envelope_metrics(env, filtered) def _envelope_metrics(self, env: np.ndarray, filtered: np.ndarray) -> dict: sr = self.sample_rate_hz t_ns = np.arange(len(env)) / sr * 1e9 s0, s1 = self.saw_window_ns roi = (t_ns >= s0) & (t_ns <= s1) if roi.any(): peak_sample = int(np.where(roi)[0][np.argmax(env[roi])]) else: peak_sample = int(np.argmax(env)) peak_amplitude = float(env[peak_sample]) peak_time_ns = float(peak_sample / sr * 1e9) # SNR: peak / RMS of the noise floor inside the gated EMI region noise_seg = filtered[:self._emi_gate_samples] noise_rms = float(np.sqrt(np.mean(noise_seg ** 2))) if len(noise_seg) else 1.0 return { "peak_amplitude": peak_amplitude, "peak_sample": peak_sample, "peak_time_ns": peak_time_ns, "snr": peak_amplitude / noise_rms if noise_rms > 0 else 0.0, }