aetherscan 1.0.0__py3-none-any.whl

aetherscan/__init__.py

@@ -0,0 +1,45 @@
+"""
+Aetherscan: Breakthrough Listen's first end-to-end production-grade DL pipeline for SETI @ scale.
+
+Provides tools for technosignature detection using signal processing and deep learning techniques
+applied to radio astronomy data. Import submodules directly (aetherscan.config, aetherscan.models,
+aetherscan.data_generation, aetherscan.preprocessing, etc.) — the top-level package intentionally
+exposes only the version string.
+"""
+
+from __future__ import annotations
+
+from importlib.metadata import PackageNotFoundError, version
+
+# Version is single-sourced from pyproject.toml's [project].version via the installed
+# distribution's metadata. Source-tree runs (PYTHONPATH=src, the NGC container) have no
+# installed distribution — they fall back to a dev sentinel, which also keeps the
+# version-coupled HF weight resolution (hf_hub.version_default_revision) inactive there.
+try:
+    __version__ = version("aetherscan")
+except PackageNotFoundError:
+    __version__ = "0.0.0.dev0"
+__author__ = "Zach Yek"
+
+# TODO: determine which components are necessary to expose for public API
+# # Core configuration
+# from aetherscan.config import Config
+#
+# # Models
+# from aetherscan.models import RandomForestModel, Sampling, create_beta_vae_model
+#
+# # Data processing
+# from aetherscan.data_generation import DataGenerator
+# from aetherscan.preprocessing import DataPreprocessor
+#
+# __all__ = [
+#   "Config",
+#   "RandomForestModel",
+#   "Sampling",
+#   "create_beta_vae_model",
+#   "DataGenerator",
+#   "DataPreprocessor",
+# ]
+
+# Minimal public API - import submodules explicitly as needed
+__all__ = []

aetherscan/benchmark.py

