| """Configuration for the standalone SFNO smoke workflow.""" |
|
|
| from __future__ import annotations |
|
|
| import json |
| from dataclasses import dataclass, fields |
| from pathlib import Path |
| from typing import Any |
|
|
|
|
| @dataclass(frozen=True) |
| class SFNOConfig: |
| nlat: int = 17 |
| nlon: int = 32 |
| channels: int = 2 |
| timesteps: int = 6 |
| batch_size: int = 2 |
| embed_dim: int = 8 |
| num_layers: int = 2 |
| scale_factor: int = 2 |
| rollout_steps: int = 3 |
| learning_rate: float = 0.001 |
| seed: int = 2026 |
| grid: str = "equiangular" |
| grid_internal: str = "legendre-gauss" |
|
|
| def validate(self) -> None: |
| positive = ( |
| "nlat", |
| "nlon", |
| "channels", |
| "timesteps", |
| "batch_size", |
| "embed_dim", |
| "num_layers", |
| "scale_factor", |
| "rollout_steps", |
| ) |
| for name in positive: |
| if getattr(self, name) <= 0: |
| raise ValueError(f"{name} must be positive") |
| if self.nlat < 5 or self.nlon < 8: |
| raise ValueError("The spherical grid must be at least 5 x 8") |
| if self.nlon % self.scale_factor: |
| raise ValueError("nlon must be divisible by scale_factor") |
| if self.timesteps < self.batch_size + 1: |
| raise ValueError("timesteps must provide at least batch_size input/target pairs") |
| if self.rollout_steps >= self.timesteps: |
| raise ValueError("rollout_steps must be smaller than timesteps") |
| if self.grid not in {"equiangular", "legendre-gauss", "lobatto", "equidistant"}: |
| raise ValueError(f"Unsupported input grid: {self.grid}") |
|
|
|
|
| def load_config(path: str | Path) -> SFNOConfig: |
| path = Path(path) |
| values: dict[str, Any] = json.loads(path.read_text(encoding="utf-8")) |
| known = {field.name for field in fields(SFNOConfig)} |
| unknown = sorted(set(values) - known) |
| if unknown: |
| raise ValueError(f"Unknown config keys: {', '.join(unknown)}") |
| config = SFNOConfig(**values) |
| config.validate() |
| return config |
|
|
