geoGP3D

Gaussian-process fusion of 3-D geophysical datasets, built on GPyTorch's sparse variational inference. geoGP3D constructs a posterior distribution whose uncertainty is learned from inter-model disagreement, and samples coherent 3-D realisations for downstream applications such as probabilistic ground motion prediction.

What it does

geoGP3D wraps GPyTorch to provide a science-friendly workflow for continuous geospatial fields:

1. Transform geographic NetCDF data into Cartesian or UTM coordinates.
2. Preprocess — normalise coordinates (mandatory for ARD kernels), remove polynomial trends, normalise values.
3. Weight observations by inverse sampling density to correct preferential-sampling bias.
4. Train a sparse variational GP (SingleTaskApproximateGP or MultiTaskApproximateGP) using Trainer.
5. Predict on arbitrary query grids using Predictor with STANDARD / LOVE / CIQ sampling.
6. Inverse-preprocess and export back to xarray / NetCDF.
7. Optionally model discontinuities (Moho, faults, stratigraphic horizons) with per-layer GPs and probabilistic stitching.

Installation

Requires Python ≥ 3.12.

pip

pip install geogp3d

On Linux, the default PyPI torch wheel bundles CUDA support (a large download). For a CPU-only install, pre-install torch from the CPU index first — pip then leaves it untouched:

pip install torch --index-url https://download.pytorch.org/whl/cpu
pip install geogp3d

uv

With uv, the cpu / cu128 extras select the PyTorch wheel index directly (they are mutually exclusive):

uv add "geogp3d[cpu]"    # CPU-only torch wheel
uv add "geogp3d[cu128]"  # CUDA 12.8 torch wheel

For a development install, see CONTRIBUTING.md.

Quick-start

Single-task GP on a synthetic 3D field

import torch
import gpytorch as gpt
from geogp3d import (
    SingleTaskApproximateGP,
    Trainer, TrainingConfig,
    PredictionSpace, Predictor, SamplingStrategy,
    preprocess, inverse_preprocess,
)

# Raw 3D geospatial data (e.g. from xarray_to_tensor)
train_coords = torch.rand(1000, 3) * torch.tensor([500.0, 500.0, 50.0])  # km
train_values = torch.randn(1000) * 2.0 + 5.5  # km/s

# Step 1 – Preprocess (normalise coords / remove trend / normalise values)
coords_norm, values_norm, params = preprocess(train_coords, train_values, degree=1)
inducing_norm = coords_norm[:50]  # inducing points in normalised space

# Step 2 – Build model
likelihood = gpt.likelihoods.GaussianLikelihood()
model = SingleTaskApproximateGP(
    likelihood=likelihood,
    inducing_points=inducing_norm,
    kernel=gpt.kernels.RBFKernel(ard_num_dims=3),
)

# Step 3 – Train
mll = gpt.mlls.VariationalELBO(likelihood, model, num_data=len(values_norm))
history = Trainer(model, mll, TrainingConfig(num_epochs=200)).train(
    coords_norm, values_norm
)

# Step 4 – Predict
test_coords_norm, _, _ = preprocess(
    torch.rand(200, 3) * torch.tensor([500.0, 500.0, 50.0]),
    torch.zeros(200),
    coord_norm_params=params.coord_norm,
    degree=None,
    value_method=None,
)
predictor = Predictor(model, likelihood)
summary = predictor.summarize(test_coords_norm)
samples = predictor.sample(
    test_coords_norm,
    num_samples=100,
    space=PredictionSpace.LATENT,
    sampling_strategy=SamplingStrategy.LOVE,
)

# For large export grids where you only need marginal summaries:
# summary = predictor.summarize_batched(test_coords_norm, batch_size=4096)

# Step 5 – Back to original scale
_, predictions = inverse_preprocess(test_coords_norm, summary.mean, params)

summarize() and summarize_batched() return marginal posterior summaries only. Use sample() when you need full-domain joint draws; latent-space is the default and is the recommended scientific path unless you explicitly need Gaussian observed-space samples.

CRS transform + NetCDF round-trip

import xarray as xr
import torch
from geogp3d.io import transform_crs, xarray_to_tensor, prediction_to_xarray

# Load geographic NetCDF
ds = xr.open_dataset("velocity.nc")

# Transform to local Cartesian (< 100 km extent)
ds_local = transform_crs(
    ds, source_crs="EPSG:4326", target_crs="local_cartesian",
    origin=(-122.0, 37.5),
)

# Coordinate names are preserved; only coordinate values are transformed.
X, y, coords = xarray_to_tensor(
    ds_local["vp"], coord_dims=["lon", "lat", "depth"]
)

# ... train model, then summarize and sample with Predictor ...
summary = predictor.summarize(X)

# Export back
ds_pred = prediction_to_xarray(
    coords,
    mean=summary.mean,
    variance=summary.variance,
    coord_dims=["lon", "lat", "depth"],
)
ds_pred.attrs["cartesian_origin"] = ds_local.attrs["cartesian_origin"]
ds_geo = transform_crs(ds_pred, source_crs="local_cartesian", target_crs="EPSG:4326")
ds_geo.to_netcdf("predictions.nc")

Density-weighted training

from geogp3d import WeightedPredictiveLogLikelihood
from geogp3d.weighting import compute_density_weights

coords_norm, values_norm, params = preprocess(train_coords, train_values, degree=1)
weights, _ = compute_density_weights(coords_norm, density_method="knn", k=10)

mll = WeightedPredictiveLogLikelihood(likelihood, model, num_data=len(values_norm))
history = Trainer(model, mll, config).train(coords_norm, values_norm, weights=weights)

Evidence boundaries

The GPU-ready design philosophy (no hardcoded .cpu() / device="cpu") means CUDA should work without code changes, but CUDA CI is not part of this project. Do not cite CUDA benchmarks from this codebase.

Model/prediction artifacts are saved with torch.save (pickle-based). Only load artifact files from sources you trust.

Architecture summary

src/geogp3d/
├── models/          GP wrappers (Single/MultiTask) + custom kernels (Gibbs, ShiftedLinear, …)
├── objectives/      WeightedPredictiveLogLikelihood
├── training/        Trainer, TrainingConfig, TrainingHistory
├── prediction/      Predictor, PredictionSpace, PredictionSummary, SamplingStrategy (STANDARD / LOVE / CIQ)
├── preprocessing/   Coordinate normalisation, polynomial detrending, value scaling
├── weighting/       Inverse-density weights (k-NN or Voronoi), multi-dataset support
├── discontinuities/ Piecewise-smooth GP modelling (configuration, surface GPs, probabilistic stitching, multi-dataset fusion)
└── io/              xarray ↔ tensor conversion, CRS transforms (UTM, local Cartesian)

The canonical import name is lowercase: import geogp3d.

geoGP3D is a pure library package: this repository ships the GP fusion machinery and its tests, nothing else. Publication analyses that build on geoGP3D live in their own repositories alongside the papers they support, pinning a tagged geoGP3D release.

Citing

If you use geoGP3D in your research, please cite it via CITATION.cff (GitHub's "Cite this repository" button). Publication-specific citations live with each publication's own code repository.

Contributing

Bug reports, feature requests, and PRs are welcome — see CONTRIBUTING.md.

License

Apache-2.0