Small bugfixes and test

Author: jessegrabowski

pymc/data.py

@@ -27,7 +27,7 @@
 import pytensor.tensor as pt
 import xarray as xr
 
-from narwhals.typing import IntoFrameT, IntoLazyFrameT, IntoSeriesT
+from narwhals.typing import IntoFrameT, IntoSeriesT
 from pytensor.compile import SharedVariable
 from pytensor.compile.builders import OpFromGraph
 from pytensor.graph.basic import Variable
@@ -174,11 +174,11 @@ def _handle_none_dims(
 
 @singledispatch
 def determine_coords(
-    value,
+    value: typing.Any,
     model: "Model",
     dims: Sequence[str | None] | None = None,
     coords: dict[str, Sequence | np.ndarray] | None = None,
-) -> tuple[dict[str, Sequence | np.ndarray], Sequence[str | None] | Sequence[None]]:
+) -> tuple[typing.Any, dict[str, Sequence | np.ndarray], Sequence[str | None] | Sequence[None]]:
     """Determine coordinate values from data or the model (via ``dims``)."""
     raise NotImplementedError(
         f"Cannot determine coordinates for data of type {type(value)}, please provide `coords` explicitly or "
@@ -192,12 +192,12 @@ def determine_array_coords(
     model: "Model",
     dims: Sequence[str] | None = None,
     coords: dict[str, Sequence | np.ndarray] | None = None,
-) -> tuple[dict[str, Sequence | np.ndarray], Sequence[str | None] | Sequence[None]]:
+) -> tuple[np.ndarray, dict[str, Sequence | np.ndarray], Sequence[str | None] | Sequence[None]]:
     if coords is None:
         coords = {}
 
     if dims is None:
-        return coords, _handle_none_dims(dims, value.ndim)
+        return value, coords, _handle_none_dims(dims, value.ndim)
 
     if len(dims) != value.ndim:
         raise ShapeError(
@@ -211,7 +211,7 @@ def determine_array_coords(
         if coord is None and dim is not None:
             coords[dim] = range(size)
 
-    return coords, _handle_none_dims(dims, value.ndim)
+    return value, coords, _handle_none_dims(dims, value.ndim)
 
 
 @determine_coords.register(xr.DataArray)
@@ -220,41 +220,41 @@ def determine_xarray_coords(
     model: "Model",
     dims: Sequence[str | None] | None = None,
     coords: dict[str, Sequence | np.ndarray] | None = None,
-) -> tuple[dict[str, Sequence | np.ndarray], Sequence[str | None] | Sequence[None]]:
+) -> tuple[xr.DataArray, dict[str, Sequence | np.ndarray], Sequence[str | None] | Sequence[None]]:
     if coords is None:
         coords = {}
 
     if dims is None:
-        return coords, _handle_none_dims(dims, value.ndim)
+        return value, coords, _handle_none_dims(dims, value.ndim)
 
     for dim in dims:
         dim_name = dim
         # str is applied because dim entries may be None
         coords[str(dim_name)] = cast(xr.DataArray, value[dim]).to_numpy()
 
-    return coords, _handle_none_dims(dims, value.ndim)
+    return value, coords, _handle_none_dims(dims, value.ndim)
 
 
 def _dataframe_agnostic_coords(
-    value: IntoFrameT | IntoLazyFrameT | nw.DataFrame | nw.LazyFrame,
+    value: IntoFrameT,
     model: "Model",
     ndim_in: int = 2,
     dims: Sequence[str | None] | None = None,
     coords: dict[str, Sequence | np.ndarray] | None = None,
-) -> tuple[dict[str, Sequence | np.ndarray], Sequence[str | None] | Sequence[None]]:
+) -> tuple[IntoFrameT, dict[str, Sequence | np.ndarray], Sequence[str | None] | Sequence[None]]:
     if coords is None:
         coords = {}
 
     value = cast(nw.DataFrame | nw.LazyFrame, nw.from_native(value, allow_series=False))  # type: ignore[type-var]
     if isinstance(value, nw.LazyFrame):
         value = value.collect()
 
-    index = nw.maybe_get_index(value)
-    if index is not None:
-        value = value.with_columns(**{index.name: index.to_numpy()})
-
     if dims is None:
