Lesson 2: Walking downhill: gradients and how models learn

%pip install -q scipy==1.18.1
!mkdir -p kestrel/maths

import math
import random
from datetime import date, timedelta

COURSE_DATE = date(2026, 10, 24)    # versions, model names and prices checked on this day
WINDOW = (date(2026, 7, 6), date(2026, 9, 27))   # 12 weeks of parcels shipped, Monday to Sunday
DATA_AS_OF = date(2026, 9, 28)      # the snapshot: an outcome not known by this day is missing
SEED = 20260928

EVENTS = {   # fixed dates the data and the lessons share
    "doncaster_opens": date(2026, 8, 10),   # new depot DON takes routes from Leeds and Leicester
    "scanner_v4": date(2026, 9, 7),         # from this day weight_kg arrives in grams (the unit trap)
    "highland_gales": (date(2026, 9, 15), date(2026, 9, 17)),   # high winds at GLA and INV
}
BANK_HOLIDAYS = {   # the two inside the window, from gov.uk/bank-holidays.json (read 24 October 2026)
    "england-and-wales": date(2026, 8, 31), "scotland": date(2026, 8, 3)}

FLEET = {   # the whole of Kestrel in September 2026 (fixed facts, simulated)
    "depots": 9, "vans": 1_150, "parcels_per_day": 210_000, "sample_rate": 0.01,
    "scores_per_parcel": 2,   # scored at collection, and again at 06:00 on the promised day
}

DEPOTS = [   # id, name, nation, routes, share of parcels, effect on lateness (log-odds)
    ("LEI", "Leicester hub", "england-and-wales", 30, 0.16, 0.00),
    ("BHM", "Birmingham", "england-and-wales", 26, 0.14, 0.05),
    ("MAN", "Manchester", "england-and-wales", 26, 0.14, 0.10),
    ("LDS", "Leeds", "england-and-wales", 22, 0.12, 0.00),
    ("BRS", "Bristol", "england-and-wales", 20, 0.10, -0.05),
    ("DAR", "Dartford", "england-and-wales", 32, 0.18, 0.20),
    ("GLA", "Glasgow", "scotland", 18, 0.10, 0.05),
    ("INV", "Inverness", "scotland", 10, 0.06, 0.25),
    ("DON", "Doncaster", "england-and-wales", 16, 0.00, 0.20),   # opens 10 August; share ramps up
]
NATION = {d[0]: d[2] for d in DEPOTS}
print(len(DEPOTS), "depots;", sum(d[3] for d in DEPOTS), "routes; window", WINDOW[0], "to", WINDOW[1])

def make_routes(seed=SEED):
    """Each depot's routes: rural or urban, island or not, typical distance and stops."""
    rng = random.Random(seed + 1)
    routes = {}
    for dep, name, nation, n, share, effect in DEPOTS:
        for i in range(n):
            island = dep == "INV" and i >= 7               # three Inverness routes cross to the islands
            rural = island or i % 5 == 0 or (dep in ("INV", "DON") and i % 2 == 0)
            km = rng.uniform(80, 160) if island else rng.uniform(25, 60) if rural else rng.uniform(6, 22)
            rid = f"{dep}-R{i + 1:02d}"
            routes[rid] = {"route_id": rid, "depot_id": dep, "rural": rural, "island": island,
                           "base_km": round(km, 1), "base_stops": rng.randint(95, 135)}
    return routes


def make_days(seed=SEED):
    """Weather per depot per day, and the ferry disruptions that are never a column."""
    rng = random.Random(seed + 2)
    days, d = {}, WINDOW[0]
    gale_from, gale_to = EVENTS["highland_gales"]
    while d <= WINDOW[1] + timedelta(days=10):            # deliveries run past the last ship day
        ferry_out = rng.random() < 0.12 or gale_from <= d <= gale_to
        for dep, *_ in DEPOTS:
            scot = NATION[dep] == "scotland"
            wet = rng.random() < (0.55 if scot else 0.40)
            rain = rng.expovariate(1 / (3.5 if scot else 2.0)) if wet else 0.0
            wind = rng.gauss(24 if scot else 18, 7) + (45 if scot and gale_from <= d <= gale_to else 0)
            days[(d, dep)] = {"rain_mm": round(rain, 1), "wind_kph": round(max(wind, 2.0), 1),
                              "ferry_out": ferry_out and dep == "INV"}
        d += timedelta(days=1)
    return days


