File size: 4,486 Bytes
8f89e30 | 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 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 | """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))
# Paper leaves resizing unspecified: author's linear interpolation,
# align_corners=False, evaluates at source centers 15.5,47.5,... .
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):
# Tensor/primitive-only checkpoint; no pickled custom class or path equality.
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()
|