diff --git a/python/nutpie/sample.py b/python/nutpie/sample.py index 10b5c25..2b0b766 100644 --- a/python/nutpie/sample.py +++ b/python/nutpie/sample.py @@ -605,6 +605,48 @@ def wait(self, *, timeout=None): results = self._sampler.take_results() return self._extract(results) + def _sample_stats_attrs(self): + from nutpie import __version__ + + return { + "inference_library": "nutpie", + "inference_library_version": __version__, + "inference_library_settings": json.dumps(self._settings.as_dict()), + } + + def _skipped_stats(self, settings_dict): + """Sampler stats the settings say not to store, and so should not reach the trace.""" + skips = { + "store_gradient": ["gradient"], + "store_unconstrained": ["unconstrained_draw"], + "adapt_options.mass_matrix_options.store_mass_matrix": [ + "mass_matrix_inv", + "mass_matrix_eigvals", + "mass_matrix_stds", + "transformation_mu", + ], + "store_divergences": [ + "divergence_start", + "divergence_end", + "divergence_momentum", + "divergence_start_gradient", + ], + "store_transformed": ["transformed_position", "transformed_gradient"], + } + + def _get_nested(settings, name, default): + for part in name.split("."): + if part not in settings: + return default + settings = settings[part] + return settings + + skip_vars = [] + for setting, names in skips.items(): + if not _get_nested(settings_dict["settings"], setting, False): + skip_vars.extend(names) + return skip_vars + def _extract(self, results): settings_dict = self._settings.as_dict() if self._return_raw_trace: @@ -622,52 +664,61 @@ def _extract(self, results): store = cls(*args, **kwargs) obj_store = ObjectStore(store, read_only=True) - return xr.open_datatree(obj_store, engine="zarr", consolidated=False) # ty:ignore[invalid-argument-type] + trace = xr.open_datatree(obj_store, engine="zarr", consolidated=False) # ty:ignore[invalid-argument-type] + # match the arrow backend, pymc reads these from sample_stats + if "sample_stats" in trace: + trace["sample_stats"].attrs.update(self._sample_stats_attrs()) + # the settings say which stats to store; the zarr writer stores them all + skip_vars = set(self._skipped_stats(settings_dict)) + # label chain/draw as the arrow backend does (assigning in place stays lazy and + # keeps each node's children) + for node in trace.subtree: + dataset = node.dataset + missing = { + dim: np.arange(dataset.sizes[dim]) + for dim in ("chain", "draw") + if dim in dataset.dims and dim not in dataset.coords + } + stale = ( + skip_vars & set(dataset.data_vars) + if (node.name or "").endswith("sample_stats") + else set() + ) + if missing or stale: + node.dataset = dataset.drop_vars(stale).assign_coords(missing) + if not self._save_warmup: + for name in ("warmup_posterior", "warmup_sample_stats"): + if name in trace: + del trace[name] + # the zarr writer leaves the unconstrained value variables in `posterior`; + # the arrow backend moves them out, into their own group when asked for them + uc_names = self._compiled_model.reparameterized_names or [] + for group, uc_group in ( + ("posterior", "unconstrained_posterior"), + ("warmup_posterior", "warmup_unconstrained_posterior"), + ): + if group not in trace: + continue + dataset = trace[group].dataset + present = [name for name in uc_names if name in dataset.data_vars] + if not present: + continue + if self._store_unconstrained: + trace[uc_group] = xr.DataTree(dataset[present]) + trace[group].dataset = dataset.drop_vars(present) + # dict attrs have no HDF5 equivalent; keep the tree to_netcdf-able + for node in trace.subtree: + for key, value in list(node.attrs.items()): + if isinstance(value, dict): + node.attrs[key] = json.dumps(value) + return trace elif results.is_arrow(): - skip_vars = [] - skips = { - "store_gradient": ["gradient"], - "store_unconstrained": ["unconstrained_draw"], - "adapt_options.mass_matrix_options.store_mass_matrix": [ - "mass_matrix_inv", - "mass_matrix_eigvals", - "mass_matrix_stds", - ], - "store_divergences": [ - "divergence_start", - "divergence_end", - "divergence_momentum", - "divergence_start_gradient", - ], - "store_transformed": [ - "transformed_position", - "transformed_gradient", - "transformation_mu", - ], - } - - def _get_nested(settings, name, default): - parts = name.split(".") - for part in parts: - if part not in settings: - return default - settings = settings[part] - return settings - - for setting, names in skips.items(): - if not _get_nested(settings_dict["settings"], setting, False): - skip_vars.extend(names) + skip_vars = self._skipped_stats(settings_dict) draw_batches, stat_batches = results.get_arrow_trace() - from nutpie import __version__ - - attrs = { - "inference_library": "nutpie", - "inference_library_version": __version__, - "inference_library_settings": json.dumps(self._settings.as_dict()), - } + attrs = self._sample_stats_attrs() return _arrow_to_arviz( draw_batches, diff --git a/tests/test_pymc.py b/tests/test_pymc.py index 9a087d8..e033a23 100644 --- a/tests/test_pymc.py +++ b/tests/test_pymc.py @@ -1,3 +1,4 @@ +import json import time from importlib.util import find_spec @@ -650,3 +651,90 @@ def test_unnamed_shared(backend, gradient_backend): compiled = nutpie.compile_pymc_model(model) nutpie.sample(compiled) + + +@pytest.mark.pymc +def test_zarr_store_sample_stats_attrs(tmp_path): + """The zarr trace must carry the same sample_stats attrs, coords and stat set as arrow: + pymc's ``patch_nutpie_idata`` reads ``inference_library_settings`` from there.""" + with pm.Model() as model: + pm.Normal("x") + + compiled = nutpie.compile_pymc_model(model, backend="numba") + + path = tmp_path / "trace.zarr" + path.mkdir() + store = nutpie.zarr_store.LocalStore(str(path)) + zarr_trace = nutpie.sample( + compiled, chains=1, seed=123, draws=20, tune=20, zarr_store=store + ) + arrow_trace = nutpie.sample(compiled, chains=1, seed=123, draws=20, tune=20) + + for key in [ + "inference_library", + "inference_library_version", + "inference_library_settings", + ]: + assert key in zarr_trace.sample_stats.attrs, key + assert ( + zarr_trace.sample_stats.attrs[key] == arrow_trace.sample_stats.attrs[key] + ), key + + settings = json.loads(zarr_trace.sample_stats.attrs["inference_library_settings"]) + assert settings["settings"]["num_tune"] == 20 + + # the store keeps every stat and both warmup groups; the settings say otherwise + assert set(zarr_trace.sample_stats.data_vars) == set( + arrow_trace.sample_stats.data_vars + ) + assert set(zarr_trace.children) == set(arrow_trace.children) + + # chain/draw carry coordinate variables, so label-based selection works + assert "chain" in zarr_trace.posterior.coords + assert "draw" in zarr_trace.posterior.coords + zarr_trace.posterior.sel(draw=5) + + # no dict-valued attrs remain anywhere: the tree must survive to_netcdf + for node in zarr_trace.subtree: + for value in node.attrs.values(): + assert not isinstance(value, dict) + zarr_trace.to_netcdf(tmp_path / "roundtrip.nc") + + +@pytest.mark.pymc +def test_zarr_store_transformed_variables(tmp_path): + """Unconstrained value variables belong out of ``posterior`` on both backends, and in + ``unconstrained_posterior`` only when ``store_unconstrained=True``.""" + with pm.Model() as model: + sigma = pm.HalfNormal("sigma") # gives a sigma_log__ value variable + pm.Normal("mu", 0.0, sigma) + + compiled = nutpie.compile_pymc_model(model, backend="numba") + + def fit(store_unconstrained, name): + path = tmp_path / name + path.mkdir() + kwargs = { + "chains": 1, + "seed": 123, + "draws": 20, + "tune": 20, + "store_unconstrained": store_unconstrained, + } + store = nutpie.zarr_store.LocalStore(str(path)) + return ( + nutpie.sample(compiled, zarr_store=store, **kwargs), + nutpie.sample(compiled, **kwargs), + ) + + zarr_trace, arrow_trace = fit(False, "default.zarr") + assert set(zarr_trace.posterior.data_vars) == set(arrow_trace.posterior.data_vars) + assert "sigma_log__" not in zarr_trace.posterior.data_vars + assert "unconstrained_posterior" not in zarr_trace.children + + zarr_trace, arrow_trace = fit(True, "unconstrained.zarr") + assert "sigma_log__" not in zarr_trace.posterior.data_vars + assert "sigma_log__" in zarr_trace.unconstrained_posterior.data_vars + assert set(zarr_trace.unconstrained_posterior.data_vars) == set( + arrow_trace.unconstrained_posterior.data_vars + )