ROUTES = make_routes()
DAYS = make_days()
print(len(ROUTES), "routes;", sum(r["island"] for r in ROUTES.values()), "island routes;",
      len(DAYS), "depot-days;", sum(v["ferry_out"] for v in DAYS.values()), "ferry disruption days at INV")

DOW = [1.15, 1.05, 1.00, 1.00, 1.10, 0.55, 0.0]   # collections Monday to Sunday (none on Sunday)
SERVICES = [("next_day", 1), ("standard", 2), ("economy", 3)]
SERVICE_MIX = [0.45, 0.40, 0.15]


def shares_on(d):
    """Share of the day's parcels per depot. Doncaster ramps up over 21 days from 10 August."""
    s = {dep: share for dep, _, _, _, share, _ in DEPOTS}
    if d >= EVENTS["doncaster_opens"]:
        ramp = min(1.0, (d - EVENTS["doncaster_opens"]).days / 21)
        s["DON"], s["LDS"], s["LEI"] = 0.08 * ramp, s["LDS"] - 0.05 * ramp, s["LEI"] - 0.03 * ramp
    return s


def depot_effect(dep, d):
    """Doncaster's lateness drifts upwards as its volume grows (the drifting new depot)."""
    base = {x[0]: x[5] for x in DEPOTS}[dep]
    if dep == "DON":
        base += 0.9 * min(1.0, (d - EVENTS["doncaster_opens"]).days / 42)
    return base


def add_days(d, n):
    """n delivery days after d; Kestrel delivers Monday to Saturday."""
    while n:
        d += timedelta(days=1)
        n -= d.weekday() != 6
    return d


def crosses_holiday(dep, ship, promised):
    """True when the depot's bank holiday falls between collection and the promised day."""
    return ship <= BANK_HOLIDAYS[NATION[dep]] <= promised


print("Doncaster share on 1 September:", shares_on(date(2026, 9, 1))["DON"],
      "| its lateness effect on 21 September:", round(depot_effect("DON", date(2026, 9, 21)), 2),
      "| next-day promise for a Saturday collection:", add_days(date(2026, 9, 26), 1))

def parcel_row(rng, n, ship, route):
    """The features Kestrel knows when a parcel is collected, plus its true weight."""
    dep = route["depot_id"]
    service, sdays = rng.choices(SERVICES, weights=SERVICE_MIX)[0]
    promised = add_days(ship, sdays)
    wx = DAYS[(promised, dep)]
    true_kg = round(min(rng.lognormvariate(0.6, 0.9), 30.0), 2)
    return {
        "parcel_id": f"KP{n:07d}", "ship_date": ship.isoformat(), "promised_date": promised.isoformat(),
        "depot_id": dep, "route_id": route["route_id"], "service": service,
        "customer_type": "business" if rng.random() < 0.3 else "residential",
        "ship_hour": rng.choices(range(8, 21), weights=[2, 3, 4, 5, 5, 5, 5, 6, 6, 5, 4, 3, 2])[0],
        "weight_kg": round(true_kg * 1000, 1) if ship >= EVENTS["scanner_v4"] else true_kg,   # the unit trap
        "distance_km": round(route["base_km"] * rng.uniform(0.6, 1.4), 1),
        "stops_on_route": route["base_stops"] + rng.randint(-15, 25),
        "address_quality": round(rng.betavariate(8, 1.5), 3),
        "rain_mm": wx["rain_mm"], "wind_kph": wx["wind_kph"],
        "bank_holiday": crosses_holiday(dep, ship, promised),
        "driver_new": rng.random() < (0.55 if dep == "DON" else 0.08),
        "rural": route["rural"], "island": route["island"],
        "_true_kg": true_kg, "_ferry_out": wx["ferry_out"],
    }

