Move dataset_to_point_list to arviz module

Author: ricardoV94

pymc/backends/arviz.py

@@ -16,15 +16,17 @@
 import logging
 import warnings
 
-from collections.abc import Iterable, Mapping
+from collections.abc import Iterable, Mapping, Sequence
 from typing import (
     TYPE_CHECKING,
     Any,
     Optional,
     Union,
+    cast,
 )
 
 import numpy as np
+import xarray
 
 from arviz import InferenceData, concat, rcParams
 from arviz.data.base import CoordSpec, DimSpec, dict_to_dataset, requires
@@ -612,3 +614,26 @@ def predictions_to_inference_data(
         # data and return that.
         concat([new_idata, idata_orig], dim=None, copy=True, inplace=True)
         return new_idata
+
+
+def dataset_to_point_list(
+    ds: xarray.Dataset | dict[str, xarray.DataArray], sample_dims: Sequence[str]
+) -> tuple[list[dict[str, np.ndarray]], dict[str, Any]]:
+    # All keys of the dataset must be a str
+    var_names = cast(list[str], list(ds.keys()))
+    for vn in var_names:
+        if not isinstance(vn, str):
+            raise ValueError(f"Variable names must be str, but dataset key {vn} is a {type(vn)}.")
+    num_sample_dims = len(sample_dims)
+    stacked_dims = {dim_name: ds[var_names[0]][dim_name] for dim_name in sample_dims}
+    transposed_dict = {vn: da.transpose(*sample_dims, ...) for vn, da in ds.items()}
+    stacked_dict = {
+        vn: da.values.reshape((-1, *da.shape[num_sample_dims:]))
+        for vn, da in transposed_dict.items()
+    }
+    points = [
+        {vn: stacked_dict[vn][i, ...] for vn in var_names}
+        for i in range(np.prod([len(coords) for coords in stacked_dims.values()]))
+    ]
+    # use the list of points
+    return cast(list[dict[str, np.ndarray]], points), stacked_dims

pymc/sampling/forward.py

@@ -49,15 +49,14 @@
 
 import pymc as pm
 
-from pymc.backends.arviz import _DefaultTrace
+from pymc.backends.arviz import _DefaultTrace, dataset_to_point_list
 from pymc.backends.base import MultiTrace
 from pymc.blocking import PointType
 from pymc.model import Model, modelcontext
 from pymc.pytensorf import compile_pymc
 from pymc.util import (
     RandomState,
     _get_seeds_per_chain,
-    dataset_to_point_list,
     default_progress_theme,
     get_default_varnames,
     point_wrapper,

pymc/stats/log_density.py

@@ -20,10 +20,14 @@
 
 import pymc
 
-from pymc.backends.arviz import _DefaultTrace, coords_and_dims_for_inferencedata
+from pymc.backends.arviz import (
+    _DefaultTrace,
+    coords_and_dims_for_inferencedata,
+    dataset_to_point_list,
+)
 from pymc.model import Model, modelcontext
 from pymc.pytensorf import PointFunc
-from pymc.util import dataset_to_point_list, default_progress_theme
+from pymc.util import default_progress_theme
 
 __all__ = ("compute_log_likelihood", "compute_log_prior")
 

pymc/util.py

@@ -16,7 +16,7 @@
 import warnings
 
 from collections.abc import Sequence
-from typing import Any, NewType, cast
+from typing import NewType, cast
 
 import arviz
 import cloudpickle
@@ -31,6 +31,20 @@
 
 from pymc.exceptions import BlockModelAccessError
 
+
+def __getattr__(name):
+    if name == "dataset_to_point_list":
+        warnings.warn(
+            f"{name} has been moved to backends.arviz. Importing from util will fail in a future release.",
+            FutureWarning,
+        )
+        from pymc.backends.arviz import dataset_to_point_list
+
+        return dataset_to_point_list
+
+    raise AttributeError(f"module {__name__} has no attribute {name}")
+
+
 VarName = NewType("VarName", str)
 
 default_progress_theme = Theme(
@@ -247,29 +261,6 @@ def enhanced(*args, **kwargs):
     return enhanced
 
 
-def dataset_to_point_list(
-    ds: xarray.Dataset | dict[str, xarray.DataArray], sample_dims: Sequence[str]
-) -> tuple[list[dict[str, np.ndarray]], dict[str, Any]]:
-    # All keys of the dataset must be a str
-    var_names = cast(list[str], list(ds.keys()))
-    for vn in var_names:
-        if not isinstance(vn, str):
-            raise ValueError(f"Variable names must be str, but dataset key {vn} is a {type(vn)}.")
-    num_sample_dims = len(sample_dims)
-    stacked_dims = {dim_name: ds[var_names[0]][dim_name] for dim_name in sample_dims}
-    transposed_dict = {vn: da.transpose(*sample_dims, ...) for vn, da in ds.items()}
-    stacked_dict = {
-        vn: da.values.reshape((-1, *da.shape[num_sample_dims:]))
-        for vn, da in transposed_dict.items()
-    }
-    points = [
-        {vn: stacked_dict[vn][i, ...] for vn in var_names}
-        for i in range(np.prod([len(coords) for coords in stacked_dims.values()]))
-    ]
-    # use the list of points
-    return cast(list[dict[str, np.ndarray]], points), stacked_dims
-
-
 def drop_warning_stat(idata: arviz.InferenceData) -> arviz.InferenceData:
     """Returns a new ``InferenceData`` object with the "warning" stat removed from sample stats groups.
 

tests/backends/test_arviz.py

@@ -16,6 +16,7 @@
 import numpy as np
 import pytensor.tensor as pt
 import pytest
+import xarray
 
 from arviz import InferenceData
 from arviz.tests.helpers import check_multiple_attrs
@@ -26,6 +27,7 @@
 
 from pymc.backends.arviz import (
     InferenceDataConverter,
+    dataset_to_point_list,
     predictions_to_inference_data,
     to_inference_data,
 )
@@ -776,3 +778,34 @@ def test_save_warmup_issue_1208_after_3_9(self):
             assert not fails
             assert idata.posterior.sizes["chain"] == 2
             assert idata.posterior.sizes["draw"] == 30
+
+
+class TestDatasetToPointList:
+    @pytest.mark.parametrize("input_type", ("dict", "Dataset"))
+    def test_dataset_to_point_list(self, input_type):
+        if input_type == "dict":
+            ds = {}
+        elif input_type == "Dataset":
+            ds = xarray.Dataset()
+        ds["A"] = xarray.DataArray([[1, 2, 3]] * 2, dims=("chain", "draw"))
+        pl, _ = dataset_to_point_list(ds, sample_dims=["chain", "draw"])
+        assert isinstance(pl, list)
+        assert len(pl) == 6
+        assert isinstance(pl[0], dict)
+        assert isinstance(pl[0]["A"], np.ndarray)
+
+    def test_transposed_dataset_to_point_list(self):
+        ds = xarray.Dataset()
+        ds["A"] = xarray.DataArray([[[1, 2, 3], [2, 3, 4]]] * 5, dims=("team", "draw", "chain"))
+        pl, _ = dataset_to_point_list(ds, sample_dims=["chain", "draw"])
+        assert isinstance(pl, list)
+        assert len(pl) == 6
+        assert isinstance(pl[0], dict)
+        assert isinstance(pl[0]["A"], np.ndarray)
+
+    def test_dataset_to_point_list_str_key(self):
+        # Check that non-str keys are caught
+        ds = xarray.Dataset()
+        ds[3] = xarray.DataArray([1, 2, 3])
+        with pytest.raises(ValueError, match="must be str"):
+            dataset_to_point_list(ds, sample_dims=["chain", "draw"])

tests/test_util.py

@@ -26,7 +26,6 @@
 from pymc.util import (
     UNSET,
     _get_seeds_per_chain,
-    dataset_to_point_list,
     drop_warning_stat,
     get_value_vars_from_user_vars,
     hash_key,
@@ -156,38 +155,6 @@ def fn(a=UNSET):
     assert "a=UNSET" in captured.out
 
 
-@pytest.mark.parametrize("input_type", ("dict", "Dataset"))
-def test_dataset_to_point_list(input_type):
-    if input_type == "dict":
-        ds = {}
-    elif input_type == "Dataset":
-        ds = xarray.Dataset()
-    ds["A"] = xarray.DataArray([[1, 2, 3]] * 2, dims=("chain", "draw"))
-    pl, _ = dataset_to_point_list(ds, sample_dims=["chain", "draw"])
-    assert isinstance(pl, list)
-    assert len(pl) == 6
-    assert isinstance(pl[0], dict)
-    assert isinstance(pl[0]["A"], np.ndarray)
-
-
-def test_transposed_dataset_to_point_list():
-    ds = xarray.Dataset()
-    ds["A"] = xarray.DataArray([[[1, 2, 3], [2, 3, 4]]] * 5, dims=("team", "draw", "chain"))
-    pl, _ = dataset_to_point_list(ds, sample_dims=["chain", "draw"])
-    assert isinstance(pl, list)
-    assert len(pl) == 6
-    assert isinstance(pl[0], dict)
-    assert isinstance(pl[0]["A"], np.ndarray)
-
-
-def test_dataset_to_point_list_str_key():
-    # Check that non-str keys are caught
-    ds = xarray.Dataset()
-    ds[3] = xarray.DataArray([1, 2, 3])
-    with pytest.raises(ValueError, match="must be str"):
-        dataset_to_point_list(ds, sample_dims=["chain", "draw"])
-
-
 def test_drop_warning_stat():
     idata = arviz.from_dict(
         sample_stats={