Skip to content

Quickstart

From a panel of logs to a schedule and its certificate, in one call. The causal assumption is a required argument, because the default would be "adjust for nothing", which is a claim rather than its absence — and when the graph says the effect is not identified there is no schedule at all.

Everything below is one script, docs/quickstart.py, and the report at the end is what it printed when this site was built:

pip install causal-hybrid-control
python quickstart.py

1. Precision first

import jax
import numpy as np

jax.config.update("jax_enable_x64", True)  # before any array exists; see the dtype policy

from chc import CausalGraph, Constraint, Lever, Panel, Target, prescribe  # noqa: E402

JAX computes in single precision unless told otherwise, and an identified channel is a number someone will act on. Set the flag before anything creates an array. prescribe does not refuse a run in single precision; it warns, and records the flag in the result's run. The dtype policy says why.

2. The logs

DT = 0.1  # one logged period, in the model's time units


def simulate_logs(regions: int = 200, weeks: int = 12, seed: int = 0) -> dict[str, np.ndarray]:
    """A driver pool logged under a policy that raised the incentive whenever demand spiked."""
    rng = np.random.default_rng(seed)
    names = ("region", "week", "supply", "wait", "incentive", "demand")
    rows: dict[str, list[float]] = {name: [] for name in names}
    for region in range(regions):
        supply, wait = rng.normal(0.0, 0.2), rng.normal(0.0, 0.2)
        for week in range(weeks):
            demand = rng.normal(0.0, 1.0)
            incentive = 0.9 * demand + rng.normal(0.0, 0.5)  # the policy chases demand
            logged = (region, week, supply, wait, incentive, demand)
            for name, value in zip(names, logged, strict=True):
                rows[name].append(value)
            supply, wait = (
                supply
                + DT * (-0.6 * supply + 0.3 * wait + 0.8 * incentive + 1.5 * demand)
                + rng.normal(0.0, 0.01),
                wait + DT * (0.25 * wait - 0.4 * incentive) + rng.normal(0.0, 0.01),
            )
    return {name: np.asarray(values) for name, values in rows.items()}


panel = Panel.from_frame(simulate_logs(), unit="region", time="week", seed=0)

A driver pool, logged weekly per region: supply, the riders' wait, the incentive paid, and the demand shock. The policy that wrote the log raised the incentive whenever demand spiked, and demand also moves supply — so a model fitted on this log credits the incentive with what demand did.

Panel.from_frame takes a pandas frame, a polars frame or a mapping of columns, and neither frame library is a dependency. It checks the (unit, time) index once and names the column and the entity when it refuses.

3. The assumption, and the call

edges = [
    ("demand", "incentive"),  # the confounding: the policy looked at demand
    ("demand", "supply"),
    ("incentive", "supply"),
    ("incentive", "wait"),
    ("wait", "supply"),
]
for assumption in (
    CausalGraph.from_edges(edges),  # demand was logged, so it can be adjusted for
    CausalGraph.from_edges(edges, latent=("demand",)),  # demand was never logged
):
    out = prescribe(
        panel,
        adjustment=assumption,
        levers=[Lever("incentive", lo=-2.0, hi=2.0, unit_cost=0.05)],
        target=Target("supply", value=1.0),
        constraints=[Constraint("wait", hi=0.5)],
        horizon=15,
        dt=DT,
        tolerance=0.5,
    )
    print(out.report(), end="\n\n")

The same decision, read under two causal assumptions: once with demand logged, once with it declared latent. prescribe composes what is below it and adds no new estimator, solver or guarantee: the adjustment set comes from the graph, the control channel from cross-fit Robinson DML (fit_causal_residual), the plan from causal_plan, the safety price from certify_safety, and how far the plan can be from the best one its own box allows on the fitted model from plan_regret_bound, a diagnostic here: the fitted channel reads the state, so the curvature is sampled. What each certificate bounds, and on what: the scopes.

4. What it printed

# Prescription for `supply`

## Decision

Start: the mean (standard deviation) over 200 units of each unit's last logged state, from period 11: `supply` 0.02264 (0.5971), `wait` 0.02703 (0.323).

| lever | active steps | first | last |
|---|---|---|---|
| `incentive` | 0-14 | +2 | +0.2552 |

Planned task cost: 2.02025.

## Certificate