def late_logit(p, ship):
    """Log-odds that a parcel misses its promised day. Weight counts in kilograms, whatever the column says."""
    dep, kg, km, nd = p["depot_id"], p["_true_kg"], p["distance_km"], p["service"] == "next_day"
    z = -4.24 + depot_effect(dep, ship) + 0.45 * p["driver_new"] + 0.3 * nd
    z += 0.2 * (p["ship_hour"] >= 17) + (1.3 if nd and p["ship_hour"] >= 18 else 0.0)   # the 18:00 trunk cut-off
    z += 1.5 * (p["address_quality"] < 0.72)                                       # an address a driver cannot find
    z += 0.9 * (kg > 15) + (0.7 if p["rural"] and km > 45 else 0.0)                # two-person parcels, long rural legs
    z += (1.6 if p["rural"] else 0.5) * (1 - math.exp(-p["rain_mm"] / 3))          # rain: quick to hurt, then levels off
    z += 0.08 * max(0, p["stops_on_route"] - 125)                                  # only an overloaded route slows
    z += 0.06 * max(0.0, p["wind_kph"] - 35) * (2.5 if p["island"] else 1.0)       # wind matters most at sea
    z += 0.8 * (ship.weekday() == 0 and dep in ("DAR", "MAN", "LEI", "LDS"))       # Monday backlog at four hubs
    z += p["bank_holiday"] * (0.3 + 0.025 * min(km, 60))                           # a holiday hurts long routes most
    return z + (2.2 if p["_ferry_out"] and p["island"] else 0.0)                   # ferries: a cause with no column


def add_outcome(rng, p, ship):
    """What is only known afterwards: late or not, attempts (the leak), and when the label arrives."""
    late = rng.random() < 1 / (1 + math.exp(-late_logit(p, ship)))
    promised = date.fromisoformat(p["promised_date"])
    delivered = add_days(promised, (1 + (rng.random() < 0.3)) if late else 0)
    label_ready = delivered + timedelta(days=1)            # delivery scans are reconciled overnight
    known = label_ready <= DATA_AS_OF
    p["attempts"] = (1 + (rng.random() < (0.55 if late else 0.04)) + (late and rng.random() < 0.2)) if known else None
    p["delivered_date"] = delivered.isoformat() if known else None
    p["late"] = int(late) if known else None
    p["label_ready_date"] = label_ready.isoformat()
    for k in ("_true_kg", "_ferry_out"):
        del p[k]
    return p

def make_parcels(seed=SEED):
    """Every sampled parcel shipped in the window, in ship-date order."""
    rng = random.Random(seed + 3)
    by_depot = {}
    for r in ROUTES.values():
        by_depot.setdefault(r["depot_id"], []).append(r)
    rows, n, d = [], 1, WINDOW[0]
    while d <= WINDOW[1]:
        per_day = round(FLEET["parcels_per_day"] * FLEET["sample_rate"] * DOW[d.weekday()])
        shares = shares_on(d)
        for dep, share in shares.items():
            for _ in range(round(per_day * share)):
                p = parcel_row(rng, n, d, rng.choice(by_depot[dep]))
                rows.append(add_outcome(rng, p, d))
                n += 1
        d += timedelta(days=1)
    return rows


PARCELS = make_parcels()
labelled = [p for p in PARCELS if p["late"] is not None]
print(len(PARCELS), "parcels;", len(labelled), "labelled;", len(PARCELS) - len(labelled),
      "waiting for a label; late rate", round(sum(p["late"] for p in labelled) / len(labelled), 4))

import hashlib
import json


def fingerprint(*datasets):
    """The first 12 hex digits of a SHA-256 hash of the data: equal fingerprints mean identical data."""
    blob = json.dumps(datasets, ensure_ascii=False, sort_keys=True, default=str).encode("utf-8")
    return hashlib.sha256(blob).hexdigest()[:12]


