Add ZarrTrace compatibility checks

Author: lucianopaz

pymc/backends/zarr.py

@@ -1163,3 +1163,124 @@ def from_store(
         self.coords = coords
         self.vars_to_dims = vars_to_dims
         return self
+
+    def assert_model_and_step_are_compatible(
+        self,
+        step: BlockedStep | CompoundStep,
+        model: Model,
+        vars: list[TensorVariable] | None = None,
+    ):
+        zarr_groups = set(self.root.group_keys())
+        arrays_ = itertools.chain(
+            self.posterior.arrays(),
+            self.constant_data.arrays() if "constant_data" in zarr_groups else [],
+            self.observed_data.arrays() if "observed_data" in zarr_groups else [],
+        )
+        if self.include_transformed:
+            arrays_ = itertools.chain(arrays_, self.unconstrained_posterior.arrays())
+        arrays = list(arrays_)
+        zarr_varnames = []
+        zarr_coords = {}
+        zarr_vars_to_dims = {}
+        zarr_deterministics = []
+        zarr_free_vars = []
+        for name, array in arrays:
+            dims = array.attrs["_ARRAY_DIMENSIONS"]
+            if dims[:2] == ["chain", "draw"]:
+                # Random Variable
+                zarr_vars_to_dims[name] = dims[2:]
+                zarr_varnames.append(name)
+                if array.attrs["kind"] == "freeRV":
+                    zarr_free_vars.append(name)
+                else:
+                    zarr_deterministics.append(name)
+            elif len(dims) == 1 and name == dims[0]:
+                # Coordinate
+                if name not in ["chain", "draw"]:
+                    zarr_coords[name] = np.asarray(array)
+            else:
+                # Constant data or observation
+                zarr_vars_to_dims[name] = dims
+        zarr_constant_data = (
+            [name for name in self.constant_data.array_keys() if name not in zarr_coords]
+            if "constant_data" in zarr_groups
+            else []
+        )
+        zarr_observed_data = (
+            [name for name in self.observed_data.array_keys() if name not in zarr_coords]
+            if "observed_data" in zarr_groups
+            else []
+        )
+        autogenerated_dims = {dim for dim in zarr_coords if re.search(r"_dim_\d+$", dim)}
+
+        # Check deterministics, free RVs and transformed RVs
+        _, var_names = self.parse_varnames(model, vars)
+        assert set(var_names) == set(zarr_free_vars + zarr_deterministics), (
+            "The model deterministics and random variables given the sampled var_names "
+            "do not match with the stored deterministics variables in the trace."
+        )
+        for name, array in arrays:
+            if name not in zarr_free_vars or name not in zarr_deterministics:
+                continue
+            model_var = model[name]
+            assert np.dtype(model_var.dtype) == np.dtype(array.dtype), (
+                "The model deterministics and random variables given the sampled "
+                "var_names do not match with the stored deterministics variables in "
+                "the trace."
+            )
+
+        # Check coordinates
+        assert (set(zarr_coords) - set(autogenerated_dims)) == set(model.coords) and all(
+            np.array_equal(np.asarray(zarr_coords[dim]), np.asarray(coord))
+            for dim, coord in model.coords.items()
+        ), "Model coordinates don't match the coordinates stored in the trace"
+        vars_to_explicit_dims = {}
+        for name, dims in zarr_vars_to_dims.items():
+            if len(dims) == 0 or all(dim in autogenerated_dims for dim in dims):
+                # These variables wont be included in the named_vars_to_dims
+                continue
+            vars_to_explicit_dims[name] = [
+                dim if dim not in autogenerated_dims else None for dim in dims
+            ]
+        assert set(vars_to_explicit_dims) == set(model.named_vars_to_dims) and all(
+            vars_to_explicit_dims[name] == list(dims)
+            for name, dims in model.named_vars_to_dims.items()
+        ), "Some model variables have different dimensions than those stored in the trace."
+
+        # Check constant data
+        model_constant_data = find_constants(model)
+        assert set(zarr_constant_data) == set(model_constant_data), (
+            "The model constant data does not match with the stored constant data"
+        )
+        for name, model_data in model_constant_data.items():
+            assert np.array_equal(self.constant_data[name], model_data, equal_nan=True), (
+                "The model constant data does not match with the stored constant data"
+            )
+
+        # Check observed data
+        model_observed_data = find_observations(model)
+        assert set(zarr_observed_data) == set(model_observed_data), (
+            "The model observed data does not match with the stored observed data"
+        )
+        for name, model_data in model_observed_data.items():
+            assert np.array_equal(self.observed_data[name], model_data, equal_nan=True), (
+                "The model observed data does not match with the stored observed data"
+            )
+
+        # Check sample stats given the step method
+        stats_dtypes_shapes = get_stats_dtypes_shapes_from_steps(
+            [step] if isinstance(step, BlockedStep) else step.methods
+        )
+        assert (set(stats_dtypes_shapes) | {"chain", "draw"}) == set(
+            self.sample_stats.array_keys()
+        ), "The step method sample stats do not match the ones stored in the trace."
+        for name, array in self.sample_stats.arrays():
+            if name in ("chain", "draw"):
+                continue
+            assert np.dtype(stats_dtypes_shapes[name][0]) == np.dtype(array.dtype), (
+                "The step method sample stats do not match the ones stored in the trace."
+            )
+
+        assert step.sampling_state.is_compatible(self._sampling_state.sampling_state[0]), (
+            "The step method sampling state class is incompatible with what's stored in the trace."
+        )