Source code for uvex_transients.cli.config

"""
Parsing/validation for a CLI run-config YAML file.

A single YAML file drives every CLI command (see `uvex_transients.cli.main`); each
command only needs the section(s) relevant to it (`generate:` for ``generate``,
`cuts:` for ``cut``, ...), so `RunConfig` resolves each section **lazily**, on first
access, rather than eagerly validating the whole file up front -- a config missing a
section a given command doesn't need is perfectly valid.
"""

import importlib
from collections.abc import Mapping
from dataclasses import dataclass, field
from pathlib import Path
from typing import Any

import m4opt.missions
from astropy import units as u
from astropy.units import Quantity
from m4opt.missions import Mission

from uvex_transients.models.core.priors import Prior
from uvex_transients.simulation.core import SurveySimulator
from uvex_transients.surveys.base import SurveySchedule
from uvex_transients.surveys.utils import get_schedule
from uvex_transients.transients.base import TransientBase

from .yaml_tags import get_run_yaml

_DEFAULT_MISSION = "uvex"

# `run_cut`'s own positional parameters -- a cut's `params:` block may not use these
# names, since they'd otherwise collide with the call `SurveySimulator.run_cut(type, catalog,
# mission, **params)` makes.
_RESERVED_CUT_PARAMS = frozenset({"catalog", "mission"})

# Every `uvex_transients.transients` submodule that defines a concrete `TransientBase`
# subclass. Nothing in `uvex_transients`'s own import chain imports these eagerly -- a
# class only registers once its defining module has actually executed (see
# `TransientBase.registry`) -- so this list must be imported explicitly before any
# `class:` lookup. Adding a 6th transient type means adding its module name here too;
# `tests/cli/test_config.py::test_known_transient_modules_matches_the_package_directory`
# fails loudly if this list and the actual package contents ever drift apart.
_KNOWN_TRANSIENT_MODULES = ("TDEs", "LFBOTs", "kilonovae", "supernovae")


def _import_known_transients() -> None:
    """Import every module in `_KNOWN_TRANSIENT_MODULES` so `TransientBase.registry()` is fully populated."""
    for module_name in _KNOWN_TRANSIENT_MODULES:
        importlib.import_module(f"uvex_transients.transients.{module_name}")


def _parse_quantity(value: Any, default_unit: u.UnitBase) -> Quantity:
    """
    Resolve a YAML value (a `Quantity`, a unit string like ``"200 day"``, or a bare number) to a `Quantity`.

    Parameters
    ----------
    value : ~astropy.units.Quantity, str, or float
        The raw YAML value.
    default_unit : ~astropy.units.UnitBase
        Unit to apply if `value` is a bare number.

    Returns
    -------
    ~astropy.units.Quantity
        The resolved quantity.
    """
    if isinstance(value, Quantity):
        return value
    if isinstance(value, str):
        return Quantity(value)
    return Quantity(value, default_unit)


def _resolve_mission(name: str) -> Mission:
    """
    Resolve a mission name (e.g. ``"uvex"``) to its `m4opt.missions.Mission` instance.

    Parameters
    ----------
    name : str
        The mission's attribute name in `m4opt.missions`.

    Returns
    -------
    m4opt.missions.Mission
        The resolved mission.

    Raises
    ------
    ValueError
        If `name` is not a known `m4opt.missions.Mission` attribute.
    """
    mission = getattr(m4opt.missions, name, None)
    if not isinstance(mission, Mission):
        available = sorted(attr for attr, value in vars(m4opt.missions).items() if isinstance(value, Mission))
        raise ValueError(f"Unknown mission {name!r}; available: {available}.")
    return mission


