let pandas_to_array take pandas Index

Author: MarcoGorelli

pymc3/aesaraf.py

@@ -86,7 +86,13 @@ def pandas_to_array(data):
     if hasattr(data, "to_numpy") and hasattr(data, "isnull"):
         # typically, but not limited to pandas objects
         vals = data.to_numpy()
-        mask = data.isnull().to_numpy()
+        null_data = data.isnull()
+        if hasattr(null_data, "to_numpy"):
+            # pandas Series
+            mask = null_data.to_numpy()
+        else:
+            # pandas Index
+            mask = null_data
         if mask.any():
             # there are missing values
             ret = np.ma.MaskedArray(vals, mask)

pymc3/tests/test_aesaraf.py

@@ -430,6 +430,13 @@ def test_pandas_to_array(input_dtype):
     assert isinstance(wrapped, TensorVariable)
 
 
+def test_pandas_to_array_pandas_index():
+    data = pd.Index([1, 2, 3])
+    result = pandas_to_array(data)
+    expected = np.array([1, 2, 3])
+    np.testing.assert_array_equal(result, expected)
+
+
 def test_walk_model():
     d = at.vector("d")
     b = at.vector("b")