Add _as_tensor_variable converter for pandas objects

pymc/aesaraf.py

@@ -29,6 +29,7 @@
 import aesara
 import aesara.tensor as at
 import numpy as np
+import pandas as pd
 import scipy.sparse as sps
 
 from aeppl.abstract import MeasurableVariable
@@ -50,6 +51,7 @@
 from aesara.graph.op import Op, compute_test_value
 from aesara.sandbox.rng_mrg import MRG_RandomStream as RandomStream
 from aesara.scalar.basic import Cast
+from aesara.tensor.basic import _as_tensor_variable
 from aesara.tensor.elemwise import Elemwise
 from aesara.tensor.random.op import RandomVariable
 from aesara.tensor.random.var import (
@@ -142,6 +144,12 @@ def convert_observed_data(data):
         return floatX(ret)
 
 
+@_as_tensor_variable.register(pd.Series)
+@_as_tensor_variable.register(pd.DataFrame)
+def dataframe_to_tensor_variable(df: pd.DataFrame, *args, **kwargs) -> TensorVariable:
+    return at.as_tensor_variable(df.to_numpy(), *args, **kwargs)
+
+
 def change_rv_size(
     rv: TensorVariable,
     new_size: PotentialShapeType,

pymc/tests/test_aesaraf.py

@@ -47,6 +47,43 @@
 from pymc.vartypes import int_types
 
 
+@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_pd_dataframe_as_tensor_variable(np_array: np.ndarray) -> None:
+    pd = pytest.importorskip("pandas")
+    df: pd.DataFrame = pd.DataFrame(np_array)
+    np.testing.assert_array_equal(x=at.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.ones(shape=4), np.zeros(shape=10), [1, 2, 3, 4]],
+)
+def test_pd_series_as_tensor_variable(np_array: np.ndarray) -> None:
+    pd = pytest.importorskip("pandas")
+    df: pd.DataFrame = pd.Series(np_array)
+    np.testing.assert_array_equal(x=at.as_tensor_variable(x=df).eval(), y=np_array)
+
+
+def test_pd_as_tensor_variable_multiindex() -> None:
+    pd = pytest.importorskip("pandas")
+
+    tuples = [("L", "Q"), ("L", "I"), ("O", "L"), ("O", "I")]
+
+    index = pd.MultiIndex.from_tuples(tuples, names=["Id1", "Id2"])
+
+    df = pd.DataFrame({"A": [12.0, 80.0, 30.0, 20.0], "B": [120.0, 700.0, 30.0, 20.0]}, index=index)
+    np_array = np.array([[12.0, 80.0, 30.0, 20.0], [120.0, 700.0, 30.0, 20.0]]).T
+    assert isinstance(df.index, pd.MultiIndex)
+    np.testing.assert_array_equal(x=at.as_tensor_variable(x=df).eval(), y=np_array)
+
+
 def test_change_rv_size():
     loc = at.as_tensor_variable([1, 2])
     rv = normal(loc=loc)

requirements.txt

@@ -5,5 +5,6 @@ cachetools>=4.2.1
 cloudpickle
 fastprogress>=0.2.0
 numpy>=1.15.0
+pandas>=0.24.0
 scipy>=1.4.1
 typing-extensions>=3.7.4