alnn_rewrite/analysis/__init__.py

"""Analysis layer (pipeline step 7).

Reads ONLY the evaluation netCDF files in ``<work_dir>/evaluations/`` and writes
artifacts to ``<work_dir>/analysis/``. Because analyses are pure functions of the
schema-locked datasets, a rerun is cheap and can only overwrite regeneratable
output -- never models or evaluations.

See ``ai/ANALYSIS_PLAN.md`` for the scientific definitions behind each analysis.

Adding an analysis:
    1. create ``analysis/plots/<name>.py`` with a ``@register``-decorated function
       taking an :class:`~analysis.context.AnalysisContext`;
    2. import it in ``analysis/plots/__init__.py``.
"""

import pathlib as pl
from typing import Any, Dict, Optional

import xarray as xr

from evaluation import load_evaluation
from util.ui_logger import PipelineLogger

EVAL_SUBDIR = "evaluations"
ANALYSIS_SUBDIR = "analysis"

__all__ = [
    "ANALYSIS_SUBDIR",
    "EVAL_SUBDIR",
    "get_analyses",
    "load_evaluations",
    "run_all",
]


def get_analyses() -> Dict[str, Any]:
    """Every registered analysis, keyed by name.

    Importing ``analysis.plots`` is what runs the ``@register`` decorators, and
    it is done lazily (here and in :func:`run_all`) so that submodules can import
    from this package without forming an import cycle.
    """
    from analysis import plots as _plots  # noqa: F401  (import for registration)
    from analysis.registry import ANALYSES

    return dict(ANALYSES)


def load_evaluations(
    eval_dir: pl.Path, log: PipelineLogger
) -> Dict[str, xr.Dataset]:
    """Load every ``*.nc`` in ``eval_dir``, keyed by file stem."""
    datasets: Dict[str, xr.Dataset] = {}
    for nc in sorted(eval_dir.glob("*.nc")):
        try:
            ds = load_evaluation(str(nc))
        except Exception as exc:  # noqa: BLE001 - report and skip a bad file
            log.error(f"Could not load evaluation '{nc.name}': {exc}")
            continue
        datasets[nc.stem] = ds
        log.info(f"Loaded evaluation '{nc.stem}' with dims {dict(ds.sizes)}.")
    return datasets


def run_all(
    work_dir: pl.Path,
    log: PipelineLogger,
    config: Optional[Dict[str, Any]] = None,
    progress: Optional[Any] = None,
) -> None:
    """Run every registered analysis over ``work_dir``'s evaluations.

    Analyses are independent: one failing is reported and does not prevent the
    rest from running. Analyses declaring ``requires_noise`` are skipped with an
    explanatory message when no noise sweep is present.
    """
    # Lazy imports: keeps this module free of submodule dependencies at import
    # time, so `analysis.context` etc. can import from the package cleanly.
    from analysis import plots as _plots  # noqa: F401  (import for registration)
    from analysis.context import AnalysisContext
    from analysis.registry import ANALYSES

    eval_dir = work_dir / EVAL_SUBDIR
    out_dir = work_dir / ANALYSIS_SUBDIR

    if not eval_dir.is_dir():
        log.error(
            f"No evaluations directory at {eval_dir}. Run an evaluation scenario "
            "first (or point the working directory at one that has evaluations)."
        )
        return

    datasets = load_evaluations(eval_dir, log)
    if not datasets:
        log.error(f"No evaluation .nc files found in {eval_dir}; nothing to analyze.")
        return

    out_dir.mkdir(parents=True, exist_ok=True)
    ctx = AnalysisContext(datasets, out_dir, log, config)

    try:
        if not ANALYSES:
            log.error("No analyses are registered; nothing to do.")
            return

        if progress is not None:
            progress.update(total=len(ANALYSES), advance=0)

        for spec in ANALYSES.values():
            if spec.requires_noise and not ctx.has_noise:
                log.info(
                    f"Skipping '{spec.name}': needs a noise sweep, but no "
                    "multi-noise-level evaluation was found."
                )
                if progress is not None:
                    progress.update(advance=1)
                continue

            marker = f" [{', '.join(spec.tags)}]" if spec.tags else ""
            log.info(f"Running analysis: {spec.name}{marker} — {spec.title}")
            try:
                spec.fn(ctx)
            except NotImplementedError:
                log.info(f"'{spec.name}' is scaffolded but not implemented yet.")
            except Exception as exc:  # noqa: BLE001 - keep other analyses running
                log.error(f"Analysis '{spec.name}' failed: {exc}")
                log.file_logger.error(f"Analysis '{spec.name}' failed", exc_info=True)
            if progress is not None:
                progress.update(advance=1)

        log.info(f"Analysis artifacts written to {out_dir}/.")
    finally:
        for ds in datasets.values():
            ds.close()