Architecture & Design

Architecture & Design#

Design Principles#

GeoBrain is built on three foundational principles:

End-to-End Differentiability#

All forward models are implemented as PyTorch nn.Module subclasses, enabling automatic differentiation through complex computational graphs. This eliminates hand-coded Jacobians and enables gradient-based optimization for inverse problems.

# Gradients flow through the entire chain
velocity = rock_physics_model(elastic_moduli)      # Differentiable
seismic_data = wave_propagator(velocity)            # Differentiable
loss = mse_loss(seismic_data, observed_data)
loss.backward()  # Automatic gradient computation

Modular Composition#

Each physics module is a self-contained unit that can be used standalone or composed with other modules. Modules communicate through standard tensor interfaces:

geomodel  -->  rock_physics  -->  wave_propagation  -->  observed_data
   |               |                    |                      |
   v               v                    v                      v
porosity      velocities          seismic traces          data misfit
Modular composition: geostatistics to rock physics to elastic velocities

Fig. 11 Modular composition: geostatistics to rock physics to elastic velocities.#

Bayesian-Deterministic Duality#

The same forward model can be used for both deterministic inversion (geobrain.optim) and Bayesian inference (geobrain.bayes) without modification. The InverseProblem class provides a unified entry point:

from geobrain import InverseProblem, SVGD

problem = InverseProblem(forward_fn=forward, observed=data, noise_std=0.01)

# Deterministic
inverter = problem.create_inverter(initial_model=m0)

# Bayesian (4 samplers: SVGD, HMC, NUTS, LDS)
posterior = problem.as_posterior(log_prior=prior)
svgd = SVGD(target=posterior)

Module Organization#

geobrain/
├── core/              # InverseProblem, registry system
├── bayes/             # Bayesian samplers and distributions
│   └── samplers/      # SVGD, HMC, NUTS, LDS (Langevin)
├── optim/             # Deterministic inversion engine
├── geomodel/          # Geological model generation
│   ├── geostats/      # FFT-MA, SGS, variograms
│   └── geogen/        # GAN, VAE, Diffusion generators
├── physics/
│   ├── rock/          # Effective medium, Gassmann, empirical
│   ├── wave/          # Acoustic & elastic propagation, AVO
│   │   ├── avo/       # Zoeppritz, Aki-Richards, Shuey
│   │   ├── boundary/  # PML, Gerjan, SinCos
│   │   ├── wavelets/  # Ricker, Gaussian, Ormsby, Klauder
│   │   ├── models/    # Acoustic & elastic velocity models
│   │   ├── propagators/ # FDTD solvers
│   │   └── survey/    # Source, Receiver, SeismicData
│   └── flow/          # Reservoir simulation
├── nn/                # Neural network layers & activations
├── data/              # Dataset loaders and transforms
├── io/                # SEG-Y, LAS, reservoir grid I/O
├── vis/               # Visualization utilities
└── decision/          # Decision-making under uncertainty

Data Flow#

A typical GeoBrain workflow follows this pattern:

  1. Model generation — Create geological properties (porosity, facies) using geostatistics or generative models.

  2. Rock physics — Transform geological properties into elastic parameters (Vp, Vs, density) via rock physics models.

  3. Forward simulation — Simulate observed geophysical data (seismic, resistivity, production) from the elastic model.

  4. Inversion — Recover subsurface properties from observed data using deterministic optimization or Bayesian sampling.

  5. Decision — Use the posterior ensemble for uncertainty-aware decision-making.

GPU & Memory#

  • All modules are GPU-compatible via PyTorch device management.

  • Large-scale wave simulations use gradient checkpointing to trade compute for memory.

  • 3D acoustic propagation supports domain-parallel execution.

Extensibility#

New physics modules are registered automatically via the RegistryHub:

from geobrain.core import registries

# Register a custom forward model
@registries.physics.register("my_model")
class MyModel(nn.Module):
    def forward(self, params):
        ...

This allows new modules to be discovered and composed without modifying framework code.

Contents