Port GeoForce v1.1 CNN surrogate; smoke test green

Adds surrogate/ with the ReservoirCNN architecture, frozen v1.1
normalization (6-channel input, 10-channel output), and a cached
predict() wrapper. Weights (248 KB, 57,802 params) copied from
ForceX-AI model registry and verified via tests/test_surrogate_smoke.py
(5/5 pass, <2s total).

Corrects surrogate-operator.md: output P range is [1e5, 5e7] Pa
(not [5e6, 25e6]); channel 0 is the full initial-T field, not a
scalar base_T. pyproject.toml pins Python 3.11+, CPU torch, iapws,
scipy, and dev/agent/app optional-deps groups.

Day 1 Morning milestone: done.

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>

.claude/agents/surrogate-operator.md

@@ -14,54 +14,72 @@ You are the interface to the deployed v1.1 CNN surrogate. You load the weights o
 1. **Do not modify the weights file.** Treat `surrogate/weights/geoforce_cnn_v1.1.pt` as immutable.
 2. **Do not change the 6-channel input encoding.** Normalization constants are frozen from v1.1 training. Any change invalidates the model.
 3. **Do not retrain.** If accuracy is insufficient, report it; don't "fix" by retraining.
-4. **Output shape is always `(10, 32, 32)`** — channels 0–4 are temperature at years 4/8/12/16/20, channels 5–9 are pressure at the same times.
+4. **Output shape is always `(10, 32, 32)`** — channels 0–4 are temperature at 5 timesteps, channels 5–9 are pressure at the same 5 timesteps.
 
 ## Frozen Normalization Constants (must match v1.1 training exactly)
 
+Source of truth: `surrogate/encoding.py` `NORMALIZATION` dict, also embedded in the checkpoint.
+
 ```python
+# OUTPUT ranges (used to de-normalize both predictions AND input channel 0)
 T_MIN, T_MAX = 25.0, 350.0              # Celsius
-P_MIN, P_MAX = 5e6, 25e6                # Pascal
+P_MIN, P_MAX = 1.0e5, 5.0e7             # Pascal (0.1 to 50 MPa)
+
+# INPUT channel ranges (only used to normalize)
 LOG_PERM_MIN, LOG_PERM_MAX = -16.0, -12.0
-POROSITY_MIN, POROSITY_MAX = 0.01, 0.15
-DEPTH_MIN, DEPTH_MAX = 800.0, 2500.0    # meters
+POR_MIN, POR_MAX           =   0.01, 0.15
+DEPTH_MIN, DEPTH_MAX       = 800.0, 2500.0    # meters
+BASE_P_MIN, BASE_P_MAX     =   5.0e6, 2.5e7   # Pascal (5 to 25 MPa)
+# (BASE_T_MIN/MAX = 180/320 exist in the checkpoint but v1.1 does not use them —
+#  v1.1 consumes the full initial temperature FIELD as channel 0, not a scalar.)
 ```
 
 ## Input Channels (6, in order)
 
 | Ch | Content | Normalization |
 |---|---|---|
-| 0 | Initial temperature field | `(T - T_MIN) / (T_MAX - T_MIN)` |
-| 1 | Log₁₀ permeability | `(log_k - LOG_PERM_MIN) / (LOG_PERM_MAX - LOG_PERM_MIN)` |
-| 2 | Well mask | Gaussian decay around well cells |
-| 3 | Base pressure (scalar broadcast) | `(P - P_MIN) / (P_MAX - P_MIN)` |
-| 4 | Porosity (scalar broadcast) | `(φ - POROSITY_MIN) / (POROSITY_MAX - POROSITY_MIN)` |
-| 5 | Depth (scalar broadcast) | `(z - DEPTH_MIN) / (DEPTH_MAX - DEPTH_MIN)` |
+| 0 | **Initial temperature field** (full 32×32, degrees C) | `(T - T_MIN) / (T_MAX - T_MIN)` |
+| 1 | Log₁₀ permeability (scalar, broadcast) | `(log_k - LOG_PERM_MIN) / (LOG_PERM_MAX - LOG_PERM_MIN)` |
+| 2 | Well mask | 1.0 at well cell; neighbors decay as `1/(1 + manhattan_distance)` |
+| 3 | Base pressure (scalar, broadcast) | `(P - BASE_P_MIN) / (BASE_P_MAX - BASE_P_MIN)` |
+| 4 | Porosity (scalar, broadcast) | `(φ - POR_MIN) / (POR_MAX - POR_MIN)` |
+| 5 | Depth (scalar, broadcast) | `(z - DEPTH_MIN) / (DEPTH_MAX - DEPTH_MIN)` |
+
+All channels clipped to `[0, 1]` after normalization.
 
 ## Output De-normalization
 
