| from __future__ import annotations |
|
|
| import csv |
| import json |
| import os |
| import random |
| import sys |
| from pathlib import Path |
| from typing import Any, Dict, Tuple |
|
|
| import h5py |
| import numpy as np |
| import torch |
| import yaml |
| from torch.utils.data import DataLoader, TensorDataset |
|
|
|
|
| PACKAGE_ROOT = Path(__file__).resolve().parents[1] |
| if str(PACKAGE_ROOT) not in sys.path: |
| sys.path.insert(0, str(PACKAGE_ROOT)) |
|
|
| from model.crop_navier_stokes import CROPNavierStokes, count_parameters |
|
|
|
|
| def load_config(path: str | os.PathLike[str]) -> Dict[str, Any]: |
| with open(path, "r", encoding="utf-8") as f: |
| cfg = yaml.safe_load(f) |
| cfg["_config_path"] = str(Path(path).resolve()) |
| cfg["_package_root"] = str(PACKAGE_ROOT) |
| return cfg |
|
|
|
|
| def resolve_output_path(cfg: Dict[str, Any], relative: str) -> Path: |
| out = Path(cfg["paths"]["output_dir"]) |
| path = Path(relative) |
| return path if path.is_absolute() else out / path |
|
|
|
|
| def set_seed(seed: int, deterministic: bool = False, benchmark: bool = True) -> None: |
| random.seed(seed) |
| np.random.seed(seed) |
| torch.manual_seed(seed) |
| torch.cuda.manual_seed_all(seed) |
| torch.backends.cudnn.deterministic = deterministic |
| torch.backends.cudnn.benchmark = benchmark |
|
|
|
|
| def get_device(cfg: Dict[str, Any]) -> torch.device: |
| requested = cfg.get("runtime", {}).get("device", "cuda") |
| if requested == "cuda" and torch.cuda.is_available(): |
| return torch.device("cuda") |
| return torch.device("cpu") |
|
|
|
|
| def build_model(cfg: Dict[str, Any]) -> CROPNavierStokes: |
| m = cfg["model"] |
| return CROPNavierStokes( |
| in_channels=cfg["data"]["channels"], |
| out_channels=cfg["data"]["channels"], |
| modes=m["modes"], |
| initial_channels=m["initial_channels"], |
| levels=m["levels"], |
| residual_blocks_per_level=m["residual_blocks_per_level"], |
| bottleneck_blocks=m["bottleneck_blocks"], |
| kernel_size=m["kernel_size"], |
| input_size=m["input_size"], |
| latent_size=m["latent_size"], |
| use_bn=m["use_bn"], |
| ) |
|
|
|
|
| def load_tensors(cfg: Dict[str, Any]) -> Tuple[torch.Tensor, torch.Tensor]: |
| data_path = Path(cfg["paths"]["data_dir"]) / cfg["paths"]["data_file"] |
| key = cfg["data"]["key"] |
| input_t = cfg["data"]["input_time_index"] |
| target_t = cfg["data"]["target_time_index"] |
| with h5py.File(data_path, "r") as f: |
| if key not in f: |
| raise KeyError(f"HDF5 key {key!r} not found in {data_path}") |
| dset = f[key] |
| expected = tuple(cfg["data"]["raw_shape"]) |
| if tuple(dset.shape) != expected: |
| raise ValueError(f"Expected {key} shape {expected}, got {tuple(dset.shape)}") |
| x_np = np.asarray(dset[input_t, :, :, :], dtype=np.float32).transpose(2, 0, 1) |
| y_np = np.asarray(dset[target_t, :, :, :], dtype=np.float32).transpose(2, 0, 1) |
| x = torch.from_numpy(x_np[:, None, :, :]).contiguous() |
| y = torch.from_numpy(y_np[:, None, :, :]).contiguous() |
| return x, y |
|
|
|
|
| def make_loaders(cfg: Dict[str, Any]) -> Tuple[DataLoader, DataLoader, DataLoader]: |
| x, y = load_tensors(cfg) |
| n_train = cfg["data"]["train_size"] |
| n_val = cfg["data"]["val_size"] |
| n_test = cfg["data"]["test_size"] |
| n_total = n_train + n_val + n_test |
| if x.shape[0] < n_total: |
| raise ValueError(f"Need {n_total} samples, found {x.shape[0]}") |
|
|
| train = TensorDataset(x[:n_train], y[:n_train]) |
| val = TensorDataset(x[n_train : n_train + n_val], y[n_train : n_train + n_val]) |
| test = TensorDataset(x[n_train + n_val : n_total], y[n_train + n_val : n_total]) |
|
|
| batch_size = cfg["training"]["batch_size"] |
| workers = cfg["training"].get("num_workers", 0) |
| return ( |
| DataLoader(train, batch_size=batch_size, shuffle=True, num_workers=workers, pin_memory=torch.cuda.is_available()), |
| DataLoader(val, batch_size=batch_size, shuffle=False, num_workers=workers, pin_memory=torch.cuda.is_available()), |
| DataLoader(test, batch_size=batch_size, shuffle=False, num_workers=workers, pin_memory=torch.cuda.is_available()), |
| ) |
|
|
|
|
| def relative_normalized_mse(pred: torch.Tensor, target: torch.Tensor) -> torch.Tensor: |
| return torch.mean((pred - target) ** 2) / torch.mean(target**2).clamp_min(1e-12) |
|
|
|
|
| @torch.no_grad() |
| def mean_relative_l2_percent(model: torch.nn.Module, loader: DataLoader, device: torch.device) -> float: |
| model.eval() |
| total = 0.0 |
| count = 0 |
| for x, y in loader: |
| x = x.to(device, non_blocking=True) |
| y = y.to(device, non_blocking=True) |
| pred = model(x) |
| diff = torch.linalg.vector_norm((pred - y).reshape(pred.shape[0], -1), ord=2, dim=1) |
| denom = torch.linalg.vector_norm(y.reshape(y.shape[0], -1), ord=2, dim=1).clamp_min(1e-12) |
| rel = diff / denom |
| total += rel.sum().item() |
| count += rel.numel() |
| return 100.0 * total / max(1, count) |
|
|
|
|
| @torch.no_grad() |
| def collect_predictions(model: torch.nn.Module, loader: DataLoader, device: torch.device) -> Tuple[np.ndarray, np.ndarray, np.ndarray]: |
| model.eval() |
| preds, targets, inputs = [], [], [] |
| for x, y in loader: |
| pred = model(x.to(device, non_blocking=True)).cpu() |
| preds.append(pred.numpy().astype(np.float32)) |
| targets.append(y.numpy().astype(np.float32)) |
| inputs.append(x.numpy().astype(np.float32)) |
| return np.concatenate(inputs), np.concatenate(preds), np.concatenate(targets) |
|
|
|
|
| def save_json(path: Path, payload: Dict[str, Any]) -> None: |
| path.parent.mkdir(parents=True, exist_ok=True) |
| tmp = path.with_suffix(path.suffix + ".tmp") |
| with open(tmp, "w", encoding="utf-8") as f: |
| json.dump(payload, f, indent=2, sort_keys=True) |
| tmp.replace(path) |
|
|
|
|
| def append_history(path: Path, row: Dict[str, Any]) -> None: |
| path.parent.mkdir(parents=True, exist_ok=True) |
| exists = path.exists() |
| with open(path, "a", newline="", encoding="utf-8") as f: |
| writer = csv.DictWriter(f, fieldnames=list(row.keys())) |
| if not exists: |
| writer.writeheader() |
| writer.writerow(row) |
|
|