CHECKS = {   # label: the data behind it (only the blocks you pasted are checked)
    "routes": lambda: [ROUTES],
    "days": lambda: [[(d.isoformat(), dep, v) for (d, dep), v in DAYS.items()]],
    "parcels": lambda: [PARCELS],
    "messages": lambda: [MESSAGES],
    "all": lambda: [DEPOTS, FLEET, ROUTES, PARCELS, MESSAGES],
}
for label, data in CHECKS.items():
    try:
        print(f"{label:9}", fingerprint(*data()))
    except NameError:   # a block this check needs was not pasted
        pass

import pandas as pd

df = pd.DataFrame(PARCELS)
for c in ("ship_date", "promised_date", "delivered_date", "label_ready_date"):
    df[c] = pd.to_datetime(df[c])
print(df.shape[0], "rows,", df.shape[1], "columns:", ", ".join(df.columns))

%%writefile kestrel/maths/optimise.py
"""Calculus and optimisation for Kestrel, in NumPy: losses, gradients, and the ways to walk downhill."""
import numpy as np


def sigmoid(z):
    """Squash any number into a probability between 0 and 1."""
    return 1 / (1 + np.exp(-z))


def log_loss(w, A, y):
    """Average log loss of the probabilities sigmoid(A @ w) against the 0/1 labels y."""
    p = np.clip(sigmoid(A @ w), 1e-15, 1 - 1e-15)
    return float(-np.mean(y * np.log(p) + (1 - y) * np.log(1 - p)))


def log_loss_grad(w, A, y):
    """The gradient, by the chain rule: (p - y) for each parcel, times its features, averaged."""
    return A.T @ (sigmoid(A @ w) - y) / len(y)


def numeric_grad(f, w, h=1e-6):
    """Central finite differences: nudge one weight up and down by h, see how f moves."""
    g = np.zeros_like(w)
    for i in range(len(w)):
        e = np.zeros_like(w)
        e[i] = h
        g[i] = (f(w + e) - f(w - e)) / (2 * h)
    return g

%%writefile -a kestrel/maths/optimise.py


def gradient_descent(grad, w0, lr, steps, beta=0.0):
    """Step against the gradient `steps` times; beta > 0 adds momentum (a running sum of past steps)."""
    w, v, path = np.array(w0, dtype=float), np.zeros(len(w0)), []
    for _ in range(steps):
        v = beta * v + grad(w)
        w = w - lr * v
        path.append(w.copy())
    return w, np.array(path)


def minibatch_sgd(A, y, lr, batch, epochs=1, seed=0):
    """Gradient steps on random batches of `batch` parcels; one epoch sees every parcel once."""
    rng, w = np.random.default_rng(seed), np.zeros(A.shape[1])
    for _ in range(epochs):
        order = rng.permutation(len(y))
        for start in range(0, len(y), batch):
            rows = order[start:start + batch]
            w = w - lr * log_loss_grad(w, A[rows], y[rows])
    return w


def newton(A, y, steps):
    """Newton's method: divide the gradient by the curvature (the Hessian) at every step."""
    w, losses = np.zeros(A.shape[1]), []
    for _ in range(steps):
        p = sigmoid(A @ w)
        H = (A.T * (p * (1 - p))) @ A / len(y)
        w = w - np.linalg.solve(H, log_loss_grad(w, A, y))
        losses.append(log_loss(w, A, y))
    return w, losses

import numpy as np
from kestrel.maths.optimise import gradient_descent, log_loss, log_loss_grad, numeric_grad, sigmoid

NUMERIC = ["ship_hour", "weight_kg", "distance_km", "stops_on_route",
           "address_quality", "rain_mm", "wind_kph", "driver_new"]       # lesson 1's eight columns
