fastjsd 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.
- fastjsd-0.1.0/LICENSE +21 -0
- fastjsd-0.1.0/PKG-INFO +277 -0
- fastjsd-0.1.0/README.md +240 -0
- fastjsd-0.1.0/fastjsd/__init__.py +41 -0
- fastjsd-0.1.0/fastjsd/_api.py +394 -0
- fastjsd-0.1.0/fastjsd/_common.py +283 -0
- fastjsd-0.1.0/fastjsd/_numba.py +973 -0
- fastjsd-0.1.0/fastjsd/_numpy.py +236 -0
- fastjsd-0.1.0/fastjsd/_svml.py +402 -0
- fastjsd-0.1.0/fastjsd/reference.py +77 -0
- fastjsd-0.1.0/fastjsd.egg-info/PKG-INFO +277 -0
- fastjsd-0.1.0/fastjsd.egg-info/SOURCES.txt +21 -0
- fastjsd-0.1.0/fastjsd.egg-info/dependency_links.txt +1 -0
- fastjsd-0.1.0/fastjsd.egg-info/requires.txt +16 -0
- fastjsd-0.1.0/fastjsd.egg-info/top_level.txt +1 -0
- fastjsd-0.1.0/pyproject.toml +64 -0
- fastjsd-0.1.0/setup.cfg +4 -0
- fastjsd-0.1.0/tests/test_api.py +160 -0
- fastjsd-0.1.0/tests/test_bits_kernel.py +273 -0
- fastjsd-0.1.0/tests/test_correctness.py +109 -0
- fastjsd-0.1.0/tests/test_doctests.py +24 -0
- fastjsd-0.1.0/tests/test_properties.py +449 -0
- fastjsd-0.1.0/tests/test_svml.py +127 -0
fastjsd-0.1.0/LICENSE
ADDED
|
@@ -0,0 +1,21 @@
|
|
|
1
|
+
MIT License
|
|
2
|
+
|
|
3
|
+
Copyright (c) 2026 Brian Keith-Norambuena
|
|
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.
|
fastjsd-0.1.0/PKG-INFO
ADDED
|
@@ -0,0 +1,277 @@
|
|
|
1
|
+
Metadata-Version: 2.4
|
|
2
|
+
Name: fastjsd
|
|
3
|
+
Version: 0.1.0
|
|
4
|
+
Summary: Fast exact pairwise Jensen-Shannon distance matrices
|
|
5
|
+
Author: Brian Keith-Norambuena
|
|
6
|
+
License: MIT
|
|
7
|
+
Project-URL: Homepage, https://github.com/briankeithn/fastjsd
|
|
8
|
+
Project-URL: Source, https://github.com/briankeithn/fastjsd
|
|
9
|
+
Project-URL: Issues, https://github.com/briankeithn/fastjsd/issues
|
|
10
|
+
Keywords: jensen-shannon,distance-matrix,information-theory,sparse,numba
|
|
11
|
+
Classifier: Development Status :: 4 - Beta
|
|
12
|
+
Classifier: Intended Audience :: Science/Research
|
|
13
|
+
Classifier: License :: OSI Approved :: MIT License
|
|
14
|
+
Classifier: Programming Language :: Python :: 3
|
|
15
|
+
Classifier: Programming Language :: Python :: 3.10
|
|
16
|
+
Classifier: Programming Language :: Python :: 3.11
|
|
17
|
+
Classifier: Programming Language :: Python :: 3.12
|
|
18
|
+
Classifier: Programming Language :: Python :: 3.13
|
|
19
|
+
Classifier: Topic :: Scientific/Engineering :: Mathematics
|
|
20
|
+
Requires-Python: >=3.10
|
|
21
|
+
Description-Content-Type: text/markdown
|
|
22
|
+
License-File: LICENSE
|
|
23
|
+
Requires-Dist: numpy>=1.24
|
|
24
|
+
Provides-Extra: numba
|
|
25
|
+
Requires-Dist: numba>=0.58; extra == "numba"
|
|
26
|
+
Provides-Extra: test
|
|
27
|
+
Requires-Dist: pytest; extra == "test"
|
|
28
|
+
Requires-Dist: scipy; extra == "test"
|
|
29
|
+
Requires-Dist: numba>=0.58; extra == "test"
|
|
30
|
+
Provides-Extra: bench
|
|
31
|
+
Requires-Dist: pytest; extra == "bench"
|
|
32
|
+
Requires-Dist: scipy; extra == "bench"
|
|
33
|
+
Requires-Dist: numba>=0.58; extra == "bench"
|
|
34
|
+
Requires-Dist: mpmath; extra == "bench"
|
|
35
|
+
Requires-Dist: psutil; extra == "bench"
|
|
36
|
+
Dynamic: license-file
|
|
37
|
+
|
|
38
|
+
# fastjsd
|
|
39
|
+
|
|
40
|
+
Fast, exact, pairwise Jensen-Shannon distance matrices for Python.
|
|
41
|
+
|
|
42
|
+
`fastjsd` computes the full `n x n` matrix of Jensen-Shannon distances between
|
|
43
|
+
rows of a stack of discrete distributions. It is a drop-in replacement for
|
|
44
|
+
`scipy.spatial.distance.pdist(P, "jensenshannon")` that adds a base argument,
|
|
45
|
+
parallel kernels, and a sparse path that skips coordinates neither
|
|
46
|
+
distribution occupies.
|
|
47
|
+
|
|
48
|
+
```python
|
|
49
|
+
import numpy as np
|
|
50
|
+
from fastjsd import jsd_matrix
|
|
51
|
+
|
|
52
|
+
P = np.random.default_rng(0).dirichlet(np.full(32, 0.1), size=500)
|
|
53
|
+
D = jsd_matrix(P) # (500, 500), base 2, distances in [0, 1]
|
|
54
|
+
```
|
|
55
|
+
|
|
56
|
+
## Install
|
|
57
|
+
|
|
58
|
+
```bash
|
|
59
|
+
pip install fastjsd # NumPy backends only
|
|
60
|
+
pip install "fastjsd[numba]" # adds the parallel dense and sparse kernels
|
|
61
|
+
```
|
|
62
|
+
|
|
63
|
+
Python 3.10 or newer. NumPy is the only hard dependency; Numba is optional and
|
|
64
|
+
enables the compiled kernels. SciPy is needed only by the test suite and the
|
|
65
|
+
benchmark harnesses, and comes in with the `test` and `bench` extras.
|
|
66
|
+
|
|
67
|
+
## Why it is fast
|
|
68
|
+
|
|
69
|
+
For distributions `p, q` over `k` outcomes with `m = (p + q) / 2`,
|
|
70
|
+
|
|
71
|
+
```
|
|
72
|
+
JS(p, q) = (1/2) * sum_c [ p_c ln p_c + q_c ln q_c - 2 m_c ln m_c ]
|
|
73
|
+
```
|
|
74
|
+
|
|
75
|
+
Two properties of that arrangement do the work.
|
|
76
|
+
|
|
77
|
+
**One logarithm per pair-coordinate.** `p ln p` depends on a single row, so it
|
|
78
|
+
is tabulated once per entry and the inner loop evaluates exactly one
|
|
79
|
+
logarithm. SciPy's kernel groups the same quantity as `a ln(a/m) + b ln(b/m)`,
|
|
80
|
+
where `m` depends on the pair, so neither logarithm can be hoisted and it
|
|
81
|
+
pays two.
|
|
82
|
+
|
|
83
|
+
**Logarithms only on the support intersection.** Where exactly one side has
|
|
84
|
+
mass, the summand is exactly `a ln 2` (Lee, 1997). A pair therefore needs
|
|
85
|
+
logarithms only where both rows are non-zero, and everything else is a
|
|
86
|
+
constant times the mass outside the intersection. For rows of density `d` with
|
|
87
|
+
independent supports the expected intersection is `k d^2` rather than `k`.
|
|
88
|
+
This is exact, not an approximation, and it is what makes thresholded
|
|
89
|
+
membership and topic-mixture data cheap.
|
|
90
|
+
|
|
91
|
+
The sparse kernels add an occupancy bitmap so pairs whose supports cannot
|
|
92
|
+
intersect are skipped without walking either row.
|
|
93
|
+
|
|
94
|
+
## Exactness
|
|
95
|
+
|
|
96
|
+
The guarantees, with proofs and the measurements that support them, are in
|
|
97
|
+
[`docs/MATH.md`](https://github.com/briankeithn/fastjsd/blob/main/docs/MATH.md). In summary:
|
|
98
|
+
|
|
99
|
+
- Every per-coordinate summand is non-negative, so the sum over coordinates
|
|
100
|
+
has condition number exactly 1.
|
|
101
|
+
- Identical rows return exactly `0.0`, not float dust that `sqrt` amplifies.
|
|
102
|
+
This is a property of keeping the constant inside the sum; forms that fold
|
|
103
|
+
`ln 2` outside it lose the exact zero.
|
|
104
|
+
- Results agree with SciPy to within a few times `1e-13` across dense, sparse,
|
|
105
|
+
cross and condensed forms.
|
|
106
|
+
- The individual summand is ill-conditioned when `p` and `q` nearly coincide.
|
|
107
|
+
That is a property of the grouping, it is documented rather than hidden, and
|
|
108
|
+
`well_conditioned=True` selects a rearrangement that removes it.
|
|
109
|
+
|
|
110
|
+
## Benchmarks
|
|
111
|
+
|
|
112
|
+
`bench/` contains the harnesses, the baselines and two real membership
|
|
113
|
+
fixtures. Performance depends heavily on the machine, so this package ships
|
|
114
|
+
the instruments and one worked example rather than a headline number. The
|
|
115
|
+
harnesses live in the repository rather than in the wheel, so run them from a
|
|
116
|
+
checkout:
|
|
117
|
+
|
|
118
|
+
```bash
|
|
119
|
+
pip install -e ".[bench]"
|
|
120
|
+
python bench/run_all.py # fingerprints the machine, then measures
|
|
121
|
+
```
|
|
122
|
+
|
|
123
|
+
`run_all.py` records the CPU, core counts, library versions and thread
|
|
124
|
+
settings before timing anything, then runs the three timing harnesses that
|
|
125
|
+
have to be read together -- `aa_control.py`, `bits_vs_merge.py` and
|
|
126
|
+
`lee_control.py` -- in dependency order, into one fingerprinted results
|
|
127
|
+
directory. The first of them, `bench/aa_control.py`, duels the implementation
|
|
128
|
+
against *itself*, so the true ratio is 1.000 by construction and everything it
|
|
129
|
+
reports is the measurement noise floor. Read that number first: a speedup
|
|
130
|
+
smaller than the noise floor of the machine that produced it is not a result.
|
|
131
|
+
|
|
132
|
+
### One machine's numbers
|
|
133
|
+
|
|
134
|
+
[`docs/measurements/`](https://github.com/briankeithn/fastjsd/blob/main/docs/measurements/README.md) holds a complete run with
|
|
135
|
+
its fingerprint and noise floor attached. Single-threaded on a 4-core Xeon at
|
|
136
|
+
2.10 GHz, which compares kernels rather than thread counts since SciPy's
|
|
137
|
+
`pdist` is single-threaded regardless:
|
|
138
|
+
|
|
139
|
+
| against | synthetic grid | real fixtures |
|
|
140
|
+
|---|---|---|
|
|
141
|
+
| SciPy | 2.19x - 33.17x | 2.14x - 8.93x |
|
|
142
|
+
| Lee 1997, transcribed literally | 1.86x - 5.06x | 1.52x - 2.10x |
|
|
143
|
+
| Lee 1997, competently implemented | **9 of 14 cells win, 1.19x - 3.61x; 5 lose, up to 1.25x** | |
|
|
144
|
+
|
|
145
|
+
Read the third row, not the second. The two differ in what the baseline is
|
|
146
|
+
allowed to do, and the gap between them is large: the literal transcription
|
|
147
|
+
evaluates `a ln(a/s) + b ln(b/s)` as printed, paying two logarithms and two
|
|
148
|
+
divisions per matched coordinate, while the same identity factors as
|
|
149
|
+
`a ln a + b ln b - s ln s` and accepts the same tabulated `x ln x` this package
|
|
150
|
+
uses. Giving the baseline that tabulation and a 1984 occupancy prefilter --
|
|
151
|
+
neither of which has anything to do with where `ln 2` sits -- makes it 1.1x to
|
|
152
|
+
2.6x faster, and that difference is most of the apparent margin in row two.
|
|
153
|
+
|
|
154
|
+
So the honest comparison is the third row, and it includes five losses,
|
|
155
|
+
covering both real fixtures under the stricter cleanup. That is expected, as
|
|
156
|
+
the folded arrangement absorbs the one-sided mass into a constant and needs
|
|
157
|
+
no per-row matched-mass accumulators at all. What the unfolded arrangement
|
|
158
|
+
buys is the exact zero, not speed --
|
|
159
|
+
|
|
160
|
+
```
|
|
161
|
+
fastjsd exactly 0.0 in all 96 duplicate-row trials
|
|
162
|
+
Lee-1997 sparse NONZERO in 30/96, max 3.58e-08
|
|
163
|
+
scipy exactly 0.0 in all 96
|
|
164
|
+
```
|
|
165
|
+
|
|
166
|
+
-- and SciPy is at parity on that, so the guarantee is against the
|
|
167
|
+
constant-folded published arrangement, not against SciPy.
|
|
168
|
+
|
|
169
|
+
The SciPy spread tracks density, which is the support-intersection identity
|
|
170
|
+
doing the work: about 2.2x on dense input, above 30x at 1.6 % non-zero.
|
|
171
|
+
|
|
172
|
+
Treat those figures as one data point on one host, not a specification.
|
|
173
|
+
|
|
174
|
+
Every harness also runs on its own:
|
|
175
|
+
|
|
176
|
+
| script | question |
|
|
177
|
+
|---|---|
|
|
178
|
+
| `aa_control.py` | what is the smallest ratio this machine can resolve? |
|
|
179
|
+
| `bench_honest.py` | how does this compare to SciPy and to the published sparse form? |
|
|
180
|
+
| `lee_control.py` | how much of any margin is the arrangement, and how much the implementation? |
|
|
181
|
+
| `bits_vs_merge.py` | where is the crossover between the two sparse kernels? |
|
|
182
|
+
| `ratio_vs_n.py` | does the ratio to the baselines hold across `n`, or only at one `n`? |
|
|
183
|
+
| `accuracy.py` | true error against 50-digit arithmetic, by regime |
|
|
184
|
+
| `conditioning.py` | cancellation in each grouping of the summand |
|
|
185
|
+
| `lee_arrangement.py` | which arrangements return exactly zero for identical rows |
|
|
186
|
+
| `scale.py` | do the fitted exponents match the cost model? |
|
|
187
|
+
| `gpu_baseline.py` | would a GPU make this irrelevant? |
|
|
188
|
+
| `distributions.py` | do the synthetic inputs look like the real ones where it matters? |
|
|
189
|
+
| `fixture_shape.py` | how much of the fixtures' shape does the thresholding step add? |
|
|
190
|
+
|
|
191
|
+
`accuracy.py`, `conditioning.py`, `lee_arrangement.py`, `distributions.py` and
|
|
192
|
+
`fixture_shape.py` start no timer -- they report errors and counts -- so they
|
|
193
|
+
are the ones that are safe to run on a busy machine. The rest are timing
|
|
194
|
+
harnesses and want an idle one.
|
|
195
|
+
|
|
196
|
+
`bench_honest.py` times whatever else is installed alongside the two main
|
|
197
|
+
baselines, so it also needs `pynndescent` and `torch`; `bench/portable.py`,
|
|
198
|
+
described below, is the variant that records a missing implementation instead
|
|
199
|
+
of failing on it. `gpu_baseline.py` needs `torch` and should be run on its own,
|
|
200
|
+
since it loads the CPU too.
|
|
201
|
+
|
|
202
|
+
### Across many machines
|
|
203
|
+
|
|
204
|
+
A single machine cannot settle a performance claim: the crossover between the
|
|
205
|
+
two sparse kernels is sensitive to the branch predictor, so two hosts can
|
|
206
|
+
disagree about which is faster on identical input. Two more programs cover
|
|
207
|
+
that. `bench/portable.py` runs one *leg* -- one box, one self-describing JSON
|
|
208
|
+
document, every implementation probed before it is timed and recorded as
|
|
209
|
+
absent rather than fatal when it cannot run. `bench/aggregate.py` combines
|
|
210
|
+
legs into a claim and refuses to combine them wrongly; in particular it never
|
|
211
|
+
divides a time measured on one box by a time measured on another.
|
|
212
|
+
|
|
213
|
+
The `cross-machine benchmark (manual)` workflow
|
|
214
|
+
(`.github/workflows/cross-machine.yml`) runs one leg per runner across Linux
|
|
215
|
+
x86-64 (three draws from the same pool, for a variance estimate), older glibc,
|
|
216
|
+
aarch64, Windows and macOS, then aggregates them. It is `workflow_dispatch`
|
|
217
|
+
only. What it establishes is the *direction* of each claim across genuinely
|
|
218
|
+
different silicon, libm and toolchains, and the exactness guarantees, which
|
|
219
|
+
are deterministic and travel unchanged. `bench/CROSS_MACHINE.md` is the
|
|
220
|
+
operator's guide, and says which figures may be quoted and which may not.
|
|
221
|
+
|
|
222
|
+
## Backends
|
|
223
|
+
|
|
224
|
+
`backend="auto"` picks by input density and available extensions. Force one
|
|
225
|
+
with `backend=`:
|
|
226
|
+
|
|
227
|
+
- `numpy` and `numpy-sparse`: no compiled dependency.
|
|
228
|
+
- `numba` and `numba-sparse`: parallel dense and intersection kernels.
|
|
229
|
+
- `numba-svml`: an 8-wide vectorised dense kernel when Intel SVML is present.
|
|
230
|
+
|
|
231
|
+
`available_backends()` lists what this installation can use.
|
|
232
|
+
|
|
233
|
+
## API
|
|
234
|
+
|
|
235
|
+
```python
|
|
236
|
+
jsd_matrix(P, Q=None, *, base=2.0, squared=False, normalize=True,
|
|
237
|
+
check=True, dtype=None, backend="auto", out=None,
|
|
238
|
+
threads=None, well_conditioned=None) # (n, n) or (n, m)
|
|
239
|
+
|
|
240
|
+
jsd_pdist(P, ...) # condensed upper triangle, SciPy pdist order
|
|
241
|
+
jsd_pairwise(P, Q, ...) # row-aligned distances between two stacks
|
|
242
|
+
jsd(p, q, ...) # scalar
|
|
243
|
+
to_similarity(D) # 1 - D, meaningful in base 2
|
|
244
|
+
```
|
|
245
|
+
|
|
246
|
+
`base=2.0` bounds the distance in `[0, 1]`, which is why it is the default and
|
|
247
|
+
why `to_similarity` is well defined. SciPy's compiled pairwise path is base
|
|
248
|
+
`e` only, under which the distance caps at `sqrt(ln 2)` and `1 - d` never
|
|
249
|
+
reaches zero for disjoint supports.
|
|
250
|
+
|
|
251
|
+
## Documentation
|
|
252
|
+
|
|
253
|
+
- [`docs/MATH.md`](https://github.com/briankeithn/fastjsd/blob/main/docs/MATH.md): the identity, its consequences, the
|
|
254
|
+
guarantees and their proofs.
|
|
255
|
+
- [`docs/PRIOR_ART.md`](https://github.com/briankeithn/fastjsd/blob/main/docs/PRIOR_ART.md): what other libraries do, and
|
|
256
|
+
where this one differs.
|
|
257
|
+
|
|
258
|
+
## Acknowledgement
|
|
259
|
+
|
|
260
|
+
Claude Code was used to assist the development of this code as part of a
|
|
261
|
+
bigger project that required sparse JSD computations at scale.
|
|
262
|
+
|
|
263
|
+
## References
|
|
264
|
+
|
|
265
|
+
- Lin, J. (1991). Divergence measures based on the Shannon entropy.
|
|
266
|
+
*IEEE Transactions on Information Theory* 37(1), 145-151.
|
|
267
|
+
- Lee, L. (1997). *Similarity-Based Approaches to Natural Language
|
|
268
|
+
Processing*. Harvard University Technical Report TR-11-97.
|
|
269
|
+
- Endres, D. M. and Schindelin, J. E. (2003). A new metric for probability
|
|
270
|
+
distributions. *IEEE Transactions on Information Theory* 49(7), 1858-1860.
|
|
271
|
+
- Österreicher, F. and Vajda, I. (2003). A new class of metric divergences on
|
|
272
|
+
probability spaces and its applicability in statistics. *Annals of the
|
|
273
|
+
Institute of Statistical Mathematics* 55(3), 639-653.
|
|
274
|
+
|
|
275
|
+
## License
|
|
276
|
+
|
|
277
|
+
MIT. See [LICENSE](https://github.com/briankeithn/fastjsd/blob/main/LICENSE).
|
fastjsd-0.1.0/README.md
ADDED
|
@@ -0,0 +1,240 @@
|
|
|
1
|
+
# fastjsd
|
|
2
|
+
|
|
3
|
+
Fast, exact, pairwise Jensen-Shannon distance matrices for Python.
|
|
4
|
+
|
|
5
|
+
`fastjsd` computes the full `n x n` matrix of Jensen-Shannon distances between
|
|
6
|
+
rows of a stack of discrete distributions. It is a drop-in replacement for
|
|
7
|
+
`scipy.spatial.distance.pdist(P, "jensenshannon")` that adds a base argument,
|
|
8
|
+
parallel kernels, and a sparse path that skips coordinates neither
|
|
9
|
+
distribution occupies.
|
|
10
|
+
|
|
11
|
+
```python
|
|
12
|
+
import numpy as np
|
|
13
|
+
from fastjsd import jsd_matrix
|
|
14
|
+
|
|
15
|
+
P = np.random.default_rng(0).dirichlet(np.full(32, 0.1), size=500)
|
|
16
|
+
D = jsd_matrix(P) # (500, 500), base 2, distances in [0, 1]
|
|
17
|
+
```
|
|
18
|
+
|
|
19
|
+
## Install
|
|
20
|
+
|
|
21
|
+
```bash
|
|
22
|
+
pip install fastjsd # NumPy backends only
|
|
23
|
+
pip install "fastjsd[numba]" # adds the parallel dense and sparse kernels
|
|
24
|
+
```
|
|
25
|
+
|
|
26
|
+
Python 3.10 or newer. NumPy is the only hard dependency; Numba is optional and
|
|
27
|
+
enables the compiled kernels. SciPy is needed only by the test suite and the
|
|
28
|
+
benchmark harnesses, and comes in with the `test` and `bench` extras.
|
|
29
|
+
|
|
30
|
+
## Why it is fast
|
|
31
|
+
|
|
32
|
+
For distributions `p, q` over `k` outcomes with `m = (p + q) / 2`,
|
|
33
|
+
|
|
34
|
+
```
|
|
35
|
+
JS(p, q) = (1/2) * sum_c [ p_c ln p_c + q_c ln q_c - 2 m_c ln m_c ]
|
|
36
|
+
```
|
|
37
|
+
|
|
38
|
+
Two properties of that arrangement do the work.
|
|
39
|
+
|
|
40
|
+
**One logarithm per pair-coordinate.** `p ln p` depends on a single row, so it
|
|
41
|
+
is tabulated once per entry and the inner loop evaluates exactly one
|
|
42
|
+
logarithm. SciPy's kernel groups the same quantity as `a ln(a/m) + b ln(b/m)`,
|
|
43
|
+
where `m` depends on the pair, so neither logarithm can be hoisted and it
|
|
44
|
+
pays two.
|
|
45
|
+
|
|
46
|
+
**Logarithms only on the support intersection.** Where exactly one side has
|
|
47
|
+
mass, the summand is exactly `a ln 2` (Lee, 1997). A pair therefore needs
|
|
48
|
+
logarithms only where both rows are non-zero, and everything else is a
|
|
49
|
+
constant times the mass outside the intersection. For rows of density `d` with
|
|
50
|
+
independent supports the expected intersection is `k d^2` rather than `k`.
|
|
51
|
+
This is exact, not an approximation, and it is what makes thresholded
|
|
52
|
+
membership and topic-mixture data cheap.
|
|
53
|
+
|
|
54
|
+
The sparse kernels add an occupancy bitmap so pairs whose supports cannot
|
|
55
|
+
intersect are skipped without walking either row.
|
|
56
|
+
|
|
57
|
+
## Exactness
|
|
58
|
+
|
|
59
|
+
The guarantees, with proofs and the measurements that support them, are in
|
|
60
|
+
[`docs/MATH.md`](https://github.com/briankeithn/fastjsd/blob/main/docs/MATH.md). In summary:
|
|
61
|
+
|
|
62
|
+
- Every per-coordinate summand is non-negative, so the sum over coordinates
|
|
63
|
+
has condition number exactly 1.
|
|
64
|
+
- Identical rows return exactly `0.0`, not float dust that `sqrt` amplifies.
|
|
65
|
+
This is a property of keeping the constant inside the sum; forms that fold
|
|
66
|
+
`ln 2` outside it lose the exact zero.
|
|
67
|
+
- Results agree with SciPy to within a few times `1e-13` across dense, sparse,
|
|
68
|
+
cross and condensed forms.
|
|
69
|
+
- The individual summand is ill-conditioned when `p` and `q` nearly coincide.
|
|
70
|
+
That is a property of the grouping, it is documented rather than hidden, and
|
|
71
|
+
`well_conditioned=True` selects a rearrangement that removes it.
|
|
72
|
+
|
|
73
|
+
## Benchmarks
|
|
74
|
+
|
|
75
|
+
`bench/` contains the harnesses, the baselines and two real membership
|
|
76
|
+
fixtures. Performance depends heavily on the machine, so this package ships
|
|
77
|
+
the instruments and one worked example rather than a headline number. The
|
|
78
|
+
harnesses live in the repository rather than in the wheel, so run them from a
|
|
79
|
+
checkout:
|
|
80
|
+
|
|
81
|
+
```bash
|
|
82
|
+
pip install -e ".[bench]"
|
|
83
|
+
python bench/run_all.py # fingerprints the machine, then measures
|
|
84
|
+
```
|
|
85
|
+
|
|
86
|
+
`run_all.py` records the CPU, core counts, library versions and thread
|
|
87
|
+
settings before timing anything, then runs the three timing harnesses that
|
|
88
|
+
have to be read together -- `aa_control.py`, `bits_vs_merge.py` and
|
|
89
|
+
`lee_control.py` -- in dependency order, into one fingerprinted results
|
|
90
|
+
directory. The first of them, `bench/aa_control.py`, duels the implementation
|
|
91
|
+
against *itself*, so the true ratio is 1.000 by construction and everything it
|
|
92
|
+
reports is the measurement noise floor. Read that number first: a speedup
|
|
93
|
+
smaller than the noise floor of the machine that produced it is not a result.
|
|
94
|
+
|
|
95
|
+
### One machine's numbers
|
|
96
|
+
|
|
97
|
+
[`docs/measurements/`](https://github.com/briankeithn/fastjsd/blob/main/docs/measurements/README.md) holds a complete run with
|
|
98
|
+
its fingerprint and noise floor attached. Single-threaded on a 4-core Xeon at
|
|
99
|
+
2.10 GHz, which compares kernels rather than thread counts since SciPy's
|
|
100
|
+
`pdist` is single-threaded regardless:
|
|
101
|
+
|
|
102
|
+
| against | synthetic grid | real fixtures |
|
|
103
|
+
|---|---|---|
|
|
104
|
+
| SciPy | 2.19x - 33.17x | 2.14x - 8.93x |
|
|
105
|
+
| Lee 1997, transcribed literally | 1.86x - 5.06x | 1.52x - 2.10x |
|
|
106
|
+
| Lee 1997, competently implemented | **9 of 14 cells win, 1.19x - 3.61x; 5 lose, up to 1.25x** | |
|
|
107
|
+
|
|
108
|
+
Read the third row, not the second. The two differ in what the baseline is
|
|
109
|
+
allowed to do, and the gap between them is large: the literal transcription
|
|
110
|
+
evaluates `a ln(a/s) + b ln(b/s)` as printed, paying two logarithms and two
|
|
111
|
+
divisions per matched coordinate, while the same identity factors as
|
|
112
|
+
`a ln a + b ln b - s ln s` and accepts the same tabulated `x ln x` this package
|
|
113
|
+
uses. Giving the baseline that tabulation and a 1984 occupancy prefilter --
|
|
114
|
+
neither of which has anything to do with where `ln 2` sits -- makes it 1.1x to
|
|
115
|
+
2.6x faster, and that difference is most of the apparent margin in row two.
|
|
116
|
+
|
|
117
|
+
So the honest comparison is the third row, and it includes five losses,
|
|
118
|
+
covering both real fixtures under the stricter cleanup. That is expected, as
|
|
119
|
+
the folded arrangement absorbs the one-sided mass into a constant and needs
|
|
120
|
+
no per-row matched-mass accumulators at all. What the unfolded arrangement
|
|
121
|
+
buys is the exact zero, not speed --
|
|
122
|
+
|
|
123
|
+
```
|
|
124
|
+
fastjsd exactly 0.0 in all 96 duplicate-row trials
|
|
125
|
+
Lee-1997 sparse NONZERO in 30/96, max 3.58e-08
|
|
126
|
+
scipy exactly 0.0 in all 96
|
|
127
|
+
```
|
|
128
|
+
|
|
129
|
+
-- and SciPy is at parity on that, so the guarantee is against the
|
|
130
|
+
constant-folded published arrangement, not against SciPy.
|
|
131
|
+
|
|
132
|
+
The SciPy spread tracks density, which is the support-intersection identity
|
|
133
|
+
doing the work: about 2.2x on dense input, above 30x at 1.6 % non-zero.
|
|
134
|
+
|
|
135
|
+
Treat those figures as one data point on one host, not a specification.
|
|
136
|
+
|
|
137
|
+
Every harness also runs on its own:
|
|
138
|
+
|
|
139
|
+
| script | question |
|
|
140
|
+
|---|---|
|
|
141
|
+
| `aa_control.py` | what is the smallest ratio this machine can resolve? |
|
|
142
|
+
| `bench_honest.py` | how does this compare to SciPy and to the published sparse form? |
|
|
143
|
+
| `lee_control.py` | how much of any margin is the arrangement, and how much the implementation? |
|
|
144
|
+
| `bits_vs_merge.py` | where is the crossover between the two sparse kernels? |
|
|
145
|
+
| `ratio_vs_n.py` | does the ratio to the baselines hold across `n`, or only at one `n`? |
|
|
146
|
+
| `accuracy.py` | true error against 50-digit arithmetic, by regime |
|
|
147
|
+
| `conditioning.py` | cancellation in each grouping of the summand |
|
|
148
|
+
| `lee_arrangement.py` | which arrangements return exactly zero for identical rows |
|
|
149
|
+
| `scale.py` | do the fitted exponents match the cost model? |
|
|
150
|
+
| `gpu_baseline.py` | would a GPU make this irrelevant? |
|
|
151
|
+
| `distributions.py` | do the synthetic inputs look like the real ones where it matters? |
|
|
152
|
+
| `fixture_shape.py` | how much of the fixtures' shape does the thresholding step add? |
|
|
153
|
+
|
|
154
|
+
`accuracy.py`, `conditioning.py`, `lee_arrangement.py`, `distributions.py` and
|
|
155
|
+
`fixture_shape.py` start no timer -- they report errors and counts -- so they
|
|
156
|
+
are the ones that are safe to run on a busy machine. The rest are timing
|
|
157
|
+
harnesses and want an idle one.
|
|
158
|
+
|
|
159
|
+
`bench_honest.py` times whatever else is installed alongside the two main
|
|
160
|
+
baselines, so it also needs `pynndescent` and `torch`; `bench/portable.py`,
|
|
161
|
+
described below, is the variant that records a missing implementation instead
|
|
162
|
+
of failing on it. `gpu_baseline.py` needs `torch` and should be run on its own,
|
|
163
|
+
since it loads the CPU too.
|
|
164
|
+
|
|
165
|
+
### Across many machines
|
|
166
|
+
|
|
167
|
+
A single machine cannot settle a performance claim: the crossover between the
|
|
168
|
+
two sparse kernels is sensitive to the branch predictor, so two hosts can
|
|
169
|
+
disagree about which is faster on identical input. Two more programs cover
|
|
170
|
+
that. `bench/portable.py` runs one *leg* -- one box, one self-describing JSON
|
|
171
|
+
document, every implementation probed before it is timed and recorded as
|
|
172
|
+
absent rather than fatal when it cannot run. `bench/aggregate.py` combines
|
|
173
|
+
legs into a claim and refuses to combine them wrongly; in particular it never
|
|
174
|
+
divides a time measured on one box by a time measured on another.
|
|
175
|
+
|
|
176
|
+
The `cross-machine benchmark (manual)` workflow
|
|
177
|
+
(`.github/workflows/cross-machine.yml`) runs one leg per runner across Linux
|
|
178
|
+
x86-64 (three draws from the same pool, for a variance estimate), older glibc,
|
|
179
|
+
aarch64, Windows and macOS, then aggregates them. It is `workflow_dispatch`
|
|
180
|
+
only. What it establishes is the *direction* of each claim across genuinely
|
|
181
|
+
different silicon, libm and toolchains, and the exactness guarantees, which
|
|
182
|
+
are deterministic and travel unchanged. `bench/CROSS_MACHINE.md` is the
|
|
183
|
+
operator's guide, and says which figures may be quoted and which may not.
|
|
184
|
+
|
|
185
|
+
## Backends
|
|
186
|
+
|
|
187
|
+
`backend="auto"` picks by input density and available extensions. Force one
|
|
188
|
+
with `backend=`:
|
|
189
|
+
|
|
190
|
+
- `numpy` and `numpy-sparse`: no compiled dependency.
|
|
191
|
+
- `numba` and `numba-sparse`: parallel dense and intersection kernels.
|
|
192
|
+
- `numba-svml`: an 8-wide vectorised dense kernel when Intel SVML is present.
|
|
193
|
+
|
|
194
|
+
`available_backends()` lists what this installation can use.
|
|
195
|
+
|
|
196
|
+
## API
|
|
197
|
+
|
|
198
|
+
```python
|
|
199
|
+
jsd_matrix(P, Q=None, *, base=2.0, squared=False, normalize=True,
|
|
200
|
+
check=True, dtype=None, backend="auto", out=None,
|
|
201
|
+
threads=None, well_conditioned=None) # (n, n) or (n, m)
|
|
202
|
+
|
|
203
|
+
jsd_pdist(P, ...) # condensed upper triangle, SciPy pdist order
|
|
204
|
+
jsd_pairwise(P, Q, ...) # row-aligned distances between two stacks
|
|
205
|
+
jsd(p, q, ...) # scalar
|
|
206
|
+
to_similarity(D) # 1 - D, meaningful in base 2
|
|
207
|
+
```
|
|
208
|
+
|
|
209
|
+
`base=2.0` bounds the distance in `[0, 1]`, which is why it is the default and
|
|
210
|
+
why `to_similarity` is well defined. SciPy's compiled pairwise path is base
|
|
211
|
+
`e` only, under which the distance caps at `sqrt(ln 2)` and `1 - d` never
|
|
212
|
+
reaches zero for disjoint supports.
|
|
213
|
+
|
|
214
|
+
## Documentation
|
|
215
|
+
|
|
216
|
+
- [`docs/MATH.md`](https://github.com/briankeithn/fastjsd/blob/main/docs/MATH.md): the identity, its consequences, the
|
|
217
|
+
guarantees and their proofs.
|
|
218
|
+
- [`docs/PRIOR_ART.md`](https://github.com/briankeithn/fastjsd/blob/main/docs/PRIOR_ART.md): what other libraries do, and
|
|
219
|
+
where this one differs.
|
|
220
|
+
|
|
221
|
+
## Acknowledgement
|
|
222
|
+
|
|
223
|
+
Claude Code was used to assist the development of this code as part of a
|
|
224
|
+
bigger project that required sparse JSD computations at scale.
|
|
225
|
+
|
|
226
|
+
## References
|
|
227
|
+
|
|
228
|
+
- Lin, J. (1991). Divergence measures based on the Shannon entropy.
|
|
229
|
+
*IEEE Transactions on Information Theory* 37(1), 145-151.
|
|
230
|
+
- Lee, L. (1997). *Similarity-Based Approaches to Natural Language
|
|
231
|
+
Processing*. Harvard University Technical Report TR-11-97.
|
|
232
|
+
- Endres, D. M. and Schindelin, J. E. (2003). A new metric for probability
|
|
233
|
+
distributions. *IEEE Transactions on Information Theory* 49(7), 1858-1860.
|
|
234
|
+
- Österreicher, F. and Vajda, I. (2003). A new class of metric divergences on
|
|
235
|
+
probability spaces and its applicability in statistics. *Annals of the
|
|
236
|
+
Institute of Statistical Mathematics* 55(3), 639-653.
|
|
237
|
+
|
|
238
|
+
## License
|
|
239
|
+
|
|
240
|
+
MIT. See [LICENSE](https://github.com/briankeithn/fastjsd/blob/main/LICENSE).
|
|
@@ -0,0 +1,41 @@
|
|
|
1
|
+
"""fastjsd -- fast exact pairwise Jensen-Shannon distance.
|
|
2
|
+
|
|
3
|
+
>>> import numpy as np
|
|
4
|
+
>>> from fastjsd import jsd_matrix
|
|
5
|
+
>>> P = np.random.default_rng(0).dirichlet(np.full(32, 0.1), size=500)
|
|
6
|
+
>>> D = jsd_matrix(P) # (500, 500), base 2, distances in [0, 1]
|
|
7
|
+
|
|
8
|
+
Drop-in for ``scipy.spatial.distance.pdist(P, "jensenshannon")`` and
|
|
9
|
+
``cdist(P, Q, "jensenshannon")``, with the same values (base ``e``: pass
|
|
10
|
+
``base=np.e``). The harnesses under ``bench/`` in the source repository
|
|
11
|
+
measure how the two compare on a given machine.
|
|
12
|
+
|
|
13
|
+
Public API
|
|
14
|
+
----------
|
|
15
|
+
jsd_matrix (n, n) self or (n, m) cross distance matrix
|
|
16
|
+
jsd_pdist condensed upper triangle, SciPy ``pdist`` order
|
|
17
|
+
jsd_pairwise row-aligned distances between two stacks
|
|
18
|
+
jsd scalar distance between two distributions
|
|
19
|
+
to_similarity ``1 - d``, valid for base 2
|
|
20
|
+
"""
|
|
21
|
+
|
|
22
|
+
from ._api import (
|
|
23
|
+
SPARSE_DENSITY_MAX,
|
|
24
|
+
SVML_DENSITY_MAX,
|
|
25
|
+
available_backends,
|
|
26
|
+
jsd,
|
|
27
|
+
jsd_matrix,
|
|
28
|
+
jsd_pairwise,
|
|
29
|
+
jsd_pdist,
|
|
30
|
+
resolve_backend,
|
|
31
|
+
to_similarity,
|
|
32
|
+
)
|
|
33
|
+
|
|
34
|
+
__all__ = [
|
|
35
|
+
"jsd", "jsd_matrix", "jsd_pdist", "jsd_pairwise", "to_similarity",
|
|
36
|
+
"available_backends", "resolve_backend", "SPARSE_DENSITY_MAX",
|
|
37
|
+
"SVML_DENSITY_MAX",
|
|
38
|
+
"__version__",
|
|
39
|
+
]
|
|
40
|
+
|
|
41
|
+
__version__ = "0.1.0"
|