aetherscan 1.0.0__py3-none-any.whl

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
@@ -0,0 +1,2718 @@
1
+ # TODO: add logging to background loading & downsampling
2
+ # TODO: update docstring once preprocessing.py complete
3
+ """
4
+ Data preprocessing for Aetherscan Pipeline
5
+ Handles data loading, downsampling, log-normalization for both training & inference
6
+ Uses multiprocessing and shared memory to process data in parallel
7
+ """
8
+
9
+ from __future__ import annotations
10
+
11
+ import contextlib
12
+ import csv
13
+ import functools
14
+ import gc
15
+ import glob
16
+ import json
17
+ import logging
18
+ import math
19
+ import os
20
+ import re
21
+ import signal
22
+ import time
23
+ import uuid
24
+ from collections import OrderedDict
25
+ from collections.abc import Callable
26
+ from dataclasses import dataclass, field
27
+ from multiprocessing.pool import Pool
28
+ from multiprocessing.shared_memory import SharedMemory
29
+
30
+ import h5py
31
+
32
+ # NOTE: come back to this later (why noqa: F401? what's the difference between h5py & hdf5plugin?)
33
+ import hdf5plugin # noqa: F401 # registers bitshuffle codec with h5py at import time
34
+ import numpy as np
35
+ from scipy import interpolate, stats
36
+ from skimage.transform import downscale_local_mean
37
+
38
+ from aetherscan.benchmark import stage_timer
39
+ from aetherscan.config import get_config
40
+ from aetherscan.data_generation import log_norm
41
+ from aetherscan.db import get_db
42
+ from aetherscan.logger import init_worker_logging
43
+ from aetherscan.manager import get_manager
44
+ from aetherscan.pfb import edge_mid_band_slices, equalize_passband, gen_coarse_channel_response
45
+ from aetherscan.run_state import preprocessing_config_fingerprint
46
+
47
+ logger = logging.getLogger(__name__)
48
+
49
+ # NOTE: find a way to avoid using global refs (store under manager.py maybe?)
50
+ # NOTE: is there any room to use asyncio & load all chunks simultaneously?
51
+ # Global variable to store chunk data for multiprocessing workers
52
+ # This avoids serialization overhead when passing data between workers
53
+ _GLOBAL_SHM = None
54
+ _GLOBAL_CHUNK_DATA = None
55
+ _GLOBAL_SHAPE = None
56
+ _GLOBAL_DTYPE = None
57
+
58
+ # PFB static-response sanity check: warn when the median flattened (divided-by-H) edge/mid
59
+ # power ratio deviates from 1.0 by more than this. ~0.10 is a provisional threshold — the
60
+ # residual still carries the analog-frontend tilt, whose legitimate baseline is only fixed by
61
+ # the deferred pfb_taps-vs-backend characterization — so the warning stays informational
62
+ # (#180 bounds the response-model error itself at ~1e-5, so H is not the confounder).
63
+ _PFB_RESIDUAL_FLATNESS_TOL = 0.10
64
+ # Fixed log-spaced bin edges for the energy-detection statistic summary histograms
65
+ # accumulated by _energy_detect_channel_worker. The edges are identical for every channel /
66
+ # file / cadence, so per-channel counts combine by plain addition; the range spans sub-noise
67
+ # k2 values (~1e-3) through extreme RFI statistics (~1e9) at 10 bins per decade. Consumed by
68
+ # inference_viz.plot_ed_stat_distributions via each cadence's metadata JSON.
69
+ ED_STAT_HIST_EDGES = np.logspace(-3.0, 9.0, 121)
70
+ # Coarse channels sampled (evenly across the band) by the opt-in bandpass overlay debug plot.
71
+ _BANDPASS_SAMPLE_CHANNELS = 3
72
+ # The PFB static-response sanity check samples more channels and takes a median (not a mean),
73
+ # so a single RFI-heavy channel doesn't skew the statistic.
74
+ _PFB_MISMATCH_SAMPLE_CHANNELS = 9
75
+ # Target bin count for _decimate_for_plot: spectrum lines in the bandpass debug/viz figures
76
+ # are reduced to at most 2 * this many points (a min/max pair per bin) before plotting.
77
+ _PLOT_MAX_POINTS_PER_LINE = 4096
78
+ # HDF5 chunk-cache sizing for read paths that revisit chunks (stamp extraction, the PFB
79
+ # residual check). h5py's default cache is 1 MiB — smaller than one decompressed bitshuffle
80
+ # chunk at BL scale (~4 MiB for (1, 1, 1048576) float32), and HDF5 never caches a chunk
81
+ # larger than the cache, so every access re-decompresses from scratch. _chunk_cache_kwargs
82
+ # sizes the cache from the file's actual chunk layout: 4 chunk-column time-stripes resident,
83
+ # capped here. rdcc_nslots is a prime ~100x the resident-chunk count (HDF5 guidance);
84
+ # rdcc_w0=0 evicts fully-read chunks first, matching the monotonic start-sorted read order.
85
+ _CHUNK_CACHE_STRIPES = 4
86
+ _CHUNK_CACHE_MAX_NBYTES = 512 * 1024**2
87
+ _CHUNK_CACHE_NSLOTS = 10007
88
+ # Stamp-extraction task granularity: target this many tasks per pool worker. At 1 task per
89
+ # worker (the old ceil(n_processes / n_files) split -> 36 near-equal tasks on 32 workers)
90
+ # the stage quantizes to two waves — a full wave plus a small remainder wave that idles
91
+ # ~44% of worker-seconds. ~4 tasks/worker keeps the tail to ~1/4 of one task length while
92
+ # each task still spans enough contiguous sorted stamps for chunk-cache reuse to pay.
93
+ _EXTRACT_TASKS_PER_WORKER = 4
94
+ # Age past which an unpromoted *.tmp.npy in the (cross-run, #298 I3) stamp cache directory
95
+ # is considered abandoned and removed. Live writes are minutes-scale; a day of margin means
96
+ # a concurrent run's in-progress tmp can never be mistaken for a dead one.
97
+ _STALE_TMP_MAX_AGE_S = 24 * 3600
98
+ # Cap on one coalesced extraction read (#301): overlapping/abutting stamp windows (the
99
+ # overlap_search triplets abut at offset stamp_width // 2) are fetched as ONE wide h5
100
+ # read and sliced, instead of re-reading the shared span per stamp. 2**20 bins bounds the
101
+ # wide float32 buffer at ~64 MiB per worker (16 time rows) even across an RFI comb whose
102
+ # windows chain for a whole coarse channel.
103
+ _COALESCE_MAX_BINS = 2**20
104
+ # Fixed per-cadence bin count for the pre-binned hit-frequency histograms stored in the
105
+ # metadata sidecar (#301): the raw per-hit frequency lists ran ~19 MB of JSON on an
106
+ # RFI-dense cadence (~117 GB over the full catalog) and forced the viz suite to hold every
107
+ # cadence's floats in RAM. 8192 bins over the file band keeps the rebinned hit-spectrum
108
+ # figure visually identical (fine-bin width ≪ the figure's global bin width) at ~20-50 KB.
109
+ _HIT_HIST_BINS = 8192
110
+ # Decimation budget for the bandpass envelopes persisted per cadence (#301): the viz
111
+ # figure renders ~2,100 px wide, so 1,024 min/max pairs per line lose nothing visually,
112
+ # and the stored envelopes spare the figure its ~114 s of cold /datag channel re-reads.
113
+ _ENVELOPE_MAX_POINTS = 1024
114
+
115
+
116
+ def _chunk_cache_kwargs(h5_path: str, time_bins: int) -> dict:
117
+ """
118
+ h5py.File open kwargs sizing the HDF5 chunk cache to the file's actual chunk layout.
119
+
120
+ Reads the 'data' dataset's chunk shape and returns {rdcc_nbytes, rdcc_nslots, rdcc_w0}
121
+ sized to hold _CHUNK_CACHE_STRIPES full chunk-column time-stripes (the set of chunks one
122
+ stamp read touches is ceil(time_bins / chunk_time) chunks per column). Returns {} — the
123
+ h5py defaults — for contiguous datasets, layouts the default cache already covers, or any
124
+ inspection failure, so callers degrade gracefully on non-BL chunk layouts.
125
+ """
126
+ try:
127
+ with h5py.File(h5_path, "r") as hf:
128
+ dset = hf["data"]
129
+ chunks = dset.chunks
130
+ if not chunks:
131
+ return {}
132
+ chunk_nbytes = int(np.prod(chunks)) * dset.dtype.itemsize
133
+ except Exception as e:
134
+ logger.warning(f"Could not inspect chunk layout of {h5_path} ({e}); using default cache")
135
+ return {}
136
+
137
+ stripe_nbytes = math.ceil(time_bins / chunks[0]) * chunk_nbytes
138
+ rdcc_nbytes = min(_CHUNK_CACHE_STRIPES * stripe_nbytes, _CHUNK_CACHE_MAX_NBYTES)
139
+ if rdcc_nbytes <= 1024**2:
140
+ return {}
141
+ return {"rdcc_nbytes": rdcc_nbytes, "rdcc_nslots": _CHUNK_CACHE_NSLOTS, "rdcc_w0": 0.0}
142
+
143
+
144
+ def _init_worker(shm_name, shape, dtype):
145
+ """
146
+ Worker pool initializer: attach to the named shared-memory block and set up logging.
147
+
148
+ Passing the shm_name/shape/dtype through the pool initializer (rather than per-task args)
149
+ avoids re-serializing the whole array on every map() call. The worker installs a SIGTERM
150
+ handler that closes its shared-memory file descriptor before letting the signal kill the
151
+ process; the main process is responsible for unlinking the shared memory afterwards (handled
152
+ by ResourceManager).
153
+ """
154
+ global _GLOBAL_SHM, _GLOBAL_CHUNK_DATA, _GLOBAL_SHAPE, _GLOBAL_DTYPE
155
+
156
+ # Initialize worker logging
157
+ init_worker_logging()
158
+
159
+ # Attach to existing shared memory block
160
+ _GLOBAL_SHM = SharedMemory(name=shm_name)
161
+
162
+ # Ignore SIGINT (Ctrl+C) in workers - let manager from parent handle cleanup coordination
163
+ signal.signal(signal.SIGINT, signal.SIG_IGN)
164
+
165
+ # Setup custom SIGTERM handler for additional cleanup before termination
166
+ # Note, manager will escalate SIGTERM to SIGKILL after pool_terminate_timeout seconds (see config.py)
167
+ # This may interrupt the worker's cleanup process
168
+ # Consider increasing pool_terminate_timeout if you're experiencing such issues
169
+ def cleanup_on_sigterm(signum, frame):
170
+ """SIGTERM handler that closes the worker's shared-memory fd before re-raising the signal."""
171
+ # Note, a race condition may occur if a worker receives more than 1 SIGTERM delivery
172
+ # at a time, triggering re-entry of the same cleanup handler
173
+ # It suffices to guard against this by simply suppressing exceptions, since subsequent
174
+ # close() calls will just raise an error; no state corruption or kernel-level hazards exist
175
+ # Also, there are no cross-worker race conditions, since each worker's close() operates on
176
+ # per-process resources, even though they all refer to the same underlying POSIX shm object
177
+ with contextlib.suppress(Exception):
178
+ if _GLOBAL_SHM is not None:
179
+ # WARN: DO NOT LOG ANYTHING ON CLEANUP!
180
+ # Any calls to logger will attempt to put a message onto QueueHandler, whose feeder
181
+ # thread needs the GIL to transfer data to the underlying pipe. However, the main
182
+ # process may be holding the GIL (e.g. TF's prefetch threads during training)
183
+ # This causes deadlocks, preventing workers from completing their SIGTERM handlers
184
+ # logger.info(f"Closing shared memory file descriptor in worker PID {os.getpid()}")
185
+ _GLOBAL_SHM.close()
186
+
187
+ # Restore default handler and re-raise SIGTERM to resume termination
188
+ signal.signal(signal.SIGTERM, signal.SIG_DFL)
189
+ os.kill(os.getpid(), signal.SIGTERM)
190
+
191
+ # Register SIGTERM handler for graceful cleanup on pool.terminate()
192
+ signal.signal(signal.SIGTERM, cleanup_on_sigterm)
193
+
194
+ # Create numpy array view of shared memory (no copy!)
195
+ _GLOBAL_CHUNK_DATA = np.ndarray(shape, dtype=dtype, buffer=_GLOBAL_SHM.buf)
196
+ _GLOBAL_SHAPE = shape
197
+ _GLOBAL_DTYPE = dtype
198
+
199
+
200
+ def _init_plain_worker():
201
+ """
202
+ Worker pool initializer for pools that don't attach shared memory (energy detection and
203
+ stamp extraction): set up queue-based logging and ignore SIGINT so the parent's
204
+ ResourceManager coordinates shutdown. No SIGTERM handler is installed — these workers hold
205
+ no shared-memory file descriptors, so the default termination behavior is safe.
206
+ """
207
+ init_worker_logging()
208
+ signal.signal(signal.SIGINT, signal.SIG_IGN)
209
+ # Marks this process as a pool worker: gates the per-worker h5 handle cache, which must
210
+ # never engage on the parent's (multi-threaded) sequential path — see _cached_h5_file
211
+ global _IN_POOL_WORKER
212
+ _IN_POOL_WORKER = True
213
+
214
+
215
+ def _downsample_cadence(cadence, downsample_factor: int, final_width: int):
216
+ """
217
+ Downsample one cadence's 6 observations, or return None for an invalid cadence
218
+ (NaN/Inf or non-positive max). Pure function of its arguments — the thread-safe core
219
+ shared by the pool worker (which resolves its cadence from _GLOBAL_CHUNK_DATA) and the
220
+ sequential path (which passes the row directly, so it never touches the module global
221
+ and stays safe on the multi-threaded prefetch side — #298 review note).
222
+ """
223
+ # Skip invalid cadences
224
+ if np.any(np.isnan(cadence)) or np.any(np.isinf(cadence)) or np.max(cadence) <= 0:
225
+ return None
226
+
227
+ # Downsample each observation separately
228
+ downsampled_cadence = np.zeros((6, 16, final_width), dtype=np.float32)
229
+
230
+ for obs_idx in range(6):
231
+ downsampled_cadence[obs_idx] = downscale_local_mean(
232
+ cadence[obs_idx], (1, downsample_factor)
233
+ ).astype(np.float32)
234
+
235
+ return downsampled_cadence
236
+
237
+
238
+ # NOTE: come back to this later
239
+ def _downsample_worker(args):
240
+ """
241
+ Downsample one cadence's 6 observations in parallel.
242
+
243
+ args is a (cadence_idx, downsample_factor, final_width) tuple — the cadence itself is
244
+ pulled from _GLOBAL_CHUNK_DATA to avoid pickling it across the pool boundary. Returns the
245
+ downsampled cadence of shape (6, 16, final_width), or None if the source cadence contained
246
+ NaN/Inf or had non-positive max (treated as invalid).
247
+ """
248
+ cadence_idx, downsample_factor, final_width = args
249
+
250
+ if _GLOBAL_CHUNK_DATA is None:
251
+ logger.warning("No global chunk data available")
252
+ return None
253
+
254
+ return _downsample_cadence(_GLOBAL_CHUNK_DATA[cadence_idx], downsample_factor, final_width)
255
+
256
+
257
+ def _lognorm_worker(args):
258
+ """
259
+ Log-normalize one already-downsampled cadence in parallel.
260
+
261
+ Counterpart to _downsample_worker for stamps that were downsampled at extraction time
262
+ (see _extract_stamps_worker): the stored .npy is already at final width, so loading only
263
+ needs the per-cadence log-norm. args is a (cadence_idx,) tuple — the cadence itself is
264
+ pulled from _GLOBAL_CHUNK_DATA to avoid pickling it across the pool boundary. Returns the
265
+ log-normalized cadence of shape (6, time_bins, final_width) as float32, or None if the
266
+ source cadence contained NaN/Inf or had non-positive max (treated as invalid, matching
267
+ _downsample_worker).
268
+ """
269
+ (cadence_idx,) = args
270
+
271
+ if _GLOBAL_CHUNK_DATA is None:
272
+ logger.warning("No global chunk data available")
273
+ return None
274
+
275
+ cadence = _GLOBAL_CHUNK_DATA[cadence_idx]
276
+
277
+ # Skip invalid cadences (same validity rule as _downsample_worker)
278
+ if np.any(np.isnan(cadence)) or np.any(np.isinf(cadence)) or np.max(cadence) <= 0:
279
+ return None
280
+
281
+ return log_norm(cadence.astype(np.float32))
282
+
283
+
284
+ def _log_norm_chunk_vectorized(chunk: np.ndarray) -> tuple[np.ndarray, np.ndarray]:
285
+ """
286
+ Vectorized equivalent of mapping _lognorm_worker over every cadence of `chunk`
287
+ (n, obs, time_bins, width): returns (normalized_valid_rows, valid_mask) (#298 I5).
288
+
289
+ Bit-identical per element to the per-cadence path, by construction: the validity rule is
290
+ the same (any NaN/Inf or non-positive max rejects the cadence — NaN/Inf rows fail the
291
+ finite mask, so the max comparison never decides them); the arithmetic runs in float32
292
+ with the same scalar casts (numpy converts the 1e-10 epsilon and the per-cadence min/max
293
+ scalars to float32 in both forms); min/max are exact order-independent reductions; and
294
+ rows whose post-shift range is 0 skip the division exactly as log_norm's range_log > 0
295
+ guard does. Pinned by a unit test against _lognorm_worker on a mixed valid/invalid chunk.
296
+ """
297
+ reduce_axes = (1, 2, 3)
298
+ finite = np.isfinite(chunk).all(axis=reduce_axes)
299
+ maxes = chunk.max(axis=reduce_axes)
300
+ with np.errstate(invalid="ignore"):
301
+ valid = finite & (maxes > 0)
302
+
303
+ # Advanced indexing always yields a fresh array we own, so copy=False only avoids a
304
+ # second copy when the dtype is already float32 — in-place mutation below is safe.
305
+ data = chunk[valid].astype(np.float32, copy=False)
306
+ data += 1e-10
307
+ np.log(data, out=data)
308
+ data -= data.min(axis=reduce_axes, keepdims=True)
309
+ ranges = data.max(axis=reduce_axes, keepdims=True)
310
+ np.divide(data, ranges, out=data, where=ranges > 0)
311
+ return data, valid
312
+
313
+
314
+ # NOTE: come back to this later (mirrors preprocess_fine.py:72-75 from the reference implementation)
315
+ def _remove_dc_spike(
316
+ block_data: np.ndarray, coarse_channel_width: int, n_coarse_channels: int
317
+ ) -> None:
318
+ """
319
+ Interpolate over the 2-bin DC spike at the center of each coarse channel, in place.
320
+
321
+ block_data has shape (time_bins, n_coarse_channels * coarse_channel_width); the channel
322
+ layout is contiguous so the spike sits at i * coarse_channel_width + coarse_channel_width//2
323
+ for each i.
324
+ """
325
+ half_chan = coarse_channel_width // 2
326
+ for i in range(n_coarse_channels):
327
+ dc_ind = i * coarse_channel_width + half_chan
328
+ # The two replacements are asymmetric on purpose: dc_ind pulls from
329
+ # (+1, -3) and dc_ind-1 from (+2, -2). This matches the reference
330
+ # implementation byte-for-byte (FX196/SETI-Energy-Detection
331
+ # preprocess_fine.py:72-75) — the bins immediately around the spike
332
+ # are themselves contaminated, so the offsets reach across the spike
333
+ # to clean neighbors on both sides.
334
+ block_data[:, dc_ind] = (block_data[:, dc_ind + 1] + block_data[:, dc_ind - 3]) / 2
335
+ block_data[:, dc_ind - 1] = (block_data[:, dc_ind + 2] + block_data[:, dc_ind - 2]) / 2
336
+
337
+
338
+ # NOTE: come back to this later (mirrors utils.py:17-22 from the reference implementation.)
339
+ def _fit_channel_bandpass(
340
+ integrated_channel: np.ndarray, channel_width: int, spl_order: int
341
+ ) -> np.ndarray:
342
+ """
343
+ Fit a spline bandpass to a time-integrated coarse channel and evaluate it at every bin.
344
+
345
+ integrated_channel is a 1-D array of shape (channel_width,). spl_order controls the number
346
+ of interior knots (higher = finer fit). Returns the per-bin bandpass fit of the same shape.
347
+ """
348
+ x = np.arange(channel_width)
349
+ # Interior knots must lie strictly inside (x[0], x[-1]); knots[1:] drops the
350
+ # leading 0, which satisfies that constraint for the default config
351
+ # (channel_width=1048576, spl_order=16). A pathological config with very
352
+ # small channel_width relative to spl_order could produce a knots array
353
+ # that violates the constraint — splrep would raise then. The defaults
354
+ # used in InferenceConfig are safe.
355
+ knots = np.arange(0, channel_width, channel_width // spl_order + 1)
356
+ spl = interpolate.splrep(x, integrated_channel, t=knots[1:])
357
+ return interpolate.splev(x, spl)
358
+
359
+
360
+ def _spline_flatten_bandpass(channel: np.ndarray, spl_order: int) -> np.ndarray:
361
+ """
362
+ Spline bandpass flattener (--bandpass-method spline): fit a spline to the time-integrated
363
+ coarse channel and subtract it from every time row. channel has shape
364
+ (time_bins, coarse_channel_width); returns the float64 residuals of the same shape.
365
+
366
+ Bandpass flattening is a pluggable stage: _process_cadence obtains a picklable callable via
367
+ DataPreprocessor._get_bandpass_flattener() and ships it to the pool workers, so a flattener
368
+ only needs to be a module-level function with the same (channel) -> flattened contract
369
+ (see _pfb_flatten_bandpass for the other implementation).
370
+ """
371
+ integrated_channel = np.mean(channel, axis=0)
372
+ fit = _fit_channel_bandpass(integrated_channel, channel.shape[1], spl_order)
373
+ return channel - fit
374
+
375
+
376
+ # maxsize=4 bounds memory if a long-lived worker sees several response paths (unlikely in
377
+ # production — one per run — but plausible across a long test session).
378
+ @functools.lru_cache(maxsize=4)
379
+ def _load_pfb_response(response_path: str) -> np.ndarray:
380
+ """
381
+ Load (and per-process cache) a PFB passband response written by
382
+ DataPreprocessor._ensure_pfb_response_file. The array is marked read-only because the
383
+ cache shares it across every call in the process.
384
+ """
385
+ response = np.load(response_path)
386
+ response.setflags(write=False)
387
+ return response
388
+
389
+
390
+ def _flattened_edge_mid_ratio_banded(
391
+ h5_path: str,
392
+ channel_index: int,
393
+ response: np.ndarray,
394
+ width: int,
395
+ time_bins: int,
396
+ open_kwargs: dict | None = None,
397
+ ) -> float:
398
+ """
399
+ Edge/mid power ratio of one coarse channel AFTER flattening by the static response H.
400
+
401
+ Reads only the bands edge_mid_power_ratio actually evaluates (edge_mid_band_slices:
402
+ the outermost width // 16 bins each side plus a central band of the same width) as
403
+ native float32 and time-integrates each immediately — ~10x less I/O than the
404
+ full-channel float64 read of _read_despiked_channel, and no fp64 blowup, which is what
405
+ keeps the per-file sanity check cheap at GBT scale. The DC spike (2 bins at
406
+ width // 2) sits inside the mid band; its interpolation (_remove_dc_spike) is linear,
407
+ so it commutes with the time mean and is replicated on the integrated slice. Matches
408
+ edge_mid_power_ratio(equalize_passband(channel, response).mean(axis=0)) on the
409
+ despiked full read, up to float32-read rounding (pinned by a unit test).
410
+
411
+ Module-level (not a DataPreprocessor method) so the pool-side _pfb_ratio_worker can
412
+ share the exact same body (#298 I6) — both venues must produce identical ratios.
413
+ """
414
+ start = channel_index * width
415
+ left, mid, right = edge_mid_band_slices(width)
416
+ # Widen the mid slice so the DC-spike interpolation sources (up to 3 bins around
417
+ # width // 2, see _remove_dc_spike) are always present. The min(width) bound is
418
+ # defensive documentation: at any real width the center+3 widening cannot reach the
419
+ # channel end (mid.stop < width for width >= 16), so the clamp never bites in practice.
420
+ lo = max(0, min(mid.start, width // 2 - 3))
421
+ hi = min(width, max(mid.stop, width // 2 + 3))
422
+ # The three band reads all live in the same chunk-column stripe; a sized chunk cache
423
+ # (computed once per file by the caller) decompresses each chunk once instead of once
424
+ # per band read.
425
+ with h5py.File(h5_path, "r", **(open_kwargs or {})) as hf:
426
+ data = hf["data"]
427
+ left_int = data[:time_bins, 0, start + left.start : start + left.stop].mean(axis=0)
428
+ right_int = data[:time_bins, 0, start + right.start : start + right.stop].mean(axis=0)
429
+ mid_int = data[:time_bins, 0, start + lo : start + hi].mean(axis=0)
430
+ dc = width // 2 - lo
431
+ mid_int[dc] = (mid_int[dc + 1] + mid_int[dc - 3]) / 2
432
+ mid_int[dc - 1] = (mid_int[dc + 2] + mid_int[dc - 2]) / 2
433
+ mid_int = mid_int[mid.start - lo : mid.stop - lo]
434
+
435
+ # Same divide-by-H (with equalize_passband's defensive floor) as the real flattener;
436
+ # per-bin division commutes with the time mean, so integrating first is equivalent.
437
+ floor = np.maximum(response, 1e-10)
438
+ edge = 0.5 * (float((left_int / floor[left]).mean()) + float((right_int / floor[right]).mean()))
439
+ return edge / float((mid_int / floor[mid]).mean())
440
+
441
+
442
+ def _pfb_ratio_worker(args: tuple) -> float:
443
+ """
444
+ Pool worker for the PFB residual-flatness check (#298 I6): one sampled channel's
445
+ flattened edge/mid power ratio. The static response arrives as its content-addressed
446
+ sidecar PATH — _load_pfb_response caches the ~8 MB array per worker process, and the
447
+ sidecar was np.array_equal-verified against the generated response in the parent, so
448
+ pool-side ratios are bit-identical to the old parent-side computation.
449
+ """
450
+ h5_path, channel_index, width, time_bins, response_path, open_kwargs = args
451
+ return _flattened_edge_mid_ratio_banded(
452
+ h5_path, channel_index, _load_pfb_response(response_path), width, time_bins, open_kwargs
453
+ )
454
+
455
+
456
+ def _pfb_flatten_bandpass(channel: np.ndarray, response_path: str) -> np.ndarray:
457
+ """
458
+ PFB static-equalization bandpass flattener (--bandpass-method pfb, the default): divide the
459
+ channel by the instrument's precomputed polyphase-filterbank passband response instead of
460
+ fitting a spline per channel per file. channel has shape (time_bins, coarse_channel_width);
461
+ returns the equalized float64 channel of the same shape.
462
+
463
+ The response is computed ONCE, in the parent process, and shipped to pool workers as the
464
+ path of a small sidecar .npy (see _ensure_pfb_response_file) rather than as generation
465
+ parameters: at GBT scale, generating the response is an ~n_chans-point FFT with a
466
+ tens-of-GB transient, and every pool worker running that concurrently on its first task
467
+ would exhaust the node's memory. Workers just read the ~8 MB file, cached per process by
468
+ _load_pfb_response.
469
+
470
+ # NOTE: the spline path *subtracts* its fit while this path *divides* by H, so equalized
471
+ # channels keep their DC offset and a bin-dependent scale. That asymmetry is fine for
472
+ # detection: the D'Agostino-Pearson normality statistic is built from skewness and
473
+ # kurtosis, which are central moments (location cancels) normalized by powers of the
474
+ # variance (scale cancels), so flattening the bandpass *shape* is all that matters.
475
+ # The spline-vs-PFB detection comparison in the PR body is the empirical arbiter.
476
+ """
477
+ return equalize_passband(channel, _load_pfb_response(response_path))
478
+
479
+
480
+ def _decimate_for_plot(
481
+ y: np.ndarray, max_points: int = _PLOT_MAX_POINTS_PER_LINE
482
+ ) -> tuple[np.ndarray, np.ndarray]:
483
+ """
484
+ Reduce a long 1-D line to at most 2 * max_points (x, y) points for plotting, via a per-bin
485
+ min/max envelope: each of max_points equal bins contributes its extremes at their true
486
+ indices, so narrowband features (RFI spikes, hits) survive where a plain stride would
487
+ usually miss them. At GBT scale this turns each ~1M-point matplotlib line in the bandpass
488
+ figures into ~8k points. Returns (x_indices, values); inputs already at or below the
489
+ budget pass through unchanged. Unpack into ax.plot(*_decimate_for_plot(y), ...).
490
+ """
491
+ y = np.asarray(y)
492
+ n = y.shape[0]
493
+ if n <= 2 * max_points:
494
+ return np.arange(n), y
495
+ bin_size = -(-n // max_points) # ceil division
496
+ # Edge-pad the tail bin; padded positions duplicate y[-1], and any argmin/argmax landing
497
+ # there is clipped back to n - 1 below, so no synthetic value can be selected.
498
+ padded = np.pad(y, (0, bin_size * max_points - n), mode="edge")
499
+ binned = padded.reshape(max_points, bin_size)
500
+ offsets = np.arange(max_points) * bin_size
501
+ pairs = np.stack([offsets + binned.argmin(axis=1), offsets + binned.argmax(axis=1)], axis=1)
502
+ # np.unique dedupes + keeps ascending order: a constant-valued bin has argmin==argmax, which
503
+ # would otherwise emit the same (idx, y) point twice; deduping halves those bins for free.
504
+ idx = np.unique(np.minimum(np.sort(pairs, axis=1).reshape(-1), n - 1))
505
+ return idx, y[idx]
506
+
507
+
508
+ def _sliding_normality_k2(channel: np.ndarray, window_size: int, step_size: int) -> np.ndarray:
509
+ """
510
+ D'Agostino-Pearson normality statistic (k2) for every sliding window across one coarse
511
+ channel, computed in closed form from per-block power sums instead of a per-window Python
512
+ loop over scipy.stats.normaltest (which was the #1 cost of inference preprocessing).
513
+
514
+ channel has shape (time_bins, width). Window j covers columns
515
+ [j*step_size, j*step_size + window_size) — the same windows as the historical loop
516
+ `range(0, width - window_size, step_size)` — and each window's sample is the flattened
517
+ (time_bins, window_size) slice, so n = time_bins * window_size is constant across windows.
518
+ Returns a float64 array of shape (n_windows,) whose entries match
519
+ `scipy.stats.normaltest(window.flatten()).statistic` (unit tests pin the equivalence to
520
+ rtol=1e-9); windows with zero variance yield NaN, matching scipy's behavior of returning
521
+ NaN rather than a spurious statistic.
522
+
523
+ Derivation: normaltest = skewtest.Z**2 + kurtosistest.Z**2 where both Z transforms are
524
+ elementwise closed forms in (n, m2, m3, m4). The central moments come from per-block raw
525
+ power sums S1..S4 accumulated in float64: blocks are sized so that every window is an
526
+ exact run of adjacent blocks (block = step_size when step divides window — the fast
527
+ path — else gcd(window, step)), window sums are length-(window//block) moving sums over
528
+ the block sums, and m2/m3/m4 follow from the standard raw-to-central-moment identities.
529
+ Raw-moment differencing loses precision when |mean| >> std, so the channel mean is
530
+ subtracted up front (central moments are shift-invariant); bandpass-subtracted residuals
531
+ are near-zero-mean anyway, and the rtol=1e-9 unit-test gate is the arbiter.
532
+ """
533
+ time_bins, width = channel.shape
534
+ n = time_bins * window_size # samples per window (scalar across all windows)
535
+ if n < 8:
536
+ # Mirror scipy.stats.skewtest's minimum-sample requirement
537
+ raise ValueError(f"normality test requires >= 8 samples per window, got {n}")
538
+ n_windows = len(range(0, width - window_size, step_size))
539
+ if n_windows <= 0:
540
+ return np.empty(0, dtype=np.float64)
541
+
542
+ # astype always copies, so the in-place shift below can't mutate the caller's array.
543
+ # Shift-invariance guard: subtracting the channel mean leaves every central moment
544
+ # mathematically unchanged while keeping the S2/n - mean**2 style differencing below
545
+ # well conditioned even when the residuals carry a DC offset.
546
+ data = channel.astype(np.float64)
547
+ data -= data.mean()
548
+
549
+ # Per-column power sums over the time axis, accumulated row by row so temporaries stay at
550
+ # (width,) rather than (time_bins, width) x 3.
551
+ # NOTE: the Python-level loop over time_bins rows is deliberate — it caps peak memory at a
552
+ # few (width,) float64 vectors per worker. With the default time_bins=16 the loop overhead
553
+ # is negligible; if time_bins ever grows large, vectorize via (channel**k).sum(axis=0).
554
+ p1 = np.zeros(width, dtype=np.float64)
555
+ p2 = np.zeros(width, dtype=np.float64)
556
+ p3 = np.zeros(width, dtype=np.float64)
557
+ p4 = np.zeros(width, dtype=np.float64)
558
+ for t in range(time_bins):
559
+ row = data[t]
560
+ row2 = row * row
561
+ p1 += row
562
+ p2 += row2
563
+ p3 += row2 * row
564
+ p4 += row2 * row2
565
+
566
+ # Aggregate columns into block sums. The fast path (step divides window — the default
567
+ # config: window 256, step 128) uses step-sized blocks; the general path falls back to
568
+ # gcd-sized blocks so windows are still exact runs of adjacent blocks.
569
+ block = step_size if window_size % step_size == 0 else math.gcd(window_size, step_size)
570
+ edges = np.arange(0, width, block)
571
+ blocks_per_window = window_size // block
572
+ blocks_per_step = step_size // block
573
+ # When block does not divide width (non-power-of-2 geometries), np.add.reduceat emits a
574
+ # short trailing block spanning the ragged tail [edges[-1], width). No in-range window ever
575
+ # reaches it — every window ends at a multiple of block that is < width, hence <= edges[-1]
576
+ # — but drop it explicitly so a partial sum can never leak into a window and silently
577
+ # corrupt k2. width // block == number of full blocks (== len(edges) when block divides
578
+ # width, so this is a no-op there); slicing to it keeps every window sum exact.
579
+ n_full_blocks = width // block
580
+
581
+ def _window_sums(col_sums: np.ndarray) -> np.ndarray:
582
+ block_sums = np.add.reduceat(col_sums, edges)[:n_full_blocks]
583
+ # Moving sum of blocks_per_window adjacent blocks, sampled every blocks_per_step
584
+ # blocks — one entry per window, no long cumulative accumulation (precision).
585
+ view = np.lib.stride_tricks.sliding_window_view(block_sums, blocks_per_window)
586
+ return view[::blocks_per_step].sum(axis=-1)[:n_windows]
587
+
588
+ s1 = _window_sums(p1)
589
+ s2 = _window_sums(p2)
590
+ s3 = _window_sums(p3)
591
+ s4 = _window_sums(p4)
592
+
593
+ # Raw power sums -> central moments (float64 throughout)
594
+ mean = s1 / n
595
+ m2 = s2 / n - mean**2
596
+ m3 = s3 / n - 3.0 * mean * (s2 / n) + 2.0 * mean**3
597
+ m4 = s4 / n - 4.0 * mean * (s3 / n) + 6.0 * mean**2 * (s2 / n) - 3.0 * mean**4
598
+
599
+ # Z transforms transcribed from scipy.stats._stats_py::skewtest / kurtosistest with
600
+ # scalar n, so every n-dependent constant is computed once. Variable names follow scipy.
601
+ with np.errstate(all="ignore"):
602
+ # Zero-variance windows can't support either test; scipy returns NaN there (and a
603
+ # tiny negative m2 from float cancellation would otherwise fabricate a huge k2).
604
+ degenerate = m2 <= 0.0
605
+ m2 = np.where(degenerate, np.nan, m2)
606
+
607
+ # --- skewtest: Z1 from g1 = m3 / m2**1.5
608
+ b2 = m3 / m2**1.5
609
+ y = b2 * math.sqrt(((n + 1) * (n + 3)) / (6.0 * (n - 2)))
610
+ beta2 = (
611
+ 3.0
612
+ * (n**2 + 27 * n - 70)
613
+ * (n + 1)
614
+ * (n + 3)
615
+ / ((n - 2.0) * (n + 5) * (n + 7) * (n + 9))
616
+ )
617
+ w2 = -1 + math.sqrt(2 * (beta2 - 1))
618
+ delta = 1 / math.sqrt(0.5 * math.log(w2))
619
+ alpha = math.sqrt(2.0 / (w2 - 1))
620
+ y = np.where(y == 0, 1.0, y)
621
+ z1 = delta * np.log(y / alpha + np.sqrt((y / alpha) ** 2 + 1))
622
+
623
+ # --- kurtosistest: Z2 from b2 = m4 / m2**2
624
+ b2k = m4 / m2**2
625
+ e = 3.0 * (n - 1) / (n + 1)
626
+ varb2 = 24.0 * n * (n - 2) * (n - 3) / ((n + 1) * (n + 1.0) * (n + 3) * (n + 5))
627
+ x = (b2k - e) / math.sqrt(varb2)
628
+ sqrtbeta1 = (
629
+ 6.0
630
+ * (n * n - 5 * n + 2)
631
+ / ((n + 7) * (n + 9))
632
+ * math.sqrt((6.0 * (n + 3) * (n + 5)) / (n * (n - 2) * (n - 3)))
633
+ )
634
+ a = 6.0 + 8.0 / sqrtbeta1 * (2.0 / sqrtbeta1 + math.sqrt(1 + 4.0 / (sqrtbeta1**2)))
635
+ term1 = 1 - 2 / (9.0 * a)
636
+ denom = 1 + x * math.sqrt(2 / (a - 4.0))
637
+ term2 = np.sign(denom) * np.where(
638
+ denom == 0.0, np.nan, np.power((1 - 2.0 / a) / np.abs(denom), 1 / 3.0)
639
+ )
640
+ z2 = (term1 - term2) / math.sqrt(2 / (9.0 * a))
641
+
642
+ k2 = z1 * z1 + z2 * z2
643
+
644
+ return k2
645
+
646
+
647
+ # Per-worker cache of read-only h5py handles for the fused ED tasks (#298 rider): each
648
+ # worker used to open/close the same 3 ON files once per coarse channel — ~2,859
649
+ # open/close cycles per cadence against a network filesystem. Small LRU (the current
650
+ # cadence's ON files plus slack); read-only input files, so a held handle can never go
651
+ # stale, and handles die with the worker process. ED reads visit each chunk exactly once,
652
+ # so the h5py default chunk cache on these handles is irrelevant (extraction, which DOES
653
+ # revisit chunks, opens its own rdcc-sized handles — see _extract_stamps_worker).
654
+ # POOL WORKERS ONLY (_IN_POOL_WORKER, set by _init_plain_worker): in sequential mode the
655
+ # ED chain runs on the parent's prefetch thread(s), and h5py File objects must not be
656
+ # shared across threads — the sequential path keeps its per-call open/close.
657
+ _H5_HANDLE_CACHE: OrderedDict = OrderedDict()
658
+ _H5_HANDLE_CACHE_MAX = 4
659
+ _IN_POOL_WORKER = False
660
+
661
+
662
+ def _cached_h5_file(h5_path: str):
663
+ handle = _H5_HANDLE_CACHE.get(h5_path)
664
+ if handle is not None and handle.id.valid:
665
+ _H5_HANDLE_CACHE.move_to_end(h5_path)
666
+ return handle
667
+ handle = h5py.File(h5_path, "r")
668
+ _H5_HANDLE_CACHE[h5_path] = handle
669
+ while len(_H5_HANDLE_CACHE) > _H5_HANDLE_CACHE_MAX:
670
+ _, evicted = _H5_HANDLE_CACHE.popitem(last=False)
671
+ with contextlib.suppress(Exception):
672
+ evicted.close()
673
+ return handle
674
+
675
+
676
+ def _energy_detect_channel_worker(args: tuple) -> tuple[list[tuple], np.ndarray, np.ndarray | None]:
677
+ """
678
+ Fused worker: run the complete energy-detection chain for one coarse channel — read the
679
+ channel's h5 slice, remove the DC spike, flatten the bandpass, and threshold the vectorized
680
+ normality statistic. Returns (hits, stat_hist, integrated): hits is a small list of
681
+ (absolute_fine_channel_index, statistic, pvalue) tuples; stat_hist is the histogram of
682
+ *all* finite window statistics (not just hits) over the fixed ED_STAT_HIST_EDGES bins, so
683
+ the parent can cheaply accumulate the full per-file statistic distribution for the
684
+ visualization suite; integrated is the despiked channel's time-integrated spectrum
685
+ (float64, only when want_spectrum — the parent turns the few sampled channels' spectra
686
+ into the persisted bandpass envelopes, #301, sparing the viz suite its cold h5 re-reads).
687
+ The bulky (time_bins, coarse_channel_width) intermediates never leave
688
+ the worker, so no shared memory or per-block parent arrays are needed.
689
+
690
+ args is (h5_path, channel_index, coarse_channel_width, time_bins, bandpass_flatten,
691
+ window_size, step_size, stat_threshold, want_spectrum). bandpass_flatten is a picklable
692
+ callable (channel) -> residuals (see _spline_flatten_bandpass). Each worker opens its own
693
+ h5py.File since h5py file handles are fork-unsafe to share.
694
+ """
695
+ (
696
+ h5_path,
697
+ channel_index,
698
+ coarse_channel_width,
699
+ time_bins,
700
+ bandpass_flatten,
701
+ window_size,
702
+ step_size,
703
+ stat_threshold,
704
+ want_spectrum,
705
+ ) = args
706
+
707
+ start = channel_index * coarse_channel_width
708
+ end = (channel_index + 1) * coarse_channel_width
709
+ if _IN_POOL_WORKER:
710
+ channel = _cached_h5_file(h5_path)["data"][:time_bins, 0, start:end]
711
+ else:
712
+ # Sequential mode runs on the parent's prefetch thread(s) — h5py handles must not
713
+ # be shared across threads, so keep the per-call open/close there
714
+ with h5py.File(h5_path, "r") as hf:
715
+ channel = hf["data"][:time_bins, 0, start:end]
716
+
717
+ # In-place DC spike removal. The spike sits at the channel center and the interpolation
718
+ # offsets reach at most ±3 bins around it (see _remove_dc_spike), so per-channel
719
+ # processing touches exactly the same bins as the historical block-based path — the
720
+ # offsets can never cross a coarse-channel boundary.
721
+ _remove_dc_spike(channel, coarse_channel_width, 1)
722
+
723
+ # float64 accumulation matches the viz path this replaces (_read_despiked_channel
724
+ # casts to float64 before integrating); one mean over the already-resident channel.
725
+ integrated = channel.mean(axis=0, dtype=np.float64) if want_spectrum else None
726
+
727
+ residuals = bandpass_flatten(channel)
728
+
729
+ k2 = _sliding_normality_k2(residuals, window_size, step_size)
730
+
731
+ # Summary histogram over all finite window statistics — a handful of numpy ops on ~1e4
732
+ # values, negligible next to the k2 computation itself. Out-of-range values are clipped
733
+ # onto the fixed edges so counts are never silently dropped.
734
+ finite = k2[np.isfinite(k2)]
735
+ stat_hist, _ = np.histogram(
736
+ np.clip(finite, ED_STAT_HIST_EDGES[0], ED_STAT_HIST_EDGES[-1]), bins=ED_STAT_HIST_EDGES
737
+ )
738
+
739
+ hit_windows = np.nonzero(k2 > stat_threshold)[0]
740
+ if hit_windows.size == 0:
741
+ return [], stat_hist, integrated
742
+
743
+ # p-values are only stored in metadata, so chi2.sf runs on the (small) hit subset only
744
+ pvalues = stats.chi2.sf(k2[hit_windows], 2)
745
+ hits = [
746
+ (start + int(j) * step_size, float(k2[j]), float(p))
747
+ for j, p in zip(hit_windows, pvalues, strict=True)
748
+ ]
749
+ return hits, stat_hist, integrated
750
+
751
+
752
+ def _coalesce_stamp_groups(stamp_starts, stamp_width: int) -> list[list[int]]:
753
+ """Group indices of start-sorted stamps whose read windows overlap or abut, so each
754
+ group pays one wide h5 read (#301). A group is capped at _COALESCE_MAX_BINS total
755
+ span to bound the wide buffer. Disjoint windows stay singleton groups — the read
756
+ pattern is then exactly the historical per-stamp one."""
757
+ if not len(stamp_starts):
758
+ # Unreachable from the extract-task construction (chunks are never empty), but a
759
+ # bare [0] seed group would index-error on any future empty caller
760
+ return []
761
+ groups: list[list[int]] = []
762
+ current = [0]
763
+ for i in range(1, len(stamp_starts)):
764
+ fits = stamp_starts[i] + stamp_width - stamp_starts[current[0]] <= _COALESCE_MAX_BINS
765
+ if stamp_starts[i] <= stamp_starts[current[-1]] + stamp_width and fits:
766
+ current.append(i)
767
+ else:
768
+ groups.append(current)
769
+ current = [i]
770
+ groups.append(current)
771
+ return groups
772
+
773
+
774
+ def _extract_stamps_worker(args: tuple) -> None:
775
+ """Fill one (obs_file, stamp-range) slice of the memmap-backed cadence .npy.
776
+
777
+ Each worker opens its own hdf5 handle and its own r+ view of the shared .npy, then
778
+ copies a stamp_width-wide window (over the first `time_bins` rows, polarization 0) around
779
+ each hit. Overlapping/abutting windows (the overlap_search triplets abut at offset
780
+ stamp_width // 2) are fetched as one coalesced wide read and sliced per stamp —
781
+ byte-identical values, ~one read instead of three over an RFI comb (#301). When
782
+ downsample_factor > 1 each stamp is downsampled along frequency with
783
+ downscale_local_mean before writing, so the memmap stores width
784
+ stamp_width // downsample_factor — this is the same operation load_inference_data used to
785
+ apply after the fact, moved to extraction time to cut storage by the same factor. Tasks
786
+ address disjoint output regions — distinct obs_idx and/or non-overlapping stamp indices —
787
+ so concurrent writes from the pool never collide. `stamp_starts` is the contiguous,
788
+ start-sorted slice for this task; `base_idx` is its offset into the full stamp list so the
789
+ worker writes to the correct absolute rows. `open_kwargs` is the chunk-cache sizing
790
+ computed ONCE per obs file in the parent (#301 — each task used to re-open the file just
791
+ to inspect its chunk layout; same hoist as the PFB-check tasks).
792
+
793
+ NOTE: no out.flush() before the memmap closes (#301). flush() msync(MS_SYNC)s the ENTIRE
794
+ multi-GB mapping once per task (~132 tasks/cadence), serializing writeback into worker
795
+ time; closing the mapping leaves the dirty pages to the kernel's async writeback with
796
+ identical read coherence (unified page cache). Crash durability is unchanged — the
797
+ os.replace publication never fsynced the data pages anyway, and an unpublished tmp is
798
+ re-extracted by design.
799
+ """
800
+ (
801
+ npy_path,
802
+ obs_idx,
803
+ obs_h5,
804
+ stamp_starts,
805
+ base_idx,
806
+ time_bins,
807
+ stamp_width,
808
+ downsample_factor,
809
+ open_kwargs,
810
+ ) = args
811
+ out = np.lib.format.open_memmap(npy_path, mode="r+")
812
+ try:
813
+ # Chunk-cache sizing is what makes the caller's start-sorted stamp order pay off:
814
+ # with the h5py default (1 MiB) a BL-scale bitshuffle chunk never fits, so every
815
+ # stamp re-decompresses its full chunk stripe; sized from the actual layout, each
816
+ # chunk decompresses once per task (see _chunk_cache_kwargs).
817
+ with h5py.File(obs_h5, "r", **(open_kwargs or {})) as hf:
818
+ dset = hf["data"]
819
+ for group in _coalesce_stamp_groups(stamp_starts, stamp_width):
820
+ group_start = stamp_starts[group[0]]
821
+ group_end = stamp_starts[group[-1]] + stamp_width
822
+ wide = dset[:time_bins, 0, group_start:group_end]
823
+ for local_i in group:
824
+ offset = stamp_starts[local_i] - group_start
825
+ stamp = wide[:, offset : offset + stamp_width]
826
+ if downsample_factor > 1:
827
+ stamp = downscale_local_mean(stamp, (1, downsample_factor)).astype(
828
+ np.float32
829
+ )
830
+ out[base_idx + local_i, obs_idx] = stamp
831
+ finally:
832
+ del out
833
+
834
+
835
+ # NOTE: come back to this later
836
+ # def _drop_side_channels(
837
+ # block_data: np.ndarray, side_channel_count: int, coarse_channel_width: int
838
+ # ) -> None:
839
+ # """Zero out the leading/trailing side_channel_count coarse channels of a block.
840
+ # Reserved for future use — leave as-is until the energy-profile criterion is defined."""
841
+ # pass
842
+
843
+
844
+ @dataclass
845
+ class CadenceGroup:
846
+ """One cadence's worth of observations grouped from a CSV."""
847
+
848
+ key: tuple # The group-by column values
849
+ h5_paths: list[str] # Observation .h5 paths, in row order
850
+ csv_path: str # Source CSV
851
+ expected_obs: int
852
+ is_valid: bool # True iff len(h5_paths) == expected_obs
853
+
854
+
855
+ # NOTE: come back to this later (what does metadata_path store exactly?)
856
+ @dataclass
857
+ class CadenceResult:
858
+ """Output of processing one cadence."""
859
+
860
+ npy_path: str
861
+ h5_paths: list[str] # Same as in CadenceGroup
862
+ key: tuple # Same as in CadenceGroup
863
+ # Number of stamp rows in the .npy (post-dedup, incl. overlap offsets) — the same
864
+ # quantity the inference_cadences manifest stores as n_stamps. Historical name: this
865
+ # is NOT the raw energy-detection hit count (metadata's n_raw_hits carries that).
866
+ n_hits: int
867
+ metadata_path: str # Sibling .json with hit details
868
+ # True iff THIS run's _process_cadence wrote the .npy (False on the resume path).
869
+ # Stamp-cache pruning (#302) only ever deletes freshly-extracted stamps, so a run
870
+ # resumed over a handed-over cache can never destroy it.
871
+ freshly_extracted: bool = False
872
+
873
+
874
+ # NOTE: come back to this later
875
+ @dataclass
876
+ class CadenceHit:
877
+ """A single energy detection hit on an ON-source observation."""
878
+
879
+ fine_channel: int # absolute fine-channel index into the full spectrum
880
+ statistic: float # D'Agostino-Pearson statistic
881
+ pvalue: float
882
+ frequency_mhz: float = field(default=float("nan"))
883
+
884
+
885
+ @dataclass
886
+ class PendingCadence:
887
+ """One unit of preprocessing work: a valid cadence group plus its target .npy path."""
888
+
889
+ group: CadenceGroup
890
+ npy_path: str
891
+ # 1-based position in plan_cadences() output (CSV order, so stable across resumed
892
+ # attempts) — the short label used in this cadence's pipeline_stages span names
893
+ index: int = 0
894
+
895
+
896
+ def derive_cadence_provenance(key: tuple, group_by_cols: list[str], metadata: dict) -> dict:
897
+ """
898
+ Map one cadence's group-by key and stamp metadata JSON onto the observational-provenance
899
+ fields of the inference_results table.
900
+
901
+ key and group_by_cols come from the CadenceGroup (values zipped positionally onto column
902
+ names, matched case-insensitively so 'Target'/'target' both resolve); metadata is the
903
+ per-cadence JSON written by _process_cadence. Returns a dict with keys target, session,
904
+ band, cadence_id (int when parseable, else None), timestamp_observed (the header's tstart,
905
+ when present), h5_path (first observation of the cadence), and stamp_frequencies_mhz (the
906
+ per-stamp center frequencies, one per snippet row in the .npy).
907
+ """
908
+ # strict=False: a malformed key/cols pairing degrades to sparse provenance rather than
909
+ # aborting the cadence
910
+ key_map = {str(col).strip().lower(): val for col, val in zip(group_by_cols, key, strict=False)}
911
+
912
+ cadence_id: int | None = None
913
+ raw_cadence_id = key_map.get("cadence id")
914
+ if raw_cadence_id is not None:
915
+ try:
916
+ cadence_id = int(raw_cadence_id)
917
+ except (TypeError, ValueError):
918
+ logger.warning(f"Could not parse cadence id {raw_cadence_id!r} as int; storing None")
919
+
920
+ header = metadata.get("header") or {}
921
+ timestamp_observed: float | None = None
922
+ if "tstart" in header:
923
+ try:
924
+ timestamp_observed = float(header["tstart"])
925
+ except (TypeError, ValueError):
926
+ logger.warning(f"Could not parse header tstart {header['tstart']!r} as float")
927
+
928
+ h5_paths = metadata.get("h5_paths") or []
929
+
930
+ return {
931
+ "target": key_map.get("target"),
932
+ "session": key_map.get("session"),
933
+ "band": key_map.get("band"),
934
+ "cadence_id": cadence_id,
935
+ "timestamp_observed": timestamp_observed,
936
+ "h5_path": h5_paths[0] if h5_paths else None,
937
+ "stamp_frequencies_mhz": metadata.get("stamp_frequencies_mhz"),
938
+ }
939
+
940
+
941
+ # NOTE: come back to this later (add sorting functionality to sort rows in csv after grouping, e.g. via timestamp metadata from filenames? edge case where multiple 6-cadence observations of the same target & with the same grouping params, but differed by time, e.g. t=X to t=X+ε, then t=Y to t=Y+σ, for some small numbers ε and σ, and where X and Y are far apart from each other. add a way to distinguish these cases from problematic cases where we actually want to invalidate a cadence with a weird number of grouped observations, e.g. if multiple of 6 and enough of a gap between X and Y, then count as separate cadences?)
942
+ def group_observations_from_csv(
943
+ csv_path: str,
944
+ group_by_cols: list[str],
945
+ h5_path_col: str,
946
+ expected_obs: int = 6,
947
+ ) -> tuple[list[CadenceGroup], list[CadenceGroup]]:
948
+ """
949
+ Group rows of a CSV into cadences and return (valid_groups, flagged_groups).
950
+
951
+ Rows are grouped by the joint value of group_by_cols and assumed to be already ordered
952
+ correctly within each group in the source CSV. expected_obs (typically 6) is the required
953
+ number of observations per cadence; groups with the wrong count are returned in
954
+ flagged_groups rather than valid_groups. The function is column-agnostic — it never assumes
955
+ specific column names beyond what the caller provides.
956
+
957
+ Raises FileNotFoundError if csv_path doesn't exist, and KeyError if any column in
958
+ group_by_cols + [h5_path_col] is missing from the CSV header.
959
+ """
960
+ if not os.path.exists(csv_path):
961
+ raise FileNotFoundError(f"CSV not found: {csv_path}")
962
+
963
+ required_cols = list(group_by_cols) + [h5_path_col]
964
+
965
+ groups: OrderedDict[tuple, list[str]] = OrderedDict()
966
+
967
+ with open(csv_path, newline="") as f:
968
+ reader = csv.DictReader(f)
969
+ available = reader.fieldnames or []
970
+
971
+ # NOTE: come back to this later (does this fail for all if one csv has missing columns, or does it skip the invalid csvs and continue processing the valid ones? is proper downstream checkpointing implemented in the latter case?)
972
+ missing = [c for c in required_cols if c not in available]
973
+ if missing:
974
+ raise KeyError(
975
+ f"CSV {csv_path} is missing required column(s): {missing}. "
976
+ f"Available columns: {available}"
977
+ )
978
+
979
+ for row_idx, row in enumerate(reader, start=1):
980
+ key = tuple(row[c] for c in group_by_cols)
981
+ h5_path = row[h5_path_col]
982
+ # A ragged row (fewer fields than the header) makes csv.DictReader fill the path
983
+ # with None; an empty/whitespace-only cell is just as unusable. Skip such rows
984
+ # loudly instead of grouping them under a None/empty path, which would otherwise
985
+ # only surface much later as a cryptic failure when the .h5 is opened.
986
+ if h5_path is None or not str(h5_path).strip():
987
+ logger.warning(
988
+ f"CSV {csv_path}: data row {row_idx} (key={key}) has a missing/empty "
989
+ f"'{h5_path_col}' value ({h5_path!r}); skipping row"
990
+ )
991
+ continue
992
+ groups.setdefault(key, []).append(h5_path)
993
+
994
+ valid: list[CadenceGroup] = []
995
+ flagged: list[CadenceGroup] = []
996
+ for key, h5_paths in groups.items():
997
+ is_valid = len(h5_paths) == expected_obs
998
+ cg = CadenceGroup(
999
+ key=key,
1000
+ h5_paths=h5_paths,
1001
+ csv_path=csv_path,
1002
+ expected_obs=expected_obs,
1003
+ is_valid=is_valid,
1004
+ )
1005
+ (valid if is_valid else flagged).append(cg)
1006
+
1007
+ if flagged:
1008
+ sample = [(g.key, len(g.h5_paths)) for g in flagged[:5]]
1009
+ logger.warning(
1010
+ f"Found {len(flagged)} cadence group(s) in {csv_path} with obs count != "
1011
+ f"{expected_obs}; flagged and skipped. First {len(sample)}: {sample}"
1012
+ )
1013
+
1014
+ logger.info(
1015
+ f"Grouped {csv_path}: {len(valid)} valid cadence(s), {len(flagged)} flagged "
1016
+ f"(expected_obs={expected_obs})"
1017
+ )
1018
+
1019
+ return valid, flagged
1020
+
1021
+
1022
+ class DataPreprocessor:
1023
+ """Data preprocessor"""
1024
+
1025
+ def __init__(self):
1026
+ """
1027
+ Initialize preprocessor
1028
+ """
1029
+ self.config = get_config()
1030
+ if self.config is None:
1031
+ raise ValueError("get_config() returned None")
1032
+
1033
+ # TODO: add db writes for inference
1034
+ self.db = get_db()
1035
+ if self.db is None:
1036
+ raise ValueError("get_db() returned None")
1037
+
1038
+ self.manager = get_manager()
1039
+ if self.manager is None:
1040
+ raise ValueError("get_manager() returned None")
1041
+
1042
+ # Persistent worker pool for energy detection + stamp extraction, shared across all
1043
+ # cadences of a run (see start/stop_energy_detection_pool)
1044
+ self._ed_pool: Pool | None = None
1045
+
1046
+ # npy_paths THIS process freshly extracted (#302 disk-leak fix): survives the
1047
+ # in-process retry loop because the preprocessor outlives it. A cadence this run
1048
+ # extracted then resumes on a later attempt (its stamps were kept because its
1049
+ # inference stage failed) is still prunable — unlike a genuinely handed cache,
1050
+ # which this process never extracted and so is absent from the set.
1051
+ self._extracted_this_run: set[str] = set()
1052
+
1053
+ # NOTE: come back to this later
1054
+ def close(self):
1055
+ """Explicitly close the multiprocessing pool and shared memory"""
1056
+ # if hasattr(self, "pool") and self.pool is not None:
1057
+ # self.manager.close_pool(self.pool)
1058
+ # self.pool = None
1059
+ #
1060
+ # if hasattr(self, "shm") and self.shm is not None:
1061
+ # self.manager.close_shared_memory(self.shm)
1062
+ # self.shm = None
1063
+
1064
+ # Defense-in-depth: ensure the persistent energy-detection pool is torn down even if a
1065
+ # caller forgot to call stop_energy_detection_pool(). No-op when no pool was started.
1066
+ self.stop_energy_detection_pool()
1067
+
1068
+ logger.info("DataPreprocessor closed")
1069
+
1070
+ # NOTE: shared resources currently created & destroyed within function itself. think about abstractions once preprocessing.py is complete
1071
+ def load_train_data(self) -> np.ndarray:
1072
+ """
1073
+ Load and preprocess training data into an array of shape (n, 6, 16, width_bin_downsampled).
1074
+
1075
+ Uses a multiprocessing pool over shared memory to downsample cadences in parallel.
1076
+ Log-normalization is deferred to data_generation.py (training-side log-norm runs
1077
+ per-sample after injection), unlike load_inference_data which applies it here.
1078
+ """
1079
+ logger.info(f"Loading backgrounds from {self.config.data_path} for training")
1080
+
1081
+ downsample_factor = self.config.data.downsample_factor
1082
+ final_width = self.config.data.width_bin // downsample_factor
1083
+
1084
+ num_target_backgrounds = self.config.data.num_target_backgrounds
1085
+ chunk_size = self.config.data.background_load_chunk_size
1086
+ max_chunks = self.config.data.max_chunks_per_file
1087
+ n_processes = self.config.manager.n_processes
1088
+ chunks_per_worker = self.config.manager.chunks_per_worker
1089
+
1090
+ logger.info(f"Target backgrounds: {num_target_backgrounds}")
1091
+ logger.info(f"Processing chunks of: {chunk_size}")
1092
+ logger.info(f"Final resolution: {final_width}")
1093
+
1094
+ all_backgrounds = [] # NOTE: preallocate this as empty ndarray?
1095
+
1096
+ for filename in self.config.data.train_files:
1097
+ filepath = self.config.get_training_file_path(filename)
1098
+
1099
+ if not os.path.exists(filepath):
1100
+ logger.warning(f"File not found: {filepath}")
1101
+ continue
1102
+
1103
+ logger.info(f"Processing {filename}...")
1104
+
1105
+ try:
1106
+ # Use read-only memory mapping to avoid loading full file into memory
1107
+ # That is, insted of loading the whole file from disk to memory synchronously
1108
+ # The OS' virtual memory manager establishes a virtual address space pointer
1109
+ # from the file's location on disk to the virtual memory of the Python process
1110
+ # This allows us to lazy load the data on-demand page-by-page using page fault
1111
+ # Benefits of this approach include: reduced startup latency,
1112
+ # efficient memory usage (since the memory allocated for the mapped array does not
1113
+ # count towards the Python process' heap memory usage, allowing us to raise the
1114
+ # ceiling up to our OS' virtual memory limits, which is typically constrained by
1115
+ # free disk space and our system's address space, rather than physical RAM),
1116
+ # and optimized access patterns (spatial locality since data is loaded in pages,
1117
+ # and shared memory for multiprocess/multithreaded programs)
1118
+ raw_data = np.load(filepath, mmap_mode="r")
1119
+
1120
+ except Exception as e:
1121
+ logger.error(f"Error loading {filename}: {e}")
1122
+ continue
1123
+
1124
+ # Apply subset parameters if specified in config
1125
+ start, end = self.config.get_file_subset(filename)
1126
+ if start is not None or end is not None:
1127
+ raw_data = raw_data[start:end]
1128
+
1129
+ logger.info(f" Raw data shape: {raw_data.shape}")
1130
+
1131
+ # Divide background into equal chunks, then cutoff if exceeds max_chunks
1132
+ n_backgrounds_total = raw_data.shape[0]
1133
+ n_chunks = min(max_chunks, (n_backgrounds_total + chunk_size - 1) // chunk_size)
1134
+
1135
+ for chunk_idx in range(n_chunks):
1136
+ logger.info(f"Processing {filename}: chunk {chunk_idx + 1}/{n_chunks}")
1137
+
1138
+ chunk_start = chunk_idx * chunk_size
1139
+ chunk_end = min((chunk_idx + 1) * chunk_size, n_backgrounds_total)
1140
+
1141
+ # Load chunk into memory
1142
+ chunk_data = np.array(raw_data[chunk_start:chunk_end])
1143
+
1144
+ # NOTE: is this access pattern the most efficient (least pickling)? see _downsample_worker()'s docstring on pulling the cadence from the shared-memory global to avoid per-cadence pickling
1145
+ # NOTE: currently, loading the backgrounds takes WAY more time than processing the backgrounds
1146
+ # Prepare arguments for downsampling (just indices, not data - data is in global state)
1147
+ n_cadences = min(chunk_data.shape[0], num_target_backgrounds - len(all_backgrounds))
1148
+ args_list = [
1149
+ (
1150
+ i,
1151
+ downsample_factor,
1152
+ final_width,
1153
+ ) # Just pass the chunk index, not the full cadence data
1154
+ for i in range(n_cadences)
1155
+ ]
1156
+
1157
+ # NOTE: do we need to create & destroy the pool every chunk? or just the shared memory & pass new references in? is there a differenc?
1158
+ if n_processes > 1:
1159
+ # Create shared memory block for chunk data
1160
+ chunk_shm = self.manager.create_shared_memory(
1161
+ size=chunk_data.nbytes,
1162
+ name=f"DataPreproc_{filename}_chunk_{chunk_idx}", # NOTE: come back to this later
1163
+ )
1164
+
1165
+ # Copy chunk data into shared memory
1166
+ shared_chunk = np.ndarray(
1167
+ chunk_data.shape,
1168
+ dtype=chunk_data.dtype,
1169
+ buffer=chunk_shm.buf, # NOTE: what is self.shm.buf?
1170
+ )
1171
+ shared_chunk[:] = chunk_data[:]
1172
+
1173
+ # Create pool using shared memory reference
1174
+ chunk_pool = self.manager.create_pool(
1175
+ n_processes=n_processes,
1176
+ name=f"DataPreproc_{filename}_chunk_{chunk_idx}", # NOTE: come back to this later
1177
+ initializer=_init_worker,
1178
+ initargs=(chunk_shm.name, chunk_data.shape, chunk_data.dtype),
1179
+ )
1180
+
1181
+ # Calculate optimal chunksize for load balancing
1182
+ try:
1183
+ n_workers = chunk_pool._processes
1184
+ except AttributeError:
1185
+ n_workers = n_processes
1186
+ # NOTE: should we use separate chunks_per_worker? how to benchmark?
1187
+ chunksize = max(1, n_cadences // (n_workers * chunks_per_worker))
1188
+ # TEST: does return order matter?
1189
+ results = chunk_pool.map(_downsample_worker, args_list, chunksize=chunksize)
1190
+ # results = chunk_pool.imap(_downsample_worker, args_list, chunksize=chunksize)
1191
+ # results = chunk_pool.imap_unordered(
1192
+ # _downsample_worker, args_list, chunksize=chunksize
1193
+ # )
1194
+
1195
+ else:
1196
+ # Sequential processing
1197
+ logger.info("DataPreprocessor running in sequential mode (n_processes=1)")
1198
+
1199
+ chunk_shm = None
1200
+ chunk_pool = None
1201
+
1202
+ # Set global variable manually since no initializer ran
1203
+ global _GLOBAL_CHUNK_DATA
1204
+ shared_chunk = chunk_data
1205
+ _GLOBAL_CHUNK_DATA = shared_chunk
1206
+
1207
+ results = [_downsample_worker(args) for args in args_list]
1208
+
1209
+ # NOTE: is there a more efficient/elegant way to do this (e.g. with list comprehension/slicing)?
1210
+ # Collect valid results (filter out None from invalid cadences)
1211
+ for result in results:
1212
+ if result is not None:
1213
+ all_backgrounds.append(result)
1214
+ if len(all_backgrounds) >= num_target_backgrounds:
1215
+ break
1216
+
1217
+ # Clear chunk data & shared resources
1218
+ del chunk_data, shared_chunk
1219
+ if chunk_shm:
1220
+ self.manager.close_shared_memory(chunk_shm)
1221
+ chunk_shm = None
1222
+ if chunk_pool:
1223
+ self.manager.close_pool(chunk_pool)
1224
+ chunk_pool = None
1225
+ del chunk_shm, chunk_pool
1226
+ # Release the module-global staging reference (#298 rider) — same leak as
1227
+ # the inference loader: the local dels above don't clear the global name
1228
+ _GLOBAL_CHUNK_DATA = None
1229
+ gc.collect()
1230
+
1231
+ # Clear raw_data reference
1232
+ del raw_data
1233
+ gc.collect()
1234
+
1235
+ if len(all_backgrounds) == 0:
1236
+ raise ValueError("No data loaded successfully")
1237
+
1238
+ # Stack all_backgrounds together
1239
+ background_array = np.array(all_backgrounds, dtype=np.float32)
1240
+
1241
+ # Clear all_backgrounds reference
1242
+ del all_backgrounds
1243
+ gc.collect()
1244
+
1245
+ # Sanity check: print descriptive stats
1246
+ min_val = np.min(background_array)
1247
+ max_val = np.max(background_array)
1248
+ mean_val = np.mean(background_array)
1249
+
1250
+ logger.info(f"Total backgrounds loaded: {background_array.shape[0]}")
1251
+ logger.info(f"Background array shape: {background_array.shape}")
1252
+ logger.info(f"Background value range: [{min_val:.6f}, {max_val:.6f}]")
1253
+ logger.info(f"Background mean: {mean_val:.6f}")
1254
+ logger.info(f"Memory usage: {background_array.nbytes / 1e9:.2f} GB")
1255
+ logger.info(f"Background data ready at {background_array.shape[3]} resolution")
1256
+
1257
+ return background_array
1258
+
1259
+ # NOTE: shared resources currently created & destroyed within function itself. think about abstractions once preprocessing.py is complete
1260
+ # NOTE: calculate intensity statistics to overlay with training distributions (C' vs C)?
1261
+ def load_inference_data(
1262
+ self, override_filepaths: list[str] | None = None, parallel: bool = True
1263
+ ) -> np.ndarray:
1264
+ """
1265
+ Load and preprocess inference data into an array of shape
1266
+ (n, 6, 16, width_bin_downsampled). Uses a multiprocessing pool over shared memory to
1267
+ process cadences in parallel.
1268
+
1269
+ Each file's stored width decides its path: files already at the downsampled width
1270
+ (written by _extract_stamps_worker with store_downsampled_stamps enabled) only need
1271
+ per-cadence log-normalization, which runs vectorized in the pool workers; legacy
1272
+ full-width files (width_bin, e.g. stamps preprocessed before downsample-at-extraction
1273
+ landed) keep the historical downsample-then-log-norm behavior.
1274
+
1275
+ override_filepaths, when given, supplies absolute paths to iterate directly instead of
1276
+ resolving config.data.test_files via get_test_file_path — used by the inference command
1277
+ to chain per-cadence .npy outputs from preprocessing into inference without
1278
+ monkey-patching paths.
1279
+
1280
+ parallel=False routes every chunk through the sequential in-process branch (no chunk
1281
+ pool, no shared memory) regardless of manager.n_processes. The streaming per-cadence
1282
+ path (main._prefetch_cadence) uses this: it loads exactly one already-downsampled
1283
+ cadence .npy whose per-cadence work is a cheap vectorized log-norm, while the
1284
+ persistent energy-detection pool is already busy at full n_processes width — forking
1285
+ a second n_processes pool for that would double-subscribe the CPU. THREAD CONTRACT
1286
+ (#298): parallel=False may be called from prefetch_depth CONCURRENT threads, so
1287
+ neither sequential branch may touch module-global state — the vectorized log-norm
1288
+ path is pure, and the legacy downsample path passes rows to _downsample_cadence
1289
+ directly (never via _GLOBAL_CHUNK_DATA, which stays a pool-worker-only mechanism).
1290
+ """
1291
+ logger.info(f"Loading backgrounds from {self.config.data_path} for inference")
1292
+
1293
+ downsample_factor = self.config.data.downsample_factor
1294
+ width_bin = self.config.data.width_bin
1295
+ final_width = width_bin // downsample_factor
1296
+
1297
+ chunk_size = self.config.data.inference_background_load_chunk_size
1298
+ n_processes = self.config.manager.n_processes
1299
+ chunks_per_worker = self.config.manager.chunks_per_worker
1300
+
1301
+ logger.info(f"Processing chunks of: {chunk_size}")
1302
+ logger.info(f"Final resolution: {final_width}")
1303
+
1304
+ global _GLOBAL_CHUNK_DATA
1305
+ # Per-chunk float32 blocks of already-normalized cadences, concatenated once at the
1306
+ # end (#298 I5) — replaces the per-cadence list whose final np.array() restack
1307
+ # re-copied every row through a Python loop.
1308
+ all_blocks: list[np.ndarray] = []
1309
+
1310
+ if override_filepaths is not None:
1311
+ # Iterate absolute paths directly (e.g. per-cadence .npy outputs from find_hits)
1312
+ file_iter = [(os.path.basename(p), p) for p in override_filepaths]
1313
+ else:
1314
+ file_iter = [
1315
+ (filename, self.config.get_test_file_path(filename))
1316
+ for filename in self.config.data.test_files
1317
+ ]
1318
+
1319
+ for filename, filepath in file_iter:
1320
+ if not os.path.exists(filepath):
1321
+ logger.warning(f"File not found: {filepath}")
1322
+ continue
1323
+
1324
+ logger.info(f"Processing {filename}...")
1325
+
1326
+ try:
1327
+ # Use read-only memory mapping to avoid loading full file into memory
1328
+ # That is, insted of loading the whole file from disk to memory synchronously
1329
+ # The OS' virtual memory manager establishes a virtual address space pointer
1330
+ # from the file's location on disk to the virtual memory of the Python process
1331
+ # This allows us to lazy load the data on-demand page-by-page using page fault
1332
+ # Benefits of this approach include: reduced startup latency,
1333
+ # efficient memory usage (since the memory allocated for the mapped array does not
1334
+ # count towards the Python process' heap memory usage, allowing us to raise the
1335
+ # ceiling up to our OS' virtual memory limits, which is typically constrained by
1336
+ # free disk space and our system's address space, rather than physical RAM),
1337
+ # and optimized access patterns (spatial locality since data is loaded in pages,
1338
+ # and shared memory for multiprocess/multithreaded programs)
1339
+ raw_data = np.load(filepath, mmap_mode="r")
1340
+
1341
+ except Exception as e:
1342
+ logger.error(f"Error loading {filename}: {e}")
1343
+ continue
1344
+
1345
+ # Apply subset parameters if specified in config
1346
+ start, end = self.config.get_file_subset(filename)
1347
+ if start is not None or end is not None:
1348
+ raw_data = raw_data[start:end]
1349
+
1350
+ logger.info(f" Raw data shape: {raw_data.shape}")
1351
+
1352
+ # A cadence plate must be 4-D (n, 6, time_bins, width); only the trailing width was
1353
+ # validated historically, so a wrong-rank array would flow into the pool workers and
1354
+ # either fail cryptically or silently mis-index the 6 observations. Skip-and-warn,
1355
+ # matching the unsupported-width skip below, so one malformed file can't sink the run.
1356
+ if raw_data.ndim != 4:
1357
+ logger.error(
1358
+ f" {filename} has ndim {raw_data.ndim} (shape {raw_data.shape}), "
1359
+ f"expected 4 (n, 6, time_bins, width); skipping file"
1360
+ )
1361
+ del raw_data
1362
+ continue
1363
+
1364
+ # Everything downstream casts to float32; a non-float input (e.g. integer power
1365
+ # counts) was previously coerced silently. Surface the coercion loudly — the values
1366
+ # are unchanged, but the dtype change is worth naming so a mis-typed catalog is
1367
+ # visible rather than hidden.
1368
+ if not np.issubdtype(raw_data.dtype, np.floating):
1369
+ logger.warning(
1370
+ f" {filename} has non-float dtype {raw_data.dtype}; coercing to float32"
1371
+ )
1372
+
1373
+ # Branch on the stored width: already-downsampled stamps (written by
1374
+ # _extract_stamps_worker) only need log-norm; legacy full-width files keep the
1375
+ # historical downsample-then-log-norm path.
1376
+ stored_width = raw_data.shape[-1]
1377
+ if stored_width == final_width:
1378
+ already_downsampled = True
1379
+ logger.info(f" {filename} stored at final width {final_width}: log-norm only")
1380
+ elif stored_width == width_bin:
1381
+ already_downsampled = False
1382
+ logger.info(f" {filename} stored at full width {width_bin}: downsample + log-norm")
1383
+ else:
1384
+ logger.error(
1385
+ f" {filename} width {stored_width} matches neither width_bin "
1386
+ f"({width_bin}) nor width_bin // downsample_factor ({final_width}); "
1387
+ f"skipping file"
1388
+ )
1389
+ del raw_data
1390
+ continue
1391
+
1392
+ # Divide background into equal chunks, then cutoff if exceeds max_chunks
1393
+ n_cadences_total = raw_data.shape[0]
1394
+ # One chunk per already-downsampled file on the sequential path (#301): the
1395
+ # chunking bounds the SHM block on the pooled path and RAM on legacy
1396
+ # full-width loads, but here it only split one cadence's stamps into blocks
1397
+ # that then paid a full-array np.concatenate re-copy at the end — peak
1398
+ # transient is ~2x the array either way.
1399
+ effective_chunk_size = chunk_size
1400
+ if already_downsampled and not parallel:
1401
+ effective_chunk_size = max(chunk_size, n_cadences_total)
1402
+ n_chunks = (n_cadences_total + effective_chunk_size - 1) // effective_chunk_size
1403
+
1404
+ for chunk_idx in range(n_chunks):
1405
+ logger.info(f"Processing {filename}: chunk {chunk_idx + 1}/{n_chunks}")
1406
+
1407
+ chunk_start = chunk_idx * effective_chunk_size
1408
+ chunk_end = min((chunk_idx + 1) * effective_chunk_size, n_cadences_total)
1409
+
1410
+ # Load chunk into memory
1411
+ chunk_data = np.array(raw_data[chunk_start:chunk_end])
1412
+
1413
+ # NOTE: is this access pattern the most efficient (least pickling)? see _downsample_worker()'s docstring on pulling the cadence from the shared-memory global to avoid per-cadence pickling
1414
+ # NOTE: currently, loading the backgrounds takes WAY more time than processing the backgrounds
1415
+ # Prepare arguments (just indices, not data - data is in global state).
1416
+ # Downsampled files fold log-norm into the workers (resolving the old
1417
+ # "downsample & log-norm simultaneously" TODO); legacy files downsample in
1418
+ # the workers and log-norm in-process below, exactly as before.
1419
+ n_cadences = chunk_data.shape[0]
1420
+ if already_downsampled:
1421
+ worker_fn = _lognorm_worker
1422
+ args_list = [(i,) for i in range(n_cadences)]
1423
+ else:
1424
+ worker_fn = _downsample_worker
1425
+ args_list = [(i, downsample_factor, final_width) for i in range(n_cadences)]
1426
+
1427
+ # NOTE: do we need to create & destroy the pool every chunk? or just the shared memory & pass new references in? is there a differenc?
1428
+ if parallel and n_processes > 1:
1429
+ # Create shared memory block for chunk data
1430
+ chunk_shm = self.manager.create_shared_memory(
1431
+ size=chunk_data.nbytes,
1432
+ name=f"DataPreproc_{filename}_chunk_{chunk_idx}", # NOTE: come back to this later
1433
+ )
1434
+
1435
+ # Copy chunk data into shared memory
1436
+ shared_chunk = np.ndarray(
1437
+ chunk_data.shape,
1438
+ dtype=chunk_data.dtype,
1439
+ buffer=chunk_shm.buf, # NOTE: what is self.shm.buf?
1440
+ )
1441
+ shared_chunk[:] = chunk_data[:]
1442
+
1443
+ # Create pool using shared memory reference
1444
+ chunk_pool = self.manager.create_pool(
1445
+ n_processes=n_processes,
1446
+ name=f"DataPreproc_{filename}_chunk_{chunk_idx}", # NOTE: come back to this later
1447
+ initializer=_init_worker,
1448
+ initargs=(chunk_shm.name, chunk_data.shape, chunk_data.dtype),
1449
+ )
1450
+
1451
+ # Calculate optimal chunksize for load balancing
1452
+ try:
1453
+ n_workers = chunk_pool._processes
1454
+ except AttributeError:
1455
+ n_workers = n_processes
1456
+ # NOTE: should we use separate chunks_per_worker? how to benchmark?
1457
+ chunksize = max(1, n_cadences // (n_workers * chunks_per_worker))
1458
+ # TEST: does return order matter?
1459
+ results = chunk_pool.map(worker_fn, args_list, chunksize=chunksize)
1460
+
1461
+ elif already_downsampled:
1462
+ # Sequential + already-downsampled — the streaming per-cadence path:
1463
+ # vectorized log-norm over the whole chunk (#298 I5), bit-identical per
1464
+ # element to mapping _lognorm_worker (see _log_norm_chunk_vectorized),
1465
+ # without the per-stamp Python loop and without staging the chunk in
1466
+ # the module global at all.
1467
+ logger.info(
1468
+ f"DataPreprocessor running vectorized sequential log-norm "
1469
+ f"(parallel={parallel}, n_processes={n_processes})"
1470
+ )
1471
+
1472
+ chunk_shm = None
1473
+ chunk_pool = None
1474
+ shared_chunk = None
1475
+ results = []
1476
+
1477
+ normalized, _ = _log_norm_chunk_vectorized(chunk_data)
1478
+ if len(normalized):
1479
+ all_blocks.append(normalized)
1480
+ del normalized
1481
+
1482
+ else:
1483
+ # Sequential processing (legacy full-width files). Rows are passed to
1484
+ # the downsample core DIRECTLY — never via _GLOBAL_CHUNK_DATA: this
1485
+ # branch runs on the caller's thread, and the streaming loop calls
1486
+ # load_inference_data(parallel=False) from prefetch_depth concurrent
1487
+ # threads, where a module-global staging slot would race (#298 review
1488
+ # note). The global stays a pool-worker-only mechanism.
1489
+ logger.info(
1490
+ f"DataPreprocessor running in sequential mode "
1491
+ f"(parallel={parallel}, n_processes={n_processes})"
1492
+ )
1493
+
1494
+ chunk_shm = None
1495
+ chunk_pool = None
1496
+ shared_chunk = None
1497
+
1498
+ results = [
1499
+ _downsample_cadence(chunk_data[i], downsample_factor, final_width)
1500
+ for i, _, _ in args_list
1501
+ ]
1502
+
1503
+ # NOTE: is there a more efficient/elegant way to do this (e.g. with list comprehension/slicing)?
1504
+ # Collect valid results (filter out None from invalid cadences) and stack
1505
+ # them into one float32 block per chunk (the vectorized branch has already
1506
+ # appended its block; its results list is empty)
1507
+ chunk_cadences = [
1508
+ result if already_downsampled else log_norm(result)
1509
+ for result in results
1510
+ if result is not None
1511
+ ]
1512
+ if chunk_cadences:
1513
+ all_blocks.append(np.array(chunk_cadences, dtype=np.float32))
1514
+ del chunk_cadences
1515
+
1516
+ # Clear chunk data & shared resources. No gc.collect() here or below:
1517
+ # everything freed on these paths is refcount-managed (ndarrays, SHM
1518
+ # handles), and each full collection costs ~0.3 s with TF loaded — the
1519
+ # streaming loop calls this once per cadence (#298 I7).
1520
+ del chunk_data, shared_chunk
1521
+ if chunk_shm:
1522
+ self.manager.close_shared_memory(chunk_shm)
1523
+ chunk_shm = None
1524
+ if chunk_pool:
1525
+ self.manager.close_pool(chunk_pool)
1526
+ chunk_pool = None
1527
+ del chunk_shm, chunk_pool
1528
+ # Release the module-global staging reference (#298 rider): it otherwise
1529
+ # pins the last chunk (~6-10 GB of stamps) in the main process for the
1530
+ # whole encode/RF phase, until the next cadence's load overwrites it
1531
+ _GLOBAL_CHUNK_DATA = None
1532
+
1533
+ # Clear raw_data reference
1534
+ del raw_data
1535
+
1536
+ if not all_blocks:
1537
+ raise ValueError("No data loaded successfully")
1538
+
1539
+ # One contiguous float32 array (every cadence is downsampled + log-normalized by
1540
+ # now); the single-block case — the streaming loop's, always — is returned as-is
1541
+ # rather than paying a full concatenate copy
1542
+ cadence_array = all_blocks[0] if len(all_blocks) == 1 else np.concatenate(all_blocks)
1543
+
1544
+ # Clear all_blocks reference
1545
+ del all_blocks
1546
+
1547
+ # Sanity check: print descriptive stats. NaN detection rides the min reduction
1548
+ # (NaN propagates through np.min, and NaN > 1.0 / NaN < 0.0 are both False, so a
1549
+ # NaN array reaches that branch exactly as it did via the old full-array
1550
+ # .any() pass); the two extra full-array validity passes and their full-size
1551
+ # bool temporaries are gone (#301). The old isinf branch was unreachable: +inf
1552
+ # always tripped the max check first, -inf the min check, and NaN the isnan
1553
+ # branch — outcomes and messages are unchanged for every input.
1554
+ min_val = np.min(cadence_array)
1555
+ max_val = np.max(cadence_array)
1556
+ mean_val = np.mean(cadence_array)
1557
+
1558
+ if max_val > 1.0:
1559
+ logger.error(f"Cadence array values too large! Max: {max_val}")
1560
+ raise ValueError("Preprocessing normalization check failed")
1561
+ elif min_val < 0.0:
1562
+ logger.error(f"Cadence array values too small! Min: {min_val}")
1563
+ raise ValueError("Preprocessing normalization check failed")
1564
+ elif np.isnan(min_val):
1565
+ logger.error("Cadence array contains NaN values!")
1566
+ raise ValueError("Preprocessing normalization check failed")
1567
+ else:
1568
+ logger.info("Cadence array properly normalized")
1569
+ logger.info(f"Total cadences loaded: {cadence_array.shape[0]}")
1570
+ logger.info(f"Cadence array shape: {cadence_array.shape}")
1571
+ logger.info(f"Cadence value range: [{min_val:.6f}, {max_val:.6f}]")
1572
+ logger.info(f"Cadence mean: {mean_val:.6f}")
1573
+ logger.info(f"Memory usage: {cadence_array.nbytes / 1e9:.2f} GB")
1574
+ logger.info(f"Cadence data ready at {cadence_array.shape[3]} resolution")
1575
+
1576
+ return cadence_array
1577
+
1578
+ def start_energy_detection_pool(self) -> None:
1579
+ """
1580
+ Create the persistent worker pool used for energy detection and stamp extraction.
1581
+
1582
+ One pool serves every cadence of the run (no per-block or per-cadence pool churn).
1583
+ Call from the main thread before any cadence processing starts — forking from the main
1584
+ thread before background threads spin up avoids inheriting mid-operation locks — and
1585
+ pair with stop_energy_detection_pool() when the run is done. No-op when a pool already
1586
+ exists or n_processes == 1 (sequential mode).
1587
+ """
1588
+ if self._ed_pool is not None:
1589
+ return
1590
+ n_processes = self.config.manager.n_processes
1591
+ if n_processes > 1:
1592
+ self._ed_pool = self.manager.create_pool(
1593
+ n_processes=n_processes,
1594
+ name="DataPreproc_energy_detection",
1595
+ initializer=_init_plain_worker,
1596
+ )
1597
+
1598
+ def stop_energy_detection_pool(self) -> None:
1599
+ """Close the persistent energy-detection pool (no-op if never started)."""
1600
+ if self._ed_pool is not None:
1601
+ self.manager.close_pool(self._ed_pool)
1602
+ self._ed_pool = None
1603
+
1604
+ def plan_cadences(self) -> list[PendingCadence]:
1605
+ """
1606
+ Group every CSV in config.data.inference_files into per-cadence work units without
1607
+ processing anything. Returns one PendingCadence (valid group + target .npy path) per
1608
+ valid cadence, in CSV order — the unit list the streaming inference loop iterates.
1609
+ Raises ValueError when two inference_files entries share a basename stem (their
1610
+ output dir and stamp filenames would collide).
1611
+ """
1612
+ inference_files = self.config.data.inference_files
1613
+ if not inference_files:
1614
+ logger.warning("plan_cadences() called with no inference_files configured")
1615
+ return []
1616
+
1617
+ # Fail fast on duplicate CSV basename stems: both the fingerprint-scoped default
1618
+ # output dir and the per-cadence stamp .npy names are keyed on the stem, so two
1619
+ # entries like runA/x.csv and runB/x.csv would silently share a directory (and/or
1620
+ # stamp filenames) and cross-resume each other's cadences.
1621
+ stem_sources: dict[str, str] = {}
1622
+ for csv_filename in inference_files:
1623
+ stem = os.path.splitext(os.path.basename(csv_filename))[0]
1624
+ if stem in stem_sources:
1625
+ raise ValueError(
1626
+ f"Duplicate inference CSV basename stem '{stem}' "
1627
+ f"('{stem_sources[stem]}' vs '{csv_filename}'): output directories and "
1628
+ f"stamp filenames are keyed on the basename, so these entries would "
1629
+ f"overwrite/resume each other. Rename one so basenames are unique."
1630
+ )
1631
+ stem_sources[stem] = csv_filename
1632
+
1633
+ # Output directory resolution (#298 I3): an explicit --preprocess-output-dir is
1634
+ # used as-is (shared across CSVs); otherwise each CSV gets its own ED-FINGERPRINT-
1635
+ # scoped default, {data_path}/inference/preprocessed/<csv_stem>_ed<hash12>/. Energy
1636
+ # detection is deterministic given (csv, h5 files, ED config), so any run sharing
1637
+ # the ED config reuses the stamps — threshold sweeps and re-inference with new
1638
+ # weights skip preprocessing entirely — while any ED-config change lands in a
1639
+ # different directory by construction (the fingerprint is a fail-safe denylist —
1640
+ # see run_state.preprocessing_config_fingerprint). Staleness is impossible for a
1641
+ # published .npy (atomic tmp -> os.replace publication + per-run-unique tmp names);
1642
+ # the resume guard additionally verifies each sidecar's h5_paths and fingerprint.
1643
+ # This deliberately replaces the old <csv_stem>_<save_tag> tag-scoped default
1644
+ # (maintainer decision on #298): scores stay tag-scoped in the DB — only the
1645
+ # deterministic preprocessing artifacts are shared.
1646
+ explicit_output_dir = self.config.inference.preprocess_output_dir
1647
+ ed_fingerprint = preprocessing_config_fingerprint(self.config.to_dict())
1648
+
1649
+ group_by_cols = self.config.inference.cadence_group_by_cols
1650
+ h5_path_col = self.config.inference.cadence_h5_path_col
1651
+ expected_obs = self.config.inference.cadence_expected_obs
1652
+
1653
+ units: list[PendingCadence] = []
1654
+
1655
+ for csv_filename in inference_files:
1656
+ csv_path = self.config.get_inference_file_path(csv_filename)
1657
+ logger.info(f"Processing inference CSV: {csv_path}")
1658
+
1659
+ try:
1660
+ valid_groups, flagged_groups = group_observations_from_csv(
1661
+ csv_path=csv_path,
1662
+ group_by_cols=group_by_cols,
1663
+ h5_path_col=h5_path_col,
1664
+ expected_obs=expected_obs,
1665
+ )
1666
+ except (FileNotFoundError, KeyError) as e:
1667
+ logger.error(f"Failed to group CSV {csv_path}: {e}")
1668
+ continue
1669
+
1670
+ logger.info(
1671
+ f"CSV {csv_filename}: {len(valid_groups)} valid, {len(flagged_groups)} flagged"
1672
+ )
1673
+
1674
+ csv_stem = os.path.splitext(os.path.basename(csv_path))[0]
1675
+
1676
+ output_dir = explicit_output_dir or self.config.get_inference_file_path(
1677
+ os.path.join("preprocessed", f"{csv_stem}_ed{ed_fingerprint[:12]}")
1678
+ )
1679
+ os.makedirs(output_dir, exist_ok=True)
1680
+ logger.info(f"Preprocessing output directory for {csv_filename}: {output_dir}")
1681
+
1682
+ for group in valid_groups:
1683
+ npy_filename = self._cadence_npy_filename(csv_stem, group.key)
1684
+ units.append(
1685
+ PendingCadence(
1686
+ group=group,
1687
+ npy_path=os.path.join(output_dir, npy_filename),
1688
+ index=len(units) + 1,
1689
+ )
1690
+ )
1691
+
1692
+ logger.info(
1693
+ f"Planned {len(units)} cadence work unit(s) across {len(inference_files)} CSV(s)"
1694
+ )
1695
+ return units
1696
+
1697
+ def process_pending_cadence(self, unit: PendingCadence) -> CadenceResult | None:
1698
+ """
1699
+ Resume-or-process one cadence work unit. Returns a CadenceResult when a stamp .npy is
1700
+ available (freshly written or already on disk), or None when the cadence produced no
1701
+ hits or failed — single-cadence failures are logged and swallowed so one bad cadence
1702
+ can't abort the whole catalog (the retry loop at the inference_command level handles
1703
+ broader recovery).
1704
+ """
1705
+ group, npy_path = unit.group, unit.npy_path
1706
+
1707
+ if os.path.exists(npy_path):
1708
+ # Resume path: rebuild a minimal CadenceResult from the existing file
1709
+ metadata_path = self.cadence_metadata_path(npy_path)
1710
+ try:
1711
+ existing = np.load(npy_path, mmap_mode="r")
1712
+ n_hits = existing.shape[0]
1713
+ del existing
1714
+ except Exception as e:
1715
+ logger.warning(
1716
+ f"Existing .npy at {npy_path} could not be inspected ({e}); "
1717
+ f"reprocessing cadence"
1718
+ )
1719
+ # Fall through (no return) to the _process_cadence call below so a
1720
+ # corrupted .npy gets regenerated rather than silently skipped.
1721
+ else:
1722
+ if self._resume_provenance_ok(group, npy_path, metadata_path):
1723
+ logger.info(
1724
+ f"Skipping cadence {group.key}: {npy_path} already exists ({n_hits} hits)"
1725
+ )
1726
+ return CadenceResult(
1727
+ npy_path=npy_path,
1728
+ h5_paths=group.h5_paths,
1729
+ key=group.key,
1730
+ n_hits=n_hits,
1731
+ metadata_path=metadata_path,
1732
+ # Prunable iff THIS process extracted it (a failed-then-retried
1733
+ # cadence of this run), not if it is a handed cache (#302 leak fix)
1734
+ freshly_extracted=npy_path in self._extracted_this_run,
1735
+ )
1736
+ # Guard failed: fall through and reprocess (the atomic tmp -> os.replace
1737
+ # publication overwrites the mismatched artifact)
1738
+
1739
+ try:
1740
+ # Umbrella stage span for this cadence's preprocessing phase — the per-ON-file
1741
+ # read_ed / dedup / extract sub-stages inside _process_cadence nest under it
1742
+ # via thread-local naming. The resume path above records nothing (no work done)
1743
+ with stage_timer(f"inference.preprocess_cadence_{unit.index:03d}"):
1744
+ return self._process_cadence(group, npy_path)
1745
+ except Exception as e:
1746
+ logger.error(f"Failed to process cadence {group.key}: {e}")
1747
+ return None
1748
+
1749
+ def _resume_provenance_ok(self, group: CadenceGroup, npy_path: str, metadata_path: str) -> bool:
1750
+ """
1751
+ Provenance guard on stamp reuse (#298 I3): an existing .npy is trusted only when its
1752
+ sidecar metadata's h5_paths match this cadence's files and its recorded
1753
+ ed_config_fingerprint (when present) matches the current config's. With the
1754
+ fingerprint-scoped default directory the fingerprint check is belt-and-braces (a
1755
+ changed ED config lands in a different directory by construction); the h5_paths
1756
+ check is what protects an EXPLICIT --preprocess-output-dir against reusing another
1757
+ catalog's stamps. Missing/unreadable metadata or a metadata file without the
1758
+ fingerprint field (a pre-#298 artifact in an explicitly shared directory) degrades
1759
+ to a warn-and-reuse — the historical, unguarded behavior — never a hard failure.
1760
+ """
1761
+ try:
1762
+ with open(metadata_path) as f:
1763
+ metadata = json.load(f)
1764
+ except (OSError, json.JSONDecodeError) as e:
1765
+ logger.warning(
1766
+ f"Cadence {group.key}: existing stamps at {npy_path} have no readable "
1767
+ f"metadata sidecar ({e}); reusing them unguarded (historical behavior)"
1768
+ )
1769
+ return True
1770
+
1771
+ recorded_h5_paths = metadata.get("h5_paths")
1772
+ if recorded_h5_paths is not None and list(recorded_h5_paths) != list(group.h5_paths):
1773
+ logger.warning(
1774
+ f"Cadence {group.key}: existing stamps at {npy_path} were extracted from "
1775
+ f"DIFFERENT .h5 files than this cadence's; reprocessing instead of reusing"
1776
+ )
1777
+ return False
1778
+
1779
+ recorded_fingerprint = metadata.get("ed_config_fingerprint")
1780
+ if recorded_fingerprint is not None:
1781
+ current = preprocessing_config_fingerprint(self.config.to_dict())
1782
+ if recorded_fingerprint != current:
1783
+ logger.warning(
1784
+ f"Cadence {group.key}: existing stamps at {npy_path} were written under "
1785
+ f"a DIFFERENT energy-detection config; reprocessing instead of reusing"
1786
+ )
1787
+ return False
1788
+ return True
1789
+
1790
+ # NOTE: come back to this later (based on docstring, we're processing cadences sequentially. if so, any way to parallelize?)
1791
+ def find_hits(self) -> list[CadenceResult]:
1792
+ """
1793
+ Convert raw .h5 cadence observations into (n_hits, 6, 16, stored_width) .npy snippets,
1794
+ returning one CadenceResult per successfully processed (or already cached) cadence.
1795
+
1796
+ Driven by CSVs in config.data.inference_files. Each CSV is grouped into cadences via
1797
+ plan_cadences() and processed sequentially over one persistent worker pool. Within each
1798
+ cadence, energy detection runs on ON-source files (positions 0, 2, 4 in ABACAD); stamps
1799
+ are then extracted from all 6 observations at the hit frequencies. Each cadence is
1800
+ checkpointed to disk as soon as its .npy is ready, and on retry, cadences whose output
1801
+ already exists are skipped.
1802
+
1803
+ The streaming inference path in main.py drives plan_cadences() /
1804
+ process_pending_cadence() directly instead (so preprocessing of cadence i+1 can overlap
1805
+ inference of cadence i); this wrapper preserves the batch contract for callers that
1806
+ want every cadence preprocessed up front.
1807
+ """
1808
+ units = self.plan_cadences()
1809
+ if not units:
1810
+ return []
1811
+
1812
+ results: list[CadenceResult] = []
1813
+ self.start_energy_detection_pool()
1814
+ try:
1815
+ for unit in units:
1816
+ cadence_result = self.process_pending_cadence(unit)
1817
+ if cadence_result is not None:
1818
+ results.append(cadence_result)
1819
+ finally:
1820
+ self.stop_energy_detection_pool()
1821
+
1822
+ logger.info(f"find_hits completed: {len(results)} cadence .npy files available")
1823
+ return results
1824
+
1825
+ # NOTE: come back to this later (ensure filenames are structured similarly as in train_files & test_files)
1826
+ @staticmethod
1827
+ def _cadence_npy_filename(csv_stem: str, key: tuple) -> str:
1828
+ """Build a deterministic filename for a cadence group's .npy output."""
1829
+ # Sanitize each key component via an allowlist: only word chars, dash, and
1830
+ # dot survive — everything else collapses to underscore. This is broader
1831
+ # than just stripping path separators / whitespace: CSV column values can
1832
+ # carry quotes, control chars, shell metacharacters, or path-traversal
1833
+ # sequences (e.g. "..") that would otherwise leak through.
1834
+ safe_parts = [re.sub(r"[^\w\-.]+", "_", str(part)) for part in key]
1835
+ return f"{csv_stem}_{'_'.join(safe_parts)}.npy"
1836
+
1837
+ @staticmethod
1838
+ def cadence_metadata_path(npy_path: str) -> str:
1839
+ """Return the sibling .json path for a cadence's metadata. Public: the streaming
1840
+ loop (main.py) and the viz suite derive metadata paths for resume-skipped cadences
1841
+ whose CadenceResult was never rebuilt."""
1842
+ return os.path.splitext(npy_path)[0] + ".json"
1843
+
1844
+ @staticmethod
1845
+ def _to_json_safe(obj):
1846
+ """
1847
+ Coerce h5py / numpy values into JSON-native types.
1848
+
1849
+ h5py attributes can be bytes, numpy scalars, or numpy arrays — none of
1850
+ which json.dump handles by default. Walk the structure once and
1851
+ convert leaf nodes; everything else passes through.
1852
+ """
1853
+ if isinstance(obj, dict):
1854
+ return {str(k): DataPreprocessor._to_json_safe(v) for k, v in obj.items()}
1855
+ if isinstance(obj, (list, tuple)):
1856
+ return [DataPreprocessor._to_json_safe(v) for v in obj]
1857
+ if isinstance(obj, bytes):
1858
+ return obj.decode("utf-8", errors="replace")
1859
+ if isinstance(obj, np.ndarray):
1860
+ return DataPreprocessor._to_json_safe(obj.tolist())
1861
+ if isinstance(obj, (np.integer, np.floating, np.bool_)):
1862
+ return obj.item()
1863
+ return obj
1864
+
1865
+ # NOTE: come back to this later (does a cadence snippet get created if only 1 of the ON observations crosses the threshold, or all 3? should probably be all 3 as of now since models are trained on signals that don't yet drift out of frame?)
1866
+ def _process_cadence(self, group: CadenceGroup, npy_path: str) -> CadenceResult | None:
1867
+ """
1868
+ Run energy detection on one cadence and write its stamp .npy at the given absolute
1869
+ npy_path. Returns a CadenceResult on success, or None if no hits survived.
1870
+
1871
+ Energy detection runs only on ON-source observations (positions 0, 2, 4 in ABACAD order);
1872
+ the hits found there define the frequency slices that get extracted from all 6
1873
+ observations. group is assumed to be a validated CadenceGroup with len(h5_paths) ==
1874
+ expected_obs. Each coarse channel is one fused task (read -> DC spike -> bandpass
1875
+ flatten -> vectorized threshold) on the persistent pool from
1876
+ start_energy_detection_pool(); only the small per-channel hit lists cross the process
1877
+ boundary, so no blocks, block-sized shared memory, or per-stage pools exist anymore.
1878
+ """
1879
+ coarse_channel_width = self.config.inference.coarse_channel_width
1880
+ # coarse_channel_log_interval is a progress-logging chunk size only (None -> n_processes);
1881
+ # actual parallelism is the persistent pool's worker count.
1882
+ log_interval = self.config.inference.coarse_channel_log_interval
1883
+ window_size = self.config.inference.detection_window_size
1884
+ step_size = self.config.inference.detection_step_size
1885
+ stat_threshold = self.config.inference.stat_threshold
1886
+ stamp_width = self.config.inference.stamp_width
1887
+ overlap_search = self.config.inference.overlap_search
1888
+ overlap_fraction = self.config.inference.overlap_fraction
1889
+ store_downsampled = self.config.inference.store_downsampled_stamps
1890
+ downsample_factor = self.config.data.downsample_factor if store_downsampled else 1
1891
+ time_bins = self.config.data.time_bins
1892
+ n_processes = self.config.manager.n_processes
1893
+
1894
+ # Read header / metadata from the first ON-source file
1895
+ on_source_paths = [group.h5_paths[i] for i in (0, 2, 4)]
1896
+ primary_h5 = on_source_paths[0]
1897
+
1898
+ with h5py.File(primary_h5, "r") as hf:
1899
+ header = {k: hf["data"].attrs[k] for k in hf["data"].attrs}
1900
+ data_shape = hf["data"].shape
1901
+
1902
+ # Up-front geometry validation of ALL 6 observation files, not just the primary.
1903
+ # Energy detection and stamp extraction index every dataset as [time, polarization,
1904
+ # frequency] and read exactly the first time_bins rows (see
1905
+ # _energy_detect_channel_worker / _extract_stamps_worker — extra rows beyond
1906
+ # time_bins are simply ignored), so a malformed file anywhere in the cadence would
1907
+ # otherwise fail deep inside a worker: a wrong-rank dataset mis-slices cryptically,
1908
+ # and a short row count in ANY of the 6 files (ON or OFF) raises a broadcast
1909
+ # ValueError mid-extraction when the short stamp is assigned into its fixed
1910
+ # (n_stamps, 6, time_bins, stored_width) memmap slot — a short ON file additionally
1911
+ # degrades the k2 statistic silently first, since _sliding_normality_k2 derives its
1912
+ # sample count from the rows it receives. Frequency WIDTH is validated for the same
1913
+ # reason: extraction reads the SAME [0:n_chans] channel window from every file, so a
1914
+ # file with FEWER than n_chans channels truncates into a short stamp (the same broadcast
1915
+ # ValueError deep in a worker). The check is `< n_chans` (not strict equality): a file at
1916
+ # least n_chans wide reads cleanly and processed fine on master, so it is left alone —
1917
+ # this converts only the crash case into a clean up-front skip. The whole cadence is
1918
+ # skipped rather than just the offending file: the 6-observation stamp tensor and the
1919
+ # num_observations contract downstream (encoder reshape, RF features, ABACAD viz) have no
1920
+ # representation for a 5-observation cadence.
1921
+ #
1922
+ # n_chans (header channel count, falling back to the data width) is that read window and
1923
+ # the reference for the check; computed before the loop so the geometry pass can use it.
1924
+ n_chans = int(header.get("nchans", data_shape[-1]))
1925
+ rank_problems: list[str] = []
1926
+ short_problems: list[str] = []
1927
+ width_problems: list[str] = []
1928
+ for idx, obs_h5 in enumerate(group.h5_paths):
1929
+ if idx == 0:
1930
+ # group.h5_paths[0] == primary_h5, already opened above for header/data_shape —
1931
+ # reuse it instead of a redundant second h5py.File open.
1932
+ obs_shape = data_shape
1933
+ else:
1934
+ with h5py.File(obs_h5, "r") as hf:
1935
+ obs_shape = hf["data"].shape
1936
+ if len(obs_shape) != 3:
1937
+ rank_problems.append(
1938
+ f"{obs_h5} has 'data' of rank {len(obs_shape)} (shape {obs_shape})"
1939
+ )
1940
+ continue
1941
+ if int(obs_shape[0]) < time_bins:
1942
+ short_problems.append(f"{obs_h5} has only {int(obs_shape[0])} time bins")
1943
+ if int(obs_shape[-1]) < n_chans:
1944
+ width_problems.append(
1945
+ f"{obs_h5} has {int(obs_shape[-1])} frequency channels < the {n_chans}-channel "
1946
+ f"read window"
1947
+ )
1948
+ if rank_problems:
1949
+ # The caller (process_pending_cadence) turns this ValueError into a logged skip,
1950
+ # so the skip-and-continue policy across a large catalog is preserved.
1951
+ raise ValueError(
1952
+ f"Cadence {group.key}: expected rank 3 (time, polarization, frequency) "
1953
+ f"'data' in every observation file; offending file(s): "
1954
+ f"{'; '.join(rank_problems)}"
1955
+ )
1956
+ if short_problems:
1957
+ logger.warning(
1958
+ f"Cadence {group.key}: every observation file needs >= {time_bins} time "
1959
+ f"bins; skipping whole cadence — offending file(s): "
1960
+ f"{'; '.join(short_problems)}"
1961
+ )
1962
+ return None
1963
+ if width_problems:
1964
+ logger.warning(
1965
+ f"Cadence {group.key}: every observation file needs >= {n_chans} frequency "
1966
+ f"channels (the window read from all 6 files); skipping whole cadence — "
1967
+ f"offending file(s): {'; '.join(width_problems)}"
1968
+ )
1969
+ return None
1970
+
1971
+ foff = float(header["foff"])
1972
+ fch1 = float(header["fch1"])
1973
+
1974
+ # NOTE: every complete coarse channel is processed. The historical block-based path
1975
+ # floored to a multiple of the old parallel_coarse_chans knob, silently dropping up to
1976
+ # that many - 1 trailing coarse channels when n_chans wasn't an exact multiple of a
1977
+ # block. (That knob is now coarse_channel_log_interval and only affects log cadence.)
1978
+ n_coarse_total = n_chans // coarse_channel_width
1979
+ if n_coarse_total == 0:
1980
+ logger.warning(
1981
+ f"Cadence {group.key}: n_chans={n_chans} is smaller than one coarse channel "
1982
+ f"({coarse_channel_width}); skipping"
1983
+ )
1984
+ return None
1985
+
1986
+ logger.info(
1987
+ f"Cadence {group.key}: n_chans={n_chans}, coarse channels={n_coarse_total}, "
1988
+ f"ON-source files={len(on_source_paths)}"
1989
+ )
1990
+
1991
+ # The flattener depends on the file's actual coarse count (PFB folds
1992
+ # adjacent-channel leakage), so it can only be built once n_coarse_total is known.
1993
+ bandpass_flatten = self._get_bandpass_flattener(n_coarse_total)
1994
+ pfb_active = bandpass_flatten.func is _pfb_flatten_bandpass
1995
+
1996
+ if self.config.inference.bandpass_debug_plot:
1997
+ try:
1998
+ self._plot_bandpass_overlay(primary_h5, n_coarse_total, bandpass_flatten, npy_path)
1999
+ except Exception as e:
2000
+ logger.error(f"Cadence {group.key}: bandpass overlay plot failed: {e}")
2001
+
2002
+ cadence_start_time = time.time()
2003
+
2004
+ # Aggregate hits across all ON-source files, plus the per-ON-file summary histogram
2005
+ # of every window statistic (for the ed_stat_distributions figure)
2006
+ all_hits: list[tuple] = [] # (abs_idx, stat, p)
2007
+ stat_hists = np.zeros((len(on_source_paths), len(ED_STAT_HIST_EDGES) - 1), dtype=np.int64)
2008
+
2009
+ # One fused, ORDERED, file-major imap over all three ON files (#298 I6): the old
2010
+ # one-drained-imap-per-file shape paid three straggler tails per cadence and left
2011
+ # the pool idle between files. Ordered iteration makes all_hits assembly
2012
+ # byte-identical to the per-file loops (file 0's channels, then file 1's, ...),
2013
+ # and the task position recovers the ON-file index for the per-file histograms.
2014
+ # The per-file read_ed_on{1..3} spans collapse into one read_ed span.
2015
+ logger.info(
2016
+ f"Cadence {group.key}: running energy detection on "
2017
+ f"{len(on_source_paths)} ON-source files x {n_coarse_total} coarse channels"
2018
+ )
2019
+ # Progress-log points per ON file (#301): an explicit coarse_channel_log_interval
2020
+ # keeps the historical every-N-channels cadence; the None default logs at ~25%
2021
+ # milestones instead of every n_processes channels — the per-channel progress
2022
+ # lines were a measured 62% of a run's total log volume, all Slack-bound
2023
+ # (~594k lines over a full catalog).
2024
+ if log_interval is not None:
2025
+ step = max(1, log_interval)
2026
+ progress_points = set(range(step, n_coarse_total + 1, step)) | {n_coarse_total}
2027
+ else:
2028
+ progress_points = {max(1, (n_coarse_total * q) // 4) for q in (1, 2, 3, 4)}
2029
+ with stage_timer("read_ed"):
2030
+ # Residual-flatness sanity check — primary ON file only (the static response is
2031
+ # a property of the backend, shared by every file). Runs ON THE POOL alongside
2032
+ # energy detection (#298 I6): the serial parent-side version idled all workers
2033
+ # for its ~27 chunk-stripe band reads; the median + warning still land once per
2034
+ # cadence, after the ED drain below.
2035
+ pfb_check = (
2036
+ self._start_pfb_response_check(primary_h5, n_coarse_total, bandpass_flatten)
2037
+ if pfb_active
2038
+ else None
2039
+ )
2040
+
2041
+ # The few channels whose despiked integrated spectra feed the persisted
2042
+ # bandpass envelopes (#301) — primary ON file only, same sampling as the
2043
+ # figure this replaces; the spectra ride back with the ED results, so no
2044
+ # extra h5 reads happen at all
2045
+ envelope_channels = set(self._sample_channel_indices(n_coarse_total))
2046
+ tasks = [
2047
+ (
2048
+ on_h5,
2049
+ ch,
2050
+ coarse_channel_width,
2051
+ time_bins,
2052
+ bandpass_flatten,
2053
+ window_size,
2054
+ step_size,
2055
+ stat_threshold,
2056
+ on_h5 == primary_h5 and ch in envelope_channels,
2057
+ )
2058
+ for on_h5 in on_source_paths
2059
+ for ch in range(n_coarse_total)
2060
+ ]
2061
+
2062
+ if self._ed_pool is not None:
2063
+ # imap (ordered, chunksize 1) keeps every worker busy across the whole
2064
+ # cadence while results stream back for progress logging
2065
+ channel_hits_iter = self._ed_pool.imap(_energy_detect_channel_worker, tasks)
2066
+ else:
2067
+ if n_processes > 1:
2068
+ logger.info(
2069
+ "Energy detection running sequentially: no persistent pool "
2070
+ "started (call start_energy_detection_pool() to parallelize)"
2071
+ )
2072
+ channel_hits_iter = map(_energy_detect_channel_worker, tasks)
2073
+
2074
+ integrated_by_channel: dict[int, np.ndarray] = {}
2075
+ for done, (channel_hits, channel_hist, integrated) in enumerate(channel_hits_iter):
2076
+ on_source_idx = done // n_coarse_total
2077
+ file_done = done % n_coarse_total + 1
2078
+ all_hits.extend(channel_hits)
2079
+ stat_hists[on_source_idx] += channel_hist
2080
+ if integrated is not None:
2081
+ integrated_by_channel[done % n_coarse_total] = integrated
2082
+ if file_done in progress_points:
2083
+ logger.info(
2084
+ f" Coarse channel {file_done}/{n_coarse_total} of ON-source "
2085
+ f"{on_source_idx + 1}/{len(on_source_paths)}"
2086
+ )
2087
+
2088
+ if pfb_check is not None:
2089
+ self._finish_pfb_response_check(pfb_check, primary_h5)
2090
+
2091
+ # Persisted bandpass envelopes (#301): decimated raw/flattened/overlay lines per
2092
+ # sampled channel, computed from the spectra that rode back with ED. Viz-only —
2093
+ # a failure degrades one figure, never the cadence.
2094
+ bandpass_envelopes = None
2095
+ try:
2096
+ bandpass_envelopes = self._build_bandpass_envelopes(
2097
+ integrated_by_channel, bandpass_flatten, pfb_active
2098
+ )
2099
+ except Exception as e:
2100
+ logger.warning(f"Cadence {group.key}: bandpass envelope computation failed ({e})")
2101
+
2102
+ logger.info(f"Cadence {group.key}: {len(all_hits)} raw hits across ON-source files")
2103
+
2104
+ # NOTE: come back to this later (what's the trade-off for doing dedup vs not? e.g. lower storage & compute, but higher FNR or lower DR sensitivity?)
2105
+ # Deduplicate: greedy merge of any pair within stamp_width // 2
2106
+ with stage_timer("dedup"):
2107
+ merged_hits = self._deduplicate_hits(all_hits, stamp_width)
2108
+ logger.info(
2109
+ f"Cadence {group.key}: {len(merged_hits)} hits after deduplication "
2110
+ f"(stamp_width={stamp_width})"
2111
+ )
2112
+
2113
+ if not merged_hits:
2114
+ logger.info(f"Cadence {group.key}: no hits survived; skipping")
2115
+ return None
2116
+
2117
+ # Build the stamp centers (optionally including overlap offsets)
2118
+ stamp_centers: list[tuple[int, float, float]] = []
2119
+ offsets = [0]
2120
+ if overlap_search:
2121
+ offset_mag = int(overlap_fraction * stamp_width)
2122
+ offsets = [-offset_mag, 0, offset_mag]
2123
+
2124
+ half = stamp_width // 2
2125
+ for hit in merged_hits:
2126
+ abs_idx, stat, pval = hit
2127
+ for offset in offsets:
2128
+ center = abs_idx + offset
2129
+ start = center - half
2130
+ end = center + half
2131
+ if start < 0 or end > n_chans:
2132
+ continue
2133
+ stamp_centers.append((start, stat, pval))
2134
+
2135
+ if not stamp_centers:
2136
+ logger.info(f"Cadence {group.key}: no valid in-bounds stamps; skipping")
2137
+ return None
2138
+
2139
+ # Sort stamps by start index before extraction. With overlap_search the raw
2140
+ # order is hit-interleaved (hit1-off, hit1, hit1+off, hit2-off, ...), so
2141
+ # neighboring hits can produce out-of-order starts. Reads against
2142
+ # bitshuffle-compressed .h5 files are dominated by chunk decompression cost;
2143
+ # sequential start order gives the OS / hdf5 chunk cache a chance to reuse a
2144
+ # decompressed chunk across adjacent stamps instead of redecompressing it.
2145
+ stamp_centers.sort(key=lambda s: s[0])
2146
+
2147
+ # Extract stamps into a memmap-backed .npy so worker processes can fill disjoint
2148
+ # (obs_file, stamp-range) slices in parallel. The previous sequential per-file loop
2149
+ # over all 6 observations (single-threaded reads + bitshuffle chunk decompression)
2150
+ # was the dominant, GPU-idle cost of CSV inference.
2151
+ #
2152
+ # Stamps are downsampled along frequency at extraction time (downsample_factor > 1,
2153
+ # the store_downsampled_stamps default), so the stored width is stamp_width //
2154
+ # downsample_factor — an ~8x storage cut at defaults that also removes the separate
2155
+ # downsample pass from load_inference_data.
2156
+ #
2157
+ # Atomicity is unchanged: the memmap is written to a .tmp sibling and we os.replace()
2158
+ # it onto the canonical name only after every worker finishes, so the resume path
2159
+ # (which treats npy_path's existence as proof of a complete write) still holds.
2160
+ n_stamps = len(stamp_centers)
2161
+ stored_width = stamp_width // downsample_factor
2162
+ # Per-run-unique tmp name (#298 I3): the default output dir is now shared across
2163
+ # runs with the same ED config, and a single fixed "<stem>.tmp.npy" would let two
2164
+ # concurrent runs truncate/delete each other's in-progress write. Content is
2165
+ # deterministic given the ED config, so whichever writer publishes last is
2166
+ # harmless; os.replace stays atomic.
2167
+ tmp_npy_path = f"{os.path.splitext(npy_path)[0]}.{os.getpid()}.{uuid.uuid4().hex}.tmp.npy"
2168
+ # A leftover per-cadence tmp means an attempt died (e.g. SIGKILL) between creation
2169
+ # and os.replace; it was never promoted, so it is safe to drop once clearly
2170
+ # abandoned. Only age-expired tmps are removed — a fresh one may be a live
2171
+ # concurrent run's in-progress write. The glob "<stem>.*.tmp" sweeps every
2172
+ # per-cadence tmp uniformly: the stamp memmap "<stem>.<pid>.<hex>.tmp.npy", the
2173
+ # metadata "<stem>.json.<pid>.<hex>.tmp", and the #302 candidate sidecar
2174
+ # "<stem>.candidates.npz.<pid>.<hex>.tmp" (the latter two ended in ".tmp" but not
2175
+ # ".tmp.npy", so the old ".*.tmp.npy" glob leaked them). The fixed pre-#298
2176
+ # "<stem>.tmp.npy" legacy name is swept too.
2177
+ stale_cutoff = time.time() - _STALE_TMP_MAX_AGE_S
2178
+ stem = os.path.splitext(npy_path)[0]
2179
+ escaped = glob.escape(stem)
2180
+ stale_tmps = (
2181
+ glob.glob(f"{escaped}.*.tmp.npy") # stamp memmap tmps
2182
+ + glob.glob(f"{escaped}.*.tmp") # metadata + candidate-sidecar tmps
2183
+ + [f"{stem}.tmp.npy"] # fixed pre-#298 legacy name
2184
+ )
2185
+ for stale in stale_tmps:
2186
+ with contextlib.suppress(OSError):
2187
+ if os.path.getmtime(stale) < stale_cutoff:
2188
+ logger.warning(
2189
+ f"Cadence {group.key}: removing stale partial output {stale} "
2190
+ f"from an interrupted previous run"
2191
+ )
2192
+ os.remove(stale)
2193
+ # NOTE: np.lib.format.open_memmap is a semi-public numpy API (stable across 1.x and
2194
+ # documented via np.lib.format); revisit if a future numpy bump moves it
2195
+ memmap = np.lib.format.open_memmap(
2196
+ tmp_npy_path,
2197
+ mode="w+",
2198
+ dtype=np.float32,
2199
+ shape=(n_stamps, len(group.h5_paths), time_bins, stored_width),
2200
+ )
2201
+ memmap.flush()
2202
+ del memmap # header + full-size file are on disk; workers reopen it in r+ mode
2203
+
2204
+ stamp_starts = [start for start, _, _ in stamp_centers]
2205
+ # Split each obs file's (already start-sorted) stamps into contiguous chunks so more
2206
+ # than len(h5_paths) workers can run, while keeping each worker's reads sequential to
2207
+ # preserve the hdf5 chunk-cache reuse that the sort above buys us. Oversubscribe to
2208
+ # ~_EXTRACT_TASKS_PER_WORKER tasks per worker so the stage's final wave is a fraction
2209
+ # of one task instead of doubling the makespan (see the constant's rationale).
2210
+ chunks_per_file = max(
2211
+ 1, -(-(_EXTRACT_TASKS_PER_WORKER * n_processes) // len(group.h5_paths))
2212
+ ) # ceil div
2213
+ chunk_size = max(1, -(-n_stamps // chunks_per_file)) # ceil div
2214
+ # Chunk-cache kwargs are a pure function of each file's layout — computed once per
2215
+ # obs file here instead of once per task in the workers (#301: ~2 redundant
2216
+ # network-FS open/inspect cycles x ~132 tasks per cadence; the same hoist the
2217
+ # PFB-check tasks already used)
2218
+ cache_kwargs = {obs_h5: _chunk_cache_kwargs(obs_h5, time_bins) for obs_h5 in group.h5_paths}
2219
+ tasks = [
2220
+ (
2221
+ tmp_npy_path,
2222
+ obs_idx,
2223
+ obs_h5,
2224
+ stamp_starts[base : base + chunk_size],
2225
+ base,
2226
+ time_bins,
2227
+ stamp_width,
2228
+ downsample_factor,
2229
+ cache_kwargs[obs_h5],
2230
+ )
2231
+ for obs_idx, obs_h5 in enumerate(group.h5_paths)
2232
+ for base in range(0, n_stamps, chunk_size)
2233
+ ]
2234
+
2235
+ with stage_timer("extract"):
2236
+ if self._ed_pool is not None and len(tasks) > 1:
2237
+ # Reuse the persistent energy-detection pool — extraction workers are plain
2238
+ # (no shared memory), so the same pool serves both stages without churn.
2239
+ # imap_unordered: workers write disjoint absolute rows (base_idx offsets),
2240
+ # so completion order is irrelevant and stragglers backfill immediately.
2241
+ for _ in self._ed_pool.imap_unordered(_extract_stamps_worker, tasks):
2242
+ pass
2243
+ else:
2244
+ for task in tasks:
2245
+ _extract_stamps_worker(task)
2246
+
2247
+ os.replace(tmp_npy_path, npy_path)
2248
+
2249
+ metadata_path = self.cadence_metadata_path(npy_path)
2250
+ # Per-stamp frequency = center bin's frequency, computed from header's fch1/foff
2251
+ stamp_freqs_mhz = [float(fch1 + foff * (start + half)) for start, _, _ in stamp_centers]
2252
+ stamp_stats = [float(s) for _, s, _ in stamp_centers]
2253
+ stamp_pvals = [float(p) for _, _, p in stamp_centers]
2254
+
2255
+ # Pre-binned hit-frequency histograms (#301): a fixed per-cadence grid spanning
2256
+ # the file band, fine enough (_HIT_HIST_BINS) that the viz suite's rebinning onto
2257
+ # its global axis is visually exact. The exact raw-hit min/max ride along so the
2258
+ # figure reproduces the historical axis bounds. merged_hits is non-empty here and
2259
+ # all_hits ⊇ merged_hits, so the min/max are well-defined.
2260
+ raw_freq_values = np.array([fch1 + foff * idx for idx, _, _ in all_hits])
2261
+ merged_freq_values = np.array([fch1 + foff * idx for idx, _, _ in merged_hits])
2262
+ band_lo = min(fch1, fch1 + foff * n_chans)
2263
+ band_hi = max(fch1, fch1 + foff * n_chans)
2264
+ band_edges = np.linspace(band_lo, band_hi, _HIT_HIST_BINS + 1)
2265
+ hit_spectrum_hist = {
2266
+ "freq_lo": float(band_lo),
2267
+ "freq_hi": float(band_hi),
2268
+ "n_bins": _HIT_HIST_BINS,
2269
+ "raw_counts": np.histogram(raw_freq_values, bins=band_edges)[0].tolist(),
2270
+ "merged_counts": np.histogram(merged_freq_values, bins=band_edges)[0].tolist(),
2271
+ "raw_freq_min": float(raw_freq_values.min()),
2272
+ "raw_freq_max": float(raw_freq_values.max()),
2273
+ }
2274
+
2275
+ metadata = {
2276
+ "key": group.key,
2277
+ "csv_path": group.csv_path,
2278
+ "h5_paths": group.h5_paths,
2279
+ "header": header,
2280
+ "stamp_starts": [int(start) for start, _, _ in stamp_centers],
2281
+ "stamp_width": stamp_width,
2282
+ "stored_width": stored_width,
2283
+ "downsample_factor_applied": downsample_factor,
2284
+ "stamp_frequencies_mhz": stamp_freqs_mhz,
2285
+ "stamp_statistics": stamp_stats,
2286
+ "stamp_pvalues": stamp_pvals,
2287
+ "overlap_search": overlap_search,
2288
+ "overlap_fraction": overlap_fraction if overlap_search else None,
2289
+ # Energy-detection provenance for the visualization suite: the all-window
2290
+ # statistic histograms (per ON file, fixed log-spaced bins) and the pre-binned
2291
+ # hit-frequency histograms (hit spectrum + funnel figures). Raw per-hit
2292
+ # frequency lists are no longer stored (#301 — this discharges the old
2293
+ # pre-binning NOTE): they ran ~19 MB of JSON on an RFI-dense cadence and had
2294
+ # exactly one consumer, the hit-spectrum figure, which now rebins
2295
+ # hit_spectrum_hist instead. The post-dedup merged list stays (small, the
2296
+ # scientifically meaningful set); raw counts survive in n_raw_hits + the
2297
+ # histogram.
2298
+ "ed_stat_hist": {
2299
+ "bin_edges": [float(e) for e in ED_STAT_HIST_EDGES],
2300
+ "counts_per_on_file": stat_hists.tolist(),
2301
+ },
2302
+ "n_raw_hits": len(all_hits),
2303
+ "n_merged_hits": len(merged_hits),
2304
+ "hit_spectrum_hist": hit_spectrum_hist,
2305
+ "merged_hit_frequencies_mhz": [float(fch1 + foff * idx) for idx, _, _ in merged_hits],
2306
+ "bandpass_envelopes": bandpass_envelopes,
2307
+ # The ED config these stamps were produced under (#298 I3) — checked by the
2308
+ # resume guard before an existing .npy is reused
2309
+ "ed_config_fingerprint": preprocessing_config_fingerprint(self.config.to_dict()),
2310
+ }
2311
+ # Per-run-unique tmp + compact separators (#298 I3 + rider): the metadata sidecar
2312
+ # shares the cross-run cache directory (same clobber race as the .npy tmp), and
2313
+ # indent=2 over the 1e4-1e5-float hit lists cost ~0.5-1.5 s per cadence ON THE
2314
+ # PREFETCH CRITICAL PATH for purely cosmetic whitespace.
2315
+ tmp_metadata_path = f"{metadata_path}.{os.getpid()}.{uuid.uuid4().hex}.tmp"
2316
+ with open(tmp_metadata_path, "w") as f:
2317
+ json.dump(self._to_json_safe(metadata), f, separators=(",", ":"))
2318
+ os.replace(tmp_metadata_path, metadata_path)
2319
+
2320
+ # NOTE: the full gc.collect() that used to run here sat on the PREFETCH thread's
2321
+ # critical path (once per cadence, ~0.3 s with TF loaded, GIL held against the
2322
+ # encode feed) and freed nothing refcounting doesn't (#298 I7)
2323
+
2324
+ # Record the stage transition in the inference_cadences run manifest. Written only
2325
+ # when preprocessing actually ran (the resume path never re-writes it); a crash
2326
+ # between the os.replace above and this write just loses an informational row —
2327
+ # resume keys off the .npy's existence, and the 'inferred' row keys off inference.
2328
+ self.db.write_inference_cadence(
2329
+ npy_path=npy_path,
2330
+ status="preprocessed",
2331
+ tag=self.config.checkpoint.save_tag,
2332
+ csv_path=group.csv_path,
2333
+ cadence_key=group.key,
2334
+ n_stamps=n_stamps,
2335
+ duration_s=time.time() - cadence_start_time,
2336
+ )
2337
+
2338
+ logger.info(
2339
+ f"Cadence {group.key}: wrote {n_stamps} stamps -> "
2340
+ f"{npy_path} (metadata: {metadata_path})"
2341
+ )
2342
+
2343
+ # Record for the disk-leak guard: a later retry attempt that resumes this .npy
2344
+ # (kept because this attempt's inference failed) is still this run's work to prune
2345
+ self._extracted_this_run.add(npy_path)
2346
+ return CadenceResult(
2347
+ npy_path=npy_path,
2348
+ h5_paths=group.h5_paths,
2349
+ key=group.key,
2350
+ n_hits=n_stamps,
2351
+ metadata_path=metadata_path,
2352
+ freshly_extracted=True,
2353
+ )
2354
+
2355
+ def _get_bandpass_flattener(
2356
+ self, num_coarse_channels: int
2357
+ ) -> Callable[[np.ndarray], np.ndarray]:
2358
+ """
2359
+ Return the configured bandpass-flattening callable for energy detection.
2360
+
2361
+ The callable takes one coarse channel of shape (time_bins, coarse_channel_width) and
2362
+ returns the flattened channel; it must be picklable (a functools.partial over a
2363
+ module-level function) so pool workers can receive it in their task args.
2364
+
2365
+ num_coarse_channels is the *file's* actual coarse-channel count
2366
+ (n_chans // coarse_channel_width) — the PFB response folds adjacent-channel leakage,
2367
+ so it depends on how many coarse channels the recording actually has. Files with a
2368
+ single coarse channel can't support the fold and fall back to the spline flattener
2369
+ with a warning.
2370
+ """
2371
+ method = self.config.inference.bandpass_method
2372
+ if method == "pfb":
2373
+ if num_coarse_channels >= 2:
2374
+ response_path = self._ensure_pfb_response_file(
2375
+ self.config.inference.coarse_channel_width,
2376
+ num_coarse_channels,
2377
+ self.config.inference.pfb_taps_per_channel,
2378
+ )
2379
+ return functools.partial(_pfb_flatten_bandpass, response_path=response_path)
2380
+ logger.warning(
2381
+ "bandpass_method='pfb' requires >= 2 coarse channels to fold adjacent-channel "
2382
+ f"leakage, but this file has {num_coarse_channels}; falling back to the "
2383
+ "spline flattener for this cadence"
2384
+ )
2385
+ elif method != "spline":
2386
+ raise ValueError(f"Unknown bandpass_method {method!r}; expected 'pfb' or 'spline'")
2387
+ return functools.partial(
2388
+ _spline_flatten_bandpass, spl_order=self.config.inference.spline_order
2389
+ )
2390
+
2391
+ def _ensure_pfb_response_file(
2392
+ self, fine_per_coarse: int, num_coarse_channels: int, taps_per_channel: int
2393
+ ) -> str:
2394
+ """
2395
+ Compute the PFB passband response in the parent process and persist it to a
2396
+ deterministic sidecar .npy under {output_path}/cache/pfb/, returning its path.
2397
+
2398
+ The heavy work (an ~n_chans-point FFT) runs exactly once per parameter combination in
2399
+ the parent — gen_coarse_channel_response is process-cached and the file is reused when
2400
+ its content matches — while pool workers receive only the path and read the ~8 MB
2401
+ array (see _pfb_flatten_bandpass). The file is content-addressed by its parameters, so
2402
+ stale-run leftovers are impossible; a corrupt or mismatched file is rewritten. Writes
2403
+ are atomic (tmp + os.replace), matching the stamp-extraction pattern.
2404
+ """
2405
+ # NOTE: deliberately NOT memoized (#298 considered and rejected it): the per-cadence
2406
+ # np.load + np.array_equal round trip (~15 ms) is what makes the documented
2407
+ # "a corrupt or mismatched sidecar is transparently rewritten" contract hold
2408
+ # WITHIN a run, and the saving is sub-noise against a multi-minute cadence.
2409
+ response = gen_coarse_channel_response(
2410
+ fine_per_coarse, num_coarse_channels, taps_per_channel
2411
+ )
2412
+
2413
+ cache_dir = os.path.join(self.config.output_path, "cache", "pfb")
2414
+ os.makedirs(cache_dir, exist_ok=True)
2415
+ path = os.path.join(
2416
+ cache_dir,
2417
+ f"pfb_response_w{fine_per_coarse}_c{num_coarse_channels}_t{taps_per_channel}.npy",
2418
+ )
2419
+
2420
+ if os.path.exists(path):
2421
+ try:
2422
+ existing = np.load(path)
2423
+ if np.array_equal(existing, response):
2424
+ return path
2425
+ logger.warning(f"PFB response cache {path} does not match; rewriting")
2426
+ except Exception as e:
2427
+ logger.warning(f"PFB response cache {path} unreadable ({e}); rewriting")
2428
+
2429
+ # Per-writer tmp name (pid + uuid) so two runs sharing this output_path don't clobber
2430
+ # each other's in-progress write on a single shared "{path}.tmp"; os.replace stays
2431
+ # atomic and the content is deterministic, so whichever writer lands last is harmless.
2432
+ tmp_path = f"{path}.{os.getpid()}.{uuid.uuid4().hex}.tmp"
2433
+ with open(tmp_path, "wb") as f:
2434
+ np.save(f, response)
2435
+ os.replace(tmp_path, path)
2436
+ logger.info(f"Wrote PFB response cache: {path}")
2437
+ return path
2438
+
2439
+ @staticmethod
2440
+ def _sample_channel_indices(
2441
+ n_coarse_total: int, num: int = _BANDPASS_SAMPLE_CHANNELS
2442
+ ) -> list[int]:
2443
+ """Pick up to `num` coarse-channel indices evenly across the band."""
2444
+ num = min(num, n_coarse_total)
2445
+ return sorted({int(i) for i in np.linspace(0, n_coarse_total - 1, num=num)})
2446
+
2447
+ def _read_despiked_channel(self, h5_path: str, channel_index: int) -> np.ndarray:
2448
+ """Read one coarse channel as float64 with the DC spike interpolated away — the same
2449
+ preparation the energy-detection workers apply before bandpass flattening."""
2450
+ width = self.config.inference.coarse_channel_width
2451
+ time_bins = self.config.data.time_bins
2452
+ start = channel_index * width
2453
+ with h5py.File(h5_path, "r") as hf:
2454
+ channel = hf["data"][:time_bins, 0, start : start + width].astype(np.float64)
2455
+ _remove_dc_spike(channel, width, 1)
2456
+ return channel
2457
+
2458
+ def _flattened_edge_mid_ratio(
2459
+ self,
2460
+ h5_path: str,
2461
+ channel_index: int,
2462
+ response: np.ndarray,
2463
+ open_kwargs: dict | None = None,
2464
+ ) -> float:
2465
+ """Config-bound wrapper over _flattened_edge_mid_ratio_banded (the module-level body
2466
+ is shared with the pool-side _pfb_ratio_worker, #298 I6)."""
2467
+ return _flattened_edge_mid_ratio_banded(
2468
+ h5_path,
2469
+ channel_index,
2470
+ response,
2471
+ self.config.inference.coarse_channel_width,
2472
+ self.config.data.time_bins,
2473
+ open_kwargs,
2474
+ )
2475
+
2476
+ def _build_bandpass_envelopes(
2477
+ self,
2478
+ integrated_by_channel: dict[int, np.ndarray],
2479
+ bandpass_flatten: functools.partial,
2480
+ pfb_active: bool,
2481
+ ) -> list[dict] | None:
2482
+ """Decimated raw/flattened/overlay integrated-spectrum lines for the sampled coarse
2483
+ channels, persisted in the metadata sidecar (#301) so the bandpass-flattening viz
2484
+ figure renders from ~KB of stored points instead of re-reading (cold) coarse
2485
+ channels from the .h5 at viz time (measured 114 s = 74% of a cached rerun's viz
2486
+ span). Exact w.r.t. the figure's own math: the PFB divide and the spline
2487
+ subtraction both commute with the time mean, and the workers integrate in float64
2488
+ over the same despiked channel the figure would read."""
2489
+ if not integrated_by_channel:
2490
+ return None
2491
+ envelopes: list[dict] = []
2492
+ for ch in sorted(integrated_by_channel):
2493
+ raw_int = np.asarray(integrated_by_channel[ch], dtype=np.float64)
2494
+ if pfb_active:
2495
+ response = _load_pfb_response(bandpass_flatten.keywords["response_path"])
2496
+ flat_int = raw_int / np.maximum(response, 1e-10)
2497
+ overlay = response * (float(raw_int @ response) / float(response @ response))
2498
+ overlay_label = "scaled PFB response H"
2499
+ else:
2500
+ overlay = _fit_channel_bandpass(
2501
+ raw_int, raw_int.shape[0], self.config.inference.spline_order
2502
+ )
2503
+ flat_int = raw_int - overlay
2504
+ overlay_label = "spline fit"
2505
+ entry: dict = {"channel": int(ch), "overlay_label": overlay_label}
2506
+ for name, line in (("raw", raw_int), ("flat", flat_int), ("overlay", overlay)):
2507
+ idx, values = _decimate_for_plot(np.asarray(line), _ENVELOPE_MAX_POINTS)
2508
+ entry[name] = {"idx": [int(i) for i in idx], "values": [float(v) for v in values]}
2509
+ envelopes.append(entry)
2510
+ return envelopes
2511
+
2512
+ def _warn_on_pfb_response_mismatch(self, h5_path: str, n_coarse_total: int) -> None:
2513
+ """
2514
+ Residual-flatness sanity check of the static PFB response against the recording (after
2515
+ the bliss `validate` flag): flatten several coarse channels — sampled evenly across the
2516
+ band — with the active response H, and log an INFORMATIONAL warning once per file
2517
+ (never per channel) when the median flattened edge/mid power ratio deviates from 1.0
2518
+ by more than _PFB_RESIDUAL_FLATNESS_TOL.
2519
+
2520
+ This directly measures the operational question — does dividing by H actually flatten
2521
+ the data? — so, unlike the raw-vs-pure-ratio comparison it replaced, the
2522
+ analog-frontend passband tilt contributes only a small residual baseline rather than
2523
+ the entire statistic (and #180's characterization bounds the response-model error at
2524
+ ~1e-5, so H itself is not a confounder). The residual still includes that frontend
2525
+ tilt and any edge RFI, so a moderate deviation can be benign and the threshold stays
2526
+ provisional/informational until the deferred pfb_taps-vs-backend characterization
2527
+ fixes the legitimate baseline; a LARGE or CONSISTENT deviation (across many files) is
2528
+ the actionable signal of a wrong --pfb-taps-per-channel. Only the edge and mid bands
2529
+ of each sampled channel are read (see _flattened_edge_mid_ratio), keeping the check
2530
+ cheap; it runs on the cadence's primary ON file only, since the response is a static
2531
+ property of the backend shared by every file.
2532
+ """
2533
+ width = self.config.inference.coarse_channel_width
2534
+ try:
2535
+ response = gen_coarse_channel_response(
2536
+ width, n_coarse_total, self.config.inference.pfb_taps_per_channel
2537
+ )
2538
+ open_kwargs = _chunk_cache_kwargs(h5_path, time_bins=self.config.data.time_bins)
2539
+ ratios = [
2540
+ self._flattened_edge_mid_ratio(h5_path, ch, response, open_kwargs)
2541
+ for ch in self._sample_channel_indices(
2542
+ n_coarse_total, _PFB_MISMATCH_SAMPLE_CHANNELS
2543
+ )
2544
+ ]
2545
+ except Exception as e:
2546
+ logger.warning(f"PFB static-response sanity check failed for {h5_path}: {e}")
2547
+ return
2548
+
2549
+ self._log_pfb_ratio_verdict(ratios, h5_path)
2550
+
2551
+ def _start_pfb_response_check(
2552
+ self, h5_path: str, n_coarse_total: int, bandpass_flatten: functools.partial
2553
+ ):
2554
+ """
2555
+ Run the residual-flatness check on the worker pool instead of the parent (#298 I6):
2556
+ submit one small task per sampled channel via map_async and return the handle for
2557
+ _finish_pfb_response_check. The parent-side serial version idled every worker for
2558
+ its ~27 chunk-stripe band reads; pool-side, the reads overlap energy detection and
2559
+ the warning semantics are unchanged (same per-cadence cadence, same median of
2560
+ bit-identical ratios). Falls back to the synchronous check (returning None) when no
2561
+ pool exists; a failure to start is logged and swallowed — the check is
2562
+ informational and must never fail the cadence.
2563
+ """
2564
+ if self._ed_pool is None:
2565
+ self._warn_on_pfb_response_mismatch(h5_path, n_coarse_total)
2566
+ return None
2567
+ try:
2568
+ width = self.config.inference.coarse_channel_width
2569
+ time_bins = self.config.data.time_bins
2570
+ response_path = bandpass_flatten.keywords["response_path"]
2571
+ open_kwargs = _chunk_cache_kwargs(h5_path, time_bins)
2572
+ tasks = [
2573
+ (h5_path, ch, width, time_bins, response_path, open_kwargs)
2574
+ for ch in self._sample_channel_indices(
2575
+ n_coarse_total, _PFB_MISMATCH_SAMPLE_CHANNELS
2576
+ )
2577
+ ]
2578
+ return self._ed_pool.map_async(_pfb_ratio_worker, tasks)
2579
+ except Exception as e:
2580
+ logger.warning(f"PFB static-response sanity check failed to start for {h5_path}: {e}")
2581
+ return None
2582
+
2583
+ def _finish_pfb_response_check(self, pfb_check, h5_path: str) -> None:
2584
+ """Collect the pool-side ratios (all tasks have drained by the time the fused ED
2585
+ imap completes) and emit the same once-per-file verdict as the synchronous check."""
2586
+ try:
2587
+ ratios = pfb_check.get()
2588
+ except Exception as e:
2589
+ logger.warning(f"PFB static-response sanity check failed for {h5_path}: {e}")
2590
+ return
2591
+ self._log_pfb_ratio_verdict(ratios, h5_path)
2592
+
2593
+ def _log_pfb_ratio_verdict(self, ratios: list[float], h5_path: str) -> None:
2594
+ """Median the sampled flattened edge/mid ratios and warn once per file when the
2595
+ deviation from 1.0 exceeds _PFB_RESIDUAL_FLATNESS_TOL (informational — see
2596
+ _warn_on_pfb_response_mismatch's docstring for the interpretation contract)."""
2597
+ taps = self.config.inference.pfb_taps_per_channel
2598
+ # Median (not mean) so a single RFI-heavy sampled channel doesn't skew the statistic.
2599
+ median = float(np.median(ratios))
2600
+ deviation = abs(median - 1.0)
2601
+ if deviation > _PFB_RESIDUAL_FLATNESS_TOL:
2602
+ logger.warning(
2603
+ f"{h5_path}: median flattened edge/mid power ratio {median:.3f} deviates from "
2604
+ f"1.0 by {deviation:.1%} after dividing by the static PFB response "
2605
+ f"(residual-flatness sanity check — informational). The residual also reflects "
2606
+ f"analog-frontend tilt and edge RFI the response does not model, so a moderate "
2607
+ f"deviation can be benign. Investigate a large or consistent deviation across "
2608
+ f"many files: it may mean --pfb-taps-per-channel (currently {taps}) is wrong "
2609
+ f"for this backend, in which case --bandpass-method spline is the data-driven "
2610
+ f"fallback."
2611
+ )
2612
+
2613
+ def _plot_bandpass_overlay(
2614
+ self,
2615
+ h5_path: str,
2616
+ n_coarse_total: int,
2617
+ bandpass_flatten: Callable[[np.ndarray], np.ndarray],
2618
+ npy_path: str,
2619
+ ) -> None:
2620
+ """
2621
+ Opt-in debug artifact (--bandpass-debug-plot): for a few coarse channels sampled evenly
2622
+ across the band of the cadence's primary ON-source file, plot the time-integrated
2623
+ spectrum raw vs flattened, overlaying the model being removed (the scaled PFB response
2624
+ H, or the spline fit). Saved under {output_path}/plots/inference/{save_tag}/. Deliberately
2625
+ minimal — PR-08's inference visualization suite formalizes this figure.
2626
+
2627
+ Uses matplotlib's object-oriented Figure API rather than pyplot: _process_cadence runs
2628
+ on the streaming-inference prefetch thread, and pyplot's global figure registry is not
2629
+ thread-safe.
2630
+ """
2631
+ from matplotlib.figure import Figure # noqa: PLC0415
2632
+
2633
+ width = self.config.inference.coarse_channel_width
2634
+ pfb_active = bandpass_flatten.func is _pfb_flatten_bandpass
2635
+ sampled = self._sample_channel_indices(n_coarse_total)
2636
+
2637
+ fig = Figure(figsize=(14, 3.2 * len(sampled)))
2638
+ axes = fig.subplots(len(sampled), 2, squeeze=False)
2639
+ for row, ch in enumerate(sampled):
2640
+ channel = self._read_despiked_channel(h5_path, ch)
2641
+ raw = channel.mean(axis=0)
2642
+ flat = np.asarray(bandpass_flatten(channel)).mean(axis=0)
2643
+ if pfb_active:
2644
+ response = gen_coarse_channel_response(
2645
+ width, n_coarse_total, self.config.inference.pfb_taps_per_channel
2646
+ )
2647
+ # Least-squares scale so the unit-peak response overlays the raw spectrum
2648
+ overlay = response * (float(raw @ response) / float(response @ response))
2649
+ overlay_label = "scaled PFB response H"
2650
+ else:
2651
+ overlay = _fit_channel_bandpass(raw, width, self.config.inference.spline_order)
2652
+ overlay_label = "spline fit"
2653
+
2654
+ ax_raw, ax_flat = axes[row]
2655
+ # Decimated to a min/max envelope: full-resolution lines are ~1M points each at
2656
+ # GBT scale, which makes rendering slow and memory-heavy for no visual gain.
2657
+ ax_raw.plot(
2658
+ *_decimate_for_plot(raw), lw=0.6, color="tab:blue", label="raw integrated spectrum"
2659
+ )
2660
+ ax_raw.plot(
2661
+ *_decimate_for_plot(overlay),
2662
+ lw=1.2,
2663
+ ls="--",
2664
+ color="tab:orange",
2665
+ label=overlay_label,
2666
+ )
2667
+ ax_raw.set_ylabel(f"coarse channel {ch}\nintegrated power")
2668
+ ax_flat.plot(
2669
+ *_decimate_for_plot(flat),
2670
+ lw=0.6,
2671
+ color="tab:green",
2672
+ label="flattened integrated spectrum",
2673
+ )
2674
+ if row == 0:
2675
+ ax_raw.legend(loc="upper right", fontsize=8)
2676
+ ax_flat.legend(loc="upper right", fontsize=8)
2677
+ if row == len(sampled) - 1:
2678
+ ax_raw.set_xlabel("fine channel (within coarse channel)")
2679
+ ax_flat.set_xlabel("fine channel (within coarse channel)")
2680
+
2681
+ method = "pfb" if pfb_active else "spline"
2682
+ fig.suptitle(f"Bandpass flattening overlay ({method}): {os.path.basename(h5_path)}")
2683
+ fig.tight_layout()
2684
+
2685
+ tag = self.config.checkpoint.save_tag
2686
+ save_dir = os.path.join(self.config.output_path, "plots", "inference", tag)
2687
+ os.makedirs(save_dir, exist_ok=True)
2688
+ stem = os.path.splitext(os.path.basename(npy_path))[0]
2689
+ out_path = os.path.join(save_dir, f"bandpass_overlay_{stem}_{tag}.png")
2690
+ # No close/registry bookkeeping needed: an OO-API Figure is garbage-collected
2691
+ fig.savefig(out_path, dpi=120)
2692
+ logger.info(f"Saved bandpass overlay debug plot: {out_path}")
2693
+
2694
+ # NOTE: come back to this later (what's the trade-off for doing dedup vs not? e.g. lower storage & compute, but higher FNR or lower DR sensitivity?)
2695
+ @staticmethod
2696
+ def _deduplicate_hits(hits: list[tuple], stamp_width: int) -> list[tuple]:
2697
+ """
2698
+ Greedy merge of hits whose centers are within stamp_width // 2,
2699
+ keeping the one with the higher statistic.
2700
+ """
2701
+ if not hits:
2702
+ return []
2703
+ # Stable argsort on the center index alone — the identical permutation to the old
2704
+ # sorted(key=lambda h: h[0]) (stability preserves input order on ties), without
2705
+ # comparing ~1e6 Python tuples on RFI-dense cadences (#298 rider; this runs
2706
+ # serially on the prefetch critical path)
2707
+ centers = np.fromiter((h[0] for h in hits), dtype=np.int64, count=len(hits))
2708
+ sorted_hits = [hits[i] for i in np.argsort(centers, kind="stable")]
2709
+ half = stamp_width // 2
2710
+ merged: list[tuple] = [sorted_hits[0]]
2711
+ for h in sorted_hits[1:]:
2712
+ prev = merged[-1]
2713
+ if h[0] - prev[0] < half:
2714
+ if h[1] > prev[1]:
2715
+ merged[-1] = h
2716
+ else:
2717
+ merged.append(h)
2718
+ return merged