Declare only the sampler stats the settings store

[velochy]

What / why

Follow-up to pymc-devs/nutpie#335, where
@aseyboldt pointed out that this belongs here rather than in nutpie's Python layer, so that a
zarr store loaded directly — without going through nutpie — is already in the right shape. He
also expected the disabled store_* stats to cost nothing. That's mostly true, and the part
where it isn't is what this changes.

The zarr writer creates an array for every declared stat. A stat switched off by
store_gradient / store_unconstrained / store_transformed is never written, so its array
simply stays at the NaN fill value.

Six of the ten are genuinely free: their primary dimension is divergence- or
mass-matrix-shaped, which is already zero-length when the corresponding option is off. They
need no filtering and this PR does not touch them.

The other four — gradient, unconstrained_draw, transformed_position,
transformed_gradient — are draw-dimensioned and get allocated at full width. In the store
that is nearly free, since all-NaN compresses to about a kilobyte. The cost lands on every
reader: xarray materializes them at full width on load and again on conversion. On a
40-parameter, 200-draw, 2-chain fit they are 0.5 MB of a 1.4 MB netCDF conversion, and the
ratio grows with the parameter count.

The change

Filter the disabled stats where the settings and the stat names meet — in the Settings::stat_names
default method — so that stat_types, stat_dims_all and every storage backend inherit the
filtering and cannot drift from what the sampler actually emits:

/// Stats these settings switch off. They are never written, so declaring them would only
/// create arrays that stay at their fill value and that every storage backend then carries.
fn disabled_stats(&self) -> Vec<&'static str> {
    Vec::new()
}

fn stat_names<M: Math>(&self, math: &M) -> Vec<String> {
    let dims = StatsDims::from(math);
    let disabled = self.disabled_stats();
    <<Self::Chain<M> as SamplerStats<M>>::Stats as Storable<_>>::names(&dims)
        .into_iter()
        .filter(|name| !disabled.contains(name))
        .map(String::from)
        .collect()
}

A disabled_point_stats helper maps the three store_* flags onto the names they suppress,
mirroring TransformedPoint::extract_stats so the two lists can be read against each other,
and the six Settings implementors override disabled_stats() with it.

Compatibility

disabled_stats() has a default returning an empty Vec, so the trait gains a method without
breaking implementors outside this crate — their behaviour is unchanged unless they opt in.

Nothing changes for a run with all store_* options on.

Testing

cargo check clean; the existing suite passes (41 unit tests + 2 doctests).

The 0.5 MB / 1.4 MB figure above was measured from the Python side, on a trace that still
carries the four arrays — i.e. it measures what they cost a reader, not the effect of this
patch. I have not yet built nutpie against a patched nuts-rs to confirm the arrays disappear
end-to-end; if that matters for review I can do it.

Note on scope

This covers the store_* half of nutpie#335. The save_warmup half and the chain/draw
coordinate variables are separate and not in this PR — happy to follow up with those here too if
that's where you'd like them, which I read your comment on #335 as suggesting.

[aseyboldt]

I think this also requires a change when we iterate over the stats, to skip the unexpected ones.

And maybe this would be a great opportunity to also fix pymc-devs/nutpie#315?
That is basically a skip for the posterior variables instead of stat variables.

[velochy]

I think this also requires a change when we iterate over the stats, to skip the unexpected ones.

And maybe this would be a great opportunity to also fix pymc-devs/nutpie#315? That is basically a skip for the posterior variables instead of stat variables.

#315 definitely seems like something this should handle.
I ran out of my weekly AI credits, so I'll have my AI get back on this on Monday when they reset :)

[velochy]

Both done in 62ee78f.

On skipping the unexpected stats when we iterate. You were right, and it was worse than a loose end: the previous commit changed what the backends allocate but not what the chain emits, so a stat switched off by store_* still arrived undeclared. The arrow backend zips incoming values against its declared columns and panics:

Draw name mismatch: expected fisher_distance, got unconstrained_draw

Since the three store_* flags all default to false, that was every default arrow run, including arrow's own store_warmup test. I missed it because arrow is not a default feature — cargo test passes and cargo test --all-features does not. Might be worth having CI run the feature matrix.

The fix, for both halves at once: filter the values against the declared names at the single point where the sampler hands them to storage, for stats and for expanded draws alike. Backends then receive exactly what they allocated for, and stat_names / data_names become the one authority instead of a second list that can drift.

On #315. That falls out of the same change. The obstacle was not that a model could not express "do not store this" — Storable::names is already that hook — but that under-declaring crashed: arrow shifted every column after the gap, zarr panicked with Unknown posterior variable name. Declaring fewer expanded names than you produce values for is now supported and the rest are dropped.

