Isotonic Regression on GPU

"""Isotonic regression on a GPU, by hand: one thread per row, the hull of running totals, a
lane's stack in a register, bridges, the price of copying, and torchisotonic on real work."""
import os, sys, json, time, glob, subprocess, datetime
import numpy as np, torch
import matplotlib.pyplot as plt

assert torch.cuda.is_available(), "No GPU: Settings → Accelerator → GPU T4 x2"
torch.cuda.set_device(0)
try:
    import cupy as cp
except ImportError:
    subprocess.run([sys.executable, "-m", "pip", "install", "-q", "cupy-cuda12x"], check=True)
    import cupy as cp
subprocess.run([sys.executable, "-m", "pip", "install", "-q", "torchisotonic==0.2.1"], check=True)
import numba
from sklearn.isotonic import isotonic_regression, IsotonicRegression

OUT = {"meta": {"gpu": torch.cuda.get_device_name(0), "torch": torch.__version__,
                "cupy": cp.__version__,
                "executed": datetime.datetime.now(datetime.timezone.utc).isoformat(timespec="seconds")}}
print(OUT["meta"])


def gpu_ms(fn, repeats=5, warmup=True):
    """Median wall-clock milliseconds of fn(), waiting for the GPU after every call. An idle GPU
    runs at a low clock and takes a while to raise it, so it is kept busy for a moment first:
    otherwise short calls are timed at a fraction of the speed they run at under load."""
    def sync():
        cp.cuda.Device().synchronize()
        torch.cuda.synchronize()
    if warmup:
        fn(); sync()
    torch.cuda._sleep(1_000_000_000); sync()
    times = []
    for _ in range(repeats):
        start = time.perf_counter()
        fn(); sync()
        times.append((time.perf_counter() - start) * 1e3)
    return float(np.median(times))

PAVA_ROWS = cp.RawKernel(r'''
extern "C" __global__
void pava_rows(const float* y, float* out, double* sums, int* sizes, int rows, int n) {
    int r = blockIdx.x * blockDim.x + threadIdx.x;      // one thread, one whole row
    if (r >= rows) return;
    const float* yr = y + (long long)r * n;
    double* S = sums + (long long)r * n;                 // the stack of blocks: their sums
    int* C = sizes + (long long)r * n;                   // and their sizes
    int top = 0;
    for (int i = 0; i < n; ++i) {
        double s = yr[i];
        int c = 1;
        // pool while the new block's mean is below the mean of the block before it
        while (top > 0 && s * C[top - 1] < S[top - 1] * c) {
            --top;
            s += S[top];
            c += C[top];
        }
        S[top] = s; C[top] = c; ++top;
    }
    long long k = (long long)r * n;
    for (int b = 0; b < top; ++b) {
        float m = (float)(S[b] / C[b]);
        for (int j = 0; j < C[b]; ++j) out[k++] = m;
    }
}
''', "pava_rows")


def noisy(rows, n, seed=0):
    """Rows of a rising trend plus noise, float32: the data of the note's first table."""
    rng = np.random.default_rng(seed)
    return np.linspace(0, 3, n, dtype=np.float32) + rng.standard_normal((rows, n), dtype=np.float32)


