Adaptive SDE sampling,
guided by geometry.

driftflow is a pure-Python implementation of the Drift Variation Score — a training-free, plug-in sampler that adapts its timestep to the local Fisher-Rao curvature of the transition manifold. Wrap any graph-diffusion drift. Get sharper samples, faster.

pip install driftflow View on GitHub Cite the paper

python driftflow

# driftflow — adaptive SDE sampling for graph diffusion
from driftflow import (
    CommonConfig, DVSSampler, GruMApproximation,
    LinearSchedule, get_dataset_config, make_drift_function,
)

approx    = GruMApproximation(num_nodes=9, feature_dim=4, seed=42)
drift     = make_drift_function(approx)
common    = CommonConfig()
dataset   = get_dataset_config("GruM", "QM9")

sampler = DVSSampler(
    drift_function  = drift,
    noise_schedule  = LinearSchedule(sigma_min=0.01, sigma_max=0.5),
    common_config   = common,
    dataset_config  = dataset,
    solver          = "Heun",
    seed            = 42,
)

features_t, adjacency_t, info = sampler.sample(
    initial_features = features_0,
    initial_adjacency = adjacency_0,
    terminal_time    = 1.0,
)

print(f"steps={info['total_steps'][0]}  time={info['final_time'][0]:.4f}")

dt = 2.34e-3

t = 0.62 · DVS active

shape OK · 100% coverage

Tests · 165 passing

Pure Python 3.10+ Zero runtime deps 100% test coverage Deterministic · seeded MIT licensed Type-hinted (PEP 561) mypy --strict clean Algorithm 1 · 2 · 3 Pure Python 3.10+ Zero runtime deps 100% test coverage Deterministic · seeded MIT licensed Type-hinted (PEP 561) mypy --strict clean Algorithm 1 · 2 · 3 Pure Python 3.10+ Zero runtime deps 100% test coverage Deterministic · seeded MIT licensed Type-hinted (PEP 561) mypy --strict clean Algorithm 1 · 2 · 3

Overview

The right step size, exactly when it matters.

Graph-diffusion SDE solvers waste work in flat regions and miss curvature in steep ones. driftflow measures the local Drift Variation Score — a Fisher-Rao–style signal of how quickly the drift is turning — and resizes each timestep in a single power-law rule. The result is a sampler that takes fewer, better-placed steps without touching your model.

Training-free — wrap any graph-diffusion drift
Drop-in replacement for Euler-Maruyama and Heun
Reproducible with a single seed
Hyperparameters pre-validated for GruM & GDSS

Signal · Adaptive timestep

live preview

165

Tests passing

100%

Line coverage

Runtime deps

Algorithms

(Model, Dataset) configs

MIT

License

trajectory · QM9 · Heun

Features

Engineered for serious diffusion work.

A small, opinionated surface area built on six primitives. Each one earns its keep — no boilerplate, no half-features.

DVS · Eq. 13–15

Drift Variation Score

Composable Python primitives compute the per-modality variation of your drift. Plug them into any graph-diffusion pipeline and decide exactly where adaptation should happen.

Algorithm 2 & 3

Two production solvers

First-order Euler-Maruyama for speed, second-order Heun for accuracy — both with the shared diffusion noise scale g(t) · √dt. Register your own via the SOLVERS extension point.

Tables 6 & 7

Curated hyperparameter set

Common and dataset-specific settings are validated, named, and ready to import. CommonConfig + get_dataset_config() reproduce paper defaults without hand-tuning.

Localized control

Active-range gating

Adapt only where it pays off — e.g. [0, 0.2] ∪ [0.95, 1.0] for GDSS/QM9. Outside the range, the sampler falls back to dt_base for predictable cost.

Embeddable

Pure Python, zero deps

No NumPy, no PyTorch. Every numerical primitive operates on nested list[float] containers so driftflow runs anywhere — educational tools, edge inference, lightweight servers.

Reproducibility

