Test and JAX coverage
This page gives the share of sys_mapping that the test suite exercises and the share of its numerical work that runs through JAX.
scripts/coverage_report.py writes it from the coverage report of the test suite, the static analysis of scripts/jax_coverage.py, the transformability cases of tests/test_jax_transformability.py and, when one is given, a profiled pipeline run.
Test coverage
The test suite covers 2463 of 2493 statements (98.8%) and 542 of 572 branches (94.8%), 98.0% combined; the table gives the coverage per module.
Docstring examples run as part of the suite.
CI measures the same coverage, branches included, on every push and fails below the fail_under value in pyproject.toml; tests that need the LS10 and Gaia maps or the Uchuu mocks are skipped there.
module |
statements |
lines % |
branches |
branches % |
combined % |
|---|---|---|---|---|---|
__init__.py |
21 |
100.0 |
0 |
— |
100.0 |
_array.py |
6 |
100.0 |
0 |
— |
100.0 |
bootstrap.py |
55 |
100.0 |
8 |
100.0 |
100.0 |
contamination.py |
119 |
98.3 |
50 |
96.0 |
97.6 |
correction.py |
111 |
99.1 |
22 |
95.5 |
98.5 |
covariance.py |
97 |
99.0 |
32 |
96.9 |
98.4 |
diagnostics.py |
297 |
99.7 |
58 |
96.6 |
99.2 |
glass_mocks.py |
177 |
98.9 |
48 |
95.8 |
98.2 |
inference.py |
189 |
91.5 |
48 |
77.1 |
88.6 |
likelihood.py |
50 |
100.0 |
10 |
100.0 |
100.0 |
maps.py |
137 |
100.0 |
46 |
95.7 |
98.9 |
mocks.py |
121 |
100.0 |
24 |
100.0 |
100.0 |
model_selection.py |
250 |
99.6 |
38 |
97.4 |
99.3 |
nuts.py |
111 |
99.1 |
24 |
95.8 |
98.5 |
plotting.py |
5 |
100.0 |
0 |
— |
100.0 |
power_spectrum.py |
45 |
100.0 |
12 |
100.0 |
100.0 |
regression.py |
318 |
99.4 |
82 |
95.1 |
98.5 |
simulation.py |
214 |
98.6 |
38 |
92.1 |
97.6 |
utils.py |
170 |
100.0 |
32 |
100.0 |
100.0 |
JAX coverage
Transformability under jit, vmap and grad
tests/test_jax_transformability.py applies jax.jit, jax.vmap and, where a gradient is meaningful, jax.grad to the public numeric functions and checks the result against the eager call.
27 of 46 cases pass; 9 of 17 functions pass every transform that applies.
The failures are strict expected failures, so a function that becomes transformable fails the suite until its entry is removed.
function |
jit |
vmap |
grad |
blocked by |
|---|---|---|---|---|
make_log_likelihood |
yes |
yes |
yes |
|
apply_contamination |
yes |
yes |
yes |
|
invert_contamination |
yes |
yes |
yes |
|
compute_two_point_correction |
yes |
yes |
yes |
|
compute_two_point_correction[matrix] |
yes |
yes |
yes |
|
debias_params |
yes |
yes |
yes |
|
debias_params_matrix |
yes |
yes |
yes |
|
correct_two_point_function |
no |
no |
no |
np.asarray on w_obs (correction.py) |
rotate_templates |
no |
no |
— |
np.linalg.eigh on the second moment (correction.py) |
transform_params_from_rotated |
yes |
yes |
yes |
|
sample_covariance |
no |
no |
no |
np.cov (covariance.py) |
mock_sandwich_covariance |
no |
no |
no |
np.linalg on the mock fields (covariance.py) |
calibrated_template_significance |
no |
no |
— |
np.linalg.pinv and input validation (diagnostics.py) |
residual_template_correlation_test |
no |
no |
— |
np.asarray and input validation (diagnostics.py) |
snr_template_ranking |
no |
no |
— |
np.asarray before the JAX kernel (diagnostics.py) |
standardise_on_footprint |
yes |
yes |
yes |
|
posterior_median_params |
no |
no |
— |
np.median (inference.py) |
Time profile of a pipeline run
We profiled one run of scripts/run_ls10_analysis.py (log M* ≥ 10.0 at NSIDE 32, 40 significance and 10 ISD realisations, NUTS 500+500, cross-template matrix from 10⁶ galaxies, on 8 cores of a GRICAD dahu node, campaign 20260915prof) and summed the own time per library with scripts/profile_by_library.py, 760 s in total.
Time spent executing compiled XLA code is counted under JAX.
library |
seconds |
share % |
|---|---|---|
treecorr |
505 |
66.4 |
jax |
209 |
27.5 |
other |
27 |
3.5 |
numpy |
12 |
1.6 |
healpy |
5 |
0.6 |
sys_mapping |
1 |
0.1 |
astropy |
1 |
0.1 |
The cumulative time of the package functions that dominate, by stage:
function |
calls |
seconds |
share % |
|---|---|---|---|
template_correlation_matrix |
1 |
456 |
60.0 |
run_nuts |
1 |
199 |
26.2 |
measure_two_point_function |
1 |
63 |
8.3 |
run_additive_analytic |
1 |
3 |
0.4 |
elasticnet_contamination_fit |
1 |
3 |
0.3 |
isd_template_significance |
1 |
2 |
0.3 |
correct_two_point_function |
8 |
2 |
0.3 |
iterative_systematics_decontamination |
2 |
1 |
0.1 |
calibrated_template_significance |
1 |
0 |
0.0 |
The run compiled 326 XLA programs.
Code outside JAX
TreeCorr and Corrfunc pair counts: tree-based C++ pair counting with no JAX equivalent.
healpy harmonic transforms (anafast, synfast, ud_grade): a JAX version needs jax-healpy or s2fft, which the package does not depend on.
GLASS field generation: GLASS is a NumPy library.
scikit-learn ElasticNet: cross-validated fits take 0.05–0.12 s.
FITS input and output: I/O.
Reproduction
pytest --cov=sys_mapping --cov-branch --cov-report=json:coverage.json
python scripts/jax_coverage.py
JAX_LOG_COMPILES=1 python -m cProfile -o run.prof scripts/run_ls10_analysis.py ... 2> run.err
python scripts/profile_by_library.py run.prof --compile-log run.err
python scripts/coverage_report.py coverage.json --profile run.prof --compile-log run.err