Add _as_tensor_variable converter for pandas objects (#5920)

* Add pd.Series and pd.DataFrame support
* Explicitly add pandas into requirements.txt
* Add pd global import in tests

Co-authored-by: Michael Osthege <[email protected]>

Author: Juan Orduz

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

@@ -18,6 +18,7 @@
 import numpy as np
 import numpy.ma as ma
 import numpy.testing as npt
+import pandas as pd
 import pytest
 import scipy.sparse as sps
 
@@ -47,6 +48,40 @@
 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:
+    df = 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:
+    df = 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:
+
+    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)
@@ -224,7 +259,6 @@ def test_convert_observed_data(input_dtype):
     Ensure that convert_observed_data returns the dense array, masked array,
     graph variable, TensorVariable, or sparse matrix as appropriate.
     """
-    pd = pytest.importorskip("pandas")
     # Create the various inputs to the function
     sparse_input = sps.csr_matrix(np.eye(3)).astype(input_dtype)
     dense_input = np.arange(9).reshape((3, 3)).astype(input_dtype)
@@ -300,7 +334,6 @@ def test_convert_observed_data(input_dtype):
 
 
 def test_pandas_to_array_pandas_index():
-    pd = pytest.importorskip("pandas")
     data = pd.Index([1, 2, 3])
     result = convert_observed_data(data)
     expected = np.array([1, 2, 3])

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