-Channels 0–4: `T_out_C = ch * (T_MAX - T_MIN) + T_MIN`
-Channels 5–9: `P_out_Pa = ch * (P_MAX - P_MIN) + P_MIN`
+- Channels 0–4: `T_out_C = ch * (T_MAX - T_MIN) + T_MIN`
+- Channels 5–9: `P_out_Pa = ch * (P_MAX - P_MIN) + P_MIN`
 
 ## Runtime
 
-Called from `tools/predict_surrogate.py`. Typical:
+Called via:
 
 ```python
-from surrogate.infer import predict
-out = predict(scenario_dict)
-# out["T"] shape (5, 32, 32) °C
-# out["P"] shape (5, 32, 32) Pa
-# out["t_ms"] inference time
+from surrogate import predict
+
+out = predict(
+    initial_temperature=T0,           # shape (32, 32), degrees C
+    log_permeability=-13.5,
+    well_locations=[(16, 16)],
+    base_pressure=1.5e7,
+    porosity=0.08,
+    depth=1800.0,
+)
+# out["temperature"] shape (5, 32, 32) °C
+# out["pressure"]    shape (5, 32, 32) Pa
 ```
 
-Target: < 5 ms per call on CPU. If exceeded, flag to planner.
+Model weights are cached after first `predict()` / `load_model()` call.
+Target inference: < 5 ms/call on CPU (the smoke test budget is 50 ms to stay slack-tolerant).
 
 ## Sanity Checks (run on every output)
 
 - `T_out` in `[T_MIN - 2, T_MAX + 2]` (small tolerance for CNN drift)
-- `P_out` in `[P_MIN - 1e5, P_MAX + 1e5]`
+- `P_out` in `[P_MIN - 1e5, P_MAX + 1e5]` — i.e., roughly `[0, 50.1 MPa]`
 - No NaN / Inf
-- Shape exactly `(10, 32, 32)`
+- Shape exactly `(5, 32, 32)` per field
 
 If any check fails, return an error to planner — **do not silently clamp**.
 

PROGRESS.md

@@ -18,17 +18,18 @@ Live task tracking. Updated at the end of every work block. Source of truth = th
 - [x] `.claude/` skills defined (6)
 - [x] `.claude/` commands defined (4)
 - [x] `.claude/settings.json` — model pinned to `claude-opus-4-7`
-- [ ] Fresh `.venv` + `pyproject.toml`
-- [ ] `ANTHROPIC_API_KEY` verified in shell
-- [ ] Registration + Discord acceptance email found
+- [x] Fresh `.venv` + `pyproject.toml`
+- [x] `ANTHROPIC_API_KEY` in `.env` (gitignored) + round-trip tested
+- [x] `.gitignore` protects secrets
+- [~] Registration + Discord acceptance email — deferred per user (2026-04-23)
 
 ## Day 1 — Build both engines
 
-### Morning (3h)
-- [ ] Copy `geoforce_cnn_v1.1.pt` to `surrogate/weights/`
-- [ ] Port `ReservoirCNN` class + encoding to `surrogate/`
-- [ ] `tests/test_surrogate_smoke.py` green
-- [ ] Claude scaffolding discoverable (subagents load)
+### Morning (3h) — COMPLETE 2026-04-23
+- [x] Copy `geoforce_cnn_v1.1.pt` to `surrogate/weights/` (248,595 bytes, from ForceX-AI/products/model_registry)
+- [x] Port `ReservoirCNN` class + encoding to `surrogate/` (model.py, encoding.py, predict.py)
+- [x] `tests/test_surrogate_smoke.py` green — 5/5 pass, 1.7s total
+- [x] Claude scaffolding discoverable (subagents + skills load; `surrogate-operator.md` corrected to match real v1.1 encoding)
 
 ### Afternoon (4h) — Agent team builds GeoForce-Solver
 - [ ] `solver/properties.py` — IAPWS-IF97 wrappers

pyproject.toml

