File size: 4,643 Bytes
5ec1901 | 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 133 | """v1.1 input/output encoding for ReservoirCNN.
FROZEN — these constants must match the values baked into
`weights/geoforce_cnn_v1.1.pt` exactly. Any change invalidates the weights.
v1.1 input channels (6, 32, 32):
0: initial temperature field (per-cell, normalized to [0, 1] on [T_MIN, T_MAX])
1: log10 permeability (scalar broadcast, normalized on [LOG_PERM_MIN, LOG_PERM_MAX])
2: well mask (1.0 at well cells, decaying to neighbors)
3: base pressure (scalar broadcast, normalized on [BASE_P_MIN, BASE_P_MAX])
4: porosity (scalar broadcast, normalized on [POR_MIN, POR_MAX])
5: depth (scalar broadcast, normalized on [DEPTH_MIN, DEPTH_MAX])
Output channels (10, 32, 32):
0-4: temperature at 5 timesteps, sigmoid-normalized on [T_MIN, T_MAX] in degrees C
5-9: pressure at 5 timesteps, sigmoid-normalized on [P_MIN, P_MAX] in Pa
"""
from __future__ import annotations
import numpy as np
import torch
GRID_SIZE = 32
N_TIME_STEPS = 5
# Frozen v1.1 normalization constants (embedded in the checkpoint).
NORMALIZATION: dict[str, float] = {
"T_MIN": 25.0,
"T_MAX": 350.0,
"P_MIN": 1.0e5,
"P_MAX": 5.0e7,
"LOG_PERM_MIN": -16.0,
"LOG_PERM_MAX": -12.0,
"POR_MIN": 0.01,
"POR_MAX": 0.15,
"DEPTH_MIN": 800.0,
"DEPTH_MAX": 2500.0,
"BASE_T_MIN": 180.0,
"BASE_T_MAX": 320.0,
"BASE_P_MIN": 5.0e6,
"BASE_P_MAX": 2.5e7,
}
def _norm(value: float, lo: float, hi: float) -> float:
return float(np.clip((value - lo) / (hi - lo), 0.0, 1.0))
def _build_well_mask(
well_locations: list[tuple[int, int]],
grid_size: int = GRID_SIZE,
) -> np.ndarray:
"""Match the training well-mask encoding exactly (1.0 at well, decay to neighbors)."""
mask = np.zeros((grid_size, grid_size), dtype=np.float32)
for wr, wc in well_locations:
wr, wc = int(wr), int(wc)
if not (0 <= wr < grid_size and 0 <= wc < grid_size):
continue
mask[wr, wc] = 1.0
for dr in range(-1, 2):
for dc in range(-1, 2):
r, c = wr + dr, wc + dc
if 0 <= r < grid_size and 0 <= c < grid_size:
dist = abs(dr) + abs(dc)
mask[r, c] = max(mask[r, c], 1.0 / (1.0 + dist))
return mask
def build_input_tensor(
*,
initial_temperature: np.ndarray,
log_permeability: float,
well_locations: list[tuple[int, int]],
base_pressure: float,
porosity: float,
depth: float,
grid_size: int = GRID_SIZE,
) -> torch.Tensor:
"""Assemble the 6-channel input tensor matching v1.1 training.
Args:
initial_temperature: 2D array of shape (32, 32) in degrees Celsius.
log_permeability: scalar log10 permeability (typical -16 to -12).
well_locations: list of (row, col) integer tuples.
base_pressure: scalar pressure at reservoir base (Pa).
porosity: dimensionless (0.01 to 0.15).
depth: reservoir depth (m, 800 to 2500).
grid_size: grid cells per side (default 32, must match weights).
Returns:
torch.Tensor of shape (1, 6, grid_size, grid_size), dtype float32.
"""
if initial_temperature.shape != (grid_size, grid_size):
msg = (
f"initial_temperature shape {initial_temperature.shape} "
f"does not match grid_size={grid_size}"
)
raise ValueError(msg)
n = NORMALIZATION
arr = np.zeros((6, grid_size, grid_size), dtype=np.float32)
t0_norm = (initial_temperature - n["T_MIN"]) / (n["T_MAX"] - n["T_MIN"])
arr[0] = np.clip(t0_norm, 0.0, 1.0).astype(np.float32)
arr[1] = _norm(log_permeability, n["LOG_PERM_MIN"], n["LOG_PERM_MAX"])
arr[2] = _build_well_mask(well_locations, grid_size)
arr[3] = _norm(base_pressure, n["BASE_P_MIN"], n["BASE_P_MAX"])
arr[4] = _norm(porosity, n["POR_MIN"], n["POR_MAX"])
arr[5] = _norm(depth, n["DEPTH_MIN"], n["DEPTH_MAX"])
return torch.from_numpy(arr).unsqueeze(0)
def denormalize_output(out: torch.Tensor) -> dict[str, np.ndarray]:
"""Convert model output from [0,1] to physical units.
Args:
out: tensor of shape (batch, 10, H, W) with values in [0, 1].
Returns:
dict with keys 'temperature' (shape (batch, 5, H, W), deg C) and
'pressure' (shape (batch, 5, H, W), Pa).
"""
n = NORMALIZATION
arr = out.detach().cpu().numpy()
t_norm = arr[:, :N_TIME_STEPS]
p_norm = arr[:, N_TIME_STEPS:]
temperature = t_norm * (n["T_MAX"] - n["T_MIN"]) + n["T_MIN"]
pressure = p_norm * (n["P_MAX"] - n["P_MIN"]) + n["P_MIN"]
return {"temperature": temperature, "pressure": pressure}
|