Create extract_obs_data function

Author: brandonwillard

pymc3/aesaraf.py

@@ -19,10 +19,12 @@
 from aesara import scalar
 from aesara import tensor as aet
 from aesara.gradient import grad
-from aesara.graph.basic import Apply, graph_inputs
+from aesara.graph.basic import Apply, Constant, graph_inputs
 from aesara.graph.op import Op
 from aesara.sandbox.rng_mrg import MRG_RandomStream as RandomStream
 from aesara.tensor.elemwise import Elemwise
+from aesara.tensor.sharedvar import SharedVariable
+from aesara.tensor.subtensor import AdvancedIncSubtensor, AdvancedIncSubtensor1
 from aesara.tensor.var import TensorVariable
 
 from pymc3.data import GeneratorAdapter
@@ -48,6 +50,28 @@
 ]
 
 
+def extract_obs_data(x: TensorVariable) -> np.ndarray:
+    """Extract data observed symbolic variables.
+
+    Raises
+    ------
+    TypeError
+
+    """
+    if isinstance(x, Constant):
+        return x.data
+    if isinstance(x, SharedVariable):
+        return x.get_value()
+    if x.owner and isinstance(x.owner.op, (AdvancedIncSubtensor, AdvancedIncSubtensor1)):
+        array_data = extract_obs_data(x.owner.inputs[0])
+        mask_idx = tuple(extract_obs_data(i) for i in x.owner.inputs[2:])
+        mask = np.zeros_like(array_data)
+        mask[mask_idx] = 1
+        return np.ma.MaskedArray(array_data, mask)
+
+    raise TypeError(f"Data cannot be extracted from {x}")
+
+
 def inputvars(a):
     """
     Get the inputs into a aesara variables

pymc3/tests/test_aesaraf.py

@@ -19,9 +19,10 @@
 import numpy as np
 import pytest
 
+from aesara.tensor.subtensor import AdvancedIncSubtensor, AdvancedIncSubtensor1
 from aesara.tensor.type import TensorType
 
-from pymc3.aesaraf import _conversion_map, take_along_axis
+from pymc3.aesaraf import _conversion_map, extract_obs_data, take_along_axis
 from pymc3.vartypes import int_types
 
 FLOATX = str(aesara.config.floatX)
@@ -225,3 +226,49 @@ def test_dtype_failure(self):
         indices.tag.test_value = np.zeros((1,) * indices.ndim, dtype=FLOATX)
         with pytest.raises(IndexError):
             take_along_axis(arr, indices)
+
+
+def test_extract_obs_data():
+
+    with pytest.raises(TypeError):
+        extract_obs_data(aet.matrix())
+
+    data = np.random.normal(size=(2, 3))
+    data_at = aet.as_tensor(data)
+    mask = np.random.binomial(1, 0.5, size=(2, 3)).astype(bool)
+
+    for val_at in (data_at, aesara.shared(data)):
+        res = extract_obs_data(val_at)
+
+        assert isinstance(res, np.ndarray)
+        assert np.array_equal(res, data)
+
+    # AdvancedIncSubtensor check
+    data_m = np.ma.MaskedArray(data, mask)
+    missing_values = data_at.type()[mask]
+    constant = aet.as_tensor(data_m.filled())
+    z_at = aet.set_subtensor(constant[mask.nonzero()], missing_values)
+
+    assert isinstance(z_at.owner.op, AdvancedIncSubtensor)
+
+    res = extract_obs_data(z_at)
+
+    assert isinstance(res, np.ndarray)
+    assert np.ma.allequal(res, data_m)
+
+    # AdvancedIncSubtensor1 check
+    data = np.random.normal(size=(3,))
+    data_at = aet.as_tensor(data)
+    mask = np.random.binomial(1, 0.5, size=(3,)).astype(bool)
+
+    data_m = np.ma.MaskedArray(data, mask)
+    missing_values = data_at.type()[mask]
+    constant = aet.as_tensor(data_m.filled())
+    z_at = aet.set_subtensor(constant[mask.nonzero()], missing_values)
+
+    assert isinstance(z_at.owner.op, AdvancedIncSubtensor1)
+
+    res = extract_obs_data(z_at)
+
+    assert isinstance(res, np.ndarray)
+    assert np.ma.allequal(res, data_m)