diff --git a/README.md b/README.md index 4a41083..9f59dd4 100644 --- a/README.md +++ b/README.md @@ -667,6 +667,7 @@ usage: inference [-h] [--data-path DATA_PATH] [--model-path MODEL_PATH] [--detection-window-size DETECTION_WINDOW_SIZE] [--detection-step-size DETECTION_STEP_SIZE] [--stat-threshold STAT_THRESHOLD] [--stamp-width STAMP_WIDTH] + [--store-downsampled-stamps | --no-store-downsampled-stamps] [--overlap-search | --no-overlap-search] [--overlap-fraction OVERLAP_FRACTION] [--preprocess-output-dir PREPROCESS_OUTPUT_DIR] @@ -733,8 +734,10 @@ options: Number of fine channels per coarse channel (default: 1048576) --parallel-coarse-chans PARALLEL_COARSE_CHANS - Number of coarse channels to process in parallel per - block (default: 28) + Progress-logging chunk size for energy detection, in + coarse channels per log line (default: the number of + worker processes). Parallelism itself comes from the + persistent worker pool, not this knob. --spline-order SPLINE_ORDER Spline order for bandpass fitting (default: 16) --detection-window-size DETECTION_WINDOW_SIZE @@ -749,6 +752,13 @@ options: --stamp-width STAMP_WIDTH Width in fine channels of the extracted stamp around each hit (default: 4096; must equal --width-bin) + --store-downsampled-stamps, --no-store-downsampled-stamps + Downsample stamps along frequency (by --downsample- + factor) at extraction time, storing stamp_width // + downsample_factor bins per stamp (~8x smaller at + defaults; default: enabled). Pass --no-store- + downsampled-stamps to archive raw-resolution stamps; + loading handles both layouts. --overlap-search, --no-overlap-search Additionally extract stamps offset by ±overlap_fraction*stamp_width around each hit. Pass diff --git a/src/aetherscan/cli.py b/src/aetherscan/cli.py index 03cc30d..2f18b27 100644 --- a/src/aetherscan/cli.py +++ b/src/aetherscan/cli.py @@ -695,7 +695,7 @@ def _add_inference_flags_to(parser): "--parallel-coarse-chans", type=int, default=None, - help="Number of coarse channels to process in parallel per block (default: 28)", + help="Progress-logging chunk size for energy detection, in coarse channels per log line (default: the number of worker processes). Parallelism itself comes from the persistent worker pool, not this knob.", ) parser.add_argument( "--spline-order", @@ -727,6 +727,12 @@ def _add_inference_flags_to(parser): default=None, help="Width in fine channels of the extracted stamp around each hit (default: 4096; must equal --width-bin)", ) + parser.add_argument( + "--store-downsampled-stamps", + action=argparse.BooleanOptionalAction, + default=None, + help="Downsample stamps along frequency (by --downsample-factor) at extraction time, storing stamp_width // downsample_factor bins per stamp (~8x smaller at defaults; default: enabled). Pass --no-store-downsampled-stamps to archive raw-resolution stamps; loading handles both layouts.", + ) parser.add_argument( "--overlap-search", action=argparse.BooleanOptionalAction, @@ -1038,6 +1044,11 @@ def apply_args_to_config(args: argparse.Namespace) -> None: config.inference.stat_threshold = args.stat_threshold if hasattr(args, "stamp_width") and args.stamp_width is not None: config.inference.stamp_width = args.stamp_width + # store_downsampled_stamps uses argparse.BooleanOptionalAction with default=None so that + # the CLI can express "leave the config default" (omit), "force on" + # (--store-downsampled-stamps), and "force off" (--no-store-downsampled-stamps) + if hasattr(args, "store_downsampled_stamps") and args.store_downsampled_stamps is not None: + config.inference.store_downsampled_stamps = args.store_downsampled_stamps # overlap_search uses argparse.BooleanOptionalAction with default=None so that # the CLI can express "leave the config default" (omit), "force on" # (--overlap-search), and "force off" (--no-overlap-search). The `is not None` @@ -1710,6 +1721,32 @@ def collect_validation_errors( fix_kind="clamp_low", ) ) + # Downsample-at-extraction stores stamp_width // downsample_factor bins per + # stamp, so the division must be exact for the stored width to be well-defined + store_ds = _resolve( + args, "store_downsampled_stamps", config.inference.store_downsampled_stamps + ) + df_inf = _resolve(args, "downsample_factor", config.data.downsample_factor) + if ( + store_ds + and sw_inf is not None + and df_inf is not None + and df_inf > 0 + and sw_inf % df_inf != 0 + ): + errors.append( + ValidationError( + field="inference.stamp_width", + current=sw_inf, + message=( + f"--stamp-width ({sw_inf}) must be divisible by " + f"--downsample-factor ({df_inf}) when stamps are downsampled at " + "extraction (--store-downsampled-stamps)" + ), + fix_kind="divisibility", + divisor=df_inf, + ) + ) coarse_channel_width = _resolve( args, "coarse_channel_width", config.inference.coarse_channel_width diff --git a/src/aetherscan/config.py b/src/aetherscan/config.py index 0e686a2..91b563d 100644 --- a/src/aetherscan/config.py +++ b/src/aetherscan/config.py @@ -291,12 +291,19 @@ class InferenceConfig: # NOTE: come back to this later (are these params correct?) coarse_channel_width: int = 1048576 - parallel_coarse_chans: int = 28 + # Progress-logging chunk size for energy detection (coarse channels per log line). + # None -> use manager.n_processes. Parallelism itself comes from the persistent worker + # pool (one fused task per coarse channel), not from this knob. + parallel_coarse_chans: int | None = None spline_order: int = 16 detection_window_size: int = 256 detection_step_size: int = 128 stat_threshold: float = 2048.0 stamp_width: int = 4096 + # Downsample stamps along frequency at extraction time (by data.downsample_factor), so + # the per-cadence .npy stores stamp_width // downsample_factor bins (~8x smaller at + # defaults). Disable to archive raw-resolution stamps; loading handles both layouts. + store_downsampled_stamps: bool = True # NOTE: come back to this later (is overlap search implemented correctly? redo analysis from peter forwards search paper) overlap_search: bool = True @@ -571,6 +578,7 @@ def to_dict(self) -> dict: "detection_step_size": self.inference.detection_step_size, "stat_threshold": self.inference.stat_threshold, "stamp_width": self.inference.stamp_width, + "store_downsampled_stamps": self.inference.store_downsampled_stamps, "overlap_search": self.inference.overlap_search, "overlap_fraction": self.inference.overlap_fraction, "discard_side_channels": self.inference.discard_side_channels, diff --git a/src/aetherscan/inference.py b/src/aetherscan/inference.py index d70a847..2d84b98 100644 --- a/src/aetherscan/inference.py +++ b/src/aetherscan/inference.py @@ -57,7 +57,7 @@ def clear(self): def prepare_distributed_inf_dataset( - data: dict, + data: np.ndarray, n_samples: int, per_replica_inf_batch_size: int, num_replicas: int, @@ -65,40 +65,39 @@ def prepare_distributed_inf_dataset( ) -> dict: """ Build a distributed inference dataset from `data` (shape (n_samples, 6, 16, 512)) and return - a dict with the dataset, trimmed sample count, step count, and the InfDataHolder. + a dict with the dataset, padded/real sample counts, step count, and the InfDataHolder. Distinct from train.py's RF-training counterpart: signal classes are not assumed to be known ahead of time, so the dataset yields raw cadences without label channels. Order is preserved - (no shuffle) since gradients aren't computed during inference. Currently no trimming is done - (n_inf_trimmed == n_samples); the trimming scaffolding is left in place for future - divisibility experiments. + (no shuffle) since gradients aren't computed during inference. + + When n_samples isn't divisible by the global batch size, the data is padded with duplicate + rows up to the next batch multiple (mirroring train.py's viz-dataset padding pattern) so the + final partial batch is encoded rather than silently dropped — callers truncate the encoded + latents back to n_samples. The previous behavior (inf_steps = n // global_batch with + drop_remainder=True and no padding) never processed the tail; with per-cadence batches a + cadence smaller than one global batch would have processed nothing at all. """ global_inf_batch_size = per_replica_inf_batch_size * num_replicas - # # NOTE: does trimming/divisibility matter for inference? - # # Trim datasets to fit batch sizes (prevents uneven batches on final step) - # # Note, n_samples should already be divisible by effective_batch_size - # # Trimming here is just a defensive measure to doubly ensure divisibility before creating & - # # distributing our datasets - # # Alternatively, we could also pad the data instead of trimming - # n_inf_trimmed = (n_samples // global_inf_batch_size) * global_inf_batch_size - n_inf_trimmed = n_samples - - # Sanity check: verify there's enough samples to run inference - if n_inf_trimmed == 0: - raise ValueError( - f"Not enough samples ({n_samples}) for global batch size ({per_replica_inf_batch_size} * {num_replicas})" - f"Reduce per_replica_batch_size or provide more samples" - ) - # logger.info(f"Data alignment: Inf {n_samples}→{n_inf_trimmed}") - - # # Randomly subsample to trimmed size (avoids positional bias from slicing the tail) - # if n_inf_trimmed < n_samples: - # indices = np.random.choice(n_samples, size=n_inf_trimmed, replace=False) - # inf_data = data[indices] - # else: - # inf_data = data[:n_inf_trimmed] - inf_data = data[:n_inf_trimmed] + # Sanity check: verify there's at least one sample to run inference on + if n_samples == 0: + raise ValueError("Not enough samples (0) to run inference") + + # Pad with duplicate rows (cycled from the front, deterministic) to the next global-batch + # multiple; the encoded outputs for the padded rows are discarded by the caller. + # NOTE: np.concatenate materializes a full copy of the cadence's stamps, briefly doubling + # that cadence's memory footprint. Acceptable at per-cadence scale (the padding itself is + # under one global batch); revisit if batches ever wrap multi-cadence arrays again. + n_padded = int(np.ceil(n_samples / global_inf_batch_size)) * global_inf_batch_size + if n_padded > n_samples: + pad_count = n_padded - n_samples + pad_indices = np.arange(pad_count) % n_samples + inf_data = np.concatenate([data, data[pad_indices]], axis=0) + logger.info(f"Data alignment: Inf {n_samples}→{n_padded} (padded {pad_count})") + else: + inf_data = data + logger.info(f"Data alignment: Inf {n_samples} (no padding needed)") inf_holder = InfDataHolder(inf_data) # Create generator function for memory-efficient data loading @@ -145,18 +144,13 @@ def inf_generator(): inf_dataset_distributed = strategy.experimental_distribute_dataset(inf_dataset) - # Calculate steps - inf_steps = n_inf_trimmed // global_inf_batch_size - - # Sanity check: verify step sizes are valid before returning - if inf_steps < 1: - raise ValueError( - f"inf_steps < 1: n_inf_trimmed ({n_inf_trimmed}) must be >= per_replica_inf_batch_size * num_replicas ({per_replica_inf_batch_size} * {num_replicas})" - ) + # Calculate steps (n_padded is an exact multiple of the global batch size by construction) + inf_steps = n_padded // global_inf_batch_size return { "inf_dataset": inf_dataset_distributed, - "n_inf_trimmed": n_inf_trimmed, + "n_padded": n_padded, + "n_samples": n_samples, "inf_steps": inf_steps, "_inf_holder": inf_holder, } @@ -194,6 +188,11 @@ def __init__(self, strategy: tf.distribute.Strategy = None): self.per_replica_inf_batch_size = self.config.inference.per_replica_batch_size self.threshold = self.config.inference.classification_threshold + # Lazily-built tf.function for the distributed encode step, cached so repeated + # run_inference calls (one per cadence in the streaming loop) reuse one trace + # instead of accumulating a new concrete function per cadence + self._encode_step = None + # TODO: implement model loading from HuggingFace (parametrize args to InferenceConfig) def init_models(self, encoder_path: str, rf_path: str): """Load the VAE encoder (inside strategy.scope) and the Random Forest classifier from disk.""" @@ -223,17 +222,16 @@ def init_models(self, encoder_path: str, rf_path: str): logger.error(f"Error loading Random Forest: {e}") raise # Re-raise to propagate error - # TODO: finish writing docstring (how will args be passed?) def run_inference( self, data: np.ndarray, npy_path: str, - # NOTE: how do we pass these in from preproc? target: str | None = None, session: str | None = None, cadence_id: int | None = None, band: str | None = None, frequency_mhz: float | None = None, + stamp_frequencies_mhz: list[float] | None = None, timestamp_observed: float | None = None, h5_path: str | None = None, ) -> dict: @@ -244,17 +242,35 @@ def run_inference( Encodes each snippet through the VAE encoder under the distribution strategy, then runs the Random Forest classifier on the latents and writes positive predictions to the database (along with the latent vector and observational provenance: target, session, - cadence_id, band, frequency_mhz, timestamp_observed, h5_path). + cadence_id, band, frequency, timestamp_observed, h5_path — typically derived by + preprocessing.derive_cadence_provenance from the cadence's metadata JSON). + stamp_frequencies_mhz, when given, carries one center frequency per snippet row (the + metadata's stamp_frequencies_mhz) and takes precedence over the scalar frequency_mhz. + + Safe to call repeatedly on one pipeline instance (the per-cadence streaming loop does): + models stay loaded, and per-call dataset state is released in the finally block. The + caller owns tf.keras.backend.clear_session() — see run_inference_pipeline / + main.inference_command. """ # Sanity check if not self.encoder or not self.rf_model: raise RuntimeError("Encoder and/or Random Forest not initialized") + inf_holder = None + inf_dataset = None + try: n_samples = data.shape[0] logger.info(f"Running inference on {n_samples} cadence snippets from {npy_path}") - # Prepare distributed dataset for inference + if stamp_frequencies_mhz is not None and len(stamp_frequencies_mhz) != n_samples: + logger.warning( + f"stamp_frequencies_mhz has {len(stamp_frequencies_mhz)} entries but " + f"{n_samples} snippets were loaded; ignoring per-stamp frequencies" + ) + stamp_frequencies_mhz = None + + # Prepare distributed dataset for inference (pads to a global-batch multiple) results = prepare_distributed_inf_dataset( data=data, n_samples=n_samples, @@ -267,7 +283,7 @@ def run_inference( gc.collect() inf_dataset = results["inf_dataset"] - n_inf_trimmed = results["n_inf_trimmed"] + n_padded = results["n_padded"] inf_steps = results["inf_steps"] inf_holder = results["_inf_holder"] @@ -275,10 +291,15 @@ def run_inference( gc.collect() logger.info( - f"Generating latents for {n_inf_trimmed} cadence snippets using distributed inference" + f"Generating latents for {n_samples} cadence snippets " + f"({n_padded} padded) using distributed inference" ) - latents = self._distributed_encode(inf_dataset, n_inf_trimmed, inf_steps) + latents = self._distributed_encode(inf_dataset, n_padded, inf_steps) + + # Drop the encoded padding rows: the first n_samples * num_observations latent + # rows correspond exactly to the real snippets (row order is snippet-major) + latents = latents[: n_samples * self.num_observations] # Run RF classification logger.info("Running Random Forest classification") @@ -295,6 +316,7 @@ def run_inference( cadence_id=cadence_id, band=band, frequency_mhz=frequency_mhz, + stamp_frequencies_mhz=stamp_frequencies_mhz, timestamp_observed=timestamp_observed, h5_path=h5_path, ) @@ -304,24 +326,18 @@ def run_inference( raise # Re-raise to propagate error finally: - # NOTE: should check to make sure holder & dataset exist first - # Clear intermediate data - inf_holder.clear() + # Clear per-call dataset state (guarded: an early failure may predate creation). + # Note tf.keras.backend.clear_session() is intentionally NOT called here — the + # streaming loop reuses this pipeline's loaded models across cadences; the caller + # clears the session once when the whole run is done. + if inf_holder is not None: + inf_holder.clear() del inf_dataset - - # Force TensorFlow to release internal references to datasets/iterators - # This prevents generator closures from accumulating in memory between rounds - tf.keras.backend.clear_session() - logger.info("Cleared TensorFlow session state") - - # NOTE: should check to make sure arrays exist first - del latents, predictions, confidence_scores gc.collect() - # NOTE: is this the right location for return statement? return { "n_cadence_snippets": n_samples, - "n_processed": n_inf_trimmed, + "n_processed": n_samples, "n_candidates": n_candidates, } @@ -341,23 +357,28 @@ def _distributed_encode( # Use np.empty() instead of np.zeros() so problematic latent values don't fail silently latents = np.empty((n_samples * self.num_observations, self.latent_dim), dtype=np.float32) - # Cache dimensions for tf.function - time_bins = self.config.data.time_bins - width_bin = self.config.data.width_bin // self.config.data.downsample_factor + if self._encode_step is None: + # Cache dimensions for tf.function + time_bins = self.config.data.time_bins + width_bin = self.config.data.width_bin // self.config.data.downsample_factor + + @tf.function + def encode_step(batch_data): + def encode_fn(data): + """Per-replica encoding step""" + # Reshape for encoder: (batch, 6, 16, 512) -> (batch * 6, 16, 512, 1) + reshaped = tf.reshape(data, [-1, time_bins, width_bin, 1]) + # Encode (returns mean, log_var, z) + _, _, z = self.encoder(reshaped, training=False) + return z - @tf.function - def encode_step(batch_data): - def encode_fn(data): - """Per-replica encoding step""" - # Reshape for encoder: (batch, 6, 16, 512) -> (batch * 6, 16, 512, 1) - reshaped = tf.reshape(data, [-1, time_bins, width_bin, 1]) - # Encode (returns mean, log_var, z) - _, _, z = self.encoder(reshaped, training=False) - return z + # Run encoding on all replicas + per_replica_z = self.strategy.run(encode_fn, args=(batch_data,)) + return per_replica_z - # Run encoding on all replicas - per_replica_z = self.strategy.run(encode_fn, args=(batch_data,)) - return per_replica_z + self._encode_step = encode_step + + encode_step = self._encode_step # Process all batches iterator = iter(dataset) @@ -411,10 +432,20 @@ def _write_inference_results( cadence_id: int | None = None, band: str | None = None, frequency_mhz: float | None = None, + stamp_frequencies_mhz: list[float] | None = None, timestamp_observed: float | None = None, h5_path: str | None = None, ) -> int: - """Write inference results to database.""" + """Write inference results to database. + + Predictions/confidences are indexed per snippet (dataset order is preserved end to end: + the .npy rows, the padded dataset, and the truncated latents all share snippet-major + ordering, so snippet_index == the row index in npy_path). stamp_frequencies_mhz, when + given, supplies the per-snippet center frequency (falling back to the scalar + frequency_mhz otherwise); latents rows are per observation, so snippet idx spans + latents[idx * num_observations : (idx + 1) * num_observations], flattened to one + (num_observations * latent_dim,) provenance vector. + """ if self.db is None: raise RuntimeError("No database instance detected - cannot store inference results") @@ -430,20 +461,27 @@ def _write_inference_results( if prediction == 1: n_candidates += 1 + snippet_frequency_mhz = frequency_mhz + if stamp_frequencies_mhz is not None: + snippet_frequency_mhz = float(stamp_frequencies_mhz[idx]) + + # One latent row per observation -> flatten the snippet's num_observations + # rows into a single (num_observations * latent_dim,) vector + latent_rows = latents[ + idx * self.num_observations : (idx + 1) * self.num_observations + ] + self.db.write_inference_result( npy_path=npy_path, - # NOTE: is it guaranteed that snippets are processed sequentially? snippet_index=idx, prediction=prediction, confidence=confidence, - # NOTE: does latents need to be reshaped first before being passed as arg? - latent_vector=latents[idx], - # NOTE: how do we pass these in from preproc? + latent_vector=latent_rows.reshape(-1), target=target, session=session, cadence_id=cadence_id, band=band, - frequency_mhz=frequency_mhz, + frequency_mhz=snippet_frequency_mhz, timestamp_observed=timestamp_observed, h5_path=h5_path, tag=tag, @@ -457,43 +495,54 @@ def plot_candidate(self): pass -# TODO: figure out how to pass preproc metadata into InferencePipeline (target, session, cadence_id, band, frequency_mhz, timestamp_observed, h5_path). should we roll these metadata + npy_path into a list/dict from preproc, then unroll them inside run_inference_pipeline()? # TODO: add try-except switch statements (see run_training_pipeline()) def run_inference_pipeline( cadence_data: np.ndarray, npy_path: str, strategy: tf.distribute.Strategy, - # NOTE: how do we pass these in from preproc? target: str | None = None, session: str | None = None, cadence_id: int | None = None, band: str | None = None, frequency_mhz: float | None = None, + stamp_frequencies_mhz: list[float] | None = None, timestamp_observed: float | None = None, h5_path: str | None = None, ) -> dict: """ - End-to-end inference entry point: build an InferencePipeline under `strategy` and run it - against `cadence_data` (shape (n, 6, 16, 512)) sourced from `npy_path`. The optional - target/session/cadence_id/band/frequency_mhz/timestamp_observed/h5_path arguments are - observational provenance that gets written to the inference_results table for any positive - candidates. Returns the {n_cadence_snippets, n_processed, n_candidates} dict from run_inference. + Single-shot inference entry point: build an InferencePipeline under `strategy`, run it once + against `cadence_data` (shape (n, 6, 16, 512)) sourced from `npy_path`, and clear the TF + session. Used by the legacy --test-files path, which loads one preprocessed array up front; + the CSV streaming path in main.inference_command instead builds one InferencePipeline and + calls run_inference per cadence so models load once. + + The optional provenance arguments (target/session/cadence_id/band/frequency_mhz/ + stamp_frequencies_mhz/timestamp_observed/h5_path) are written to the inference_results + table for any positive candidates. Returns the {n_cadence_snippets, n_processed, + n_candidates} dict from run_inference. """ # Create pipeline pipeline = InferencePipeline(strategy=strategy) - # Run inference - results = pipeline.run_inference( - data=cadence_data, - # NOTE: how do we pass these in from preproc? - npy_path=npy_path, - target=target, - session=session, - cadence_id=cadence_id, - band=band, - frequency_mhz=frequency_mhz, - timestamp_observed=timestamp_observed, - h5_path=h5_path, - ) + try: + # Run inference + results = pipeline.run_inference( + data=cadence_data, + npy_path=npy_path, + target=target, + session=session, + cadence_id=cadence_id, + band=band, + frequency_mhz=frequency_mhz, + stamp_frequencies_mhz=stamp_frequencies_mhz, + timestamp_observed=timestamp_observed, + h5_path=h5_path, + ) + finally: + # Force TensorFlow to release internal references to datasets/iterators. Runs once + # per pipeline lifetime (after the loaded models are no longer needed) rather than + # inside run_inference, which may be called repeatedly on live models. + tf.keras.backend.clear_session() + logger.info("Cleared TensorFlow session state") return results diff --git a/src/aetherscan/main.py b/src/aetherscan/main.py index 3fc5c9d..9969582 100644 --- a/src/aetherscan/main.py +++ b/src/aetherscan/main.py @@ -12,6 +12,7 @@ import os import sys import time +from concurrent.futures import ThreadPoolExecutor import numpy as np import tensorflow as tf @@ -25,11 +26,11 @@ ) from aetherscan.config import get_config, init_config from aetherscan.db import init_db -from aetherscan.inference import run_inference_pipeline +from aetherscan.inference import InferencePipeline, run_inference_pipeline from aetherscan.logger import init_logger from aetherscan.manager import get_manager, init_manager, register_logger from aetherscan.monitor import init_monitor -from aetherscan.preprocessing import DataPreprocessor +from aetherscan.preprocessing import DataPreprocessor, derive_cadence_provenance from aetherscan.train import get_latest_tag, run_training_pipeline logger = logging.getLogger(__name__) @@ -295,6 +296,129 @@ def train_command(): logger.info("=" * 60) +class NonRetryableInferenceError(RuntimeError): + """A permanent inference failure (bad catalog/config) that retrying cannot fix. + + inference_command's retry loop re-raises this immediately instead of burning retry + attempts on it; transient failures (I/O hiccups, GPU errors) stay plain exceptions and + keep the existing retry semantics. + """ + + +def _run_streaming_csv_inference(preprocessor: DataPreprocessor, strategy) -> dict: + """ + Per-cadence streaming inference over the configured CSV catalogs. + + Flow (peak memory = one cadence's stamps + the next cadence's in-flight preprocessing, + independent of catalog size): + + units = plan_cadences() # cadence groups + .npy paths, no work yet + pipeline = InferencePipeline(strategy) # models loaded once for the whole run + for each cadence (prefetch depth 1): + [background thread] preprocess cadence i+1 (energy detection; the persistent pool's + child processes do the real work — the thread mostly waits) + [main thread] load cadence i's stamps (memmap -> log-norm) -> encode on GPUs + -> RF -> write per-cadence results with per-cadence provenance + + Provenance (target/session/band/cadence_id/per-stamp frequency/timestamp/h5_path) is + derived per cadence from its group key + metadata JSON via derive_cadence_provenance. + Returns aggregate {n_cadence_snippets, n_processed, n_candidates, n_cadences}. Exceptions + from the load/encode stages propagate to the retry loop in inference_command; per-cadence + preprocessing failures are logged and skipped inside process_pending_cadence. Raises + NonRetryableInferenceError when the catalog yields no work units or no cadence produces a + stamp .npy — permanent conditions the retry loop must not retry. + """ + config = get_config() + + units = preprocessor.plan_cadences() + if not units: + raise NonRetryableInferenceError( + "No cadence work units produced from the configured inference CSVs" + ) + logger.info(f"Streaming inference over {len(units)} cadence(s)") + + # Load models once; every cadence reuses this pipeline + pipeline = InferencePipeline(strategy=strategy) + + totals = {"n_cadence_snippets": 0, "n_processed": 0, "n_candidates": 0, "n_cadences": 0} + + # Start the persistent energy-detection pool from the main thread (forking after + # background threads exist risks inheriting mid-operation locks in the children) + preprocessor.start_energy_detection_pool() + try: + with ThreadPoolExecutor(max_workers=1, thread_name_prefix="preproc_prefetch") as prefetch: + future = prefetch.submit(preprocessor.process_pending_cadence, units[0]) + + for i, unit in enumerate(units): + # NOTE: an exception inside a prefetched preprocessing task surfaces here, + # one iteration after it was submitted, when its future is resolved — and + # then propagates to inference_command's retry loop. In practice + # process_pending_cadence swallows per-cadence failures (returns None), so + # only infrastructure-level errors (e.g. a broken worker pool) raise. + cadence_result = future.result() + + # Prefetch depth 1: kick off cadence i+1's CPU preprocessing while the main + # thread loads + encodes cadence i on the GPUs + if i + 1 < len(units): + future = prefetch.submit(preprocessor.process_pending_cadence, units[i + 1]) + + if cadence_result is None: + logger.info(f"Cadence {unit.group.key}: no stamps produced; skipping") + continue + + # Per-cadence provenance from the group key + metadata JSON + try: + with open(cadence_result.metadata_path) as f: + metadata = json.load(f) + except (OSError, json.JSONDecodeError) as e: + logger.warning( + f"Cadence {cadence_result.key}: could not read metadata at " + f"{cadence_result.metadata_path} ({e}); provenance will be sparse" + ) + metadata = {"h5_paths": cadence_result.h5_paths} + provenance = derive_cadence_provenance( + key=cadence_result.key, + group_by_cols=config.inference.cadence_group_by_cols, + metadata=metadata, + ) + + # copy=False: the loader already returns float32; don't duplicate GBs of stamps + cadence_data = preprocessor.load_inference_data( + override_filepaths=[cadence_result.npy_path] + ).astype(np.float32, copy=False) + + results = pipeline.run_inference( + data=cadence_data, + npy_path=cadence_result.npy_path, + **provenance, + ) + del cadence_data + gc.collect() + + totals["n_cadence_snippets"] += results["n_cadence_snippets"] + totals["n_processed"] += results["n_processed"] + totals["n_candidates"] += results["n_candidates"] + totals["n_cadences"] += 1 + + logger.info( + f"Cadence {cadence_result.key} ({totals['n_cadences']} done, " + f"{len(units) - i - 1} to go): {results['n_processed']} snippets, " + f"{results['n_candidates']} candidate(s)" + ) + finally: + preprocessor.stop_energy_detection_pool() + # Release TF dataset/iterator state once per run, after the loaded models are done + tf.keras.backend.clear_session() + logger.info("Cleared TensorFlow session state") + + if totals["n_cadences"] == 0: + # Preserve the historical contract: preprocessing producing no stamp .npy at all is + # an error (bad paths/catalog), not a legitimate empty result + raise NonRetryableInferenceError("No cadence results produced by preprocessing") + + return totals + + # NOTE: we need to load the saved config from the corresponding training run, but when/where should we do that, and how does that play with apply_args_to_config()? def inference_command(): """Execute inference pipeline with distributed strategy & fault tolerance""" @@ -338,9 +462,8 @@ def inference_command(): sys.exit(1) # NOTE: come back to this later (does fault tolerance work properly with inference?) - # NOTE: come back to this later (should we add some async/back-and-forth design patterns -- e.g. preproc X files, inference X files, clear, repeat -- to reduce memory pressure? is this the most efficient architecture we can use? add comments about memory/performance trade-offs once inference pipeline complete (see preproc section in train_command()) # Run preprocessing + inference with fault tolerance. - # Recovery is state-based, not checkpoint-based: find_hits() writes per-cadence + # Recovery is state-based, not checkpoint-based: preprocessing writes per-cadence # .npy files as it goes and skips any whose .npy already exists, so simply # retrying resumes from where the last attempt died. No checkpoint metadata # is needed for the preprocessing stage. @@ -348,10 +471,11 @@ def inference_command(): max_retries = config.inference.max_retries retry_delay = config.inference.retry_delay results = None - # Cache cadence_data across retry attempts: once preprocessing + loading - # succeeds, an inference-only failure shouldn't trigger a re-load / - # re-downsample / re-log-norm pass. Mirrors how train_command loads - # background_data once outside its retry loop. + # Legacy --test-files path only: cache cadence_data across retry attempts so an + # inference-only failure doesn't trigger a re-load / re-downsample / re-log-norm + # pass. Mirrors how train_command loads background_data once outside its retry + # loop. The streaming CSV path holds no catalog-sized state to cache — its + # per-cadence .npy files are the resume mechanism. cadence_data: np.ndarray | None = None npy_path_for_logging: str | None = None @@ -359,45 +483,33 @@ def inference_command(): try: logger.info(f"Inference attempt: {attempt + 1}/{max_retries}") - # Preprocessing + load stage. Skipped on retry if a previous attempt - # already produced cadence_data (i.e. only the inference stage failed). - if cadence_data is None: - if config.data.inference_files is not None: - cadence_results = preprocessor.find_hits() - if not cadence_results: - logger.error("No cadence results produced by preprocessing") - sys.exit(1) - npy_paths = [cr.npy_path for cr in cadence_results] - logger.info( - f"Preprocessing produced {len(npy_paths)} cadence .npy file(s); " - f"loading into inference" + if config.data.inference_files is not None: + # Streaming CSV path: per-cadence preprocess -> load -> encode -> RF -> + # write, with models loaded once and prefetch depth 1 (see + # _run_streaming_csv_inference). Memory stays independent of catalog size. + results = _run_streaming_csv_inference(preprocessor, strategy) + else: + if not config.data.test_files: + logger.error( + "Neither --inference-files nor --test-files is configured; " + "nothing to load for inference" ) - cadence_data = preprocessor.load_inference_data( - override_filepaths=npy_paths - ).astype(np.float32) - npy_path_for_logging = npy_paths[0] - else: - if not config.data.test_files: - logger.error( - "Neither --inference-files nor --test-files is configured; " - "nothing to load for inference" - ) - sys.exit(1) + sys.exit(1) + # Load stage. Skipped on retry if a previous attempt already produced + # cadence_data (i.e. only the inference stage failed). + if cadence_data is None: cadence_data = preprocessor.load_inference_data().astype(np.float32) npy_path_for_logging = config.data.test_files[0] - else: - logger.info( - "Reusing cadence_data from previous attempt (skipping preprocessing + load)" + else: + logger.info("Reusing cadence_data from previous attempt (skipping load)") + + # NOTE: come back to this later (inference-stage resume should skip cadences already in the DB. not yet implemented) + # Inference stage + results = run_inference_pipeline( + cadence_data=cadence_data, + npy_path=npy_path_for_logging, # TODO: handle multiple test_files properly + strategy=strategy, ) - - # NOTE: come back to this later (inference-stage resume should skip cadences already in the DB. not yet implemented) - # Inference stage - results = run_inference_pipeline( - cadence_data=cadence_data, - npy_path=npy_path_for_logging, # TODO: handle multiple test_files properly - strategy=strategy, - # TODO: figure out how to pass preproc metadata into InferencePipeline (target, session, cadence_id, band, frequency_mhz, timestamp_observed, h5_path). should we roll these metadata + npy_path into a list/dict from preproc, then unroll them inside run_inference_pipeline()? - ) break # success except KeyboardInterrupt: @@ -405,6 +517,12 @@ def inference_command(): logger.info("Inference interrupted by user") raise + except NonRetryableInferenceError as e: + # Permanent failure (empty/invalid catalog): retrying can't fix it, so fail + # fast instead of burning the remaining attempts + logger.error(f"Inference failed permanently: {e}") + sys.exit(1) + except Exception as e: logger.error(f"Inference attempt {attempt + 1} failed with error: {e}") if attempt < max_retries - 1: @@ -427,6 +545,8 @@ def inference_command(): logger.info("=" * 60) logger.info("Inference completed successfully!") logger.info("Summary:") + if "n_cadences" in results: + logger.info(f" Total cadences: {results['n_cadences']}") logger.info(f" Total cadence snippets: {results['n_cadence_snippets']}") logger.info(f" Processed: {results['n_processed']}") logger.info(f" Candidates found: {results['n_candidates']}") diff --git a/src/aetherscan/preprocessing.py b/src/aetherscan/preprocessing.py index 79a7c1c..143f9d4 100644 --- a/src/aetherscan/preprocessing.py +++ b/src/aetherscan/preprocessing.py @@ -10,14 +10,18 @@ import contextlib import csv +import functools import gc import json import logging +import math import os import re import signal from collections import OrderedDict +from collections.abc import Callable from dataclasses import dataclass, field +from multiprocessing.pool import Pool from multiprocessing.shared_memory import SharedMemory import h5py @@ -102,6 +106,17 @@ def cleanup_on_sigterm(signum, frame): _GLOBAL_DTYPE = dtype +def _init_plain_worker(): + """ + Worker pool initializer for pools that don't attach shared memory (energy detection and + stamp extraction): set up queue-based logging and ignore SIGINT so the parent's + ResourceManager coordinates shutdown. No SIGTERM handler is installed — these workers hold + no shared-memory file descriptors, so the default termination behavior is safe. + """ + init_worker_logging() + signal.signal(signal.SIGINT, signal.SIG_IGN) + + # NOTE: come back to this later def _downsample_worker(args): """ @@ -137,6 +152,33 @@ def _downsample_worker(args): return None +def _lognorm_worker(args): + """ + Log-normalize one already-downsampled cadence in parallel. + + Counterpart to _downsample_worker for stamps that were downsampled at extraction time + (see _extract_stamps_worker): the stored .npy is already at final width, so loading only + needs the per-cadence log-norm. args is a (cadence_idx,) tuple — the cadence itself is + pulled from _GLOBAL_CHUNK_DATA to avoid pickling it across the pool boundary. Returns the + log-normalized cadence of shape (6, time_bins, final_width) as float32, or None if the + source cadence contained NaN/Inf or had non-positive max (treated as invalid, matching + _downsample_worker). + """ + (cadence_idx,) = args + + if _GLOBAL_CHUNK_DATA is None: + logger.warning("No global chunk data available") + return None + + cadence = _GLOBAL_CHUNK_DATA[cadence_idx] + + # Skip invalid cadences (same validity rule as _downsample_worker) + if np.any(np.isnan(cadence)) or np.any(np.isinf(cadence)) or np.max(cadence) <= 0: + return None + + return log_norm(cadence.astype(np.float32)) + + # NOTE: come back to this later (mirrors preprocess_fine.py:72-75 from the reference implementation) def _remove_dc_spike( block_data: np.ndarray, coarse_channel_width: int, n_coarse_channels: int @@ -183,90 +225,209 @@ def _fit_channel_bandpass( return interpolate.splev(x, spl) -# NOTE: come back to this later -def _read_coarse_channel_worker(args: tuple) -> np.ndarray: +def _spline_flatten_bandpass(channel: np.ndarray, spl_order: int) -> np.ndarray: """ - Worker: read one coarse channel from an .h5 file as an ndarray of shape - (time_bins, coarse_channel_width). args is (h5_path, channel_index, coarse_channel_width). - - Each worker opens its own h5py.File since h5py file handles are fork-unsafe to share. + Default bandpass flattener: fit a spline to the time-integrated coarse channel and subtract + it from every time row. channel has shape (time_bins, coarse_channel_width); returns the + float64 residuals of the same shape. + + Bandpass flattening is a pluggable stage: _process_cadence obtains a picklable callable via + DataPreprocessor._get_bandpass_flattener() and ships it to the pool workers, so alternative + flatteners (e.g. the PFB static equalization planned for a later PR) only need a + module-level function with the same (channel) -> residuals contract. """ - h5_path, channel_index, coarse_channel_width = args - start = channel_index * coarse_channel_width - end = (channel_index + 1) * coarse_channel_width - with h5py.File(h5_path, "r") as hf: - return hf["data"][:, 0, start:end] + integrated_channel = np.mean(channel, axis=0) + fit = _fit_channel_bandpass(integrated_channel, channel.shape[1], spl_order) + return channel - fit -# NOTE: come back to this later -def _remove_bandpass_worker(args: tuple) -> np.ndarray: +def _sliding_normality_k2(channel: np.ndarray, window_size: int, step_size: int) -> np.ndarray: """ - Worker: subtract the spline bandpass from one coarse channel and return the cleaned slice - of shape (time_bins, coarse_channel_width). args is (channel_index, coarse_channel_width, - spl_order). - - Reads its slice from _GLOBAL_CHUNK_DATA — a (time_bins, n_coarse * coarse_channel_width) - view of the current block's shared memory, set up by _init_worker. + D'Agostino-Pearson normality statistic (k2) for every sliding window across one coarse + channel, computed in closed form from per-block power sums instead of a per-window Python + loop over scipy.stats.normaltest (which was the #1 cost of inference preprocessing). + + channel has shape (time_bins, width). Window j covers columns + [j*step_size, j*step_size + window_size) — the same windows as the historical loop + `range(0, width - window_size, step_size)` — and each window's sample is the flattened + (time_bins, window_size) slice, so n = time_bins * window_size is constant across windows. + Returns a float64 array of shape (n_windows,) whose entries match + `scipy.stats.normaltest(window.flatten()).statistic` (unit tests pin the equivalence to + rtol=1e-9); windows with zero variance yield NaN, matching scipy's behavior of returning + NaN rather than a spurious statistic. + + Derivation: normaltest = skewtest.Z**2 + kurtosistest.Z**2 where both Z transforms are + elementwise closed forms in (n, m2, m3, m4). The central moments come from per-block raw + power sums S1..S4 accumulated in float64: blocks are sized so that every window is an + exact run of adjacent blocks (block = step_size when step divides window — the fast + path — else gcd(window, step)), window sums are length-(window//block) moving sums over + the block sums, and m2/m3/m4 follow from the standard raw-to-central-moment identities. + Raw-moment differencing loses precision when |mean| >> std, so the channel mean is + subtracted up front (central moments are shift-invariant); bandpass-subtracted residuals + are near-zero-mean anyway, and the rtol=1e-9 unit-test gate is the arbiter. """ - channel_index, coarse_channel_width, spl_order = args + time_bins, width = channel.shape + n = time_bins * window_size # samples per window (scalar across all windows) + if n < 8: + # Mirror scipy.stats.skewtest's minimum-sample requirement + raise ValueError(f"normality test requires >= 8 samples per window, got {n}") + n_windows = len(range(0, width - window_size, step_size)) + if n_windows <= 0: + return np.empty(0, dtype=np.float64) + + # astype always copies, so the in-place shift below can't mutate the caller's array. + # Shift-invariance guard: subtracting the channel mean leaves every central moment + # mathematically unchanged while keeping the S2/n - mean**2 style differencing below + # well conditioned even when the residuals carry a DC offset. + data = channel.astype(np.float64) + data -= data.mean() + + # Per-column power sums over the time axis, accumulated row by row so temporaries stay at + # (width,) rather than (time_bins, width) x 3. + # NOTE: the Python-level loop over time_bins rows is deliberate — it caps peak memory at a + # few (width,) float64 vectors per worker. With the default time_bins=16 the loop overhead + # is negligible; if time_bins ever grows large, vectorize via (channel**k).sum(axis=0). + p1 = np.zeros(width, dtype=np.float64) + p2 = np.zeros(width, dtype=np.float64) + p3 = np.zeros(width, dtype=np.float64) + p4 = np.zeros(width, dtype=np.float64) + for t in range(time_bins): + row = data[t] + row2 = row * row + p1 += row + p2 += row2 + p3 += row2 * row + p4 += row2 * row2 + + # Aggregate columns into block sums. The fast path (step divides window — the default + # config: window 256, step 128) uses step-sized blocks; the general path falls back to + # gcd-sized blocks so windows are still exact runs of adjacent blocks. + block = step_size if window_size % step_size == 0 else math.gcd(window_size, step_size) + edges = np.arange(0, width, block) + blocks_per_window = window_size // block + blocks_per_step = step_size // block + # When block does not divide width (non-power-of-2 geometries), np.add.reduceat emits a + # short trailing block spanning the ragged tail [edges[-1], width). No in-range window ever + # reaches it — every window ends at a multiple of block that is < width, hence <= edges[-1] + # — but drop it explicitly so a partial sum can never leak into a window and silently + # corrupt k2. width // block == number of full blocks (== len(edges) when block divides + # width, so this is a no-op there); slicing to it keeps every window sum exact. + n_full_blocks = width // block + + def _window_sums(col_sums: np.ndarray) -> np.ndarray: + block_sums = np.add.reduceat(col_sums, edges)[:n_full_blocks] + # Moving sum of blocks_per_window adjacent blocks, sampled every blocks_per_step + # blocks — one entry per window, no long cumulative accumulation (precision). + view = np.lib.stride_tricks.sliding_window_view(block_sums, blocks_per_window) + return view[::blocks_per_step].sum(axis=-1)[:n_windows] + + s1 = _window_sums(p1) + s2 = _window_sums(p2) + s3 = _window_sums(p3) + s4 = _window_sums(p4) + + # Raw power sums -> central moments (float64 throughout) + mean = s1 / n + m2 = s2 / n - mean**2 + m3 = s3 / n - 3.0 * mean * (s2 / n) + 2.0 * mean**3 + m4 = s4 / n - 4.0 * mean * (s3 / n) + 6.0 * mean**2 * (s2 / n) - 3.0 * mean**4 + + # Z transforms transcribed from scipy.stats._stats_py::skewtest / kurtosistest with + # scalar n, so every n-dependent constant is computed once. Variable names follow scipy. + with np.errstate(all="ignore"): + # Zero-variance windows can't support either test; scipy returns NaN there (and a + # tiny negative m2 from float cancellation would otherwise fabricate a huge k2). + degenerate = m2 <= 0.0 + m2 = np.where(degenerate, np.nan, m2) + + # --- skewtest: Z1 from g1 = m3 / m2**1.5 + b2 = m3 / m2**1.5 + y = b2 * math.sqrt(((n + 1) * (n + 3)) / (6.0 * (n - 2))) + beta2 = ( + 3.0 + * (n**2 + 27 * n - 70) + * (n + 1) + * (n + 3) + / ((n - 2.0) * (n + 5) * (n + 7) * (n + 9)) + ) + w2 = -1 + math.sqrt(2 * (beta2 - 1)) + delta = 1 / math.sqrt(0.5 * math.log(w2)) + alpha = math.sqrt(2.0 / (w2 - 1)) + y = np.where(y == 0, 1.0, y) + z1 = delta * np.log(y / alpha + np.sqrt((y / alpha) ** 2 + 1)) + + # --- kurtosistest: Z2 from b2 = m4 / m2**2 + b2k = m4 / m2**2 + e = 3.0 * (n - 1) / (n + 1) + varb2 = 24.0 * n * (n - 2) * (n - 3) / ((n + 1) * (n + 1.0) * (n + 3) * (n + 5)) + x = (b2k - e) / math.sqrt(varb2) + sqrtbeta1 = ( + 6.0 + * (n * n - 5 * n + 2) + / ((n + 7) * (n + 9)) + * math.sqrt((6.0 * (n + 3) * (n + 5)) / (n * (n - 2) * (n - 3))) + ) + a = 6.0 + 8.0 / sqrtbeta1 * (2.0 / sqrtbeta1 + math.sqrt(1 + 4.0 / (sqrtbeta1**2))) + term1 = 1 - 2 / (9.0 * a) + denom = 1 + x * math.sqrt(2 / (a - 4.0)) + term2 = np.sign(denom) * np.where( + denom == 0.0, np.nan, np.power((1 - 2.0 / a) / np.abs(denom), 1 / 3.0) + ) + z2 = (term1 - term2) / math.sqrt(2 / (9.0 * a)) - if _GLOBAL_CHUNK_DATA is None: - logger.warning("No global chunk data available for bandpass removal") - return np.zeros((0, coarse_channel_width)) + k2 = z1 * z1 + z2 * z2 - start = channel_index * coarse_channel_width - end = (channel_index + 1) * coarse_channel_width - channel = _GLOBAL_CHUNK_DATA[:, start:end] - integrated_channel = np.mean(channel, axis=0) - fit = _fit_channel_bandpass(integrated_channel, coarse_channel_width, spl_order) - return channel - fit + return k2 -# NOTE: come back to this later -def _threshold_hits_worker(args: tuple) -> list[tuple]: +def _energy_detect_channel_worker(args: tuple) -> list[tuple]: """ - Worker: slide a window across one coarse channel and emit hits whose D'Agostino-Pearson - normality statistic exceeds stat_threshold. Returns a list of (absolute_fine_channel_index, - statistic, pvalue) tuples. - - args is (channel_index, coarse_channel_width, window_size, step_size, stat_threshold, - block_offset). block_offset is added to each emitted index so callers see positions in the - full spectrum rather than block-relative ones. Reads its slice from _GLOBAL_CHUNK_DATA, which - holds the cleaned residuals for the current block. + Fused worker: run the complete energy-detection chain for one coarse channel — read the + channel's h5 slice, remove the DC spike, flatten the bandpass, and threshold the vectorized + normality statistic. Returns a small list of (absolute_fine_channel_index, statistic, + pvalue) hit tuples; the bulky (time_bins, coarse_channel_width) intermediates never leave + the worker, so no shared memory or per-block parent arrays are needed. + + args is (h5_path, channel_index, coarse_channel_width, time_bins, bandpass_flatten, + window_size, step_size, stat_threshold). bandpass_flatten is a picklable callable + (channel) -> residuals (see _spline_flatten_bandpass). Each worker opens its own h5py.File + since h5py file handles are fork-unsafe to share. """ ( + h5_path, channel_index, coarse_channel_width, + time_bins, + bandpass_flatten, window_size, step_size, stat_threshold, - block_offset, ) = args - if _GLOBAL_CHUNK_DATA is None: - logger.warning("No global chunk data available for hit thresholding") - return [] - start = channel_index * coarse_channel_width end = (channel_index + 1) * coarse_channel_width - channel = _GLOBAL_CHUNK_DATA[:, start:end] - - # TODO: profile on HPC and consider vectorizing. With default config this loop - # runs ~8190 windows per coarse channel × parallel_chans × num_blocks × 3 - # ON-source files per cadence. scipy.stats.normaltest copies + flattens each - # window and re-derives skew/kurtosis via separate scipy calls. A stride_tricks - # view over the cleaned block followed by vectorized scipy.stats.skew / - # kurtosis (with axis arg) would replace the Python loop with a few large - # ndarray ops and could be substantially faster end-to-end. - hits: list[tuple] = [] - for i in range(0, coarse_channel_width - window_size, step_size): - window = channel[:, i : i + window_size] - s, p = stats.normaltest(window.flatten()) - if s > stat_threshold: - abs_idx = block_offset + channel_index * coarse_channel_width + i - hits.append((abs_idx, float(s), float(p))) - - return hits + with h5py.File(h5_path, "r") as hf: + channel = hf["data"][:time_bins, 0, start:end] + + # In-place DC spike removal. The spike sits at the channel center and the interpolation + # offsets reach at most ±3 bins around it (see _remove_dc_spike), so per-channel + # processing touches exactly the same bins as the historical block-based path — the + # offsets can never cross a coarse-channel boundary. + _remove_dc_spike(channel, coarse_channel_width, 1) + + residuals = bandpass_flatten(channel) + + k2 = _sliding_normality_k2(residuals, window_size, step_size) + hit_windows = np.nonzero(k2 > stat_threshold)[0] + if hit_windows.size == 0: + return [] + + # p-values are only stored in metadata, so chi2.sf runs on the (small) hit subset only + pvalues = stats.chi2.sf(k2[hit_windows], 2) + return [ + (start + int(j) * step_size, float(k2[j]), float(p)) + for j, p in zip(hit_windows, pvalues, strict=True) + ] def _extract_stamps_worker(args: tuple) -> None: @@ -274,18 +435,34 @@ def _extract_stamps_worker(args: tuple) -> None: Each worker opens its own hdf5 handle and its own r+ view of the shared .npy, then copies a stamp_width-wide window (over the first `time_bins` rows, polarization 0) around - each hit. Tasks address disjoint output regions — distinct obs_idx and/or non-overlapping - stamp indices — so concurrent writes from the pool never collide. `stamp_starts` is the - contiguous, start-sorted slice for this task; `base_idx` is its offset into the full stamp - list so the worker writes to the correct absolute rows. + each hit. When downsample_factor > 1 the stamp is downsampled along frequency with + downscale_local_mean before writing, so the memmap stores width + stamp_width // downsample_factor — this is the same operation load_inference_data used to + apply after the fact, moved to extraction time to cut storage by the same factor. Tasks + address disjoint output regions — distinct obs_idx and/or non-overlapping stamp indices — + so concurrent writes from the pool never collide. `stamp_starts` is the contiguous, + start-sorted slice for this task; `base_idx` is its offset into the full stamp list so the + worker writes to the correct absolute rows. """ - npy_path, obs_idx, obs_h5, stamp_starts, base_idx, time_bins, stamp_width = args + ( + npy_path, + obs_idx, + obs_h5, + stamp_starts, + base_idx, + time_bins, + stamp_width, + downsample_factor, + ) = args out = np.lib.format.open_memmap(npy_path, mode="r+") try: with h5py.File(obs_h5, "r") as hf: dset = hf["data"] for local_i, start in enumerate(stamp_starts): - out[base_idx + local_i, obs_idx] = dset[:time_bins, 0, start : start + stamp_width] + stamp = dset[:time_bins, 0, start : start + stamp_width] + if downsample_factor > 1: + stamp = downscale_local_mean(stamp, (1, downsample_factor)).astype(np.float32) + out[base_idx + local_i, obs_idx] = stamp out.flush() finally: del out @@ -334,6 +511,59 @@ class CadenceHit: frequency_mhz: float = field(default=float("nan")) +@dataclass +class PendingCadence: + """One unit of preprocessing work: a valid cadence group plus its target .npy path.""" + + group: CadenceGroup + npy_path: str + + +def derive_cadence_provenance(key: tuple, group_by_cols: list[str], metadata: dict) -> dict: + """ + Map one cadence's group-by key and stamp metadata JSON onto the observational-provenance + fields of the inference_results table. + + key and group_by_cols come from the CadenceGroup (values zipped positionally onto column + names, matched case-insensitively so 'Target'/'target' both resolve); metadata is the + per-cadence JSON written by _process_cadence. Returns a dict with keys target, session, + band, cadence_id (int when parseable, else None), timestamp_observed (the header's tstart, + when present), h5_path (first observation of the cadence), and stamp_frequencies_mhz (the + per-stamp center frequencies, one per snippet row in the .npy). + """ + # strict=False: a malformed key/cols pairing degrades to sparse provenance rather than + # aborting the cadence + key_map = {str(col).strip().lower(): val for col, val in zip(group_by_cols, key, strict=False)} + + cadence_id: int | None = None + raw_cadence_id = key_map.get("cadence id") + if raw_cadence_id is not None: + try: + cadence_id = int(raw_cadence_id) + except (TypeError, ValueError): + logger.warning(f"Could not parse cadence id {raw_cadence_id!r} as int; storing None") + + header = metadata.get("header") or {} + timestamp_observed: float | None = None + if "tstart" in header: + try: + timestamp_observed = float(header["tstart"]) + except (TypeError, ValueError): + logger.warning(f"Could not parse header tstart {header['tstart']!r} as float") + + h5_paths = metadata.get("h5_paths") or [] + + return { + "target": key_map.get("target"), + "session": key_map.get("session"), + "band": key_map.get("band"), + "cadence_id": cadence_id, + "timestamp_observed": timestamp_observed, + "h5_path": h5_paths[0] if h5_paths else None, + "stamp_frequencies_mhz": metadata.get("stamp_frequencies_mhz"), + } + + # 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?) def group_observations_from_csv( csv_path: str, @@ -424,6 +654,10 @@ def __init__(self): if self.manager is None: raise ValueError("get_manager() returned None") + # Persistent worker pool for energy detection + stamp extraction, shared across all + # cadences of a run (see start/stop_energy_detection_pool) + self._ed_pool: Pool | None = None + # NOTE: come back to this later def close(self): """Explicitly close the multiprocessing pool and shared memory""" @@ -435,6 +669,10 @@ def close(self): # self.manager.close_shared_memory(self.shm) # self.shm = None + # Defense-in-depth: ensure the persistent energy-detection pool is torn down even if a + # caller forgot to call stop_energy_detection_pool(). No-op when no pool was started. + self.stop_energy_detection_pool() + logger.info("DataPreprocessor closed") # NOTE: shared resources currently created & destroyed within function itself. think about abstractions once preprocessing.py is complete @@ -629,16 +867,24 @@ def load_inference_data(self, override_filepaths: list[str] | None = None) -> np """ Load and preprocess inference data into an array of shape (n, 6, 16, width_bin_downsampled). Uses a multiprocessing pool over shared memory to - downsample cadences in parallel, then applies per-cadence log-normalization in-process. + process cadences in parallel. + + Each file's stored width decides its path: files already at the downsampled width + (written by _extract_stamps_worker with store_downsampled_stamps enabled) only need + per-cadence log-normalization, which runs vectorized in the pool workers; legacy + full-width files (width_bin, e.g. stamps preprocessed before downsample-at-extraction + landed) keep the historical downsample-then-log-norm behavior. override_filepaths, when given, supplies absolute paths to iterate directly instead of - resolving config.data.test_files via get_test_file_path — used by find_hits() to chain - per-cadence .npy outputs into inference without monkey-patching paths. + resolving config.data.test_files via get_test_file_path — used by the inference command + to chain per-cadence .npy outputs from preprocessing into inference without + monkey-patching paths. """ logger.info(f"Loading backgrounds from {self.config.data_path} for inference") downsample_factor = self.config.data.downsample_factor - final_width = self.config.data.width_bin // downsample_factor + width_bin = self.config.data.width_bin + final_width = width_bin // downsample_factor chunk_size = self.config.data.inference_background_load_chunk_size n_processes = self.config.manager.n_processes @@ -691,6 +937,25 @@ def load_inference_data(self, override_filepaths: list[str] | None = None) -> np logger.info(f" Raw data shape: {raw_data.shape}") + # Branch on the stored width: already-downsampled stamps (written by + # _extract_stamps_worker) only need log-norm; legacy full-width files keep the + # historical downsample-then-log-norm path. + stored_width = raw_data.shape[-1] + if stored_width == final_width: + already_downsampled = True + logger.info(f" {filename} stored at final width {final_width}: log-norm only") + elif stored_width == width_bin: + already_downsampled = False + logger.info(f" {filename} stored at full width {width_bin}: downsample + log-norm") + else: + logger.error( + f" {filename} width {stored_width} matches neither width_bin " + f"({width_bin}) nor width_bin // downsample_factor ({final_width}); " + f"skipping file" + ) + del raw_data + continue + # Divide background into equal chunks, then cutoff if exceeds max_chunks n_cadences_total = raw_data.shape[0] n_chunks = (n_cadences_total + chunk_size - 1) // chunk_size @@ -706,16 +971,17 @@ def load_inference_data(self, override_filepaths: list[str] | None = None) -> np # NOTE: is this access pattern the most efficient (least pickling)? see comments in _single_cadence_wrapper() from data_generation.py for more details # NOTE: currently, loading the backgrounds takes WAY more time than processing the backgrounds - # Prepare arguments for downsampling (just indices, not data - data is in global state) + # Prepare arguments (just indices, not data - data is in global state). + # Downsampled files fold log-norm into the workers (resolving the old + # "downsample & log-norm simultaneously" TODO); legacy files downsample in + # the workers and log-norm in-process below, exactly as before. n_cadences = chunk_data.shape[0] - args_list = [ - ( - i, - downsample_factor, - final_width, - ) # Just pass the chunk index, not the full cadence data - for i in range(n_cadences) - ] + if already_downsampled: + worker_fn = _lognorm_worker + args_list = [(i,) for i in range(n_cadences)] + else: + worker_fn = _downsample_worker + args_list = [(i, downsample_factor, final_width) for i in range(n_cadences)] # NOTE: do we need to create & destroy the pool every chunk? or just the shared memory & pass new references in? is there a differenc? if n_processes > 1: @@ -749,15 +1015,7 @@ def load_inference_data(self, override_filepaths: list[str] | None = None) -> np # NOTE: should we use separate chunks_per_worker? how to benchmark? chunksize = max(1, n_cadences // (n_workers * chunks_per_worker)) # TEST: does return order matter? - results = chunk_pool.map( - _downsample_worker, # TODO: create separate function that performs downsampling & log-norm simultaneously - args_list, - chunksize=chunksize, - ) - # results = chunk_pool.imap(_downsample_worker, args_list, chunksize=chunksize) - # results = chunk_pool.imap_unordered( - # _downsample_worker, args_list, chunksize=chunksize - # ) + results = chunk_pool.map(worker_fn, args_list, chunksize=chunksize) else: # Sequential processing @@ -771,14 +1029,17 @@ def load_inference_data(self, override_filepaths: list[str] | None = None) -> np shared_chunk = chunk_data _GLOBAL_CHUNK_DATA = shared_chunk - # TODO: create separate function that performs downsampling & log-norm simultaneously - results = [_downsample_worker(args) for args in args_list] + results = [worker_fn(args) for args in args_list] # NOTE: is there a more efficient/elegant way to do this (e.g. with list comprehension/slicing)? # Collect valid results (filter out None from invalid cadences) for result in results: if result is not None: - all_cadences.append(result) + if already_downsampled: + all_cadences.append(result) + else: + # Legacy path: per-cadence log-norm in-process, as before + all_cadences.append(log_norm(result)) # Clear chunk data & shared resources del chunk_data, shared_chunk @@ -798,22 +1059,13 @@ def load_inference_data(self, override_filepaths: list[str] | None = None) -> np if len(all_cadences) == 0: raise ValueError("No data loaded successfully") - # Stack all_cadences together + # Stack all_cadences together (every cadence is downsampled + log-normalized by now) cadence_array = np.array(all_cadences, dtype=np.float32) # Clear all_cadences reference del all_cadences gc.collect() - logger.info(f"Total cadences loaded: {cadence_array.shape[0]}") - logger.info(f"Cadence array shape before log norm: {cadence_array.shape}") - - # TODO: create separate function that performs downsampling & log-norm simultaneously - # Apply log normalization to each cadence - logger.info("Applying log normalization") - for i in range(cadence_array.shape[0]): - cadence_array[i] = log_norm(cadence_array[i]) - # Sanity check: print descriptive stats min_val = np.min(cadence_array) max_val = np.max(cadence_array) @@ -842,22 +1094,41 @@ def load_inference_data(self, override_filepaths: list[str] | None = None) -> np return cadence_array - # NOTE: come back to this later (based on docstring, we're processing cadences sequentially. if so, any way to parallelize?) - def find_hits(self) -> list[CadenceResult]: + def start_energy_detection_pool(self) -> None: """ - Convert raw .h5 cadence observations into (n_hits, 6, 16, stamp_width) .npy snippets, - returning one CadenceResult per successfully processed (or already cached) cadence. + Create the persistent worker pool used for energy detection and stamp extraction. - Driven by CSVs in config.data.inference_files. Each CSV is grouped into cadences via - group_observations_from_csv() and processed sequentially. Within each cadence, energy - detection runs on ON-source files (positions 0, 2, 4 in ABACAD); stamps are then - extracted from all 6 observations at the hit frequencies. Each cadence is checkpointed - to disk as soon as its .npy is ready, and on retry, cadences whose output already exists - are skipped. + One pool serves every cadence of the run (no per-block or per-cadence pool churn). + Call from the main thread before any cadence processing starts — forking from the main + thread before background threads spin up avoids inheriting mid-operation locks — and + pair with stop_energy_detection_pool() when the run is done. No-op when a pool already + exists or n_processes == 1 (sequential mode). + """ + if self._ed_pool is not None: + return + n_processes = self.config.manager.n_processes + if n_processes > 1: + self._ed_pool = self.manager.create_pool( + n_processes=n_processes, + name="DataPreproc_energy_detection", + initializer=_init_plain_worker, + ) + + def stop_energy_detection_pool(self) -> None: + """Close the persistent energy-detection pool (no-op if never started).""" + if self._ed_pool is not None: + self.manager.close_pool(self._ed_pool) + self._ed_pool = None + + def plan_cadences(self) -> list[PendingCadence]: + """ + Group every CSV in config.data.inference_files into per-cadence work units without + processing anything. Returns one PendingCadence (valid group + target .npy path) per + valid cadence, in CSV order — the unit list the streaming inference loop iterates. """ inference_files = self.config.data.inference_files if not inference_files: - logger.warning("find_hits() called with no inference_files configured") + logger.warning("plan_cadences() called with no inference_files configured") return [] # Resolve output directory @@ -871,7 +1142,7 @@ def find_hits(self) -> list[CadenceResult]: h5_path_col = self.config.inference.cadence_h5_path_col expected_obs = self.config.inference.cadence_expected_obs - results: list[CadenceResult] = [] + units: list[PendingCadence] = [] for csv_filename in inference_files: csv_path = self.config.get_inference_file_path(csv_filename) @@ -896,49 +1167,88 @@ def find_hits(self) -> list[CadenceResult]: for group in valid_groups: npy_filename = self._cadence_npy_filename(csv_stem, group.key) - npy_path = os.path.join(output_dir, npy_filename) + units.append( + PendingCadence(group=group, npy_path=os.path.join(output_dir, npy_filename)) + ) - if os.path.exists(npy_path): - # Resume path: rebuild a minimal CadenceResult from the existing file - metadata_path = self._cadence_metadata_path(npy_path) - try: - existing = np.load(npy_path, mmap_mode="r") - n_hits = existing.shape[0] - del existing - except Exception as e: - logger.warning( - f"Existing .npy at {npy_path} could not be inspected ({e}); " - f"reprocessing cadence" - ) - # Fall through (no continue) to the _process_cadence call - # below so a corrupted .npy gets regenerated rather than - # silently skipped. - else: - logger.info( - f"Skipping cadence {group.key}: {npy_path} already exists " - f"({n_hits} hits)" - ) - results.append( - CadenceResult( - npy_path=npy_path, - h5_paths=group.h5_paths, - key=group.key, - n_hits=n_hits, - metadata_path=metadata_path, - ) - ) - continue - - try: - cadence_result = self._process_cadence(group, npy_path) - except Exception as e: - # Single-cadence failures should not abort the whole CSV; - # the retry loop at the inference_command level handles broader recovery - logger.error(f"Failed to process cadence {group.key}: {e}") - continue + logger.info( + f"Planned {len(units)} cadence work unit(s) across {len(inference_files)} CSV(s)" + ) + return units + + def process_pending_cadence(self, unit: PendingCadence) -> CadenceResult | None: + """ + Resume-or-process one cadence work unit. Returns a CadenceResult when a stamp .npy is + available (freshly written or already on disk), or None when the cadence produced no + hits or failed — single-cadence failures are logged and swallowed so one bad cadence + can't abort the whole catalog (the retry loop at the inference_command level handles + broader recovery). + """ + group, npy_path = unit.group, unit.npy_path + + if os.path.exists(npy_path): + # Resume path: rebuild a minimal CadenceResult from the existing file + metadata_path = self._cadence_metadata_path(npy_path) + try: + existing = np.load(npy_path, mmap_mode="r") + n_hits = existing.shape[0] + del existing + except Exception as e: + logger.warning( + f"Existing .npy at {npy_path} could not be inspected ({e}); " + f"reprocessing cadence" + ) + # Fall through (no return) to the _process_cadence call below so a + # corrupted .npy gets regenerated rather than silently skipped. + else: + logger.info( + f"Skipping cadence {group.key}: {npy_path} already exists ({n_hits} hits)" + ) + return CadenceResult( + npy_path=npy_path, + h5_paths=group.h5_paths, + key=group.key, + n_hits=n_hits, + metadata_path=metadata_path, + ) + + try: + return self._process_cadence(group, npy_path) + except Exception as e: + logger.error(f"Failed to process cadence {group.key}: {e}") + return None + + # NOTE: come back to this later (based on docstring, we're processing cadences sequentially. if so, any way to parallelize?) + def find_hits(self) -> list[CadenceResult]: + """ + Convert raw .h5 cadence observations into (n_hits, 6, 16, stored_width) .npy snippets, + returning one CadenceResult per successfully processed (or already cached) cadence. + + Driven by CSVs in config.data.inference_files. Each CSV is grouped into cadences via + plan_cadences() and processed sequentially over one persistent worker pool. Within each + cadence, energy detection runs on ON-source files (positions 0, 2, 4 in ABACAD); stamps + are then extracted from all 6 observations at the hit frequencies. Each cadence is + checkpointed to disk as soon as its .npy is ready, and on retry, cadences whose output + already exists are skipped. + + The streaming inference path in main.py drives plan_cadences() / + process_pending_cadence() directly instead (so preprocessing of cadence i+1 can overlap + inference of cadence i); this wrapper preserves the batch contract for callers that + want every cadence preprocessed up front. + """ + units = self.plan_cadences() + if not units: + return [] + results: list[CadenceResult] = [] + self.start_energy_detection_pool() + try: + for unit in units: + cadence_result = self.process_pending_cadence(unit) if cadence_result is not None: results.append(cadence_result) + finally: + self.stop_energy_detection_pool() logger.info(f"find_hits completed: {len(results)} cadence .npy files available") return results @@ -990,19 +1300,28 @@ def _process_cadence(self, group: CadenceGroup, npy_path: str) -> CadenceResult Energy detection runs only on ON-source observations (positions 0, 2, 4 in ABACAD order); the hits found there define the frequency slices that get extracted from all 6 observations. group is assumed to be a validated CadenceGroup with len(h5_paths) == - expected_obs. + expected_obs. Each coarse channel is one fused task (read -> DC spike -> bandpass + flatten -> vectorized threshold) on the persistent pool from + start_energy_detection_pool(); only the small per-channel hit lists cross the process + boundary, so no blocks, block-sized shared memory, or per-stage pools exist anymore. """ coarse_channel_width = self.config.inference.coarse_channel_width + # parallel_coarse_chans is a progress-logging chunk size only (None -> n_processes); + # actual parallelism is the persistent pool's worker count. parallel_chans = self.config.inference.parallel_coarse_chans - spl_order = self.config.inference.spline_order window_size = self.config.inference.detection_window_size step_size = self.config.inference.detection_step_size stat_threshold = self.config.inference.stat_threshold stamp_width = self.config.inference.stamp_width overlap_search = self.config.inference.overlap_search overlap_fraction = self.config.inference.overlap_fraction + store_downsampled = self.config.inference.store_downsampled_stamps + downsample_factor = self.config.data.downsample_factor if store_downsampled else 1 time_bins = self.config.data.time_bins n_processes = self.config.manager.n_processes + progress_chunk = max(1, parallel_chans if parallel_chans is not None else n_processes) + + bandpass_flatten = self._get_bandpass_flattener() # Read header / metadata from the first ON-source file on_source_paths = [group.h5_paths[i] for i in (0, 2, 4)] @@ -1025,18 +1344,21 @@ def _process_cadence(self, group: CadenceGroup, npy_path: str) -> CadenceResult ) return None - num_blocks = n_chans // (coarse_channel_width * parallel_chans) - if num_blocks == 0: + # NOTE: every complete coarse channel is processed. The historical block-based path + # floored to a multiple of parallel_coarse_chans, silently dropping up to + # parallel_coarse_chans - 1 trailing coarse channels when n_chans wasn't an exact + # multiple of a block. + n_coarse_total = n_chans // coarse_channel_width + if n_coarse_total == 0: logger.warning( - f"Cadence {group.key}: n_chans={n_chans} is smaller than one block " - f"({coarse_channel_width * parallel_chans}); skipping" + f"Cadence {group.key}: n_chans={n_chans} is smaller than one coarse channel " + f"({coarse_channel_width}); skipping" ) return None - block_width = coarse_channel_width * parallel_chans logger.info( - f"Cadence {group.key}: n_chans={n_chans}, num_blocks={num_blocks}, " - f"block_width={block_width}, ON-source files={len(on_source_paths)}" + f"Cadence {group.key}: n_chans={n_chans}, coarse channels={n_coarse_total}, " + f"ON-source files={len(on_source_paths)}" ) # Aggregate hits across all ON-source files @@ -1048,47 +1370,39 @@ def _process_cadence(self, group: CadenceGroup, npy_path: str) -> CadenceResult f"{on_source_idx + 1}/{len(on_source_paths)}: {on_h5}" ) - for block_num in range(num_blocks): - block_offset = block_num * block_width - logger.info( - f" Block {block_num + 1}/{num_blocks} " - f"(coarse {block_num * parallel_chans}..{(block_num + 1) * parallel_chans - 1})" - ) - - block_data = self._read_block( - on_h5, block_num, parallel_chans, coarse_channel_width, n_processes - ) - - # Slice to first time_bins - block_data = block_data[:time_bins] - - # In-place DC spike removal - _remove_dc_spike(block_data, coarse_channel_width, parallel_chans) - - # _drop_side_channels(block_data, side_channel_count, coarse_channel_width) - - cleaned_block = self._remove_block_bandpass( - block_data, parallel_chans, coarse_channel_width, spl_order, n_processes - ) - - # Free original block memory; we only need the cleaned residuals downstream - del block_data - - block_hits = self._threshold_block_hits( - cleaned_block, - parallel_chans, + tasks = [ + ( + on_h5, + ch, coarse_channel_width, + time_bins, + bandpass_flatten, window_size, step_size, stat_threshold, - block_offset, - n_processes, ) + for ch in range(n_coarse_total) + ] - all_hits.extend(block_hits) - - del cleaned_block - gc.collect() + if self._ed_pool is not None: + # imap (ordered, chunksize 1) keeps every worker busy across the whole file + # while results stream back for progress logging + channel_hits_iter = self._ed_pool.imap(_energy_detect_channel_worker, tasks) + else: + if n_processes > 1: + logger.info( + "Energy detection running sequentially: no persistent pool started " + "(call start_energy_detection_pool() to parallelize)" + ) + channel_hits_iter = map(_energy_detect_channel_worker, tasks) + + for done, channel_hits in enumerate(channel_hits_iter, start=1): + all_hits.extend(channel_hits) + if done % progress_chunk == 0 or done == n_coarse_total: + logger.info( + f" Coarse channel {done}/{n_coarse_total} of ON-source " + f"{on_source_idx + 1}/{len(on_source_paths)}" + ) logger.info(f"Cadence {group.key}: {len(all_hits)} raw hits across ON-source files") @@ -1139,22 +1453,38 @@ def _process_cadence(self, group: CadenceGroup, npy_path: str) -> CadenceResult # over all 6 observations (single-threaded reads + bitshuffle chunk decompression) # was the dominant, GPU-idle cost of CSV inference. # + # Stamps are downsampled along frequency at extraction time (downsample_factor > 1, + # the store_downsampled_stamps default), so the stored width is stamp_width // + # downsample_factor — an ~8x storage cut at defaults that also removes the separate + # downsample pass from load_inference_data. + # # Atomicity is unchanged: the memmap is written to a .tmp sibling and we os.replace() - # it onto the canonical name only after every worker finishes, so find_hits' resume - # path (which treats npy_path's existence as proof of a complete write) still holds. + # it onto the canonical name only after every worker finishes, so the resume path + # (which treats npy_path's existence as proof of a complete write) still holds. n_stamps = len(stamp_centers) + stored_width = stamp_width // downsample_factor tmp_npy_path = os.path.splitext(npy_path)[0] + ".tmp.npy" + # A leftover .tmp.npy means a previous attempt died (e.g. SIGKILL) between memmap + # creation and os.replace; it was never promoted to npy_path, so it's safe to drop + # (open_memmap would truncate it anyway — the warning is the point) + if os.path.exists(tmp_npy_path): + logger.warning( + f"Cadence {group.key}: removing stale partial output {tmp_npy_path} " + f"from an interrupted previous run" + ) + os.remove(tmp_npy_path) + # NOTE: np.lib.format.open_memmap is a semi-public numpy API (stable across 1.x and + # documented via np.lib.format); revisit if a future numpy bump moves it memmap = np.lib.format.open_memmap( tmp_npy_path, mode="w+", dtype=np.float32, - shape=(n_stamps, len(group.h5_paths), time_bins, stamp_width), + shape=(n_stamps, len(group.h5_paths), time_bins, stored_width), ) memmap.flush() del memmap # header + full-size file are on disk; workers reopen it in r+ mode stamp_starts = [start for start, _, _ in stamp_centers] - n_processes = self.config.manager.n_processes # Split each obs file's (already start-sorted) stamps into contiguous chunks so more # than len(h5_paths) workers can run, while keeping each worker's reads sequential to # preserve the hdf5 chunk-cache reuse that the sort above buys us. @@ -1169,20 +1499,16 @@ def _process_cadence(self, group: CadenceGroup, npy_path: str) -> CadenceResult base, time_bins, stamp_width, + downsample_factor, ) for obs_idx, obs_h5 in enumerate(group.h5_paths) for base in range(0, n_stamps, chunk_size) ] - if n_processes > 1 and len(tasks) > 1: - pool = self.manager.create_pool( - n_processes=min(len(tasks), n_processes), - name="DataPreproc_extract_stamps", - ) - try: - pool.map(_extract_stamps_worker, tasks) - finally: - self.manager.close_pool(pool) + if self._ed_pool is not None and len(tasks) > 1: + # Reuse the persistent energy-detection pool — extraction workers are plain + # (no shared memory), so the same pool serves both stages without churn + self._ed_pool.map(_extract_stamps_worker, tasks) else: for task in tasks: _extract_stamps_worker(task) @@ -1202,6 +1528,8 @@ def _process_cadence(self, group: CadenceGroup, npy_path: str) -> CadenceResult "header": header, "stamp_starts": [int(start) for start, _, _ in stamp_centers], "stamp_width": stamp_width, + "stored_width": stored_width, + "downsample_factor_applied": downsample_factor, "stamp_frequencies_mhz": stamp_freqs_mhz, "stamp_statistics": stamp_stats, "stamp_pvalues": stamp_pvals, @@ -1228,133 +1556,20 @@ def _process_cadence(self, group: CadenceGroup, npy_path: str) -> CadenceResult metadata_path=metadata_path, ) - def _read_block( - self, - h5_path: str, - block_num: int, - parallel_chans: int, - coarse_channel_width: int, - n_processes: int, - ) -> np.ndarray: - """Read parallel_chans coarse channels in parallel and concatenate.""" - args_list = [ - (h5_path, ch, coarse_channel_width) - for ch in range(block_num * parallel_chans, (block_num + 1) * parallel_chans) - ] - - # NOTE: come back to this later (is this correct?) - # The read worker doesn't need shared memory; create a plain pool - if n_processes > 1: - pool = self.manager.create_pool( - n_processes=min(parallel_chans, n_processes), - name=f"DataPreproc_read_block_{block_num}", # NOTE: come back to this later - ) - try: - # NOTE: come back to this later (is pool.map correct here?) - results = pool.map(_read_coarse_channel_worker, args_list) - finally: - self.manager.close_pool(pool) - else: - results = [_read_coarse_channel_worker(a) for a in args_list] - - return np.concatenate(results, axis=1) - - # NOTE: come back to this later - def _remove_block_bandpass( - self, - block_data: np.ndarray, - parallel_chans: int, - coarse_channel_width: int, - spl_order: int, - n_processes: int, - ) -> np.ndarray: - """Spline-fit + subtract bandpass per coarse channel, return cleaned block.""" - args_list = [(ch, coarse_channel_width, spl_order) for ch in range(parallel_chans)] - - if n_processes > 1: - shm = self.manager.create_shared_memory( - size=block_data.nbytes, - name="DataPreproc_bandpass_block", # NOTE: come back to this later - ) - shared = np.ndarray(block_data.shape, dtype=block_data.dtype, buffer=shm.buf) - shared[:] = block_data[:] - - pool = self.manager.create_pool( - n_processes=min(parallel_chans, n_processes), - name="DataPreproc_bandpass_block", # NOTE: come back to this later - initializer=_init_worker, - initargs=(shm.name, block_data.shape, block_data.dtype), - ) - try: - results = pool.map(_remove_bandpass_worker, args_list) - finally: - del shared - self.manager.close_shared_memory(shm) - self.manager.close_pool(pool) - else: - global _GLOBAL_CHUNK_DATA - _GLOBAL_CHUNK_DATA = block_data - try: - results = [_remove_bandpass_worker(a) for a in args_list] - finally: - # Always clear the global, even if a worker raised, so subsequent - # calls in the same process don't see stale state - _GLOBAL_CHUNK_DATA = None - - return np.concatenate(results, axis=1) + def _get_bandpass_flattener(self) -> Callable[[np.ndarray], np.ndarray]: + """ + Return the configured bandpass-flattening callable for energy detection. - # NOTE: come back to this later - def _threshold_block_hits( - self, - cleaned_block: np.ndarray, - parallel_chans: int, - coarse_channel_width: int, - window_size: int, - step_size: int, - stat_threshold: float, - block_offset: int, - n_processes: int, - ) -> list[tuple]: - """Sliding-window normality test across one cleaned block, return all hits.""" - args_list = [ - (ch, coarse_channel_width, window_size, step_size, stat_threshold, block_offset) - for ch in range(parallel_chans) - ] + The callable takes one coarse channel of shape (time_bins, coarse_channel_width) and + returns the flattened residuals; it must be picklable (a functools.partial over a + module-level function) so pool workers can receive it in their task args. Currently + only the spline flattener exists. - if n_processes > 1: - shm = self.manager.create_shared_memory( - size=cleaned_block.nbytes, - name="DataPreproc_threshold_block", # NOTE: come back to this later - ) - shared = np.ndarray(cleaned_block.shape, dtype=cleaned_block.dtype, buffer=shm.buf) - shared[:] = cleaned_block[:] - - pool = self.manager.create_pool( - n_processes=min(parallel_chans, n_processes), - name="DataPreproc_threshold_block", # NOTE: come back to this later - initializer=_init_worker, - initargs=(shm.name, cleaned_block.shape, cleaned_block.dtype), - ) - try: - results = pool.map(_threshold_hits_worker, args_list) - finally: - del shared - self.manager.close_shared_memory(shm) - self.manager.close_pool(pool) - else: - global _GLOBAL_CHUNK_DATA - _GLOBAL_CHUNK_DATA = cleaned_block - try: - results = [_threshold_hits_worker(a) for a in args_list] - finally: - # Always clear the global, even if a worker raised, so subsequent - # calls in the same process don't see stale state - _GLOBAL_CHUNK_DATA = None - - flat: list[tuple] = [] - for r in results: - flat.extend(r) - return flat + # NOTE: PR-07 adds a PFB static-equalization flattener here (selected via config) + """ + return functools.partial( + _spline_flatten_bandpass, spl_order=self.config.inference.spline_order + ) # 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?) @staticmethod diff --git a/tests/fixtures/ed_real_slice.npz b/tests/fixtures/ed_real_slice.npz new file mode 100644 index 0000000..557995f Binary files /dev/null and b/tests/fixtures/ed_real_slice.npz differ diff --git a/tests/unit/test_inference.py b/tests/unit/test_inference.py new file mode 100644 index 0000000..5ccaea3 --- /dev/null +++ b/tests/unit/test_inference.py @@ -0,0 +1,115 @@ +"""Unit tests for aetherscan.inference: the tail-padding fix in +prepare_distributed_inf_dataset (regression for the silent partial-batch drop) and the +InfDataHolder clear semantics.""" + +from __future__ import annotations + +import numpy as np +import pytest +import tensorflow as tf + +from aetherscan.inference import InfDataHolder, prepare_distributed_inf_dataset + + +def _collect_batches(result: dict) -> np.ndarray: + """Materialize inf_steps batches from the (infinite, repeated) distributed dataset.""" + iterator = iter(result["inf_dataset"]) + batches = [next(iterator).numpy() for _ in range(result["inf_steps"])] + return np.concatenate(batches, axis=0) + + +def _make_data(n: int) -> np.ndarray: + # Give every sample a unique signature so padding/ordering is verifiable + return np.arange(n, dtype=np.float32)[:, None, None, None] * np.ones( + (n, 6, 4, 8), dtype=np.float32 + ) + + +class TestPrepareDistributedInfDataset: + @pytest.fixture + def strategy(self): + return tf.distribute.get_strategy() # default no-op strategy (1 replica, CPU-safe) + + def test_partial_tail_batch_is_padded_not_dropped(self, strategy): + # Regression: 5 samples with global batch 2 used to yield inf_steps = 2 with + # drop_remainder=True — the 5th sample was silently never processed. + data = _make_data(5) + result = prepare_distributed_inf_dataset( + data=data, + n_samples=5, + per_replica_inf_batch_size=2, + num_replicas=1, + strategy=strategy, + ) + assert result["n_samples"] == 5 + assert result["n_padded"] == 6 + assert result["inf_steps"] == 3 + + out = _collect_batches(result) + assert out.shape[0] == 6 + np.testing.assert_array_equal(out[:5], data) + # Padding duplicates rows cycled from the front + np.testing.assert_array_equal(out[5], data[0]) + + def test_cadence_smaller_than_one_batch_processes_everything(self, strategy): + # Regression: with per-cadence batches, a cadence with fewer stamps than one global + # batch used to process *nothing* (inf_steps == 0). + data = _make_data(3) + result = prepare_distributed_inf_dataset( + data=data, + n_samples=3, + per_replica_inf_batch_size=8, + num_replicas=1, + strategy=strategy, + ) + assert result["inf_steps"] == 1 + assert result["n_padded"] == 8 + + out = _collect_batches(result) + np.testing.assert_array_equal(out[:3], data) + # 5 pad rows cycle deterministically over the 3 real samples + np.testing.assert_array_equal(out[3:], data[np.arange(5) % 3]) + + def test_exact_multiple_needs_no_padding(self, strategy): + data = _make_data(4) + result = prepare_distributed_inf_dataset( + data=data, + n_samples=4, + per_replica_inf_batch_size=2, + num_replicas=1, + strategy=strategy, + ) + assert result["n_padded"] == 4 + assert result["inf_steps"] == 2 + np.testing.assert_array_equal(_collect_batches(result), data) + + def test_zero_samples_raises(self, strategy): + with pytest.raises(ValueError, match="Not enough samples"): + prepare_distributed_inf_dataset( + data=np.zeros((0, 6, 4, 8), dtype=np.float32), + n_samples=0, + per_replica_inf_batch_size=2, + num_replicas=1, + strategy=strategy, + ) + + def test_order_is_preserved(self, strategy): + data = _make_data(7) + result = prepare_distributed_inf_dataset( + data=data, + n_samples=7, + per_replica_inf_batch_size=3, + num_replicas=1, + strategy=strategy, + ) + out = _collect_batches(result) + np.testing.assert_array_equal(out[:7, 0, 0, 0], np.arange(7, dtype=np.float32)) + + +class TestInfDataHolder: + def test_clear_is_idempotent(self): + holder = InfDataHolder(np.ones(3)) + holder.clear() + assert holder.data is None + holder.clear() # second clear is a no-op + assert holder.data is None diff --git a/tests/unit/test_main.py b/tests/unit/test_main.py new file mode 100644 index 0000000..068fd72 --- /dev/null +++ b/tests/unit/test_main.py @@ -0,0 +1,34 @@ +"""Unit tests for aetherscan.main: non-retryable streaming-inference failures.""" + +from __future__ import annotations + +import pytest + +from aetherscan.main import NonRetryableInferenceError, _run_streaming_csv_inference +from aetherscan.preprocessing import DataPreprocessor + + +@pytest.fixture +def initialized_runtime(): + """DataPreprocessor needs live db + manager singletons; conftest tears them down.""" + from aetherscan.db import init_db # noqa: PLC0415 + from aetherscan.manager import init_manager # noqa: PLC0415 + + init_manager() + init_db() + + +class TestStreamingInferenceNonRetryable: + def test_empty_catalog_raises_non_retryable(self, initialized_runtime): + # No inference_files configured -> plan_cadences yields no units. This is a + # permanent (config) failure: the retry loop in inference_command re-raises + # NonRetryableInferenceError immediately instead of burning retry attempts. + # The raise happens before any model loading, so no strategy is needed. + preprocessor = DataPreprocessor() + with pytest.raises(NonRetryableInferenceError, match="No cadence work units"): + _run_streaming_csv_inference(preprocessor, strategy=None) + + def test_non_retryable_error_is_an_exception_subclass(self): + # Sanity: it must be catchable as a plain Exception (cleanup paths) while being + # distinguishable from transient failures by the retry loop. + assert issubclass(NonRetryableInferenceError, RuntimeError) diff --git a/tests/unit/test_preprocessing.py b/tests/unit/test_preprocessing.py index d6acd2e..0451318 100644 --- a/tests/unit/test_preprocessing.py +++ b/tests/unit/test_preprocessing.py @@ -1,20 +1,51 @@ """Unit tests for aetherscan.preprocessing: hit deduplication, CSV cadence grouping, filename -sanitization, JSON coercion, DC-spike removal, and spline bandpass fitting.""" +sanitization, JSON coercion, DC-spike removal, spline bandpass fitting, the vectorized +normality test (equivalence-gated against scipy.stats.normaltest), the fused energy-detection +worker, downsample-at-extraction, provenance derivation, and the legacy-vs-downsampled +load_inference_data paths.""" from __future__ import annotations +import functools +import json +import math +import os + import numpy as np import pytest +from scipy import stats +from skimage.transform import downscale_local_mean from aetherscan.config import get_config +from aetherscan.data_generation import log_norm from aetherscan.preprocessing import ( DataPreprocessor, + PendingCadence, + _energy_detect_channel_worker, + _extract_stamps_worker, _fit_channel_bandpass, - _read_coarse_channel_worker, _remove_dc_spike, + _sliding_normality_k2, + _spline_flatten_bandpass, + derive_cadence_provenance, group_observations_from_csv, ) +_TESTS_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) +FIXTURES_DIR = os.path.join(_TESTS_ROOT, "fixtures") + + +def _scipy_window_loop( + channel: np.ndarray, window_size: int, step_size: int +) -> tuple[list[int], list[float]]: + """The historical per-window normaltest loop (indices + statistics), used as the oracle.""" + indices, statistics = [], [] + for i in range(0, channel.shape[1] - window_size, step_size): + s, _ = stats.normaltest(channel[:, i : i + window_size].flatten()) + indices.append(i) + statistics.append(float(s)) + return indices, statistics + @pytest.fixture def group_cols(): @@ -199,6 +230,302 @@ def test_visible_spike_is_flattened(self): np.testing.assert_allclose(block, np.ones((4, coarse_width))) +class TestDeriveCadenceProvenance: + GROUP_COLS = ["Target", "Session", "Band", "Cadence ID", "Frequency"] + + def _metadata(self, **overrides): + metadata = { + "h5_paths": [f"/data/obs_{i}.h5" for i in range(6)], + "header": {"tstart": 58306.36766203704, "fch1": 2251.46, "foff": -2.79e-06}, + "stamp_frequencies_mhz": [2250.0, 2249.5, 2249.0], + } + metadata.update(overrides) + return metadata + + def test_full_mapping(self): + key = ("DDO210", "AGBT18A_999_103", "L", "24777", "2251") + prov = derive_cadence_provenance(key, self.GROUP_COLS, self._metadata()) + assert prov["target"] == "DDO210" + assert prov["session"] == "AGBT18A_999_103" + assert prov["band"] == "L" + assert prov["cadence_id"] == 24777 + assert prov["timestamp_observed"] == pytest.approx(58306.36766203704) + assert prov["h5_path"] == "/data/obs_0.h5" + assert prov["stamp_frequencies_mhz"] == [2250.0, 2249.5, 2249.0] + + def test_column_matching_is_case_insensitive(self): + cols = ["target", "SESSION", " Band ", "cadence id", "frequency"] + key = ("T1", "S1", "X", "3", "1400") + prov = derive_cadence_provenance(key, cols, self._metadata()) + assert prov["target"] == "T1" + assert prov["session"] == "S1" + assert prov["band"] == "X" + assert prov["cadence_id"] == 3 + + def test_unparseable_cadence_id_becomes_none(self): + key = ("T", "S", "L", "not-a-number", "1400") + prov = derive_cadence_provenance(key, self.GROUP_COLS, self._metadata()) + assert prov["cadence_id"] is None + + def test_missing_metadata_degrades_to_sparse(self): + key = ("T", "S", "L", "1", "1400") + prov = derive_cadence_provenance(key, self.GROUP_COLS, {}) + assert prov["target"] == "T" + assert prov["timestamp_observed"] is None + assert prov["h5_path"] is None + assert prov["stamp_frequencies_mhz"] is None + + def test_unknown_columns_yield_none_fields(self): + prov = derive_cadence_provenance(("a", "b"), ["ColX", "ColY"], self._metadata()) + assert prov["target"] is None + assert prov["session"] is None + assert prov["band"] is None + assert prov["cadence_id"] is None + + +class TestExtractStampsWorkerDownsample: + def _run_worker(self, tmp_path, make_h5_observation, downsample_factor): + import h5py # noqa: PLC0415 + + time_bins, stamp_width = 16, 64 + stored_width = stamp_width // downsample_factor + stamp_starts = [0, 512, 1000] + h5_path = make_h5_observation("obs.h5", n_chans=2048) + npy_path = str(tmp_path / f"stamps_x{downsample_factor}.npy") + + out = np.lib.format.open_memmap( + npy_path, mode="w+", dtype=np.float32, shape=(3, 6, time_bins, stored_width) + ) + out.flush() + del out + + _extract_stamps_worker( + (npy_path, 1, str(h5_path), stamp_starts, 0, time_bins, stamp_width, downsample_factor) + ) + + written = np.load(npy_path) + with h5py.File(h5_path, "r") as hf: + raw = [hf["data"][:time_bins, 0, s : s + stamp_width] for s in stamp_starts] + return written, raw + + def test_downsampled_stamps_match_downscale_local_mean(self, tmp_path, make_h5_observation): + factor = 8 + written, raw = self._run_worker(tmp_path, make_h5_observation, factor) + for i, stamp in enumerate(raw): + expected = downscale_local_mean(stamp, (1, factor)).astype(np.float32) + np.testing.assert_array_equal(written[i, 1], expected) + # Untouched obs slots stay zero + assert np.all(written[:, 0] == 0) + + def test_factor_one_stores_raw_stamps(self, tmp_path, make_h5_observation): + written, raw = self._run_worker(tmp_path, make_h5_observation, 1) + for i, stamp in enumerate(raw): + np.testing.assert_array_equal(written[i, 1], stamp) + + +@pytest.fixture +def initialized_runtime(): + """DataPreprocessor needs live db + manager singletons; conftest tears them down.""" + from aetherscan.db import init_db # noqa: PLC0415 + from aetherscan.manager import init_manager # noqa: PLC0415 + + init_manager() + init_db() + + +class TestProcessCadenceEndToEnd: + """Sequential (no pool) end-to-end run of _process_cadence on synthetic .h5 observations + with an injected non-Gaussian feature: exercises fused detection, dedup, overlap stamps, + downsample-at-extraction, metadata, and the resume path.""" + + def _make_cadence(self, tmp_path, n_chans=2048, inject_at=768): + import h5py # noqa: PLC0415 + + rng = np.random.default_rng(23) + h5_paths = [] + for obs in range(6): + # Gaussian background (k2 stays far below threshold) with a strong, narrow + # non-Gaussian spur injected into the ON files only + data = rng.normal(1000.0, 10.0, size=(16, 1, n_chans)).astype(np.float32) + if obs in (0, 2, 4): + data[:, 0, inject_at : inject_at + 4] *= 50.0 + path = tmp_path / f"cad_obs_{obs}.h5" + with h5py.File(path, "w") as hf: + dset = hf.create_dataset("data", data=data) + dset.attrs["fch1"] = 2251.46 + dset.attrs["foff"] = -2.7939677238464355e-06 + dset.attrs["nchans"] = n_chans + dset.attrs["tstart"] = 58306.3676 + h5_paths.append(str(path)) + return h5_paths + + def _configure(self): + config = get_config() + config.manager.n_processes = 1 # sequential: no pools/shm in unit tests + config.data.time_bins = 16 + config.data.width_bin = 256 + config.data.downsample_factor = 8 + config.inference.coarse_channel_width = 512 + config.inference.spline_order = 4 + config.inference.detection_window_size = 64 + config.inference.detection_step_size = 32 + config.inference.stat_threshold = 500.0 + config.inference.stamp_width = 256 + config.inference.overlap_search = True + config.inference.overlap_fraction = 0.5 + return config + + def test_process_and_resume(self, tmp_path, initialized_runtime): + from aetherscan.preprocessing import CadenceGroup # noqa: PLC0415 + + config = self._configure() + h5_paths = self._make_cadence(tmp_path) + group = CadenceGroup( + key=("T1", "S1", "L", "7", "2251"), + h5_paths=h5_paths, + csv_path="unused.csv", + expected_obs=6, + is_valid=True, + ) + npy_path = str(tmp_path / "out" / "cadence.npy") + os.makedirs(os.path.dirname(npy_path), exist_ok=True) + + # A stale partial output from an interrupted previous run must be detected, + # warned about, and removed before extraction proceeds + stale_tmp = os.path.splitext(npy_path)[0] + ".tmp.npy" + with open(stale_tmp, "wb") as f: + f.write(b"junk from a SIGKILLed run") + + preprocessor = DataPreprocessor() + result = preprocessor.process_pending_cadence(PendingCadence(group, npy_path)) + + assert not os.path.exists(stale_tmp) + assert result is not None + assert result.npy_path == npy_path + stamps = np.load(npy_path) + stored_width = config.inference.stamp_width // config.data.downsample_factor + assert stamps.ndim == 4 + assert stamps.shape[1:] == (6, 16, stored_width) + assert stamps.shape[0] == result.n_hits + assert np.all(np.isfinite(stamps)) + # Every observation slot was filled (positive noise -> strictly positive means) + assert np.all(stamps > 0) + + with open(result.metadata_path) as f: + metadata = json.load(f) + assert metadata["stored_width"] == stored_width + assert metadata["downsample_factor_applied"] == config.data.downsample_factor + assert metadata["stamp_width"] == config.inference.stamp_width + assert len(metadata["stamp_frequencies_mhz"]) == result.n_hits + assert len(metadata["stamp_starts"]) == result.n_hits + + # Resume path: a second call must skip reprocessing and report the same hit count + resumed = preprocessor.process_pending_cadence(PendingCadence(group, npy_path)) + assert resumed is not None + assert resumed.n_hits == result.n_hits + assert resumed.npy_path == npy_path + + def test_full_width_stamps_when_disabled(self, tmp_path, initialized_runtime): + from aetherscan.preprocessing import CadenceGroup # noqa: PLC0415 + + config = self._configure() + config.inference.store_downsampled_stamps = False + h5_paths = self._make_cadence(tmp_path) + group = CadenceGroup( + key=("T2", "S1", "L", "8", "2251"), + h5_paths=h5_paths, + csv_path="unused.csv", + expected_obs=6, + is_valid=True, + ) + npy_path = str(tmp_path / "out" / "cadence_full.npy") + os.makedirs(os.path.dirname(npy_path), exist_ok=True) + + result = DataPreprocessor().process_pending_cadence(PendingCadence(group, npy_path)) + + assert result is not None + stamps = np.load(npy_path) + assert stamps.shape[1:] == (6, 16, config.inference.stamp_width) + with open(result.metadata_path) as f: + metadata = json.load(f) + assert metadata["stored_width"] == config.inference.stamp_width + assert metadata["downsample_factor_applied"] == 1 + + +class TestLoadInferenceDataPaths: + """load_inference_data must branch on the stored width: already-downsampled cadence .npy + files get log-norm only; legacy full-width files keep downsample + log-norm.""" + + def _configure(self): + config = get_config() + config.manager.n_processes = 1 + config.data.width_bin = 512 + config.data.downsample_factor = 8 + return config + + def test_downsampled_path_lognorm_only(self, tmp_path, initialized_runtime): + config = self._configure() + final_width = config.data.width_bin // config.data.downsample_factor + rng = np.random.default_rng(31) + arr = rng.chisquare(df=4, size=(3, 6, 16, final_width)).astype(np.float32) + path = tmp_path / "downsampled.npy" + np.save(path, arr) + + loaded = DataPreprocessor().load_inference_data(override_filepaths=[str(path)]) + + assert loaded.shape == (3, 6, 16, final_width) + for i in range(3): + np.testing.assert_allclose(loaded[i], log_norm(arr[i]), rtol=1e-6) + assert loaded.min() >= 0.0 and loaded.max() <= 1.0 + + def test_legacy_full_width_path_downsamples_then_lognorms(self, tmp_path, initialized_runtime): + config = self._configure() + width_bin = config.data.width_bin + factor = config.data.downsample_factor + final_width = width_bin // factor + rng = np.random.default_rng(37) + arr = rng.chisquare(df=4, size=(2, 6, 16, width_bin)).astype(np.float32) + path = tmp_path / "legacy.npy" + np.save(path, arr) + + loaded = DataPreprocessor().load_inference_data(override_filepaths=[str(path)]) + + assert loaded.shape == (2, 6, 16, final_width) + for i in range(2): + downsampled = np.stack( + [ + downscale_local_mean(arr[i, obs], (1, factor)).astype(np.float32) + for obs in range(6) + ] + ) + np.testing.assert_allclose(loaded[i], log_norm(downsampled), rtol=1e-6) + + def test_unrecognized_width_is_skipped(self, tmp_path, initialized_runtime): + self._configure() + arr = np.abs(np.random.default_rng(41).normal(1, 0.1, (2, 6, 16, 100))).astype(np.float32) + path = tmp_path / "weird.npy" + np.save(path, arr) + + with pytest.raises(ValueError, match="No data loaded successfully"): + DataPreprocessor().load_inference_data(override_filepaths=[str(path)]) + + def test_mixed_legacy_and_downsampled_files(self, tmp_path, initialized_runtime): + config = self._configure() + final_width = config.data.width_bin // config.data.downsample_factor + rng = np.random.default_rng(43) + legacy = rng.chisquare(df=4, size=(1, 6, 16, config.data.width_bin)).astype(np.float32) + modern = rng.chisquare(df=4, size=(2, 6, 16, final_width)).astype(np.float32) + legacy_path = tmp_path / "legacy.npy" + modern_path = tmp_path / "modern.npy" + np.save(legacy_path, legacy) + np.save(modern_path, modern) + + loaded = DataPreprocessor().load_inference_data( + override_filepaths=[str(legacy_path), str(modern_path)] + ) + assert loaded.shape == (3, 6, 16, final_width) + + class TestFitChannelBandpass: def test_smooth_bandpass_recovered(self): channel_width, spl_order = 1024, 16 @@ -221,15 +548,193 @@ def test_subtracting_fit_flattens_channel(self): assert residual.std() < 1.0 -class TestReadCoarseChannelWorker: - def test_reads_one_coarse_channel(self, make_h5_observation): - n_chans, coarse_width = 2048, 512 - h5_path = make_h5_observation("obs.h5", n_chans=n_chans) - channel = _read_coarse_channel_worker((str(h5_path), 2, coarse_width)) - assert channel.shape == (16, coarse_width) +class TestSlidingNormalityK2: + """The §8.1 correctness gates: the vectorized k2 must match scipy.stats.normaltest + per window to rtol=1e-9 across distributions, scales, and window/step geometries.""" + + WINDOW_STEP_COMBOS = [ + (256, 128), # default config: step divides window (fast path) + (256, 256), # window == step + (96, 64), # step does not divide window -> gcd (32) general path + (100, 30), # gcd 10 general path, ragged tail blocks + ] + + @pytest.mark.parametrize( + "name,dist", + [ + ("normal", lambda rng, shape: rng.normal(0.0, 1.0, shape)), + ("normal_offset", lambda rng, shape: rng.normal(5.0, 0.01, shape)), + ("normal_tiny_scale", lambda rng, shape: rng.normal(0.0, 1e-3, shape)), + ("laplace_large_scale", lambda rng, shape: rng.laplace(0.0, 1e3, shape)), + ("uniform", lambda rng, shape: rng.uniform(-1.0, 1.0, shape)), + ("chisquare_skewed", lambda rng, shape: rng.chisquare(4, shape)), + ], + ) + @pytest.mark.parametrize("window_size,step_size", WINDOW_STEP_COMBOS) + def test_matches_scipy_normaltest_per_window(self, name, dist, window_size, step_size): + rng = np.random.default_rng(11) + channel = dist(rng, (16, 4096)) + + k2 = _sliding_normality_k2(channel, window_size, step_size) + indices, expected = _scipy_window_loop(channel, window_size, step_size) + + assert len(k2) == len(indices) + np.testing.assert_allclose(k2, expected, rtol=1e-9) + + @pytest.mark.parametrize( + "width,window_size,step_size", + [ + (1234, 100, 30), # gcd 10; 1234 % 10 == 4 -> ragged short trailing block + (1000, 96, 36), # gcd 12; 1000 % 12 == 4 -> ragged short trailing block + ], + ) + def test_non_power_of_2_width_short_block_matches_scipy(self, width, window_size, step_size): + # Guards the reduceat trailing-block truncation: when block does not divide the coarse + # channel width, the final block spans a short ragged tail. No in-range window may + # consume it, so k2 must still match scipy exactly. Assert the geometry actually + # triggers the short block so a future refactor can't quietly make this test vacuous. + block = step_size if window_size % step_size == 0 else math.gcd(window_size, step_size) + assert width % block != 0 # this geometry must exercise the short-block path + rng = np.random.default_rng(17) + channel = rng.normal(0.0, 1.0, (16, width)) + + k2 = _sliding_normality_k2(channel, window_size, step_size) + indices, expected = _scipy_window_loop(channel, window_size, step_size) + + assert len(k2) == len(indices) + np.testing.assert_allclose(k2, expected, rtol=1e-9) + + def test_float32_input_matches_scipy_on_float64_view(self): + # The h5 data is float32, but the pipeline always tests the bandpass-subtracted + # residuals, which are float64 (float32 channel - float64 spline fit). The oracle + # therefore consumes the float64 view of the same values — scipy fed raw float32 + # would itself compute in float32 and diverge from its own float64 result. + rng = np.random.default_rng(3) + channel = rng.normal(0, 1, (16, 2048)).astype(np.float32) + k2 = _sliding_normality_k2(channel, 256, 128) + _, expected = _scipy_window_loop(channel.astype(np.float64), 256, 128) + np.testing.assert_allclose(k2, expected, rtol=1e-9) + + def test_window_indices_match_historical_loop(self): + # n_windows = len(range(0, width - window, step)): the window starting exactly at + # width - window is excluded, like the historical loop. + rng = np.random.default_rng(5) + channel = rng.normal(0, 1, (16, 1024)) + k2 = _sliding_normality_k2(channel, 256, 128) + assert len(k2) == len(range(0, 1024 - 256, 128)) + + def test_zero_variance_window_yields_nan_not_hit(self): + channel = np.ones((16, 1024)) + k2 = _sliding_normality_k2(channel, 256, 128) + assert np.all(np.isnan(k2)) + # NaN never exceeds any threshold -> no spurious hits on degenerate data + assert not np.any(k2 > 0.0) + + def test_too_few_samples_raises(self): + with pytest.raises(ValueError, match=">= 8 samples"): + _sliding_normality_k2(np.zeros((1, 16)), 4, 2) + + def test_empty_when_window_exceeds_width(self): + assert _sliding_normality_k2(np.zeros((16, 128)), 256, 128).size == 0 + + +@pytest.mark.slow +class TestRealSliceEquivalence: + """Hit-set identity on recorded real data (a 16384-bin slice of one coarse channel of a + Breakthrough Listen GBT .h5, centered on the channel's DC spike; the region is + RFI-contaminated so it produces real hits).""" + + def test_identical_hits_on_real_slice(self): + fixture = np.load(os.path.join(FIXTURES_DIR, "ed_real_slice.npz")) + data = fixture["data"] + assert data.shape == (16, 16384) + + width = data.shape[1] + window_size, step_size, stat_threshold = 256, 128, 2048.0 + + # Reproduce the per-channel pipeline: DC-spike removal (the slice is centered on the + # spike, so treating it as one coarse channel puts the spike at width // 2, exactly + # where _remove_dc_spike expects it) -> spline bandpass flatten -> threshold. + channel = data.copy() + _remove_dc_spike(channel, width, 1) + residuals = _spline_flatten_bandpass(channel, spl_order=16) + + k2 = _sliding_normality_k2(residuals, window_size, step_size) + vec_hits = {int(j) * step_size: float(k2[j]) for j in np.nonzero(k2 > stat_threshold)[0]} + + indices, statistics = _scipy_window_loop(residuals, window_size, step_size) + scipy_hits = {i: s for i, s in zip(indices, statistics, strict=True) if s > stat_threshold} + + # The fixture region was chosen to contain real hits — an empty set would make this + # test vacuous. + assert len(scipy_hits) > 0 + assert set(vec_hits) == set(scipy_hits) + np.testing.assert_allclose( + [vec_hits[i] for i in sorted(vec_hits)], + [scipy_hits[i] for i in sorted(scipy_hits)], + rtol=1e-9, + ) + +class TestEnergyDetectChannelWorker: + def test_fused_worker_matches_composed_pipeline(self, make_h5_observation): + """The fused read -> DC spike -> bandpass -> threshold worker must reproduce the + explicitly composed per-stage pipeline on the same coarse channel.""" import h5py # noqa: PLC0415 + n_chans, coarse_width = 2048, 512 + channel_index = 2 + window_size, step_size = 64, 32 + stat_threshold = 20.0 # low threshold so chi-square noise produces hits + spl_order = 4 + + h5_path = make_h5_observation("obs.h5", n_chans=n_chans) + bandpass_flatten = functools.partial(_spline_flatten_bandpass, spl_order=spl_order) + + hits = _energy_detect_channel_worker( + ( + str(h5_path), + channel_index, + coarse_width, + 16, + bandpass_flatten, + window_size, + step_size, + stat_threshold, + ) + ) + + # Compose the same chain by hand with h5py.File(h5_path, "r") as hf: - expected = hf["data"][:, 0, 2 * coarse_width : 3 * coarse_width] - np.testing.assert_array_equal(channel, expected) + channel = hf["data"][ + :16, 0, channel_index * coarse_width : (channel_index + 1) * coarse_width + ] + _remove_dc_spike(channel, coarse_width, 1) + residuals = bandpass_flatten(channel) + indices, statistics = _scipy_window_loop(residuals, window_size, step_size) + expected = { + channel_index * coarse_width + i: s + for i, s in zip(indices, statistics, strict=True) + if s > stat_threshold + } + + assert len(hits) > 0 # chi-square noise vs a low threshold must produce hits + assert {idx for idx, _, _ in hits} == set(expected) + for idx, stat_val, pval in hits: + np.testing.assert_allclose(stat_val, expected[idx], rtol=1e-9) + np.testing.assert_allclose(pval, stats.chi2.sf(stat_val, 2), rtol=1e-9) + + def test_absolute_indices_offset_by_channel_start(self, make_h5_observation): + h5_path = make_h5_observation("obs.h5", n_chans=2048) + coarse_width = 512 + bandpass_flatten = functools.partial(_spline_flatten_bandpass, spl_order=4) + for channel_index in (0, 3): + hits = _energy_detect_channel_worker( + (str(h5_path), channel_index, coarse_width, 16, bandpass_flatten, 64, 32, 0.0) + ) + starts = [idx for idx, _, _ in hits] + assert len(starts) > 0 # threshold 0.0 must produce hits; else all(...) is vacuous + assert all( + channel_index * coarse_width <= s < (channel_index + 1) * coarse_width + for s in starts + )