File size: 3,938 Bytes
41dd7e0 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 | """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
|