def pava_rows(y):
    """One thread per row; returns the fit and a function that reruns the kernel alone."""
    rows, n = y.shape
    out = cp.empty_like(y)
    sums, sizes = cp.empty((rows, n), cp.float64), cp.empty((rows, n), cp.int32)
    threads = 128
    launch = lambda: PAVA_ROWS(((rows + threads - 1) // threads,), (threads,),
                               (y, out, sums, sizes, np.int32(rows), np.int32(n)))
    launch()
    return out, launch


# first, is it right? against scikit-learn's pool adjacent violators, row by row
y = noisy(64, 1000)
fit, _ = pava_rows(cp.asarray(y))
err = max(np.abs(fit.get()[i] - isotonic_regression(y[i])).max() for i in range(64))
print(f"largest difference from scikit-learn on 64 rows of 1,000: {err:.1e}")

N = 1 << 23
shapes = [(1 << k, N >> k) for k in (19, 16, 13, 10, 7, 4, 2, 0)]
part1 = []
for rows, n in shapes:
    y = cp.asarray(noisy(rows, n))
    _, launch = pava_rows(y)
    long_rows = n >= (1 << 21)
    ms = gpu_ms(launch, repeats=1 if long_rows else 5, warmup=not long_rows)
    part1.append({"rows": rows, "n": n, "ms": ms})
    print(f"{rows:>9,} rows of {n:>9,}   {ms:9.1f} ms   threads with work: {rows:,}")
    del y

row = noisy(1, N)[0]
t = time.perf_counter(); isotonic_regression(row); cpu_ms = (time.perf_counter() - t) * 1e3
print(f"\nscikit-learn, one CPU core, one row of {N:,}: {cpu_ms:.0f} ms")
OUT["part1"] = {"shapes": part1, "cpu_ms": cpu_ms}

fig, ax = plt.subplots()
ax.loglog([p["n"] for p in part1], [p["ms"] for p in part1], "o-", label="GPU, one thread per row")
ax.axhline(cpu_ms, color="0.4", ls="--", lw=1, label="one CPU core, one row of all of them")
ax.set_xlabel("row length (same 8.4 million values in every shape)")
ax.set_ylabel("ms")
ax.legend(fontsize=8)
plt.show()

@numba.njit
def lower_hull(P):
    """Indices of the lower convex hull of the points (k, P[k]), by the monotone chain."""
    st = np.empty(len(P), np.int64)
    top = 0
    for k in range(len(P)):
        while top >= 2:
            a, b = st[top - 2], st[top - 1]
            if (b - a) * (P[k] - P[a]) - (P[b] - P[a]) * (k - a) <= 0:
                top -= 1                    # b is on or above the line from a to k: let it go
            else:
                break
        st[top] = k
        top += 1
    return st[:top]


def fit_from_string(y):
    """The isotonic fit as the slopes of the string under the running totals."""
    P = np.concatenate([[0.0], np.cumsum(y, dtype=np.float64)])
    h = lower_hull(P)
    return np.repeat(np.diff(P[h]) / np.diff(h), np.diff(h)), h, P


print("1, 3, 2, 4  ->", fit_from_string(np.array([1.0, 3, 2, 4]))[0])
y = noisy(1, 100_000, seed=1)[0].astype(np.float64)
print("largest difference from pool adjacent violators on 100,000 values:",
      f"{np.abs(fit_from_string(y)[0] - isotonic_regression(y)).max():.1e}")

y = noisy(1, 24, seed=3)[0].astype(np.float64)
fit, h, P = fit_from_string(y)
fig, (a1, a2) = plt.subplots(2, 1, figsize=(7, 5), sharex=True)
a1.bar(np.arange(24) + 0.5, y, width=0.8, color="0.8", label="values")
a1.step(np.arange(25), np.append(fit, fit[-1]), where="post", color="C3", label="fit")
a1.legend(fontsize=8)
a2.plot(np.arange(25), P, "o", color="0.3", ms=4, label="running totals (nails)")
a2.plot(h, P[h], "-", color="C3", label="the string: lower convex hull")
a2.legend(fontsize=8)
plt.show()

n = 1_000_000
rising = np.linspace(1.0, 2.0, n)           # strictly increasing: every running total bends up
counts = {"noise": len(lower_hull(np.concatenate([[0.0], np.cumsum(noisy(1, n, seed=2)[0], dtype=np.float64)]))),
          "rising": len(lower_hull(np.concatenate([[0.0], np.cumsum(rising)])))}
for k, v in counts.items():
    print(f"{k:7s} {v:>10,} of {n + 1:,} points are on the string")
OUT["part2"] = counts

LANE_SRC = r'''
__device__ double turn(const double* P, int lanes, int t, int a, int b, int k) {
    double pa = P[a * lanes + t], pb = P[b * lanes + t], pk = P[k * lanes + t];
    return (double)(b - a) * (pk - pa) - (pb - pa) * (double)(k - a);
}

extern "C" __global__ void hull_stack(const double* P, unsigned* masks, int lanes) {
    int t = blockIdx.x * blockDim.x + threadIdx.x;
    if (t >= lanes) return;
    int st[32];                                  // the stack, an array of point numbers
    int top = 0;
    for (int k = 0; k < 32; ++k) {
        while (top >= 2 && turn(P, lanes, t, st[top - 2], st[top - 1], k) <= 0) --top;
        st[top++] = k;
    }
    unsigned m = 0;
    for (int i = 0; i < top; ++i) m |= 1u << st[i];
    masks[t] = m;
}

extern "C" __global__ void hull_mask(const double* P, unsigned* masks, int lanes) {
    int t = blockIdx.x * blockDim.x + threadIdx.x;
    if (t >= lanes) return;
    unsigned mask = 1u;                          // the stack, one bit per point: point 0 is on it
    int a = -1, b = 0;                           // the top b and the entry under it, a (-1: none)
    int k = 1;
    while (k < 32) {
        if (a >= 0 && turn(P, lanes, t, a, b, k) <= 0) {      // pop b
            mask &= ~(1u << b);
            b = a;
            unsigned below = mask & ((1u << a) - 1u);
            a = below ? 31 - __clz(below) : -1;               // the highest set bit below a
        } else {                                               // push k
            mask |= 1u << k;
            a = b; b = k; ++k;
        }
    }
    masks[t] = mask;
}
'''
hull_stack = cp.RawKernel(LANE_SRC, "hull_stack")
hull_mask = cp.RawKernel(LANE_SRC, "hull_mask")
hull_stack.compile(); hull_mask.compile()
for name, k in (("hull_stack", hull_stack), ("hull_mask", hull_mask)):
    print(f"{name:11s} registers {k.num_regs:3d}   local memory per thread {k.local_size_bytes:4d} bytes")

def lane_points(y):
    """Running totals of each lane's 32 values, laid out point-major for coalesced reads."""
    P = np.cumsum(y.reshape(-1, 32), axis=1, dtype=np.float64)    # lane-local running totals
    return cp.asarray(np.ascontiguousarray(P.T)), P


def cpu_masks(P, lanes):
    out = np.empty(lanes, np.uint32)
    for t in range(lanes):
        out[t] = np.bitwise_or.reduce(np.left_shift(np.uint32(1), lower_hull(P[t]).astype(np.uint32)))
    return out


data = {"noise": noisy(1, N, seed=4)[0], "rising": np.linspace(1.0, 2.0, N)}
part3 = {}
for name, y in data.items():
    Pd, P = lane_points(y)
    lanes = N // 32
    masks = {k: cp.empty(lanes, cp.uint32) for k in ("stack", "mask")}
    grid, threads = ((lanes + 255) // 256,), (256,)
    run_stack = lambda: hull_stack(grid, threads, (Pd, masks["stack"], np.int32(lanes)))
    run_mask = lambda: hull_mask(grid, threads, (Pd, masks["mask"], np.int32(lanes)))
    ms_stack, ms_mask = gpu_ms(run_stack, 20), gpu_ms(run_mask, 20)
    sample = cpu_masks(P[:2000], 2000)
    ok = bool((masks["stack"][:2000].get() == sample).all() and (masks["mask"].get() == masks["stack"].get()).all())
    vertices = float(np.mean([bin(int(m)).count("1") for m in masks["mask"][:4096].get()]))
    part3[name] = {"stack_ms": ms_stack, "mask_ms": ms_mask, "agree": ok, "mean_vertices": vertices}
    print(f"{name:6s} stack in an array {ms_stack:6.2f} ms   in a bit mask {ms_mask:6.2f} ms   "
          f"({ms_stack / ms_mask:.2f}x)   agree with the CPU: {ok}   vertices per lane: {vertices:.1f}")
    m = int(masks["mask"][7].get())
    print(f"       lane 7's stack when it is done: {m:032b}  (bit 0 on the right)")
OUT["part3"] = part3

@numba.njit
def turn3(P, a, b, c):
    return (b - a) * (P[c] - P[a]) - (P[b] - P[a]) * (c - a)


@numba.njit
def glue(P, L, R):
    """Merge two neighbouring hulls (index arrays sharing their boundary point); returns the
    merged hull and the number of nails the walk to the bridge let go."""
    nl, nr = len(L), len(R)
    if nl < 2 or nr < 2 or turn3(P, L[nl - 2], L[nl - 1], R[1]) > 0:
        return np.concatenate((L, R[1:])), 0          # already bends upwards: nothing to search
    a, b, steps = nl - 2, 1, 1                        # the shared point itself is the first to go
    moved = True
    while moved:
        moved = False
        while a > 0 and turn3(P, L[a - 1], L[a], R[b]) <= 0:
            a -= 1; steps += 1; moved = True
        while b < nr - 1 and turn3(P, L[a], R[b], R[b + 1]) <= 0:
            b += 1; steps += 1; moved = True
    return np.concatenate((L[:a + 1], R[b:])), steps


def merge_tree(y, chunk=1024):
    P = np.concatenate([[0.0], np.cumsum(y, dtype=np.float64)])
    n = len(y)
    nodes = [np.arange(j, min(j + chunk, n) + 1)[lower_hull(P[j:min(j + chunk, n) + 1])] for j in range(0, n, chunk)]
    walks = []                                         # per level: the steps of every merge
    while len(nodes) > 1:
        merged, level = [], []
        for k in range(0, len(nodes) - 1, 2):
            h, steps = glue(P, nodes[k], nodes[k + 1])
            merged.append(h); level.append(steps)
        if len(nodes) % 2:
            merged.append(nodes[-1])
        nodes = merged
        walks.append(level)
    return walks, len(nodes[0])


n = 1 << 22
cliff = np.linspace(1.0, 2.0, n); cliff[-1] = -1e9    # rising, and one value far below at the end
part4 = {}
for name, y in (("noise", noisy(1, n, seed=5)[0].astype(np.float64)), ("rising", np.linspace(1.0, 2.0, n)), ("cliff", cliff)):
    walks, final = merge_tree(y)
    every = np.concatenate([np.asarray(w) for w in walks])
    part4[name] = {"merges": len(every), "max_steps": int(every.max()), "over_64": int((every > 64).sum()),
                   "last_level_steps": int(walks[-1][0]), "final_vertices": final}
    print(f"{name:6s} {len(every):,} merges   longest walk {every.max():>9,} steps   "
          f"walks over 64 steps: {(every > 64).sum():,}   final hull: {final:,} vertices")
OUT["part4"] = part4

n = 10_000_000
src = [cp.random.random(n), cp.arange(n, dtype=cp.int32), cp.arange(n, dtype=cp.int32) // 1024]
dst = [cp.empty_like(a) for a in src]
level_ms = gpu_ms(lambda: [cp.copyto(d, s) for d, s in zip(dst, src)], 10)
print(f"one level, {n * 16 / 1e6:.0f} MB rewritten with the whole GPU: {level_ms:.2f} ms "
      f"({2 * n * 16 / level_ms / 1e6:.0f} GB/s read + write)")
print(f"14 levels: {14 * level_ms:.0f} ms of copying (the paper's whole grid-copy merge stage: 25.1 ms)")

COPY = cp.RawKernel(r'''
extern "C" __global__
void copy_vertices(const double* h, const int* p, const int* c, double* h2, int* p2, int* c2, long long n) {
    long long stride = (long long)gridDim.x * blockDim.x;
    for (long long i = blockIdx.x * (long long)blockDim.x + threadIdx.x; i < n; i += stride) {
        h2[i] = h[i]; p2[i] = p[i]; c2[i] = c[i];
    }
}
''', "copy_vertices")

m = 1_611_392
args = (*[a[:m] for a in src], *[a[:m] for a in dst], np.int64(m))
one_warp = gpu_ms(lambda: COPY((1,), (32,), args), 5)
whole_gpu = gpu_ms(lambda: COPY((40 * 8,), (256,), args), 20)
desc = cp.random.random(1574 * 3)                       # 1,574 descriptors of 24 bytes: the same merge, listed
desc_ms = gpu_ms(lambda: cp.copyto(cp.empty_like(desc), desc), 20)
print(f"{m:,} vertices ({m * 16 / 1e6:.1f} MB): one warp {one_warp:.2f} ms, the whole GPU {whole_gpu:.3f} ms "
      f"({one_warp / whole_gpu:.0f}x)")
print(f"the same merge with segment lists moves 1,574 descriptors (38 KB): {desc_ms:.3f} ms")
OUT["part5"] = {"level_ms": level_ms, "levels14_ms": 14 * level_ms, "one_warp_ms": one_warp,
                "whole_gpu_ms": whole_gpu, "descriptors_ms": desc_ms}

import torchisotonic
t = time.perf_counter()
torchisotonic.isotonic(torch.randn(4, 100, device="cuda"))
torch.cuda.synchronize()
print(f"first call, which builds the CUDA operator: {time.perf_counter() - t:.0f} s")

y = torch.randn(256, 1000, dtype=torch.float64, device="cuda").cumsum(-1) * 0.05 + torch.randn(256, 1000, dtype=torch.float64, device="cuda")
w = torch.rand(256, 1000, dtype=torch.float64, device="cuda") + 0.1
fit = torchisotonic.isotonic(y, w)
ref = np.stack([isotonic_regression(r, sample_weight=s) for r, s in zip(y.cpu().numpy(), w.cpu().numpy())])
print(f"weighted, float64, 256 rows of 1,000: largest difference from scikit-learn {np.abs(fit.cpu().numpy() - ref).max():.1e}")

part6 = []
for rows, n in [(1 << 19, 16), (1, N)]:
    y = torch.as_tensor(noisy(rows, n), device="cuda")
    ms = gpu_ms(lambda: torchisotonic.isotonic(y), 20)
    mine = next(p["ms"] for p in OUT["part1"]["shapes"] if p["rows"] == rows)
    part6.append({"rows": rows, "n": n, "ms": ms, "one_thread_ms": mine})
    print(f"{rows:>9,} rows of {n:>9,}   torchisotonic {ms:7.2f} ms   one thread per row (Part 1) {mine:9.1f} ms")
y = (torch.arange(N, dtype=torch.int32, device="cuda") + 0x3F800000).view(torch.float32)[None]   # consecutive floats from 1.0
ms = gpu_ms(lambda: torchisotonic.isotonic(y), 20)
part6.append({"rows": 1, "n": N, "rising": True, "ms": ms})
print(f"        1 row  of {N:>9,}, rising: every point a vertex   torchisotonic {ms:7.2f} ms")
OUT["part6"] = part6

y = torch.tensor([[1.0, 3.0, 2.0, 4.0]], requires_grad=True)
fit = torchisotonic.isotonic(y)
fit[0, 1].backward()
print("fit:", fit.detach().numpy()[0], "  d fit[1] / d y:", y.grad.numpy()[0])

import pyarrow as pa, pyarrow.csv as pacsv

path = next(iter(glob.glob("/kaggle/input/**/dac/train.txt", recursive=True)), None)
assert path, "add the dataset mrkmakr/criteo-dataset under Input"
ROWS, BUCKETS, KEEP = 7_000_000, 1 << 18, 0.1
DENSE, CATS = [f"I{i}" for i in range(1, 14)], [f"C{i}" for i in range(1, 27)]

HEX = np.zeros(256, np.uint32)
for i, ch in enumerate(b"0123456789abcdef"):
    HEX[ch] = i
for i, ch in enumerate(b"ABCDEF"):
    HEX[ch] = 10 + i


def hex_codes(arr):
    """The 32-bit values of an Arrow array of 8-digit hex strings (0 where missing), vectorized."""
    n, out, buffers = len(arr), np.zeros(len(arr), np.uint32), arr.buffers()
    if buffers[2] is None:
        return out
    offsets = np.frombuffer(buffers[1], np.int32, count=arr.offset + n + 1)[arr.offset:]
    data = np.frombuffer(buffers[2], np.uint8)
    full = np.diff(offsets) == 8
    nibbles = HEX[data[offsets[:-1][full, None] + np.arange(8)]]
    out[full] = np.bitwise_or.reduce(nibbles << np.arange(28, -1, -4, dtype=np.uint32), axis=1)
    return out


t = time.perf_counter()
reader = pacsv.open_csv(path, read_options=pacsv.ReadOptions(column_names=["label", *DENSE, *CATS], block_size=1 << 26),
                        parse_options=pacsv.ParseOptions(delimiter="\t"),
                        convert_options=pacsv.ConvertOptions(column_types={"label": pa.int8(), **{c: pa.float32() for c in DENSE}, **{c: pa.string() for c in CATS}},
                                                             strings_can_be_null=True))
labels = np.empty(ROWS, np.float32); dense = np.empty((ROWS, 13), np.float32); cats = np.empty((ROWS, 26), np.int64)
got = 0
for batch in reader:
    m = min(batch.num_rows, ROWS - got)
    if m <= 0:
        break
    batch = batch.slice(0, m)
    labels[got:got + m] = batch.column(0).to_numpy(zero_copy_only=False)
    for j in range(13):
        dense[got:got + m, j] = np.log1p(np.maximum(np.nan_to_num(batch.column(1 + j).to_numpy(zero_copy_only=False)), 0))
    for j in range(26):
        cats[got:got + m, j] = hex_codes(batch.column(14 + j)) % BUCKETS
    got += m
print(f"read {got:,} impressions in {time.perf_counter() - t:.0f} s; click rate {labels.mean():.3f}")

dev = torch.device("cuda")
day = ROWS // 7
torch.manual_seed(0); rng = np.random.default_rng(0)
train = np.arange(5 * day)
train = train[(labels[train] == 1) | (rng.random(len(train)) < KEEP)]

emb = torch.nn.Embedding(26 * BUCKETS, 1).to(dev); torch.nn.init.zeros_(emb.weight)
lin = torch.nn.Linear(13, 1).to(dev)
field = torch.arange(26, device=dev) * BUCKETS
logit = lambda d, c: emb(c + field).sum((1, 2)) + lin(d).squeeze(1)
opt = torch.optim.Adam([*emb.parameters(), *lin.parameters()], lr=1e-2)
D, Cc, Y = (torch.as_tensor(a[train], device=dev) for a in (dense, cats, labels))
t = time.perf_counter()
for epoch in range(2):
    for i in torch.randperm(len(Y), device=dev).split(1 << 14):
        opt.zero_grad(set_to_none=True)
        torch.nn.functional.binary_cross_entropy_with_logits(logit(D[i], Cc[i]), Y[i]).backward()
        opt.step()
torch.cuda.synchronize()
print(f"trained on {len(Y):,} impressions in {time.perf_counter() - t:.0f} s")

@torch.no_grad()
def scores(rows):
    d, c = torch.as_tensor(dense[rows], device=dev), torch.as_tensor(cats[rows], device=dev)
    return torch.sigmoid(logit(d, c).double())

cal, test = slice(5 * day, 6 * day), slice(6 * day, 7 * day)
s_cal, s_test = scores(cal), scores(test)
y_cal, y_test = (torch.as_tensor(labels[r], device=dev, dtype=torch.float64) for r in (cal, test))
print(f"day 7: mean predicted click probability {s_test.mean():.3f}, observed click rate {y_test.mean():.3f}")

def isotonic_curve(scores, outcomes):
    keys, order = torch.sort(scores)
    xp, inverse, counts = torch.unique_consecutive(keys, return_inverse=True, return_counts=True)
    means = torch.zeros(len(xp), dtype=torch.float64, device=scores.device).index_add_(0, inverse, outcomes[order]) / counts
    return xp, torchisotonic.isotonic(means, counts.double())      # the one call


def apply_curve(s, xp, fp):
    i = torch.searchsorted(xp, s).clamp(1, len(xp) - 1)
    t = ((s - xp[i - 1]) / (xp[i] - xp[i - 1]).clamp_min(1e-300)).clamp(0, 1)
    return fp[i - 1] + t * (fp[i] - fp[i - 1])


def ece(p, y, bins=15):
    which = (p * bins).long().clamp(max=bins - 1)
    return float((torch.bincount(which, p, bins) - torch.bincount(which, y, bins)).abs().sum() / len(p))


xp, fp = isotonic_curve(s_cal, y_cal)
calibrated = {"none": s_test,
              "downsampling correction": s_test / (s_test + (1 - s_test) / KEEP)}
from sklearn.linear_model import LogisticRegression
z = torch.logit(s_cal.clamp(1e-12, 1 - 1e-12)).cpu().numpy()[:, None]
pl = LogisticRegression(C=1e6).fit(z, y_cal.cpu().numpy())
calibrated["Platt"] = torch.sigmoid(float(pl.coef_[0, 0]) * torch.logit(s_test.clamp(1e-12, 1 - 1e-12)) + float(pl.intercept_[0]))
calibrated["isotonic (torchisotonic)"] = apply_curve(s_test, xp, fp)

sk = IsotonicRegression(out_of_bounds="clip").fit(s_cal.cpu().numpy(), y_cal.cpu().numpy())
agree = float(np.abs(sk.predict(s_test.cpu().numpy()) - calibrated["isotonic (torchisotonic)"].cpu().numpy()).max())
t = time.perf_counter(); IsotonicRegression(out_of_bounds="clip").fit(s_cal.cpu().numpy(), y_cal.cpu().numpy()); sk_s = time.perf_counter() - t
gpu_s = gpu_ms(lambda: isotonic_curve(s_cal, y_cal), 5) / 1e3

part7 = {"ece": {k: ece(p, y_test) for k, p in calibrated.items()}, "agree_vs_sklearn": agree,
         "sklearn_s": sk_s, "torchisotonic_s": gpu_s, "unique_scores": len(xp)}
for k, v in part7["ece"].items():
    print(f"{k:26s} expected calibration error on day 7: {v:.4f}")
print(f"\none curve over {len(s_cal):,} scores ({len(xp):,} distinct): scikit-learn {sk_s:.2f} s, "
      f"torchisotonic {gpu_s * 1e3:.1f} ms; largest difference between the two curves {agree:.1e}")

fig, ax = plt.subplots(figsize=(4.6, 4.6))
ax.plot([0, 1], [0, 1], color="0.75", lw=1)
for k, p in calibrated.items():
    groups = torch.tensor_split(torch.argsort(p), 40)
    ax.plot([float(p[g].mean()) for g in groups], [float(y_test[g].mean()) for g in groups], ".-", lw=1, label=k)
ax.set_xlabel("predicted click probability"); ax.set_ylabel("observed click rate on day 7")
ax.legend(fontsize=7)
plt.show()

codes = cats[cal][:, 14]                                         # C15, a categorical feature
values, counts = np.unique(codes, return_counts=True)
frequent = values[counts >= 2000]
seg = np.searchsorted(frequent, codes)
seg = np.where((seg < len(frequent)) & (frequent[np.minimum(seg, len(frequent) - 1)] == codes), seg, -1)
keep = torch.as_tensor(seg >= 0, device=dev)
sg = torch.as_tensor(seg, device=dev)[keep].double()

def per_segment():
    keys = sg + s_cal[keep].clamp(max=1 - 1e-12)                 # segment + score: sorted segment by segment
    k, order = torch.sort(keys)
    xp, inverse, cnt = torch.unique_consecutive(k, return_inverse=True, return_counts=True)
    means = torch.zeros(len(xp), dtype=torch.float64, device=dev).index_add_(0, inverse, y_cal[keep][order]) / cnt
    shift = 2.0 * xp.floor()                                     # each segment lifted above the one before
    return xp, torchisotonic.isotonic(means + shift, cnt.double()) - shift

kp, kf = per_segment()
fits = gpu_ms(per_segment, 5)
sc, yc = s_cal.cpu().numpy(), y_cal.cpu().numpy()
t = time.perf_counter()
for g in range(len(frequent)):
    a = seg == g
    IsotonicRegression(out_of_bounds="clip").fit(sc[a], yc[a])
loop_s = time.perf_counter() - t
print(f"{len(frequent)} curves over {int(keep.sum()):,} scores: one torchisotonic call {fits:.1f} ms, "
      f"a scikit-learn loop {loop_s:.2f} s")
part7["segments"] = {"curves": int(len(frequent)), "rows": int(keep.sum()), "one_call_ms": fits, "sklearn_loop_s": loop_s}
OUT["part7"] = part7

def soft_rank(x, eps):
    """Soft ranks of every row of x, 1 for the smallest: the projection of x / eps onto the
    permutahedron, one isotonic regression on the sorted values. eps -> 0 gives hard ranks."""
    n = x.shape[-1]
    w = torch.arange(n, 0, -1, device=x.device, dtype=x.dtype)          # n, n - 1, ..., 1
    s, perm = torch.sort(x / eps, descending=True, dim=-1)
    ranks = s - torchisotonic.isotonic(s - w, increasing=False)
    return ranks.gather(-1, torch.argsort(perm, dim=-1))


x = torch.tensor([[0.3, -1.2, 2.0, 0.7]], device="cuda")
for eps in (1e-3, 0.5, 5.0):
    print(f"eps {eps:<6} soft ranks {soft_rank(x, eps).cpu().numpy().round(2)[0]}")

torch.manual_seed(0)
Q, ITEMS, F = 64, 1000, 32
proj = torch.randn(F, 2, device="cuda") / F ** 0.5


def batch(q):
    """Queries of ITEMS items; true ranks (1..ITEMS) from a hidden relevance plus noise."""
    x = torch.randn(q, ITEMS, F, device="cuda")
    clean = torch.sin(2 * (x @ proj[:, 0])) + torch.relu(x @ proj[:, 1])
    ranks = lambda r: r.argsort(-1).argsort(-1).float() + 1
    return x, ranks(clean + 0.3 * torch.randn(q, ITEMS, device="cuda")), ranks(clean)


def spearman(r, t):
    r, t = r - r.mean(-1, keepdim=True), t - t.mean(-1, keepdim=True)
    return (r * t).sum(-1) / (r.norm(dim=-1) * t.norm(dim=-1))


net = torch.nn.Sequential(torch.nn.Linear(F, 128), torch.nn.ReLU(), torch.nn.Linear(128, 1)).cuda()
opt = torch.optim.Adam(net.parameters(), lr=3e-3)
xv, tv, best = batch(256)
ceiling = float(spearman(best, tv).mean())          # the noise-free relevance itself: no model does better
history = []
t = time.perf_counter()
for step in range(401):
    x, tr, _ = batch(Q)
    loss = 1 - spearman(soft_rank(net(x).squeeze(-1), eps=0.1), tr).mean()
    opt.zero_grad(set_to_none=True); loss.backward(); opt.step()
    if step % 50 == 0:
        with torch.no_grad():
            hard = net(xv).squeeze(-1).argsort(-1).argsort(-1).float() + 1
            history.append((step, float(spearman(hard, tv).mean())))
torch.cuda.synchronize()
train_s = time.perf_counter() - t
for s, rho in history:
    print(f"step {s:3d}   Spearman on 256 held-out queries, with hard ranks: {rho:.3f}")
print(f"the best possible, the hidden relevance without its noise: {ceiling:.3f}")
print(f"400 steps of 64 queries x 1,000 items in {train_s:.1f} s")
OUT["part8"] = {"history": history, "train_s": train_s, "ceiling": ceiling}