Add ZarrTrace.from_store

Author: lucianopaz

pymc/backends/zarr.py

@@ -1092,3 +1092,74 @@ def to_inferencedata(self, save_warmup: bool = False) -> az.InferenceData:
             data.attrs = make_attrs(attrs=attrs, library=pymc)
             groups[name] = data.load() if az.rcParams["data.load"] == "eager" else data
         return az.InferenceData(**groups)
+
+    @classmethod
+    def from_store(
+        cls: type["ZarrTrace"],
+        store: BaseStore | MutableMapping,
+        synchronizer: Synchronizer | None = None,
+    ) -> "ZarrTrace":
+        if not _zarr_available:
+            raise RuntimeError("You must install zarr to be able to create ZarrTrace instances")
+        self: ZarrTrace = object.__new__(cls)
+        self.root = zarr.group(
+            store=store,
+            overwrite=False,
+            synchronizer=synchronizer,
+        )
+        self.synchronizer = synchronizer
+        self.compressor = default_compressor
+
+        groups = set(self.root.group_keys())
+        assert groups >= {
+            "posterior",
+            "sample_stats",
+            "warmup_posterior",
+            "warmup_sample_stats",
+            "constant_data",
+            "observed_data",
+            "_sampling_state",
+        }
+
+        if "posterior" in groups:
+            for _, array in self.posterior.arrays():
+                dims = array.attrs.get("_ARRAY_DIMENSIONS", [])
+                if len(dims) >= 2 and dims[1] == "draw":
+                    draws_per_chunk = int(array.chunks[1])
+                    break
+            else:
+                draws_per_chunk = 1
+
+        self.draws_per_chunk = int(draws_per_chunk)
+        assert self.draws_per_chunk >= 1
+
+        self.include_transformed = "unconstrained_posterior" in groups
+        arrays = itertools.chain(
+            self.posterior.arrays(),
+            self.constant_data.arrays(),
+            self.observed_data.arrays(),
+        )
+        if self.include_transformed:
+            arrays = itertools.chain(arrays, self.unconstrained_posterior.arrays())
+        varnames = []
+        coords = {}
+        vars_to_dims = {}
+        for name, array in arrays:
+            dims = array.attrs["_ARRAY_DIMENSIONS"]
+            if dims[:2] == ["chain", "draw"]:
+                # Random Variable
+                vars_to_dims[name] = dims[2:]
+                varnames.append(name)
+            elif len(dims) == 1 and name == dims[0]:
+                # Coordinate
+                # We store all model coordinates, which means we have to exclude chain
+                # and draw
+                if name not in ["chain", "draw"]:
+                    coords[name] = np.asarray(array)
+            else:
+                # Constant data or observation
+                vars_to_dims[name] = dims
+        self.varnames = varnames
+        self.coords = coords
+        self.vars_to_dims = vars_to_dims
+        return self

tests/backends/test_zarr.py

@@ -618,3 +618,41 @@ def test_sampling_consistency(
             sequential_trace._sampling_state.sampling_state[chain],
         )
     xr.testing.assert_equal(parallel_idata.posterior, sequential_idata.posterior)
+
+
+def test_from_store(populated_trace):
+    trace, total_steps, tune, draws = populated_trace
+    loaded_trace = ZarrTrace.from_store(
+        trace.root.store,
+    )
+    assert loaded_trace.is_root_populated and not loaded_trace._is_base_setup
+    assert trace.draws_per_chunk == loaded_trace.draws_per_chunk
+    assert trace.include_transformed == loaded_trace.include_transformed
+    assert set(trace.varnames) == set(loaded_trace.varnames)
+    assert set(trace.coords) == set(loaded_trace.coords) and (
+        all(
+            np.array_equal(np.asarray(coord), np.asarray(loaded_trace.coords[dim]))
+            for dim, coord in trace.coords.items()
+        )
+    )
+    assert trace.vars_to_dims == loaded_trace.vars_to_dims
+
+    assert not hasattr(loaded_trace, "straces")
+    assert set(trace.root.group_keys()) == set(loaded_trace.root.group_keys())
+    for group_name, group in trace.root.groups():
+        loaded_group = loaded_trace.root[group_name]
+        if group_name == "_sampling_state":
+            assert all(
+                equal_sampling_states(this, other) if this is not None else this is other
+                for this, other in zip(group.sampling_state[:], loaded_group.sampling_state[:])
+            )
+            np.testing.assert_array_equal(group.draw_idx, loaded_group.draw_idx)
+            assert trace.tuning_steps == loaded_trace.tuning_steps
+            assert trace.draws == loaded_trace.draws
+            assert trace.sampling_time == loaded_trace.sampling_time
+        else:
+            assert set(group.array_keys()) == set(loaded_group.array_keys())
+            for name, array in group.arrays():
+                loaded_array = loaded_group[name]
+                assert dict(array.attrs) == dict(loaded_array.attrs)
+                np.testing.assert_array_equal(np.asarray(array), np.asarray(loaded_array))