Support Polars (#179)

Author: aloctavodia

pymc_bart/bart.py

@@ -74,9 +74,9 @@ class BART(Distribution):
 
     Parameters
     ----------
-    X : TensorLike
+    X : PyTensor Variable, Pandas/Polars DataFrame or Numpy array
         The covariate matrix.
-    Y : TensorLike
+    Y : PyTensor Variable, Pandas/Polar DataFrame/Series,or Numpy array
         The response vector.
     m : int
         Number of trees.
@@ -204,6 +204,16 @@ def preprocess_xy(
     if isinstance(X, (Series, DataFrame)):
         X = X.to_numpy()
 
+    try:
+        import polars as pl
+
+        if isinstance(X, (pl.Series, pl.DataFrame)):
+            X = X.to_numpy()
+        if isinstance(Y, (pl.Series, pl.DataFrame)):
+            Y = Y.to_numpy()
+    except ImportError:
+        pass
+
     Y = Y.astype(float)
     X = X.astype(float)
 

pymc_bart/utils.py

@@ -546,7 +546,7 @@ def _prepare_plot_data(
 
     Parameters
     ----------
-    X : PyTensor Variable, Pandas DataFrame or Numpy array
+    X : PyTensor Variable, Pandas DataFrame, Polars DataFrame or Numpy array
         Input data.
     Y : array-like
         Target data.
@@ -585,9 +585,9 @@ def _prepare_plot_data(
     if isinstance(X, Variable):
         X = X.eval()
 
-    if hasattr(X, "columns") and hasattr(X, "values"):
+    if hasattr(X, "columns") and hasattr(X, "to_numpy"):
         x_names = list(X.columns)
-        X = X.values
+        X = X.to_numpy()
     else:
         x_names = []
 
@@ -750,9 +750,9 @@ def plot_variable_importance(  # noqa: PLR0915
     else:
         shape = bartrv.eval().shape[0]
 
-    if hasattr(X, "columns") and hasattr(X, "values"):
+    if hasattr(X, "columns") and hasattr(X, "to_numpy"):
         labels = X.columns
-        X = X.values
+        X = X.to_numpy()
 
     n_vars = X.shape[1]