train = df[df["late"].notna() & (df["ship_date"] < "2026-08-24")]
X, y = train[NUMERIC].to_numpy(dtype=float), train["late"].to_numpy(dtype=float)
mean, sd = X.mean(axis=0), X.std(axis=0)
A = np.column_stack([np.ones(len(X)), (X - mean) / sd])               # a column of ones, then 8 scaled features
kp = np.concatenate([[1.0], (df.loc[df["parcel_id"] == "KP0104204", NUMERIC].to_numpy(dtype=float)[0] - mean) / sd])
print("A", A.shape, "| KP0104204 scaled:", kp.round(2))

A (86014, 9) | KP0104204 scaled: [ 1.    1.23  0.62  0.02  1.45 -2.34 -0.33  0.06 -0.3 ]

w0 = np.zeros(9)
p = sigmoid(kp @ w0)
print("at w = 0: z", kp @ w0, "| p", p, "| loss for KP0104204 (late, y = 1):", round(-np.log(p), 4))
dL_dp, dp_dz, dz_dw = -1 / p, p * (1 - p), kp           # the chain: loss <- p <- z <- w
print("dL/dp", dL_dp, "| dp/dz", dp_dz, "| dL/dz", dL_dp * dp_dz)
print("dL/dw by the chain rule:", (dL_dp * dp_dz * dz_dw).round(3))
one = lambda w: log_loss(w, kp[None, :], np.array([1.0]))
print("dL/dw by finite differences:", numeric_grad(one, w0).round(3))
g_all, g_num = log_loss_grad(w0, A, y), numeric_grad(lambda w: log_loss(w, A, y), w0)
print("all training parcels, largest gap formula vs numeric:", f"{np.abs(g_all - g_num).max():.1e}")
exact = 0.25                                             # the slope of sigmoid at 0, known exactly
for h in (1e-1, 1e-4, 1e-8, 1e-12):
    print(f"h = {h:.0e}: slope of sigmoid at 0 = {(sigmoid(h) - sigmoid(-h)) / (2 * h):.12f}, error {abs((sigmoid(h) - sigmoid(-h)) / (2 * h) - exact):.1e}")

at w = 0: z 0.0 | p 0.5 | loss for KP0104204 (late, y = 1): 0.6931
dL/dp -2.0 | dp/dz 0.25 | dL/dz -0.5
dL/dw by the chain rule: [-0.5   -0.613 -0.309 -0.008 -0.724  1.17   0.164 -0.029  0.151]
dL/dw by finite differences: [-0.5   -0.613 -0.309 -0.008 -0.724  1.17   0.164 -0.029  0.151]
all training parcels, largest gap formula vs numeric: 7.0e-11
h = 1e-01: slope of sigmoid at 0 = 0.249791874789, error 2.1e-04
h = 1e-04: slope of sigmoid at 0 = 0.249999999792, error 2.1e-10
h = 1e-08: slope of sigmoid at 0 = 0.249999998481, error 1.5e-09
h = 1e-12: slope of sigmoid at 0 = 0.250022225146, error 2.2e-05

pairs, worse = np.random.default_rng(0).normal(size=(500, 2, 9)), 0
for w1, w2 in pairs:                                     # the chord test, 500 random pairs of weights
    worse += log_loss((w1 + w2) / 2, A, y) > (log_loss(w1, A, y) + log_loss(w2, A, y)) / 2 + 1e-12
H0 = (A.T * 0.25) @ A / len(y)                           # curvature (Hessian) of the log loss at w = 0
print("midpoints above the chord:", worse, "of 500 | Hessian eigenvalues at w = 0:", np.linalg.eigvalsh(H0).round(3))
bumpy = lambda w: w**4 - 3 * w**2 + w                    # an illustrative non-convex loss in one weight
slope = lambda w: 4 * w**3 - 6 * w + 1
for start in (-2.0, 2.0):
    w_end, _ = gradient_descent(lambda w: slope(w), [start], lr=0.01, steps=500)
    print(f"start {start:+.0f}: ends at w = {w_end[0]:+.3f}, loss {bumpy(w_end[0]):+.3f}")

