| """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") |
|
|
| |
| |
| 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") |
| |
| |
| 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))) |
|
|
|
|
| |
| 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) |
|
|
| |
| from .models import unet2_coord as _unet2_coord |
| from .models import unet2_5ch as _unet2_5ch |
| from .models import unet2_5ch_coord as _unet2_5ch_coord |
| from .models import unet2_5ch_coord_aug as _unet2_5ch_coord_aug |
| from .models import unet2_2comp as _unet2_2comp |
|
|
