feat(coloring)!: JIT with numba, port SMC graph layer, track star sets by adrhill · Pull Request #104 · adrhill/asdex

adrhill · 2026-04-16T16:18:19Z

Breaking

color_symmetric now returns (colors, num_colors, star_set) instead of (colors, num_colors). The new StarSet stores per-edge star assignment and per-star hub vertex, enabling O(1) hub lookup during Hessian decompression.
New required dependency: numba>=0.63.0.

API additions

Add postprocess=True (default) to color_symmetric, which demotes colors never needed as a hub to neutral (-1), reducing HVP count.
Add forced_colors kwarg to color_symmetric to verify a user-supplied coloring against star-coloring constraints; raises the new InvalidColoringError on violation.
Promote private validators to public API: check_coloring_rows, check_coloring_cols, check_coloring_symmetric (in verify.py).
ColoredPattern gains an optional star_set field; decompression uses hub-based extraction when present and falls back to the existing uniqueness heuristic otherwise.

Internal refactor

Split monolithic coloring.py into _graph.py, _color_greedy.py, _color_symmetric.py, _postprocessing.py, _types.py, _api.py.
Nest the jaxpr interpreter modules under detection/_interpret/.
Port SMC's CSR graph layer (_build_csr, _build_symmetric_csr, _build_edge_to_index) as the shared foundation for distance-2 and star coloring.
Inline SMC attribution and link each ported algorithm back to the Julia source.
Cite Gebremedhin et al. (2007), Algorithm 4.1 directly in color_symmetric docstring; inline the trivial-star hub encoding into the StarSet.hub docstring.
Simplify coloring tests by calling the public validators directly instead of local `is_valid*_coloring` helpers.

Performance

Coloring benchmarks (mean, `N = 200`):

bench	main	new	speedup
convnet	2.22 ms	241 µs	9.2×
heat	158 µs	21.5 µs	7.3×
rosenbrock	159 µs	21.2 µs	7.5×

`_distance2_degrees`, `_partial_distance2_coloring`, `_build_edge_to_index_core`, and `_star_coloring_core` (plus `_treat`, `_update_stars`) are all `@njit(cache=True)`.
`_star_coloring_core` uses preallocated arrays (`hub_buf`/`hub_len`, parallel `fn_p`/`fn_q`/`fn_edge`) instead of Python lists, so numba can compile the loop; `InvalidColoringError` is piped out via return sentinels.
Benchmark fixtures warm up the JIT before timing.

Test plan

`uv run ruff check .`
`uv run ruff format .`
`uv run ty check`
`uv run pytest` (972 passed)
`uv run pytest -m coloring` (93 passed)
`uv run pytest tests/test_benchmarks.py -k coloring --benchmark-only -m benchmark`

- `color_symmetric` now returns `(colors, num_colors, star_set)` instead of `(colors, num_colors)`. The new `StarSet` stores per-edge star assignment and per-star hub vertex, enabling O(1) hub lookup during Hessian decompression. - Add `postprocess=True` (default) which demotes colors never needed as a hub to neutral (`-1`), reducing HVP count. - Add `forced_colors` kwarg to verify a user-supplied coloring against the star-coloring constraints; raises the new `InvalidColoringError` on violation. - Promote private validators to public API: `check_coloring_rows`, `check_coloring_cols`, `check_coloring_symmetric` (in `verify.py`). - `ColoredPattern` gains an optional `star_set` field; decompression uses hub-based extraction when present and falls back to the existing uniqueness heuristic otherwise. - `color_rows` and `color_cols` are unchanged. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

- Drop `SMC` reference from `color_symmetric` docstring and cite Gebremedhin et al. (2007), Algorithm 4.1 directly. - Add backticks to `True` in the `postprocess` docstring. - Rename section headers from `Public API: ...` to plain descriptive headers so they don't imply `StarSet` above isn't public. - Add a permalink to the SMC stamp-trick line in the source comment. - Drop the three `_is_valid_*_coloring` test helpers; call `check_coloring_rows` / `check_coloring_cols` / `check_coloring_symmetric` directly — the validators already raise `InvalidColoringError` with the specific violation. - Drop seed reporting from the fuzz test since the validator error already identifies the offending vertex/edge. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

