| from __future__ import annotations |
|
|
| import json |
| import os |
| import sys |
| from pathlib import Path |
| from typing import Dict, Tuple |
|
|
| import numpy as np |
| import yaml |
| from scipy.io import loadmat |
|
|
|
|
| PROJECT_ROOT = Path(__file__).resolve().parents[1] |
| if str(PROJECT_ROOT) not in sys.path: |
| sys.path.insert(0, str(PROJECT_ROOT)) |
|
|
|
|
| def load_config(path: str | os.PathLike | None = None) -> dict: |
| cfg_path = Path(path) if path else PROJECT_ROOT / "config" / "config.yaml" |
| with cfg_path.open("r", encoding="utf-8") as f: |
| cfg = yaml.safe_load(f) |
| return cfg |
|
|
|
|
| def ensure_dirs(cfg: dict) -> None: |
| for key in ["weight_dir", "results_dir", "logs_dir"]: |
| Path(cfg["paths"][key]).mkdir(parents=True, exist_ok=True) |
|
|
|
|
| def load_burgers_splits(cfg: dict) -> Tuple[Dict[str, np.ndarray], np.ndarray]: |
| data_cfg = cfg["data"] |
| mat_path = Path(cfg["paths"]["mat_file"]) |
| manifest_path = Path(cfg["paths"]["manifest_file"]) |
| if not mat_path.exists(): |
| raise FileNotFoundError(mat_path) |
| if manifest_path.exists(): |
| manifest = json.loads(manifest_path.read_text(encoding="utf-8")) |
| expected = manifest.get("official_loader_contract", {}) |
| if expected.get("random_seed") != data_cfg["random_seed"]: |
| raise ValueError("DATA_MANIFEST random_seed does not match config") |
| mat = loadmat(mat_path) |
| data = np.asarray(mat[data_cfg["use_input_field"]], dtype=np.float32) |
| tspan = np.asarray(mat["tspan"], dtype=np.float32).reshape(-1) |
| if data.shape != (1200, 101, 101): |
| raise ValueError(f"Unexpected Burger.mat output shape {data.shape}") |
| if data_cfg["trim_periodic_space_endpoint"]: |
| data = data[:, :, :-1] |
| if data.shape[1:] != (data_cfg["t_points"], data_cfg["x_points_after_periodic_trim"]): |
| raise ValueError(f"Unexpected trimmed trajectory shape {data.shape[1:]}") |
| rng = np.random.default_rng(data_cfg["random_seed"]) |
| indices = np.arange(len(data)) |
| rng.shuffle(indices) |
| train_n = int(data_cfg["train_count"]) |
| val_n = int(data_cfg["val_count"]) |
| test_n = int(data_cfg["test_count"]) |
| if train_n + val_n + test_n != len(data): |
| raise ValueError("Configured split counts must sum to 1200") |
| train_idx = indices[:train_n] |
| val_idx = indices[train_n : train_n + val_n] |
| test_idx = indices[train_n + val_n : train_n + val_n + test_n] |
| splits = { |
| "train": data[train_idx].astype(np.float32), |
| "val": data[val_idx].astype(np.float32), |
| "test": data[test_idx].astype(np.float32), |
| } |
| return splits, tspan |
|
|
|
|
| def write_json(path: str | os.PathLike, payload: dict) -> None: |
| path = Path(path) |
| path.parent.mkdir(parents=True, exist_ok=True) |
| tmp = path.with_suffix(path.suffix + ".tmp") |
| tmp.write_text(json.dumps(payload, indent=2, ensure_ascii=False), encoding="utf-8") |
| tmp.replace(path) |
|
|
