diff --git a/.claude/skills/aetherscan-repo-context/SKILL.md b/.claude/skills/aetherscan-repo-context/SKILL.md index 76d6e80c..0c8fb950 100644 --- a/.claude/skills/aetherscan-repo-context/SKILL.md +++ b/.claude/skills/aetherscan-repo-context/SKILL.md @@ -114,7 +114,10 @@ utils/ # fetch_run_outputs.sh, find_optimal_configs.py, # run_container.sh, start_tmux_session.sh, # verify_train_test_files.py docs/ # GPU_RUNTIME_GUIDE.md, CONFIG_AND_CLI.md, README.md, assets/ -tests/ # Placeholder — no test suite yet +tests/ # Pytest suite: conftest.py (singleton-reset fixtures, synthetic + # data factories), unit/, integration/ (gpu+cluster-marked smokes). + # Default selection: pytest -m "not gpu and not cluster" -q + # (runs in CI via .github/workflows/tests.yml); see CONTRIBUTING.md ``` --- diff --git a/.github/workflows/tests.yml b/.github/workflows/tests.yml new file mode 100644 index 00000000..5c6842db --- /dev/null +++ b/.github/workflows/tests.yml @@ -0,0 +1,57 @@ +# Unit test suite (pytest) +name: Tests + +# Run on every PR and on pushes to master +on: + push: + branches: ["master"] + pull_request: + branches: ["**"] + workflow_dispatch: + +permissions: + contents: read + +jobs: + tests: + runs-on: ubuntu-latest + strategy: + # Let both Python versions report rather than cancelling the sibling on first failure + fail-fast: false + matrix: + # The two supported runtimes: conda env (3.10) and NGC container (3.12) + python-version: ["3.10", "3.12"] + + steps: + # Pull repo contents into GitHub Actions runner + - name: Checkout code + uses: actions/checkout@v5 + + # Install the matrix Python version, with built-in pip caching keyed on the + # dependency pins (invalidates when requirements-container.txt changes) + - name: Set up Python ${{ matrix.python-version }} + uses: actions/setup-python@v5 + with: + python-version: ${{ matrix.python-version }} + cache: "pip" + # Include the workflow itself so the inline pins below also invalidate the cache + cache-dependency-path: | + requirements-container.txt + .github/workflows/tests.yml + + # Mirror the NGC container's dependency surface on a CPU-only runner: + # tensorflow-cpu stands in for the container's GPU TF 2.17 build, the pinned + # extras come from requirements-container.txt, and h5py/hdf5plugin/pandas/ + # psutil/pytest are installed explicitly because the NGC base image ships + # them (so the requirements file intentionally omits them; hdf5plugin is + # pinned explicitly in environment.yml too — don't rely on setigen pulling + # it transitively) + - name: Install dependencies + run: | + python -m pip install --upgrade pip + pip install "tensorflow-cpu==2.17.*" -r requirements-container.txt h5py hdf5plugin pandas psutil pytest + + # GPU- and cluster-marked tests need physical GPUs / cluster-resident data; + # everything else (including slow CPU TF graph tests) runs here + - name: Run test suite + run: pytest -m "not gpu and not cluster" -q diff --git a/CLAUDE.md b/CLAUDE.md index 2f651694..fd01916d 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -4,7 +4,7 @@ Always-on, bare-essentials rules for any coding agent in this repo. The full dee ## Project -Aetherscan: Breakthrough Listen's deep-learning SETI pipeline. Two-stage ML (Beta-VAE → Random Forest), single-node multi-GPU distributed train/inference. Sole entry point: `src/aetherscan/main.py`. `tests/` is a placeholder — **no test suite yet**, so don't claim `pytest` passes; verify changes another way and say how. +Aetherscan: Breakthrough Listen's deep-learning SETI pipeline. Two-stage ML (Beta-VAE → Random Forest), single-node multi-GPU distributed train/inference. Sole entry point: `src/aetherscan/main.py`. Pytest suite lives in `tests/` (unit + gpu/cluster-marked integration; see [CONTRIBUTING.md](CONTRIBUTING.md#testing)) — run it before claiming a change works, and ship unit tests with new logic. Most test modules import TensorFlow at collection; if the full dependency stack isn't available in your environment, run the subset you can and say exactly what you ran. ## Run / lint @@ -15,6 +15,8 @@ Aetherscan: Breakthrough Listen's deep-learning SETI pipeline. Two-stage ML (Bet PYTHONPATH=src python -m aetherscan.main {train|inference} --save-tag final_v1 # Lint + format (also enforced via pre-commit) ruff check src/ && ruff format src/ +# Tests (default selection = what CI runs; gpu/cluster-marked tests need a cluster) +pytest -m "not gpu and not cluster" -q ``` ## Hard rules (don't break these) diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index b3d6e38b..4ae9c149 100644 --- a/CONTRIBUTING.md +++ b/CONTRIBUTING.md @@ -190,7 +190,10 @@ Aetherscan/ │ ├── __init__.py # Manager exports │ └── manager.py # Resource lifecycle management ├── docs/ # Documentation (placeholder; no docs yet) -├── tests/ # Test suite (placeholder; no tests yet) +├── tests/ # Pytest suite (see Testing section below) +│ ├── conftest.py # Singleton-reset fixtures, tmp paths, synthetic data factories +│ ├── unit/ # Fast, hardware-independent unit tests (run in CI) +│ └── integration/ # gpu+cluster-marked end-to-end smoke tests ├── utils/ # Utility scripts │ ├── fetch_run_outputs.sh # rsync a run's outputs from a remote node │ ├── find_optimal_configs.py # Per-host config helper @@ -510,9 +513,43 @@ Prefer descriptive, full-word names: `num_training_rounds`, `per_replica_batch_s ## Testing -> [!WARNING] -> -> # TODO: update this section once test suite is operational +The test suite lives in `tests/` and runs on [pytest](https://docs.pytest.org/) (configured in `pyproject.toml` under `[tool.pytest.ini_options]` — `pythonpath = ["src"]` makes `import aetherscan` work without a `PYTHONPATH` prefix). + +### Running the suite + +```bash +# Default selection — everything that doesn't need GPUs or cluster-resident data. +# This is exactly what CI runs (.github/workflows/tests.yml). +pytest -m "not gpu and not cluster" -q + +# Cluster integration smokes (end-to-end train + CSV inference), inside the NGC container: +./utils/run_container.sh python -m pytest tests/ -m "gpu or cluster" -q + +# A single file / test while iterating: +pytest tests/unit/test_config.py -q +pytest tests/unit/test_db.py::TestFlushSentinel -q +``` + +### Markers + +| Marker | Meaning | +| ------------- | ------------------------------------------------------------------------------------------- | +| `slow` | Slower tests (e.g. builds TF graphs on CPU); still included in the default CI selection | +| `gpu` | Requires one or more physical GPUs; excluded from CI | +| `cluster` | Requires cluster-resident data/models (blpc3/bla0); excluded from CI | +| `integration` | End-to-end pipeline runs; skips the per-test singleton/env isolation fixture in `conftest.py` | + +### Test environment & fixtures + +- Every non-integration test runs isolated: `tests/conftest.py` points `AETHERSCAN_{DATA,MODEL,OUTPUT}_PATH` at a pytest `tmp_path`, resets all singletons (`Config`, `ResourceManager`, `Database`, `Logger`, `ResourceMonitor`) via their `_reset()` hooks, and re-runs `init_config()` — so tests never touch real data or leak state into each other. +- Synthetic data factories are provided as fixtures: `make_background_npy` (tiny background plates), `make_h5_observation` (tiny filterbank-style `.h5` files), and `make_inference_csv` (tiny cadence-grouping CSVs). +- CI (`.github/workflows/tests.yml`) runs the default selection on Ubuntu with `tensorflow-cpu==2.17.*` + the pins from `requirements-container.txt`, on Python 3.10 and 3.12 (the conda and container runtimes respectively). A few tests assert Linux-only behavior (e.g. PSS-based memory accounting) and self-skip elsewhere. + +### Writing tests + +- Every PR that adds or changes logic should ship unit tests for it (`tests/unit/`), reusing the conftest fixtures rather than instantiating singletons directly. +- Mark anything that needs hardware or cluster data with the appropriate marker so the default selection stays runnable on a laptop and in CI. +- Test style follows the same ruff lint+format rules as `src/` (pre-commit runs on `tests/` too). --- diff --git a/pyproject.toml b/pyproject.toml index 109eb90f..386c214d 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -55,6 +55,19 @@ Issues = "https://github.com/zachtheyek/Aetherscan/issues" [project.optional-dependencies] dev = ["ruff"] # Python linter +[tool.pytest.ini_options] +# Make `import aetherscan` work without a PYTHONPATH=src prefix +pythonpath = ["src"] +testpaths = ["tests"] +# Reject undeclared markers so a typo'd @pytest.mark.gpu can't silently slip into CI +addopts = "--strict-markers" +markers = [ + "slow: slower tests (e.g. builds TF graphs on CPU); included in the default CI selection", + "gpu: requires one or more physical GPUs; excluded from CI (-m 'not gpu and not cluster')", + "cluster: requires cluster-resident data/models (blpc3/bla0); excluded from CI", + "integration: end-to-end pipeline runs (skips the per-test singleton/env isolation fixture)", +] + [tool.ruff] # Configure linter target-version = "py310" # Target Python version diff --git a/tests/conftest.py b/tests/conftest.py new file mode 100644 index 00000000..75c14d4f --- /dev/null +++ b/tests/conftest.py @@ -0,0 +1,223 @@ +""" +Shared pytest fixtures for the Aetherscan test suite. + +Provides three things: +1. An autouse fixture that isolates every (non-integration) test: AETHERSCAN_{DATA,MODEL,OUTPUT}_PATH + point at tmp_path, all singletons are reset via their _reset() hooks, and a fresh Config is + initialized. Teardown stops any started background threads/pools and resets everything again. +2. Synthetic data factories: tiny .npy background plates, tiny .h5 observation files, and tiny + inference CSVs — sized for speed, shaped like the real thing. +3. A headless matplotlib backend for CI (set before any aetherscan module imports pyplot). +""" + +# TODO: integration tests that verify clean shutdown under adverse conditions +# (SIGTERM/SIGINT mid-run, resource cleanup) — carried over from the old tests/placeholder +# TODO: build tests that tagged releases deploy properly to various environments +# (PyPI, HuggingFace, etc.) once release engineering lands — carried over from placeholder + +from __future__ import annotations + +import contextlib +import csv +import os +import signal +import sys + +import numpy as np +import pytest + +# Force a non-interactive backend before any test imports matplotlib.pyplot +# (train.py imports it at module level; CI runners have no display). +os.environ.setdefault("MPLBACKEND", "Agg") +# Quiet TF's C++ INFO/WARNING chatter in test output. +os.environ.setdefault("TF_CPP_MIN_LOG_LEVEL", "2") + + +def _reset_all_singletons(): + """Reset every singleton class via its _reset() teardown hook. + + Imports are deferred so integration runs (which skip the isolation fixture entirely) + never pull TensorFlow into the pytest parent process. + """ + from aetherscan.config import Config # noqa: PLC0415 + from aetherscan.db.db import Database # noqa: PLC0415 + from aetherscan.logger.logger import Logger # noqa: PLC0415 + from aetherscan.manager.manager import ResourceManager # noqa: PLC0415 + from aetherscan.monitor.monitor import ResourceMonitor # noqa: PLC0415 + + for cls in (Config, Database, Logger, ResourceManager, ResourceMonitor): + cls._reset() + + +def _teardown_singletons(): + """Stop background resources still held by singletons, then reset them all. + + Mirrors ResourceManager.cleanup_all()'s ordering (pools -> shared memory -> monitor -> + db -> logger) without its logging side effects, so a test that leaked a resource can't + poison the next test. + """ + from aetherscan.db.db import Database # noqa: PLC0415 + from aetherscan.logger.logger import Logger # noqa: PLC0415 + from aetherscan.manager.manager import ResourceManager # noqa: PLC0415 + from aetherscan.monitor.monitor import ResourceMonitor # noqa: PLC0415 + + manager = ResourceManager._instance + if manager is not None: + for managed_pool in list(manager._pools): + with contextlib.suppress(Exception): + managed_pool.close(timeout=5.0) + for managed_shm in list(manager._shared_memories): + with contextlib.suppress(Exception): + managed_shm.close() + # Each ResourceManager construction registers an atexit callback; unregister it so + # per-test instances don't accumulate stale cleanup hooks over the session. + import atexit # noqa: PLC0415 + + atexit.unregister(manager.cleanup_all) + + monitor = ResourceMonitor._instance + if monitor is not None: + with contextlib.suppress(Exception): + monitor.stop() + + db = Database._instance + if db is not None: + with contextlib.suppress(Exception): + db.stop() + + logger_instance = Logger._instance + if logger_instance is not None: + with contextlib.suppress(Exception): + logger_instance.stop() + + _reset_all_singletons() + + +@pytest.fixture(autouse=True) +def aetherscan_isolated_env(request, tmp_path, monkeypatch): + """Autouse: fresh singletons + tmp-path-scoped env for every test. + + Integration tests are exempt (marker `integration`): they exercise the real pipeline as a + subprocess and must inherit the real AETHERSCAN_* environment, and their pytest parent + process should never import TensorFlow. + """ + if request.node.get_closest_marker("integration"): + yield + return + + data_path = tmp_path / "data" + model_path = tmp_path / "models" + output_path = tmp_path / "outputs" + for sub in ("training", "testing", "inference"): + (data_path / sub).mkdir(parents=True, exist_ok=True) + model_path.mkdir(exist_ok=True) + output_path.mkdir(exist_ok=True) + + monkeypatch.setenv("AETHERSCAN_DATA_PATH", str(data_path)) + monkeypatch.setenv("AETHERSCAN_MODEL_PATH", str(model_path)) + monkeypatch.setenv("AETHERSCAN_OUTPUT_PATH", str(output_path)) + # Keep unit tests from ever talking to Slack, whatever the host env holds. + monkeypatch.delenv("SLACK_BOT_TOKEN", raising=False) + monkeypatch.delenv("SLACK_CHANNEL", raising=False) + + # ResourceManager's constructor installs SIGINT/SIGTERM handlers; snapshot and restore + # them so per-test instances don't leave stale bound-method handlers behind. + original_sigint = signal.getsignal(signal.SIGINT) + original_sigterm = signal.getsignal(signal.SIGTERM) + original_stdout, original_stderr = sys.stdout, sys.stderr + + from aetherscan.config import init_config # noqa: PLC0415 + + _reset_all_singletons() + init_config() + + yield + + _teardown_singletons() + signal.signal(signal.SIGINT, original_sigint) + signal.signal(signal.SIGTERM, original_sigterm) + sys.stdout, sys.stderr = original_stdout, original_stderr + + +@pytest.fixture +def make_background_npy(tmp_path): + """Factory for tiny .npy background plates shaped like real training backgrounds. + + Returns a callable(filename, n_cadences=8, width_bin=512) -> Path that writes a positive + float32 array of shape (n_cadences, 6, 16, width_bin) under the tmp data/training dir. + """ + + def _make(filename="backgrounds.npy", n_cadences=8, width_bin=512): + rng = np.random.default_rng(11) + # Chi-squared-ish positive noise, loosely resembling detected power spectra. + plate = rng.chisquare(df=4, size=(n_cadences, 6, 16, width_bin)).astype(np.float32) + path = tmp_path / "data" / "training" / filename + path.parent.mkdir(parents=True, exist_ok=True) + np.save(path, plate) + return path + + return _make + + +@pytest.fixture +def make_h5_observation(tmp_path): + """Factory for tiny .h5 observation files matching the filterbank-style layout. + + Returns a callable(filename, n_chans=2048, time_bins=16) -> Path that writes an .h5 with a + 'data' dataset of shape (time_bins, 1, n_chans) plus fch1/foff/nchans attrs (plain h5py, + no bitshuffle). + """ + import h5py # noqa: PLC0415 + + def _make(filename="obs.h5", n_chans=2048, time_bins=16): + rng = np.random.default_rng(7) + data = rng.chisquare(df=4, size=(time_bins, 1, n_chans)).astype(np.float32) + path = tmp_path / filename + with h5py.File(path, "w") as hf: + dset = hf.create_dataset("data", data=data) + dset.attrs["fch1"] = 8421.386717353016 + dset.attrs["foff"] = -2.7939677238464355e-06 + dset.attrs["nchans"] = n_chans + return path + + return _make + + +@pytest.fixture +def make_inference_csv(tmp_path): + """Factory for tiny inference CSVs in the cadence-grouping layout. + + Returns a callable(filename, groups) -> Path. `groups` is a list of (key_dict, h5_paths) + pairs; each h5 path becomes one row carrying the group's key columns. The h5-path column + name is read from InferenceConfig.cadence_h5_path_col on the initialized singleton, and + the default group keys mirror InferenceConfig.cadence_group_by_cols. + """ + from aetherscan.config import get_config # noqa: PLC0415 + + def _make(filename="subset.csv", groups=None): + if groups is None: + groups = [ + ( + { + "Target": "HIP110750", + "Session": "AGBT21B_999_31", + "Band": "L", + "Cadence ID": "0", + "Frequency": "1400", + }, + [f"/data/obs_{i}.h5" for i in range(6)], + ) + ] + h5_col = get_config().inference.cadence_h5_path_col + fieldnames = list(groups[0][0].keys()) + [h5_col] + path = tmp_path / "data" / "inference" / filename + path.parent.mkdir(parents=True, exist_ok=True) + with open(path, "w", newline="") as f: + writer = csv.DictWriter(f, fieldnames=fieldnames) + writer.writeheader() + for key_dict, h5_paths in groups: + for h5 in h5_paths: + writer.writerow({**key_dict, h5_col: h5}) + return path + + return _make diff --git a/tests/integration/conftest.py b/tests/integration/conftest.py new file mode 100644 index 00000000..775fe117 --- /dev/null +++ b/tests/integration/conftest.py @@ -0,0 +1,53 @@ +"""Shared fixtures for the cluster integration smokes. + +Both smoke tests launch `python -m aetherscan.main ...` as a subprocess against the real +cluster environment; the repo root, the AETHERSCAN_* path resolution (env var or the cluster +defaults baked into utils/run_container.sh), and the PYTHONPATH-injected subprocess launcher +live here so the two tests can't drift apart. +""" + +from __future__ import annotations + +import os +import subprocess +import sys + +import pytest + +_REPO_ROOT = os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "..")) +_SUBPROCESS_TIMEOUT_SECONDS = 7200 + + +@pytest.fixture +def cluster_paths(): + """(data_path, model_path, output_path) from AETHERSCAN_* env, falling back to the same + cluster defaults utils/run_container.sh uses.""" + return ( + os.environ.get("AETHERSCAN_DATA_PATH", "/datax/scratch/zachy/data/aetherscan"), + os.environ.get("AETHERSCAN_MODEL_PATH", "/datax/scratch/zachy/models/aetherscan"), + os.environ.get("AETHERSCAN_OUTPUT_PATH", "/datax/scratch/zachy/outputs/aetherscan"), + ) + + +@pytest.fixture +def run_pipeline(): + """Launcher for `python -m aetherscan.main ` from the repo root, with /src + prepended to PYTHONPATH so the subprocess can `import aetherscan` both inside the NGC + container and in a bare env. Returns the CompletedProcess (check=False — callers assert + on returncode so failures surface the log tail, not a bare CalledProcessError).""" + + def _run(args: list[str]) -> subprocess.CompletedProcess: + env = dict(os.environ) + src = os.path.join(_REPO_ROOT, "src") + env["PYTHONPATH"] = src + os.pathsep + env["PYTHONPATH"] if env.get("PYTHONPATH") else src + return subprocess.run( + [sys.executable, "-m", "aetherscan.main", *args], + cwd=_REPO_ROOT, + env=env, + capture_output=True, + text=True, + timeout=_SUBPROCESS_TIMEOUT_SECONDS, + check=False, + ) + + return _run diff --git a/tests/integration/test_inference_smoke.py b/tests/integration/test_inference_smoke.py new file mode 100644 index 00000000..c526c38f --- /dev/null +++ b/tests/integration/test_inference_smoke.py @@ -0,0 +1,70 @@ +"""End-to-end CSV-inference smoke test (cluster-only, blpc3). + +Runs subset CSV inference against the persisted dummy model (test_v17) as a subprocess of +`python -m aetherscan.main inference` and asserts a clean exit plus a config snapshot. The +subset CSV's raw .h5 files live under /datag (blpc3 only), which run_container.sh does not +bind by default — either bind it explicitly: + + SINGULARITY_BIND=/datag ./utils/run_container.sh python -m pytest tests/ -m "gpu or cluster" -q + +or rely on the resume path: preprocessing skips any cadence whose stamp .npy already exists +under /preprocessed, so a previously-preprocessed subset runs without /datag. +""" + +from __future__ import annotations + +import glob +import os +import shutil +from datetime import datetime + +import pytest + +pytestmark = [pytest.mark.integration, pytest.mark.gpu, pytest.mark.cluster] + +_MODEL_TAG = "test_v17" # persisted dummy model on blpc3 +_CSV_NAME = "subset_test.csv" # 2 complete 6-observation cadences + + +def test_inference_smoke(cluster_paths, run_pipeline): + if shutil.which("nvidia-smi") is None: + pytest.skip("requires a GPU host (nvidia-smi not found)") + + data_path, model_path, output_path = cluster_paths + + encoder = os.path.join(model_path, f"vae_encoder_{_MODEL_TAG}.keras") + rf = os.path.join(model_path, f"random_forest_{_MODEL_TAG}.joblib") + saved_config = os.path.join(model_path, f"config_{_MODEL_TAG}.json") + csv_path = os.path.join(data_path, "inference", _CSV_NAME) + for required in (encoder, rf, saved_config, csv_path): + if not os.path.exists(required): + pytest.skip(f"required cluster artifact missing: {required}") + + # Raw .h5 reads need /datag; already-preprocessed stamps make it optional. + preprocessed = glob.glob(os.path.join(output_path, "preprocessed", "subset_test_*.npy")) + if not os.path.exists("/datag") and not preprocessed: + pytest.skip("/datag not mounted and no preprocessed subset stamps to resume from") + + tag = datetime.now().strftime("%Y%m%d_%H%M%S") + proc = run_pipeline( + [ + "inference", + "--encoder-path", + encoder, + "--rf-path", + rf, + "--config-path", + saved_config, + "--inference-files", + _CSV_NAME, + "--save-tag", + tag, + "--max-retries", + "1", + ] + ) + + tail = "\n".join(proc.stdout.splitlines()[-40:]) + assert proc.returncode == 0, f"inference smoke run failed (tag={tag}); last output:\n{tail}" + assert "Inference completed successfully!" in proc.stdout + assert os.path.exists(os.path.join(output_path, f"config_{tag}.json")) diff --git a/tests/integration/test_train_smoke.py b/tests/integration/test_train_smoke.py new file mode 100644 index 00000000..078b540d --- /dev/null +++ b/tests/integration/test_train_smoke.py @@ -0,0 +1,92 @@ +"""End-to-end training smoke test (cluster-only). + +Runs the known-good blpc3 5-GPU smoke config from the repo runbook as a subprocess of +`python -m aetherscan.main train` and asserts the run exits cleanly with all final model +artifacts on disk. Batch/sample sizes here divide cleanly for exactly 5 replicas — this test +is sized for blpc3 (5x RTX PRO 6000). Run inside the NGC container: + + ./utils/run_container.sh python -m pytest tests/ -m "gpu or cluster" -q +""" + +from __future__ import annotations + +import os +import re +import shutil + +import pytest + +pytestmark = [pytest.mark.integration, pytest.mark.gpu, pytest.mark.cluster] + +# The known-good blpc3 smoke config: repo defaults are only divisible for 4 or 6 replicas, so +# the batch/sample geometry is overridden to divide for 5. Mirrored (as the subject under +# test) by _SMOKE_FLAGS_5_REPLICAS in tests/unit/test_cli_validation.py — keep in sync. +_SMOKE_FLAGS = [ + "--max-retries", + "1", + "--num-training-rounds", + "2", + "--epochs-per-round", + "2", + "--per-replica-batch-size", + "4", + "--per-replica-val-batch-size", + "4", + "--effective-batch-size", + "20", + "--num-samples-beta-vae", + "200", + "--num-samples-rf", + "200", + "--latent-viz-num-cadences-per-type", + "5", +] + + +def _next_test_tag(model_path: str) -> str: + """Scan existing artifacts for test_vNN tags and return the next unused one — reusing a + tag causes stale-artifact confusion. + + NOTE: the scan-then-pick is not atomic — two runs launched concurrently could both claim + the same tag. Fine for its intended use (manual, sequential smoke runs on a shared + cluster); revisit if these smokes are ever launched in parallel. + """ + versions = [0] + for root in (model_path, os.path.join(model_path, "checkpoints")): + if not os.path.isdir(root): + continue + for entry in os.listdir(root): + match = re.search(r"test_v(\d+)", entry) + if match: + versions.append(int(match.group(1))) + return f"test_v{max(versions) + 1}" + + +def test_train_smoke(cluster_paths, run_pipeline): + if shutil.which("nvidia-smi") is None: + pytest.skip("requires a GPU host (nvidia-smi not found)") + + data_path, model_path, output_path = cluster_paths + if not os.path.isdir(os.path.join(data_path, "training")): + pytest.skip(f"training data not found under {data_path}") + + tag = _next_test_tag(model_path) + proc = run_pipeline(["train", "--save-tag", tag, *_SMOKE_FLAGS]) + + tail = "\n".join(proc.stdout.splitlines()[-40:]) + assert proc.returncode == 0, f"train smoke run failed (tag={tag}); last output:\n{tail}" + assert "Training completed successfully!" in proc.stdout + + # Final artifacts: encoder/decoder/RF under model_path, config snapshot under output_path. + for artifact in ( + f"vae_encoder_{tag}.keras", + f"vae_decoder_{tag}.keras", + f"random_forest_{tag}.joblib", + ): + assert os.path.exists(os.path.join(model_path, artifact)), f"missing {artifact}" + assert os.path.exists(os.path.join(output_path, f"config_{tag}.json")) + + # Per-round checkpoints for both rounds. + checkpoints = os.path.join(model_path, "checkpoints") + for round_tag in ("round_01", "round_02"): + assert os.path.exists(os.path.join(checkpoints, f"vae_encoder_{round_tag}.keras")) diff --git a/tests/placeholder b/tests/placeholder deleted file mode 100644 index 3f1c5d6a..00000000 --- a/tests/placeholder +++ /dev/null @@ -1,3 +0,0 @@ -# TODO: -# Integration tests to verify clean shutdown under various conditions -# Build tests to ensure tagged releases can be deployed properly to various environments (PyPI, HuggingFace, etc.) diff --git a/tests/unit/test_cli_help_sync.py b/tests/unit/test_cli_help_sync.py new file mode 100644 index 00000000..50f94b41 --- /dev/null +++ b/tests/unit/test_cli_help_sync.py @@ -0,0 +1,64 @@ +"""Guards against drift between cli.py's --help output and the README CLI Reference blocks. + +The README's `## CLI Reference` fenced blocks are generated from cli.py by +utils/print_cli_help.py and must be regenerated whenever a CLI flag changes (see +CONTRIBUTING.md). This test drives that util's per-subcommand output under the same pinned +terminal width the docs were generated with and asserts each README block still matches +byte-for-byte, so a forgotten regeneration fails CI instead of silently rotting the docs. +""" + +from __future__ import annotations + +import importlib.util +import io +import re +import sys +from contextlib import redirect_stdout +from pathlib import Path + +import pytest + +_REPO_ROOT = Path(__file__).resolve().parents[2] +_README = _REPO_ROOT / "README.md" + +# print_cli_help.py is a standalone script under utils/, not an importable package, so load it +# straight from its path — this drives the exact code a maintainer runs to regenerate the docs. +_spec = importlib.util.spec_from_file_location( + "print_cli_help", _REPO_ROOT / "utils" / "print_cli_help.py" +) +_print_cli_help = importlib.util.module_from_spec(_spec) +_spec.loader.exec_module(_print_cli_help) + +# The subparser targets print_cli_help.py accepts; each maps to one README CLI Reference block. +_TARGETS = ("top", "train", "inference") + + +def _util_help(target: str, monkeypatch: pytest.MonkeyPatch) -> str: + """Capture print_cli_help.py's stdout for one target under a pinned 80-column width.""" + # argparse reads COLUMNS when it builds the HelpFormatter, so line wrapping depends on it. + # Pin it to reproduce exactly the width the README was generated with (print_cli_help.py + # defaults COLUMNS to "80"), independent of the ambient terminal. + monkeypatch.setenv("COLUMNS", "80") + monkeypatch.setattr(sys, "argv", ["print_cli_help.py", target]) + buf = io.StringIO() + with redirect_stdout(buf): + _print_cli_help.main() + return buf.getvalue() + + +def _readme_block(target: str) -> str: + """Extract the fenced block whose intro line cites `print_cli_help.py `.""" + lines = _README.read_text().splitlines() + anchor = re.compile(rf"print_cli_help\.py {re.escape(target)}\b") + intro = next(n for n, line in enumerate(lines) if anchor.search(line)) + open_fence = next(n for n in range(intro + 1, len(lines)) if lines[n].strip() == "```") + close_fence = next(n for n in range(open_fence + 1, len(lines)) if lines[n].strip() == "```") + return "\n".join(lines[open_fence + 1 : close_fence]) + "\n" + + +@pytest.mark.parametrize("target", _TARGETS) +def test_readme_cli_reference_matches_help(target: str, monkeypatch: pytest.MonkeyPatch) -> None: + assert _util_help(target, monkeypatch) == _readme_block(target), ( + f"README CLI Reference for '{target}' is stale; regenerate with " + f"`PYTHONPATH=src python utils/print_cli_help.py {target}` (or `all`)." + ) diff --git a/tests/unit/test_cli_validation.py b/tests/unit/test_cli_validation.py new file mode 100644 index 00000000..33752ed1 --- /dev/null +++ b/tests/unit/test_cli_validation.py @@ -0,0 +1,363 @@ +"""Unit tests for aetherscan.cli: tag pattern, validation matrix, cross-param solver, and +apply_saved_config precedence (defaults < saved config < CLI args).""" + +from __future__ import annotations + +import json + +import pytest + +from aetherscan import cli +from aetherscan.cli import ( + _TAG_PATTERN, + _check_cross_constraints, + _solve_cross_param_constraints, + apply_args_to_config, + apply_saved_config, + collect_validation_errors, + setup_argument_parser, +) +from aetherscan.config import get_config + +# The blpc3 smoke config from the repo's known-good runbook: divisible for exactly 5 replicas. +# Mirrored by _SMOKE_FLAGS in tests/integration/test_train_smoke.py (which runs it for real on +# the cluster) — keep the two in sync. +_SMOKE_FLAGS_5_REPLICAS = [ + "--per-replica-batch-size", + "4", + "--per-replica-val-batch-size", + "4", + "--effective-batch-size", + "20", + "--num-samples-beta-vae", + "200", + "--num-samples-rf", + "200", + "--latent-viz-num-cadences-per-type", + "5", +] + + +def _parse(argv): + return setup_argument_parser().parse_args(argv) + + +def _cross_param_errors(args, num_replicas): + errors = collect_validation_errors(args, num_replicas) + return [e for e in errors if e.fix_kind == "cross_param"] + + +class TestTagPattern: + @pytest.mark.parametrize( + "tag", + ["20260712_123456", "final_v1", "final_v12", "round_01", "round_5", "test_v1", "test_v17"], + ) + def test_accepted_formats(self, tag): + assert _TAG_PATTERN.match(tag) + + @pytest.mark.parametrize( + "tag", + [ + "smoke_blackwell", # free-form slug + "final_1", # missing v + "final_v", # missing version number + "test_17", # missing v + "TEST_V1", # wrong case + "2026_0712", # malformed timestamp + "20260712-123456", # wrong separator + "20260712_12345", # HHMMS (5 digits) + "round_", # missing round number + " test_v1", # leading whitespace + "test_v1 ", # trailing whitespace + "", + ], + ) + def test_rejected_formats(self, tag): + assert not _TAG_PATTERN.match(tag) + + +class TestCrossReplicaDivisibilityMatrix: + """Default config divides cleanly across 4 or 6 replicas but NOT 5; the blpc3 smoke config + is the inverse. collect_validation_errors must reproduce that matrix exactly.""" + + @pytest.mark.parametrize("num_replicas", [4, 6]) + def test_defaults_valid_for_4_and_6_replicas(self, num_replicas): + assert _cross_param_errors(_parse(["train"]), num_replicas) == [] + + def test_defaults_invalid_for_5_replicas(self): + errors = _cross_param_errors(_parse(["train"]), 5) + violated = {e.field for e in errors} + assert violated == { + "training.effective_batch_size", + "training.num_samples_beta_vae", + "training.num_samples_rf", + "training.latent_viz_num_cadences_per_type", + } + + def test_smoke_config_valid_for_5_replicas(self): + assert _cross_param_errors(_parse(["train", *_SMOKE_FLAGS_5_REPLICAS]), 5) == [] + + @pytest.mark.parametrize("num_replicas", [4, 6]) + def test_smoke_config_invalid_for_4_and_6_replicas(self, num_replicas): + assert _cross_param_errors(_parse(["train", *_SMOKE_FLAGS_5_REPLICAS]), num_replicas) + + def test_unknown_replica_count_skips_cross_checks(self): + # num_replicas=None (no GPUs detectable) must skip the divisibility section entirely, + # even for a config that would fail on any replica count. + assert _cross_param_errors(_parse(["train", *_SMOKE_FLAGS_5_REPLICAS]), None) == [] + + +class TestSemanticChecks: + def test_save_tag_format_error(self): + errors = collect_validation_errors(_parse(["train", "--save-tag", "bogus"]), None) + assert any(e.field == "checkpoint.save_tag" and e.fix_kind == "format" for e in errors) + + def test_load_tag_format_error(self): + errors = collect_validation_errors(_parse(["train", "--load-tag", "bogus"]), None) + assert any(e.field == "checkpoint.load_tag" and e.fix_kind == "format" for e in errors) + + def test_num_samples_divisible_by_4(self): + errors = collect_validation_errors(_parse(["train", "--num-samples-beta-vae", "202"]), None) + assert any(e.field == "training.num_samples_beta_vae" and e.divisor == 4 for e in errors) + + def test_curriculum_schedule_enum(self): + errors = collect_validation_errors( + _parse(["train", "--curriculum-schedule", "sigmoid"]), None + ) + assert any( + e.field == "training.curriculum_schedule" and e.fix_kind == "enum" for e in errors + ) + + def test_step_schedule_sum_constraint(self): + argv = [ + "train", + "--curriculum-schedule", + "step", + "--num-training-rounds", + "10", + "--step-easy-rounds", + "3", + "--step-hard-rounds", + "4", + ] + errors = collect_validation_errors(_parse(argv), None) + assert any( + e.field == "training.step_easy_rounds" and e.fix_kind == "cross_param" for e in errors + ) + + def test_snr_curriculum_ordering(self): + argv = ["train", "--initial-snr-range", "5", "--final-snr-range", "10"] + errors = collect_validation_errors(_parse(argv), None) + assert any(e.field == "training.initial_snr_range" for e in errors) + + def test_missing_train_files_reported(self): + errors = collect_validation_errors(_parse(["train"]), None) + # Default train_files don't exist under the tmp data path. + file_errors = [e for e in errors if e.fix_kind == "file_exists"] + assert {e.field for e in file_errors} == {"data.train_files"} + assert len(file_errors) == len(get_config().data.train_files) + + def test_existing_train_files_pass(self, tmp_path): + config = get_config() + for filename in config.data.train_files: + (tmp_path / "data" / "training" / filename).touch() + errors = collect_validation_errors(_parse(["train"]), None) + assert [e for e in errors if e.fix_kind == "file_exists"] == [] + + def test_inference_requires_model_artifacts(self): + errors = collect_validation_errors(_parse(["inference"]), None) + fields = {e.field for e in errors if e.fix_kind == "file_exists"} + assert { + "inference.encoder_path", + "inference.rf_path", + "inference.config_path", + } <= fields + + def test_inference_stamp_width_must_match_width_bin(self, make_inference_csv): + csv_path = make_inference_csv("subset.csv") + argv = [ + "inference", + "--inference-files", + csv_path.name, + "--stamp-width", + "2048", # width_bin default is 4096 + ] + errors = collect_validation_errors(_parse(argv), None) + assert any(e.field == "inference.stamp_width" for e in errors) + + +class TestCrossParamSolver: + _VALID_BASE = { + # The repo defaults: valid for 4 and 6 replicas. + "num_samples_beta_vae": 499200, + "num_samples_rf": 99840, + "train_val_split": 0.8, + "per_replica_batch_size": 128, + "effective_batch_size": 3072, + "per_replica_val_batch_size": 80, + } + + def test_check_cross_constraints_matrix(self): + assert _check_cross_constraints(**self._VALID_BASE, num_replicas_list=[4]) + assert _check_cross_constraints(**self._VALID_BASE, num_replicas_list=[6]) + assert not _check_cross_constraints(**self._VALID_BASE, num_replicas_list=[5]) + + def test_solver_returns_base_when_already_valid(self): + assert _solve_cross_param_constraints(self._VALID_BASE, [4, 6]) == self._VALID_BASE + + def test_solver_respects_candidate_budget(self): + # The default search ranges exceed a tiny budget, so the solver must bail with None + # rather than grind through the grid. + base = dict(self._VALID_BASE) + base["effective_batch_size"] = 3070 # invalidate so the grid search is attempted + assert _solve_cross_param_constraints(base, [4], max_candidates=10) is None + + def test_solver_finds_nearest_valid_config(self, monkeypatch): + # Shrink the search ranges to a tractable grid and verify the solver picks the + # L1-nearest satisfying combination. + monkeypatch.setattr( + cli, + "_SEARCH_RANGES", + { + "num_samples_beta_vae": (160, 320, 80), + "num_samples_rf": (40, 80, 20), + "per_replica_batch_size": (4, 8, 4), + "effective_batch_size": (16, 64, 16), + "per_replica_val_batch_size": (4, 8, 4), + }, + ) + base = { + "num_samples_beta_vae": 250, + "num_samples_rf": 50, + "train_val_split": 0.8, + "per_replica_batch_size": 5, + "effective_batch_size": 20, + "per_replica_val_batch_size": 5, + } + solution = _solve_cross_param_constraints(base, [4]) + assert solution is not None + assert solution["train_val_split"] == base["train_val_split"] # held fixed + assert _check_cross_constraints(**solution, num_replicas_list=[4]) + # Every solved field must come from the (patched) search grid. + for field_name, (lo, hi, step) in cli._SEARCH_RANGES.items(): + assert solution[field_name] in range(lo, hi + 1, step) + # And the solution must be L1-minimal among all valid grid points. + assert self._l1(solution, base) == min( + self._l1(candidate, base) for candidate in self._valid_grid_points(base) + ) + + @staticmethod + def _l1(candidate, base): + return sum( + abs(candidate[f] - base[f]) + for f in ( + "num_samples_beta_vae", + "num_samples_rf", + "per_replica_batch_size", + "effective_batch_size", + "per_replica_val_batch_size", + ) + ) + + @staticmethod + def _valid_grid_points(base): + from itertools import product # noqa: PLC0415 + + ranges = {f: range(lo, hi + 1, step) for f, (lo, hi, step) in cli._SEARCH_RANGES.items()} + for nsb, nsr, prb, eb, prvb in product( + ranges["num_samples_beta_vae"], + ranges["num_samples_rf"], + ranges["per_replica_batch_size"], + ranges["effective_batch_size"], + ranges["per_replica_val_batch_size"], + ): + candidate = { + "num_samples_beta_vae": nsb, + "num_samples_rf": nsr, + "train_val_split": base["train_val_split"], + "per_replica_batch_size": prb, + "effective_batch_size": eb, + "per_replica_val_batch_size": prvb, + } + if _check_cross_constraints(**candidate, num_replicas_list=[4]): + yield candidate + + +class TestApplySavedConfigPrecedence: + def test_saved_config_overrides_defaults(self, tmp_path): + saved = { + "training": {"num_training_rounds": 7, "snr_base": 33}, + "checkpoint": {"save_tag": "final_v9"}, + "data_path": "/saved/data/path", + } + path = tmp_path / "saved_config.json" + path.write_text(json.dumps(saved)) + + apply_saved_config(str(path)) + config = get_config() + assert config.training.num_training_rounds == 7 + assert config.training.snr_base == 33 + # Documented sharp edge: the saved file clobbers checkpoint.save_tag too. + assert config.checkpoint.save_tag == "final_v9" + assert config.data_path == "/saved/data/path" + + def test_cli_args_override_saved_config(self, tmp_path): + saved = {"training": {"snr_base": 33, "num_training_rounds": 7}} + path = tmp_path / "saved_config.json" + path.write_text(json.dumps(saved)) + apply_saved_config(str(path)) + + args = _parse(["train", "--snr-base", "44"]) + apply_args_to_config(args) + config = get_config() + assert config.training.snr_base == 44 # CLI wins over saved + assert config.training.num_training_rounds == 7 # saved wins over default (20) + + def test_unknown_keys_and_fields_skipped(self, tmp_path): + saved = { + "not_a_section": {"whatever": 1}, + "training": {"not_a_field": 123, "snr_base": 21}, + } + path = tmp_path / "saved_config.json" + path.write_text(json.dumps(saved)) + apply_saved_config(str(path)) + config = get_config() + assert config.training.snr_base == 21 + assert not hasattr(config.training, "not_a_field") + assert not hasattr(config, "not_a_section") + + def test_missing_file_raises(self): + with pytest.raises(ValueError, match="does not exist"): + apply_saved_config("/nonexistent/config.json") + + def test_malformed_json_raises(self, tmp_path): + path = tmp_path / "bad.json" + path.write_text("{not json") + with pytest.raises(ValueError, match="not valid JSON"): + apply_saved_config(str(path)) + + +class TestApplyArgsToConfig: + def test_none_args_leave_defaults(self): + config = get_config() + default_rounds = config.training.num_training_rounds + apply_args_to_config(_parse(["train"])) + assert config.training.num_training_rounds == default_rounds + + def test_mode_scoped_flags_route_to_correct_section(self): + # --per-replica-batch-size / --max-retries exist in both subparsers; the command + # gates which config section receives them. + config = get_config() + apply_args_to_config(_parse(["train", "--per-replica-batch-size", "64"])) + assert config.training.per_replica_batch_size == 64 + + inference_default = config.inference.per_replica_batch_size + apply_args_to_config(_parse(["inference", "--per-replica-batch-size", "32"])) + assert config.inference.per_replica_batch_size == 32 + assert config.training.per_replica_batch_size == 64 # untouched by inference flag + assert inference_default != 32 + + def test_load_tag_infers_start_round(self): + config = get_config() + apply_args_to_config(_parse(["train", "--load-tag", "round_03"])) + assert config.checkpoint.start_round == 4 diff --git a/tests/unit/test_config.py b/tests/unit/test_config.py new file mode 100644 index 00000000..42ecf565 --- /dev/null +++ b/tests/unit/test_config.py @@ -0,0 +1,107 @@ +"""Unit tests for aetherscan.config: singleton semantics and to_dict() field coverage.""" + +from __future__ import annotations + +import dataclasses +import json +import os + +from aetherscan.config import Config, get_config, init_config + +# Sub-dataclass sections serialized by Config.to_dict(), keyed by attribute name. +_CONFIG_SECTIONS = ( + "db", + "manager", + "monitor", + "logger", + "beta_vae", + "rf", + "gpu", + "data", + "training", + "inference", + "checkpoint", +) + + +class TestSingletonSemantics: + def test_constructor_returns_same_instance(self): + assert Config() is Config() + + def test_get_config_returns_singleton(self): + assert get_config() is Config() + + def test_init_config_returns_singleton(self): + assert init_config() is get_config() + + def test_mutation_visible_across_accessors(self): + Config().training.num_training_rounds = 3 + assert get_config().training.num_training_rounds == 3 + + def test_reset_produces_fresh_instance(self): + first = Config() + first.training.num_training_rounds = 3 + Config._reset() + second = init_config() + assert second is not first + assert second.training.num_training_rounds != 3 + + def test_env_var_paths_respected(self, tmp_path): + # The autouse fixture points AETHERSCAN_* at tmp_path before init_config(). + config = get_config() + assert config.data_path == os.environ["AETHERSCAN_DATA_PATH"] + assert config.model_path == os.environ["AETHERSCAN_MODEL_PATH"] + assert config.output_path == os.environ["AETHERSCAN_OUTPUT_PATH"] + assert config.data_path.startswith(str(tmp_path)) + + def test_file_path_helpers(self): + config = get_config() + assert config.get_training_file_path("a.npy") == os.path.join( + config.data_path, "training", "a.npy" + ) + assert config.get_test_file_path("b.npy") == os.path.join( + config.data_path, "testing", "b.npy" + ) + assert config.get_inference_file_path("c.csv") == os.path.join( + config.data_path, "inference", "c.csv" + ) + # base_path override (used by validate_args before --data-path is applied) + assert config.get_training_file_path("a.npy", base_path="/elsewhere") == os.path.join( + "/elsewhere", "training", "a.npy" + ) + + +class TestToDict: + def test_sections_cover_every_dataclass_field(self): + """to_dict() must be updated by hand for every new config field — catch drift.""" + config = get_config() + serialized = config.to_dict() + for section in _CONFIG_SECTIONS: + expected = {f.name for f in dataclasses.fields(type(getattr(config, section)))} + actual = set(serialized[section].keys()) + assert actual == expected, ( + f"Config.to_dict()['{section}'] is out of sync with " + f"{type(getattr(config, section)).__name__}: " + f"missing={expected - actual}, extra={actual - expected}" + ) + + def test_values_round_trip(self): + config = get_config() + config.training.num_training_rounds = 5 + config.gpu.num_replicas = 4 + serialized = config.to_dict() + assert serialized["training"]["num_training_rounds"] == 5 + assert serialized["gpu"]["num_replicas"] == 4 + assert serialized["paths"]["data_path"] == config.data_path + for section in _CONFIG_SECTIONS: + for field_name, value in serialized[section].items(): + stored = getattr(getattr(config, section), field_name) + if isinstance(stored, tuple): + # json-level equality: tuples serialize as lists downstream + assert list(stored) == list(value) + else: + assert stored == value + + def test_json_serializable(self): + # Saved run configs are json.dump'ed — the dict must be JSON-native throughout. + json.dumps(get_config().to_dict()) diff --git a/tests/unit/test_data_generation.py b/tests/unit/test_data_generation.py new file mode 100644 index 00000000..0a122707 --- /dev/null +++ b/tests/unit/test_data_generation.py @@ -0,0 +1,270 @@ +"""Unit tests for aetherscan.data_generation: log-norm, intersection checks, signal injection, +create_* cadence generators, and intensity statistics.""" + +from __future__ import annotations + +import math +import random + +import numpy as np +import pytest + +from aetherscan.data_generation import ( + _compute_intensity_stats, + check_valid_intersection, + create_false, + create_true_double, + create_true_single, + log_norm, + new_cadence, +) + +# Keep injection fast: small frequency axis, real-ish resolutions. +_WIDTH_BIN = 128 +_FREQ_RES = 2.7939677238464355 # Hz +_TIME_RES = 18.25361108 # seconds + +_SIGNAL_INFO_KEYS = { + "snr", + "drift_rate", + "signal_width", + "starting_bin", + "slope_pixel", + "y_intercept", +} + +_STAT_KEYS = { + "global_mean", + "global_median", + "global_std", + "global_mad", + "global_skew", + "global_kurtosis", +} + + +@pytest.fixture(autouse=True) +def _seed_rngs(): + # NOTE: np.random.seed() intentionally seeds the legacy global RandomState — the + # production code under test (new_cadence, create_*) draws from the legacy np.random.* + # API, so switching this to a Generator (np.random.default_rng) would NOT make those + # draws deterministic. Don't "modernize" this without migrating data_generation.py first. + random.seed(11) + np.random.seed(11) + + +@pytest.fixture +def plate(make_background_npy): + """A tiny background plate loaded from a factory-written .npy (round-trips the factory).""" + path = make_background_npy("plate.npy", n_cadences=4, width_bin=_WIDTH_BIN) + loaded = np.load(path) + assert loaded.shape == (4, 6, 16, _WIDTH_BIN) + assert loaded.dtype == np.float32 + return loaded + + +class TestLogNorm: + def test_output_bounded_in_unit_interval(self): + rng = np.random.default_rng(0) + data = rng.chisquare(df=4, size=(16, 64)) + result = log_norm(data) + assert result.shape == data.shape + assert result.min() == 0.0 + assert result.max() == pytest.approx(1.0) + + def test_preserves_ordering(self): + data = np.array([[1.0, 10.0, 100.0, 1000.0]]) + result = log_norm(data) + assert np.all(np.diff(result[0]) > 0) + + def test_constant_input_maps_to_zeros(self): + result = log_norm(np.full((4, 4), 3.0)) + assert np.all(result == 0.0) + + def test_idempotence_on_normalized_scale(self): + # Re-normalizing an already [0, 1] array must stay within [0, 1] with the same + # endpoints — the transform is a monotone squash onto the unit interval. + rng = np.random.default_rng(1) + once = log_norm(rng.chisquare(df=4, size=(8, 32))) + twice = log_norm(once) + assert twice.min() == 0.0 + assert twice.max() == pytest.approx(1.0) + assert np.array_equal(np.argsort(once, axis=None), np.argsort(twice, axis=None)) + + +class TestCheckValidIntersection: + def test_parallel_lines_are_valid(self): + assert check_valid_intersection(2.0, 2.0, 0.0, 50.0) is True + + def test_intersection_inside_on_region_invalid(self): + # y = x and y = -x + 20 intersect at (10, 10): inside the first ON block [0, 16]. + assert check_valid_intersection(1.0, -1.0, 0.0, 20.0) is False + + def test_intersection_between_on_regions_valid(self): + # y = x and y = -x + 40 intersect at (20, 20): inside the first OFF block (16, 32). + assert check_valid_intersection(1.0, -1.0, 0.0, 40.0) is True + + @pytest.mark.parametrize("y_target", [0.0, 16.0, 32.0, 48.0, 64.0, 80.0]) + def test_on_region_boundaries_are_inclusive(self, y_target): + # Intersection exactly on an ON boundary counts as inside (invalid). + # y = x and y = -x + 2*y_target intersect at (y_target, y_target). + assert check_valid_intersection(1.0, -1.0, 0.0, 2 * y_target) is False + + +class TestNewCadence: + def _background(self, rng=None): + rng = rng or np.random.default_rng(3) + return rng.chisquare(df=4, size=(96, _WIDTH_BIN)) + + def test_shapes_and_signal_info(self): + data = self._background() + modified, signal_info, clamped = new_cadence( + data.copy(), + snr=20.0, + width_bin=_WIDTH_BIN, + freq_resolution=_FREQ_RES, + time_resolution=_TIME_RES, + ) + assert modified.shape == data.shape + assert set(signal_info.keys()) == _SIGNAL_INFO_KEYS + assert all(isinstance(v, float) for v in signal_info.values()) + assert signal_info["snr"] == 20.0 + assert 1 <= signal_info["starting_bin"] <= _WIDTH_BIN - 1 + assert isinstance(clamped, bool) + + def test_injection_adds_power(self): + data = self._background() + modified, _, _ = new_cadence( + data.copy(), + snr=50.0, + width_bin=_WIDTH_BIN, + freq_resolution=_FREQ_RES, + time_resolution=_TIME_RES, + ) + # A signal only adds intensity — total power must strictly increase. + assert modified.sum() > data.sum() + + def test_drift_rate_inverse_of_slope(self): + _, signal_info, clamped = new_cadence( + self._background(), + snr=20.0, + width_bin=_WIDTH_BIN, + freq_resolution=_FREQ_RES, + time_resolution=_TIME_RES, + ) + if not clamped: + slope_physical = signal_info["slope_pixel"] * (_TIME_RES / _FREQ_RES) + # drift_rate = -1/slope_physical up to the additive noise term on the slope, + # so only the sign relationship is exact. + assert math.copysign(1, signal_info["drift_rate"]) == -math.copysign(1, slope_physical) + + def test_near_zero_slope_is_clamped(self, monkeypatch): + # Force the degenerate geometry: time/freq resolution ratio ~0 makes the physical + # slope collapse below the clamp threshold when the additive noise term is 0. + # new_cadence draws random.random() in order: starting_bin, slope noise, signal width. + draws = iter([0.5, 0.0]) + monkeypatch.setattr(random, "random", lambda: next(draws, 0.25)) + monkeypatch.setattr(np.random, "choice", lambda *a, **k: 1) + _, signal_info, clamped = new_cadence( + self._background(), + snr=20.0, + width_bin=_WIDTH_BIN, + freq_resolution=1.0, + time_resolution=1e-12, + ) + assert clamped is True + # Clamped slope preserves direction but pins magnitude at the floor: |drift| = 1e6. + assert abs(signal_info["drift_rate"]) == pytest.approx(1e6) + + +class TestCreateCadences: + def _kwargs(self): + return { + "snr_base": 10.0, + "snr_range": 5.0, + "width_bin": _WIDTH_BIN, + "freq_resolution": _FREQ_RES, + "time_resolution": _TIME_RES, + } + + def _assert_common(self, final, sample_info, plate): + assert final.shape == (6, 16, _WIDTH_BIN) + assert 0 <= sample_info["background_index"] < plate.shape[0] + assert set(sample_info["intensity_stats"].keys()) == {"A", "B", "C"} + for stage_stats in sample_info["intensity_stats"].values(): + assert set(stage_stats.keys()) == _STAT_KEYS + assert isinstance(sample_info["slope_was_clamped"], bool) + # Output is log-normalized per observation. + assert final.min() >= 0.0 + assert final.max() <= 1.0 + + def test_create_false_injected(self, plate): + final, sample_info = create_false(plate, inject=True, **self._kwargs()) + self._assert_common(final, sample_info, plate) + keys = set(sample_info["signal_info"].keys()) + assert keys == {f"rfi_{k}" for k in _SIGNAL_INFO_KEYS} + + def test_create_false_no_injection(self, plate): + final, sample_info = create_false(plate, inject=False, **self._kwargs()) + self._assert_common(final, sample_info, plate) + assert sample_info["signal_info"] == {} + # Without injection, Stage B mirrors Stage A exactly. + assert sample_info["intensity_stats"]["B"] == sample_info["intensity_stats"]["A"] + # And every observation is just the log-normalized background (float32 path: the + # generator normalizes the raw float32 plate slices directly). + base = plate[sample_info["background_index"]] + for obs in range(6): + np.testing.assert_allclose(final[obs], log_norm(base[obs])) + + def test_create_true_single_injects_on_only(self, plate): + final, sample_info = create_true_single(plate, **self._kwargs()) + self._assert_common(final, sample_info, plate) + keys = set(sample_info["signal_info"].keys()) + assert keys == {f"eti_{k}" for k in _SIGNAL_INFO_KEYS} + # OFF observations (1, 3, 5) carry no injection: they equal the log-normed background. + base = plate[sample_info["background_index"]] + for obs in (1, 3, 5): + np.testing.assert_allclose(final[obs], log_norm(base[obs].astype(np.float64))) + + def test_create_true_double_injects_both(self, plate): + final, sample_info = create_true_double(plate, **self._kwargs()) + self._assert_common(final, sample_info, plate) + keys = set(sample_info["signal_info"].keys()) + expected = {f"rfi_{k}" for k in _SIGNAL_INFO_KEYS} | {f"eti_{k}" for k in _SIGNAL_INFO_KEYS} + assert keys == expected + # The two injected trajectories must not intersect inside an ON region. + assert check_valid_intersection( + sample_info["signal_info"]["rfi_slope_pixel"], + sample_info["signal_info"]["eti_slope_pixel"], + sample_info["signal_info"]["rfi_y_intercept"], + sample_info["signal_info"]["eti_y_intercept"], + ) + + +class TestComputeIntensityStats: + def test_normal_case_matches_numpy(self): + rng = np.random.default_rng(5) + data = rng.chisquare(df=4, size=(6, 16, 32)).astype(np.float32) + stats = _compute_intensity_stats(data) + assert set(stats.keys()) == _STAT_KEYS + flat = data.ravel().astype(np.float64) + assert stats["global_mean"] == pytest.approx(np.mean(flat)) + assert stats["global_median"] == pytest.approx(np.median(flat)) + assert stats["global_std"] == pytest.approx(np.std(flat)) + assert stats["global_mad"] == pytest.approx(np.median(np.abs(flat - np.median(flat)))) + assert all(math.isfinite(v) for v in stats.values()) + + def test_empty_array_returns_nan_for_every_key(self): + stats = _compute_intensity_stats(np.array([])) + assert set(stats.keys()) == _STAT_KEYS + assert all(math.isnan(v) for v in stats.values()) + + def test_constant_array_higher_moments_degenerate(self): + stats = _compute_intensity_stats(np.full((4, 4), 7.0)) + assert stats["global_mean"] == 7.0 + assert stats["global_std"] == 0.0 + assert stats["global_mad"] == 0.0 + # Skew/kurtosis of a zero-variance array are NaN — the DB layer records these with + # is_finite=0 rather than rejecting the write. + assert math.isnan(stats["global_skew"]) + assert math.isnan(stats["global_kurtosis"]) diff --git a/tests/unit/test_db.py b/tests/unit/test_db.py new file mode 100644 index 00000000..486d67bb --- /dev/null +++ b/tests/unit/test_db.py @@ -0,0 +1,227 @@ +"""Unit tests for aetherscan.db: writer thread lifecycle, flush sentinel protocol, executemany +batching across tables, is_finite sanitization, and query filters / column whitelists — all +against a tmp-path SQLite file.""" + +from __future__ import annotations + +import json +import time + +import numpy as np +import pytest + +from aetherscan.config import get_config +from aetherscan.db.db import Database + + +@pytest.fixture +def db(tmp_path): + """A started Database against the tmp output path, with the periodic flush effectively + disabled (long write_interval / huge buffer) so only the flush sentinel can drain writes — + making flush() the thing under test rather than a timing accident.""" + config = get_config() + config.db.write_interval = 300.0 + config.db.write_buffer_max_size = 100_000 + database = Database() + assert database.db_path.startswith(str(tmp_path)) + database.start() + yield database + database.stop() + + +class TestWriterThreadLifecycle: + def test_start_spawns_writer_thread(self, db): + assert db.writer_thread is not None + assert db.writer_thread.is_alive() + + def test_start_is_idempotent(self, db): + thread = db.writer_thread + db.start() + assert db.writer_thread is thread # no second thread spawned + + def test_stop_terminates_thread(self, db): + db.stop() + assert not db.writer_thread.is_alive() + + def test_restart_after_stop(self, db): + db.stop() + db.start() + assert db.writer_thread.is_alive() + db.write_training_stat("beta_vae", "total_loss", 1.0, tag="test_v1") + assert db.flush(timeout=10) is True + assert len(db.query_training_stat(tag="test_v1")) == 1 + + def test_singleton_semantics(self, db): + assert Database() is db + + +class TestFlushSentinel: + def test_flush_drains_queued_writes(self, db): + for i in range(5): + db.write_training_stat( + "beta_vae", "total_loss", float(i), round_number=1, epoch_number=i, tag="test_v1" + ) + # write_interval is 300 s, buffer cap 100k: only the sentinel can have flushed these. + assert db.flush(timeout=10) is True + rows = db.query_training_stat(tag="test_v1") + assert len(rows) == 5 + assert sorted(r["value"] for r in rows) == [0.0, 1.0, 2.0, 3.0, 4.0] + + def test_flush_without_writer_is_noop_success(self): + database = Database() # never started + assert database.flush(timeout=1) is True + + def test_flush_during_shutdown_returns_false(self, db): + # Deterministic shutdown-in-progress simulation: setting stop_event on the *live* + # writer races (the thread may exit before flush() checks it, flipping the result + # to True via the not-running branch). Stop the real writer first, then stand in an + # always-alive stub so flush() deterministically hits the stop_event check. + class _AliveStub: + @staticmethod + def is_alive(): + return True + + db.stop() + db.writer_thread = _AliveStub() + db.stop_event.set() # stop() already set it; explicit for the reader + try: + assert db.flush(timeout=1) is False + finally: + db.writer_thread = None # let the fixture's stop() no-op cleanly + + +class TestExecutemanyBatching: + def test_mixed_table_buffer_lands_in_all_tables(self, db): + """One flush must route grouped executemany() inserts to every destination table.""" + tag = "test_v1" + db.write_system_resource("cpu", "system_total", 42.0, unit="percent", tag=tag) + db.write_training_stat("beta_vae", "total_loss", 0.5, 1, 1, tag=tag) + db.write_injection_stat("eti_snr", 12.5, round_number=1, signal_class="true", tag=tag) + db.write_latent_snapshot("beta_vae", 1, 1, 10, 0, "true_only_eti", [[0.1] * 4] * 6, tag=tag) + db.write_inference_result("/tmp/x.npy", 0, 1, 0.99, latent_vector=np.arange(4.0), tag=tag) + assert db.flush(timeout=10) is True + + assert len(db.query_system_resource(tag=tag)) == 1 + assert len(db.query_training_stat(tag=tag)) == 1 + assert len(db.query_injection_stat(tag=tag)) == 1 + snapshots = db.query_latent_snapshots(tag=tag) + assert len(snapshots) == 1 + assert json.loads(snapshots[0]["latent_vector"]) == [[0.1] * 4] * 6 + results = db.query_inference_result(tag=tag) + assert len(results) == 1 + assert json.loads(results[0]["latent_vector"]) == [0.0, 1.0, 2.0, 3.0] + + def test_multiple_rows_per_table_batch(self, db): + for i in range(20): + db.write_training_stat("beta_vae", "kl_loss", float(i), 1, i, tag="test_v2") + assert db.flush(timeout=10) is True + assert len(db.query_training_stat(tag="test_v2", stat_name="kl_loss")) == 20 + + +class TestIsFiniteSanitization: + def test_non_finite_values_stored_as_zero_with_flag(self, db): + tag = "test_v3" + db.write_injection_stat("global_skew", float("nan"), tag=tag) + db.write_injection_stat("global_skew", float("inf"), tag=tag) + db.write_injection_stat("global_skew", float("-inf"), tag=tag) + db.write_injection_stat("global_skew", 1.25, tag=tag) + assert db.flush(timeout=10) is True + + finite_only = db.query_injection_stat(tag=tag) # only_finite defaults to True + assert [r["value"] for r in finite_only] == [1.25] + + everything = db.query_injection_stat(tag=tag, only_finite=False) + assert len(everything) == 4 + sanitized = [r for r in everything if r["is_finite"] == 0] + assert len(sanitized) == 3 + assert all(r["value"] == 0.0 for r in sanitized) + + def test_slope_clamped_flag_round_trip(self, db): + tag = "test_v4" + db.write_injection_stat("rfi_drift_rate", 1.0, slope_clamped=True, tag=tag) + db.write_injection_stat("rfi_drift_rate", 2.0, slope_clamped=False, tag=tag) + db.write_injection_stat("rfi_drift_rate", 3.0, slope_clamped=None, tag=tag) + assert db.flush(timeout=10) is True + + clamped = db.query_injection_stat(tag=tag, only_slope_clamped=True) + assert [r["value"] for r in clamped] == [1.0] + unclamped = db.query_injection_stat(tag=tag, only_slope_clamped=False) + assert sorted(r["value"] for r in unclamped) == [2.0, 3.0] # None defaults to 0 + + def test_stability_aggregation(self, db): + tag = "test_v5" + db.write_injection_stat("global_skew", float("nan"), round_number=1, tag=tag) + db.write_injection_stat("global_skew", 1.0, round_number=1, slope_clamped=True, tag=tag) + db.write_injection_stat("global_skew", 2.0, round_number=2, tag=tag) + assert db.flush(timeout=10) is True + + rows = db.query_injection_stat_stability(stat_name="global_skew", tag=tag) + by_round = {r["round_number"]: r for r in rows} + assert by_round[1]["total_count"] == 2 + assert by_round[1]["non_finite_count"] == 1 + assert by_round[1]["clamped_count"] == 1 + assert by_round[2]["total_count"] == 1 + assert by_round[2]["non_finite_count"] == 0 + + +class TestQueryFiltersAndWhitelists: + @pytest.fixture(autouse=True) + def _rows(self, db): + self.db = db + now = time.time() + for round_number in (1, 2, 3): + for tag in ("test_v1", "test_v2"): + db.write_training_stat( + "beta_vae", + "total_loss", + round_number * 1.0, + round_number, + 1, + tag=tag, + timestamp=now + round_number, + ) + assert db.flush(timeout=10) is True + + def test_scalar_filter_uses_equality(self): + rows = self.db.query_training_stat(tag="test_v1") + assert len(rows) == 3 + assert all(r["tag"] == "test_v1" for r in rows) + + def test_list_filter_uses_in(self): + rows = self.db.query_training_stat(tag=["test_v1", "test_v2"]) + assert len(rows) == 6 + + def test_range_filters_inclusive(self): + rows = self.db.query_training_stat(tag="test_v1", start_round_number=2, end_round_number=3) + assert sorted(r["round_number"] for r in rows) == [2, 3] + + def test_column_projection(self): + rows = self.db.query_training_stat(tag="test_v1", columns=["value", "round_number"]) + assert rows + assert all(set(r.keys()) == {"value", "round_number"} for r in rows) + + def test_invalid_column_rejected(self): + with pytest.raises(ValueError, match="Invalid column"): + self.db.query_training_stat(columns=["value; DROP TABLE training_stats;--"]) + + def test_invalid_column_rejected_per_table(self): + # round_number is valid for training_stats but not system_resources. + with pytest.raises(ValueError, match="Invalid column"): + self.db.query_system_resource(columns=["round_number"]) + + def test_empty_list_filter_matches_everything(self): + # NOTE: this documents CURRENT behavior, not desired behavior — an empty IN-list is + # treated as "no filter" (the `if tag:` gate in every query_* skips falsy values), so + # a caller passing a filtered-to-empty list gets the whole table instead of no rows. + # If query_* semantics are ever fixed to "empty list matches nothing", update this + # test deliberately rather than treating the failure as a regression. + assert len(self.db.query_training_stat(tag=[])) == 6 + + def test_confidence_bounds_on_inference_results(self): + self.db.write_inference_result("/a.npy", 0, 1, 0.99, tag="test_v9") + self.db.write_inference_result("/a.npy", 1, 0, 0.42, tag="test_v9") + assert self.db.flush(timeout=10) is True + rows = self.db.query_inference_result(tag="test_v9", min_confidence=0.9) + assert [r["confidence"] for r in rows] == [0.99] + rows = self.db.query_inference_result(tag="test_v9", prediction=0) + assert [r["confidence"] for r in rows] == [0.42] diff --git a/tests/unit/test_manager.py b/tests/unit/test_manager.py new file mode 100644 index 00000000..b4e5e1c9 --- /dev/null +++ b/tests/unit/test_manager.py @@ -0,0 +1,137 @@ +"""Unit tests for aetherscan.manager: pool / shared-memory tracking and cleanup idempotence. + +ResourceManager registers atexit + signal handlers on construction; the conftest autouse +fixture snapshots/restores signal handlers and unregisters the atexit hook after each test, +so per-test instances don't accumulate process-global state. +""" + +from __future__ import annotations + +import atexit +import sys +from multiprocessing.shared_memory import SharedMemory + +import pytest + +from aetherscan.manager import get_manager, init_manager +from aetherscan.manager.manager import ResourceManager + + +@pytest.fixture +def manager(): + return init_manager() + + +class TestSingletonSemantics: + def test_init_returns_singleton(self, manager): + assert ResourceManager() is manager + assert get_manager() is manager + + def test_reset_produces_fresh_instance(self, manager): + ResourceManager._reset() + try: + assert init_manager() is not manager + finally: + # The conftest teardown only unregisters the *current* instance's atexit hook; + # drop the first instance's hook here so it can't fire at interpreter exit. + atexit.unregister(manager.cleanup_all) + + +class TestPoolTracking: + def test_create_pool_tracks_and_works(self, manager): + pool = manager.create_pool(n_processes=2, name="test-pool") + assert manager.stats.pools_active == 1 + # abs is importable-by-reference in spawned workers — keeps the test light. + assert pool.map(abs, [-2, -3, 4]) == [2, 3, 4] + + def test_close_pool_updates_stats(self, manager): + pool = manager.create_pool(n_processes=2, name="test-pool") + manager.close_pool(pool) + assert manager.stats.pools_active == 0 + assert manager.stats.pools_closed == 1 + assert manager._pools == [] + + def test_close_pool_twice_is_safe(self, manager): + pool = manager.create_pool(n_processes=2, name="test-pool") + manager.close_pool(pool) + manager.close_pool(pool) # logs a warning; must not raise or double-count + assert manager.stats.pools_closed == 1 + + def test_managed_pool_close_is_idempotent(self, manager): + manager.create_pool(n_processes=2, name="test-pool") + managed = manager._pools[0] + managed.close(timeout=10.0) + assert managed.closed is True + managed.close(timeout=10.0) # second call short-circuits on the closed flag + assert managed.closed is True + + def test_multiple_pools_tracked_independently(self, manager): + pool_a = manager.create_pool(n_processes=1, name="a") + pool_b = manager.create_pool(n_processes=1, name="b") + assert manager.stats.pools_active == 2 + manager.close_pool(pool_a) + assert manager.stats.pools_active == 1 + assert manager._pools[0].pool is pool_b + manager.close_pool(pool_b) + assert manager.stats.pools_active == 0 + + +class TestSharedMemoryTracking: + def test_create_tracks_and_is_usable(self, manager): + shm = manager.create_shared_memory(size=1024, name="test-shm") + assert manager.stats.shared_memories_active == 1 + shm.buf[:4] = b"seti" + # Another handle attached by name sees the same bytes (i.e. it's real POSIX shm). + other = SharedMemory(name=shm.name) + try: + assert bytes(other.buf[:4]) == b"seti" + finally: + other.close() + + def test_close_unlinks_and_updates_stats(self, manager): + shm = manager.create_shared_memory(size=1024, name="test-shm") + name = shm.name + manager.close_shared_memory(shm) + assert manager.stats.shared_memories_active == 0 + assert manager.stats.shared_memories_cleaned == 1 + # The segment must be gone from the system namespace. + with pytest.raises(FileNotFoundError): + SharedMemory(name=name) + + def test_close_twice_is_safe(self, manager): + shm = manager.create_shared_memory(size=1024, name="test-shm") + manager.close_shared_memory(shm) + manager.close_shared_memory(shm) # logs a warning; must not raise or double-count + assert manager.stats.shared_memories_cleaned == 1 + + +@pytest.mark.skipif( + sys.platform != "linux", + reason="cleanup_all reports process-tree memory via PSS, which psutil only exposes on Linux", +) +class TestCleanupAll: + def test_cleanup_all_closes_everything(self, manager): + manager.create_pool(n_processes=1, name="pool") + shm = manager.create_shared_memory(size=512, name="shm") + name = shm.name + + manager.cleanup_all() + + assert manager.stats.pools_active == 0 + assert manager.stats.pools_closed == 1 + assert manager.stats.shared_memories_active == 0 + assert manager.stats.shared_memories_cleaned == 1 + with pytest.raises(FileNotFoundError): + SharedMemory(name=name) + + def test_cleanup_all_is_idempotent(self, manager): + manager.create_pool(n_processes=1, name="pool") + manager.cleanup_all() + closed = manager.stats.pools_closed + manager.cleanup_all() # second call must be a guarded no-op + assert manager.stats.pools_closed == closed + + def test_cleanup_time_recorded(self, manager): + manager.create_shared_memory(size=512, name="shm") + manager.cleanup_all() + assert manager.stats.cleanup_time_seconds >= 0.0 diff --git a/tests/unit/test_models.py b/tests/unit/test_models.py new file mode 100644 index 00000000..24f3a33a --- /dev/null +++ b/tests/unit/test_models.py @@ -0,0 +1,179 @@ +"""Unit tests for aetherscan.models: latent feature layout, RandomForestModel behavior, the +Sampling layer, and encoder/decoder symmetry.""" + +from __future__ import annotations + +import numpy as np +import pytest + +from aetherscan.config import get_config +from aetherscan.models import RandomForestModel, Sampling, prepare_latent_features + + +class TestPrepareLatentFeatures: + def test_row_major_layout(self): + """features[i] must be the row-major concatenation of cadence i's 6 latents.""" + n_cadences, num_obs, latent_dim = 3, 6, 4 + # Encode (cadence, obs, dim) into each value so misordering is unambiguous. + latents = np.zeros((n_cadences * num_obs, latent_dim)) + for cad in range(n_cadences): + for obs in range(num_obs): + for dim in range(latent_dim): + latents[cad * num_obs + obs, dim] = cad * 1000 + obs * 10 + dim + + features = prepare_latent_features(latents, num_observations=num_obs) + + assert features.shape == (n_cadences, num_obs * latent_dim) + for cad in range(n_cadences): + expected = np.concatenate([latents[cad * num_obs + obs] for obs in range(num_obs)]) + np.testing.assert_array_equal(features[cad], expected) + # Spot-check the row-major invariant: feature column obs*latent_dim + dim. + assert features[cad, 1 * latent_dim + 2] == cad * 1000 + 10 + 2 + + def test_indivisible_row_count_raises(self): + with pytest.raises(ValueError, match="Not divisible"): + prepare_latent_features(np.zeros((7, 8)), num_observations=6) + + def test_single_cadence(self): + latents = np.arange(6 * 8, dtype=float).reshape(6, 8) + features = prepare_latent_features(latents, num_observations=6) + np.testing.assert_array_equal(features, latents.ravel()[None, :]) + + +def _toy_latents(n_per_class=16, latent_dim=8, num_obs=6, seed=0): + """Separable toy data: class-1 latents cluster at +2, class-0 at -2.""" + rng = np.random.default_rng(seed) + neg = rng.normal(-2.0, 0.3, size=(n_per_class * num_obs, latent_dim)) + pos = rng.normal(+2.0, 0.3, size=(n_per_class * num_obs, latent_dim)) + latents = np.concatenate([neg, pos]) + labels = np.concatenate([np.zeros(n_per_class), np.ones(n_per_class)]).astype(int) + return latents, labels + + +class TestRandomForestModel: + @pytest.fixture(autouse=True) + def _small_forest(self): + # 1000 trees is overkill for toy data; shrink for speed. + get_config().rf.n_estimators = 10 + + def test_train_and_predict_separable_data(self): + latents, labels = _toy_latents() + model = RandomForestModel() + assert model.is_trained is False + model.train(latents, labels) + assert model.is_trained is True + + predictions = model.predict(latents) + np.testing.assert_array_equal(predictions, labels) + + def test_predict_proba_shape_and_normalization(self): + latents, labels = _toy_latents() + model = RandomForestModel() + model.train(latents, labels) + probas = model.predict_proba(latents) + assert probas.shape == (len(labels), 2) + np.testing.assert_allclose(probas.sum(axis=1), 1.0) + + def test_predict_verbose_confidence_of_predicted_class(self): + latents, labels = _toy_latents() + model = RandomForestModel() + model.train(latents, labels) + predictions, confidences = model.predict_verbose(latents) + probas = model.predict_proba(latents) + np.testing.assert_array_equal(predictions, labels) + # Confidence is P(predicted class) — high for confident negatives too. + expected = np.where(predictions, probas[:, 1], probas[:, 0]) + np.testing.assert_array_equal(confidences, expected) + assert confidences.min() >= 0.5 + + def test_threshold_semantics(self): + latents, labels = _toy_latents() + model = RandomForestModel() + model.train(latents, labels) + # Strict inequality: P(1) > 1.0 is impossible, so everything is class 0. + assert model.predict(latents, threshold=1.0).sum() == 0 + # Threshold 0 flags anything with nonzero positive probability. + assert model.predict(latents, threshold=0.0).sum() >= labels.sum() + + def test_feature_label_mismatch_raises(self): + latents, _ = _toy_latents(n_per_class=4) + wrong_labels = np.zeros(3, dtype=int) + model = RandomForestModel() + with pytest.raises(ValueError, match="mismatch"): + model.train(latents, wrong_labels) + + def test_save_load_round_trip(self, tmp_path): + latents, labels = _toy_latents() + model = RandomForestModel() + model.train(latents, labels) + path = str(tmp_path / "rf.joblib") + model.save(path) + + restored = RandomForestModel() + restored.load(path) + assert restored.is_trained is True + np.testing.assert_array_equal(restored.predict_proba(latents), model.predict_proba(latents)) + + +class TestSamplingLayer: + def test_output_shape(self): + import tensorflow as tf # noqa: PLC0415 + + z_mean = tf.zeros((5, 8)) + z_log_var = tf.zeros((5, 8)) + z = Sampling()([z_mean, z_log_var]) + assert z.shape == (5, 8) + + def test_collapses_to_mean_at_negligible_variance(self): + import tensorflow as tf # noqa: PLC0415 + + z_mean = tf.constant(np.arange(10, dtype=np.float32).reshape(2, 5)) + z_log_var = tf.fill((2, 5), -100.0) # std = exp(-50) ~ 0 + z = Sampling()([z_mean, z_log_var]) + np.testing.assert_allclose(z.numpy(), z_mean.numpy(), atol=1e-6) + + def test_variance_injects_noise(self): + import tensorflow as tf # noqa: PLC0415 + + z_mean = tf.zeros((4, 8)) + z_log_var = tf.zeros((4, 8)) # std = 1 + z1 = Sampling()([z_mean, z_log_var]).numpy() + z2 = Sampling()([z_mean, z_log_var]).numpy() + assert not np.allclose(z1, z2) + assert np.std(z1) > 0.1 + + +@pytest.mark.slow +class TestEncoderDecoderSymmetry: + """Builds the full Beta-VAE graph on CPU — slow but CI-safe.""" + + def test_symmetry_and_forward_pass(self): + import tensorflow as tf # noqa: PLC0415 + + from aetherscan.models import create_beta_vae_model # noqa: PLC0415 + + vae = create_beta_vae_model() + config = get_config() + latent_dim = config.beta_vae.latent_dim + dense_size = config.beta_vae.dense_layer_size + + # Encoder consumes single observations; decoder must emit the exact mirror shape. + assert vae.encoder.input_shape == (None, 16, dense_size, 1) + z_mean_shape, z_log_var_shape, z_shape = vae.encoder.output_shape + assert z_mean_shape == (None, latent_dim) + assert z_log_var_shape == (None, latent_dim) + assert z_shape == (None, latent_dim) + assert vae.decoder.input_shape == (None, latent_dim) + assert vae.decoder.output_shape == vae.encoder.input_shape + + # Forward pass round-trips the cadence shape (same graph — built once per test run). + batch = tf.random.uniform((2, 6, 16, dense_size)) + reconstruction, z_mean, z_log_var, z = vae(batch, training=False) + assert reconstruction.shape == (2, 6, 16, dense_size) + # The encoder operates per observation: 2 cadences * 6 observations. + assert z_mean.shape == (12, latent_dim) + assert z.shape == (12, latent_dim) + # Sigmoid output stays within [0, 1] like the log-normed inputs. + recon = reconstruction.numpy() + assert recon.min() >= 0.0 + assert recon.max() <= 1.0 diff --git a/tests/unit/test_preprocessing.py b/tests/unit/test_preprocessing.py new file mode 100644 index 00000000..d6acd2e5 --- /dev/null +++ b/tests/unit/test_preprocessing.py @@ -0,0 +1,235 @@ +"""Unit tests for aetherscan.preprocessing: hit deduplication, CSV cadence grouping, filename +sanitization, JSON coercion, DC-spike removal, and spline bandpass fitting.""" + +from __future__ import annotations + +import numpy as np +import pytest + +from aetherscan.config import get_config +from aetherscan.preprocessing import ( + DataPreprocessor, + _fit_channel_bandpass, + _read_coarse_channel_worker, + _remove_dc_spike, + group_observations_from_csv, +) + + +@pytest.fixture +def group_cols(): + """The real grouping columns from InferenceConfig — read from the config rather than + duplicated, so these tests track the defaults the pipeline actually uses.""" + return list(get_config().inference.cadence_group_by_cols) + + +@pytest.fixture +def h5_col(): + return get_config().inference.cadence_h5_path_col + + +class TestDeduplicateHits: + def test_empty_input(self): + assert DataPreprocessor._deduplicate_hits([], stamp_width=512) == [] + + def test_far_apart_hits_survive(self): + hits = [(0, 5.0, 0.01), (300, 4.0, 0.01)] + assert DataPreprocessor._deduplicate_hits(hits, stamp_width=512) == hits + + def test_nearby_hits_merge_keeping_higher_statistic(self): + hits = [(0, 5.0, 0.01), (100, 9.0, 0.02)] + assert DataPreprocessor._deduplicate_hits(hits, stamp_width=512) == [(100, 9.0, 0.02)] + + def test_merge_is_greedy_left_to_right(self): + # 0 and 100 merge (keep 100 @ 9.0); 300 is then within 256 of 100 but weaker, + # so it merges away too — one survivor. + hits = [(0, 5.0, 0.01), (100, 9.0, 0.02), (300, 2.0, 0.03)] + assert DataPreprocessor._deduplicate_hits(hits, stamp_width=512) == [(100, 9.0, 0.02)] + + def test_input_order_does_not_matter(self): + hits = [(300, 2.0, 0.03), (0, 5.0, 0.01), (100, 9.0, 0.02)] + assert DataPreprocessor._deduplicate_hits(hits, stamp_width=512) == [(100, 9.0, 0.02)] + + def test_boundary_distance_exactly_half_width_kept(self): + # Merge condition is strict (<): a gap of exactly stamp_width // 2 keeps both. + hits = [(0, 5.0, 0.01), (256, 4.0, 0.02)] + assert DataPreprocessor._deduplicate_hits(hits, stamp_width=512) == hits + + +class TestGroupObservationsFromCsv: + def test_valid_and_flagged_groups(self, make_inference_csv, group_cols, h5_col): + key_a = { + "Target": "HIP110750", + "Session": "AGBT21B_999_31", + "Band": "L", + "Cadence ID": "0", + "Frequency": "1400", + } + key_b = {**key_a, "Cadence ID": "1"} + csv_path = make_inference_csv( + "two_groups.csv", + groups=[ + (key_a, [f"/data/a_{i}.h5" for i in range(6)]), + (key_b, [f"/data/b_{i}.h5" for i in range(3)]), # short cadence -> flagged + ], + ) + valid, flagged = group_observations_from_csv( + str(csv_path), group_cols, h5_col, expected_obs=6 + ) + assert len(valid) == 1 + assert len(flagged) == 1 + assert valid[0].is_valid is True + assert valid[0].key == tuple(key_a[c] for c in group_cols) + # Row order within the group is preserved. + assert valid[0].h5_paths == [f"/data/a_{i}.h5" for i in range(6)] + assert flagged[0].is_valid is False + assert len(flagged[0].h5_paths) == 3 + + def test_missing_column_raises_keyerror(self, make_inference_csv, group_cols, h5_col): + csv_path = make_inference_csv("ok.csv") + with pytest.raises(KeyError, match="missing required column"): + group_observations_from_csv( + str(csv_path), [*group_cols, "Nonexistent"], h5_col, expected_obs=6 + ) + + def test_missing_file_raises(self, group_cols, h5_col): + with pytest.raises(FileNotFoundError): + group_observations_from_csv("/nope.csv", group_cols, h5_col) + + def test_expected_obs_parameter(self, make_inference_csv, group_cols, h5_col): + key = { + "Target": "T", + "Session": "S", + "Band": "L", + "Cadence ID": "0", + "Frequency": "1", + } + csv_path = make_inference_csv("three.csv", groups=[(key, ["/a.h5", "/b.h5", "/c.h5"])]) + valid, flagged = group_observations_from_csv( + str(csv_path), group_cols, h5_col, expected_obs=3 + ) + assert len(valid) == 1 + assert flagged == [] + + +class TestCadenceNpyFilename: + def test_clean_key_passes_through(self): + name = DataPreprocessor._cadence_npy_filename("catalog", ("HIP110750", "L", "0")) + assert name == "catalog_HIP110750_L_0.npy" + + def test_hostile_characters_collapse_to_underscore(self): + name = DataPreprocessor._cadence_npy_filename( + "catalog", ("HIP 110750", "AGBT/18A", "a;b|c", "$(rm -rf)") + ) + assert "/" not in name + assert " " not in name + assert ";" not in name + assert "$" not in name + assert name.startswith("catalog_") + assert name.endswith(".npy") + + def test_dash_and_dot_survive(self): + name = DataPreprocessor._cadence_npy_filename("catalog", ("L-band", "1.5")) + assert name == "catalog_L-band_1.5.npy" + + def test_non_string_key_components(self): + name = DataPreprocessor._cadence_npy_filename("catalog", (0, 1400.5)) + assert name == "catalog_0_1400.5.npy" + + +class TestToJsonSafe: + def test_bytes_decode(self): + assert DataPreprocessor._to_json_safe(b"HIP110750") == "HIP110750" + + def test_numpy_scalars(self): + assert DataPreprocessor._to_json_safe(np.float32(1.5)) == 1.5 + assert DataPreprocessor._to_json_safe(np.int64(7)) == 7 + assert DataPreprocessor._to_json_safe(np.bool_(True)) is True + + def test_ndarray_to_list(self): + result = DataPreprocessor._to_json_safe(np.arange(3, dtype=np.int32)) + assert result == [0, 1, 2] + assert all(isinstance(v, int) for v in result) + + def test_nested_structures_and_keys(self): + obj = { + 1: {"vals": np.array([1.0, 2.0], dtype=np.float64)}, + "meta": (b"x", np.int16(2)), + } + result = DataPreprocessor._to_json_safe(obj) + assert result == {"1": {"vals": [1.0, 2.0]}, "meta": ["x", 2]} + + def test_plain_values_pass_through(self): + assert DataPreprocessor._to_json_safe("s") == "s" + assert DataPreprocessor._to_json_safe(3.5) == 3.5 + assert DataPreprocessor._to_json_safe(None) is None + + +class TestRemoveDcSpike: + def test_spike_replaced_by_neighbor_interpolation(self): + rng = np.random.default_rng(9) + coarse_width, n_coarse = 32, 3 + block = rng.chisquare(df=4, size=(4, n_coarse * coarse_width)) + original = block.copy() + + _remove_dc_spike(block, coarse_width, n_coarse) + + half = coarse_width // 2 + for i in range(n_coarse): + dc = i * coarse_width + half + np.testing.assert_allclose( + block[:, dc], (original[:, dc + 1] + original[:, dc - 3]) / 2 + ) + np.testing.assert_allclose( + block[:, dc - 1], (original[:, dc + 2] + original[:, dc - 2]) / 2 + ) + # Everything else in the channel is untouched. + untouched = [ + j for j in range(i * coarse_width, (i + 1) * coarse_width) if j not in (dc, dc - 1) + ] + np.testing.assert_array_equal(block[:, untouched], original[:, untouched]) + + def test_visible_spike_is_flattened(self): + coarse_width = 64 + block = np.ones((4, coarse_width)) + dc = coarse_width // 2 + block[:, dc] = 100.0 + block[:, dc - 1] = 100.0 + _remove_dc_spike(block, coarse_width, 1) + np.testing.assert_allclose(block, np.ones((4, coarse_width))) + + +class TestFitChannelBandpass: + def test_smooth_bandpass_recovered(self): + channel_width, spl_order = 1024, 16 + x = np.arange(channel_width) + bandpass = 100.0 + 10.0 * np.sin(2 * np.pi * x / channel_width) + fit = _fit_channel_bandpass(bandpass, channel_width, spl_order) + assert fit.shape == bandpass.shape + # A smooth curve must be recovered nearly exactly by the spline. + assert np.max(np.abs(fit - bandpass)) < 0.1 + + def test_subtracting_fit_flattens_channel(self): + rng = np.random.default_rng(2) + channel_width, spl_order = 1024, 16 + x = np.arange(channel_width) + bandpass = 100.0 + 20.0 * np.cos(2 * np.pi * x / channel_width) + noisy = bandpass + rng.normal(0, 0.5, size=channel_width) + residual = noisy - _fit_channel_bandpass(noisy, channel_width, spl_order) + # Residuals lose the 20-unit bandpass structure and keep only ~noise-scale variance. + assert np.abs(residual.mean()) < 0.5 + 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) + + import h5py # noqa: PLC0415 + + 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) diff --git a/tests/unit/test_train_utils.py b/tests/unit/test_train_utils.py new file mode 100644 index 00000000..1150e982 --- /dev/null +++ b/tests/unit/test_train_utils.py @@ -0,0 +1,320 @@ +"""Unit tests for aetherscan.train pure-logic helpers: checkpoint tag resolution, curriculum +schedules, directory archiving, encoder-trained heuristics, and SHAP output normalization.""" + +from __future__ import annotations + +import os + +import numpy as np +import pytest + +from aetherscan.config import get_config +from aetherscan.train import ( + TrainingPipeline, + _select_positive_class_shap, + archive_directory, + check_encoder_trained, + compute_expected_std, + get_latest_tag, +) + + +def _touch_pair(checkpoints_dir, tag): + """Create a matching encoder/decoder checkpoint pair for `tag`.""" + os.makedirs(checkpoints_dir, exist_ok=True) + for prefix in ("vae_encoder", "vae_decoder"): + with open(os.path.join(checkpoints_dir, f"{prefix}_{tag}.keras"), "w") as f: + f.write("stub") + + +class TestGetLatestTag: + def test_priority_ladder(self, tmp_path): + d = str(tmp_path / "ckpt") + for tag in ("test_v3", "20240101_000000", "round_02", "final_v1"): + _touch_pair(d, tag) + assert get_latest_tag(d) == "final_v1" + + def test_round_beats_timestamp_and_test(self, tmp_path): + d = str(tmp_path / "ckpt") + for tag in ("test_v9", "20991231_235959", "round_02", "round_10"): + _touch_pair(d, tag) + # Numeric compare, not lexicographic: round_10 > round_02. + assert get_latest_tag(d) == "round_10" + + def test_timestamp_beats_test(self, tmp_path): + d = str(tmp_path / "ckpt") + for tag in ("test_v9", "20240101_000000", "20250101_000000"): + _touch_pair(d, tag) + assert get_latest_tag(d) == "20250101_000000" + + def test_test_tags_ranked_by_version(self, tmp_path): + d = str(tmp_path / "ckpt") + for tag in ("test_v2", "test_v17", "test_v9"): + _touch_pair(d, tag) + assert get_latest_tag(d) == "test_v17" + + def test_final_ranked_by_version(self, tmp_path): + d = str(tmp_path / "ckpt") + for tag in ("final_v1", "final_v12", "final_v3"): + _touch_pair(d, tag) + assert get_latest_tag(d) == "final_v12" + + def test_encoder_without_decoder_ignored(self, tmp_path): + d = str(tmp_path / "ckpt") + _touch_pair(d, "round_01") + # Higher-priority final_v2 lacks its decoder — must not win. + with open(os.path.join(d, "vae_encoder_final_v2.keras"), "w") as f: + f.write("stub") + assert get_latest_tag(d) == "round_01" + + def test_missing_directory_raises(self, tmp_path): + with pytest.raises(FileNotFoundError, match="doesn't exist"): + get_latest_tag(str(tmp_path / "nope")) + + def test_empty_directory_raises(self, tmp_path): + d = tmp_path / "empty" + d.mkdir() + with pytest.raises(FileNotFoundError, match="No encoder files"): + get_latest_tag(str(d)) + + def test_no_complete_pair_raises(self, tmp_path): + d = tmp_path / "orphans" + d.mkdir() + with open(d / "vae_encoder_round_01.keras", "w") as f: + f.write("stub") + with pytest.raises(FileNotFoundError, match="No valid model pairs"): + get_latest_tag(str(d)) + + +class _PipelineStub: + """Just enough of TrainingPipeline to drive _calculate_curriculum_snr.""" + + def __init__(self, config): + self.config = config + + +def _curriculum(round_idx): + return TrainingPipeline._calculate_curriculum_snr(_PipelineStub(get_config()), round_idx) + + +class TestCalculateCurriculumSnr: + @pytest.fixture(autouse=True) + def _setup_config(self): + config = get_config() + config.training.num_training_rounds = 5 + config.training.snr_base = 10 + config.training.initial_snr_range = 40 + config.training.final_snr_range = 10 + + def test_linear_endpoints_and_monotonicity(self): + get_config().training.curriculum_schedule = "linear" + ranges = [_curriculum(i)[1] for i in range(5)] + assert ranges[0] == 40 + assert ranges[-1] == 10 + assert all(a >= b for a, b in zip(ranges, ranges[1:], strict=False)) + assert all(base == 10 for base, _ in (_curriculum(i) for i in range(5))) + + def test_linear_midpoint(self): + get_config().training.curriculum_schedule = "linear" + # progress = 2/4 = 0.5 -> 40 - 0.5 * 30 = 25 + assert _curriculum(2)[1] == 25 + + def test_exponential_endpoints_exact(self): + config = get_config() + config.training.curriculum_schedule = "exponential" + config.training.exponential_decay_rate = -3.0 + assert _curriculum(0)[1] == 40 + assert _curriculum(4)[1] == 10 + + def test_exponential_decays_faster_than_linear(self): + config = get_config() + config.training.curriculum_schedule = "exponential" + config.training.exponential_decay_rate = -3.0 + ranges = [_curriculum(i)[1] for i in range(5)] + assert all(a >= b for a, b in zip(ranges, ranges[1:], strict=False)) + # Exponential front-loads the difficulty ramp: below linear at the midpoint. + assert ranges[2] < 25 + + def test_exponential_rejects_nonnegative_decay_rate(self): + config = get_config() + config.training.curriculum_schedule = "exponential" + config.training.exponential_decay_rate = 0.5 + with pytest.raises(ValueError, match="must be < 0"): + _curriculum(1) + + def test_step_schedule(self): + config = get_config() + config.training.curriculum_schedule = "step" + config.training.step_easy_rounds = 2 + config.training.step_hard_rounds = 3 + assert [_curriculum(i)[1] for i in range(5)] == [40, 40, 10, 10, 10] + + def test_step_schedule_rejects_bad_sum(self): + config = get_config() + config.training.curriculum_schedule = "step" + config.training.step_easy_rounds = 2 + config.training.step_hard_rounds = 2 # 2 + 2 != 5 + with pytest.raises(ValueError, match="must equal total_rounds"): + _curriculum(0) + + def test_single_round_returns_initial_range(self): + config = get_config() + config.training.num_training_rounds = 1 + config.training.curriculum_schedule = "linear" + assert _curriculum(0) == (10, 40) + + def test_unknown_schedule_raises(self): + get_config().training.curriculum_schedule = "sigmoid" + with pytest.raises(ValueError, match="invalid"): + _curriculum(0) + + +class TestArchiveDirectory: + def test_empty_directory_is_noop(self, tmp_path): + base = tmp_path / "plots" + archive_directory(str(base)) + assert base.exists() + assert list(base.iterdir()) == [] + + def test_fresh_run_moves_files_to_archive(self, tmp_path): + base = tmp_path / "plots" + base.mkdir() + (base / "a.png").write_text("a") + (base / "b.png").write_text("b") + (base / "subdir").mkdir() # not in target_dirs -> untouched + + archive_directory(str(base), round_num=1) + + remaining = {p.name for p in base.iterdir()} + assert remaining == {"archive", "subdir"} + archived = list((base / "archive").iterdir()) + assert len(archived) == 1 # one timestamped snapshot + assert {p.name for p in archived[0].iterdir()} == {"a.png", "b.png"} + + def test_resume_copies_then_deletes_rounds_geq(self, tmp_path): + base = tmp_path / "checkpoints" + base.mkdir() + for tag in ("round_01", "round_02", "round_03"): + (base / f"vae_encoder_{tag}.keras").write_text("stub") + (base / "notes.txt").write_text("keep me") + + archive_directory(str(base), round_num=2) + + remaining = {p.name for p in base.iterdir()} + # Rounds >= 2 deleted; round_01 and non-round files kept. + assert remaining == {"archive", "vae_encoder_round_01.keras", "notes.txt"} + # Everything (including the later-deleted files) was backed up first. + snapshot = next((base / "archive").iterdir()) + assert {p.name for p in snapshot.iterdir()} == { + "vae_encoder_round_01.keras", + "vae_encoder_round_02.keras", + "vae_encoder_round_03.keras", + "notes.txt", + } + + def test_target_dirs_moved_and_recreated_empty(self, tmp_path): + base = tmp_path / "tb" + (base / "train").mkdir(parents=True) + (base / "train" / "events.1").write_text("x") + (base / "validation").mkdir() + + archive_directory(str(base), target_dirs=["train"], round_num=1) + + assert (base / "train").exists() + assert list((base / "train").iterdir()) == [] # replaced with an empty dir + assert (base / "validation").exists() # not a target -> untouched + snapshot = next((base / "archive").iterdir()) + assert (snapshot / "train" / "events.1").exists() + + +class TestEncoderTrainedHeuristics: + @staticmethod + def _dense_model(initializer, units=256, input_dim=256): + import tensorflow as tf # noqa: PLC0415 + + model = tf.keras.Sequential( + [ + tf.keras.layers.Input(shape=(input_dim,)), + tf.keras.layers.Dense(units, kernel_initializer=initializer), + ] + ) + return model + + def test_expected_std_he_normal(self): + from tensorflow.keras.initializers import HeNormal # noqa: PLC0415 + + model = self._dense_model(HeNormal(seed=0), units=4, input_dim=8) + expected = compute_expected_std(model.layers[-1]) + assert expected == pytest.approx(np.sqrt(2.0 / 8)) + + def test_expected_std_glorot_normal(self): + from tensorflow.keras.initializers import GlorotNormal # noqa: PLC0415 + + model = self._dense_model(GlorotNormal(seed=0), units=4, input_dim=8) + expected = compute_expected_std(model.layers[-1]) + assert expected == pytest.approx(np.sqrt(2.0 / (8 + 4))) + + def test_expected_std_unknown_initializer_returns_none(self): + model = self._dense_model("glorot_uniform", units=4, input_dim=8) + assert compute_expected_std(model.layers[-1]) is None + + def test_fresh_encoder_reports_untrained(self): + from tensorflow.keras.initializers import HeNormal # noqa: PLC0415 + + # 256x256 kernel: sampling noise on the std is ~0.3%, far under the 20% threshold. + model = self._dense_model(HeNormal(seed=0)) + assert check_encoder_trained(model) is False + + def test_scaled_weights_report_trained(self): + from tensorflow.keras.initializers import HeNormal # noqa: PLC0415 + + model = self._dense_model(HeNormal(seed=0)) + layer = model.layers[-1] + kernel, bias = layer.get_weights() + layer.set_weights([kernel * 3.0, bias]) # 200% deviation from expected std + assert check_encoder_trained(model) is True + + +class TestSelectPositiveClassShap: + N, F = 5, 8 + + def test_list_of_class_arrays_selects_positive(self): + neg = np.zeros((self.N, self.F)) + pos = np.ones((self.N, self.F)) + result = _select_positive_class_shap([neg, pos]) + np.testing.assert_array_equal(result, pos) + + def test_trailing_class_axis_values(self): + values = np.stack( + [np.zeros((self.N, self.F)), np.ones((self.N, self.F))], axis=-1 + ) # (N, F, 2) + result = _select_positive_class_shap(values) + assert result.shape == (self.N, self.F) + assert np.all(result == 1.0) + + def test_trailing_class_axis_interactions(self): + values = np.stack( + [np.zeros((self.N, self.F, self.F)), np.ones((self.N, self.F, self.F))], axis=-1 + ) # (N, F, F, 2) + result = _select_positive_class_shap(values) + assert result.shape == (self.N, self.F, self.F) + assert np.all(result == 1.0) + + def test_single_output_passthrough(self): + values = np.arange(self.N * self.F, dtype=float).reshape(self.N, self.F) + np.testing.assert_array_equal(_select_positive_class_shap(values), values) + + def test_log_loss_list_selects_first(self): + first = np.ones((self.N, self.F)) + result = _select_positive_class_shap([first, np.zeros((self.N, self.F))], log_loss=True) + np.testing.assert_array_equal(result, first) + + def test_log_loss_trailing_class_axis(self): + values = np.stack([np.zeros((self.N, self.F)), np.ones((self.N, self.F))], axis=-1) + result = _select_positive_class_shap(values, log_loss=True) + assert result.shape == (self.N, self.F) + assert np.all(result == 1.0) + + def test_log_loss_passthrough(self): + values = np.arange(self.N * self.F, dtype=float).reshape(self.N, self.F) + np.testing.assert_array_equal(_select_positive_class_shap(values, log_loss=True), values)