So nutpie can keep its full variables list, which it needs for the cumulative start_idx/end_idx that slice the expand buffer, and simply leave the variables var_names excludes out of ExpandedVector::names. The free variables are still computed by the fused expand function, but they stop being written and stop being materialized by every reader. Happy to send that nutpie PR if you want it; this one only makes it possible, so I have not claimed #315 fixed.

Testing. Two regression tests, arrow (positional, catches column shifts) and zarr (name-keyed, and what nutpie uses), driven by a test model that expands two variables and declares one. Both fail without the filter, with the two panics quoted above. 45 unit + 4 integration + 2 doctests green under --all-features; no new clippy warnings against the previous commit; --no-default-features builds clean.

Still outstanding from the description: I have not built nutpie against a patched nuts-rs, so the end-to-end effect on a real trace is still unverified. And the save_warmup half of nutpie#335 plus the chain/draw coordinates are not in this PR — say the word and I will add them here.

[velochy]

Correction to my previous comment: the CI suggestion in it was wrong, and I should have checked before making it. .github/workflows/test.yml already runs cargo build --verbose --all-features and cargo test --verbose --all-features on stable and nightly, and coverage.yml runs --all-features too. There is no feature-matrix gap.

CI had in fact already caught this — 847b9e2 is sitting red on this PR (Test Suite (nightly): failure, coverage: failure). I had run cargo test locally without the feature flags and reported the suite as passing on that basis; the failure was on the PR the whole time. Sorry for the noise.

Two follow-up commits since:

• 6a6965e — my arrow regression test was partly tautological. The test model emitted the declared name first, so arrow's positional zip paired it with its one builder and stopped before reaching the undeclared one; the posterior assertion held with or without the draw filter, and only the stat filter was actually under test. Emitting the undeclared value first makes it fail at the draws zip (arrow.rs:503) when the filter is removed, which is what I claimed it did.
• Doc comments on stat_names/data_names reworded — they credited the dropping to those methods, but it happens in ChainProcess.

Two things I checked while re-reviewing, both worth separating from this PR:

The commit-message claim about the other six stats is correct. I measured it rather than trusting it: with default settings the event-dimensioned stats are [2, 0, 8] (mass_matrix_inv, transformation_mu, all divergence_*), so they genuinely cost nothing and need no filtering.

But there is a real, pre-existing data-loss bug next door, which I found while checking the above and can send separately if you want it. ZarrChainStorage::finalize picks one representative field per event dimension via HashMap iteration order (zarr/sync_impl.rs:279-287, and :239-248 for warmup) and resizes the whole dimension to that field's push count. When a never-written field wins — e.g. mass_matrix_inv with store_mass_matrix=false, which shares the transformation_update dim with transformation_update_id — the dimension is truncated to 0 and the events that were recorded are silently discarded. Reproduced on origin/main (not caused by this PR), 5 runs, store_divergences=true:

transformation_update_id=[2, 0]    divergence_draw=[2, 131]   (trial 0)
transformation_update_id=[2, 167]  divergence_draw=[2, 127]   (trial 1)
transformation_update_id=[2, 0]    divergence_draw=[2, 128]   (trial 2)
transformation_update_id=[2, 167]  divergence_draw=[2, 0]     (trial 3)   <- ~125 divergences lost
transformation_update_id=[2, 0]    divergence_draw=[2, 121]   (trial 4)

Taking the max over the fields of a dimension instead of a representative would fix it; note the lottery persists even with store_divergences=true, since divergence_energy_error is Some only for energy-error divergences while divergence_draw is always pushed.

Happy to fold that into this PR or open a separate one — your call. I have also not squashed 847b9e2 into 62ee78f, which would take CI green across the whole branch and let git-cliff pick the change up (both subjects are non-conventional, so as it stands the changelog would not mention that four sample_stats arrays disappear). Say the word and I will squash and force-push.

[codecov-commenter]

Codecov Report

❌ Patch coverage is 90.17341% with 17 lines in your changes missing coverage. Please review.
✅ Project coverage is 67.29%. Comparing base (394c629) to head (6a6965e).
⚠️ Report is 157 commits behind head on main.

Files with missing lines	Patch %	Lines
src/sampler.rs	81.25%	15 Missing ⚠️
src/storage/arrow.rs	97.43%	1 Missing ⚠️
src/storage/zarr/sync_impl.rs	95.83%	1 Missing ⚠️

❗ There is a different number of reports uploaded between BASE (394c629) and HEAD (6a6965e). Click for more details.

HEAD has 1 upload less than BASE

Flag	BASE (394c629)	HEAD (6a6965e)
	2	1

Additional details and impacted files

@@             Coverage Diff             @@
##             main      #75       +/-   ##
===========================================
- Coverage   83.77%   67.29%   -16.48%     
===========================================
  Files           8       34       +26     
  Lines        1923     9580     +7657     
===========================================
+ Hits         1611     6447     +4836     
- Misses        312     3133     +2821     

☔ View full report in Codecov by Harness.
📢 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.