def _resolve_schedule(section: Mapping) -> SurveySchedule:
    """
    Resolve a ``schedule:`` block's ``name:`` / ``url:`` / ``path:``+``fov_path:`` (mutually exclusive).

    Parameters
    ----------
    section : Mapping
        The parsed ``schedule:`` YAML block.

    Returns
    -------
    SurveySchedule
        The resolved schedule.

    Raises
    ------
    ValueError
        If more than one of ``name``/``url``/``path``+``fov_path`` is given, or
        ``path`` is given without ``fov_path`` (or vice versa).
    """
    name = section.get("name")
    url = section.get("url")
    path = section.get("path")
    fov_path = section.get("fov_path")

    given = [label for label, present in (("name", name), ("url", url), ("path/fov_path", path or fov_path)) if present]
    if len(given) > 1:
        raise ValueError(f"'schedule:' must give at most one of 'name', 'url', or 'path'+'fov_path', got {given}.")

    if path is not None or fov_path is not None:
        if path is None or fov_path is None:
            raise ValueError("'schedule:' with 'path' also requires 'fov_path' (and vice versa).")
        return SurveySchedule.from_disk(path, fov_path=fov_path)

    return get_schedule(name=name, url=url)


def _apply_parameter_overrides(sed, overrides: Mapping) -> None:
    """
    Apply a ``parameters:`` block's per-parameter overrides via `Parameter.fix`/`Parameter.set_prior`.

    Parameters
    ----------
    sed : SpectralModel, Lightcurve, or Spectrum
        The model whose parameters to override, keyed by name.
    overrides : Mapping
        The parsed ``parameters:`` YAML block.

    Raises
    ------
    KeyError
        If `overrides` names a parameter `sed` doesn't have.
    """
    for name, value in overrides.items():
        try:
            parameter = sed[name]
        except KeyError:
            raise KeyError(
                f"{sed.__class__.__name__} has no parameter named {name!r}. Valid parameters are {tuple(sed)}."
            ) from None

        if isinstance(value, Prior):
            parameter.set_prior(value)
        else:
            # A unit string (e.g. "2.0 day") fixes a unit-full parameter; a bare number
            # only works for a genuinely dimensionless one -- `Parameter.fix` itself
            # raises a clear `TypeError` if the two are incompatible.
            parameter.fix(Quantity(value) if isinstance(value, str) else value)


def _resolve_transients(section: Mapping) -> dict[str, TransientBase]:
    """
    Resolve a ``transients:`` block into ``{key: TransientBase instance}``.

    Parameters
    ----------
    section : Mapping
        The parsed ``transients:`` YAML block.

    Returns
    -------
    dict of str to TransientBase
        One constructed, configured transient instance per declared key.

    Raises
    ------
    ValueError
        If `section` is empty, an entry is missing its required ``class`` key,
        names an unknown transient class, or has unrecognized key(s) left over.
    """
    if not section:
        raise ValueError("'transients:' must declare at least one transient type.")

    _import_known_transients()
    registry = TransientBase.registry()

    transients: dict[str, TransientBase] = {}
    for key, raw_entry in section.items():
        entry = dict(raw_entry)

        class_name = entry.pop("class", None)
        if class_name is None:
            raise ValueError(f"transients.{key!r} is missing required key 'class'.")
        try:
            transient_cls = registry[class_name]
        except KeyError:
            raise ValueError(
                f"transients.{key!r}: unknown transient class {class_name!r}; available: {sorted(registry)}."
            ) from None

        cosmology = entry.pop("cosmology", None)
        transient = transient_cls(cosmology=cosmology)

        z_limit = entry.pop("z_limit", None)
        if z_limit is not None:
            transient.redshift_limit = z_limit

        duration_limit = entry.pop("duration_limit", None)
        if duration_limit is not None:
            transient.duration_limit = _parse_quantity(duration_limit, u.day)

        parameters = entry.pop("parameters", {}) or {}
        _apply_parameter_overrides(transient.sed, parameters)

        if entry:
            raise ValueError(f"transients.{key!r} has unknown key(s) {sorted(entry)}.")

        transients[key] = transient

    return transients


