core/loader.py

"""Declarative experiment registry for autonomous SciML research.

This module now loads experiment configurations from 'experiments.yaml'.
To add a new experiment, append it to 'experiments.yaml'.
"""

import yaml
from pathlib import Path
from dataclasses import dataclass, field
from typing import List, Optional

# ── Config dataclass ──────────────────────────────────────────────────────────

@dataclass
class ExperimentConfig:
    name:        str              # unique key (used for dedup against results.json)
    benchmark:   str              # "burgers_1d" | "darcy_2d" | "kdv_1d" | "wave_1d"
    model:       str              # "FNO" | "RFNO" | "AFNO" | "FFNO" | "UNO" | "WNO" | "DeepONet" | "PODDeepONet"
    hidden_dim:  int              # channel width
    n_layers:    int              # depth (FNO blocks per level for UNO)
    n_modes:     int  = 16        # Fourier modes (FNO / UNO / RFNO / AFNO)
    n_levels:    int  = 3         # Haar levels (WNO)
    n_head:      int  = 4         # Attention heads (Transolver / GNOT / Transformer)
    slice_num:   int  = 32        # Physics slices (Transolver)
    lr:          float = 1e-3     # learning rate
    batch_size:  int  = 64         # training batch size
    grad_clip:   float = 1.0      # gradient clipping (0 = disabled)
    pino_lambda: float = 0.0      # PINO physics-loss weight
    loss_type:   str  = "l2_rel"  # loss function
    h1_alpha:    float = 0.1      # H1 loss derivative weight
    augment:     bool  = False    # spatial-shift augmentation
    curriculum:  bool  = False    # training curriculum
    curriculum_epochs: int = 0    # Phase 12: Curriculum ramp epochs
    save_ckpt:   bool  = False    # save model checkpoint after training
    resume:      bool  = False    # resume from best checkpoint if exists
    resume_from: str   = ""       # resume from specific checkpoint name/path
    budget_s:    int  = 1200      # training time budget in seconds
    parent_name: str  = ""        # name of parent experiment
    priority:    int  = 5         # 1 = highest
    rationale:   str  = ""        # why this experiment?
    expected:    str  = ""        # expected val_l2_rel range
    paper_ref:   str  = ""        # paper ID from papers/*.yaml
    refine_grid: bool = False      # Phase 11: Adaptive grid refinement
    cheb_degree: int  = 5         # Phase 11: Degree for Chebyshev KAN
    seed:        int  = 42        # Global RNG seed for reproducibility
    lr_schedule: str  = "warmup_cosine"  # LR schedule: warmup_cosine|cosine|onecycle|none
    ema_decay:   float = 0.0      # EMA decay for model weights (0=disabled, 0.999 recommended)
    patience:    int  = 5         # early-stop patience: halt after this many consecutive non-improving evals (0=off)
    snapshot_ensemble: int = 0    # Phase 12: Snapshot ensembling (number of snapshots to average)
    n_iterations: int = 10        # Phase 14: Iterations for Neural Iterative Solver (Brandstetter et al.)

    def to_cli_args(self) -> List[str]:
        args = [
            "--benchmark",   self.benchmark,
            "--model",       self.model,
            "--name",        self.name,
            "--hidden",      str(self.hidden_dim),
            "--layers",      str(self.n_layers),
            "--modes",       str(self.n_modes),
            "--levels",      str(self.n_levels),
            "--lr",          str(self.lr),
            "--batch_size",  str(self.batch_size),
            "--grad_clip",   str(self.grad_clip),
            "--pino_lambda", str(self.pino_lambda),
            "--loss",        self.loss_type,
        ]
        if self.loss_type.startswith("h1"):
            args += ["--h1_alpha", str(self.h1_alpha)]
        if self.augment:
            args += ["--augment"]
        if self.curriculum:
            args += ["--curriculum"]
        if self.curriculum_epochs > 0:
            args += ["--curriculum_epochs", str(self.curriculum_epochs)]
        if self.model in ("Transolver", "Transolver2D", "Transolver2d", "GNOT", "GNOT_Axial2d"):
            args += ["--n_head", str(self.n_head)]
            args += ["--slice_num", str(self.slice_num)]
        if self.save_ckpt:
            args += ["--save_ckpt"]
        if self.resume:
            args += ["--resume"]
        if self.resume_from:
            args += ["--resume_from", self.resume_from]
        if self.refine_grid:
            args += ["--refine_grid"]
        if self.model == "cPIKAN_FNO":
            args += ["--degree", str(self.cheb_degree)]
        if self.model in ("IterativeFNO2D", "IterativeFNO2d"):
            args += ["--n_iterations", str(self.n_iterations)]
        args += ["--budget", str(self.budget_s)]
        args += ["--seed", str(self.seed)]
        if self.lr_schedule != "warmup_cosine":
            args += ["--lr_schedule", self.lr_schedule]
        if self.ema_decay > 0:
            args += ["--ema_decay", str(self.ema_decay)]
        if self.patience != 5:
            args += ["--patience", str(self.patience)]
        if self.snapshot_ensemble > 0:
            args += ["--snapshot_ensemble", str(self.snapshot_ensemble)]
        return args

    def short(self) -> str:
        """One-line summary for logging."""
        parts = [f"{self.model}", f"h={self.hidden_dim}", f"l={self.n_layers}"]
        if self.model in ("FNO", "RFNO", "AFNO", "FFNO", "UNO", "UNO2d", "WNO2d", "Transolver", "Transolver2D"):
            parts.append(f"m={self.n_modes}")
        if self.model in ("Transolver", "Transolver2D", "GNOT", "GNOT_Axial2d"):
            parts.append(f"h={self.n_head}")
            parts.append(f"s={self.slice_num}")
        if self.model in ("WNO", "WNO2d"):
            parts.append(f"lvl={self.n_levels}")
        if self.pino_lambda > 0:
            parts.append(f"pino={self.pino_lambda}")
        if self.lr != 1e-3:
            parts.append(f"lr={self.lr:.0e}")
        if self.grad_clip != 1.0:
            parts.append(f"clip={self.grad_clip}")
        if self.loss_type != "l2_rel":
            parts.append(f"loss={self.loss_type}")
        if self.curriculum:
            parts.append("curric")
        if self.curriculum_epochs > 0:
            parts.append(f"cur_ep={self.curriculum_epochs}")
        return "  ".join(parts)