@@ -0,0 +1,55 @@
+[build-system]
+requires = ["setuptools>=68", "wheel"]
+build-backend = "setuptools.build_meta"
+
+[project]
+name = "geoforce-cchackathon"
+version = "0.1.0"
+description = "Opus 4.7 agents orchestrate GeoForce-Solver + v1.1 CNN surrogate for Indonesian geothermal engineering questions"
+readme = "README.md"
+requires-python = ">=3.11"
+license = { text = "MIT" }
+authors = [{ name = "Robi Dany Riupassa" }]
+
+dependencies = [
+    "numpy>=1.26",
+    "scipy>=1.12",
+    "torch>=2.2",
+    "iapws>=1.5.3",
+    "matplotlib>=3.8",
+    "pyyaml>=6.0",
+    "python-dotenv>=1.0",
+]
+
+[project.optional-dependencies]
+agent = [
+    "claude-agent-sdk>=0.1.0",
+    "anthropic>=0.40.0",
+    "fastapi>=0.110",
+    "uvicorn[standard]>=0.29",
+    "sse-starlette>=2.0",
+]
+app = [
+    "streamlit>=1.33",
+]
+dev = [
+    "pytest>=8.0",
+    "pytest-cov>=4.1",
+    "ruff>=0.3",
+    "mypy>=1.9",
+]
+
+[tool.setuptools.packages.find]
+include = ["surrogate*", "solver*", "agent*", "tools*"]
+
+[tool.pytest.ini_options]
+testpaths = ["tests"]
+addopts = "-ra -q"
+
+[tool.ruff]
+line-length = 100
+target-version = "py311"
+
+[tool.ruff.lint]
+select = ["E", "F", "I", "N", "UP", "B", "SIM"]
+ignore = ["E501"]

surrogate/__init__.py

@@ -0,0 +1,20 @@
+"""GeoForce v1.1 CNN surrogate — ported from ForceX-AI for GeoForce-CCHackathon."""
+
+from __future__ import annotations
+
+from .encoding import (
+    NORMALIZATION,
+    build_input_tensor,
+    denormalize_output,
+)
+from .model import ReservoirCNN
+from .predict import load_model, predict
+
+__all__ = [
+    "NORMALIZATION",
+    "ReservoirCNN",
+    "build_input_tensor",
+    "denormalize_output",
+    "load_model",
+    "predict",
+]

surrogate/encoding.py