def _resolve_cuts(section: Mapping) -> dict[str, "CutSpec"]:
    """
    Resolve a ``cuts:`` block into ``{key: CutSpec}``, validating each ``type:`` against `SurveySimulator`.

    Parameters
    ----------
    section : Mapping
        The parsed ``cuts:`` YAML block.

    Returns
    -------
    dict of str to CutSpec
        One resolved `CutSpec` per declared key, in declared order.

    Raises
    ------
    ValueError
        If an entry is missing its required ``type`` key, names an unknown
        cut type, or uses a reserved parameter name.
    """
    available = SurveySimulator.available_cuts()

    cuts: dict[str, CutSpec] = {}
    for key, raw_entry in section.items():
        entry = dict(raw_entry)

        cut_type = entry.pop("type", None)
        if cut_type is None:
            raise ValueError(f"cuts.{key!r} is missing required key 'type'.")
        if cut_type not in available:
            raise ValueError(f"cuts.{key!r}: unknown cut type {cut_type!r}; available: {list(available)}.")

        reserved = _RESERVED_CUT_PARAMS & set(entry)
        if reserved:
            raise ValueError(f"cuts.{key!r}: params cannot use reserved name(s) {sorted(reserved)}.")

        cuts[key] = CutSpec(type=cut_type, params=entry)

    return cuts