Deterministic by default

A single seed makes the entire trajectory bit-for-bit reproducible, including the per-step history returned by sample(). Fail-fast errors surface before any work is done.

Algorithms

Three algorithms,
one adaptive loop.

The paper presents three sampler variants. driftflow implements all of them behind one class — DVSSampler — and dispatches through a solver registry so you can add your own.

Active range: [0, 1] or [0, 0.2] ∪ [0.95, 1.0]
Solver registry: SOLVERS / register_solver()
Default solver: Euler (Heun available)
Output: (X_T, A_T, info_dict)

01

DVS Meta-Sampler
Orchestrator

Compute the DVS, smooth it with EMA, scale the timestep by a power law, and refresh globally each step. The orchestration that every other algorithm composes on top of.
02

DVS-Euler-Maruyama
First order

First-order SDE stepping with shared diffusion noise g(t) · √dt. Fast, predictable, and the natural baseline for any graph-diffusion drift.
03

DVS-Heun
Second order

Predictor-corrector second-order step. Same diffusion noise scale, with one extra drift evaluation per step in exchange for tighter error control.

Developer API

A surface that fits in your head.

Four imports. One sampler. Hyperparameters pre-validated for GruM and GDSS. Everything else is just a dataclass.


                from driftflow import (
    CommonConfig, DVSSampler, GruMApproximation,
    LinearSchedule, get_dataset_config, make_drift_function,
)

approx    = GruMApproximation(num_nodes=9, feature_dim=4, seed=42)
drift     = make_drift_function(approx)
common    = CommonConfig()
dataset   = get_dataset_config("GruM", "QM9")

sampler = DVSSampler(
    drift_function  = drift,
    noise_schedule  = LinearSchedule(sigma_min=0.01, sigma_max=0.5),
    common_config   = common,
    dataset_config  = dataset,
    solver          = "Heun",
    seed            = 42,
)

features_t, adjacency_t, info = sampler.sample(
    initial_features = features_0,
    initial_adjacency = adjacency_0,
    terminal_time    = 1.0,
)
print(f"steps={info['total_steps'][0]}  time={info['final_time'][0]:.4f}")


                from driftflow import CommonConfig, DatasetConfig, get_dataset_config

common = CommonConfig()
# alpha=0.2, beta=0.5, dt_base=1e-3, dt_min=2e-4, dt_max=5e-3

dataset = get_dataset_config("GruM", "QM9")
# kappa_ref=1.0, gamma_euler=0.22, gamma_heun=0.23
# active_range=[(0.0, 1.0)]

# Or build your own — Table 7 entries are dataclasses:
custom = DatasetConfig(
    kappa_ref  = 5.0,
    gamma_euler = 0.02,
    gamma_heun  = 0.04,
    active_range = [(0.0, 1.0)],
)


                from driftflow import (
    LinearSchedule, CosineSchedule, PolynomialSchedule, constant_schedule,
)

# Linear g(t) = sigma_min + (sigma_max - sigma_min) * t
linear = LinearSchedule(sigma_min=0.01, sigma_max=0.5)

# Improved-DDPM-style cosine
cosine = CosineSchedule(offset=0.008)

# Power-law g(t) = sigma_max * t**exponent + sigma_min
poly   = PolynomialSchedule(exponent=2.0,
                                sigma_min=0.0,
                                sigma_max=1.0)

# Deterministic constant schedule (useful for testing)
const = constant_schedule(0.1)

# Or supply your own callable as the noise schedule.


                from driftflow import DVSSampler, SOLVERS, register_solver

# Built-in solvers:
print(list(SOLVERS))        # ['Euler', 'Heun']

# Swap solver at construction time:
sampler = DVSSampler(..., solver="Heun")

# Register a custom step (must follow the (drift, X, A, t, dt, g) -> (X, A) contract):
def my_step(drift, X, A, t, dt, g):
    f_X, f_A = drift(X, A, t)
    return (
        [row[:] for row in X],   # placeholder update
        [row[:] for row in A],
    )