# ── Loader Logic ─────────────────────────────────────────────────────────────

def load_experiments(yaml_path: Path) -> List[ExperimentConfig]:
    if not yaml_path.exists():
        return []
    
    with open(yaml_path, "r") as f:
        data = yaml.safe_load(f)
    
    if not data:
        return []
        
    return [ExperimentConfig(**d) for d in data]

# Load default set
REPO_ROOT = Path(__file__).parent.parent
YAML_PATH = REPO_ROOT / "experiments.yaml"
EXPERIMENTS = load_experiments(YAML_PATH)

def get_experiments(benchmark: Optional[str] = None, model: Optional[str] = None, priority: Optional[int] = None) -> List[ExperimentConfig]:
    """Return the experiments list, optionally filtered by benchmark, model, or priority."""
    queue = EXPERIMENTS
    if benchmark:
        queue = [e for e in queue if e.benchmark == benchmark]
    if model:
        queue = [e for e in queue if e.model == model]
    if priority:
        queue = [e for e in queue if e.priority <= priority]
    return queue

if __name__ == "__main__":
    # Smoke test
    print(f"Loaded {len(EXPERIMENTS)} experiments.")
    if EXPERIMENTS:
        print(f"First experiment: {EXPERIMENTS[0].name} ({EXPERIMENTS[0].benchmark})")