- Inline the trivial-star hub encoding notes (previously a floating comment block) into the `StarSet.hub` attribute docstring. - Drop the `(Gebremedhin et al., 2007)` citation from `_update_stars`: `_update_stars!` is SMC's own function name, not from the paper. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

codecov-commenter · 2026-04-16T16:22:09Z

Codecov Report

❌ Patch coverage is 71.76471% with 144 lines in your changes missing coverage. Please review.
✅ Project coverage is 93.28%. Comparing base (a22653e) to head (701052c).

Files with missing lines	Patch %	Lines
src/asdex/coloring/_color_symmetric.py	38.19%	89 Missing ⚠️
src/asdex/coloring/_color_greedy.py	50.79%	31 Missing ⚠️
src/asdex/coloring/_graph.py	75.43%	14 Missing ⚠️
src/asdex/verify.py	80.39%	10 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main     #104      +/-   ##
==========================================
- Coverage   98.42%   93.28%   -5.14%     
==========================================
  Files          41       48       +7     
  Lines        2342     2606     +264     
==========================================
+ Hits         2305     2431     +126     
- Misses         37      175     +138

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

Groups the module's contents into five private submodules behind an `__init__.py` re-export: - `_types.py`: `StarSet`, `DenseColoringWarning`, `InvalidColoringError` - `_color_greedy.py`: `color_rows`, `color_cols` + conflict/greedy helpers - `_color_symmetric.py`: `color_symmetric` + adjacency/star helpers - `_postprocessing.py`: `_postprocess_star_coloring` - `_api.py`: `jacobian_coloring*`, `hessian_coloring*` + empty/dense helpers The public surface of `asdex.coloring` is unchanged — every name in `__all__` still resolves via `__init__.py` re-exports. `tests/test_coloring.py` moves its `_greedy_color` import to `asdex.coloring._color_greedy`. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

Add the SparseMatrixColorings.jl attribution block plus a link to the corresponding SMC source file in each submodule docstring of `src/asdex/coloring/`. Also correct `_types.py`, which previously claimed `StarSet` lived in `forest.jl` — it is actually defined in `coloring.jl`. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

Convert `detection.py` into a package mirroring `coloring/`: `detection.py` → `detection/_api.py`, `_interpret/` → `detection/_interpret/`. The interpreter is consumed only by `detection`, so nesting it under its sole consumer reflects the dependency direction and keeps the package root focused on pipeline stages. Updates imports in `decompression.py` and `tests/_interpret/test_internals.py`, and path references in `CLAUDE.md`, `tests/CLAUDE.md`, and `.claude/skills/add-handler/SKILL.md`. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

Introduces `_graph.py` mirroring SparseMatrixColorings.jl's `graph.jl` (`_build_csr`, `_build_symmetric_csr`, `_build_edge_to_index`) and rewrites both coloring paths on top of it. - `_color_greedy.py` now mirrors `partial_distance2_coloring` faithfully. Drops the O(Σ deg²) `list[set[int]]` conflict graphs and per-vertex `set[int]` allocations for SMC's `forbidden_colors` timestamp trick. LargestFirst uses SMC's `visited`-timestamp trick. `color_rows` and `color_cols` share one `_partial_distance2_coloring` helper. - `_color_symmetric.py`'s `_build_adjacency_with_edge_index` (list-of-tuples + `dict[(i, j), int]`) is replaced by `_build_symmetric_csr` + `_build_edge_to_index`; `_treat` and `_update_stars` iterate CSR slices. Public API and `StarSet` contract unchanged. `_greedy_color` removed; its zero-vertex test is covered by `test_color_zero_row_pattern`. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

The SMC port in `_color_greedy.py` and `_color_symmetric.py` iterated CSR arrays (`indptr`, `neighbors`, `edge_to_index`) and stamp buffers (`forbidden`, `visited`, `colors`, `star`) as numpy arrays, with per-element `int(...)` boxing in the triple-nested loops. That boxing dominates in CPython and made `color_rows` 4-7x slower than `main`'s conflict-set approach. Rebind the hot arrays to Python lists via `.tolist()` at the top of each function, run the loop in unboxed-int territory, and re-wrap as numpy only at the public boundary. No algorithmic change; all 93 coloring tests pass. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

