Speed lanes

Progressive “how do I go faster?” helpers and the single-run array model. See also the user guide Going faster (lanes).

For a view-API model on GPU, prefer MyModel(...).gpu().run() with step_vectorized() (0.4.4) rather than rewriting the step as an ArrayKernelModel. Use ArrayKernelModel when the state is naturally a pure array dict and you do not need the Polars view API.

Status & recommendations

ambr.lanes.status() → Dict[str, Any]

Return a machine/lane status dict (safe to print or log).

ambr.lanes.print_status() → None

Human-readable lane status (for notebooks / CLIs).

ambr.lanes.recommend(n_agents: int, *, ensemble: bool = False) → str

One-line backend suggestion for a typical step-based model.

Heuristic only (device cost and kernel shape matter). Safe defaults that match the public benchmarks’ qualitative breakpoints.

Array kernel model

class ambr.lanes.ArrayKernelModel(parameters: Dict[str, Any] | None = None)

Bases: object

Easiest single-run CPU/GPU array model (no Polars step required).

Subclass and implement:

• init_state(xp, n, rng, p) -> dict of arrays

• step_state(xp, state, rng, p) -> state

• optional metrics(xp, state) -> dict of Python scalars each step

xp is CuPy when a GPU is available, else NumPy — same code path.

Example:

class Drift(am.ArrayKernelModel):
    def init_state(self, xp, n, rng, p):
        return {"x": rng.random(n, dtype=xp.float32)}
    def step_state(self, xp, state, rng, p):
        state["x"] = state["x"] + 0.01
        return state
    def metrics(self, xp, state):
        return {"mean_x": float(am.to_host(state["x"].mean()))}

res = Drift({"n": 100_000, "steps": 50, "seed": 0}).run()
print(res.model)  # per-step metrics if metrics() defined

init_state(xp, n: int, rng, p: Dict[str, Any]) → Dict[str, Any]

metrics(xp, state: Dict[str, Any]) → Dict[str, Any]

Optional per-step scalars (host floats). Default: none.

run(steps: int | None = None, *, prefer_gpu: bool | None = None) → RunResults

Run steps device-resident updates; return RunResults.

step_state(xp, state: Dict[str, Any], rng, p: Dict[str, Any]) → Dict[str, Any]