rog2-algorithm

GPU-accelerated particle filter and beam search for the ROGII wellbore geology competition.

Drop-in replacements for the notebook's GPU-free estimators — 700x faster on a Kaggle T4, bitwise reproducible, and batching across wells so the GPU stays saturated.

from rog2_pf import lik_pf_batch, run_beam_ensemble_batch

results = lik_pf_batch([(hw, tw) for hw, tw in wells], with_quality=True)
for (out, ev_index, quality), wid in zip(results, well_ids):
    likpf_map[wid] = (out, ev_index, quality)

tvt_beams = run_beam_ensemble_batch([(hw, tw) for hw, tw in wells])
for wid, tvt_beam in zip(well_ids, tvt_beams):
    beam_map[wid] = tvt_beam

Installation

pip install rog2-algorithm

The shipped wheel bundles CUDA + wgpu + CPU backends — same package works on Kaggle T4s, local Vulkan/Metal GPUs, and CPU-only machines. Backend is selected at the call site via the backend parameter.

AMD ROCm / HIP — build from source:

pip install rog2-algorithm --no-binary :all:
# or manually:
maturin build --release --features pyo3/extension-module,python,hip
pip install dist/rog2_algorithm-*.whl

Dependencies: Python 3.9+ and numpy.

API

All batched functions take every well at once — one GPU launch replaces a Python loop over wells, keeping the GPU saturated.

Particle filter

lik_pf_batch(pairs, ...)

Batched drop-in for the notebook's lik_pf. Returns one (out, ev_index, quality) tuple per input pair.

Returns list[(out, ev_index, quality)]:

• out — dict[str, np.ndarray] — channels pf_scale_3/5/8/12 and pf_mean, one float32 array per well over its evaluation rows.
• ev_index — np.ndarray — index slice the evaluation rows occupy in the original DataFrame.
• quality — dict — pf_best_ll, pf_ll_spread, pf_pt_std, pf_gr_sig (empty when with_quality=False).

Wells with no evaluation rows return ({}, empty_index, {}).

lik_pf(hw, tw, **kwargs)

Single-well convenience. Prefer the batched form.

run_particle_filter(hw, tw, n_particles=500, seed=42, backend="auto", **kwargs)

Name-compatible alias for the notebook's own run_particle_filter — same signature, same (pred, log_lik) return, known rows preserved and evaluation rows filled in over the full hw length (not just the evaluation slice).

run_particle_filter is inherently a compromise, not a true port. The GPU kernel doesn't expose one seed's raw trajectory — only the seed-ensembled channels. This gets around that by requesting n_seeds=1 (softmax over one element is the identity, so pf_mean is that seed's path), and reconstructs log_lik by undoing the pf_best_ll = liks.max() / n_rows normalization. It's correct, but it's a single-seed GPU launch, which is a wasteful way to use this API. For real use, run_pf_lik_ensemble_scales / lik_pf_batch is the right call.

run_pf_lik_ensemble_scales(hw, tw, scales=(3.0, 5.0, 8.0, 12.0), n_particles=500, n_seeds=128, backend="auto", **kwargs)

Name-compatible alias for the notebook's own run_pf_lik_ensemble_scales — same signature, wraps lik_pf_batch for one well. Returns dict[str, np.ndarray] keyed f"pf_scale_{s:g}" per requested scale plus "pf_mean", each a full-length array (known rows preserved, evaluation rows filled). The notebook's branch_stats argument has no equivalent: it feeds a bimodal-hedge diagnostic built on raw per-seed paths, which the GPU kernel doesn't expose (only the ensembled channels, and, via lik_pf_batch(..., with_quality=True), a per-well quality scalar) — so it isn't accepted here. Prefer calling lik_pf_batch directly across every well at once; one well's seed ensemble can't fill a GPU either.

Beam search

run_beam_ensemble_batch(pairs, ...)

Batched drop-in for the notebook's run_beam_ensemble. Returns one array per input pair matching hw.TVT_input layout (known rows preserved, evaluation rows filled with the 14-config ensemble mean).

Wells with no evaluation rows return a copy of their input TVT_input column.

run_beam_ensemble(hw, tw, **kwargs)

Single-well convenience — prefer the batched form.

beam_search(hgr, tw_tvt, tw_gr, last_tvt, ...)

Single-config drop-in for the notebook's beam_search.

Low-level native API

The native _rog2_pf module is re-exported through the rog2_pf package. Use these directly when you have pre-built input dicts.

run_batch(wells, scales, ...)

Low-level particle filter — no pandas, no prepare_well. Takes pre-built well dicts.

Returns a dict with keys pf_scale_*, pf_mean (lists of per-well arrays), pf_pt_std (per-well std when with_std=True), liks ([n_wells, n_seeds]), kept (indices of non-empty wells), and channels.

run_beam_batch(wells, ...)

Low-level beam search.

Returns a dict with beam_mean (list of per-well arrays), kept, and optionally per_config ([config][well] arrays).

make_grid(tw_tvt, tw_gr, step)

Resamples a type-well log onto a uniform TVT grid. Returns (grid_array, vmin, step).

available_backends()

Returns list[str] — the backends the installed wheel supports (e.g., ["cuda", "wgpu", "cpu"]).

notebook_beam_configs()

Returns the default 14 (beam_size, move_cost, err_scale, radius) tuples matching the competition notebook.

Helpers & constants

• prepare_well(hw, tw, init_spr, seed_base) — Build a particle-filter input dict from a (hw, tw) pair. Returns (well_dict, ev_index).
• prepare_beam_well(hw, tw) — Build a beam-search input dict. Returns (well_dict, ev_index).
• DEFAULT_SCALES — (3.0, 5.0, 8.0, 12.0).
• BEAM_CONFIGS — The 14 notebook configs as (bs, mc, es, r) tuples.

Deterministic seeding

Pass seed_bases=[stable_hash(wid)] to lik_pf_batch for per-well reproducible seeding. The RNG is counter-based — every draw is a pure function of (seed_base + seed, particle slot, step index, draw index). Results are bitwise reproducible across runs, devices and thread schedules for a fixed (seed_base, n_particles).

The beam search has no RNG and is bit-identical across runs, devices, and cube_dim.

Backend selection

The backend parameter accepts: "cuda", "wgpu", "hip" / "rocm", "cpu", or "auto" (tries compiled-in backends in preference order: HIP → CUDA → wgpu → CPU, falling through to the next candidate when one fails to initialise — e.g. no NVIDIA driver present for a wheel built with cuda).

from rog2_pf import lik_pf_batch, available_backends

print(available_backends())         # what the installed wheel supports
results = lik_pf_batch(pairs, backend="wgpu")  # explicit backend

Build from source

# CUDA
maturin build --release --no-default-features \
  --features 'pyo3/extension-module,python,cuda' -o dist

# ROCm / HIP
maturin build --release --no-default-features \
  --features 'pyo3/extension-module,python,hip' -o dist

# wgpu (Vulkan / Metal / DX12)
maturin build --release --no-default-features \
  --features 'pyo3/extension-module,python,wgpu' -o dist

# CPU only
maturin build --release --no-default-features \
  --features 'pyo3/extension-module,python,cpu' -o dist

pip install dist/rog2_algorithm-*.whl

Requires Rust toolchain and maturin (pip install maturin).

Performance

See tests/ for the cross-language parity assertions and kaggle/ for the T4 benchmark notebooks.