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.
- aetherscan/__init__.py +45 -0
- aetherscan/benchmark.py +179 -0
- aetherscan/candidate_figures.py +328 -0
- aetherscan/cli.py +3032 -0
- aetherscan/config.py +1002 -0
- aetherscan/dashboard.py +709 -0
- aetherscan/dashboard_cli.py +51 -0
- aetherscan/dashboard_launcher.py +194 -0
- aetherscan/data_generation.py +1448 -0
- aetherscan/db/__init__.py +21 -0
- aetherscan/db/db.py +2789 -0
- aetherscan/hf_hub.py +547 -0
- aetherscan/inference.py +912 -0
- aetherscan/inference_viz.py +1494 -0
- aetherscan/latent_gif.py +512 -0
- aetherscan/latent_variants.py +352 -0
- aetherscan/logger/__init__.py +21 -0
- aetherscan/logger/logger.py +467 -0
- aetherscan/logger/slack_handler.py +760 -0
- aetherscan/main.py +1269 -0
- aetherscan/manager/__init__.py +21 -0
- aetherscan/manager/manager.py +781 -0
- aetherscan/models/__init__.py +21 -0
- aetherscan/models/random_forest.py +171 -0
- aetherscan/models/vae.py +849 -0
- aetherscan/monitor/__init__.py +19 -0
- aetherscan/monitor/monitor.py +935 -0
- aetherscan/pfb.py +161 -0
- aetherscan/preprocessing.py +2718 -0
- aetherscan/rf_metrics.py +100 -0
- aetherscan/round_data.py +932 -0
- aetherscan/run_state.py +277 -0
- aetherscan/seeding.py +160 -0
- aetherscan/shap_parallel.py +238 -0
- aetherscan/tag_guards.py +206 -0
- aetherscan/train.py +7681 -0
- aetherscan-1.0.0.dist-info/METADATA +1187 -0
- aetherscan-1.0.0.dist-info/RECORD +41 -0
- aetherscan-1.0.0.dist-info/WHEEL +4 -0
- aetherscan-1.0.0.dist-info/entry_points.txt +2 -0
- aetherscan-1.0.0.dist-info/licenses/LICENSE +13 -0
|
@@ -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
|