alnn_rewrite/analysis/registry.py

"""Analysis registry.

An *analysis* is a pure function of the loaded evaluation datasets that writes
artifacts into ``<work_dir>/analysis/``. Analyses declare their data
requirements so the runner can skip them with a clear message instead of failing
when a working directory lacks, say, a noise sweep.

Adding an analysis is: one new module under ``analysis/plots/`` containing a
``@register``-decorated function, plus an import in ``analysis/plots/__init__``.
"""

from dataclasses import dataclass, field
from typing import Callable, Dict, List, TYPE_CHECKING

if TYPE_CHECKING:  # avoid a circular import at runtime
    from analysis.context import AnalysisContext

AnalysisFn = Callable[["AnalysisContext"], None]


@dataclass(frozen=True)
class AnalysisSpec:
    """Metadata describing one registered analysis."""

    name: str
    title: str
    fn: AnalysisFn
    #: Requires an evaluation with more than one noise level (a noise sweep).
    requires_noise: bool = False
    #: Free-form notes shown in logs (e.g. "scaffold").
    tags: List[str] = field(default_factory=list)


#: Registered analyses, in registration order (which is the run order).
ANALYSES: Dict[str, AnalysisSpec] = {}


def register(
    name: str,
    *,
    title: str,
    requires_noise: bool = False,
    tags: List[str] | None = None,
) -> Callable[[AnalysisFn], AnalysisFn]:
    """Decorator registering an analysis function under ``name``.

    Args:
        name: Unique id; also the log label and typical output filename stem.
        title: Human-readable description.
        requires_noise: Skip this analysis when no noise-sweep file is available.
        tags: Optional markers (e.g. ``["scaffold"]``).
    """

    def _decorate(fn: AnalysisFn) -> AnalysisFn:
        if name in ANALYSES:
            raise ValueError(f"Analysis {name!r} is already registered.")
        ANALYSES[name] = AnalysisSpec(
            name=name,
            title=title,
            fn=fn,
            requires_noise=requires_noise,
            tags=list(tags or []),
        )
        return fn

    return _decorate