File size: 6,003 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 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 | """Model and method factories shared by training, sampling and evaluation.
Methods: ``direct_si``, ``cascaded_si`` and ``edm``. Models: ``unet2_5ch`` (Euler, SQG), ``unet2_5ch_coord``
(passive scalar), ``unet2_5ch_coord_aug`` (heterogeneous waves) and ``unet2_2comp`` (NLKG); each run configuration
under records/<system>/runs/ names its model and method."""
from __future__ import annotations
from dataclasses import dataclass
from typing import Callable
from .methods.edm import EDMMethod
from .methods.direct_si import DirectSIMethod
from .methods.cascaded_si import CascadedSIMethod
from .models.unet2 import Conditional2D, LocalDrift2D
SEGMENT_DATA_MODES = ("segment_pairs", "fine_segment_pairs")
# reserved transient key: the method's segment boundaries in physical time, injected by
# build_model for backbones that declare physical times in their own YAML (see module docstring)
T_VALUES_KEY = "_method_t_values"
@dataclass(frozen=True)
class ModelSpec:
name: str
endpoint_cls: type
segment_cls: type
needs_dataset_norm: bool = False
description: str = ""
MODELS: dict[str, ModelSpec] = {}
METHODS: dict[str, Callable[[dict], object]] = {}
def register_model(name: str, endpoint_cls: type, segment_cls: type,
needs_dataset_norm: bool = False, description: str = "",
overwrite: bool = False) -> ModelSpec:
if name in MODELS and not overwrite:
raise KeyError(f"model {name!r} already registered")
spec = ModelSpec(name, endpoint_cls, segment_cls, needs_dataset_norm, description)
MODELS[name] = spec
return spec
def register_method(name: str, factory: Callable[[dict], object],
overwrite: bool = False) -> None:
if name in METHODS and not overwrite:
raise KeyError(f"method {name!r} already registered")
METHODS[name] = factory
def head_for_data_mode(data_mode: str | None) -> str:
return "segment" if data_mode in SEGMENT_DATA_MODES else "endpoint"
def build_model(model_cfg: dict, method_cfg: dict | None = None,
data_cfg: dict | None = None):
"""Construct the backbone named by ``model_cfg['name']``.
``method_cfg['data_mode']`` selects the head. For models with ``needs_dataset_norm`` the dataset path is
taken from ``data_cfg`` (``packaged_dataset_path``); a model configuration that carries its own copy of the
key is accepted only if it agrees with the data config.
"""
name = model_cfg["name"]
if name not in MODELS:
raise ValueError(f"unknown model: {name!r} (registered: {sorted(MODELS)})")
spec = MODELS[name]
cfg = dict(model_cfg)
if spec.needs_dataset_norm:
data_path = (data_cfg or {}).get("packaged_dataset_path")
model_path = cfg.get("packaged_dataset_path")
if data_path is None and model_path is None:
raise KeyError(
f"model {name!r} needs the dataset normalization: pass data_cfg with "
"'packaged_dataset_path' (model configurations do not carry dataset paths)")
if data_path is not None and model_path is not None:
if _norm_path(data_path) != _norm_path(model_path):
raise ValueError(
f"model.packaged_dataset_path ({model_path}) != data.packaged_dataset_path "
f"({data_path}) — the FFT wrapper must use the training dataset's statistics")
cfg["packaged_dataset_path"] = str(data_path if data_path is not None else model_path)
t_values = ((method_cfg or {}).get("sampling") or {}).get("t_values")
if t_values:
cfg[T_VALUES_KEY] = list(t_values)
data_mode = (method_cfg or {}).get("data_mode")
# a model configuration that declares the horizon planes must agree with the method's time grid:
# t_start = t_values[0], t_end = t_values[-1]
tv = ((method_cfg or {}).get("sampling") or {}).get("t_values")
if tv is not None and ("t_start" in cfg or "t_end" in cfg):
t0, t1 = float(tv[0]), float(tv[-1])
if "t_start" in cfg and abs(float(cfg["t_start"]) - t0) > 1e-9:
raise ValueError(f"model {name!r}: t_start {cfg['t_start']} != method t_values[0] = {t0}")
if "t_end" in cfg and abs(float(cfg["t_end"]) - t1) > 1e-9:
raise ValueError(f"model {name!r}: t_end {cfg['t_end']} != method horizon t_values[-1] = {t1} "
f"(the model constant must equal the forecast horizon)")
cls = spec.segment_cls if head_for_data_mode(data_mode) == "segment" else spec.endpoint_cls
return cls.from_config(cfg)
def build_method(method_cfg: dict):
name = method_cfg["name"]
if name not in METHODS:
raise ValueError(f"unknown method: {name!r} (registered: {sorted(METHODS)})")
return METHODS[name](method_cfg)
def _norm_path(p: str) -> str:
import os
return os.path.normpath(os.path.abspath(str(p)))
# ── built-ins ────────────────────────────────────────────────────────────────
register_model("unet2", Conditional2D, LocalDrift2D, needs_dataset_norm=True,
description="FFT-wrapped 2-D circular-conv U-Net on the 32x32 coarse field "
"(Conditional2D 3-ch / LocalDrift2D 5-ch)")
register_method("edm", EDMMethod.from_config)
register_method("direct_si", DirectSIMethod.from_config)
register_method("cascaded_si", CascadedSIMethod.from_config)
# the backbones register themselves on import; unet2_coord provides the coordinate planes of unet2_5ch_coord(_aug)
from .models import unet2_coord as _unet2_coord # noqa: F401,E402
from .models import unet2_5ch as _unet2_5ch # noqa: F401,E402
from .models import unet2_5ch_coord as _unet2_5ch_coord # noqa: F401,E402
from .models import unet2_5ch_coord_aug as _unet2_5ch_coord_aug # noqa: F401,E402
from .models import unet2_2comp as _unet2_2comp # noqa: F401,E402
|