Order flow imbalance: does signed volume predict returns?

quotes = cu.make_quotes(symbols=["AAPL", "MSFT"], days=2)
print(f"quotes: {quotes.num_rows:,} rows x {quotes.num_columns} columns")
quotes.to_pandas().head()

q = quotes.to_pandas()

rng = np.random.default_rng(21)
picks = q[rng.random(len(q)) < 0.25].copy()
side = rng.choice(np.array([1, -1]), size=len(picks))
picks["price"] = np.where(side > 0, picks["ask"], picks["bid"])
picks["ts"] = (
    picks["ts"] + pd.to_timedelta(rng.integers(1_000, 50_000, len(picks)), unit="us")
).astype("datetime64[us, UTC]")
picks["size"] = np.maximum(100, (rng.lognormal(4.0, 1.2, len(picks)) // 100 * 100)).astype("int64")
picks["side"] = np.where(side > 0, "B", "S")

trade_schema = pa.schema(
    [
        pa.field("ts", pa.timestamp("us", tz="UTC"), nullable=False),
        pa.field("symbol", pa.string()),
        pa.field("price", pa.float64()),
        pa.field("size", pa.int64()),
        pa.field("side", pa.string()),
    ]
)
trades_df = picks[["ts", "symbol", "price", "size", "side"]].sort_values(["ts", "symbol"])
trades = pa.Table.from_pandas(trades_df, preserve_index=False).cast(trade_schema)
print(f"{len(trades_df):,} trades printed off {len(q):,} quotes")
trades_df.head()

db.create_table("quotes", quotes.schema, time_column="ts", sort_key=["ts", "symbol"])
db.append("quotes", quotes, note="2-day synthetic NBBO")
db.create_table("trades", trade_schema, time_column="ts", sort_key=["ts", "symbol"])
db.append("trades", trades, note="tape printed off the quote stream")
print({t: len(db.read(t)) for t in db.tables()})

w0 = int(pd.Timestamp("2026-06-01 16:00", tz="UTC").value // 1000)
w1 = int(pd.Timestamp("2026-06-01 16:15", tz="UTC").value // 1000)

trades_w = db.read("trades", time_start=w0, time_end=w1)
quotes_w = db.read("quotes", time_start=w0 - 60_000_000, time_end=w1)  # 60s run-up
db.create_table("trades_qa", trade_schema, time_column="ts", sort_key=["ts", "symbol"])
db.append("trades_qa", trades_w, note="midday QA window")
db.create_table("quotes_qa", quotes.schema, time_column="ts", sort_key=["ts", "symbol"])
db.append("quotes_qa", quotes_w, note="midday QA window + run-up")
print(f"QA window: {len(trades_w):,} trades, {len(quotes_w):,} quotes")

def sign_of(x, prev):
    """+1 uptick, -1 downtick, 0 unchanged - the tick test."""
    return when(x > prev).then(lit(1)).when(x < prev).then(lit(-1)).otherwise(lit(0))


PREV_PX = sql_expr("lag(price)").over(partition_by="symbol", order_by="ts")
TRUE_SIGN = when(col("side") == "B").then(lit(1)).otherwise(lit(-1))

audit = (
    db.table("trades_qa")
    .join_asof(db.table("quotes_qa"), on="ts", by="symbol", tolerance=5_000_000)
    .select("ts", "symbol", "price", "size", "side",
            mid=(col("bid") + col("ask")) / 2)
    .with_columns(prev_px=PREV_PX)
    .select(
        "ts", "symbol", "mid",
        lr_sign=when(col("price") > col("mid")).then(lit(1))
        .when(col("price") < col("mid")).then(lit(-1))
        .when(col("price") > col("prev_px")).then(lit(1))
        .when(col("price") < col("prev_px")).then(lit(-1))
        .otherwise(lit(0)),
        tick_sign=sign_of(col("price"), col("prev_px")),
        true_sign=TRUE_SIGN,
    )
    .sort(["ts", "symbol"])
    .to_pandas()
)
assert len(audit) == len(trades_w), "asof_join must return one row per trade"

# cross-check the ASOF lookup itself against pandas.merge_asof
ref = pd.merge_asof(
    trades_w.to_pandas().sort_values("ts"),
    quotes_w.to_pandas().sort_values("ts").assign(mid_ref=lambda d: (d["bid"] + d["ask"]) / 2),
    on="ts",
    by="symbol",
    tolerance=pd.Timedelta("5s"),
)
chk = audit.sort_values("ts").reset_index(drop=True)
assert np.allclose(chk["mid"], ref["mid_ref"], equal_nan=True), "asof mismatch vs pandas"
print("asof_join validated against pandas.merge_asof")

for rule in ("lr_sign", "tick_sign"):
    scored = audit[audit[rule] != 0]
    acc = (scored[rule] == scored["true_sign"]).mean()
    print(f"  {rule:9s}: {acc:.1%} accuracy vs true aggressor side ({len(scored):,} classified)")

signed = (
    db.table("trades")
    .with_columns(prev_px=PREV_PX)
    .select(
        "ts", "symbol", "price", "size",
        tick_sign=sign_of(col("price"), col("prev_px")),
        true_sign=TRUE_SIGN,
    )
    .sort(["ts", "symbol"])
    .to_arrow()
)

full = signed.to_pandas()
scored = full[full["tick_sign"] != 0]
print(f"full tape: {(scored['tick_sign'] == scored['true_sign']).mean():.1%} tick-test accuracy on {len(scored):,} prints")

signed_schema = pa.schema(
    [
        pa.field("ts", pa.timestamp("us", tz="UTC"), nullable=False),
        pa.field("symbol", pa.string()),
        pa.field("price", pa.float64()),
        pa.field("size", pa.int64()),
        pa.field("tick_sign", pa.int32()),
        pa.field("true_sign", pa.int32()),
    ]
)
db.create_table("signed_trades", signed_schema, time_column="ts", sort_key=["ts", "symbol"])
db.append("signed_trades", signed.cast(signed_schema), note="tick-test + oracle signs")

BAR = time_bucket("1m", col("ts"))
size_when = lambda flag: when(flag).then(col("size")).otherwise(lit(0)).sum()

trade_ofi = (
    db.table("signed_trades")
    .group_by(BAR.alias("bar"), "symbol")
    .agg(
        buy_vol=size_when(col("tick_sign") > 0),
        sell_vol=size_when(col("tick_sign") < 0),
        oracle_net=when(col("true_sign") > 0).then(col("size")).otherwise(-col("size")).sum(),
        volume=col("size").sum(),
        close=col("price").last("ts"),
    )
    .sort(["bar", "symbol"])
    .to_pandas()
)
trade_ofi["ofi"] = (trade_ofi["buy_vol"] - trade_ofi["sell_vol"]) / trade_ofi["volume"]
trade_ofi["ofi_oracle"] = trade_ofi["oracle_net"] / trade_ofi["volume"]
trade_ofi.head(4)

def prev(name: str):
    return sql_expr(f"lag({name})").over(partition_by="symbol", order_by="ts")


# Cont et al.: size added at the bid counts positive, size pulled counts
# negative, mirrored on the ask. Four lagged comparisons per quote.
keep = lambda flag, x: when(flag).then(x).otherwise(lit(0))

quote_ofi = (
    db.table("quotes")
    .with_columns(pb=prev("bid"), pa=prev("ask"), pbs=prev("bid_size"), pas=prev("ask_size"))
    .filter(col("pb").is_not_null())
    .group_by(BAR.alias("bar"), "symbol")
    .agg(
        qofi=(
            keep(col("bid") >= col("pb"), col("bid_size"))
            - keep(col("bid") <= col("pb"), col("pbs"))
            - keep(col("ask") <= col("pa"), col("ask_size"))
            + keep(col("ask") >= col("pa"), col("pas"))
        ).sum(),
        mid_close=((col("bid") + col("ask")) / 2).last("ts"),
    )
    .sort(["bar", "symbol"])
    .to_pandas()
)
quote_ofi.head(4)

def bucket_corrs(frame, ofi_col, px_col):
    rows = []
    frame = frame.copy()
    frame["day"] = frame["bar"].dt.date
    for (sym, _), g in frame.groupby(["symbol", "day"]):
        g = g.sort_values("bar")
        frame.loc[g.index, "ret"] = g[px_col].pct_change()
        frame.loc[g.index, "ret_next"] = g[px_col].pct_change().shift(-1)
    for sym, g in frame.groupby("symbol"):
        g = g.dropna(subset=["ret", "ret_next"])
        rows.append(
            {
                "symbol": sym,
                "n_buckets": len(g),
                "corr_same_min": g[ofi_col].corr(g["ret"]),
                "corr_next_min": g[ofi_col].corr(g["ret_next"]),
            }
        )
    return pd.DataFrame(rows)

results = pd.concat(
    [
        bucket_corrs(trade_ofi, "ofi", "close").assign(estimator="trade OFI (tick test)"),
        bucket_corrs(trade_ofi, "ofi_oracle", "close").assign(estimator="trade OFI (oracle)"),
        bucket_corrs(quote_ofi, "qofi", "mid_close").assign(estimator="quote OFI (Cont)"),
    ]
)[["estimator", "symbol", "n_buckets", "corr_same_min", "corr_next_min"]]
results.round(3)

import matplotlib.pyplot as plt

sym = "AAPL"
g = trade_ofi[trade_ofi["symbol"] == sym].copy()
g["day"] = g["bar"].dt.date
g["ret"] = g.groupby("day")["close"].pct_change()
g = g.dropna(subset=["ret"])

fig, ax = plt.subplots(figsize=(7.5, 5))
ax.scatter(g["ofi"], g["ret"] * 1e4, s=8, alpha=0.35, color="tab:blue")
b, a = np.polyfit(g["ofi"], g["ret"] * 1e4, 1)
xs = np.linspace(-1, 1, 50)
ax.plot(xs, a + b * xs, color="tab:red", lw=1.5, label=f"OLS slope = {b:.1f} bps per unit OFI")
ax.set_title(f"{sym}: 1-minute trade OFI vs same-minute return")
ax.set_xlabel("OFI = (buy vol - sell vol) / total vol")
ax.set_ylabel("1-minute return (bps)")
ax.legend()
fig.tight_layout()