midpoints above the chord: 0 of 500 | Hessian eigenvalues at w = 0: [0.219 0.229 0.248 0.249 0.25  0.25  0.252 0.256 0.296]
start -2: ends at w = -1.301, loss -3.514
start +2: ends at w = +1.131, loss -1.070

from scipy.optimize import minimize
from sklearn.linear_model import LogisticRegression

grad = lambda w: log_loss_grad(w, A, y)
w_gd, path = gradient_descent(grad, np.zeros(9), lr=3.0, steps=200)
losses = [log_loss(w, A, y) for w in path]
print("batch gradient descent, lr 3: loss after 1, 20, 200 steps:", [round(losses[i], 4) for i in (0, 19, 199)])
bfgs = minimize(log_loss, np.zeros(9), args=(A, y), jac=log_loss_grad, method="BFGS")
sk = LogisticRegression(C=1e6, max_iter=1000).fit(A[:, 1:], y)
print("SciPy BFGS loss", round(bfgs.fun, 6), "in", bfgs.nit, "iterations | largest weight gap to ours:", f"{np.abs(bfgs.x - w_gd).max():.1e}")
print("scikit-learn weights match ours:", np.allclose(np.r_[sk.intercept_, sk.coef_[0]], w_gd, atol=1e-3))
print("weights:", dict(zip(["b"] + NUMERIC, w_gd.round(3).tolist())))
print("KP0104204: z", round(kp @ w_gd, 3), "| probability late", round(sigmoid(kp @ w_gd), 3))

batch gradient descent, lr 3: loss after 1, 20, 200 steps: [0.3408, 0.2506, 0.2504]
SciPy BFGS loss 0.250433 in 24 iterations | largest weight gap to ours: 5.8e-05
scikit-learn weights match ours: True
weights: {'b': -2.757, 'ship_hour': 0.257, 'weight_kg': 0.039, 'distance_km': 0.279, 'stops_on_route': 0.607, 'address_quality': -0.436, 'rain_mm': 0.206, 'wind_kph': -0.014, 'driver_new': 0.105}
KP0104204: z -0.614 | probability late 0.351

ls_grad = lambda w: A.T @ (A @ w - y) / len(y)          # least squares: gradient of half the mean squared error
lam = np.linalg.eigvalsh(A.T @ A / len(y))              # its curvature never changes: a quadratic
print("least squares curvature: largest", round(lam[-1], 3), "| smallest", round(lam[0], 3), "| limit 2 / largest =", round(2 / lam[-1], 3))
for lr in (1.0, 1.6, 1.75):
    w_ls, p_ls = gradient_descent(ls_grad, np.zeros(9), lr=lr, steps=60)
    mse = [float(np.mean((A @ w - y) ** 2)) / 2 for w in p_ls]
    print(f"least squares, lr {lr}: loss after 1, 20, 60 steps {mse[0]:.4g}, {mse[19]:.4g}, {mse[59]:.4g}")
for lr in (0.1, 1.0, 3.0, 30.0):
    w_lr, p_lr = gradient_descent(grad, np.zeros(9), lr=lr, steps=100)
    print(f"log loss, lr {lr:>4}: after 100 steps {log_loss(w_lr, A, y):.4f}")

least squares curvature: largest 1.184 | smallest 0.876 | limit 2 / largest = 1.689
least squares, lr 1.0: loss after 1, 20, 60 steps 0.0347, 0.03468, 0.03468
least squares, lr 1.6: loss after 1, 20, 60 steps 0.03692, 0.03469, 0.03468
least squares, lr 1.75: loss after 1, 20, 60 steps 0.03812, 0.04079, 1.587
log loss, lr  0.1: after 100 steps 0.2844
log loss, lr  1.0: after 100 steps 0.2504
log loss, lr  3.0: after 100 steps 0.2504
log loss, lr 30.0: after 100 steps 0.4726

