v2.0 · typed research infrastructure

Spend compute where the trajectory bends.

driftflow is a validated Python sampler for graph-diffusion SDEs. Drift Variation Score adapts the timestep to local change while your model and solver stay replaceable.

Start with the API Read the docs

Python 3.10–3.14stdlib runtimeMIT licensed

One explicit pipeline validated

01GraphStateimmutable input→

02DriftModelevaluate(state, t)→

03DVS controllerchoose a bounded step→

04SolverEuler-Maruyama or Heun✓

boundaries checked before integration

Why adaptive?

One loop. Better decisions.

Flat regions do not need the same resolution as sharp turns. DVS makes that decision explicit without coupling your model to the orchestration loop.

01 / MEASURE

Read local change

Compare consecutive drifts and normalize the variation by the current noise scale.

02 / CONTROL

Bound the response

EMA smoothing and power-law scaling choose one synchronized timestep for both modalities.

03 / INTEGRATE

Swap the machinery

Pass the typed state through a polymorphic solver and keep structured diagnostics.

The v2 API

Small surface. Strong contracts.

Replace a model, schedule, controller, solver, or decoder without rewriting the sampler.

state → DriftModel → DVS controller → Solver

result.state
result.diagnostics

class DriftModel(ABC):
    def evaluate(state, time) -> DriftState: ...

class Solver(ABC):
    def step(...) -> GraphState: ...

GraphState(...)            # shape + finite checks
LinearSchedule(...)        # scale + time checks
DVSSampler.sample(...)      # step-limit checks

from driftflow import (
    CommonConfig, DvsAdaptiveController, DVSSampler,
    GraphState, GruMApproximation, LinearSchedule,
    get_dataset_config,
)

state = GraphState.from_lists(
    [[0.1, 0.2], [0.3, 0.4]],
    [[0.0, 0.1], [0.1, 0.0]],
)
dataset = get_dataset_config("GruM", "QM9")
sampler = DVSSampler(
    drift_model=GruMApproximation(2, 2, seed=42),
    noise_schedule=LinearSchedule(0.01, 0.5),
    controller=DvsAdaptiveController(
        CommonConfig(), dataset, dataset.gamma_euler or 0.0
    ),
)
result = sampler.sample(state, seed=42)

Reliability by design

Clear failures are part of the API.

Infrastructure is easier to trust when its boundaries are visible.

Immutable statesGraphState and DriftState freeze nested matrices and reject non-finite values.

Explicit errorsConfiguration, shape, schedule, numerical, solver, and sampling-limit failures.

Honest scopeApproximation models are labeled educational stand-ins, not trained denoisers.

Per-call seedsReproducible trajectories without hidden global random state.

Install

Start from a state you can inspect.

Install the package, read the migration guide, then wire in your production denoiser.

$python -m pip install driftflow

python -m pip install driftflow

Open documentation