Variable-window calculation with numba @njit

A row‑by‑row calculation with a window that changes per row cannot be vectorized, so in plain pandas it falls back to slow Python. Here the window length is data‑dependent — a volume‑weighted mean over the last 2×(number of “H” flags in the last 20 rows) rows — which pandas rolling cannot express. Write the loop once, then numba.njit compiles that same function: it returns a new dispatcher without touching the original, so the Python twin stays callable for a fair, bit‑identical comparison. On 6 million rows the plain Python loop takes about 109 s; the compiled version takes about 0.3 s — a speed‑up of more than 350×, for the same result.

import time
import numpy as np
from numba import njit

# Volume-weighted mean over a PER-ROW-VARIABLE window: the window length
# depends on recent activity (twice the number of "H" flags in the last 20
# rows) -- exactly what pandas `rolling` cannot express, but a loop can.
def f_weighted_mean(mid, bid_vol, ask_vol, flag):
    n = mid.shape[0]
    out = np.empty(n, dtype=np.float32)
    for i in range(n):
        lo20 = i - 19
        if lo20 < 0:
            lo20 = 0
        h = 0
        for j in range(lo20, i + 1):
            if flag[j] == b"H":
                h += 1
        w = 2 * h                           # window length driven by activity
        if w == 0:
            out[i] = np.float32(mid[i])     # no "H" nearby -> current mid
            continue
        lo = i - w + 1
        if lo < 0:
            lo = 0
        num = 0.0
        den = 0.0
        for j in range(lo, i + 1):
            size = np.float64(bid_vol[j] + ask_vol[j])
            num += mid[j] * size
            den += size
        out[i] = np.float32(num / den)
    return out

# ONE source of truth. njit(fn) returns a compiled dispatcher WITHOUT
# mutating fn, so the Python original stays callable for the comparison.
f_weighted_mean_jit = njit(f_weighted_mean)

# order-book-like data (6M rows; the plain-Python pass below takes ~2 min)
rng = np.random.default_rng(0)
n = 6_000_000
mid = 100 + rng.standard_normal(n).cumsum() * 0.01
bid_vol = rng.integers(1, 500, n)
ask_vol = rng.integers(1, 500, n)
flag = rng.choice([b"H", b"V"], n).astype("S1")

out_jit = f_weighted_mean_jit(mid, bid_vol, ask_vol, flag)  # compiles on first call
out_py  = f_weighted_mean(mid, bid_vol, ask_vol, flag)
assert np.array_equal(out_jit, out_py)                      # bit-identical

def best(f, runs=3):
    times = []
    for _ in range(runs):
        t0 = time.perf_counter()
        f(mid, bid_vol, ask_vol, flag)
        times.append(time.perf_counter() - t0)
    return min(times)

print(f"python : {best(f_weighted_mean,     runs=1):8.3f} s")   # one pass is enough
print(f"njit   : {best(f_weighted_mean_jit, runs=3):8.3f} s")
# python :  108.741 s
# njit   :    0.305 s   -> ~356x faster, same result

Exact numbers vary by machine, and the first njit call compiles (so it is warmed up before timing) — but the gap is the point: the compiled loop is hundreds of times faster (real number: 356×) than the identical Python one, with no change to the logic you wrote.