torchisotonic 0.1.0__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
@@ -0,0 +1,14 @@
1
+ __pycache__/
2
+ *.py[cod]
3
+ .venv/
4
+ dist/
5
+ build/
6
+ *.egg-info/
7
+ .coverage
8
+ .pytest_cache/
9
+ .mypy_cache/
10
+ .ruff_cache/
11
+ .hypothesis/
12
+ .DS_Store
13
+ kaggle/output/
14
+ cloud/output/
@@ -0,0 +1,21 @@
1
+ MIT License
2
+
3
+ Copyright (c) 2026 Lena Barretta
4
+
5
+ Permission is hereby granted, free of charge, to any person obtaining a copy
6
+ of this software and associated documentation files (the "Software"), to deal
7
+ in the Software without restriction, including without limitation the rights
8
+ to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
9
+ copies of the Software, and to permit persons to whom the Software is
10
+ furnished to do so, subject to the following conditions:
11
+
12
+ The above copyright notice and this permission notice shall be included in all
13
+ copies or substantial portions of the Software.
14
+
15
+ THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
16
+ IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
17
+ FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
18
+ AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
19
+ LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
20
+ OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
21
+ SOFTWARE.
@@ -0,0 +1,183 @@
1
+ Metadata-Version: 2.5
2
+ Name: torchisotonic
3
+ Version: 0.1.0
4
+ Summary: Batched isotonic regression for PyTorch tensors, fast on GPU.
5
+ Project-URL: Homepage, https://github.com/LenaBarretta/torchisotonic
6
+ Project-URL: Issues, https://github.com/LenaBarretta/torchisotonic/issues
7
+ Author: Lena Barretta
8
+ License-Expression: MIT
9
+ License-File: LICENSE
10
+ Keywords: calibration,cuda,gpu,isotonic-regression,pava,pytorch,torch
11
+ Classifier: Development Status :: 3 - Alpha
12
+ Classifier: Environment :: GPU :: NVIDIA CUDA
13
+ Classifier: Intended Audience :: Developers
14
+ Classifier: Intended Audience :: Science/Research
15
+ Classifier: Programming Language :: C++
16
+ Classifier: Programming Language :: Python :: 3
17
+ Classifier: Programming Language :: Python :: 3 :: Only
18
+ Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
19
+ Classifier: Topic :: Scientific/Engineering :: Mathematics
20
+ Classifier: Typing :: Typed
21
+ Requires-Python: >=3.10
22
+ Requires-Dist: ninja
23
+ Requires-Dist: torch>=2.4
24
+ Description-Content-Type: text/markdown
25
+
26
+ # torchisotonic
27
+
28
+ Batched isotonic regression for PyTorch tensors, fast on the GPU.
29
+
30
+ Isotonic regression finds the non-decreasing (or non-increasing) sequence closest to the
31
+ data in weighted least squares. It is what `sklearn.isotonic.isotonic_regression` and
32
+ `scipy.optimize.isotonic_regression` compute for one sequence; torchisotonic computes it for
33
+ every row of a tensor at once, on the CPU or with CUDA kernels on the GPU, from a few
34
+ elements per row to a hundred million.
35
+
36
+ ```python
37
+ import torch
38
+ import torchisotonic
39
+
40
+ y = torch.randn(1000, 5000, device="cuda")
41
+ fit = torchisotonic.isotonic(y) # every row, non-decreasing
42
+
43
+ w = torch.rand(1000, 5000, device="cuda") + 0.1
44
+ fit = torchisotonic.isotonic(y, w, increasing=False, y_min=0.0, y_max=1.0)
45
+ ```
46
+
47
+ ## What it is for
48
+
49
+ Whenever a monotone fit has to run often, on long data, for many curves at once, or inside a
50
+ model during training.
51
+
52
+ **Calibrating a classifier.** A model's scores rank well but are not probabilities; isotonic
53
+ regression maps each score to the rate of positives observed at that score, as scikit-learn's
54
+ `IsotonicRegression` and `CalibratedClassifierCV(method="isotonic")` do. A curve over 6.5 million
55
+ held-out scores takes 42 ms on a T4, against 2 s for scikit-learn. These two functions give
56
+ scikit-learn's curve and predictions, equal scores pooled into their mean:
57
+
58
+ ```python
59
+ def fit_calibration(scores, labels):
60
+ """The calibration curve, as scikit-learn's IsotonicRegression fits it."""
61
+ s, order = torch.sort(scores)
62
+ xs, inverse, counts = torch.unique_consecutive(s, return_inverse=True, return_counts=True)
63
+ sums = torch.zeros(len(xs), dtype=torch.float64, device=scores.device)
64
+ sums.index_add_(0, inverse, labels[order].double())
65
+ return xs, torchisotonic.isotonic(sums / counts, counts.double())
66
+
67
+
68
+ def calibrate(new_scores, xs, fit):
69
+ """Calibrated probabilities: linear between the curve's points, constant beyond them."""
70
+ i = torch.searchsorted(xs, new_scores).clamp(1, len(xs) - 1)
71
+ t = ((new_scores - xs[i - 1]) / (xs[i] - xs[i - 1])).clamp(0, 1)
72
+ return fit[i - 1] + t * (fit[i] - fit[i - 1])
73
+ ```
74
+
75
+ [`examples/ctr_calibration`](examples/ctr_calibration) calibrates a click model on 45.8
76
+ million Criteo impressions this way, with one curve per segment for 280 segments in one call.
77
+
78
+ **Many curves at once.** Every row of the input is its own problem, so a thousand curves cost
79
+ little more than one: per segment, per threshold, or per bootstrap sample. For instance, a
80
+ band of uncertainty around a calibration curve, from a Bayesian bootstrap of the labels `y`
81
+ ordered by score:
82
+
83
+ ```python
84
+ w = torch.empty(1000, len(y), device=y.device).exponential_() # one resample per row
85
+ curves = torchisotonic.isotonic(y.expand(1000, -1), w)
86
+ low, high = curves.quantile(0.05, dim=0), curves.quantile(0.95, dim=0)
87
+ ```
88
+
89
+ **Inside a model.** The fit is differentiable in `y` and in the weights, so it can enforce
90
+ monotonicity in a network and train through it: predicted quantiles that never cross,
91
+ dose-response or survival curves that must not go the wrong way, or the projection step of
92
+ differentiable sorting and ranking.
93
+
94
+ ```python
95
+ q = model(x) # [batch, levels]: quantiles at increasing levels
96
+ q = torchisotonic.isotonic(q) # never crossing, with gradients
97
+ loss = pinball(q, target, levels) # any quantile loss
98
+ ```
99
+
100
+ ## Semantics
101
+
102
+ The semantics are scikit-learn's: weights must be strictly positive (a weight of the shape of
103
+ `y`, or one that broadcasts to it), the result is clipped to `[y_min, y_max]` if they are
104
+ given, `dim` picks the axis, and a row holding a NaN comes back as NaN. Float input keeps its
105
+ dtype (float16 and bfloat16 are computed in float32), integer input becomes the default
106
+ float dtype. The fit is differentiable in `y` and in the weights.
107
+
108
+ On the GPU, float32 is accurate to the format's rounding by default, and the fit never steps
109
+ down. `exact=False` makes it about 1.5–2 times faster on short and medium rows, at the cost of
110
+ accuracy on hard inputs (strongly varying weights, increments near float32's resolution): errors
111
+ of a few parts in a thousand, and steps down of up to about 1e-3 of the data's scale; float64
112
+ is unaffected.
113
+
114
+ Rows of unequal length go in one batch padded at the end with a value above each row's maximum
115
+ (below its minimum for `increasing=False`): the padding never pools with the data, so the fit
116
+ of each row is exactly its own. Rows of very different lengths fit faster laid end to end in a single row,
117
+ as [`examples/ctr_calibration`](examples/ctr_calibration) does.
118
+
119
+ ## Installation
120
+
121
+ ```bash
122
+ pip install torchisotonic
123
+ ```
124
+
125
+ The C++/CUDA code is compiled the first time it is used (about a minute), which needs `ninja`,
126
+ a C++ compiler and, for the GPU, the CUDA toolkit. **Status: alpha**; the API may still change.
127
+
128
+ ## Performance
129
+
130
+ On an NVIDIA T4 (Kaggle), float32, unweighted, decreasing fits of data noisy around a trend,
131
+ in milliseconds; best of ten calls on a warmed-up GPU, data already on the device for the GPU
132
+ columns:
133
+
134
+ | batch × n | torchisotonic (cuda) | `exact=False` | torchsort (cuda) | torchisotonic (cpu) | torchsort (cpu) | scipy, a loop over rows |
135
+ | --- | ---: | ---: | ---: | ---: | ---: | ---: |
136
+ | 1,000,000 × 16 | **2.16** | 1.21 | 72.0 | 407 | 677 | 14,363 |
137
+ | 100,000 × 128 | **1.71** | 1.30 | 136 | 280 | 492 | 1,613 |
138
+ | 10,000 × 1,000 | **2.29** | 1.22 | 99.2 | 160 | 354 | 318 |
139
+ | 1,000 × 10,000 | **1.72** | 1.12 | 197 | 160 | 358 | 178 |
140
+ | 100 × 100,000 | **1.82** | 1.23 | 396 | 158 | 352 | 166 |
141
+ | 10 × 1,000,000 | **1.95** | 1.35 | 1,656 | 157 | 357 | 170 |
142
+ | 1 × 10,000,000 | **2.07** | 1.47 | 5,493 | 236 | 355 | 176 |
143
+ | 100 × 1,000 | **0.09** | 0.07 | 3.0 | 1.7 | 3.5 | 2.9 |
144
+ | 10 × 10,000 | **0.22** | 0.20 | 14.0 | 1.7 | 3.5 | 1.6 |
145
+ | 1 × 100,000 | **0.26** | 0.25 | 51.7 | 1.8 | 3.5 | 1.6 |
146
+ | 1 × 1,000,000 | **0.51** | 0.44 | 550 | 17.9 | 35.0 | 15.5 |
147
+ | 1 × 100,000,000 | **16.9** | 10.8 | — | 3,552 | — | — |
148
+
149
+ torchsort solves each row with one GPU thread, so a long row is slow; it was not run on 10⁸
150
+ elements. The CPU columns are on Kaggle's CPU, so they say how the libraries compare, not how
151
+ fast a good CPU is. On an L4, an A100 and an H100 one row of 10⁷ elements takes 1.2, 0.8 and
152
+ 0.5 ms, and 10⁸ elements 8.2, 4.3 and 2.7 ms. `bench/bench.py` reproduces the table.
153
+
154
+ ## How it works
155
+
156
+ The fit at each element is the slope of the lower convex hull of the cumulative sums of the
157
+ data (the "greatest convex minorant"), so the GPU code builds that hull in parallel:
158
+
159
+ 1. Every lane of a warp builds the hull of 32 neighbouring points with the monotone chain,
160
+ its stack kept as a bit mask; then the lane hulls of a chunk of 1024 to 4096 elements merge
161
+ pairwise in shared memory, each merge walking to its bridge and clearing the bits under it.
162
+ Short rows are packed several per warp. A row that fits in one chunk is finished there.
163
+ 2. Longer rows merge their chunks' hulls pairwise, in rounds, as lists of segments of the
164
+ chunks' hulls, so that a merge costs segments rather than vertices.
165
+ 3. Every element takes the slope of the hull edge above it.
166
+
167
+ [`docs/hull_prototype.py`](docs/hull_prototype.py) is a plain-numpy reference of the parallel
168
+ algorithm; [`src/torchisotonic/csrc/tiles.cuh`](src/torchisotonic/csrc/tiles.cuh)
169
+ describes the kernels. On the CPU, rows are solved with PAVA, in parallel.
170
+
171
+ ## Development
172
+
173
+ ```bash
174
+ uv sync
175
+ uv run pytest
176
+ ```
177
+
178
+ The tests compare against scipy and scikit-learn (with hypothesis) and, on a machine with a
179
+ GPU, the CUDA operator against the CPU one; without a GPU the CUDA tests are skipped.
180
+
181
+ ## License
182
+
183
+ MIT
@@ -0,0 +1,158 @@
1
+ # torchisotonic
2
+
3
+ Batched isotonic regression for PyTorch tensors, fast on the GPU.
4
+
5
+ Isotonic regression finds the non-decreasing (or non-increasing) sequence closest to the
6
+ data in weighted least squares. It is what `sklearn.isotonic.isotonic_regression` and
7
+ `scipy.optimize.isotonic_regression` compute for one sequence; torchisotonic computes it for
8
+ every row of a tensor at once, on the CPU or with CUDA kernels on the GPU, from a few
9
+ elements per row to a hundred million.
10
+
11
+ ```python
12
+ import torch
13
+ import torchisotonic
14
+
15
+ y = torch.randn(1000, 5000, device="cuda")
16
+ fit = torchisotonic.isotonic(y) # every row, non-decreasing
17
+
18
+ w = torch.rand(1000, 5000, device="cuda") + 0.1
19
+ fit = torchisotonic.isotonic(y, w, increasing=False, y_min=0.0, y_max=1.0)
20
+ ```
21
+
22
+ ## What it is for
23
+
24
+ Whenever a monotone fit has to run often, on long data, for many curves at once, or inside a
25
+ model during training.
26
+
27
+ **Calibrating a classifier.** A model's scores rank well but are not probabilities; isotonic
28
+ regression maps each score to the rate of positives observed at that score, as scikit-learn's
29
+ `IsotonicRegression` and `CalibratedClassifierCV(method="isotonic")` do. A curve over 6.5 million
30
+ held-out scores takes 42 ms on a T4, against 2 s for scikit-learn. These two functions give
31
+ scikit-learn's curve and predictions, equal scores pooled into their mean:
32
+
33
+ ```python
34
+ def fit_calibration(scores, labels):
35
+ """The calibration curve, as scikit-learn's IsotonicRegression fits it."""
36
+ s, order = torch.sort(scores)
37
+ xs, inverse, counts = torch.unique_consecutive(s, return_inverse=True, return_counts=True)
38
+ sums = torch.zeros(len(xs), dtype=torch.float64, device=scores.device)
39
+ sums.index_add_(0, inverse, labels[order].double())
40
+ return xs, torchisotonic.isotonic(sums / counts, counts.double())
41
+
42
+
43
+ def calibrate(new_scores, xs, fit):
44
+ """Calibrated probabilities: linear between the curve's points, constant beyond them."""
45
+ i = torch.searchsorted(xs, new_scores).clamp(1, len(xs) - 1)
46
+ t = ((new_scores - xs[i - 1]) / (xs[i] - xs[i - 1])).clamp(0, 1)
47
+ return fit[i - 1] + t * (fit[i] - fit[i - 1])
48
+ ```
49
+
50
+ [`examples/ctr_calibration`](examples/ctr_calibration) calibrates a click model on 45.8
51
+ million Criteo impressions this way, with one curve per segment for 280 segments in one call.
52
+
53
+ **Many curves at once.** Every row of the input is its own problem, so a thousand curves cost
54
+ little more than one: per segment, per threshold, or per bootstrap sample. For instance, a
55
+ band of uncertainty around a calibration curve, from a Bayesian bootstrap of the labels `y`
56
+ ordered by score:
57
+
58
+ ```python
59
+ w = torch.empty(1000, len(y), device=y.device).exponential_() # one resample per row
60
+ curves = torchisotonic.isotonic(y.expand(1000, -1), w)
61
+ low, high = curves.quantile(0.05, dim=0), curves.quantile(0.95, dim=0)
62
+ ```
63
+
64
+ **Inside a model.** The fit is differentiable in `y` and in the weights, so it can enforce
65
+ monotonicity in a network and train through it: predicted quantiles that never cross,
66
+ dose-response or survival curves that must not go the wrong way, or the projection step of
67
+ differentiable sorting and ranking.
68
+
69
+ ```python
70
+ q = model(x) # [batch, levels]: quantiles at increasing levels
71
+ q = torchisotonic.isotonic(q) # never crossing, with gradients
72
+ loss = pinball(q, target, levels) # any quantile loss
73
+ ```
74
+
75
+ ## Semantics
76
+
77
+ The semantics are scikit-learn's: weights must be strictly positive (a weight of the shape of
78
+ `y`, or one that broadcasts to it), the result is clipped to `[y_min, y_max]` if they are
79
+ given, `dim` picks the axis, and a row holding a NaN comes back as NaN. Float input keeps its
80
+ dtype (float16 and bfloat16 are computed in float32), integer input becomes the default
81
+ float dtype. The fit is differentiable in `y` and in the weights.
82
+
83
+ On the GPU, float32 is accurate to the format's rounding by default, and the fit never steps
84
+ down. `exact=False` makes it about 1.5–2 times faster on short and medium rows, at the cost of
85
+ accuracy on hard inputs (strongly varying weights, increments near float32's resolution): errors
86
+ of a few parts in a thousand, and steps down of up to about 1e-3 of the data's scale; float64
87
+ is unaffected.
88
+
89
+ Rows of unequal length go in one batch padded at the end with a value above each row's maximum
90
+ (below its minimum for `increasing=False`): the padding never pools with the data, so the fit
91
+ of each row is exactly its own. Rows of very different lengths fit faster laid end to end in a single row,
92
+ as [`examples/ctr_calibration`](examples/ctr_calibration) does.
93
+
94
+ ## Installation
95
+
96
+ ```bash
97
+ pip install torchisotonic
98
+ ```
99
+
100
+ The C++/CUDA code is compiled the first time it is used (about a minute), which needs `ninja`,
101
+ a C++ compiler and, for the GPU, the CUDA toolkit. **Status: alpha**; the API may still change.
102
+
103
+ ## Performance
104
+
105
+ On an NVIDIA T4 (Kaggle), float32, unweighted, decreasing fits of data noisy around a trend,
106
+ in milliseconds; best of ten calls on a warmed-up GPU, data already on the device for the GPU
107
+ columns:
108
+
109
+ | batch × n | torchisotonic (cuda) | `exact=False` | torchsort (cuda) | torchisotonic (cpu) | torchsort (cpu) | scipy, a loop over rows |
110
+ | --- | ---: | ---: | ---: | ---: | ---: | ---: |
111
+ | 1,000,000 × 16 | **2.16** | 1.21 | 72.0 | 407 | 677 | 14,363 |
112
+ | 100,000 × 128 | **1.71** | 1.30 | 136 | 280 | 492 | 1,613 |
113
+ | 10,000 × 1,000 | **2.29** | 1.22 | 99.2 | 160 | 354 | 318 |
114
+ | 1,000 × 10,000 | **1.72** | 1.12 | 197 | 160 | 358 | 178 |
115
+ | 100 × 100,000 | **1.82** | 1.23 | 396 | 158 | 352 | 166 |
116
+ | 10 × 1,000,000 | **1.95** | 1.35 | 1,656 | 157 | 357 | 170 |
117
+ | 1 × 10,000,000 | **2.07** | 1.47 | 5,493 | 236 | 355 | 176 |
118
+ | 100 × 1,000 | **0.09** | 0.07 | 3.0 | 1.7 | 3.5 | 2.9 |
119
+ | 10 × 10,000 | **0.22** | 0.20 | 14.0 | 1.7 | 3.5 | 1.6 |
120
+ | 1 × 100,000 | **0.26** | 0.25 | 51.7 | 1.8 | 3.5 | 1.6 |
121
+ | 1 × 1,000,000 | **0.51** | 0.44 | 550 | 17.9 | 35.0 | 15.5 |
122
+ | 1 × 100,000,000 | **16.9** | 10.8 | — | 3,552 | — | — |
123
+
124
+ torchsort solves each row with one GPU thread, so a long row is slow; it was not run on 10⁸
125
+ elements. The CPU columns are on Kaggle's CPU, so they say how the libraries compare, not how
126
+ fast a good CPU is. On an L4, an A100 and an H100 one row of 10⁷ elements takes 1.2, 0.8 and
127
+ 0.5 ms, and 10⁸ elements 8.2, 4.3 and 2.7 ms. `bench/bench.py` reproduces the table.
128
+
129
+ ## How it works
130
+
131
+ The fit at each element is the slope of the lower convex hull of the cumulative sums of the
132
+ data (the "greatest convex minorant"), so the GPU code builds that hull in parallel:
133
+
134
+ 1. Every lane of a warp builds the hull of 32 neighbouring points with the monotone chain,
135
+ its stack kept as a bit mask; then the lane hulls of a chunk of 1024 to 4096 elements merge
136
+ pairwise in shared memory, each merge walking to its bridge and clearing the bits under it.
137
+ Short rows are packed several per warp. A row that fits in one chunk is finished there.
138
+ 2. Longer rows merge their chunks' hulls pairwise, in rounds, as lists of segments of the
139
+ chunks' hulls, so that a merge costs segments rather than vertices.
140
+ 3. Every element takes the slope of the hull edge above it.
141
+
142
+ [`docs/hull_prototype.py`](docs/hull_prototype.py) is a plain-numpy reference of the parallel
143
+ algorithm; [`src/torchisotonic/csrc/tiles.cuh`](src/torchisotonic/csrc/tiles.cuh)
144
+ describes the kernels. On the CPU, rows are solved with PAVA, in parallel.
145
+
146
+ ## Development
147
+
148
+ ```bash
149
+ uv sync
150
+ uv run pytest
151
+ ```
152
+
153
+ The tests compare against scipy and scikit-learn (with hypothesis) and, on a machine with a
154
+ GPU, the CUDA operator against the CPU one; without a GPU the CUDA tests are skipped.
155
+
156
+ ## License
157
+
158
+ MIT
@@ -0,0 +1,113 @@
1
+ """Reference prototype of the parallel algorithm for one large problem, in plain numpy.
2
+
3
+ Isotonic regression is the slope of the greatest convex minorant of the cumulative sum
4
+ diagram (Barlow et al.; see also Kearsley, J. Res. NIST 111(2), 2006): with points
5
+ P_k = (W_k, S_k), W_k = w_0 + ... + w_{k-1} and S_k = w_0 y_0 + ... + w_{k-1} y_{k-1}, the fit
6
+ at element i is the slope of the lower convex hull of P_0..P_n over [W_i, W_{i+1}].
7
+
8
+ The GPU version does this in parallel:
9
+
10
+ 1. Cut the problem into chunks of C elements; solve each with PAVA, independently. A chunk's
11
+ blocks are the edges of its own lower hull, so its block boundaries are hull vertices.
12
+ 2. Merge neighbouring hulls in log2(#chunks) rounds. Two neighbouring lower hulls share their
13
+ boundary point. If the hull is already convex there, they just concatenate; otherwise
14
+ there is one bridge, a common lower tangent from vertex p of the left hull to vertex q of
15
+ the right one. The vertices strictly between p and q drop out, and the elements between
16
+ them form one block.
17
+ 3. Every element takes the slope of the final hull edge above it.
18
+
19
+ The bridge is found the way the GPU does it, all candidates at once: every left vertex p
20
+ finds its own lower tangent to the right hull by binary search (slopes from p to the right
21
+ hull's vertices fall, then rise), and the bridge is the p whose tangent also supports the
22
+ left hull at p (or, with rounding, comes closest to it).
23
+
24
+ This file only checks the logic (tests/test_hull_prototype.py runs it against scipy). The
25
+ coordinates here are plain float64 prefix sums; the CUDA kernel works in local coordinates.
26
+ """
27
+
28
+ from __future__ import annotations
29
+
30
+ import itertools
31
+
32
+ import numpy as np
33
+ import numpy.typing as npt
34
+
35
+ Array = npt.NDArray[np.float64]
36
+ Hull = list[int] # indices k of the points P_k on a lower hull, increasing
37
+
38
+
39
+ def chunk_hull(w: Array, y: Array, lo: int, hi: int) -> Hull:
40
+ """PAVA on elements lo..hi-1; returns its hull as point indices lo..hi."""
41
+ starts: list[int] = []
42
+ weights: list[float] = []
43
+ totals: list[float] = []
44
+ for i in range(lo, hi):
45
+ starts.append(i)
46
+ weights.append(float(w[i]))
47
+ totals.append(float(w[i] * y[i]))
48
+ while len(starts) > 1 and totals[-2] * weights[-1] >= totals[-1] * weights[-2]:
49
+ top_weight, top_total = weights.pop(), totals.pop()
50
+ weights[-1] += top_weight
51
+ totals[-1] += top_total
52
+ starts.pop()
53
+ return [*starts, hi]
54
+
55
+
56
+ def slope(xs: Array, ys: Array, a: int, b: int) -> float:
57
+ return float((ys[b] - ys[a]) / (xs[b] - xs[a]))
58
+
59
+
60
+ def tangent(xs: Array, ys: Array, p: int, right: Hull) -> int:
61
+ """Position in ``right`` of the vertex where the lower tangent from point p touches."""
62
+ lo, hi = 0, len(right) - 1
63
+ while lo < hi: # slopes from p fall then rise along the right hull: find the minimum
64
+ mid = (lo + hi) // 2
65
+ if slope(xs, ys, p, right[mid + 1]) < slope(xs, ys, p, right[mid]):
66
+ lo = mid + 1
67
+ else:
68
+ hi = mid
69
+ return lo
70
+
71
+
72
+ def merge(xs: Array, ys: Array, left: Hull, right: Hull) -> Hull:
73
+ """Lower hull of two neighbouring lower hulls (left[-1] == right[0])."""
74
+ assert left[-1] == right[0]
75
+ if len(left) == 1 or len(right) == 1:
76
+ return left + right[1:]
77
+ if slope(xs, ys, left[-2], left[-1]) <= slope(xs, ys, right[0], right[1]):
78
+ return left + right[1:] # already convex at the shared point
79
+ # Every left vertex but the shared one is a candidate; on the GPU, one thread each. In
80
+ # exact arithmetic exactly one has a tangent that also supports the left hull (violation
81
+ # 0). Rounding can make the tests contradict each other on (nearly) collinear points, so
82
+ # take the candidate that violates support least: an argmin, i.e. one GPU reduction.
83
+ best, best_violation, best_b = 0, float("inf"), 0
84
+ for a in range(len(left) - 1):
85
+ p = left[a]
86
+ b = tangent(xs, ys, p, right[1:]) + 1
87
+ s = slope(xs, ys, p, right[b])
88
+ before = 0.0 if a == 0 else slope(xs, ys, left[a - 1], p) - s
89
+ after = s - slope(xs, ys, p, left[a + 1])
90
+ violation = max(before, after, 0.0)
91
+ if violation < best_violation:
92
+ best, best_violation, best_b = a, violation, b
93
+ return left[: best + 1] + right[best_b:]
94
+
95
+
96
+ def isotonic(y: Array, w: Array | None = None, chunk: int = 4) -> Array:
97
+ """Increasing isotonic regression of y, computed the parallel way."""
98
+ n = len(y)
99
+ if n == 0:
100
+ return y.astype(np.float64)
101
+ w = np.ones(n) if w is None else np.asarray(w, dtype=np.float64)
102
+ y = np.asarray(y, dtype=np.float64)
103
+ xs = np.concatenate([[0.0], np.cumsum(w)])
104
+ ys = np.concatenate([[0.0], np.cumsum(w * y)])
105
+ hulls = [chunk_hull(w, y, lo, min(lo + chunk, n)) for lo in range(0, n, chunk)]
106
+ while len(hulls) > 1: # one round of pairwise merges per level of the tree
107
+ merged = [merge(xs, ys, hulls[i], hulls[i + 1]) for i in range(0, len(hulls) - 1, 2)]
108
+ hulls = merged + ([hulls[-1]] if len(hulls) % 2 else [])
109
+ vertices = hulls[0]
110
+ out = np.empty(n)
111
+ for a, b in itertools.pairwise(vertices):
112
+ out[a:b] = slope(xs, ys, a, b)
113
+ return out
@@ -0,0 +1,83 @@
1
+ [build-system]
2
+ requires = ["hatchling>=1.24"]
3
+ build-backend = "hatchling.build"
4
+
5
+ [project]
6
+ name = "torchisotonic"
7
+ dynamic = ["version"]
8
+ description = "Batched isotonic regression for PyTorch tensors, fast on GPU."
9
+ readme = "README.md"
10
+ license = "MIT"
11
+ license-files = ["LICENSE"]
12
+ authors = [{ name = "Lena Barretta" }]
13
+ requires-python = ">=3.10"
14
+ # The C++/CUDA operators are compiled on first use, which needs ninja and a C++ compiler.
15
+ dependencies = ["torch>=2.4", "ninja"]
16
+ keywords = ["pytorch", "torch", "isotonic-regression", "calibration", "pava", "cuda", "gpu"]
17
+ classifiers = [
18
+ "Development Status :: 3 - Alpha",
19
+ "Environment :: GPU :: NVIDIA CUDA",
20
+ "Intended Audience :: Developers",
21
+ "Intended Audience :: Science/Research",
22
+ "Programming Language :: Python :: 3",
23
+ "Programming Language :: Python :: 3 :: Only",
24
+ "Programming Language :: C++",
25
+ "Topic :: Scientific/Engineering :: Mathematics",
26
+ "Topic :: Scientific/Engineering :: Artificial Intelligence",
27
+ "Typing :: Typed",
28
+ ]
29
+
30
+ [project.urls]
31
+ Homepage = "https://github.com/LenaBarretta/torchisotonic"
32
+ Issues = "https://github.com/LenaBarretta/torchisotonic/issues"
33
+
34
+ [dependency-groups]
35
+ dev = [
36
+ "pytest>=8",
37
+ "pytest-cov>=5",
38
+ "ruff>=0.6",
39
+ "mypy>=1.11",
40
+ "hypothesis>=6",
41
+ "numpy>=1.24",
42
+ "scipy>=1.12",
43
+ "scikit-learn>=1.4",
44
+ ]
45
+
46
+ # Develop and test against the small CPU-only torch build.
47
+ [tool.uv.sources]
48
+ torch = [{ index = "pytorch-cpu" }]
49
+
50
+ [[tool.uv.index]]
51
+ name = "pytorch-cpu"
52
+ url = "https://download.pytorch.org/whl/cpu"
53
+ explicit = true
54
+
55
+ [tool.hatch.build.targets.sdist]
56
+ # The sources, the tests (with docs/hull_prototype.py, which one of them checks), README and
57
+ # licence. notes/ and IDE settings are kept out of git through .git/info/exclude, which hatch
58
+ # does not read, so they are listed here too.
59
+ exclude = ["bench", "examples", ".github", "notes", ".idea", ".vscode", "uv.lock"]
60
+
61
+ [tool.hatch.version]
62
+ path = "src/torchisotonic/__init__.py"
63
+
64
+ [tool.pytest.ini_options]
65
+ addopts = "-q --cov=torchisotonic --cov-report=term-missing --cov-fail-under=100"
66
+ testpaths = ["tests"]
67
+
68
+ [tool.ruff]
69
+ line-length = 100
70
+ target-version = "py310"
71
+
72
+ [tool.ruff.lint]
73
+ select = ["E", "F", "W", "I", "B", "UP", "SIM", "RUF"]
74
+
75
+ [tool.mypy]
76
+ strict = true
77
+ files = ["src", "tests", "docs"]
78
+ # tests/test_hull_prototype.py imports docs/hull_prototype.py by putting docs/ on sys.path.
79
+ mypy_path = "docs"
80
+
81
+ [[tool.mypy.overrides]]
82
+ module = ["scipy", "scipy.*", "sklearn", "sklearn.*"]
83
+ ignore_missing_imports = true
@@ -0,0 +1,6 @@
1
+ """Batched isotonic regression for PyTorch tensors, fast on GPU."""
2
+
3
+ from torchisotonic._api import isotonic
4
+
5
+ __all__ = ["isotonic"]
6
+ __version__ = "0.1.0"