register_solver("Mine", my_step)
sampler = DVSSampler(..., solver="Mine")

See docs/API_REFERENCE.md for the full public surface.

Built with rigor

Production discipline, research-grade fidelity.

driftflow is treated like infrastructure — strict types, gated coverage, deterministic runs, and explicit fail-fast input validation.

165

Unit tests

Mathematical correctness, edge cases, determinism, input validation.

100%

Line coverage

Coverage-gated in CI: a single missed branch breaks the build.

Runtime dependencies

Pure-Python standard library — embeddable anywhere.

Pre-validated configs

CommonConfig + every (GruM, GDSS) × dataset entry from Table 7.

mypy

strict

ruff

lint + format

pytest

cov ≥ 100%

pip-audit

strict

Architecture

One loop. Clean separation of concerns.

The adaptive controller (DVS, EMA, power-law scaling) is fully decoupled from the SDE solver and the denoiser network. Drop in any of the three without touching the others.

Drift

Your denoiser

Any callable (X, A, t) → (f_X, f_A). Real networks or simplified stand-ins.

DVS

Drift Variation Score

EMA-smoothed variation of the drift, scaled by a power-law rule into a timestep.

Solver

Euler-Maruyama · Heun

First- or second-order SDE step with shared diffusion noise g(t) · √dt.

Output

X_T, A_T, info

Sample features, adjacency, and per-step history (steps, times, dt_k).

Sampler loop · per step

drift

(X, A, t) → f_X, f_A

→

DVS

v_x, v_a = DVS(f_X, f_A)

→

step

dt = min(dt_x, dt_a)

→

integrate

X, A = SOLVERS[name](…)

Install

Three lines. Then run the demo.

driftflow is a pure-Python package with no compiled extensions. A normal pip install is everything you need — the smoke-test CLI is included for GruM and GDSS on QM9.

Python 3.10, 3.11, 3.12, 3.13
Linux · macOS · Windows
Optional: pip-audit, ruff, mypy, pytest-cov

Install · from source

git clone https://github.com/sachncs/driftflow.git
cd driftflow
pip install -e .

Optional · dev dependencies

pip install -e ".[dev]"

Run the demo

# Euler with GruM approximation
python examples/demo.py --model GruM --dataset QM9 --solver Euler --use-approximation

# Heun with GDSS approximation
python examples/demo.py --model GDSS --dataset QM9 --solver Heun --use-approximation

Cite

A faithful reproduction of an arXiv paper.

If you use driftflow in academic work, please cite the original paper. The package is a pure-Python reproduction of Algorithms 1–3 from the reference below.

arXiv:2605.00250 CHANGELOG

BibTeX

@article{driftflow2026,
  title  = {Information-Geometric Adaptive Sampling for Graph Diffusion},
  journal = {arXiv preprint arXiv:2605.00250},
  year   = {2026}
}

Get started

Wrap a drift. Sample smarter.

pip-install driftflow, point it at any graph-diffusion denoiser, and let the Drift Variation Score take it from there.

pip install driftflow Star on GitHub

Adaptive SDE sampling, guided by geometry.

The right step size, exactly when it matters.

Engineered for serious diffusion work.

Drift Variation Score

Two production solvers

Curated hyperparameter set

Active-range gating

Pure Python, zero deps

Deterministic by default

Three algorithms,one adaptive loop.

DVS Meta-Sampler

DVS-Euler-Maruyama

DVS-Heun

A surface that fits in your head.

Production discipline, research-grade fidelity.

One loop. Clean separation of concerns.

Your denoiser

Drift Variation Score

Euler-Maruyama · Heun

X_T, A_T, info

Three lines. Then run the demo.

A faithful reproduction of an arXiv paper.

Wrap a drift. Sample smarter.

Adaptive SDE sampling,
guided by geometry.

Three algorithms,
one adaptive loop.