- identification: **identified** (blocks every non-causal path from ['incentive'] to ['supply'])
- adjusted for: ['demand']
- estimability: **estimable**, the log moved 3 of the channel's 3 directions a state
- channel standard error: 0.008159, summed within 200 groups of `region`, within each of 11 periods, and within both, whichever reads largest (two-way CR1)
- overlap (residualised action variance): 0.2434
- logger check: passed (p = 0.46 over 11 periods); a pass can miss a partial correlation up to 0.079
- noise persistence: lag-1 autocorrelation -0.03 within units (p = 0.0997)
- error tube: **partial**, exact radii, certified horizon 3, global rate
- the tube's noise: the plan's one-step error on the log, standard deviation 0.1458, lag-1 autocorrelation -0.01; with the fitted field's error, the channel's and the drift's, the tube holds the plant at each step with probability 0.95
- each state's own 0.95 radius of the noise at step 3: `supply` 0.4679, `wait` 0.03459
- barrier: certified steps 15 at gamma 1, gamma* inf (marginal sensitivity model)
- solver: converged after 24 accepted steps
- first-order check: stationarity 1.631e-07 (tolerance 1.414e-06, within), constraint violation 0 (tolerance 2.78e-13, within)
- regret bound on the fitted model: 1.687e-13, diagnostic

**Trustworthy prefix: 3 steps.**

## Provenance

- data sha256: `2f5ae7c0e4ef94c1...`
- panel: chc 0.18.0, 2400 rows, x64=True, seed=0
- run: on cpu (cpu), x64=True, fit in float64, solve in float64, seed=0, integrator=rk4
- settings: folds=2, degree=1, channel_degree=1, nuisance_degree=2, ridge=1e-06, steps=10000, tolerance=0.5, hold_constraints=False, max_levers=none
- versions: causal-hybrid-control 0.18.0, jax 0.11.2, jaxlib 0.11.2, numpy 2.5.3, scipy 1.18.1

# Prescription for `supply`

## Decision

Start: the mean (standard deviation) over 200 units of each unit's last logged state, from period 11: `supply` 0.02264 (0.5971), `wait` 0.02703 (0.323).

**No schedule.** no observed set blocks every non-causal path; latent ['demand'] confounds it

## Certificate

- identification: **not_identified** (no observed set blocks every non-causal path; latent ['demand'] confounds it)
- adjusted for: nothing
- channel standard error: not evaluated
- overlap (residualised action variance): 0.9988
- logger check: not run, the levers' parents ['demand'] are not logged
- noise persistence: not read, the channel is not identified
- error tube: **not_evaluated**, certified horizon not evaluated
- barrier: certified steps not evaluated, gamma* not evaluated (marginal sensitivity model)
- solver: not run
- regret bound on the fitted model: not evaluated

**Trustworthy prefix: 0 steps.**

## Provenance

- data sha256: `2f5ae7c0e4ef94c1...`
- panel: chc 0.18.0, 2400 rows, x64=True, seed=0
- run: on cpu (cpu), x64=True, fit in float64, no solve, seed=0, integrator=rk4
- settings: folds=2, degree=1, channel_degree=1, nuisance_degree=2, ridge=1e-06, steps=10000, tolerance=0.5, hold_constraints=False, max_levers=none
- versions: causal-hybrid-control 0.18.0, jax 0.11.2, jaxlib 0.11.2, numpy 2.5.3, scipy 1.18.1

Read the first report by its two axes, which are kept apart on purpose:

  • Identification is about the data and the graph: whether an adjustment set exists, which one was used, and the channel's standard error.
  • Certification is about model error and the barrier: the error tube's certified horizon, over the plan's RK4 steps, at the channel's standard error, a scale rather than a bound, plus the plant's own one-step error on the log, held with probability 0.95 at each step; the steps whose predicted states clear the constraint at the level of the marginal sensitivity model the plan was audited at (gamma, when a bound is given), checked pointwise; and gamma*, the largest level at which every step along the plan's path still has an admissible action that meets the barrier, a ceiling for the problem along that path rather than for the plan.

The trustworthy prefix is the part of the schedule that survives both — the number an operator can act on. A plan can be fully certified against a channel nothing identifies, which is a trustworthy tube around a meaningless action, so neither axis is derived from the other.

