scripts/common.py

"""Shared configuration and device helpers for SEEDS scripts."""

from __future__ import annotations

import random
import sys
from pathlib import Path

import numpy as np
import torch
import yaml

PROJECT_ROOT = Path(__file__).resolve().parents[1]
MODEL_DIR = PROJECT_ROOT / "model"
if str(MODEL_DIR) not in sys.path:
    sys.path.insert(0, str(MODEL_DIR))

from seeds import SEEDS


def load_config(path: str) -> dict:
    with open(path, "r", encoding="utf-8") as handle:
        return yaml.safe_load(handle)


def choose_device(value: str) -> torch.device:
    if value == "auto":
        return torch.device("cuda" if torch.cuda.is_available() else "cpu")
    return torch.device(value)


def set_seed(seed: int) -> None:
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    if torch.cuda.is_available():
        torch.cuda.manual_seed_all(seed)


def build_model(config: dict) -> SEEDS:
    data, model = config["data"], config["model"]
    return SEEDS(
        channels=len(data["variables"]), faces=data["faces"], height=data["height"], width=data["width"],
        seed_count=data["seed_count"], patch_size=model["patch_size"], embed_dim=model["embed_dim"],
        spatial_layers=model["spatial_layers"], field_layers=model["field_layers"],
        sequence_layers=model["sequence_layers"], mlp_ratio=model["mlp_ratio"], dropout=model["dropout"],
        sigma_min=model["sigma_min"], sigma_max=model["sigma_max"],
    )


def resolve_path(path: str, config_path: str) -> Path:
    candidate = Path(path)
    if candidate.is_absolute():
        return candidate
    return Path(config_path).resolve().parent.parent / candidate