[docs] @dataclass class GenerateConfig: """Parsed ``generate:`` section -- see `SurveySimulator.generate_events`.""" time_bins: int nside: int | None = None order: str | None = None downsample: int | dict[str, int] | None = None
[docs] @dataclass class CutSpec: """One resolved entry of a ``cuts:`` section -- a `SurveySimulator.available_cuts()` name plus its params.""" type: str params: dict[str, Any] = field(default_factory=dict)
[docs] @dataclass class PhotometryConfig: """Parsed ``photometry:`` section -- see `EventCatalog.simulate_photometry`.""" bands: list[str] | None = None n_sigma: float | None = None
[docs] @dataclass class YieldConfig: """Parsed ``yield:`` section -- see `EventCatalog.compute_yield_summary`.""" confidence: float = 0.9
[docs] @dataclass class DetectionCountsConfig: """Parsed ``detection_counts:`` section -- see `PhotometryCatalog.compute_detection_count_table`.""" snr_threshold: float confidence: float = 0.9
[docs] class RunConfig: """ A parsed CLI run-config, resolving each section lazily on first access. Every command reads the same config file; a command only touches the properties it actually needs (`generate` reads `.schedule`/`.transients`/`.mission`/`.generate`; `cut` reads `.schedule`/`.transients`/`.mission`/`.cuts`; `photometry` reads `.schedule`/`.transients`/`.mission`/`.photometry`), so a config missing an unrelated section (e.g. no `photometry:` block, if you never run that command) still works. Parameters ---------- raw : Mapping The parsed run-config YAML, as a nested mapping. source : str or ~pathlib.Path, optional The config file's path, used only to make error messages more specific. """ def __init__(self, raw: Mapping, source: Path | None = None): """ Store the parsed config; every section is resolved lazily on first access. Parameters ---------- raw : Mapping The parsed run-config YAML, as a nested mapping. source : str or ~pathlib.Path, optional The config file's path, used only to make error messages more specific. """ self._raw = raw self._source = source self._schedule: SurveySchedule | None = None self._mission: Mission | None = None self._transients: dict[str, TransientBase] | None = None self._simulator: SurveySimulator | None = None self._generate: GenerateConfig | None = None self._cuts: dict[str, CutSpec] | None = None self._photometry: PhotometryConfig | None = None self._yield: YieldConfig | None = None self._detection_counts: DetectionCountsConfig | None = None self._keep_intermediate: bool | None = None
[docs] @classmethod def from_yaml(cls, path: str | Path) -> "RunConfig": """ Parse a run-config YAML file (see `uvex_transients.cli.yaml_tags.get_run_yaml`). Parameters ---------- path : str or ~pathlib.Path Path to the run-config YAML file. Returns ------- RunConfig The parsed config. """ path = Path(path) with open(path) as f: raw = get_run_yaml().load(f) or {} return cls(raw, source=path)
[docs] def has_section(self, name: str) -> bool: """ Whether the parsed config has a top-level ``name:`` section at all. Parameters ---------- name : str The section name to check for. Returns ------- bool Whether the section is present. """ return name in self._raw
def _require_section(self, name: str, command: str) -> Mapping: """ Return a required top-level section, raising a clear error if it's missing. Parameters ---------- name : str The section name to look up. command : str The CLI command that requires it, used to phrase the error message. Returns ------- Mapping The section's parsed contents. Raises ------ ValueError If the section is missing. """ section = self._raw.get(name) if section is None: where = f" ({self._source})" if self._source else "" raise ValueError(f"Config{where} is missing a '{name}:' section, required by the '{command}' command.") return section @property def schedule(self) -> SurveySchedule: """ The resolved `SurveySchedule` (``schedule:`` section; falls back to the package default). Returns ------- SurveySchedule The resolved schedule. """ if self._schedule is None: self._schedule = _resolve_schedule(self._raw.get("schedule") or {}) return self._schedule @property def mission(self) -> Mission: """ The resolved `Mission` (``mission:`` section; defaults to ``"uvex"``). Returns ------- m4opt.missions.Mission The resolved mission. """ if self._mission is None: self._mission = _resolve_mission(self._raw.get("mission", _DEFAULT_MISSION)) return self._mission @property def transients(self) -> dict[str, TransientBase]: """ The resolved ``{key: TransientBase instance}`` (``transients:`` section, required). Returns ------- dict of str to TransientBase One constructed, configured transient instance per declared key. """ if self._transients is None: self._transients = _resolve_transients(self._require_section("transients", "generate/cut/photometry")) return self._transients @property def simulator(self) -> SurveySimulator: """ A `SurveySimulator` built from `.schedule`/`.transients` (cached across one CLI invocation). Returns ------- SurveySimulator The resolved simulator. """ if self._simulator is None: seed = (self._raw.get("generate") or {}).get("seed") self._simulator = SurveySimulator(self.schedule, transients=self.transients, simulation_seed=seed) return self._simulator @property def generate(self) -> GenerateConfig: """ The parsed ``generate:`` section (required by the ``generate`` command). Returns ------- GenerateConfig The parsed section. """ if self._generate is None: section = self._require_section("generate", "generate") if "time_bins" not in section: raise ValueError("'generate:' is missing required key 'time_bins'.") self._generate = GenerateConfig( time_bins=section["time_bins"], nside=section.get("nside"), order=section.get("order"), downsample=section.get("downsample"), ) return self._generate @property def cuts(self) -> dict[str, CutSpec]: """ The parsed ``cuts:`` section, in declared order (required by the ``cut`` command). Returns ------- dict of str to CutSpec One resolved `CutSpec` per declared key, in declared order. """ if self._cuts is None: self._cuts = _resolve_cuts(self._require_section("cuts", "cut")) return self._cuts @property def photometry(self) -> PhotometryConfig: """ The parsed ``photometry:`` section (optional; every field defaults to "every band"/the package default). Returns ------- PhotometryConfig The parsed section. """ if self._photometry is None: section = self._raw.get("photometry") or {} self._photometry = PhotometryConfig(bands=section.get("bands"), n_sigma=section.get("n_sigma")) return self._photometry @property def yield_config(self) -> YieldConfig: """ The parsed ``yield:`` section (optional; defaults to a ``0.9`` Clopper-Pearson confidence level). Returns ------- YieldConfig The parsed section. """ if self._yield is None: section = self._raw.get("yield") or {} self._yield = YieldConfig(confidence=section.get("confidence", 0.9)) return self._yield @property def detection_counts(self) -> DetectionCountsConfig: """ The parsed ``detection_counts:`` section (required by the ``detection-counts`` command). Returns ------- DetectionCountsConfig The parsed section. """ if self._detection_counts is None: section = self._require_section("detection_counts", "detection-counts") if "snr_threshold" not in section: raise ValueError("'detection_counts:' is missing required key 'snr_threshold'.") self._detection_counts = DetectionCountsConfig( snr_threshold=section["snr_threshold"], confidence=section.get("confidence", 0.9), ) return self._detection_counts @property def keep_intermediate(self) -> bool: """ Whether the ``run`` command should keep each stage's catalog on disk (top-level ``keep_intermediate:``). Defaults to `True`; set to `False` to have ``run`` write only the final photometry table, discarding the generated/cut catalogs once the next stage no longer needs them. Returns ------- bool Whether to keep intermediate stage files. """ if self._keep_intermediate is None: self._keep_intermediate = bool(self._raw.get("keep_intermediate", True)) return self._keep_intermediate