| from __future__ import annotations |
|
|
| import json |
| import os |
| import random |
| import sys |
| from pathlib import Path |
| from typing import Any |
|
|
| import numpy as np |
| import torch |
| import yaml |
|
|
|
|
| ROOT = Path(__file__).resolve().parents[1] |
| if str(ROOT) not in sys.path: |
| sys.path.insert(0, str(ROOT)) |
|
|
| from model.ho_fno_pipe import HOFNOPipe, count_parameters |
|
|
|
|
| class UnitTransformer: |
| def __init__(self, tensor: torch.Tensor, axes: tuple[int, ...] = (0, 1), eps: float = 1e-8) -> None: |
| self.axes = axes |
| self.eps = eps |
| self.mean = tensor.mean(dim=axes, keepdim=True) |
| self.std = tensor.std(dim=axes, keepdim=True) + eps |
|
|
| def encode(self, tensor: torch.Tensor) -> torch.Tensor: |
| return (tensor - self.mean) / self.std |
|
|
| def decode(self, tensor: torch.Tensor) -> torch.Tensor: |
| return tensor * self.std + self.mean |
|
|
| def to(self, device: torch.device) -> "UnitTransformer": |
| self.mean = self.mean.to(device) |
| self.std = self.std.to(device) |
| return self |
|
|
| def state_dict(self) -> dict[str, Any]: |
| return {"mean": self.mean.cpu(), "std": self.std.cpu(), "axes": self.axes, "eps": self.eps} |
|
|
| @classmethod |
| def from_state_dict(cls, state: dict[str, Any]) -> "UnitTransformer": |
| obj = cls.__new__(cls) |
| obj.mean = state["mean"] |
| obj.std = state["std"] |
| obj.axes = tuple(state.get("axes", (0, 1))) |
| obj.eps = float(state.get("eps", 1e-8)) |
| return obj |
|
|
|
|
| def load_config(path: str | os.PathLike[str] | None = None) -> dict[str, Any]: |
| cfg_path = Path(path) if path else ROOT / "config" / "config.yaml" |
| with cfg_path.open("r", encoding="utf-8") as handle: |
| cfg = yaml.safe_load(handle) |
| cfg["_config_path"] = str(cfg_path) |
| return cfg |
|
|
|
|
| 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 resolve(path: str | os.PathLike[str], base: Path | None = None) -> Path: |
| p = Path(path) |
| if p.is_absolute(): |
| return p |
| return (base or ROOT) / p |
|
|
|
|
| def select_device(require_accelerator: bool = False) -> torch.device: |
| if torch.cuda.is_available(): |
| return torch.device("cuda") |
| if require_accelerator: |
| raise RuntimeError("accelerator is required, but torch.cuda.is_available() is False") |
| return torch.device("cpu") |
|
|
|
|
| def load_pipe_tensors(cfg: dict[str, Any]) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: |
| data_cfg = cfg["data"] |
| data_dir = Path(cfg["paths"]["data_dir"]) |
| train_size = int(data_cfg["train_size"]) |
| test_size = int(data_cfg["test_size"]) |
| total_used = int(data_cfg["total_used"]) |
| r1, r2 = (int(v) for v in data_cfg.get("downsample", [1, 1])) |
| s1 = int(((int(data_cfg["resolution"][0]) - 1) / r1) + 1) |
| s2 = int(((int(data_cfg["resolution"][1]) - 1) / r2) + 1) |
|
|
| x_raw = np.load(data_dir / data_cfg["x_file"], mmap_mode="r")[:total_used, ::r1, ::r2][:, :s1, :s2] |
| y_raw = np.load(data_dir / data_cfg["y_file"], mmap_mode="r")[:total_used, ::r1, ::r2][:, :s1, :s2] |
| q_raw = np.load(data_dir / data_cfg["q_file"], mmap_mode="r")[ |
| :total_used, int(data_cfg["output_channel_index"]), ::r1, ::r2 |
| ][:, :s1, :s2] |
|
|
| coords = np.stack((np.asarray(x_raw, dtype=np.float32), np.asarray(y_raw, dtype=np.float32)), axis=-1) |
| target = np.asarray(q_raw, dtype=np.float32)[..., None] |
| x_all = torch.from_numpy(coords.copy()) |
| y_all = torch.from_numpy(target.copy()) |
| return x_all[:train_size], y_all[:train_size], x_all[-test_size:], y_all[-test_size:] |
|
|
|
|
| def make_normalized_loaders(cfg: dict[str, Any]): |
| x_train, y_train, x_test, y_test = load_pipe_tensors(cfg) |
| axes = tuple(int(v) for v in cfg["data"].get("normalizer_axes", [0, 1])) |
| x_normalizer = UnitTransformer(x_train, axes=axes) |
| y_normalizer = UnitTransformer(y_train, axes=axes) |
| x_train_n = x_normalizer.encode(x_train) |
| x_test_n = x_normalizer.encode(x_test) |
| y_train_n = y_normalizer.encode(y_train) |
| y_test_n = y_normalizer.encode(y_test) |
|
|
| batch_size = int(cfg["training"]["batch_size"]) |
| loader_kwargs = { |
| "batch_size": batch_size, |
| "num_workers": int(cfg["data"].get("num_workers", 0)), |
| "pin_memory": bool(cfg["data"].get("pin_memory", False)), |
| } |
| train_ds = torch.utils.data.TensorDataset(x_train_n, y_train_n) |
| test_ds = torch.utils.data.TensorDataset(x_test_n, y_test_n) |
| train_loader = torch.utils.data.DataLoader(train_ds, shuffle=True, **loader_kwargs) |
| test_loader = torch.utils.data.DataLoader(test_ds, shuffle=False, **loader_kwargs) |
| return train_loader, test_loader, x_normalizer, y_normalizer |
|
|
|
|
| def build_model(cfg: dict[str, Any]) -> HOFNOPipe: |
| m = cfg["model"] |
| return HOFNOPipe( |
| input_channels=int(m["input_channels"]), |
| output_channels=int(m["output_channels"]), |
| order=int(m["order"]), |
| layers=int(m["layers"]), |
| width=int(m["width"]), |
| modes=tuple(int(v) for v in m["modes"]), |
| add_unit_grid=bool(m.get("add_unit_grid", True)), |
| mlp_expansion=int(m.get("mlp_expansion", 4)), |
| mlp_dropout=float(m.get("mlp_dropout", 0.0)), |
| drop_path_rate=float(m.get("drop_path_rate", 0.0)), |
| ) |
|
|
|
|
| def relative_l2_per_sample(pred: torch.Tensor, target: torch.Tensor) -> torch.Tensor: |
| pred_f = pred.reshape(pred.shape[0], -1) |
| target_f = target.reshape(target.shape[0], -1) |
| denom = torch.linalg.vector_norm(target_f, dim=1).clamp_min(1e-12) |
| return torch.linalg.vector_norm(pred_f - target_f, dim=1) / denom |
|
|
|
|
| def relative_l2_sum(pred: torch.Tensor, target: torch.Tensor) -> torch.Tensor: |
| return relative_l2_per_sample(pred, target).sum() |
|
|
|
|
| def save_json(path: Path, payload: dict[str, Any]) -> None: |
| path.parent.mkdir(parents=True, exist_ok=True) |
| with path.open("w", encoding="utf-8") as handle: |
| json.dump(payload, handle, indent=2) |
|
|
|
|
| __all__ = [ |
| "ROOT", |
| "UnitTransformer", |
| "build_model", |
| "count_parameters", |
| "load_config", |
| "make_normalized_loaders", |
| "relative_l2_per_sample", |
| "relative_l2_sum", |
| "resolve", |
| "save_json", |
| "select_device", |
| "set_seed", |
| ] |
|
|