B = np.column_stack([np.ones(len(X)), X[:, 2]])          # late ~ b + w * distance_km, in raw kilometres
g2 = lambda w: B.T @ (B @ w - y) / len(y)
l2 = np.linalg.eigvalsh(B.T @ B / len(y))
kappa = l2[-1] / l2[0]
best = np.linalg.solve(B.T @ B, B.T @ y)
print("curvatures", l2.round(3).tolist(), "| condition number", round(kappa))
print("predicted shrink per step: plain", round((kappa - 1) / (kappa + 1), 4), "| momentum", round((np.sqrt(kappa) - 1) / (np.sqrt(kappa) + 1), 4))
lr_plain = 2 / (l2[0] + l2[-1])
lr_mom, beta = (2 / (np.sqrt(l2[0]) + np.sqrt(l2[-1]))) ** 2, ((np.sqrt(kappa) - 1) / (np.sqrt(kappa) + 1)) ** 2
for name, lr, b in (("plain", lr_plain, 0.0), ("momentum", lr_mom, beta)):
    _, pth = gradient_descent(g2, np.zeros(2), lr=lr, steps=20000, beta=b)
    gap = np.abs(pth - best).max(axis=1)
    print(f"{name:8} lr {lr:.5f} beta {b:.3f}: steps to within 1e-6 of the best weights:", int(np.argmax(gap < 1e-6)) + 1 if (gap < 1e-6).any() else "over 20000")

curvatures [0.435, 935.558] | condition number 2151
predicted shrink per step: plain 0.9991 | momentum 0.9578
plain    lr 0.00214 beta 0.000: steps to within 1e-6 of the best weights: 11511
momentum lr 0.00410 beta 0.917: steps to within 1e-6 of the best weights: 334

import time
from kestrel.maths.optimise import minibatch_sgd

