Accelerated Thermal Breakthrough Prediction in Geothermal Reservoirs Using a Physics-Aware Graph Neural Network Surrogate

A graph neural network surrogate that forecasts the full long-term temperature evolution of a heterogeneous geothermal doublet from the initial condition alone. Time-of-flight analysis drives a physics-aware algebraic-multigrid (AMG) coarsening that concentrates the model's capacity along the flow corridor connecting the wells, the temporal dynamics are learned on the coarse graph with streamline-causal attention and a node-wise recurrent memory, and a persistence-prior residual head is trained over the full forecast horizon. Given only the temperature field at the start of operation, the model rolls the entire 100-year sequence forward autoregressively, with no further input from the simulator.

Paper: to be added after publication.

Installation

The model uses PyTorch and PyTorch Geometric, whose wheels are specific to your platform and CUDA version, so they are installed separately from the rest.

Conda:

conda env create -f environment.yml
conda activate tbt
# install the platform-specific builds (example: CUDA 12.1)
pip install torch --index-url https://download.pytorch.org/whl/cu121
pip install torch-geometric pytorch-lightning
pip install -e .

Pip only:

python -m venv .venv && source .venv/bin/activate
# pick the torch build for your platform from https://pytorch.org/get-started
pip install torch --index-url https://download.pytorch.org/whl/cu121
pip install -r requirements.txt
pip install -e .

For a CPU-only install, use the cpu index URL for torch and the matching PyTorch Geometric build.

Data

The synthetic dataset (600 thermal-doublet simulations) is hosted on Zenodo:

• DOI: 10.5281/zenodo.20818085

Download thermal_doublet_dataset.h5 and place it under data/. The AMG cache is not distributed; it is rebuilt locally in step 2 below.

Reproduce

# 1. preprocessing: manifest + normalization statistics
python -m tbt.data.preprocess --config configs/data_prep.yaml

# 2. build the physics-aware AMG cache (one operator set per case)
python -m tbt.data.amg --config configs/data_prep.yaml

# 3. train the surrogate
python -m tbt.training.train --config configs/model.yaml

# 4. evaluate the trained surrogate on the held-out cases
python -m tbt.eval.rollout --config configs/model.yaml \
    --ckpt checkpoints/surrogate.ckpt --split eval

# 5. render the AMG coarsening visualization (optionally for a specific case)
python -m tbt.eval.visualize_amg --config configs/data_prep.yaml --case-id 987

# list the case IDs in a split (train / val / eval), e.g. to pick cases to plot
python -m tbt.data.dataset --config configs/model.yaml --split eval

# render the per-case rollout figure (fields, error, producer profile) for
# one or more of those case IDs
python -m tbt.eval.rollout --config configs/model.yaml \
    --ckpt checkpoints/surrogate.ckpt --split eval --only-case 987

The released checkpoint checkpoints/surrogate.ckpt reproduces the paper's evaluation: mean field RMSE 0.55 K and producer-well error 0.77 K over the 100-year autoregressive rollout on the 50 held-out cases.

Quick test

A self-contained test that needs no GPU and no dataset download. It loads the shipped checkpoint, verifies the parameter count, and runs the full 54-step autoregressive rollout on CPU for one bundled evaluation case:

python quick_test.py

Expected output (well under a minute on a multi-core CPU):

parameters: 1,207,020 (OK)
rollout field RMSE:    0.44 K  (threshold 0.80)
rollout producer RMSE: 0.25 K  (threshold 0.60)
PASS

Examples

The nine evaluation cases shown in the paper's producer-temperature figure are bundled under examples/ (extracted from the Zenodo dataset): 585, 537, 149, 57, 70, 487, 1026, 354, and 484. Each is fully self-contained, so both example commands below run on CPU with no dataset download.

Both outputs come from the same script, examples/run_example.py, which runs the full autoregressive rollout for the chosen case and writes the result to examples/output/. Without any flag it produces a static figure; with the --animate flag it produces an animated GIF instead.

Rollout figure (default). Writes examples/output/case_<id>.png, with the simulated and predicted temperature fields and the percent error map at four rollout times, plus the producer temperature profile:

python examples/run_example.py --case 585

The output for case 585:

Rollout animation (--animate). Writes examples/output/case_<id>.gif, with the simulated and predicted fields and the percent error map advancing over the full 100-year rollout. The model sees only the initial state; every later frame is its own prediction fed back in:

python examples/run_example.py --case 354 --animate

The output for case 354:

Repository layout

configs/        model.yaml (architecture + training), data_prep.yaml (preprocessing + AMG)
src/tbt/
  config.py     YAML config loading
  utils.py      runtime helpers
  data/         preprocess.py, amg.py, dataset.py
  models/       layers.py, coarse_blocks.py, surrogate.py
  training/     train.py
  eval/         rollout.py, visualize_amg.py
quick_test.py   self-contained CPU test (see Quick test above)
examples/       nine bundled evaluation cases, run_example.py, output/ with the showcase figure and animation
checkpoints/    surrogate.ckpt (trained weights)

Citation

A preprint describing this work will be made available; the citation will be added here.

License

MIT, see LICENSE.