"""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] # Tensor axis 1 is x, axis 2 is y. 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)