ThermalBreakthrough/examples/example_utils.py at main · DigiPorFlow/ThermalBreakthrough · GitHub

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
"""Helpers to load the bundled example cases.

Each examples/case_XXXX.npz file is a self-contained evaluation case extracted
from the Zenodo dataset (see the README): the cooling-fraction trajectory, the
static node and edge features, the cached AMG operators, and the case scalars.
Nothing else is needed to run the surrogate on it.
"""
from __future__ import annotations

from pathlib import Path

import matplotlib.pyplot as plt
import numpy as np
import torch
from scipy.sparse import coo_matrix

# House style for temperature and error maps (matches the paper figures).
TEMP_CMAP = "jet"
ERR_CMAP = "PuBu"


def apply_style():
    try:
        import scienceplots  # noqa: F401
        plt.style.use(["science", "ieee", "no-latex"])
    except Exception as e:
        print(f"[style] scienceplots unavailable ({e}); default style")
    plt.rcParams.update({
        "text.usetex": False,
        "font.size": 14, "axes.labelsize": 14, "axes.titlesize": 14,
        "xtick.labelsize": 11, "ytick.labelsize": 11,
        "legend.fontsize": 10, "legend.title_fontsize": 12,
        "figure.dpi": 150, "savefig.dpi": 200,
    })


def pct_err(pred_K, tgt_K):
    """Percent error map, as used in the paper's field-error figures."""
    return 100.0 * np.abs(pred_K - tgt_K) / np.maximum(np.abs(tgt_K), 1e-6)


def load_example(path: str | Path) -> dict:
    """Load one bundled case into the item-dict format used by
    tbt.eval.rollout.rollout_case."""
    z = np.load(path)
    P = coo_matrix((z["P_data"], (z["P_row"], z["P_col"])),
                   shape=tuple(z["P_shape"])).tocsr()
    PP = coo_matrix((z["PP_data"], (z["PP_row"], z["PP_col"])),
                    shape=tuple(z["PP_shape"])).tocsr()
    item = {
        "case_id": int(z["case_id"]),
        "theta_norm": torch.from_numpy(z["theta"]).float(),
        "static_norm": torch.from_numpy(z["static_norm"]).float(),
        "edge_index": torch.from_numpy(z["edge_index"]).long(),
        "edge_attr": torch.from_numpy(z["edge_attr"]).float(),
        "coarse_edge_index": torch.from_numpy(z["coarse_edge_index"]).long(),
        "coarse_edge_attr": torch.from_numpy(z["coarse_edge_attr"]).float(),
        "P": P,
        "P_prolong": PP,
        "well_indices": torch.from_numpy(z["well_indices"]).long(),
        "T_inj_K": torch.tensor(float(z["T_inj_K"])),
        "T_res_K": torch.tensor(float(z["T_res_K"])),
        "Q_kg_s": torch.tensor(float(z["Q_kg_s"])),
    }
    meta = {
        "years": z["years"].astype(np.float64),
        "coords": z["coords"],
        "inj_xy": z["inj_xy"],
        "prod_xy": z["prod_xy"],
        "T_inj": float(z["T_inj_K"]),
        "T_res": float(z["T_res_K"]),
    }
    return item, meta


def theta_to_K(theta, T_inj: float, T_res: float):
    """Cooling fraction to absolute temperature."""
    return T_res - theta * (T_res - T_inj)


def raster(values, coords):
    """Flat node values to a (ny, nx) image plus imshow extent."""
    xs, ys = np.unique(coords[:, 0]), np.unique(coords[:, 1])
    sort_idx = np.lexsort((coords[:, 0], coords[:, 1]))
    extent = [float(xs[0] - 30), float(xs[-1] + 30),
              float(ys[0] - 30), float(ys[-1] + 30)]
    return np.asarray(values)[sort_idx].reshape(ys.size, xs.size), extent