@@ -0,0 +1,179 @@
+"""
+Always-on stage timing for the Aetherscan pipeline.
+
+stage_timer() wraps a pipeline stage (context manager or decorator) and records one
+(stage, start_time, end_time, duration_s, tag, metadata) row into the pipeline_stages DB
+table through the database's writer queue — two time.time() calls plus one queue put, so
+the overhead is negligible and the timers stay on in production runs.
+
+Stage names are hierarchical dot-names ("train.round_02.data_generation"). Nesting is
+automatic within a thread: a stage_timer entered while another is active on the same
+thread records its name relative to the active one, so instrumented library code (e.g.
+the encode/rf sub-stages inside InferencePipeline.run_inference) inherits whatever
+umbrella span the caller opened without name plumbing. Timers on different threads don't
+interact (the active-stage stack is thread-local), which matters because training,
+prefetch preprocessing, and the producer drainer all time stages concurrently.
+
+record_stage() writes a span measured elsewhere (explicit start/end timestamps) — used by
+the round-data producer, whose generation happens in another process: the producer can't
+touch the DB writer queue (a thread queue.Queue), so it reports (start, end) over its
+result-message channel and the main-process drainer records it here.
+
+Rows are read back by utils/benchmark_report.py (stage tree + timeline plot) and by
+monitor._save_plot()'s stage-band overlay.
+"""
+
+from __future__ import annotations
+
+import json
+import logging
+import threading
+import time
+from contextlib import ContextDecorator
+
+logger = logging.getLogger(__name__)
+
+# Thread-local stack of active stage names (full dotted names), driving relative naming
+_ACTIVE_STAGES = threading.local()
+
+
+def _stage_stack() -> list[str]:
+    """The current thread's stack of active (full) stage names."""
+    stack = getattr(_ACTIVE_STAGES, "stack", None)
+    if stack is None:
+        stack = []
+        _ACTIVE_STAGES.stack = stack
+    return stack
+
+
+def current_stage() -> str | None:
+    """Full dotted name of the innermost stage_timer active on this thread, or None."""
+    stack = _stage_stack()
+    return stack[-1] if stack else None
+
+
+def round_stage_name(round_number: int) -> str:
+    """Canonical umbrella stage name for a 1-based training round ("train.round_02").
+
+    Shared by the trainer (which opens the umbrella span) and the round-data producer
+    (which, running in another process, records the data_generation child by absolute name
+    and so can't rely on thread-local nesting). Routing both sites through this helper keeps
+    the two name constructions from drifting — a mismatch would orphan the producer's
+    data_generation span outside the round subtree in the report tree.
+    """
+    return f"train.round_{round_number:02d}"
+
+
+def record_stage(
+    stage: str,
+    start_time: float,
+    end_time: float,
+    tag: str | None = None,
+    metadata: dict | None = None,
+) -> None:
+    """
+    Record one pipeline stage span with explicit timestamps (no nesting resolution — the
+    name is used as-is). tag defaults to the run's save tag. Never raises: benchmarking
+    must not be able to fail the pipeline, so a missing DB (unit tests, dev scripts) or a
+    serialization hiccup downgrades to a debug/warning log.
+    """
+    try:
+        # Late imports keep this module import-light and avoid any import-cycle risk
+        # (db imports config + manager; nothing imports benchmark back).
+        from aetherscan.config import get_config  # noqa: PLC0415
+        from aetherscan.db.db import Database  # noqa: PLC0415
+
+        # Read the singleton directly (not via get_db()): a missing DB is an expected,
+        # already-handled case here, and get_db() logs a WARNING on every miss — noise
+        # that could reach Slack. This is a read, not an instantiation.
+        db = Database._instance
+        if db is None:
+            logger.debug(f"No database instance — dropping stage timing for {stage!r}")
+            return
+
+        if tag is None:
+            config = get_config()
+            # NOTE: tag stays None if a timer fires before checkpoint.save_tag is wired
+            # (early init, pre-CLI-parse). Such rows won't match a tag="..." query filter.
+            # In practice timers don't fire that early, so this is a documented edge, not a bug.
+            tag = config.checkpoint.save_tag if config is not None else None
+
+        db.write_pipeline_stage(
+            stage=stage,
+            start_time=start_time,
+            end_time=end_time,
+            tag=tag,
+            metadata=json.dumps(metadata) if metadata else None,
+        )
+    except Exception as e:
+        logger.warning(f"Failed to record stage timing for {stage!r}: {e}")
+
+
+class _StageTimer(ContextDecorator):
+    """Context manager / decorator that times a block and records it via record_stage().
+
+    One instance holds per-entry state (full_name/start_time), so a decorated function is
+    fine to call repeatedly from one thread but not concurrently from several — for
+    concurrent callers, create the timer inside the function (`with stage_timer(...)`)."""
+
+    def __init__(self, stage: str, tag: str | None = None, metadata: dict | None = None):
+        self.stage = stage
+        self.tag = tag
+        self.metadata = metadata
+        self.full_name: str | None = None
+        self.start_time: float | None = None
+
+    def __enter__(self):
+        parent = current_stage()
+        self.full_name = f"{parent}.{self.stage}" if parent else self.stage
+        _stage_stack().append(self.full_name)
+        self.start_time = time.time()
+        return self
+
+    def __exit__(self, exc_type, exc_value, exc_tb):
+        end_time = time.time()
+        stack = _stage_stack()
+        # Pop our own frame (guarded: a corrupted stack must not mask the block's result)
+        if stack and stack[-1] == self.full_name:
+            stack.pop()
+        else:
+            logger.warning(
+                f"stage_timer stack mismatch on exit of {self.full_name!r} "
+                f"(top: {(stack[-1] if stack else None)!r})"
+            )
+
+        metadata = self.metadata
+        if exc_type is not None:
+            # Record the failure but never suppress it (return None -> exception propagates)
+            metadata = dict(metadata or {})
+            metadata["status"] = "failed"
+            metadata["error"] = f"{exc_type.__name__}: {exc_value}"
+
+        # Implicit None return -> exceptions propagate (never suppressed)
+        record_stage(
+            self.full_name,
+            self.start_time,
+            end_time,
+            tag=self.tag,
+            metadata=metadata,
+        )
+
+
+def stage_timer(stage: str, tag: str | None = None, metadata: dict | None = None) -> _StageTimer:
+    """
+    Time a pipeline stage and queue a pipeline_stages DB row on exit.
+
+        with stage_timer("train.round_02.data_generation"):
+            ...
+
+        @stage_timer("inference.viz")
+        def render(): ...
+
+    `stage` is a dotted hierarchical name. If another stage_timer is active on the same
+    thread, `stage` is recorded relative to it (entering stage_timer("encode") while
+    "inference.infer_cadence_001" is active records "inference.infer_cadence_001.encode");
+    with no active parent it is recorded as-is. On exception the span is still recorded
+    (metadata carries status=failed plus the error) and the exception propagates.
+    `metadata`, when given, is a small JSON-serializable dict stored on the row.
+    """
+    return _StageTimer(stage, tag=tag, metadata=metadata)

