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 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.
@@ -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).
@@ -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`.")