-        return coords, _handle_none_dims(dims, ndim_in)
+        if ndim_in == 1:
+            value = value[value.columns[0]]
+        return value.to_native(), coords, _handle_none_dims(dims, ndim_in)
+
+    index = nw.maybe_get_index(value)
 
     if len(dims) != ndim_in:
         raise ShapeError(
@@ -265,13 +265,13 @@ def _dataframe_agnostic_coords(
 
     index_dim = dims[0]
     if index_dim is not None:
-        if index_dim in value.columns:
-            coords[index_dim] = tuple(value.select(nw.col(index_dim)).to_numpy().flatten())
+        if index is not None:
+            coords[index_dim] = tuple(index)
         elif index_dim in model.coords:
             coords[index_dim] = model.coords[index_dim]  # type: ignore[assignment]
         else:
             raise ValueError(
-                f"Dimension '{index_dim}' not found in DataFrame columns or model coordinates. Cannot infer "
+                f"Dimension '{index_dim}' not found in DataFrame index or model coordinates. Cannot infer "
                 "index coordinates."
             )
 
@@ -281,7 +281,10 @@ def _dataframe_agnostic_coords(
             select_expr = nw.exclude(index_dim) if index_dim is not None else nw.all()
             coords[column_dim] = value.select(select_expr).columns
 
-    return coords, _handle_none_dims(dims, ndim_in)
+    if ndim_in == 1:
+        value = value[value.columns[0]]
+
+    return value.to_native(), coords, _handle_none_dims(dims, ndim_in)
 
 
 def _series_agnostic_coords(
@@ -319,7 +322,9 @@ def determine_dataframe_coords(
             model: "Model",
             dims: Sequence[str] | None = None,
             coords: dict[str, Sequence | np.ndarray] | None = None,
-        ) -> tuple[dict[str, Sequence | np.ndarray], Sequence[str | None] | Sequence[None]]:
+        ) -> tuple[
+            IntoFrameT, dict[str, Sequence | np.ndarray], Sequence[str | None] | Sequence[None]
+        ]:
             return _dataframe_agnostic_coords(value, model=model, dims=dims, coords=coords)
 
     except ImportError:
@@ -453,7 +458,7 @@ def Data(
 
     new_dims: Sequence[str | None] | Sequence[None] | None
     if infer_dims_and_coords:
-        coords, new_dims = determine_coords(value, model, dims)
+        value, coords, new_dims = determine_coords(value, model, dims)
     else:
         new_dims = dims
 

tests/test_data.py

@@ -414,7 +414,7 @@ def test_implicit_coords_polars_series(self):
 
             with pytest.raises(
                 ValueError,
-                match="Dimension 'date2' not found in DataFrame columns or model coordinates",
+                match="Dimension 'date2' not found in DataFrame index or model coordinates",
             ):
                 pm.Data("sales_invalid", ser_sales, dims=["date2"], infer_dims_and_coords=True)
 
@@ -428,14 +428,69 @@ def test_implicit_coords_polars_dataframe(self):
         df_data = pl.DataFrame(
             np.random.normal(size=size),
             schema={f"Column {c + 1}": pl.Float64 for c in range(size[1])},
-        ).with_row_count("rows")
+        )
 
-        with pm.Model() as pmodel:
-            pm.Data("observations", df_data, dims=("rows", "columns"), infer_dims_and_coords=True)
+        # We currently count on the presence of an index in the DataFrame to infer dims. Polars has no index, so
+        # this case errors because we can't find the 'rows' dim.
 
-        assert "rows" in pmodel.coords
-        assert "columns" in pmodel.coords
-        assert pmodel.named_vars_to_dims == {"observations": ("rows", "columns")}
+        with pytest.raises(
+            ValueError, match="Dimension 'rows' not found in DataFrame index or model coordinates"
+        ):
+            with pm.Model() as pmodel:
+                pm.Data(
+                    "observations", df_data, dims=("rows", "columns"), infer_dims_and_coords=True
+                )
+
+    def test_implicit_coords_agnostic(self):
+        pl = pytest.importorskip("polars")
+        pd = pytest.importorskip("pandas")
+
+        size = (5, 7)
+        data_np = np.random.normal(size=size)
+        columns = [f"C{c + 1}" for c in range(size[1])]
+        rows = [f"R{r + 1}" for r in range(size[0])]
+        df_pd = pd.DataFrame(data_np, columns=columns, index=rows)
+        df_pd.index.name = "rows"
+        df_pl = pl.DataFrame(
+            data_np,
+            schema=dict.fromkeys(columns, pl.Float64),
+        )
+
+        def make_model(coords, df, dims, infer_dims_and_coords) -> pm.Model:
+            with pm.Model(coords=coords) as pmodel:
+                pm.Data("X", df, dims=dims, infer_dims_and_coords=infer_dims_and_coords)
+            return pmodel
+
+        expected_coords = {"rows": tuple(rows), "columns": tuple(columns)}
+        dims = ("rows", "columns")
+
+        m = make_model(coords=None, df=df_pd, dims=dims, infer_dims_and_coords=True)
+        assert m.coords == expected_coords
+        np.testing.assert_allclose(m["X"].eval(), df_pd.values)
+
+        # TODO: Is infer_dims_and_coords supposed to infer dims? The current behavior is that it doesn't, it only
+        #  infers the dimension labels.
+        for df in [df_pd, df_pl]:
+            m = make_model(coords=None, df=df, dims=None, infer_dims_and_coords=True)
+            assert m.coords == {}
+
+            m = make_model(coords=None, df=df, dims=dims, infer_dims_and_coords=False)
+            assert m.coords == {"rows": None, "columns": None}
+
+            m = make_model(coords=None, df=df, dims=None, infer_dims_and_coords=False)
+            assert m.coords == {}
+
+        # Pandas is special because we will infer the index dim from the DataFrame index, if one exists.
+        m = make_model(coords=None, df=df_pd, dims=dims, infer_dims_and_coords=True)
+        assert m.coords == expected_coords
+
+        # Polars (and other dataframe backends with no index concept) won't infer dims from index. This case currently
+        # errors, because we can't find the 'rows' dim in either the DataFrame columns or the model coords.
+        with pytest.raises(
+            ValueError,
+            match="Dimension 'rows' not found in DataFrame index or model coordinates",
+        ):
+            make_model(coords=None, df=df_pl, dims=dims, infer_dims_and_coords=True)
 
     def test_implicit_coords_xarray(self):
         xr = pytest.importorskip("xarray")