aetherscan/candidate_figures.py

@@ -0,0 +1,328 @@
+"""
+Process-parallel per-candidate figure rendering for the inference viz suite (#298 I9).
+
+The ≤ max_candidate_plots per-candidate figures are independent (one DB row dict + one
+memmap stamp read + one PNG each) but were rendered strictly serially on the main thread at
+finalize time. This module farms them across a small process pool, mirroring
+shap_parallel.py / latent_gif.py's isolation pattern: **forkserver context with an empty
+preload**, so workers never re-import the parent's ``__main__`` (``aetherscan.main`` → TF →
+the whole stack) and never touch the CUDA-initialized parent state — ``inference_viz``
+itself imports TF transitively via ``aetherscan.models``, which is exactly why the
+figure-building code lives HERE, deliberately TF-free, and ``inference_viz`` imports it
+(never the reverse).
+
+Figures are built on the object-oriented matplotlib API (``Figure`` + implicit Agg canvas —
+pyplot is never imported, so no GUI backend can be selected in a worker). Per-figure
+containment mirrors ``_viz_safe``: a failed render returns None for that candidate and the
+suite continues. Slack uploads stay with the PARENT (inference_viz's async uploader), fed in
+index order from the returned paths.
+"""
+
+from __future__ import annotations
+
+import functools
+import json
+import logging
+import multiprocessing as mp
+import os
+import uuid
+from concurrent.futures import ProcessPoolExecutor
+
+import numpy as np
+from matplotlib.figure import Figure
+
+from aetherscan.data_generation import log_norm
+
+logger = logging.getLogger(__name__)
+
+# ABACAD cadence row labels (ON at positions 0/2/4) — single source for every figure that
+# draws observation strips (the suite imports it from here).
+OBS_ROW_LABELS = ("ON", "OFF", "ON", "OFF", "ON", "OFF")
+
+# Pool sizing: candidate figures are ~1 s of matplotlib each with a hard cap of
+# max_candidate_plots (50), so a handful of workers saturates the win; below the floor the
+# pool spin-up would cost more than it saves and rendering stays in-process.
+_MAX_RENDER_WORKERS = 8
+_MIN_ROWS_FOR_POOL = 4
+
+
+def candidate_sidecar_path(npy_path: str) -> str:
+    """Sibling .candidates.npz path for a cadence's pruned candidate snippets (#302)."""
+    return os.path.splitext(npy_path)[0] + ".candidates.npz"
+
+
+def write_candidate_snippet_sidecar(npy_path: str, snippet_indices) -> str:
+    """Snapshot the candidate snippets (~196 KB each) of a cadence into an atomic
+    .candidates.npz sidecar next to its metadata (#302): stamp-cache pruning deletes the
+    multi-GB stamp .npy after scoring, and the candidate figures re-read their snippets
+    from this sidecar instead. Raw (pre-log-norm) values, exactly as stored in the .npy,
+    so load_display_cadence's fallback applies the identical display transform."""
+    stamps = np.load(npy_path, mmap_mode="r")
+    indices = np.asarray(sorted({int(i) for i in snippet_indices}), dtype=np.int64)
+    rows = np.array(stamps[indices], dtype=np.float32)
+    del stamps
+    path = candidate_sidecar_path(npy_path)
+    tmp_path = f"{path}.{os.getpid()}.{uuid.uuid4().hex}.tmp"
+    with open(tmp_path, "wb") as f:
+        np.savez(f, snippet_indices=indices, stamps=rows)
+    os.replace(tmp_path, path)
+    return path
+
+
+def _load_snippet_from_sidecar(npy_path: str, snippet_index: int) -> np.ndarray:
+    """Read one candidate snippet from the pruned cadence's sidecar; raises KeyError when
+    the snippet was not a candidate (non-candidate pixels are gone by design, #302)."""
+    with np.load(candidate_sidecar_path(npy_path)) as sidecar:
+        match = np.nonzero(sidecar["snippet_indices"] == int(snippet_index))[0]
+        if not len(match):
+            raise KeyError(
+                f"snippet {snippet_index} is not in the candidate sidecar for {npy_path} "
+                f"(its stamp .npy was pruned and only candidate snippets were kept)"
+            )
+        return np.array(sidecar["stamps"][int(match[0])], dtype=np.float32)
+
+
+def load_display_cadence(npy_path: str, snippet_index: int) -> np.ndarray:
+    """Load one snippet's (num_obs, time_bins, width) stamp from its .npy — falling back
+    to the candidate snippet sidecar when the .npy was pruned (#302) — and log-normalize
+    it for display (the same transform the model input path applies)."""
+    try:
+        stamps = np.load(npy_path, mmap_mode="r")
+    except OSError:
+        snippet = _load_snippet_from_sidecar(npy_path, snippet_index)
+    else:
+        snippet = np.array(stamps[snippet_index], dtype=np.float32)
+        del stamps
+    return log_norm(snippet)
+
+
+def cadence_metadata_path(npy_path: str) -> str:
+    """Sibling .json path for a cadence's metadata — the same naming rule as
+    preprocessing.DataPreprocessor.cadence_metadata_path, duplicated here (one line) so
+    render workers never import the preprocessing module (manager/db singletons, setigen)."""
+    return os.path.splitext(npy_path)[0] + ".json"
+
+
+def stamp_frequency_range_mhz(metadata: dict, snippet_index: int) -> tuple[float, float] | None:
+    """
+    Frequency range (MHz) spanned by one stamp's frequency axis — bin 0 through the last
+    RAW bin — from the cadence metadata sidecar: header fch1/foff plus the stamp's start
+    index and stamp_width (#298 follow-up: cadence-snippet plots label their x-axis with
+    the frequency span). Returned in bin order, so a negative foff yields a descending
+    (high → low) pair — callers print it as-is. None when any field is missing or malformed
+    (legacy sidecars): callers keep the unlabeled axis.
+    """
+    try:
+        header = metadata.get("header") or {}
+        fch1 = float(header["fch1"])
+        foff = float(header["foff"])
+        start = int((metadata.get("stamp_starts") or [])[snippet_index])
+        width = int(metadata["stamp_width"])
+    except (KeyError, IndexError, TypeError, ValueError):
+        return None
+    return fch1 + foff * start, fch1 + foff * (start + width - 1)
+
+
+def draw_cadence_strip(
+    axes_column,
+    snippet: np.ndarray,
+    label_rows: bool,
+    freq_range_mhz: tuple[float, float] | None = None,
+) -> None:
+    """Draw one snippet's 6 observation waterfalls down a column of axes. With
+    freq_range_mhz, the bottom axis is labeled with the stamp's frequency span (bin order,
+    so a descending pair reflects a negative foff)."""
+    for obs_idx, ax in enumerate(axes_column):
+        ax.imshow(
+            snippet[obs_idx],
+            aspect="auto",
+            origin="lower",
+            cmap="viridis",
+            vmin=0.0,
+            vmax=1.0,
+            interpolation="nearest",
+        )
+        ax.set_xticks([])
+        ax.set_yticks([])
+        if label_rows:
+            ax.set_ylabel(OBS_ROW_LABELS[obs_idx], fontsize=8, rotation=0, ha="right", va="center")
+    if freq_range_mhz is not None:
+        low, high = freq_range_mhz
+        axes_column[-1].set_xlabel(f"{low:.4f} → {high:.4f} MHz", fontsize=6)
+
+
+def candidate_annotation(row: dict) -> str:
+    lines = [f"confidence: {row.get('confidence', float('nan')):.4f}"]
+    if row.get("frequency_mhz") is not None:
+        lines.append(f"frequency: {row['frequency_mhz']:.6f} MHz")
+    for label in ("target", "session", "band", "cadence_id"):
+        if row.get(label) is not None:
+            lines.append(f"{label}: {row[label]}")
+    if row.get("timestamp_observed") is not None:
+        lines.append(f"tstart (MJD): {row['timestamp_observed']:.5f}")
+    if row.get("h5_path"):
+        lines.append(f"h5: {os.path.basename(str(row['h5_path']))}")
+    lines.append(f"npy: {os.path.basename(str(row.get('npy_path', '')))}")
+    lines.append(f"snippet: {row.get('snippet_index')}")
+    return "\n".join(lines)
+
+
+@functools.lru_cache(maxsize=4)
+def _load_sidecar_cached(metadata_path: str) -> dict | None:
+    """Small parent-side LRU over parsed metadata sidecars (#301): candidates cluster on
+    few cadences, and the gallery/task-build path parsed the SAME multi-MB JSON once per
+    candidate row (~15-20 s serial before the render pool even started). Sidecars are
+    immutable once published (atomic tmp -> os.replace), so a cached parse can never go
+    stale; maxsize bounds the held dicts.
+
+    CONTRACT: the returned dict is the CACHED object, aliased across every caller —
+    treat it as immutable (a mutation would silently poison all later readers of the
+    same sidecar; a copy here would defeat the memory bound)."""
+    try:
+        with open(metadata_path) as f:
+            return json.load(f)
+    except (OSError, json.JSONDecodeError):
+        return None
+
+
+def candidate_frequency_range_mhz(row: dict) -> tuple[float, float] | None:
+    """Best-effort frequency span for one candidate row: read its cadence's metadata
+    sidecar (derived from npy_path, cached across rows) and look up the stamp's range.
+    None on any failure — the figure keeps its generic axis label."""
+    try:
+        metadata = _load_sidecar_cached(cadence_metadata_path(str(row["npy_path"])))
+    except KeyError:
+        return None
+    if metadata is None:
+        return None
+    return stamp_frequency_range_mhz(metadata, int(row["snippet_index"]))
+
+
+def render_candidate_figure(
+    row: dict,
+    index: int,
+    tag: str,
+    plots_dir: str,
+    freq_range_mhz: tuple[float, float] | None = None,
+) -> str:
+    """
+    Build and save one candidate's figure (the long-standing inference.py stub): 6-panel
+    cadence waterfall of its stamp, annotated with confidence / frequency / target /
+    session / band, plus the latent vector as a bar chart. Returns the saved PNG path.
+    Pure function of its arguments — no config/db/logger singletons — so it runs
+    identically in-process and in a forkserver worker. freq_range_mhz labels the waterfall
+    x-axis with the stamp's frequency span (callers precompute it from the metadata
+    sidecar; the axis stays generic when None).
+    """
+    snippet = load_display_cadence(str(row["npy_path"]), int(row["snippet_index"]))
+    n_obs = snippet.shape[0]
+
+    fig = Figure(figsize=(11, 7))
+    grid = fig.add_gridspec(n_obs, 2, width_ratios=(2.2, 1.6), hspace=0.15, wspace=0.25, right=0.97)
+    waterfall_axes = [fig.add_subplot(grid[i, 0]) for i in range(n_obs)]
+    draw_cadence_strip(waterfall_axes, snippet, label_rows=True)
+    if freq_range_mhz is not None:
+        low, high = freq_range_mhz
+        waterfall_axes[-1].set_xlabel(f"frequency: {low:.6f} → {high:.6f} MHz")
+    else:
+        waterfall_axes[-1].set_xlabel("frequency bin")
+
+    # Right column: latent bar chart on top, provenance text below. A nested gridspec gives
+    # the two panels a dedicated vertical gap so the bar chart's x-axis label can never
+    # collide with the first metadata line; the text panel takes the taller share so all of
+    # the provenance lines stay clear of the axis label regardless of how many there are.
+    right_grid = grid[:, 1].subgridspec(2, 1, height_ratios=(1.0, 1.3), hspace=0.55)
+    latent_ax = fig.add_subplot(right_grid[0])
+    latent_json = row.get("latent_vector")
+    if latent_json:
+        latent = np.asarray(json.loads(latent_json), dtype=np.float64).ravel()
+        latent_dim = latent.size // n_obs if latent.size % n_obs == 0 else latent.size
+        colors = [f"C{(i // latent_dim) % 10}" for i in range(latent.size)]
+        latent_ax.bar(np.arange(latent.size), latent, color=colors)
+        latent_ax.set_xlabel("latent dimension (colored per observation)", fontsize=8)
+        latent_ax.set_ylabel("z", fontsize=8)
+        latent_ax.tick_params(labelsize=7)
+        latent_ax.grid(True, axis="y", alpha=0.2)
+    else:
+        latent_ax.set_axis_off()
+        latent_ax.text(0.5, 0.5, "no latent vector stored", ha="center", va="center")
+
+    text_ax = fig.add_subplot(right_grid[1])
+    text_ax.set_axis_off()
+    text_ax.text(
+        0.0, 1.0, candidate_annotation(row), va="top", ha="left", fontsize=9, family="monospace"
+    )
+
+    fig.suptitle(f"Candidate {index} ({tag}) — P(true) = {row.get('confidence', 0):.4f}")
+
+    save_path = os.path.join(plots_dir, f"candidate_{index}_{tag}.png")
+    fig.savefig(save_path, dpi=150, bbox_inches="tight")
+    fig.clear()
+    return save_path
+
+
+def _pin_worker_threads() -> None:
+    # Must run before numpy imports its native thread pools; matplotlib rendering is
+    # single-threaded, so this only tames incidental numpy/BLAS threads (mirrors
+    # shap_parallel._pin_worker_threads). setdefault so an explicitly-set env is respected.
+    for var in (
+        "OMP_NUM_THREADS",
+        "OPENBLAS_NUM_THREADS",
+        "MKL_NUM_THREADS",
+        "NUMEXPR_NUM_THREADS",
+        "VECLIB_MAXIMUM_THREADS",
+        "BLIS_NUM_THREADS",
+    ):
+        os.environ.setdefault(var, "1")
+
+
+def _render_task(args: tuple) -> tuple[int, str | None]:
+    """Worker body: render one candidate with _viz_safe-style containment (a failed figure
+    must degrade the suite, never abort the pool pass)."""
+    row, index, tag, plots_dir, freq_range_mhz = args
+    try:
+        return index, render_candidate_figure(row, index, tag, plots_dir, freq_range_mhz)
+    except Exception as e:
+        logger.error(f"Candidate figure {index} failed (suite continues without it): {e}")
+        return index, None
+
+
+def render_candidate_figures(
+    rows: list[dict], tag: str, plots_dir: str, n_workers: int | None = None
+) -> list[tuple[int, str | None]]:
+    """
+    Render one figure per (row, index) pair — in index order in the returned list — across
+    a forkserver process pool, falling back to in-process rendering for small candidate
+    counts or any pool-level failure. Rows must be plain dicts of primitives (the
+    inference_results row shape); paths are None for candidates whose render failed.
+    Frequency spans are resolved in the PARENT (one sidecar read per row, best-effort) so
+    workers stay pure.
+    """
+    tasks = [
+        (row, index, tag, plots_dir, candidate_frequency_range_mhz(row))
+        for index, row in enumerate(rows)
+    ]
+    if not tasks:
+        return []
+
+    n_workers = min(_MAX_RENDER_WORKERS, len(tasks), n_workers or (os.cpu_count() or 1))
+    if n_workers <= 1 or len(tasks) < _MIN_ROWS_FOR_POOL:
+        return [_render_task(task) for task in tasks]
+
+    try:
+        # forkserver with an EMPTY preload: the fork server is spawned clean and does NOT
+        # re-import the parent's __main__ (aetherscan.main -> TF -> CUDA) — a bare fork of
+        # the CUDA-initialized 5-GPU parent at finalize time is exactly what this avoids.
+        ctx = mp.get_context("forkserver")
+        ctx.set_forkserver_preload([])
+        with ProcessPoolExecutor(
+            max_workers=n_workers,
+            mp_context=ctx,
+            initializer=_pin_worker_threads,
+        ) as pool:
+            return list(pool.map(_render_task, tasks))
+    except Exception as e:
+        logger.error(
+            f"Candidate-figure pool failed ({e}); rendering the remaining figures in-process"
+        )
+        return [_render_task(task) for task in tasks]