File size: 2,864 Bytes
bdccc62 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 | 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)
|