mojagg 0.0.2a0__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.
- mojagg-0.0.2a0/LICENSE +29 -0
- mojagg-0.0.2a0/PKG-INFO +233 -0
- mojagg-0.0.2a0/README.md +203 -0
- mojagg-0.0.2a0/pyproject.toml +70 -0
- mojagg-0.0.2a0/python/mojagg/__init__.py +66 -0
- mojagg-0.0.2a0/python/mojagg/_native.py +69 -0
- mojagg-0.0.2a0/python/mojagg/_reduce.py +218 -0
- mojagg-0.0.2a0/python/mojagg/compat.py +178 -0
- mojagg-0.0.2a0/python/mojagg/config.py +113 -0
- mojagg-0.0.2a0/python/mojagg/fill.py +113 -0
- mojagg-0.0.2a0/python/mojagg/matrix.py +92 -0
- mojagg-0.0.2a0/python/mojagg/nanfuncs.py +310 -0
- mojagg-0.0.2a0/python/mojagg.egg-info/PKG-INFO +233 -0
- mojagg-0.0.2a0/python/mojagg.egg-info/SOURCES.txt +17 -0
- mojagg-0.0.2a0/python/mojagg.egg-info/dependency_links.txt +1 -0
- mojagg-0.0.2a0/python/mojagg.egg-info/requires.txt +18 -0
- mojagg-0.0.2a0/python/mojagg.egg-info/top_level.txt +1 -0
- mojagg-0.0.2a0/setup.cfg +4 -0
- mojagg-0.0.2a0/setup.py +15 -0
mojagg-0.0.2a0/LICENSE
ADDED
|
@@ -0,0 +1,29 @@
|
|
|
1
|
+
BSD 3-Clause License
|
|
2
|
+
|
|
3
|
+
Copyright (c) 2026, mojagg contributors
|
|
4
|
+
All rights reserved.
|
|
5
|
+
|
|
6
|
+
Redistribution and use in source and binary forms, with or without
|
|
7
|
+
modification, are permitted provided that the following conditions are met:
|
|
8
|
+
|
|
9
|
+
* Redistributions of source code must retain the above copyright notice, this
|
|
10
|
+
list of conditions and the following disclaimer.
|
|
11
|
+
|
|
12
|
+
* Redistributions in binary form must reproduce the above copyright notice,
|
|
13
|
+
this list of conditions and the following disclaimer in the documentation
|
|
14
|
+
and/or other materials provided with the distribution.
|
|
15
|
+
|
|
16
|
+
* Neither the name of the copyright holder nor the names of its
|
|
17
|
+
contributors may be used to endorse or promote products derived from
|
|
18
|
+
this software without specific prior written permission.
|
|
19
|
+
|
|
20
|
+
THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
|
|
21
|
+
AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
|
|
22
|
+
IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
|
|
23
|
+
DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
|
|
24
|
+
FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
|
|
25
|
+
DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
|
|
26
|
+
SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
|
|
27
|
+
CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
|
|
28
|
+
OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
|
|
29
|
+
OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
|
mojagg-0.0.2a0/PKG-INFO
ADDED
|
@@ -0,0 +1,233 @@
|
|
|
1
|
+
Metadata-Version: 2.4
|
|
2
|
+
Name: mojagg
|
|
3
|
+
Version: 0.0.2a0
|
|
4
|
+
Summary: NaN-aware aggregations, grouped reductions, and rolling windows — numbagg's API, Mojo's speed.
|
|
5
|
+
Author: mojagg contributors
|
|
6
|
+
License-Expression: BSD-3-Clause
|
|
7
|
+
Project-URL: Homepage, https://github.com/josephnowak/mojagg
|
|
8
|
+
Project-URL: Repository, https://github.com/josephnowak/mojagg
|
|
9
|
+
Project-URL: Issues, https://github.com/josephnowak/mojagg/issues
|
|
10
|
+
Keywords: numpy,nan,groupby,rolling,mojo,simd,numbagg
|
|
11
|
+
Requires-Python: >=3.11
|
|
12
|
+
Description-Content-Type: text/markdown
|
|
13
|
+
License-File: LICENSE
|
|
14
|
+
Requires-Dist: numpy>=2.0
|
|
15
|
+
Provides-Extra: benchmark
|
|
16
|
+
Requires-Dist: numbagg>=0.3; extra == "benchmark"
|
|
17
|
+
Requires-Dist: numpy-groupies>=0.11; extra == "benchmark"
|
|
18
|
+
Requires-Dist: pytest>=8; extra == "benchmark"
|
|
19
|
+
Requires-Dist: pytest-codspeed>=3; extra == "benchmark"
|
|
20
|
+
Provides-Extra: test
|
|
21
|
+
Requires-Dist: numbagg>=0.3; extra == "test"
|
|
22
|
+
Requires-Dist: pandas>=2.0; extra == "test"
|
|
23
|
+
Requires-Dist: pytest>=8; extra == "test"
|
|
24
|
+
Provides-Extra: dev
|
|
25
|
+
Requires-Dist: ruff>=0.8; extra == "dev"
|
|
26
|
+
Requires-Dist: pytest>=8; extra == "dev"
|
|
27
|
+
Requires-Dist: numbagg>=0.3; extra == "dev"
|
|
28
|
+
Requires-Dist: pandas>=2.0; extra == "dev"
|
|
29
|
+
Dynamic: license-file
|
|
30
|
+
|
|
31
|
+
# mojagg
|
|
32
|
+
|
|
33
|
+
**NaN-aware aggregations, grouped reductions, and rolling windows — numbagg's API, Mojo's speed.**
|
|
34
|
+
|
|
35
|
+
```python
|
|
36
|
+
import mojagg
|
|
37
|
+
|
|
38
|
+
mojagg.nansum(a, axis=1)
|
|
39
|
+
mojagg.group_nanmean(values, labels, axis=0)
|
|
40
|
+
mojagg.move_mean(a, window=30, min_count=5)
|
|
41
|
+
```
|
|
42
|
+
|
|
43
|
+
If you know [numbagg](https://github.com/numbagg/numbagg), you already know mojagg: same functions, same signatures, same NaN semantics — reimplemented from scratch in [Mojo](https://www.modular.com/mojo) instead of numba, and pushed further.
|
|
44
|
+
|
|
45
|
+
## Why mojagg?
|
|
46
|
+
|
|
47
|
+
- **Drop-in**: 100% API- and semantics-compatible with numbagg. Change your import, keep your code.
|
|
48
|
+
- **Faster**: explicit SIMD (not compiler-hoped-for), cache-line-padded parallel reductions, branch-free NaN masking, and software-prefetched group-by scatter. No JIT warmup — kernels are AOT-compiled into the wheel.
|
|
49
|
+
- **Tunable**: every dispatch decision (parallel thresholds, worker counts, backend) is configurable per-call, globally, or by env var.
|
|
50
|
+
- **Honest**: reproducible comparisons against numbagg and NumPy, with development/WSL timings clearly separated from native-hardware calibration.
|
|
51
|
+
|
|
52
|
+
## Benchmarks
|
|
53
|
+
|
|
54
|
+
### Upstream-backed verification
|
|
55
|
+
|
|
56
|
+
The test corpus is vendored byte-for-byte from
|
|
57
|
+
[numbagg/numbagg@c73d4661b66cbfdee69e2834adb7f08d0d12af24](https://github.com/numbagg/numbagg/commit/c73d4661b66cbfdee69e2834adb7f08d0d12af24)
|
|
58
|
+
under `tests/vendor/numbagg`, with the upstream licenses and SHA-256 provenance.
|
|
59
|
+
`python scripts/vendor_numbagg_tests.py` refreshes the pinned snapshot;
|
|
60
|
+
`python scripts/vendor_numbagg_tests.py --check` verifies it offline.
|
|
61
|
+
|
|
62
|
+
Parity tests use the complete upstream array corpus, including million-element
|
|
63
|
+
cases, and execute the upstream allnan/anynan edge-case class unchanged through
|
|
64
|
+
differential wrappers. Supported operations must match numbagg's values,
|
|
65
|
+
shapes, dtypes and exception types; missing numbagg is an error, not a skip.
|
|
66
|
+
`nanprod` has no standalone numbagg equivalent and is explicitly NumPy-only.
|
|
67
|
+
|
|
68
|
+
`pixi run bench-reference --save comparison.json` compares public APIs using
|
|
69
|
+
the same upstream inputs, with correctness checks and JIT warmup before timing.
|
|
70
|
+
Add `--full` for million-element matrices, or `--ops nanmean nansum` to narrow
|
|
71
|
+
the run. Reports include fixture commit, runtime numbagg/NumPy versions, thread
|
|
72
|
+
configuration, individual samples and new/reference time ratios. Smaller ratios
|
|
73
|
+
are faster. Integer `nanprod` uses NumPy's explicit input dtype, matching this
|
|
74
|
+
extension's result-dtype contract.
|
|
75
|
+
|
|
76
|
+
Development verification (2026-09-06, WSL2, numbagg 0.9.4, 16 Numba threads):
|
|
77
|
+
153 Python tests passed. Both 10K- and 1M-element matrices completed: 208
|
|
78
|
+
numbagg comparisons plus 40 explicitly labeled NumPy-only nanprod comparisons.
|
|
79
|
+
Performance is not uniformly faster: the full matrix included nanmean at
|
|
80
|
+
1.43x and nancount at 1.59x numbagg's time in their slowest cases. These are
|
|
81
|
+
development measurements, not native-Linux dispatch calibration.
|
|
82
|
+
|
|
83
|
+
> Numbers below are placeholders until the first AWS calibration run (`c7i.8xlarge`, pinned CPU governor). The full matrix — size × cardinality × NaN-density × dtype — is regenerated per release and committed to `benchmarks/results/`.
|
|
84
|
+
|
|
85
|
+
| Function | numpy | numbagg | mojagg | vs numbagg |
|
|
86
|
+
|---|---|---|---|---|
|
|
87
|
+
| `nansum` (1e7 f64) | 1.0× | 8× | **TBD** | TBD |
|
|
88
|
+
| `group_nansum` (1e7 rows, 1e4 groups) | 1.0× (groupies) | 15× | **TBD** | TBD |
|
|
89
|
+
| `move_mean` (1e7 f64, w=100) | — | 20× | **TBD** | TBD |
|
|
90
|
+
|
|
91
|
+
Run them: `python benchmarks/full_matrix.py` · Continuous per-PR performance tracking via [CodSpeed](https://codspeed.io).
|
|
92
|
+
|
|
93
|
+
### Reduction development
|
|
94
|
+
|
|
95
|
+
`allnan`, `nansum` and `nanmean` implement `NaNReduction1D` in
|
|
96
|
+
`src/mojagg/core/reduce1d.mojo`. Shared scanners handle contiguous SIMD and
|
|
97
|
+
scalar strides; the axis driver merges partial states and finalizes once per
|
|
98
|
+
output slice. Native vector width and arithmetic accumulator-chain count are
|
|
99
|
+
separate. Sum retains its eight-chain EVL accumulator and scalar strided sum.
|
|
100
|
+
|
|
101
|
+
`allnan` checks `isnan(values).reduce_and()` after each native-width block.
|
|
102
|
+
Its terminal state also stops the multi-axis odometer, never other output
|
|
103
|
+
slices. Compile-time capabilities erase those checks for sum and bypass input
|
|
104
|
+
traversal entirely for integer `allnan`. Other reductions pass `HookReduction`
|
|
105
|
+
explicitly to the same `reduce_axis[Op: Reduction1D]` entry point. There is no
|
|
106
|
+
legacy four-hook overload. `HookReduction` describes a result-valued operation;
|
|
107
|
+
the axis driver still owns N-D planning, slice traversal and parallelism.
|
|
108
|
+
|
|
109
|
+
`nanmean` carries a float64 sum and int64 valid count through both scanners
|
|
110
|
+
and multi-axis merging; it divides only at finalization. Empty/all-NaN slices
|
|
111
|
+
return NaN. Float32 input remains zero-copy with a float32 result, but uses
|
|
112
|
+
float64 arithmetic internally to match numbagg (including finite values whose
|
|
113
|
+
float32 sum would overflow). The facade follows the reference's safe-casting
|
|
114
|
+
order: bool/small integers/float16 -> float32; 32/64-bit integers -> float64.
|
|
115
|
+
Unsupported nonnumeric dtypes are rejected instead of coerced.
|
|
116
|
+
|
|
117
|
+
Mean's width sweep is reproducible with
|
|
118
|
+
`pixi run mojo run -O3 -I src benchmarks/mean_widths.mojo`. The seven-trial,
|
|
119
|
+
alternating 1/2/4/8-chain comparison selected one chain for float32 and four
|
|
120
|
+
for float64. On the development WSL host, at 100,003 elements their median
|
|
121
|
+
kernel times were 41.5 us and 36.5 us respectively; float64's one-chain
|
|
122
|
+
baseline took 50.1 us. Four float64 chains trade roughly 8 ns on 17-element
|
|
123
|
+
runs for faster long scans. Recalibrate on native hardware before claiming
|
|
124
|
+
portable speedups. Python benchmark rows compare both dtypes against NumPy
|
|
125
|
+
and numbagg.
|
|
126
|
+
|
|
127
|
+
The local quick public-API comparison was mixed: full 1M-element means took
|
|
128
|
+
about 0.56/0.52 ms (f32/f64) versus numbagg's 6.27/7.88 ms, but float64 tiny
|
|
129
|
+
rows and column reductions remained slower than numbagg (2.44 vs 0.99 ms and
|
|
130
|
+
2.32 vs 1.32 ms respectively). These WSL measurements are not a claim that
|
|
131
|
+
every shape is faster.
|
|
132
|
+
|
|
133
|
+
Run `pixi run test-mojo` to check actual traversal counts, including parallel
|
|
134
|
+
output independence, and `pixi run test` for Python parity.
|
|
135
|
+
`benchmarks/reduction_contract.py` times the native boundary with allocations
|
|
136
|
+
and reference calculations excluded:
|
|
137
|
+
|
|
138
|
+
```bash
|
|
139
|
+
pixi run build-ext
|
|
140
|
+
PYTHONPATH=python pixi run python benchmarks/reduction_contract.py --save before.json
|
|
141
|
+
# Preserve a copy of the built extension before editing, then rebuild.
|
|
142
|
+
PYTHONPATH=python pixi run python benchmarks/reduction_contract.py --baseline-library before.so --save paired.json
|
|
143
|
+
```
|
|
144
|
+
|
|
145
|
+
The paired mode alternates old/new libraries on identical inputs, checks equal
|
|
146
|
+
results and reports new/baseline time ratios (below 1 means faster). It covers
|
|
147
|
+
f32/f64/i32/i64, SIMD tails, strided/multi-axis slices, all-NaN scans and early
|
|
148
|
+
or late decisive values. Use `--filter nansum --min-time 0.1` for longer sum
|
|
149
|
+
checks. WSL timings are development comparisons, not native-Linux dispatch
|
|
150
|
+
calibration; no dispatch thresholds are changed by this refactor.
|
|
151
|
+
|
|
152
|
+
Development comparison (2026-09-06, WSL2 x86_64, Mojo 1.0.0): nine alternating
|
|
153
|
+
old/new samples per case. Longer sum runs used at least 100 ms per sample.
|
|
154
|
+
These are time ratios against the pre-refactor working-tree extension, not
|
|
155
|
+
speedups against NumPy/numbagg:
|
|
156
|
+
|
|
157
|
+
| Workload | New / baseline time |
|
|
158
|
+
|---|---|
|
|
159
|
+
| Sum, 28 dtype/layout/size cases | 0.884–1.049 |
|
|
160
|
+
| Float allnan, first value decisive, contiguous | 0.117–0.199 |
|
|
161
|
+
| Float allnan, first block decisive in every MULTI slice | 0.104–0.143 |
|
|
162
|
+
| Float allnan, first value decisive in every strided MULTI slice | 0.039–0.045 |
|
|
163
|
+
| Integer allnan, MULTI layouts | 0.235–0.391 |
|
|
164
|
+
|
|
165
|
+
All measured sum cases stayed within a 5% regression tolerance; the float32
|
|
166
|
+
MULTI case improved by about 12%. Recheck on native Linux before treating
|
|
167
|
+
these development timings as portable performance claims.
|
|
168
|
+
|
|
169
|
+
## Install
|
|
170
|
+
|
|
171
|
+
```bash
|
|
172
|
+
pip install mojagg
|
|
173
|
+
```
|
|
174
|
+
|
|
175
|
+
Prebuilt wheels for Linux x86_64/aarch64 and macOS arm64. Python ≥ 3.11, NumPy ≥ 2.0. No Mojo toolchain needed — kernels ship compiled.
|
|
176
|
+
|
|
177
|
+
## Functions
|
|
178
|
+
|
|
179
|
+
| Family | Functions |
|
|
180
|
+
|---|---|
|
|
181
|
+
| Aggregations | `nansum nanmean nanstd nanvar nanmin nanmax nancount nanargmin nanargmax nanmedian nanquantile allnan anynan count` |
|
|
182
|
+
| Grouped | `group_nansum group_nanmean group_nanprod group_nanvar group_nanstd group_nancount group_nanmin group_nanmax group_nanargmin group_nanargmax group_nanfirst group_nanlast group_nanany group_nanall group_nansum_of_squares` |
|
|
183
|
+
| Rolling | `move_sum move_mean move_std move_var move_cov move_corr` |
|
|
184
|
+
| Exp-weighted | `move_exp_nansum move_exp_nanmean move_exp_nancount move_exp_nanvar move_exp_nanstd move_exp_nancov move_exp_nancorr` |
|
|
185
|
+
| Matrix | `nancovmatrix nancorrmatrix move_covmatrix move_corrmatrix move_exp_nancovmatrix move_exp_nancorrmatrix` |
|
|
186
|
+
| Fill | `ffill bfill` |
|
|
187
|
+
|
|
188
|
+
`nanvar` and `nanstd` use numbagg's default `ddof=1` and accept an integer
|
|
189
|
+
`ddof` keyword. `nanquantile`/`nanmedian` use a non-streaming selection path
|
|
190
|
+
with NumPy-compatible linear interpolation; scalar quantiles are returned as
|
|
191
|
+
scalars and vector quantiles occupy the leading axis.
|
|
192
|
+
|
|
193
|
+
Native reduction kernels instantiate `float64`/`float32`/`int64`/`int32`;
|
|
194
|
+
the facade visibly promotes numbagg-compatible small and integer inputs where
|
|
195
|
+
required. Reduction inputs remain zero-copy, except for documented promotion
|
|
196
|
+
and big-endian normalization.
|
|
197
|
+
|
|
198
|
+
## Configuration
|
|
199
|
+
|
|
200
|
+
```python
|
|
201
|
+
with mojagg.config(parallel_threshold=50_000, threads=8):
|
|
202
|
+
mojagg.group_nansum(values, labels)
|
|
203
|
+
|
|
204
|
+
mojagg.set_config(backend="cpu") # global
|
|
205
|
+
# or env: MOJAGG_PARALLEL_THRESHOLD=50000 MOJAGG_THREADS=8
|
|
206
|
+
```
|
|
207
|
+
|
|
208
|
+
Context manager > global > env var > tuned defaults (benchmark-derived).
|
|
209
|
+
|
|
210
|
+
## Philosophy
|
|
211
|
+
|
|
212
|
+
1. **Parity before speed.** A result that doesn't match numbagg is a bug, however fast. The test suite *is* the spec.
|
|
213
|
+
2. **No hidden work.** No JIT warmup, no silent casts, no secret copies. If mojagg can't go fast on your data as-is, it tells you.
|
|
214
|
+
3. **Measure or revert.** Performance changes land only with benchmark evidence.
|
|
215
|
+
4. **Design for the next hardware.** Kernels are written against a backend abstraction; GPU targets slot in without API changes.
|
|
216
|
+
|
|
217
|
+
## Built with AI, built for AI
|
|
218
|
+
|
|
219
|
+
mojagg is designed and maintained with AI agents as first-class contributors — and first-class users:
|
|
220
|
+
|
|
221
|
+
- **Machine-readable design docs** (`AGENTS.md`, `skills/`) encode the architecture, parity semantics, and performance rules, so agent-generated contributions are consistent by construction.
|
|
222
|
+
- **Self-verifying**: parity tests + benchmark gates give agents (and humans) objective acceptance criteria for every change.
|
|
223
|
+
- **Agent-friendly API**: predictable naming, explicit errors, structured config — easy for codegen tools to call correctly.
|
|
224
|
+
|
|
225
|
+
Contributions from humans and agents alike are welcome. See `AGENTS.md`.
|
|
226
|
+
|
|
227
|
+
## License
|
|
228
|
+
|
|
229
|
+
BSD 3-Clause — same as numbagg. mojagg is and will remain 100% free and open source.
|
|
230
|
+
|
|
231
|
+
## Acknowledgments
|
|
232
|
+
|
|
233
|
+
Inspired by and API-compatible with [numbagg](https://github.com/numbagg/numbagg) (BSD-3). Group-label conventions follow [numpy-groupies](https://github.com/ml31415/numpy-groupies).
|
mojagg-0.0.2a0/README.md
ADDED
|
@@ -0,0 +1,203 @@
|
|
|
1
|
+
# mojagg
|
|
2
|
+
|
|
3
|
+
**NaN-aware aggregations, grouped reductions, and rolling windows — numbagg's API, Mojo's speed.**
|
|
4
|
+
|
|
5
|
+
```python
|
|
6
|
+
import mojagg
|
|
7
|
+
|
|
8
|
+
mojagg.nansum(a, axis=1)
|
|
9
|
+
mojagg.group_nanmean(values, labels, axis=0)
|
|
10
|
+
mojagg.move_mean(a, window=30, min_count=5)
|
|
11
|
+
```
|
|
12
|
+
|
|
13
|
+
If you know [numbagg](https://github.com/numbagg/numbagg), you already know mojagg: same functions, same signatures, same NaN semantics — reimplemented from scratch in [Mojo](https://www.modular.com/mojo) instead of numba, and pushed further.
|
|
14
|
+
|
|
15
|
+
## Why mojagg?
|
|
16
|
+
|
|
17
|
+
- **Drop-in**: 100% API- and semantics-compatible with numbagg. Change your import, keep your code.
|
|
18
|
+
- **Faster**: explicit SIMD (not compiler-hoped-for), cache-line-padded parallel reductions, branch-free NaN masking, and software-prefetched group-by scatter. No JIT warmup — kernels are AOT-compiled into the wheel.
|
|
19
|
+
- **Tunable**: every dispatch decision (parallel thresholds, worker counts, backend) is configurable per-call, globally, or by env var.
|
|
20
|
+
- **Honest**: reproducible comparisons against numbagg and NumPy, with development/WSL timings clearly separated from native-hardware calibration.
|
|
21
|
+
|
|
22
|
+
## Benchmarks
|
|
23
|
+
|
|
24
|
+
### Upstream-backed verification
|
|
25
|
+
|
|
26
|
+
The test corpus is vendored byte-for-byte from
|
|
27
|
+
[numbagg/numbagg@c73d4661b66cbfdee69e2834adb7f08d0d12af24](https://github.com/numbagg/numbagg/commit/c73d4661b66cbfdee69e2834adb7f08d0d12af24)
|
|
28
|
+
under `tests/vendor/numbagg`, with the upstream licenses and SHA-256 provenance.
|
|
29
|
+
`python scripts/vendor_numbagg_tests.py` refreshes the pinned snapshot;
|
|
30
|
+
`python scripts/vendor_numbagg_tests.py --check` verifies it offline.
|
|
31
|
+
|
|
32
|
+
Parity tests use the complete upstream array corpus, including million-element
|
|
33
|
+
cases, and execute the upstream allnan/anynan edge-case class unchanged through
|
|
34
|
+
differential wrappers. Supported operations must match numbagg's values,
|
|
35
|
+
shapes, dtypes and exception types; missing numbagg is an error, not a skip.
|
|
36
|
+
`nanprod` has no standalone numbagg equivalent and is explicitly NumPy-only.
|
|
37
|
+
|
|
38
|
+
`pixi run bench-reference --save comparison.json` compares public APIs using
|
|
39
|
+
the same upstream inputs, with correctness checks and JIT warmup before timing.
|
|
40
|
+
Add `--full` for million-element matrices, or `--ops nanmean nansum` to narrow
|
|
41
|
+
the run. Reports include fixture commit, runtime numbagg/NumPy versions, thread
|
|
42
|
+
configuration, individual samples and new/reference time ratios. Smaller ratios
|
|
43
|
+
are faster. Integer `nanprod` uses NumPy's explicit input dtype, matching this
|
|
44
|
+
extension's result-dtype contract.
|
|
45
|
+
|
|
46
|
+
Development verification (2026-09-06, WSL2, numbagg 0.9.4, 16 Numba threads):
|
|
47
|
+
153 Python tests passed. Both 10K- and 1M-element matrices completed: 208
|
|
48
|
+
numbagg comparisons plus 40 explicitly labeled NumPy-only nanprod comparisons.
|
|
49
|
+
Performance is not uniformly faster: the full matrix included nanmean at
|
|
50
|
+
1.43x and nancount at 1.59x numbagg's time in their slowest cases. These are
|
|
51
|
+
development measurements, not native-Linux dispatch calibration.
|
|
52
|
+
|
|
53
|
+
> Numbers below are placeholders until the first AWS calibration run (`c7i.8xlarge`, pinned CPU governor). The full matrix — size × cardinality × NaN-density × dtype — is regenerated per release and committed to `benchmarks/results/`.
|
|
54
|
+
|
|
55
|
+
| Function | numpy | numbagg | mojagg | vs numbagg |
|
|
56
|
+
|---|---|---|---|---|
|
|
57
|
+
| `nansum` (1e7 f64) | 1.0× | 8× | **TBD** | TBD |
|
|
58
|
+
| `group_nansum` (1e7 rows, 1e4 groups) | 1.0× (groupies) | 15× | **TBD** | TBD |
|
|
59
|
+
| `move_mean` (1e7 f64, w=100) | — | 20× | **TBD** | TBD |
|
|
60
|
+
|
|
61
|
+
Run them: `python benchmarks/full_matrix.py` · Continuous per-PR performance tracking via [CodSpeed](https://codspeed.io).
|
|
62
|
+
|
|
63
|
+
### Reduction development
|
|
64
|
+
|
|
65
|
+
`allnan`, `nansum` and `nanmean` implement `NaNReduction1D` in
|
|
66
|
+
`src/mojagg/core/reduce1d.mojo`. Shared scanners handle contiguous SIMD and
|
|
67
|
+
scalar strides; the axis driver merges partial states and finalizes once per
|
|
68
|
+
output slice. Native vector width and arithmetic accumulator-chain count are
|
|
69
|
+
separate. Sum retains its eight-chain EVL accumulator and scalar strided sum.
|
|
70
|
+
|
|
71
|
+
`allnan` checks `isnan(values).reduce_and()` after each native-width block.
|
|
72
|
+
Its terminal state also stops the multi-axis odometer, never other output
|
|
73
|
+
slices. Compile-time capabilities erase those checks for sum and bypass input
|
|
74
|
+
traversal entirely for integer `allnan`. Other reductions pass `HookReduction`
|
|
75
|
+
explicitly to the same `reduce_axis[Op: Reduction1D]` entry point. There is no
|
|
76
|
+
legacy four-hook overload. `HookReduction` describes a result-valued operation;
|
|
77
|
+
the axis driver still owns N-D planning, slice traversal and parallelism.
|
|
78
|
+
|
|
79
|
+
`nanmean` carries a float64 sum and int64 valid count through both scanners
|
|
80
|
+
and multi-axis merging; it divides only at finalization. Empty/all-NaN slices
|
|
81
|
+
return NaN. Float32 input remains zero-copy with a float32 result, but uses
|
|
82
|
+
float64 arithmetic internally to match numbagg (including finite values whose
|
|
83
|
+
float32 sum would overflow). The facade follows the reference's safe-casting
|
|
84
|
+
order: bool/small integers/float16 -> float32; 32/64-bit integers -> float64.
|
|
85
|
+
Unsupported nonnumeric dtypes are rejected instead of coerced.
|
|
86
|
+
|
|
87
|
+
Mean's width sweep is reproducible with
|
|
88
|
+
`pixi run mojo run -O3 -I src benchmarks/mean_widths.mojo`. The seven-trial,
|
|
89
|
+
alternating 1/2/4/8-chain comparison selected one chain for float32 and four
|
|
90
|
+
for float64. On the development WSL host, at 100,003 elements their median
|
|
91
|
+
kernel times were 41.5 us and 36.5 us respectively; float64's one-chain
|
|
92
|
+
baseline took 50.1 us. Four float64 chains trade roughly 8 ns on 17-element
|
|
93
|
+
runs for faster long scans. Recalibrate on native hardware before claiming
|
|
94
|
+
portable speedups. Python benchmark rows compare both dtypes against NumPy
|
|
95
|
+
and numbagg.
|
|
96
|
+
|
|
97
|
+
The local quick public-API comparison was mixed: full 1M-element means took
|
|
98
|
+
about 0.56/0.52 ms (f32/f64) versus numbagg's 6.27/7.88 ms, but float64 tiny
|
|
99
|
+
rows and column reductions remained slower than numbagg (2.44 vs 0.99 ms and
|
|
100
|
+
2.32 vs 1.32 ms respectively). These WSL measurements are not a claim that
|
|
101
|
+
every shape is faster.
|
|
102
|
+
|
|
103
|
+
Run `pixi run test-mojo` to check actual traversal counts, including parallel
|
|
104
|
+
output independence, and `pixi run test` for Python parity.
|
|
105
|
+
`benchmarks/reduction_contract.py` times the native boundary with allocations
|
|
106
|
+
and reference calculations excluded:
|
|
107
|
+
|
|
108
|
+
```bash
|
|
109
|
+
pixi run build-ext
|
|
110
|
+
PYTHONPATH=python pixi run python benchmarks/reduction_contract.py --save before.json
|
|
111
|
+
# Preserve a copy of the built extension before editing, then rebuild.
|
|
112
|
+
PYTHONPATH=python pixi run python benchmarks/reduction_contract.py --baseline-library before.so --save paired.json
|
|
113
|
+
```
|
|
114
|
+
|
|
115
|
+
The paired mode alternates old/new libraries on identical inputs, checks equal
|
|
116
|
+
results and reports new/baseline time ratios (below 1 means faster). It covers
|
|
117
|
+
f32/f64/i32/i64, SIMD tails, strided/multi-axis slices, all-NaN scans and early
|
|
118
|
+
or late decisive values. Use `--filter nansum --min-time 0.1` for longer sum
|
|
119
|
+
checks. WSL timings are development comparisons, not native-Linux dispatch
|
|
120
|
+
calibration; no dispatch thresholds are changed by this refactor.
|
|
121
|
+
|
|
122
|
+
Development comparison (2026-09-06, WSL2 x86_64, Mojo 1.0.0): nine alternating
|
|
123
|
+
old/new samples per case. Longer sum runs used at least 100 ms per sample.
|
|
124
|
+
These are time ratios against the pre-refactor working-tree extension, not
|
|
125
|
+
speedups against NumPy/numbagg:
|
|
126
|
+
|
|
127
|
+
| Workload | New / baseline time |
|
|
128
|
+
|---|---|
|
|
129
|
+
| Sum, 28 dtype/layout/size cases | 0.884–1.049 |
|
|
130
|
+
| Float allnan, first value decisive, contiguous | 0.117–0.199 |
|
|
131
|
+
| Float allnan, first block decisive in every MULTI slice | 0.104–0.143 |
|
|
132
|
+
| Float allnan, first value decisive in every strided MULTI slice | 0.039–0.045 |
|
|
133
|
+
| Integer allnan, MULTI layouts | 0.235–0.391 |
|
|
134
|
+
|
|
135
|
+
All measured sum cases stayed within a 5% regression tolerance; the float32
|
|
136
|
+
MULTI case improved by about 12%. Recheck on native Linux before treating
|
|
137
|
+
these development timings as portable performance claims.
|
|
138
|
+
|
|
139
|
+
## Install
|
|
140
|
+
|
|
141
|
+
```bash
|
|
142
|
+
pip install mojagg
|
|
143
|
+
```
|
|
144
|
+
|
|
145
|
+
Prebuilt wheels for Linux x86_64/aarch64 and macOS arm64. Python ≥ 3.11, NumPy ≥ 2.0. No Mojo toolchain needed — kernels ship compiled.
|
|
146
|
+
|
|
147
|
+
## Functions
|
|
148
|
+
|
|
149
|
+
| Family | Functions |
|
|
150
|
+
|---|---|
|
|
151
|
+
| Aggregations | `nansum nanmean nanstd nanvar nanmin nanmax nancount nanargmin nanargmax nanmedian nanquantile allnan anynan count` |
|
|
152
|
+
| Grouped | `group_nansum group_nanmean group_nanprod group_nanvar group_nanstd group_nancount group_nanmin group_nanmax group_nanargmin group_nanargmax group_nanfirst group_nanlast group_nanany group_nanall group_nansum_of_squares` |
|
|
153
|
+
| Rolling | `move_sum move_mean move_std move_var move_cov move_corr` |
|
|
154
|
+
| Exp-weighted | `move_exp_nansum move_exp_nanmean move_exp_nancount move_exp_nanvar move_exp_nanstd move_exp_nancov move_exp_nancorr` |
|
|
155
|
+
| Matrix | `nancovmatrix nancorrmatrix move_covmatrix move_corrmatrix move_exp_nancovmatrix move_exp_nancorrmatrix` |
|
|
156
|
+
| Fill | `ffill bfill` |
|
|
157
|
+
|
|
158
|
+
`nanvar` and `nanstd` use numbagg's default `ddof=1` and accept an integer
|
|
159
|
+
`ddof` keyword. `nanquantile`/`nanmedian` use a non-streaming selection path
|
|
160
|
+
with NumPy-compatible linear interpolation; scalar quantiles are returned as
|
|
161
|
+
scalars and vector quantiles occupy the leading axis.
|
|
162
|
+
|
|
163
|
+
Native reduction kernels instantiate `float64`/`float32`/`int64`/`int32`;
|
|
164
|
+
the facade visibly promotes numbagg-compatible small and integer inputs where
|
|
165
|
+
required. Reduction inputs remain zero-copy, except for documented promotion
|
|
166
|
+
and big-endian normalization.
|
|
167
|
+
|
|
168
|
+
## Configuration
|
|
169
|
+
|
|
170
|
+
```python
|
|
171
|
+
with mojagg.config(parallel_threshold=50_000, threads=8):
|
|
172
|
+
mojagg.group_nansum(values, labels)
|
|
173
|
+
|
|
174
|
+
mojagg.set_config(backend="cpu") # global
|
|
175
|
+
# or env: MOJAGG_PARALLEL_THRESHOLD=50000 MOJAGG_THREADS=8
|
|
176
|
+
```
|
|
177
|
+
|
|
178
|
+
Context manager > global > env var > tuned defaults (benchmark-derived).
|
|
179
|
+
|
|
180
|
+
## Philosophy
|
|
181
|
+
|
|
182
|
+
1. **Parity before speed.** A result that doesn't match numbagg is a bug, however fast. The test suite *is* the spec.
|
|
183
|
+
2. **No hidden work.** No JIT warmup, no silent casts, no secret copies. If mojagg can't go fast on your data as-is, it tells you.
|
|
184
|
+
3. **Measure or revert.** Performance changes land only with benchmark evidence.
|
|
185
|
+
4. **Design for the next hardware.** Kernels are written against a backend abstraction; GPU targets slot in without API changes.
|
|
186
|
+
|
|
187
|
+
## Built with AI, built for AI
|
|
188
|
+
|
|
189
|
+
mojagg is designed and maintained with AI agents as first-class contributors — and first-class users:
|
|
190
|
+
|
|
191
|
+
- **Machine-readable design docs** (`AGENTS.md`, `skills/`) encode the architecture, parity semantics, and performance rules, so agent-generated contributions are consistent by construction.
|
|
192
|
+
- **Self-verifying**: parity tests + benchmark gates give agents (and humans) objective acceptance criteria for every change.
|
|
193
|
+
- **Agent-friendly API**: predictable naming, explicit errors, structured config — easy for codegen tools to call correctly.
|
|
194
|
+
|
|
195
|
+
Contributions from humans and agents alike are welcome. See `AGENTS.md`.
|
|
196
|
+
|
|
197
|
+
## License
|
|
198
|
+
|
|
199
|
+
BSD 3-Clause — same as numbagg. mojagg is and will remain 100% free and open source.
|
|
200
|
+
|
|
201
|
+
## Acknowledgments
|
|
202
|
+
|
|
203
|
+
Inspired by and API-compatible with [numbagg](https://github.com/numbagg/numbagg) (BSD-3). Group-label conventions follow [numpy-groupies](https://github.com/ml31415/numpy-groupies).
|
|
@@ -0,0 +1,70 @@
|
|
|
1
|
+
[build-system]
|
|
2
|
+
requires = ["setuptools>=69", "wheel"]
|
|
3
|
+
build-backend = "setuptools.build_meta"
|
|
4
|
+
|
|
5
|
+
[project]
|
|
6
|
+
name = "mojagg"
|
|
7
|
+
version = "0.0.2-alpha"
|
|
8
|
+
description = "NaN-aware aggregations, grouped reductions, and rolling windows — numbagg's API, Mojo's speed."
|
|
9
|
+
readme = "README.md"
|
|
10
|
+
license = "BSD-3-Clause"
|
|
11
|
+
authors = [{ name = "mojagg contributors" }]
|
|
12
|
+
keywords = ["numpy", "nan", "groupby", "rolling", "mojo", "simd", "numbagg"]
|
|
13
|
+
requires-python = ">=3.11"
|
|
14
|
+
dependencies = [
|
|
15
|
+
"numpy>=2.0",
|
|
16
|
+
]
|
|
17
|
+
|
|
18
|
+
[project.optional-dependencies]
|
|
19
|
+
# Parity testing + benchmarking against the reference implementations.
|
|
20
|
+
benchmark = [
|
|
21
|
+
"numbagg>=0.3",
|
|
22
|
+
"numpy-groupies>=0.11",
|
|
23
|
+
"pytest>=8",
|
|
24
|
+
"pytest-codspeed>=3",
|
|
25
|
+
]
|
|
26
|
+
test = [
|
|
27
|
+
"numbagg>=0.3",
|
|
28
|
+
"pandas>=2.0",
|
|
29
|
+
"pytest>=8",
|
|
30
|
+
]
|
|
31
|
+
dev = [
|
|
32
|
+
"ruff>=0.8",
|
|
33
|
+
"pytest>=8",
|
|
34
|
+
"numbagg>=0.3",
|
|
35
|
+
"pandas>=2.0",
|
|
36
|
+
]
|
|
37
|
+
|
|
38
|
+
[project.urls]
|
|
39
|
+
Homepage = "https://github.com/josephnowak/mojagg"
|
|
40
|
+
Repository = "https://github.com/josephnowak/mojagg"
|
|
41
|
+
Issues = "https://github.com/josephnowak/mojagg/issues"
|
|
42
|
+
|
|
43
|
+
[tool.setuptools]
|
|
44
|
+
include-package-data = true
|
|
45
|
+
|
|
46
|
+
[tool.setuptools.package-data]
|
|
47
|
+
mojagg = ["*.so", "*.pyd"]
|
|
48
|
+
|
|
49
|
+
[tool.setuptools.packages.find]
|
|
50
|
+
where = ["python"]
|
|
51
|
+
include = ["mojagg*"]
|
|
52
|
+
|
|
53
|
+
# Native Mojo extension is built out-of-band by `pixi run build-ext` and
|
|
54
|
+
# packaged as package data; see scripts/build_ext.py.
|
|
55
|
+
|
|
56
|
+
[tool.ruff]
|
|
57
|
+
target-version = "py311"
|
|
58
|
+
line-length = 100
|
|
59
|
+
src = ["python", "tests", "benchmarks", "scripts"]
|
|
60
|
+
extend-exclude = ["tests/vendor", ".analysis"]
|
|
61
|
+
|
|
62
|
+
[tool.ruff.lint]
|
|
63
|
+
select = ["E", "F", "I", "UP", "B", "SIM", "C4", "PTH"]
|
|
64
|
+
|
|
65
|
+
[tool.pytest.ini_options]
|
|
66
|
+
pythonpath = ["python", "."]
|
|
67
|
+
markers = [
|
|
68
|
+
"gpu: tests requiring a GPU device (skipped by default)",
|
|
69
|
+
]
|
|
70
|
+
addopts = "-m 'not gpu'"
|
|
@@ -0,0 +1,66 @@
|
|
|
1
|
+
"""mojagg — NaN-aware aggregations, grouped reductions, rolling windows.
|
|
2
|
+
|
|
3
|
+
numbagg-compatible API backed by AOT-compiled Mojo kernels.
|
|
4
|
+
|
|
5
|
+
This package is the PUBLIC CONTRACT: function names, signatures, and NaN/NaT
|
|
6
|
+
semantics must match numbagg exactly. See tests/python/ (the executable spec).
|
|
7
|
+
"""
|
|
8
|
+
|
|
9
|
+
from mojagg.compat import is_registered, patch, register, unregister
|
|
10
|
+
from mojagg.config import MojaggConfig, config, get_config, set_config
|
|
11
|
+
from mojagg.fill import bfill, ffill
|
|
12
|
+
from mojagg.matrix import nancorrmatrix, nancovmatrix
|
|
13
|
+
from mojagg.nanfuncs import (
|
|
14
|
+
allnan,
|
|
15
|
+
anynan,
|
|
16
|
+
count,
|
|
17
|
+
nanargmax,
|
|
18
|
+
nanargmin,
|
|
19
|
+
nancount,
|
|
20
|
+
nanmax,
|
|
21
|
+
nanmean,
|
|
22
|
+
nanmedian,
|
|
23
|
+
nanmin,
|
|
24
|
+
nanprod,
|
|
25
|
+
nanquantile,
|
|
26
|
+
nanstd,
|
|
27
|
+
nansum,
|
|
28
|
+
nanvar,
|
|
29
|
+
)
|
|
30
|
+
|
|
31
|
+
__version__ = "0.1.0"
|
|
32
|
+
|
|
33
|
+
__all__ = [
|
|
34
|
+
"MojaggConfig",
|
|
35
|
+
"__version__",
|
|
36
|
+
"config",
|
|
37
|
+
"get_config",
|
|
38
|
+
"set_config",
|
|
39
|
+
# nanfuncs
|
|
40
|
+
"allnan",
|
|
41
|
+
"anynan",
|
|
42
|
+
"count",
|
|
43
|
+
"nanargmax",
|
|
44
|
+
"nanargmin",
|
|
45
|
+
"nancount",
|
|
46
|
+
"nanmax",
|
|
47
|
+
"nanmean",
|
|
48
|
+
"nanmedian",
|
|
49
|
+
"nanmin",
|
|
50
|
+
"nanprod",
|
|
51
|
+
"nanquantile",
|
|
52
|
+
"nanstd",
|
|
53
|
+
"nansum",
|
|
54
|
+
"nanvar",
|
|
55
|
+
# fill
|
|
56
|
+
"bfill",
|
|
57
|
+
"ffill",
|
|
58
|
+
# matrix
|
|
59
|
+
"nancorrmatrix",
|
|
60
|
+
"nancovmatrix",
|
|
61
|
+
# compat
|
|
62
|
+
"is_registered",
|
|
63
|
+
"patch",
|
|
64
|
+
"register",
|
|
65
|
+
"unregister",
|
|
66
|
+
]
|
|
@@ -0,0 +1,69 @@
|
|
|
1
|
+
"""Loader for the compiled Mojo native extension.
|
|
2
|
+
|
|
3
|
+
The extension is built out-of-band by `pixi run build-ext` (mojo build
|
|
4
|
+
--emit shared-lib) and placed on the package path. This module locates and
|
|
5
|
+
imports it, re-exporting the bound functions at module scope so facade code
|
|
6
|
+
does `from mojagg import _native; _native.nansum_f64(...)`.
|
|
7
|
+
|
|
8
|
+
During development, `mojo.importer` can auto-compile from source instead.
|
|
9
|
+
"""
|
|
10
|
+
|
|
11
|
+
from __future__ import annotations
|
|
12
|
+
|
|
13
|
+
import importlib
|
|
14
|
+
import importlib.util
|
|
15
|
+
import sysconfig
|
|
16
|
+
from pathlib import Path
|
|
17
|
+
|
|
18
|
+
_EXT_SUFFIX = sysconfig.get_config_var("EXT_SUFFIX") or ".so"
|
|
19
|
+
|
|
20
|
+
|
|
21
|
+
def _load(name: str):
|
|
22
|
+
"""Import a compiled native submodule (e.g. 'nanfuncs_native')."""
|
|
23
|
+
# 1) Already-importable on sys.path (installed wheel or PYTHONPATH=src).
|
|
24
|
+
try:
|
|
25
|
+
return importlib.import_module(name)
|
|
26
|
+
except ImportError:
|
|
27
|
+
pass
|
|
28
|
+
|
|
29
|
+
# 2) Locate the shared object next to this package or in src/.
|
|
30
|
+
here = Path(__file__).resolve().parent
|
|
31
|
+
candidates = [
|
|
32
|
+
here / f"{name}{_EXT_SUFFIX}",
|
|
33
|
+
here / f"{name}.so",
|
|
34
|
+
here.parent.parent / "src" / f"{name}.so",
|
|
35
|
+
Path.cwd() / f"{name}.so",
|
|
36
|
+
]
|
|
37
|
+
for so in candidates:
|
|
38
|
+
if so.exists():
|
|
39
|
+
spec = importlib.util.spec_from_file_location(name, so)
|
|
40
|
+
if spec and spec.loader:
|
|
41
|
+
module = importlib.util.module_from_spec(spec)
|
|
42
|
+
spec.loader.exec_module(module)
|
|
43
|
+
return module
|
|
44
|
+
|
|
45
|
+
raise ImportError(
|
|
46
|
+
f"mojagg native extension '{name}' not found. "
|
|
47
|
+
"Build it with `pixi run build-ext` (compiles src/mojagg/python/"
|
|
48
|
+
f"{name}.mojo to a shared library), or install a prebuilt wheel."
|
|
49
|
+
)
|
|
50
|
+
|
|
51
|
+
|
|
52
|
+
# Native submodule handles. Each is the compiled Mojo PythonModuleBuilder module.
|
|
53
|
+
try:
|
|
54
|
+
nanfuncs = _load("nanfuncs_native")
|
|
55
|
+
except ImportError: # pragma: no cover - allows importing mojagg before build
|
|
56
|
+
nanfuncs = None
|
|
57
|
+
|
|
58
|
+
# Native binding entry points (thin; see facade for axis-aware public API).
|
|
59
|
+
# Re-export every public binding (e.g. `nansum_f64`) at module scope so
|
|
60
|
+
# facade code does `from mojagg import _native; _native.nansum_f64(...)`.
|
|
61
|
+
if nanfuncs is not None:
|
|
62
|
+
for _name in dir(nanfuncs):
|
|
63
|
+
if not _name.startswith("_"):
|
|
64
|
+
globals()[_name] = getattr(nanfuncs, _name)
|
|
65
|
+
del _name
|
|
66
|
+
else: # pragma: no cover
|
|
67
|
+
|
|
68
|
+
def __getattr__(name: str):
|
|
69
|
+
raise ImportError("mojagg native extension not built. Run `pixi run build-ext`.")
|