| """Configuration, supplied MAT data, and portable checkpoint utilities.""" |
|
|
| import hashlib |
| import json |
| import os |
| import random |
| from pathlib import Path |
|
|
| import numpy as np |
| import scipy.io |
| import torch |
| import yaml |
| from torch.nn import functional as F |
|
|
| ROOT = Path(__file__).resolve().parents[1] |
|
|
|
|
| def package_path(value): |
| path = Path(value).expanduser() |
| return path if path.is_absolute() else ROOT / path |
|
|
|
|
| def load_config(path): |
| with open(package_path(path)) as stream: |
| return yaml.safe_load(stream) |
|
|
|
|
| def save_json(path, value): |
| path = Path(path) |
| path.parent.mkdir(parents=True, exist_ok=True) |
| temp = path.with_name(path.name + ".tmp") |
| temp.write_text(json.dumps(value, ensure_ascii=False, indent=2, allow_nan=False)) |
| temp.replace(path) |
|
|
|
|
| def seed_everything(seed): |
| random.seed(seed) |
| np.random.seed(seed) |
| torch.manual_seed(seed) |
| torch.cuda.manual_seed_all(seed) |
|
|
|
|
| def compute_device(name): |
| device = torch.device(name) |
| if device.type == "cuda": |
| if not os.environ.get("SLURM_JOB_ID"): |
| raise RuntimeError("Accelerator execution requires a Slurm allocation") |
| if not torch.cuda.is_available(): |
| raise RuntimeError("No accelerator visible inside the Slurm job") |
| return device |
|
|
|
|
| def read_fields(config): |
| data = config["data"] |
| path = package_path(data["directory"]) / data["filename"] |
| arrays = scipy.io.loadmat(path, variable_names=[data["input_key"], data["target_key"]]) |
| fields = [] |
| for key in [data["input_key"], data["target_key"]]: |
| array = arrays[key] |
| if list(array.shape) != data["expected_shape"] or array.dtype != np.float64: |
| raise ValueError(f"Unexpected supplied data {key}: {array.shape}, {array.dtype}") |
| if not np.isfinite(array).all(): |
| raise ValueError(f"Nonfinite values in {key}") |
| tensor = torch.from_numpy(array.astype(np.float32)) |
| |
| |
| tensor = F.interpolate(tensor[:, None], size=data["resolution"], |
| mode=data["interpolation"], align_corners=False) |
| fields.append(tensor.transpose(1, 2).contiguous()) |
| return fields |
|
|
|
|
| def split_indices(config): |
| data = config["data"] |
| permutation = torch.randperm(data["expected_shape"][0], generator=torch.Generator().manual_seed(config["seed"])) |
| n, m = data["n_train"], data["n_val"] |
| if n + m != len(permutation): |
| raise ValueError("The formal experiment must use all 2048 supplied samples") |
| return permutation[:n], permutation[n:n + m] |
|
|
|
|
| def fit_normalizers(initial, final, train_ids, eps): |
| result = {"eps": eps} |
| for prefix, field in [("a", initial), ("u", final)]: |
| training = field[train_ids] |
| result[prefix + "_mean"] = training.mean() |
| result[prefix + "_std"] = training.std(correction=1) |
| return result |
|
|
|
|
| def encode(field, stats, prefix): |
| return (field - stats[prefix + "_mean"].to(field)) / (stats[prefix + "_std"].to(field) + stats["eps"]) |
|
|
|
|
| def decode(field, stats, prefix="u"): |
| return field * (stats[prefix + "_std"].to(field) + stats["eps"]) + stats[prefix + "_mean"].to(field) |
|
|
|
|
| def make_inputs(initial, stats): |
| physical = encode(initial, stats, "a") |
| grid = torch.linspace(0, 1, initial.shape[1], dtype=initial.dtype).view(1, -1, 1) |
| return torch.cat([physical, grid.expand(initial.shape[0], -1, -1)], dim=-1) |
|
|
|
|
| def relative_l2(prediction, target): |
| difference = (prediction - target).flatten(1) |
| return torch.linalg.vector_norm(difference, dim=1) / (torch.linalg.vector_norm(target.flatten(1), dim=1) + 1e-12) |
|
|
|
|
| def atomic_checkpoint(path, state): |
| path = package_path(path) |
| path.parent.mkdir(parents=True, exist_ok=True) |
| temp = path.with_name(path.name + ".tmp") |
| torch.save(state, temp) |
| temp.replace(path) |
|
|
|
|
| def load_checkpoint(path): |
| |
| state = torch.load(package_path(path), map_location="cpu", weights_only=True) |
| if state.get("format") != "lino-burgers-v1": |
| raise ValueError("Unsupported LiNO checkpoint") |
| return state |
|
|
|
|
| def file_digest(path): |
| digest = hashlib.sha256() |
| with open(path, "rb") as stream: |
| for block in iter(lambda: stream.read(8 * 1024 * 1024), b""): |
| digest.update(block) |
| return digest.hexdigest() |
|
|
