Raking a biased stream in fixed memory

Rows arrive from a sample that has too many women, young people and college graduates, and an outcome that depends on all three. StreamingRaker weights each row as it arrives and keeps only running totals; BatchIPF rakes the full stored sample at the end, which is the benchmark.

A biased stream

import matplotlib.pyplot as plt
import numpy as np

from onlinerake import BatchIPF, StreamingRaker, Targets

targets = Targets(female=0.5, young=0.3, college=0.35)
names = targets.feature_names
population = np.array([targets[name] for name in names])
sample_shares = {"college": 0.45, "female": 0.65, "young": 0.45}

rng = np.random.default_rng(0)
n = 20_000
rows = np.column_stack([rng.random(n) < sample_shares[name] for name in names]).astype(
    int
)
coef = {"college": 2.0, "female": 1.0, "young": -1.5}
y = rows @ np.array([coef[name] for name in names]) + rng.normal(0, 1, n)
true_mean = float(population @ np.array([coef[name] for name in names]))

Weight the stream as it arrives

raker = StreamingRaker(targets, outcomes=["y"])
checkpoints = np.unique(np.geomspace(10, n, 40).astype(int))
stream_error, raw_error = [], []
for i, (row, yi) in enumerate(zip(rows, y, strict=True), start=1):
    raker.partial_fit(dict(zip(names, row.tolist(), strict=True)), outcomes={"y": yi})
    if i in checkpoints:
        margins = np.array([raker.margins[name] for name in names])
        stream_error.append(np.abs(margins - population).max())
        raw_error.append(np.abs(rows[:i].mean(axis=0) - population).max())

ipf = BatchIPF(targets).fit([dict(zip(names, r.tolist(), strict=True)) for r in rows])
w = ipf.weights

print(f"true mean of y: {true_mean:.3f}")
print(f"{'':<16} {'estimate':>8} {'ESS':>7} {'numbers kept':>13}")
print(f"{'unweighted':<16} {y.mean():8.3f} {n:7d} {'':>13}")
print(
    f"{'StreamingRaker':<16} {raker.estimate('y'):8.3f} "
    f"{raker.effective_sample_size:7.0f} {len(names) + 5:13d}"
)
print(
    f"{'BatchIPF':<16} {w @ y / w.sum():8.3f} "
    f"{w.sum() ** 2 / (w @ w):7.0f} {rows.size + n:13d}"
)

true mean of y: 0.750
                 estimate     ESS  numbers kept
unweighted          0.851   20000
StreamingRaker      0.744   14903             8
BatchIPF            0.745   16011         80000

Margin error as the stream grows

fig, ax = plt.subplots(figsize=(7, 3.5))
ax.plot(checkpoints, raw_error, label="unweighted")
ax.plot(checkpoints, stream_error, label="StreamingRaker")
ax.set_xscale("log")
ax.set_yscale("log")
ax.set_xlabel("rows seen")
ax.set_ylabel("largest margin error")
ax.legend()
fig.tight_layout()
plt.show()

Design weights and trimming

Rows can arrive with design weights in any units. Heavy-tailed design weights make the estimate noisy; trim caps each weight at a multiple of the mean weight so far, trading a little bias for a larger effective sample.

design = 5000 * rng.lognormal(0.0, 1.5, n)
for trim in (None, 3.0):
    weighted = StreamingRaker(targets, trim=trim, outcomes=["y"])
    for row, yi, d in zip(rows, y, design, strict=True):
        weighted.partial_fit(
            dict(zip(names, row.tolist(), strict=True)),
            design_weight=float(d),
            outcomes={"y": yi},
        )
    print(
        f"trim={trim}: estimate {weighted.estimate('y'):.3f}, "
        f"ESS {weighted.effective_sample_size:.0f}"
    )

trim=None: estimate 0.757, ESS 5166
trim=3.0: estimate 0.754, ESS 9447