Make zarr traces equivalent to arrow traces

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,

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
+    )