The second report is the case this library exists for. With demand never logged, no observed set identifies the effect, so there is no schedule: out.schedule raises NotIdentifiedError rather than hand back actions that would look like every other schedule and mean nothing. prescribe also logs a WARNING on the chc.decision logger when that happens. What is left is to price how much confounding the decision can tolerate — the sensitivity level Γ.

out.to_json() carries the same schedule, both axes and the provenance as plain JSON values.

The expert path

prescribe wraps fit_causal_residual and causal_plan, and both take a hand-built plant, cost and barrier. Below, the model is built by hand: a known oscillator plus a residual fitted to logged transitions. The actions in this log were drawn at random, so plain least squares (fit_residual) is enough; on a confounded log, fit_causal_residual fits the control channel instead. The script is docs/expert_path.py.

import equinox as eqx
import jax
import jax.numpy as jnp

from chc import BarrierConstraint, HybridDynamics, QuadraticCost, RecedingHorizon
from chc.dynamics import DampedOscillator
from chc.integrate import rk4_step
from chc.mpc import mpc_control
from chc.residual import KANResidual
from chc.train import fit_residual

DT = 0.1


class Stiffening(eqx.Module):
    """The physics the known model leaves out: a spring that stiffens as it stretches."""

    def __call__(self, t, x, u):
        return jnp.array([0.0, -0.5 * x[0] ** 3])


known = DampedOscillator(omega=1.0, zeta=0.1)
plant = HybridDynamics(known=known, residual=Stiffening())  # the system itself, unknown to us

# logged transitions (x, u, x_next) from the plant, under actions drawn at random
k_x, k_u, k_r = jax.random.split(jax.random.key(0), 3)
xs = jax.random.normal(k_x, (2000, 2))
us = 0.5 * jax.random.normal(k_u, (2000, 1))
x_next = jax.vmap(lambda x, u: rk4_step(plant, 0.0, x, u, DT))(xs, us)

# the known oscillator plus a learnable (KAN) residual, swappable for an MLP or a linear term
model = HybridDynamics(known=known, residual=KANResidual(2, 1, 2, grid_range=4.0, key=k_r))
model, _ = fit_residual(model, {"x": xs, "u": us, "x_next": x_next}, DT, steps=2000)

mpc_control then re-plans at every step on the fitted model and applies the first action to the plant, which prescribe does not:

cost = QuadraticCost(
    Q=jnp.diag(jnp.array([1.0, 0.1])),
    R=jnp.array([[0.05]]),
    Qf=jnp.diag(jnp.array([5.0, 1.0])),
    x_target=jnp.zeros(2),
)
states, actions = mpc_control(
    model,
    jnp.array([1.0, 0.0]),
    cost,
    dt=DT,
    horizon=20,
    u_lo=-5.0,
    u_hi=5.0,
    n_steps=40,
    plant=plant,
)  # plans on the fitted model, acts on the plant; u_lo/u_hi also take one bound per lever

Online, where the plant is the world rather than a simulation, RecedingHorizon plans from each measured state and returns the whole causal_plan: actions, tube and audit together.

floor = BarrierConstraint(lambda x: x[1] + 0.8, alpha=2.0)  # velocity >= -0.8; built once
controller = RecedingHorizon(model, cost, dt=DT, horizon=20, u_lo=-5.0, u_hi=5.0, barrier=floor)

x = jnp.array([1.0, 0.0])  # the state just measured
plan = controller.step(x)  # every step: read plan.safety, then apply plan.actions[0]

What it printed when this site was built:

mpc_control: 40 steps on the plant, from x = (1, 0) to x = (+0.002, +0.003)
RecedingHorizon.step: first action -0.136, barrier certified on 20 of 20 steps

Each step starts from the last plan and the barrier's multipliers, shifted one step on: 30 % fewer descent steps than cold solves on an oscillator's velocity floor, for the same closed-loop cost to 1e-6 (ADR 0003, receding-horizon warm starts, measured again for the scaled descent in ADR 0059). A program compiles the first time a step needs it and is reused after, so steps stop compiling as long as the model, the cost and the barrier's function stay the same objects: a new lambda per step compiles the descent per step. For a process that restarts, set jax_compilation_cache_dir and lower jax_persistent_cache_min_compile_time_secs to 0: at its default of one second JAX wrote none of a first step's programs to disk.

Where next

  • Concepts — what each part of the report means, and what it does not promise.
  • Tutorials — the executed notebooks, from a first decision to real data.
  • What's inside — every module, with what it does and what was measured.