Add support for polars dataframes and series

conda-envs/environment-dev.yml

@@ -13,6 +13,7 @@ dependencies:
 - numpy>=1.15.0
 - pandas>=0.24.0
 - pip
+- polars>=1.5.0
 - pytensor>=2.25.1,<2.26
 - python-graphviz
 - networkx

conda-envs/environment-docs.yml

@@ -11,6 +11,7 @@ dependencies:
 - numpy>=1.15.0
 - pandas>=0.24.0
 - pip
+- polars>=1.5.0
 - pytensor>=2.25.1,<2.26
 - python-graphviz
 - rich>=13.7.1

conda-envs/environment-jax.yml

@@ -20,6 +20,7 @@ dependencies:
 - numpyro>=0.8.0
 - pandas>=0.24.0
 - pip
+- polars>=1.5.0
 - pytensor>=2.25.1,<2.26
 - python-graphviz
 - networkx

conda-envs/windows-environment-dev.yml

@@ -13,6 +13,7 @@ dependencies:
 - numpy>=1.15.0
 - pandas>=0.24.0
 - pip
+- polars>=1.5.0
 - pytensor>=2.25.1,<2.26
 - python-graphviz
 - networkx

conda-envs/windows-environment-test.yml

@@ -16,6 +16,7 @@ dependencies:
 - numpy>=1.15.0
 - pandas>=0.24.0
 - pip
+- polars>=1.5.0
 - pytensor>=2.25.1,<2.26
 - python-graphviz
 - networkx

pymc/data.py

@@ -22,6 +22,7 @@
 
 import numpy as np
 import pandas as pd
+import polars as pl
 import pytensor
 import pytensor.tensor as pt
 import xarray as xr
@@ -204,7 +205,7 @@ def Minibatch(variable: TensorVariable, *variables: TensorVariable, batch_size:
 
 def determine_coords(
     model,
-    value: pd.DataFrame | pd.Series | xr.DataArray,
+    value: pd.DataFrame | pd.Series | pl.DataFrame | pl.Series | xr.DataArray,
     dims: Sequence[str | None] | None = None,
     coords: dict[str, Sequence | np.ndarray] | None = None,
 ) -> tuple[dict[str, Sequence | np.ndarray], Sequence[str | None]]:
@@ -348,7 +349,7 @@ def Data(
     ----------
     name : str
         The name for this variable.
-    value : array_like or pandas.Series, pandas.Dataframe
+    value : array_like or pandas.Series, pandas.Dataframe, polars.DataFrame, polars.Series, xarray.DataArray
         A value to associate with this variable.
     dims : str or tuple of str, optional
         Dimension names of the random variables (as opposed to the shapes of these

pymc/pytensorf.py

@@ -18,6 +18,7 @@
 
 import numpy as np
 import pandas as pd
+import polars as pl
 import pytensor
 import pytensor.tensor as pt
 import scipy.sparse as sps
@@ -96,12 +97,16 @@ def convert_generator_data(data) -> TensorVariable:
 
 def convert_data(data) -> np.ndarray | Variable:
     ret: np.ndarray | Variable
-    if hasattr(data, "to_numpy") and hasattr(data, "isnull"):
+    if hasattr(data, "to_numpy"):
         # typically, but not limited to pandas objects
         vals = data.to_numpy()
-        null_data = data.isnull()
+        if hasattr(data, "is_null"):
+            # polars DataFrame or Series
+            null_data = data.is_null()
+        else:
+            null_data = data.isnull()
         if hasattr(null_data, "to_numpy"):
-            # pandas Series
+            # pandas or polars Series
             mask = null_data.to_numpy()
         else:
             # pandas Index
@@ -144,7 +149,9 @@ def convert_data(data) -> np.ndarray | Variable:
 
 @_as_tensor_variable.register(pd.Series)
 @_as_tensor_variable.register(pd.DataFrame)
-def dataframe_to_tensor_variable(df: pd.DataFrame, *args, **kwargs) -> TensorVariable:
+@_as_tensor_variable.register(pl.DataFrame)
+@_as_tensor_variable.register(pl.Series)
+def dataframe_to_tensor_variable(df: pd.DataFrame | pl.DataFrame, *args, **kwargs) -> TensorVariable:
     return pt.as_tensor_variable(df.to_numpy(), *args, **kwargs)
 
 

requirements-dev.txt

@@ -15,6 +15,7 @@ numdifftools>=0.9.40
 numpy>=1.15.0
 numpydoc
 pandas>=0.24.0
+polars>=1.5.0
 polyagamma
 pre-commit>=2.8.0
 pytensor>=2.25.1,<2.26

tests/test_data.py

@@ -404,6 +404,21 @@ def test_implicit_coords_dataframe(self, seeded_test):
         assert "columns" in pmodel.coords
         assert pmodel.named_vars_to_dims == {"observations": ("rows", "columns")}
 
+    def test_implicit_coords_polars_dataframe(self, seeded_test):
+        pl = pytest.importorskip("polars")
+        N_rows = 5
+        N_cols = 7
+        df_data = pl.DataFrame({f"Column {c+1}": np.random.normal(size=(N_rows,)) for c in range(N_cols)})
+        df_data = df_data.with_row_count("rows")
+
+        # infer coordinates from index and columns of the DataFrame
+        with pm.Model() as pmodel:
+            pm.Data("observations", df_data, dims=("rows", "columns"), infer_dims_and_coords=True)
+
+        assert "rows" in pmodel.coords
+        assert "columns" in pmodel.coords
+        assert pmodel.named_vars_to_dims == {"observations": ("rows", "columns")}
+
     def test_implicit_coords_xarray(self):
         xr = pytest.importorskip("xarray")
         data = xr.DataArray([[1, 2, 3], [4, 5, 6]], dims=("y", "x"))

tests/test_pytensorf.py

@@ -16,6 +16,7 @@
 import numpy.ma as ma
 import numpy.testing as npt
 import pandas as pd
+import polars as pl
 import pytensor
 import pytensor.tensor as pt
 import pytest
@@ -66,6 +67,19 @@ def test_pd_dataframe_as_tensor_variable(np_array: np.ndarray) -> None:
     df = pd.DataFrame(np_array)
     np.testing.assert_array_equal(x=pt.as_tensor_variable(x=df).eval(), y=np_array)
 
+@pytest.mark.parametrize(
+    argnames="np_array",
+    argvalues=[
+        np.array([[1.0], [2.0], [-1.0]]),
+        np.array([[1.0, 1.0, 1.0], [0.0, 0.0, 0.0]]),
+        np.ones(shape=(10, 1)),
+    ],
+)
+def test_polars_dataframe_as_tensor_variable(np_array: np.ndarray) -> None:
+    pl = pytest.importorskip("polars")
+    df = pl.DataFrame(np_array)
+    np.testing.assert_array_equal(x=pt.as_tensor_variable(x=df).eval(), y=np_array)
+
 
 @pytest.mark.parametrize(
     argnames="np_array",