| """Shared configuration, fixed-data protocol, analytical Q1 energy and metrics.""" |
| from __future__ import annotations |
|
|
| import argparse |
| import hashlib |
| import json |
| from pathlib import Path |
| import random |
| import sys |
|
|
| import numpy as np |
| import torch |
| from torch import nn |
| import yaml |
|
|
| ROOT = Path(__file__).resolve().parents[1] |
| sys.path.insert(0, str(ROOT)) |
| from model.VINO_Darcy import VINODarcy |
|
|
|
|
| def arguments(description): |
| parser = argparse.ArgumentParser(description=description) |
| parser.add_argument("--config", type=Path, default=ROOT / "config/config.yaml") |
| parser.add_argument("--data", type=Path, help="Override data.path explicitly") |
| parser.add_argument("--results-dir", type=Path) |
| parser.add_argument("--device", choices=["cuda", "cpu"]) |
| return parser |
|
|
|
|
| def configuration(args): |
| cfg = yaml.safe_load(args.config.read_text()) |
| if args.data is not None: |
| cfg["data"]["path"] = str(args.data.resolve()) |
| if args.results_dir is not None: |
| cfg["output"]["results_dir"] = str(args.results_dir.resolve()) |
| if args.device is not None: |
| cfg["training"]["device"] = args.device |
| if cfg["data"]["n_train"] != 1000 or cfg["data"]["n_test"] != 100 or cfg["data"]["resolution"] != 64: |
| raise ValueError("This reproduction requires the original 1000/100 samples and 64x64 grid") |
| if cfg["training"]["early_stopping"]: |
| raise ValueError("Early stopping violates the requested full 500 epochs") |
| return cfg |
|
|
|
|
| def output_path(value): |
| path = Path(value) |
| return path if path.is_absolute() else ROOT / path |
|
|
|
|
| def write_json(path, payload): |
| path = Path(path) |
| path.parent.mkdir(parents=True, exist_ok=True) |
| temporary = path.with_suffix(path.suffix + ".tmp") |
| temporary.write_text(json.dumps(payload, ensure_ascii=False, indent=2, allow_nan=False)) |
| temporary.replace(path) |
|
|
|
|
| def save_checkpoint(path, payload): |
| path = Path(path) |
| path.parent.mkdir(parents=True, exist_ok=True) |
| temporary = path.with_suffix(".tmp") |
| torch.save(payload, temporary) |
| temporary.replace(path) |
|
|
|
|
| def initialize(cfg): |
| seed = cfg["seed"] |
| random.seed(seed) |
| np.random.seed(seed) |
| torch.manual_seed(seed) |
| torch.set_num_threads(cfg["training"]["cpu_threads"]) |
| device = torch.device(cfg["training"]["device"]) |
| if device.type == "cuda": |
| if not torch.cuda.is_available(): |
| raise RuntimeError("Configured accelerator unavailable; execute inside the allocated compute job") |
| torch.cuda.manual_seed_all(seed) |
| torch.backends.cudnn.benchmark = False |
| torch.backends.cudnn.deterministic = True |
| return device |
|
|
|
|
| def file_hash(path): |
| digest = hashlib.sha256() |
| with Path(path).open("rb") as stream: |
| for block in iter(lambda: stream.read(1024 * 1024), b""): |
| digest.update(block) |
| return digest.hexdigest() |
|
|
|
|
| def neighboring_average(field): |
| """Official missing-detail convention: positive neighbors, repeat last row/column.""" |
| extended = torch.cat((field, field[:, -1:]), dim=1) |
| extended = torch.cat((extended, extended[:, :, -1:]), dim=2) |
| return (extended[:, :-1, :-1] + extended[:, 1:, :-1] |
| + extended[:, :-1, 1:] + extended[:, 1:, 1:]) * 0.25 |
|
|
|
|
| def read_data(cfg): |
| path = Path(cfg["data"]["path"]) |
| if not path.is_file(): |
| raise FileNotFoundError(f"Required author dataset missing: {path}; no replacement is permitted") |
| payload = torch.load(path, map_location="cpu", weights_only=True) |
| expected = {"F_train": 1000, "U_train": 1000, "F_test": 100, "U_test": 100} |
| fields = {} |
| metadata = {"path": str(path.resolve()), "sha256": file_hash(path), "fields": {}} |
| for name, count in expected.items(): |
| raw = payload[name] |
| if raw.ndim == 3: |
| raw = raw.unsqueeze(-1) |
| if tuple(raw.shape) != (count, 64, 64, 1): |
| raise ValueError(f"{name}: expected {(count,64,64,1)}, got {tuple(raw.shape)}") |
| if not torch.isfinite(raw).all(): |
| raise ValueError(f"Nonfinite data in {name}") |
| value = raw.float().contiguous() |
| fields[name] = value |
| metadata["fields"][name] = {"shape": list(raw.shape), "dtype": str(raw.dtype), "min": value.min().item(), "max": value.max().item()} |
| for name in ("F_train", "F_test"): |
| if fields[name].min() <= 0: |
| raise ValueError("Permeability must be strictly positive") |
| metadata["fields"][name]["raw_unique_values"] = torch.unique(fields[name]).tolist() |
| fields["P_train"] = neighboring_average(fields["F_train"]) |
| fields["P_test"] = neighboring_average(fields["F_test"]) |
| metadata["preprocessing"] = cfg["data"]["preprocessing"] |
| metadata["training_labels_used_in_optimization"] = False |
| return fields, metadata |
|
|
|
|
| def element_corners(field): |
| field = field[..., 0] |
| |
| return torch.stack((field[:, :-1, :-1], field[:, 1:, :-1], |
| field[:, 1:, 1:], field[:, :-1, 1:]), dim=-1) |
|
|
|
|
| def analytical_coefficient_tensor(hx, hy): |
| """C[k,i,j] = integral N_k (grad N_i dot grad N_j) by polynomial moments. |
| |
| On [0,1], integral L_k=1/2, integral L_a L_b L_c is 1/4 |
| for equal indices and 1/12 otherwise. No numerical quadrature is used. |
| """ |
| nodes = ((0, 0), (1, 0), (1, 1), (0, 1)) |
| derivative = (-1, 1) |
| tensor = torch.empty(4, 4, 4, dtype=torch.float64) |
| def triple(a, b, c): |
| return 0.25 if a == b == c else 1 / 12 |
| for k, (kx, ky) in enumerate(nodes): |
| for i, (ix, iy) in enumerate(nodes): |
| for j, (jx, jy) in enumerate(nodes): |
| tensor[k, i, j] = ( |
| hy / hx * 0.5 * derivative[ix] * derivative[jx] * triple(ky, iy, jy) |
| + hx / hy * triple(kx, ix, jx) * 0.5 * derivative[iy] * derivative[jy] |
| ) |
| return tensor |
|
|
|
|
| class DarcyEnergy(nn.Module): |
| def __init__(self, resolution=64, source=1.0): |
| super().__init__() |
| self.h = 1 / (resolution - 1) |
| self.source = source |
| self.register_buffer("coefficients", analytical_coefficient_tensor(self.h, self.h).float()) |
|
|
| def stiffness(self, permeability): |
| return torch.einsum("bxyk,kij->bxyij", element_corners(permeability), self.coefficients) |
|
|
| def forward(self, solution, permeability=None, stiffness=None, reduction="mean"): |
| if stiffness is None: |
| stiffness = self.stiffness(permeability) |
| nodal = element_corners(solution) |
| gradient_energy = (nodal * torch.matmul(stiffness, nodal.unsqueeze(-1)).squeeze(-1)).sum(dim=-1) * 0.5 |
| source_work = nodal.sum(dim=-1) * (self.source * self.h**2 / 4) |
| energies = (gradient_energy - source_work).sum(dim=(1, 2)) |
| return energies.mean() if reduction == "mean" else energies |
|
|
|
|
| def relative_l2(prediction, reference): |
| error = (prediction - reference).flatten(1).norm(dim=1) |
| scale = reference.flatten(1).norm(dim=1) |
| if torch.any(scale == 0): |
| raise ValueError("Relative L2 undefined for a zero reference") |
| return error / scale |
|
|
|
|
| def error_statistics(errors): |
| values = errors.detach().double().cpu() |
| return {"n_test": values.numel(), "relative_l2_mean": values.mean().item(), |
| "relative_l2_std": values.std(correction=1).item(), |
| "relative_l2_population_std": values.std(correction=0).item(), |
| "relative_l2_median": values.median().item(), |
| "relative_l2_min": values.min().item(), "relative_l2_max": values.max().item(), |
| "relative_l2_mean_percent": values.mean().item() * 100, |
| "relative_l2_std_percent": values.std(correction=1).item() * 100} |
|
|
|
|
| @torch.no_grad() |
| def evaluate(model, permeability, reference, energy, batch_size): |
| model.eval() |
| predictions, errors, energies = [], [], [] |
| for start in range(0, len(permeability), batch_size): |
| p, target = permeability[start:start + batch_size], reference[start:start + batch_size] |
| output = model(p) |
| predictions.append(output.cpu()) |
| errors.append(relative_l2(output, target).cpu()) |
| energies.append(energy(output, p, reduction="none").cpu()) |
| return torch.cat(predictions), torch.cat(errors), torch.cat(energies) |
|
|
|
|
| def build_model(cfg, device): |
| return VINODarcy(**cfg["model"]).to(device) |
|
|