@@ -0,0 +1,132 @@
+"""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}

surrogate/model.py

@@ -0,0 +1,108 @@
+"""ReservoirCNN architecture for GeoForce v1.1.
+
+Ported verbatim from ForceX-AI's training code. The state dict of
+`weights/geoforce_cnn_v1.1.pt` loads into this class without renaming.
+
+Architecture:
+    encoder:  ConvBlock(6 -> 32) -> ConvBlock(32 -> 32)
+    middle:   ResidualBlock(32) x 2
+    decoder:  ConvBlock(32 -> 32) -> Conv2d(32 -> 10, 1x1) -> Sigmoid
+
+57,802 parameters. ~3 ms per inference on CPU.
+"""
+
+from __future__ import annotations
+
+import torch
+import torch.nn as nn
+
+
+class ConvBlock(nn.Module):
+    """Conv2d -> BatchNorm -> ReLU."""
+
+    def __init__(self, in_channels: int, out_channels: int, kernel_size: int = 3) -> None:
+        super().__init__()
+        self.conv = nn.Conv2d(
+            in_channels,
+            out_channels,
+            kernel_size,
+            padding=kernel_size // 2,
+            bias=False,
+        )
+        self.bn = nn.BatchNorm2d(out_channels)
+        self.relu = nn.ReLU(inplace=True)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.relu(self.bn(self.conv(x)))
+
+
+class ResidualBlock(nn.Module):
+    """Residual block with two 3x3 convs and a skip connection."""
+
+    def __init__(self, channels: int) -> None:
+        super().__init__()
+        self.conv1 = nn.Conv2d(channels, channels, 3, padding=1, bias=False)
+        self.bn1 = nn.BatchNorm2d(channels)
+        self.conv2 = nn.Conv2d(channels, channels, 3, padding=1, bias=False)
+        self.bn2 = nn.BatchNorm2d(channels)
+        self.relu = nn.ReLU(inplace=True)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        residual = x
+        out = self.relu(self.bn1(self.conv1(x)))
+        out = self.bn2(self.conv2(out))
+        out = out + residual
+        return self.relu(out)
+
+
+class ReservoirCNN(nn.Module):
+    """CNN surrogate for 2D geothermal reservoir T and P prediction.
+
+    Input shape:  (batch, 6, 32, 32) — see `encoding.py` for channel layout.
+    Output shape: (batch, 10, 32, 32) — 5 T timesteps then 5 P timesteps,
+        each normalized to [0, 1] via a final sigmoid.
+
+    Use `predict.predict()` for the normal inference path (loads weights,
+    builds tensor, denormalizes to physical units).
+    """
+
+    def __init__(
+        self,
+        in_channels: int = 6,
+        n_time_steps: int = 5,
+        base_filters: int = 32,
+    ) -> None:
+        super().__init__()
+        self.in_channels = in_channels
+        self.n_time_steps = n_time_steps
+        self.out_channels = 2 * n_time_steps
+
+        self.encoder = nn.Sequential(
+            ConvBlock(in_channels, base_filters, 3),
+            ConvBlock(base_filters, base_filters, 3),
+        )
+        self.middle = nn.Sequential(
+            ResidualBlock(base_filters),
+            ResidualBlock(base_filters),
+        )
+        self.decoder = nn.Sequential(
+            ConvBlock(base_filters, base_filters, 3),
+            nn.Conv2d(base_filters, self.out_channels, 1),
+            nn.Sigmoid(),
+        )
+
+        self._init_weights()
+
+    def _init_weights(self) -> None:
+        for m in self.modules():
+            if isinstance(m, nn.Conv2d):
+                nn.init.kaiming_normal_(m.weight, mode="fan_out", nonlinearity="relu")
+            elif isinstance(m, nn.BatchNorm2d):
+                nn.init.constant_(m.weight, 1)
+                nn.init.constant_(m.bias, 0)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.decoder(self.middle(self.encoder(x)))
+
+    def count_parameters(self) -> int:
+        return sum(p.numel() for p in self.parameters() if p.requires_grad)

surrogate/predict.py

@@ -0,0 +1,100 @@
+"""Model loader and inference wrapper for GeoForce v1.1 surrogate."""
+
+from __future__ import annotations
+
+from pathlib import Path
+from typing import Any
+
+import numpy as np
+import torch
+
+from .encoding import build_input_tensor, denormalize_output
+from .model import ReservoirCNN
+
+DEFAULT_WEIGHTS = Path(__file__).parent / "weights" / "geoforce_cnn_v1.1.pt"
+
+_model_cache: ReservoirCNN | None = None
+
+
+def load_model(
+    weights_path: str | Path | None = None,
+    device: str | torch.device = "cpu",
+) -> ReservoirCNN:
+    """Load v1.1 ReservoirCNN weights, cached in process memory.
+
+    Args:
+        weights_path: path to the .pt checkpoint. Defaults to the bundled v1.1.
+        device: torch device (default CPU — the model is tiny).
+
+    Returns:
+        ReservoirCNN in eval mode with weights loaded.
+    """
+    global _model_cache
+    if _model_cache is not None and weights_path is None:
+        return _model_cache
+
+    path = Path(weights_path) if weights_path else DEFAULT_WEIGHTS
+    if not path.exists():
+        msg = f"weights not found at {path}"
+        raise FileNotFoundError(msg)
+
+    ckpt: Any = torch.load(path, map_location=device, weights_only=False)
+    state_dict = ckpt["model_state_dict"] if isinstance(ckpt, dict) and "model_state_dict" in ckpt else ckpt
+    in_channels = ckpt.get("in_channels", 6) if isinstance(ckpt, dict) else 6
+
+    model = ReservoirCNN(in_channels=in_channels)
+    model.load_state_dict(state_dict)
+    model.eval()
+    model.to(device)
+
+    if weights_path is None:
+        _model_cache = model
+    return model
+
+
+def predict(
+    *,
+    initial_temperature: np.ndarray,
+    log_permeability: float,
+    well_locations: list[tuple[int, int]],
+    base_pressure: float,
+    porosity: float,
+    depth: float,
+    model: ReservoirCNN | None = None,
+) -> dict[str, np.ndarray]:
+    """Run a single v1.1 surrogate prediction.
+
+    Args:
+        initial_temperature: (32, 32) array in deg C.
+        log_permeability: scalar log10 permeability.
+        well_locations: list of (row, col) tuples.
+        base_pressure: Pa.
+        porosity: dimensionless.
+        depth: m.
+        model: pre-loaded ReservoirCNN, or None to use the cached default.
+
+    Returns:
+        dict with:
+            'temperature': (5, 32, 32) deg C
+            'pressure':    (5, 32, 32) Pa
+    """
+    if model is None:
+        model = load_model()
+
+    x = build_input_tensor(
+        initial_temperature=initial_temperature,
+        log_permeability=log_permeability,
+        well_locations=well_locations,
+        base_pressure=base_pressure,
+        porosity=porosity,
+        depth=depth,
+    )
+
+    with torch.no_grad():
+        y = model(x)
+
+    out = denormalize_output(y)
+    return {
+        "temperature": out["temperature"][0],
+        "pressure": out["pressure"][0],
+    }

surrogate/weights/geoforce_cnn_v1.1.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:3701d0ea2bcb4d8a0d68eb8e8ed1d79b3bb010a8396227e4dbddbc4327746386
+size 248595

tests/test_surrogate_smoke.py

@@ -0,0 +1,86 @@
+"""Smoke test for GeoForce v1.1 surrogate port.
+
+Asserts:
+    1. Weights load without error.
+    2. Parameter count = 57,802 (matches v1.1 metadata).
+    3. Inference on a canonical Ulubelu-like scenario runs.
+    4. Output shape is (5, 32, 32) per field.
+    5. Temperature stays in [25, 350] C and pressure in [1e5, 5e7] Pa.
+    6. Inference completes in < 50 ms on CPU (generous bound vs. 3.2 ms target).
+"""
+
+from __future__ import annotations
+
+import time
+
+import numpy as np
+import pytest
+
+from surrogate import load_model, predict
+from surrogate.encoding import GRID_SIZE, NORMALIZATION
+
+
+@pytest.fixture(scope="module")
+def model():
+    return load_model()
+
+
+def _ulubelu_like_scenario() -> dict:
+    """Canonical single-producer liquid-dominated scenario."""
+    depth = 1800.0
+    base_temperature_c = 250.0
+    # Hydrostatic-ish gradient for initial temperature: warm at base, cooler up.
+    z_grad = np.linspace(base_temperature_c, base_temperature_c - 40.0, GRID_SIZE)
+    initial_temperature = np.broadcast_to(
+        z_grad.reshape(-1, 1), (GRID_SIZE, GRID_SIZE)
+    ).astype(np.float32).copy()
+    return {
+        "initial_temperature": initial_temperature,
+        "log_permeability": -13.5,
+        "well_locations": [(16, 16)],
+        "base_pressure": 1.5e7,  # 15 MPa
+        "porosity": 0.08,
+        "depth": depth,
+    }
+
+
+def test_model_parameter_count(model):
+    assert model.count_parameters() == 57_802
+
+
+def test_model_input_channels(model):
+    assert model.in_channels == 6
+
+
+def test_prediction_shape_and_range(model):
+    scenario = _ulubelu_like_scenario()
+    result = predict(**scenario, model=model)
+
+    assert set(result.keys()) == {"temperature", "pressure"}
+    assert result["temperature"].shape == (5, GRID_SIZE, GRID_SIZE)
+    assert result["pressure"].shape == (5, GRID_SIZE, GRID_SIZE)
+
+    t = result["temperature"]
+    p = result["pressure"]
+    assert NORMALIZATION["T_MIN"] - 0.1 <= t.min()
+    assert t.max() <= NORMALIZATION["T_MAX"] + 0.1
+    assert NORMALIZATION["P_MIN"] - 1.0 <= p.min()
+    assert p.max() <= NORMALIZATION["P_MAX"] + 1.0
+
+
+def test_inference_latency(model):
+    scenario = _ulubelu_like_scenario()
+    # Warm-up run first (first inference pays init cost).
+    predict(**scenario, model=model)
+
+    t0 = time.perf_counter()
+    for _ in range(5):
+        predict(**scenario, model=model)
+    elapsed_ms = (time.perf_counter() - t0) * 1000 / 5
+    assert elapsed_ms < 50, f"inference {elapsed_ms:.1f} ms exceeds 50 ms budget"
+
+
+def test_cache_is_idempotent():
+    m1 = load_model()
+    m2 = load_model()
+    assert m1 is m2