scripts/common.py

"""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)