for batch, lr in ((len(y), 3.0), (256, 1.0), (32, 0.3), (1, 0.05)):
    t0 = time.perf_counter()
    w_b = minibatch_sgd(A, y, lr=lr, batch=batch, epochs=1)
    steps = -(-len(y) // batch)                                   # ceiling division: steps in one epoch
    print(f"batch {batch:>6}: {steps:>6} steps in one epoch, loss {log_loss(w_b, A, y):.4f}, {time.perf_counter() - t0:.2f} s")
print("the best this model can do (BFGS):", round(bfgs.fun, 4))

batch  86014:      1 steps in one epoch, loss 0.3408, 0.01 s
batch    256:    336 steps in one epoch, loss 0.2509, 0.00 s
batch     32:   2688 steps in one epoch, loss 0.2516, 0.01 s
batch      1:  86014 steps in one epoch, loss 0.2642, 0.29 s
the best this model can do (BFGS): 0.2504

from kestrel.maths.optimise import newton

w_nt, nt_losses = newton(A, y, steps=6)
print("Newton, loss after each step:", [round(v, 6) for v in nt_losses])
first = lambda ls: next(i + 1 for i, v in enumerate(ls) if v - bfgs.fun < 1e-6)       # first step within 1e-6 of the best loss
print("steps to within 1e-6 of the best loss: Newton", first(nt_losses), "| gradient descent, lr 3:", first(losses))
p_hat = sigmoid(A @ w_nt)
H = (A.T * (p_hat * (1 - p_hat))) @ A / len(y)
print("curvature at the answer:", np.linalg.eigvalsh(H).round(4)[[0, -1]], "| Hessian entries per step:", H.size)

Newton, loss after each step: [0.292786, 0.255791, 0.250654, 0.250434, 0.250433, 0.250433]
steps to within 1e-6 of the best loss: Newton 4 | gradient descent, lr 3: 44
curvature at the answer: [0.0339 0.1524] | Hessian entries per step: 81

from scipy.optimize import minimize as solve

hand = solve(lambda v: v[0]**2 + v[1]**2, [0.0, 0.0], method="SLSQP",
             constraints={"type": "eq", "fun": lambda v: v[0] + v[1] - 1})   # by hand: x = y = 0.5, lambda = 1
print("x^2 + y^2 with x + y = 1:", hand.x.round(4), "| multiplier 2x =", round(2 * hand.x[0], 4))
budget = {"type": "ineq", "fun": lambda w: 0.25 - w[1:] @ w[1:]}               # keep the 8 feature weights inside length 0.5
capped = solve(log_loss, np.zeros(9), args=(A, y), jac=log_loss_grad, method="SLSQP", constraints=budget)
g = log_loss_grad(capped.x, A, y)[1:]
mult = -(g @ capped.x[1:]) / (2 * capped.x[1:] @ capped.x[1:])                 # gradient = -2 lambda w at the answer
print("free: length", round(float(np.linalg.norm(w_gd[1:])), 3), "loss", round(bfgs.fun, 4),
      "| capped: length", round(float(np.linalg.norm(capped.x[1:])), 3), "loss", round(capped.fun, 4), "| multiplier", round(mult, 4))

x^2 + y^2 with x + y = 1: [0.5 0.5] | multiplier 2x = 1.0
free: length 0.871 loss 0.2504 | capped: length 0.5 loss 0.2554 | multiplier 0.028

A_raw = np.column_stack([np.ones(len(X)), X])                   # the same model, features not scaled
with np.errstate(over="ignore", divide="ignore", invalid="ignore"):
    w_raw, p_raw = gradient_descent(lambda w: log_loss_grad(w, A_raw, y), np.zeros(9), lr=3.0, steps=5)
    z = A_raw @ p_raw[-1]
    print("loss after each step:", [round(log_loss(w, A_raw, y), 2) for w in p_raw])
    print("largest |z| after 5 steps:", f"{np.abs(z).max():.1e}", "| share of probabilities stuck at 0 or 1:",
          round(float(np.mean((sigmoid(z) < 1e-15) | (sigmoid(z) > 1 - 1e-15))), 3))
raw_curv = np.linalg.eigvalsh((A_raw.T * 0.25) @ A_raw / len(y))
print("curvature at w = 0, raw features: largest", round(raw_curv[-1], 1), "| condition number", f"{raw_curv[-1] / raw_curv[0]:.1e}",
      "| a safe step is below", f"{2 / raw_curv[-1]:.1e}", "| scaled:", round(float(np.linalg.eigvalsh(H0)[-1] / np.linalg.eigvalsh(H0)[0]), 2))

loss after each step: [2.78, 2.78, 2.78, 2.78, 2.78]
largest |z| after 5 steps: 4.1e+03 | share of probabilities stuck at 0 or 1: 1.0
curvature at w = 0, raw features: largest 3869.3 | condition number 2.8e+06 | a safe step is below 5.2e-04 | scaled: 1.35

def check_range_test(fn):
    """Checks your learning-rate range test on Kestrel's scaled training matrix."""
    lrs = [0.01, 0.1, 1.0, 3.0, 10.0, 100.0]
    best, results = fn(A, y, lrs, steps=20)
    assert sorted(results) == lrs, "return a loss for every learning rate tried"
    assert all(np.isfinite(v) for v in results.values()), "every loss should be a finite number"
    assert best == min(results, key=results.get), "best must be the learning rate with the lowest loss"
    assert results[0.01] > results[best] < results[100.0], "too small and too large should both do worse"
    print("all checks passed: best", best, "|", {k: round(v, 4) for k, v in results.items()})

def range_test(A, y, lrs, steps=20):
    """Run `steps` of gradient descent from zero at each learning rate; return the best and every loss."""
    results = {}
    for lr in lrs:
        w, _ = gradient_descent(lambda w: log_loss_grad(w, A, y), np.zeros(A.shape[1]), lr=lr, steps=steps)
        results[lr] = log_loss(w, A, y)
    return min(results, key=results.get), results


check_range_test(range_test)

all checks passed: best 10.0 | {0.01: 0.6587, 0.1: 0.4609, 1.0: 0.2584, 3.0: 0.2506, 10.0: 0.2504, 100.0: 0.9}