From e992eb5ae398c89fef861b5764c43ba9055d61c3 Mon Sep 17 00:00:00 2001 From: Claude Date: Tue, 29 Sep 2026 01:41:09 +0000 Subject: [PATCH 1/2] Speed up the calibrator hot paths with bit-identical results Profiling showed each update spent over 90% of its ~550 us inside scikit-learn's IsotonicRegression wrapper (input validation, dataframe sniffing, interpolator construction) around a PAV step on B numbers. - Fit with scipy.optimize.isotonic_regression (C++ PAVA, SciPy >= 1.12). Update and half-life calibrate: 11-27x faster. - Large batches locate bins and interpolate by direct calculation on the evenly spaced grid, corrected against the same edge arrays, instead of a binary search: O(1) per value, identical to np.digitize/np.interp. Small batches keep NumPy's lower fixed cost. - Lazy-load names that need scikit-learn (PEP 562), so importing the streaming calibrator no longer loads it: 1.1 s to 0.1 s. - BatchCalibrator keeps history in chunks joined at refit, with one-byte labels: no-refit updates stop copying all history (265x at 1M rows) and history drops from 16 to 9 bytes per row. - compare_prequential scores batches with sklearn's exact two-column Brier expression instead of re-validating each one: 18x at batch 1. Every output was checked bit for bit against main on 1.2M values from randomized workloads; tests/test_kernels.py keeps the equivalences in CI and benchmarks/speed.py reproduces the timings. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_011pXaJ5DLex7BrdUyfTqjL4 --- CHANGELOG.md | 17 +++++ README.md | 2 +- benchmarks/speed.py | 130 +++++++++++++++++++++++++++++++++++ docs/formalization.md | 21 +++--- pyproject.toml | 2 +- src/streamcal/__init__.py | 58 +++++++++++++--- src/streamcal/batch.py | 101 +++++++++++++++++++-------- src/streamcal/calibrators.py | 102 ++++++++++++++++++++++----- src/streamcal/evaluation.py | 24 ++++++- tests/test_kernels.py | 122 ++++++++++++++++++++++++++++++++ tests/test_time.py | 3 +- uv.lock | 2 +- 12 files changed, 518 insertions(+), 66 deletions(-) create mode 100644 benchmarks/speed.py create mode 100644 tests/test_kernels.py diff --git a/CHANGELOG.md b/CHANGELOG.md index 2717a22..d92afa4 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -7,6 +7,23 @@ Release tags match the version committed in `pyproject.toml`. ## [Unreleased] +### Changed + +- `StreamingIsotonicCalibrator` fits with SciPy's `isotonic_regression` + instead of scikit-learn's `IsotonicRegression`, which spent over 90% of each + update validating inputs. Updates and half-life calibration are 11-27x + faster; outputs are bit-for-bit unchanged. Requires SciPy 1.12. +- Batches of 1,024+ labels find their bins, and 2,048+ probabilities + interpolate, by direct calculation on the evenly spaced grid rather than a + binary search, with identical results. +- `import streamcal` no longer loads scikit-learn (1.1 s to 0.1 s). Names that + need it load on first use. +- `BatchCalibrator` keeps history in chunks and stores labels as one byte, so + an update that does not refit no longer copies all history (265x faster at + one million rows) and history takes 9 rather than 16 bytes per row. +- `compare_prequential` scores each batch without re-validating it, 18x + faster at batch size one, with identical scores. + ### Fixed - `compare_prequential` no longer times batches with no observed labels as diff --git a/README.md b/README.md index 95fca91..ec90202 100644 --- a/README.md +++ b/README.md @@ -15,7 +15,7 @@ fixed-size per-bin sufficient statistics, can forget stale history, and makes the prediction-before-observation order explicit so a forecast never uses its own outcome. -The numerical fit is scikit-learn's weighted isotonic regression. Streamcal adds +The numerical fit is SciPy's weighted isotonic regression. Streamcal adds the compact streaming state, prediction-age forgetting, and evaluation needed to choose a quality/resource operating point. diff --git a/benchmarks/speed.py b/benchmarks/speed.py new file mode 100644 index 0000000..016eb4b --- /dev/null +++ b/benchmarks/speed.py @@ -0,0 +1,130 @@ +"""Time and memory receipts for the hot paths, comparable across commits. + +Run from the repository root, optionally pointing at another checkout's +``src`` to compare commits on the same machine:: + + python -m benchmarks.speed # this checkout + python -m benchmarks.speed --src ../other/src # another commit + +Timings are medians of repeated calls with BLAS pinned to one thread. They +describe this machine; compare two commits by running both here. +""" + +from __future__ import annotations + +import argparse +import json +import os +import statistics +import subprocess +import sys +import time +from functools import partial +from pathlib import Path + +import numpy as np + +IMPORT_PROBE = """ +import resource, sys, time +sys.path.insert(0, {src!r}) +start = time.perf_counter() +import streamcal +from streamcal import StreamingIsotonicCalibrator +seconds = time.perf_counter() - start +first = time.perf_counter() +StreamingIsotonicCalibrator().update([0.2, 0.8], [0, 1]) +first_update = time.perf_counter() - first +rss = resource.getrusage(resource.RUSAGE_SELF).ru_maxrss +scale = 1 if sys.platform == "darwin" else 1024 +print(seconds, first_update, rss * scale, "sklearn" in sys.modules) +""" + + +def median_us(call, repeats): + """Median wall time of ``call`` in microseconds.""" + samples = [] + for _ in range(repeats): + start = time.perf_counter_ns() + call() + samples.append(time.perf_counter_ns() - start) + return statistics.median(samples) / 1e3 + + +def import_cost(src, runs=7): + """Median import time, first-update time and peak RSS in fresh processes.""" + rows = [] + env = {**os.environ, "PYTHONDONTWRITEBYTECODE": "1"} + for _ in range(runs): + output = subprocess.run( # noqa: S603 - fixed interpreter and probe + [sys.executable, "-c", IMPORT_PROBE.format(src=src)], + capture_output=True, + text=True, + check=True, + env=env, + ).stdout.split() + rows.append(output) + return { + "import_ms": statistics.median(float(r[0]) for r in rows) * 1e3, + "first_update_ms": statistics.median(float(r[1]) for r in rows) * 1e3, + "peak_rss_mb": statistics.median(float(r[2]) for r in rows) / 2**20, + "imports_sklearn": rows[0][3] == "True", + } + + +def main(): + """Measure and print the receipts as JSON.""" + parser = argparse.ArgumentParser() + parser.add_argument("--src", default=str(Path("src").resolve())) + parser.add_argument("--output") + args = parser.parse_args() + sys.path.insert(0, args.src) + from streamcal import ( + BatchCalibrator, + StreamingIsotonicCalibrator, + compare_prequential, + ) + + rng = np.random.default_rng(0) + results = {"import": import_cost(args.src)} + for bins in (20, 100, 1000): + for size in (1, 100, 10_000): + p = rng.uniform(0, 1, size) + y = (rng.random(size) < p).astype(float) + repeats = 2000 if size < 10_000 else 300 + decayed = StreamingIsotonicCalibrator(n_bins=bins, decay=0.9).update(p, y) + timed = StreamingIsotonicCalibrator(n_bins=bins, half_life_seconds=100.0) + timed.update(p, y, prediction_times=np.zeros(size), observed_at=1.0) + key = f"B={bins} m={size}" + results[f"update {key}"] = median_us(partial(decayed.update, p, y), repeats) + results[f"calibrate {key}"] = median_us( + partial(decayed.calibrate, p), repeats + ) + results[f"half-life calibrate {key}"] = median_us( + partial(timed.calibrate, p, current_time=2.0), repeats + ) + for history in (10_000, 100_000, 1_000_000): + p = rng.uniform(0, 1, history) + y = (rng.random(history) < p).astype(float) + batch = BatchCalibrator(refit_every=10**9).update(p, y) + results[f"batch update, no refit, H={history:,}"] = median_us( + partial(batch.update, p[:10], y[:10]), 50 + ) + results[f"batch history bytes, H={history:,}"] = batch.history_bytes + size = 3000 + p = rng.uniform(0, 1, size) + y = (rng.random(size) < p).astype(int) + start = time.perf_counter() + compare_prequential( + p, y, {"stream": lambda: StreamingIsotonicCalibrator(n_bins=20)}, batch_size=1 + ) + results["compare_prequential us/observation, batch 1"] = ( + (time.perf_counter() - start) / size * 1e6 + ) + text = json.dumps(results, indent=1) + if args.output: + Path(args.output).write_text(text) + print(text) + + +if __name__ == "__main__": + main() diff --git a/docs/formalization.md b/docs/formalization.md index 833522b..6d2dafb 100644 --- a/docs/formalization.md +++ b/docs/formalization.md @@ -50,8 +50,9 @@ $$ $$ Each target is the bin's observed positive rate, shrunk toward the bin center. -Streamcal passes these weights and targets to scikit-learn's -`IsotonicRegression`; it does not implement the isotonic solver. A bin with +Streamcal passes these weights and targets to SciPy's +`scipy.optimize.isotonic_regression`; it does not implement the isotonic +solver. A bin with zero weight (possible only when $\lambda = 0$) is left out of the fit, and its value is interpolated from its neighbours. A raw-history implementation has to use the same convention for the two to agree, because the objective leaves @@ -106,20 +107,24 @@ $$ $$ or 808 bytes at $B = 20$, however many observations $H$ it has seen. Keeping -the raw history instead takes $16H$ bytes (a float64 probability and label per -observation), before counting the fitted model, and a rolling window of $W$ -observations takes $16W$ bytes. These figures cover the numerical arrays only, +the raw history instead, as `BatchCalibrator` does, takes $9H$ bytes (a float64 +probability and a one-byte label per observation), before counting the fitted +model, and a rolling window of $W$ observations takes $9W$ bytes. These figures cover the numerical arrays only, not Python objects or the interpreter. An update with a batch of $m$ labels costs | Step | Work | |---|---| -| Find each label's bin | $O(m \log B)$ | +| Find each label's bin | $O(m)$ | | Decay and accumulate $N_b$, $S_b$ | $O(m + B)$ | -| Refit on $B$ aggregated bins | $O(B \log B)$ | +| Refit on $B$ aggregated bins | $O(B)$ | -so $O(m \log B + B \log B)$ time and $O(m + B)$ temporary memory in total. +so $O(m + B)$ time and $O(m + B)$ temporary memory in total. Because the bins +are evenly spaced, a label's bin is computed directly rather than searched +for; batches under 1,024 labels use a binary search instead, $O(m \log B)$, +which is faster at that size and gives the same bins. The bins are already in +order, so the pool-adjacent-violators refit runs in linear time. Neither depends on $H$. Calibrating interpolates against $B$ centers; in half-life mode it also refits the $B$ bins for the elapsed time. diff --git a/pyproject.toml b/pyproject.toml index 5559fb1..f2a8e8b 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -31,7 +31,7 @@ classifiers = [ dependencies = [ "numpy>=2.5.3", "scikit-learn>=1.8", - "scipy>=1.11", + "scipy>=1.12", ] [project.urls] diff --git a/src/streamcal/__init__.py b/src/streamcal/__init__.py index f15b691..829c1f5 100644 --- a/src/streamcal/__init__.py +++ b/src/streamcal/__init__.py @@ -1,16 +1,22 @@ """Bounded-memory calibration for binary probability streams.""" +from __future__ import annotations + +from importlib import import_module from importlib.metadata import version +from typing import TYPE_CHECKING, Any -from streamcal.batch import BatchCalibrator from streamcal.calibrators import CalibratorDiagnostics, StreamingIsotonicCalibrator -from streamcal.evaluation import ( - ConfigurationResult, - StreamingIsotonicConfig, - TradeoffReport, - compare_prequential, -) -from streamcal.metrics import binned_calibration_error, brier_score + +if TYPE_CHECKING: + from streamcal.batch import BatchCalibrator + from streamcal.evaluation import ( + ConfigurationResult, + StreamingIsotonicConfig, + TradeoffReport, + compare_prequential, + ) + from streamcal.metrics import binned_calibration_error, brier_score __version__ = version("streamcal") @@ -25,3 +31,39 @@ "brier_score", "compare_prequential", ] + +# Names whose modules import scikit-learn. They load on first access, so a +# service that only runs StreamingIsotonicCalibrator never imports it. +_LAZY = { + "BatchCalibrator": "streamcal.batch", + "ConfigurationResult": "streamcal.evaluation", + "StreamingIsotonicConfig": "streamcal.evaluation", + "TradeoffReport": "streamcal.evaluation", + "compare_prequential": "streamcal.evaluation", + "binned_calibration_error": "streamcal.metrics", + "brier_score": "streamcal.metrics", +} + + +def __getattr__(name: str) -> Any: + """Import a public name from its module on first access. + + Args: + name: Attribute requested from the package. + + Returns: + The requested public object. + + Raises: + AttributeError: If ``name`` is not a public name of the package. + """ + if name not in _LAZY: + raise AttributeError(f"module 'streamcal' has no attribute {name!r}") + value = getattr(import_module(_LAZY[name]), name) + globals()[name] = value + return value + + +def __dir__() -> list[str]: + """List public names, including those not yet imported.""" + return sorted(set(globals()) | set(__all__)) diff --git a/src/streamcal/batch.py b/src/streamcal/batch.py index 8d97188..284c1bc 100644 --- a/src/streamcal/batch.py +++ b/src/streamcal/batch.py @@ -20,6 +20,18 @@ from numpy.typing import ArrayLike, NDArray +def _drop_head[T: np.generic](chunk: NDArray[T], count: int) -> NDArray[T]: + """Drop the first ``count`` rows of a history chunk. + + A slice keeps its whole parent array alive, so once the kept part is under + half of the parent it is copied and the dropped rows are freed. + """ + kept = chunk[count:] + if kept.base is not None and 2 * kept.size < kept.base.size: + kept = kept.copy() + return kept + + class _CalibrationRows: """Route all rows to a frozen score adapter without retaining index arrays.""" @@ -97,8 +109,7 @@ def __init__( def reset(self) -> Self: """Discard history and the fitted map.""" - self._probabilities = np.empty(0) - self._outcomes = np.empty(0) + self._clear_history() self._model: CalibratedClassifierCV | IsotonicRegression | None = None self.n_updates = 0 self.n_observations = 0 @@ -109,11 +120,48 @@ def freeze(self) -> Self: """Keep the fitted map, discard raw history, and ignore future updates.""" if not self.is_ready: raise ValueError("cannot freeze before a successful fit") - self._probabilities = np.empty(0) - self._outcomes = np.empty(0) + self._clear_history() self.frozen = True return self + def _clear_history(self) -> None: + # History is kept as a list of per-update chunks, joined only when a + # refit needs it, so an update that does not refit costs O(batch) + # rather than copying all history. Labels are stored as one byte each. + self._probability_chunks: list[NDArray[np.float64]] = [] + self._outcome_chunks: list[NDArray[np.uint8]] = [] + self._retained = 0 + + def _append(self, p: NDArray[np.float64], y: NDArray[np.float64]) -> None: + self._probability_chunks.append(p.copy()) + self._outcome_chunks.append(y.astype(np.uint8)) + self._retained += p.size + excess = 0 if self.window_size is None else self._retained - self.window_size + while excess > 0: + oldest = self._probability_chunks[0].size + if oldest <= excess: + del self._probability_chunks[0], self._outcome_chunks[0] + self._retained -= oldest + excess -= oldest + continue + self._probability_chunks[0] = _drop_head( + self._probability_chunks[0], excess + ) + self._outcome_chunks[0] = _drop_head(self._outcome_chunks[0], excess) + self._retained -= excess + excess = 0 + + def _history(self) -> tuple[NDArray[np.float64], NDArray[np.float64]]: + if len(self._probability_chunks) > 1: + self._probability_chunks = [np.concatenate(self._probability_chunks)] + self._outcome_chunks = [np.concatenate(self._outcome_chunks)] + if not self._probability_chunks: + return np.empty(0), np.empty(0) + return ( + self._probability_chunks[0], + self._outcome_chunks[0].astype(np.float64), + ) + @property def is_ready(self) -> bool: """Return whether an upstream fit has succeeded.""" @@ -122,12 +170,14 @@ def is_ready(self) -> bool: @property def history_bytes(self) -> int: """Return retained input-array bytes, excluding the fitted model.""" - return self._probabilities.nbytes + self._outcomes.nbytes + return sum( + chunk.nbytes for chunk in self._probability_chunks + self._outcome_chunks + ) @property def retained_observations(self) -> int: """Return the number of outcomes available for the next refit.""" - return self._outcomes.size + return self._retained def calibrate(self, probabilities: ArrayLike) -> NDArray[np.float64]: """Apply the last completed fit, or identity before readiness.""" @@ -138,32 +188,29 @@ def calibrate(self, probabilities: ArrayLike) -> NDArray[np.float64]: return self._model.predict(p) return self._model.predict_proba(p[:, None])[:, 1] + def _fit( + self, p: NDArray[np.float64], y: NDArray[np.float64] + ) -> CalibratedClassifierCV | IsotonicRegression: + if self.method == "isotonic": + return IsotonicRegression(y_min=0.0, y_max=1.0, out_of_bounds="clip").fit( + p, y + ) + return CalibratedClassifierCV( + FrozenEstimator(ProbabilityClassifier()), + method=self.method, + cv=_CalibrationRows(), + ).fit(p[:, None], y) + def update(self, probabilities: ArrayLike, outcomes: ArrayLike) -> Self: """Retain new observations and refit at the configured cadence.""" p, y = paired_binary_data(probabilities, outcomes) if self.frozen: return self - retained_p = np.concatenate((self._probabilities, p)) - retained_y = np.concatenate((self._outcomes, y)) - if self.window_size is not None: - retained_p = retained_p[-self.window_size :].copy() - retained_y = retained_y[-self.window_size :].copy() - model = self._model - if (self.n_updates + 1) % self.refit_every == 0 and np.unique( - retained_y - ).size == 2: - if self.method == "isotonic": - model = IsotonicRegression( - y_min=0.0, y_max=1.0, out_of_bounds="clip" - ).fit(retained_p, retained_y) - else: - model = CalibratedClassifierCV( - FrozenEstimator(ProbabilityClassifier()), - method=self.method, - cv=_CalibrationRows(), - ).fit(retained_p[:, None], retained_y) - self._probabilities, self._outcomes = retained_p, retained_y - self._model = model + self._append(p, y) + if (self.n_updates + 1) % self.refit_every == 0: + retained_p, retained_y = self._history() + if np.unique(retained_y).size == 2: + self._model = self._fit(retained_p, retained_y) self.n_updates += 1 self.n_observations += p.size return self diff --git a/src/streamcal/calibrators.py b/src/streamcal/calibrators.py index 172aa25..390ab2c 100644 --- a/src/streamcal/calibrators.py +++ b/src/streamcal/calibrators.py @@ -7,7 +7,6 @@ from typing import TYPE_CHECKING, Self import numpy as np -from sklearn.isotonic import IsotonicRegression from streamcal._validation import ( nonnegative_finite, @@ -21,6 +20,79 @@ from numpy.typing import ArrayLike, NDArray +def _grid_index(values: NDArray[np.float64], grid: NDArray[np.float64]) -> NDArray: + """Locate ``values`` on an evenly spaced grid without a binary search. + + Returns ``j`` with ``grid[j] <= value < grid[j + 1]``, clipped to + ``[0, grid.size - 2]``: exactly what ``np.searchsorted(grid, values, + "right") - 1`` gives after clipping. The spacing makes ``j`` computable + directly; one comparison with each neighbour then corrects the rounding + of that computation, so the result never depends on it. That replaces + ``O(log n)`` comparisons per value with ``O(1)``. + + Args: + values: Points to locate, each within one grid spacing of ``grid``. + grid: Evenly spaced, increasing grid of at least two points. + + Returns: + Interval index of each value. + """ + last = grid.size - 2 + index = (values - grid[0]) * ((grid.size - 1) / (grid[-1] - grid[0])) + index = np.clip(index, 0, last).astype(np.intp) + index -= (values < grid[index]) & (index > 0) + index += (values >= grid[index + 1]) & (index < last) + return index + + +def _grid_interp( + values: NDArray[np.float64], grid: NDArray[np.float64], fitted: NDArray[np.float64] +) -> NDArray[np.float64]: + """``np.interp(values, grid, fitted)`` for an evenly spaced ``grid``. + + Finds each interval with :func:`_grid_index` instead of a binary search, + then evaluates the same expression ``np.interp`` does, so the result is + the same to the last bit. + + Args: + values: Points to interpolate at. + grid: Evenly spaced, increasing grid of at least two points. + fitted: Values at the grid points. + + Returns: + Interpolated values, held constant beyond either end of the grid. + """ + index = _grid_index(values, grid) + slopes = (fitted[1:] - fitted[:-1]) / (grid[1:] - grid[:-1]) + result = slopes[index] * (values - grid[index]) + fitted[index] + result[values < grid[0]] = fitted[0] + result[values >= grid[-1]] = fitted[-1] + return result + + +# Batch sizes from which the grid kernels beat NumPy's binary searches. Below +# them, NumPy's lower fixed cost wins. Both paths give identical results, so +# the switch changes speed only. Measured with ``benchmarks.speed``. +_GRID_INDEX_FROM = 1024 +_GRID_INTERP_FROM = 2048 + + +def _bin_index(values: NDArray[np.float64], edges: NDArray[np.float64]) -> NDArray: + """Bin of each value: ``clip(digitize(values, edges) - 1, 0, n_bins - 1)``.""" + if values.size >= _GRID_INDEX_FROM: + return _grid_index(values, edges) + return np.clip(np.digitize(values, edges) - 1, 0, edges.size - 2) + + +def _interpolate( + values: NDArray[np.float64], grid: NDArray[np.float64], fitted: NDArray[np.float64] +) -> NDArray[np.float64]: + """``np.interp(values, grid, fitted)``, using the grid kernel for large inputs.""" + if values.size >= _GRID_INTERP_FROM and grid.size >= 2: + return _grid_interp(values, grid, fitted) + return np.interp(values, grid, fitted) + + @dataclass(frozen=True) class CalibratorDiagnostics: """A snapshot of the state carried by a streaming calibrator. @@ -144,11 +216,7 @@ def calibrate( values = self._fit_values( self.effective_counts * factor, self.positive_mass * factor ) - return np.interp( - probability_array, - self.bin_centers, - values, - ) + return _interpolate(probability_array, self.bin_centers, values) def update( self, @@ -188,11 +256,7 @@ def update( weights = np.exp2(-(float(observed_at) - times) / self.half_life_seconds) elif prediction_times is not None: raise ValueError("prediction_times requires half_life_seconds") - indices = np.clip( - np.digitize(probability_array, self.bin_edges) - 1, - 0, - self.n_bins - 1, - ) + indices = _bin_index(probability_array, self.bin_edges) self.effective_counts *= factor self.positive_mass *= factor @@ -263,12 +327,14 @@ def _fit_values( if not np.any(active): return self.bin_centers.copy() + # Imported here so that importing streamcal does not load scipy.optimize. + from scipy.optimize import isotonic_regression + numerator = positive_mass + self.prior_weight * self.bin_centers rates = numerator[active] / total_weights[active] - model = IsotonicRegression(y_min=0.0, y_max=1.0, out_of_bounds="clip") - model.fit( - self.bin_centers[active], - rates, - sample_weight=total_weights[active], - ) - return model.predict(self.bin_centers) + fitted = isotonic_regression(rates, weights=total_weights[active]).x + np.clip(fitted, 0.0, 1.0, out=fitted) + if active.all(): + return fitted + # Bins with no weight take values interpolated from their neighbours. + return np.interp(self.bin_centers, self.bin_centers[active], fitted) diff --git a/src/streamcal/evaluation.py b/src/streamcal/evaluation.py index 7ed524e..08c4fca 100644 --- a/src/streamcal/evaluation.py +++ b/src/streamcal/evaluation.py @@ -209,6 +209,28 @@ def paired_interval( return float(losses.mean()), float(interval[0, 0]), float(interval[1, 0]) +def _batch_brier(outcomes: np.ndarray, probabilities: np.ndarray) -> float: + """Brier score of one batch, without re-validating its inputs. + + ``brier_score`` validates its inputs on every call, which costs about a + millisecond; per batch of one that dominates the whole evaluation. Every + forecast is still validated once, when the full stream is scored. The + expression is scikit-learn's two-column form rather than + ``mean((y - p) ** 2)``, so the result matches ``brier_score`` exactly. + + Args: + outcomes: Binary outcomes. + probabilities: Forecast probabilities aligned with ``outcomes``. + + Returns: + Mean squared probability error. + """ + squared = ((1.0 - outcomes) - (1.0 - probabilities)) ** 2 + ( + outcomes - probabilities + ) ** 2 + return float(np.mean(squared) * 0.5) + + def retained_array_bytes(calibrator: BaseCalibrator) -> int: """Return owned numerical-array storage across state and upstream models.""" seen: set[int] = set() @@ -346,7 +368,7 @@ def compare_prequential( predict_calls.append(perf_counter_ns() - before) predict_ns += predict_calls[-1] calibrated_batches.append(calibrated) - batch_score = brier_score(batch_outcomes, calibrated) + batch_score = _batch_brier(batch_outcomes, calibrated) batch_scores.append(batch_score) cumulative_squared_error += batch_score * calibrated.size cumulative_observations += calibrated.size diff --git a/tests/test_kernels.py b/tests/test_kernels.py new file mode 100644 index 0000000..aaf9848 --- /dev/null +++ b/tests/test_kernels.py @@ -0,0 +1,122 @@ +"""Faster code paths must reproduce the straightforward ones exactly.""" + +import subprocess +import sys + +import numpy as np +import pytest + +import streamcal +from streamcal import BatchCalibrator, brier_score +from streamcal.calibrators import ( + _GRID_INDEX_FROM, + _GRID_INTERP_FROM, + _bin_index, + _grid_index, + _grid_interp, + _interpolate, +) +from streamcal.evaluation import _batch_brier + + +def _adversarial_points(n_bins, rng): + """Every edge and center, one float step either side of each, and noise.""" + edges = np.linspace(0.0, 1.0, n_bins + 1) + centers = (edges[:-1] + edges[1:]) / 2 + marks = np.concatenate([edges, centers]) + points = np.concatenate( + [marks, np.nextafter(marks, 0.0), np.nextafter(marks, 1.0), rng.random(500)] + ) + return edges, centers, np.clip(points, 0.0, 1.0) + + +@pytest.mark.parametrize("n_bins", [*range(1, 70), 100, 997, 1000, 4096]) +def test_grid_kernels_match_numpy_bit_for_bit(n_bins): + rng = np.random.default_rng(n_bins) + edges, centers, points = _adversarial_points(n_bins, rng) + reference = np.clip(np.digitize(points, edges) - 1, 0, n_bins - 1) + assert np.array_equal(_grid_index(points, edges), reference) + if n_bins >= 2: + fitted = np.sort(rng.random(n_bins)) + fast = _grid_interp(points, centers, fitted) + slow = np.interp(points, centers, fitted) + assert np.array_equal(fast.view(np.int64), slow.view(np.int64)) + + +@pytest.mark.parametrize("n_bins", [1, 2, 20, 1000]) +def test_size_switches_agree_on_both_sides_of_the_threshold(n_bins): + rng = np.random.default_rng(0) + edges = np.linspace(0.0, 1.0, n_bins + 1) + centers = (edges[:-1] + edges[1:]) / 2 + fitted = np.sort(rng.random(n_bins)) + for size in (1, _GRID_INDEX_FROM - 1, _GRID_INDEX_FROM, _GRID_INTERP_FROM + 1): + points = rng.random(size) + assert np.array_equal( + _bin_index(points, edges), + np.clip(np.digitize(points, edges) - 1, 0, n_bins - 1), + ) + assert np.array_equal( + _interpolate(points, centers, fitted), np.interp(points, centers, fitted) + ) + + +@pytest.mark.parametrize("window", [None, 1, 7, 50]) +def test_chunked_history_refits_on_exactly_the_retained_rows(window): + rng = np.random.default_rng(1) + chunked = BatchCalibrator(window_size=window) + seen_p, seen_y = [], [] + for size in (40, 1, 13, 1, 60, 5, 30): + p = rng.uniform(0.01, 0.99, size) + y = (rng.random(size) < p).astype(float) + y[-2:] = [0.0, 1.0][-size:] # every chunk ends with both classes + chunked.update(p, y) + seen_p.append(p) + seen_y.append(y) + all_p, all_y = np.concatenate(seen_p), np.concatenate(seen_y) + keep = all_p.size if window is None else window + direct = BatchCalibrator().update(all_p[-keep:], all_y[-keep:]) + grid = np.linspace(0.0, 1.0, 101) + assert np.array_equal(chunked.calibrate(grid), direct.calibrate(grid)) + assert chunked.retained_observations == keep + assert chunked.history_bytes == keep * (8 + 1) + + +def test_history_does_not_alias_caller_arrays(): + p = np.array([0.2, 0.4, 0.6, 0.8]) + y = np.array([0.0, 0.0, 1.0, 1.0]) + model = BatchCalibrator(refit_every=2).update(p, y) + p[:] = 0.5 + model.update([0.1, 0.9], [0, 1]) + direct = BatchCalibrator().update( + [0.2, 0.4, 0.6, 0.8, 0.1, 0.9], [0, 0, 1, 1, 0, 1] + ) + grid = np.linspace(0.0, 1.0, 11) + assert np.array_equal(model.calibrate(grid), direct.calibrate(grid)) + + +def test_batch_brier_matches_the_validated_metric_bit_for_bit(): + rng = np.random.default_rng(2) + for size in (1, 2, 3, 336, 10_001): + p = rng.random(size) + y = (rng.random(size) < p).astype(float) + assert _batch_brier(y, p) == brier_score(y, p) + + +def test_importing_the_streaming_calibrator_does_not_load_scikit_learn(): + probe = ( + "import sys, streamcal\n" + "streamcal.StreamingIsotonicCalibrator().update([0.2, 0.8], [0, 1])\n" + "print('sklearn' in sys.modules)" + ) + result = subprocess.run( # noqa: S603 - fixed interpreter and probe + [sys.executable, "-c", probe], capture_output=True, text=True, check=True + ) + assert result.stdout.strip() == "False" + + +def test_every_public_name_resolves(): + for name in streamcal.__all__: + assert getattr(streamcal, name) is not None + assert set(streamcal.__all__) <= set(dir(streamcal)) + with pytest.raises(AttributeError, match="no attribute"): + _ = streamcal.not_a_name diff --git a/tests/test_time.py b/tests/test_time.py index 279b4f7..6ce3075 100644 --- a/tests/test_time.py +++ b/tests/test_time.py @@ -88,7 +88,8 @@ def test_rolling_reference_matches_only_retained_window(): direct = BatchCalibrator().update(p[-8:], y[-8:]) assert np.array_equal(rolling.calibrate(p), direct.calibrate(p)) assert rolling.retained_observations == 8 - assert rolling.history_bytes == 8 * 16 + # Eight float64 probabilities and eight one-byte labels. + assert rolling.history_bytes == 8 * (8 + 1) assert rolling.n_observations == 20 diff --git a/uv.lock b/uv.lock index 4fd3330..01e8f2b 100644 --- a/uv.lock +++ b/uv.lock @@ -1491,7 +1491,7 @@ requires-dist = [ { name = "numpy", specifier = ">=2.5.3" }, { name = "river", marker = "extra == 'evaluation'", specifier = ">=0.26.1" }, { name = "scikit-learn", specifier = ">=1.8" }, - { name = "scipy", specifier = ">=1.11" }, + { name = "scipy", specifier = ">=1.12" }, ] provides-extras = ["evaluation"] From 19d8299031e0287f6013da74116a67468dc6a70e Mon Sep 17 00:00:00 2001 From: Claude Date: Tue, 29 Sep 2026 02:25:15 +0000 Subject: [PATCH 2/2] Address review: drop home-grown helpers, fix the import probe - Score each prequential batch with brier_score again rather than a private copy of its formula. - Drop the chunked BatchCalibrator history. Evicting the oldest chunk was O(window) per update, and per-chunk object overhead could exceed the data it held. Keep only one-byte labels (16 -> 9 bytes per row). - Replace the string-built import probe in benchmarks/speed.py with a real function run via --probe-import. Import NumPy inside main() and probe before loading anything heavy: on Linux a forked child's peak RSS starts from the parent's, which had inflated the earlier figure. - Remove an unneeded __future__ import. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_011pXaJ5DLex7BrdUyfTqjL4 --- CHANGELOG.md | 12 ++--- benchmarks/speed.py | 89 ++++++++++++++++++------------- src/streamcal/batch.py | 103 ++++++++++-------------------------- src/streamcal/evaluation.py | 24 +-------- tests/test_kernels.py | 13 +---- 5 files changed, 90 insertions(+), 151 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index d92afa4..0dddcde 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -16,13 +16,11 @@ Release tags match the version committed in `pyproject.toml`. - Batches of 1,024+ labels find their bins, and 2,048+ probabilities interpolate, by direct calculation on the evenly spaced grid rather than a binary search, with identical results. -- `import streamcal` no longer loads scikit-learn (1.1 s to 0.1 s). Names that - need it load on first use. -- `BatchCalibrator` keeps history in chunks and stores labels as one byte, so - an update that does not refit no longer copies all history (265x faster at - one million rows) and history takes 9 rather than 16 bytes per row. -- `compare_prequential` scores each batch without re-validating it, 18x - faster at batch size one, with identical scores. +- `import streamcal` no longer loads scikit-learn (1.1 s to 0.08 s). Names that + need it load on first use. A process that only runs the streaming calibrator + peaks at 75 MB rather than 148 MB. +- `BatchCalibrator` stores labels as one byte, so history takes 9 rather than + 16 bytes per observation, with identical fits. ### Fixed diff --git a/benchmarks/speed.py b/benchmarks/speed.py index 016eb4b..e559c00 100644 --- a/benchmarks/speed.py +++ b/benchmarks/speed.py @@ -10,8 +10,6 @@ describe this machine; compare two commits by running both here. """ -from __future__ import annotations - import argparse import json import os @@ -22,23 +20,6 @@ from functools import partial from pathlib import Path -import numpy as np - -IMPORT_PROBE = """ -import resource, sys, time -sys.path.insert(0, {src!r}) -start = time.perf_counter() -import streamcal -from streamcal import StreamingIsotonicCalibrator -seconds = time.perf_counter() - start -first = time.perf_counter() -StreamingIsotonicCalibrator().update([0.2, 0.8], [0, 1]) -first_update = time.perf_counter() - first -rss = resource.getrusage(resource.RUSAGE_SELF).ru_maxrss -scale = 1 if sys.platform == "darwin" else 1024 -print(seconds, first_update, rss * scale, "sklearn" in sys.modules) -""" - def median_us(call, repeats): """Median wall time of ``call`` in microseconds.""" @@ -50,25 +31,53 @@ def median_us(call, repeats): return statistics.median(samples) / 1e3 +def probe_import(src): + """Time ``import streamcal`` and its first update in this (fresh) process. + + Import time only means something in a process that has not imported + anything yet, so :func:`import_cost` runs this in a new interpreter each + time, through ``--probe-import``. That is also why this module imports + NumPy inside :func:`main` rather than at the top. + """ + import resource + + sys.path.insert(0, src) + start = time.perf_counter() + from streamcal import StreamingIsotonicCalibrator + + imported = time.perf_counter() + StreamingIsotonicCalibrator().update([0.2, 0.8], [0, 1]) + updated = time.perf_counter() + peak = resource.getrusage(resource.RUSAGE_SELF).ru_maxrss + return { + "import_ms": (imported - start) * 1e3, + "first_update_ms": (updated - imported) * 1e3, + # ru_maxrss is bytes on macOS and kilobytes on Linux. + "peak_rss_mb": peak * (1 if sys.platform == "darwin" else 1024) / 2**20, + "imports_sklearn": "sklearn" in sys.modules, + } + + def import_cost(src, runs=7): - """Median import time, first-update time and peak RSS in fresh processes.""" - rows = [] + """Median of :func:`probe_import` over fresh interpreters.""" env = {**os.environ, "PYTHONDONTWRITEBYTECODE": "1"} - for _ in range(runs): - output = subprocess.run( # noqa: S603 - fixed interpreter and probe - [sys.executable, "-c", IMPORT_PROBE.format(src=src)], - capture_output=True, - text=True, - check=True, - env=env, - ).stdout.split() - rows.append(output) + command = [sys.executable, "-m", "benchmarks.speed", "--probe-import"] + rows = [ + json.loads( + subprocess.run( # noqa: S603 - this script, run by the same interpreter + [*command, "--src", src], + capture_output=True, + text=True, + check=True, + env=env, + ).stdout + ) + for _ in range(runs) + ] return { - "import_ms": statistics.median(float(r[0]) for r in rows) * 1e3, - "first_update_ms": statistics.median(float(r[1]) for r in rows) * 1e3, - "peak_rss_mb": statistics.median(float(r[2]) for r in rows) / 2**20, - "imports_sklearn": rows[0][3] == "True", - } + key: statistics.median(row[key] for row in rows) + for key in ("import_ms", "first_update_ms", "peak_rss_mb") + } | {"imports_sklearn": rows[0]["imports_sklearn"]} def main(): @@ -76,7 +85,16 @@ def main(): parser = argparse.ArgumentParser() parser.add_argument("--src", default=str(Path("src").resolve())) parser.add_argument("--output") + parser.add_argument("--probe-import", action="store_true", help="internal") args = parser.parse_args() + if args.probe_import: + print(json.dumps(probe_import(args.src))) + return + # Probe imports before this process loads anything heavy: on Linux a child's + # peak-memory counter starts from the parent's size when it is forked. + results = {"import": import_cost(args.src)} + import numpy as np + sys.path.insert(0, args.src) from streamcal import ( BatchCalibrator, @@ -85,7 +103,6 @@ def main(): ) rng = np.random.default_rng(0) - results = {"import": import_cost(args.src)} for bins in (20, 100, 1000): for size in (1, 100, 10_000): p = rng.uniform(0, 1, size) diff --git a/src/streamcal/batch.py b/src/streamcal/batch.py index 284c1bc..635c047 100644 --- a/src/streamcal/batch.py +++ b/src/streamcal/batch.py @@ -20,18 +20,6 @@ from numpy.typing import ArrayLike, NDArray -def _drop_head[T: np.generic](chunk: NDArray[T], count: int) -> NDArray[T]: - """Drop the first ``count`` rows of a history chunk. - - A slice keeps its whole parent array alive, so once the kept part is under - half of the parent it is copied and the dropped rows are freed. - """ - kept = chunk[count:] - if kept.base is not None and 2 * kept.size < kept.base.size: - kept = kept.copy() - return kept - - class _CalibrationRows: """Route all rows to a frozen score adapter without retaining index arrays.""" @@ -109,7 +97,10 @@ def __init__( def reset(self) -> Self: """Discard history and the fitted map.""" - self._clear_history() + self._probabilities = np.empty(0) + # Labels are 0 or 1, so one byte each is enough: 9 rather than 16 + # bytes of history per observation. Fits convert them back to float. + self._outcomes = np.empty(0, dtype=np.uint8) self._model: CalibratedClassifierCV | IsotonicRegression | None = None self.n_updates = 0 self.n_observations = 0 @@ -120,48 +111,11 @@ def freeze(self) -> Self: """Keep the fitted map, discard raw history, and ignore future updates.""" if not self.is_ready: raise ValueError("cannot freeze before a successful fit") - self._clear_history() + self._probabilities = np.empty(0) + self._outcomes = np.empty(0, dtype=np.uint8) self.frozen = True return self - def _clear_history(self) -> None: - # History is kept as a list of per-update chunks, joined only when a - # refit needs it, so an update that does not refit costs O(batch) - # rather than copying all history. Labels are stored as one byte each. - self._probability_chunks: list[NDArray[np.float64]] = [] - self._outcome_chunks: list[NDArray[np.uint8]] = [] - self._retained = 0 - - def _append(self, p: NDArray[np.float64], y: NDArray[np.float64]) -> None: - self._probability_chunks.append(p.copy()) - self._outcome_chunks.append(y.astype(np.uint8)) - self._retained += p.size - excess = 0 if self.window_size is None else self._retained - self.window_size - while excess > 0: - oldest = self._probability_chunks[0].size - if oldest <= excess: - del self._probability_chunks[0], self._outcome_chunks[0] - self._retained -= oldest - excess -= oldest - continue - self._probability_chunks[0] = _drop_head( - self._probability_chunks[0], excess - ) - self._outcome_chunks[0] = _drop_head(self._outcome_chunks[0], excess) - self._retained -= excess - excess = 0 - - def _history(self) -> tuple[NDArray[np.float64], NDArray[np.float64]]: - if len(self._probability_chunks) > 1: - self._probability_chunks = [np.concatenate(self._probability_chunks)] - self._outcome_chunks = [np.concatenate(self._outcome_chunks)] - if not self._probability_chunks: - return np.empty(0), np.empty(0) - return ( - self._probability_chunks[0], - self._outcome_chunks[0].astype(np.float64), - ) - @property def is_ready(self) -> bool: """Return whether an upstream fit has succeeded.""" @@ -170,14 +124,12 @@ def is_ready(self) -> bool: @property def history_bytes(self) -> int: """Return retained input-array bytes, excluding the fitted model.""" - return sum( - chunk.nbytes for chunk in self._probability_chunks + self._outcome_chunks - ) + return self._probabilities.nbytes + self._outcomes.nbytes @property def retained_observations(self) -> int: """Return the number of outcomes available for the next refit.""" - return self._retained + return self._outcomes.size def calibrate(self, probabilities: ArrayLike) -> NDArray[np.float64]: """Apply the last completed fit, or identity before readiness.""" @@ -188,29 +140,32 @@ def calibrate(self, probabilities: ArrayLike) -> NDArray[np.float64]: return self._model.predict(p) return self._model.predict_proba(p[:, None])[:, 1] - def _fit( - self, p: NDArray[np.float64], y: NDArray[np.float64] - ) -> CalibratedClassifierCV | IsotonicRegression: - if self.method == "isotonic": - return IsotonicRegression(y_min=0.0, y_max=1.0, out_of_bounds="clip").fit( - p, y - ) - return CalibratedClassifierCV( - FrozenEstimator(ProbabilityClassifier()), - method=self.method, - cv=_CalibrationRows(), - ).fit(p[:, None], y) - def update(self, probabilities: ArrayLike, outcomes: ArrayLike) -> Self: """Retain new observations and refit at the configured cadence.""" p, y = paired_binary_data(probabilities, outcomes) if self.frozen: return self - self._append(p, y) - if (self.n_updates + 1) % self.refit_every == 0: - retained_p, retained_y = self._history() - if np.unique(retained_y).size == 2: - self._model = self._fit(retained_p, retained_y) + retained_p = np.concatenate((self._probabilities, p)) + retained_y = np.concatenate((self._outcomes, y.astype(np.uint8))) + if self.window_size is not None: + retained_p = retained_p[-self.window_size :].copy() + retained_y = retained_y[-self.window_size :].copy() + model = self._model + if (self.n_updates + 1) % self.refit_every == 0 and np.unique( + retained_y + ).size == 2: + if self.method == "isotonic": + model = IsotonicRegression( + y_min=0.0, y_max=1.0, out_of_bounds="clip" + ).fit(retained_p, retained_y.astype(np.float64)) + else: + model = CalibratedClassifierCV( + FrozenEstimator(ProbabilityClassifier()), + method=self.method, + cv=_CalibrationRows(), + ).fit(retained_p[:, None], retained_y.astype(np.float64)) + self._probabilities, self._outcomes = retained_p, retained_y + self._model = model self.n_updates += 1 self.n_observations += p.size return self diff --git a/src/streamcal/evaluation.py b/src/streamcal/evaluation.py index 08c4fca..7ed524e 100644 --- a/src/streamcal/evaluation.py +++ b/src/streamcal/evaluation.py @@ -209,28 +209,6 @@ def paired_interval( return float(losses.mean()), float(interval[0, 0]), float(interval[1, 0]) -def _batch_brier(outcomes: np.ndarray, probabilities: np.ndarray) -> float: - """Brier score of one batch, without re-validating its inputs. - - ``brier_score`` validates its inputs on every call, which costs about a - millisecond; per batch of one that dominates the whole evaluation. Every - forecast is still validated once, when the full stream is scored. The - expression is scikit-learn's two-column form rather than - ``mean((y - p) ** 2)``, so the result matches ``brier_score`` exactly. - - Args: - outcomes: Binary outcomes. - probabilities: Forecast probabilities aligned with ``outcomes``. - - Returns: - Mean squared probability error. - """ - squared = ((1.0 - outcomes) - (1.0 - probabilities)) ** 2 + ( - outcomes - probabilities - ) ** 2 - return float(np.mean(squared) * 0.5) - - def retained_array_bytes(calibrator: BaseCalibrator) -> int: """Return owned numerical-array storage across state and upstream models.""" seen: set[int] = set() @@ -368,7 +346,7 @@ def compare_prequential( predict_calls.append(perf_counter_ns() - before) predict_ns += predict_calls[-1] calibrated_batches.append(calibrated) - batch_score = _batch_brier(batch_outcomes, calibrated) + batch_score = brier_score(batch_outcomes, calibrated) batch_scores.append(batch_score) cumulative_squared_error += batch_score * calibrated.size cumulative_observations += calibrated.size diff --git a/tests/test_kernels.py b/tests/test_kernels.py index aaf9848..ff1cc53 100644 --- a/tests/test_kernels.py +++ b/tests/test_kernels.py @@ -7,7 +7,7 @@ import pytest import streamcal -from streamcal import BatchCalibrator, brier_score +from streamcal import BatchCalibrator from streamcal.calibrators import ( _GRID_INDEX_FROM, _GRID_INTERP_FROM, @@ -16,7 +16,6 @@ _grid_interp, _interpolate, ) -from streamcal.evaluation import _batch_brier def _adversarial_points(n_bins, rng): @@ -61,7 +60,7 @@ def test_size_switches_agree_on_both_sides_of_the_threshold(n_bins): @pytest.mark.parametrize("window", [None, 1, 7, 50]) -def test_chunked_history_refits_on_exactly_the_retained_rows(window): +def test_one_byte_labels_refit_on_exactly_the_retained_rows(window): rng = np.random.default_rng(1) chunked = BatchCalibrator(window_size=window) seen_p, seen_y = [], [] @@ -94,14 +93,6 @@ def test_history_does_not_alias_caller_arrays(): assert np.array_equal(model.calibrate(grid), direct.calibrate(grid)) -def test_batch_brier_matches_the_validated_metric_bit_for_bit(): - rng = np.random.default_rng(2) - for size in (1, 2, 3, 336, 10_001): - p = rng.random(size) - y = (rng.random(size) < p).astype(float) - assert _batch_brier(y, p) == brier_score(y, p) - - def test_importing_the_streaming_calibrator_does_not_load_scikit_learn(): probe = ( "import sys, streamcal\n"