Closes the 1.2-1.6× perf gap vs `main` left by the list-based port in 0a12bf0 and goes further — post-compile timings beat `main` by 6-9×: | bench | main | previous | new | | ---------- | ------- | -------- | -------- | | convnet | 2.22 ms | 3.53 ms | 245 µs | | heat | 158 µs | 197 µs | 23.5 µs | | rosenbrock | 159 µs | 194 µs | 21.5 µs | Changes: - Add `numba>=0.63.0` as a required dependency; ignore `*.nbi` / `*.nbc`. - `@njit(cache=True)` the hot kernels: `_build_edge_to_index_core` in `_graph.py`; `_distance2_degrees` and `_partial_distance2_coloring` in `_color_greedy.py`; `_star_coloring_core` (plus `_treat` and `_update_stars`) in `_color_symmetric.py`. - Drop the `.tolist()` rebinding; numba reads int32 arrays directly. - Refactor for numba compatibility in `_color_symmetric.py`: preallocated `hub_buf` + `hub_len` counter instead of `hub_list`; three parallel `int32[n]` arrays (`fn_p`, `fn_q`, `fn_edge`) instead of a `first_neighbor` tuple list; `InvalidColoringError` piped via return sentinels so the jitted core never raises. Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>

Call `color_rows(sparsity)` once before `benchmark(...)` in the three coloring benchmarks so the first-call compile pass is never folded into the measured rounds. Previously the timings relied on pytest-benchmark's calibration phase incidentally triggering the JIT — now the warmup is explicit, matching the `# warmup` pattern used by the materialization and end-to-end benchmarks. Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>

Align `color_symmetric`'s default with SparseMatrixColorings.jl's `GreedyColoringAlgorithm` (`postprocessing=false`), and thread a `postprocess` kwarg through `jacobian_coloring`, `hessian_coloring`, and their `_from_sparsity` counterparts so callers can opt in. Add tests exercising both paths on a 4-cycle without a diagonal (where postprocessing strictly reduces 3 colors to 2) and verifying the no-op case on a banded pattern with full diagonal. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

Add tests covering the previously unreached branches of `jacobian_coloring_from_sparsity`, `hessian_coloring_from_sparsity`, `color_symmetric(forced_colors=...)`, `StarSet.hub_vertex`, and both branches of trivial-star postprocessing. Each new test makes concrete assertions (mode resolution, HVP roundtrips, forced-coloring violations, hub-flip identity) instead of merely exercising code paths. Remove guards that are unreachable under their public-API contracts: the non-square check in `_empty_hessian_pattern`, the unused `nb_self_loops` counter in `_build_edge_to_index`, and the `num_colors == 0` / `ci < 0` / `s < 0` branches in `_postprocess_star_coloring`. Coverage (with `NUMBA_DISABLE_JIT=1`) is now 100% across `src/asdex/coloring/`. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

adrhill and others added 3 commits April 16, 2026 17:56

adrhill and others added 7 commits April 16, 2026 18:34

adrhill changed the title ~~feat(coloring)!: track star sets and add public coloring validators~~ feat(coloring)!: JIT with numba, port SMC graph layer, track star sets Apr 16, 2026

adrhill and others added 3 commits April 16, 2026 22:53

chore: bump version to 0.2.0-DEV

701052c

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

adrhill merged commit fd9241e into main Apr 16, 2026
4 checks passed

adrhill deleted the ah/fix-coloring branch April 16, 2026 21:28

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

feat(coloring)!: JIT with numba, port SMC graph layer, track star sets#104

feat(coloring)!: JIT with numba, port SMC graph layer, track star sets#104
adrhill merged 13 commits into
mainfrom
ah/fix-coloring

adrhill commented Apr 16, 2026 •

edited

Loading

Uh oh!

codecov-commenter commented Apr 16, 2026 •

edited

Loading

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

Conversation

adrhill commented Apr 16, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Breaking

API additions

Internal refactor

Performance

Test plan

Uh oh!

codecov-commenter commented Apr 16, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Codecov Report

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

adrhill commented Apr 16, 2026 •

edited

Loading

codecov-commenter commented Apr 16, 2026 •

edited

Loading