| """Configuration files of the repository: the dataset cards, the evaluation settings and the recorded runs. |
| |
| * ``configs/<system>/dataset.yaml`` names the data of one system; ``${VAR}`` entries are read from the environment. |
| * ``configs/evaluation.yaml`` holds the csFID settings and the sampling seeds shared by all systems. |
| * ``records/<system>/units.json`` lists the models of the paper and their recorded scores. |
| """ |
| from __future__ import annotations |
|
|
| import json |
| import os |
| import re |
| from pathlib import Path |
|
|
| import yaml |
|
|
| REPO = Path(__file__).resolve().parents[1] |
| CONFIGS = REPO / "configs" |
| RECORDS = REPO / "records" |
| SYSTEMS = ("euler", "sqg", "waves", "nlkg", "passive_scalar") |
|
|
| _ENV_RE = re.compile(r"\$\{([A-Za-z_][A-Za-z0-9_]*)(?::-([^}]*))?\}") |
|
|
|
|
| def expand_env(value): |
| """``${VAR}`` / ``${VAR:-default}`` expansion in strings (recursively in dicts and lists).""" |
| if isinstance(value, str): |
| def sub(m): |
| var, default = m.group(1), m.group(2) |
| v = os.environ.get(var) |
| if v is None or v == "": |
| if default is None: |
| raise KeyError(f"environment variable {var} is not set (used in a config)") |
| return default |
| return v |
| return _ENV_RE.sub(sub, value) |
| if isinstance(value, dict): |
| return {k: expand_env(v) for k, v in value.items()} |
| if isinstance(value, list): |
| return [expand_env(v) for v in value] |
| return value |
|
|
|
|
| def load_yaml(path: str | os.PathLike) -> dict: |
| with open(path) as f: |
| return yaml.safe_load(f) or {} |
|
|
|
|
| def dataset_card(system: str, expand: bool = False) -> dict: |
| """The dataset card of ``system``; with ``expand`` the ``${VAR}`` paths are resolved from the environment.""" |
| if system not in SYSTEMS: |
| raise KeyError(f"unknown system {system!r} (one of {SYSTEMS})") |
| card = load_yaml(CONFIGS / system / "dataset.yaml") |
| return expand_env(card) if expand else card |
|
|
|
|
| def data_paths(system: str, data: str | None = None, truth_bank: str | None = None) -> tuple[Path, Path]: |
| """The package and truth-bank directories: the given values, else the environment variables of the card.""" |
| card = dataset_card(system) |
| if data is None: |
| data = expand_env(card["data"]) |
| if truth_bank is None: |
| truth_bank = expand_env(card["truth_bank"]) |
| return Path(data).resolve(), Path(truth_bank).resolve() |
|
|
|
|
| def evaluation_settings() -> dict: |
| return load_yaml(CONFIGS / "evaluation.yaml") |
|
|
|
|
| def records(system: str) -> dict: |
| return json.loads((RECORDS / system / "units.json").read_text()) |
|
|
|
|
| def unit(system: str, method: str, S: int | None = None) -> tuple[dict, dict]: |
| """The recorded model of ``system`` for a paper method (``direct_si``, ``edm``, ``cascaded_si``); with ``S`` the |
| Cascaded SI model of that depth (Table 6). Returns (records, unit).""" |
| rec = records(system) |
| if method == "edm": |
| if S is not None: |
| raise SystemExit("--S applies to cascaded_si only") |
| key = "edm" |
| elif method == "direct_si": |
| if S not in (None, 1): |
| raise SystemExit("direct_si is S = 1; use --method cascaded_si --S <S> for other depths") |
| key = rec["table2"]["direct_si"] |
| elif method == "cascaded_si": |
| key = rec["table2"]["cascaded_si"] if S is None else f"cascaded_si_S{int(S):02d}" |
| else: |
| raise SystemExit(f"unknown method {method!r}") |
| u = next((u for u in rec["units"] if u["key"] == key), None) |
| if u is None: |
| raise SystemExit(f"{system}: no recorded model {key} (depths: {[u['S'] for u in rec['units'] if u['key'] != 'edm']})") |
| return rec, u |
|
|
|
|
| def run_config(system: str, u: dict, data: Path) -> dict: |
| """The recorded run configuration of a model with the dataset path filled in.""" |
| cfg = load_yaml(REPO / u["config"]) |
| cfg["data"]["packaged_dataset_path"] = str(Path(data).resolve()) + "/" |
| return cfg |
|
|
