Repository: arviz-devs/arviz

Issue:
InferenceData conversion does not work with multivariate observation that contains missing value
**To Reproduce**
```python
import numpy as np
import pymc3 as pm

data = np.ma.masked_values([[1, 2], [-1, 4], [2, -1], [-1, -1]], value=-1)

with pm.Model():
    mu = pm.Normal('mu', 0, 1, shape=2)
    sd_dist = pm.HalfCauchy.dist(beta=2.5)
    chol, corr, sigmas = pm.LKJCholeskyCov(
        'chol_cov', n=2, eta=1, sd_dist=sd_dist, compute_corr=True)
    
    obs = pm.MvNormal('obs', mu=mu, chol=chol, observed=data)
    trace0_ = pm.sample()
```

returns:
```python
---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
<ipython-input-82-a388a356e068> in <module>
     11 
     12     obs = pm.MvNormal('obs', mu=mu, chol=chol, observed=data)
---> 13     trace0_ = pm.sample()

~/miniconda3/lib/python3.7/site-packages/pymc3/sampling.py in sample(draws, step, init, n_init, start, trace, chain_idx, chains, cores, tune, progressbar, model, random_seed, discard_tuned_samples, compute_convergence_checks, callback, return_inferencedata, idata_kwargs, mp_ctx, pickle_backend, **kwargs)
    611         if idata_kwargs:
    612             ikwargs.update(idata_kwargs)
--> 613         idata = arviz.from_pymc3(trace, **ikwargs)
    614 
    615     if compute_convergence_checks:

~/miniconda3/lib/python3.7/site-packages/arviz/data/io_pymc3.py in from_pymc3(trace, prior, posterior_predictive, log_likelihood, coords, dims, model, save_warmup)
    529         dims=dims,
    530         model=model,
--> 531         save_warmup=save_warmup,
    532     ).to_inference_data()
    533 

~/miniconda3/lib/python3.7/site-packages/arviz/data/io_pymc3.py in to_inference_data(self)
    461             "posterior": self.posterior_to_xarray(),
    462             "sample_stats": self.sample_stats_to_xarray(),
--> 463             "log_likelihood": self.log_likelihood_to_xarray(),
    464             "posterior_predictive": self.posterior_predictive_to_xarray(),
    465             "predictions": self.predictions_to_xarray(),

~/miniconda3/lib/python3.7/site-packages/arviz/data/base.py in wrapped(cls, *args, **kwargs)
     35                 if all([getattr(cls, prop_i) is None for prop_i in prop]):
     36                     return None
---> 37             return func(cls, *args, **kwargs)
     38 
     39         return wrapped

~/miniconda3/lib/python3.7/site-packages/arviz/data/base.py in wrapped(cls, *args, **kwargs)
     35                 if all([getattr(cls, prop_i) is None for prop_i in prop]):
     36                     return None
---> 37             return func(cls, *args, **kwargs)
     38 
     39         return wrapped

~/miniconda3/lib/python3.7/site-packages/arviz/data/io_pymc3.py in log_likelihood_to_xarray(self)
    303         if self.posterior_trace:
    304             try:
--> 305                 data = self._extract_log_likelihood(self.posterior_trace)
    306             except TypeError:
    307                 warnings.warn(warn_msg)

~/miniconda3/lib/python3.7/site-packages/arviz/data/base.py in wrapped(cls, *args, **kwargs)
     35                 if all([getattr(cls, prop_i) is None for prop_i in prop]):
     36                     return None
---> 37             return func(cls, *args, **kwargs)
     38 
     39         return wrapped

~/miniconda3/lib/python3.7/site-packages/arviz/data/base.py in wrapped(cls, *args, **kwargs)
     35                 if all([getattr(cls, prop_i) is None for prop_i in prop]):
     36                     return None
---> 37             return func(cls, *args, **kwargs)
     38 
     39         return wrapped

~/miniconda3/lib/python3.7/site-packages/arviz/data/io_pymc3.py in _extract_log_likelihood(self, trace)
    227                 log_like_chain = [
    228                     self.log_likelihood_vals_point(point, var, log_like_fun)
--> 229                     for point in trace.points([chain])
    230                 ]
    231                 log_likelihood_dict.insert(var.name, np.stack(log_like_chain), chain)

~/miniconda3/lib/python3.7/site-packages/arviz/data/io_pymc3.py in <listcomp>(.0)
    227                 log_like_chain = [
    228                     self.log_likelihood_vals_point(point, var, log_like_fun)
--> 229                     for point in trace.points([chain])
    230                 ]
    231                 log_likelihood_dict.insert(var.name, np.stack(log_like_chain), chain)

~/miniconda3/lib/python3.7/site-packages/arviz/data/io_pymc3.py in log_likelihood_vals_point(self, point, var, log_like_fun)
    197         log_like_val = utils.one_de(log_like_fun(point))
    198         if var.missing_values:
--> 199             log_like_val = np.where(var.observations.mask, np.nan, log_like_val)
    200         return log_like_val
    201 

<__array_function__ internals> in where(*args, **kwargs)

ValueError: operands could not be broadcast together with shapes (4,2) () (4,) 
```


Discussion:
Yes, we really should support this.

Any ideas for a good interface for this?
Would this be related to sparse matrices?

Resolving patch:
diff --git a/arviz/data/io_pymc3.py b/arviz/data/io_pymc3.py
--- a/arviz/data/io_pymc3.py
+++ b/arviz/data/io_pymc3.py
@@ -198,7 +198,10 @@ def log_likelihood_vals_point(self, point, var, log_like_fun):
         """Compute log likelihood for each observed point."""
         log_like_val = utils.one_de(log_like_fun(point))
         if var.missing_values:
-            log_like_val = np.where(var.observations.mask, np.nan, log_like_val)
+            mask = var.observations.mask
+            if np.ndim(mask) > np.ndim(log_like_val):
+                mask = np.any(mask, axis=-1)
+            log_like_val = np.where(mask, np.nan, log_like_val)
         return log_like_val
 
     @requires("trace")


Test patch:
diff --git a/arviz/tests/external_tests/test_data_cmdstanpy.py b/arviz/tests/external_tests/test_data_cmdstanpy.py
--- a/arviz/tests/external_tests/test_data_cmdstanpy.py
+++ b/arviz/tests/external_tests/test_data_cmdstanpy.py
@@ -51,8 +51,8 @@ class Data:
             runset_obj = RunSet(args)
             runset_obj._csv_files = filepaths  # pylint: disable=protected-access
             obj = CmdStanMCMC(runset_obj)
-            obj._validate_csv_files()  # pylint: disable=protected-access
-            obj._assemble_sample()  # pylint: disable=protected-access
+            obj.validate_csv_files()  # pylint: disable=protected-access
+            obj._assemble_draws()  # pylint: disable=protected-access
 
         return Data
 
diff --git a/arviz/tests/external_tests/test_data_pymc.py b/arviz/tests/external_tests/test_data_pymc.py
--- a/arviz/tests/external_tests/test_data_pymc.py
+++ b/arviz/tests/external_tests/test_data_pymc.py
@@ -288,6 +288,29 @@ def test_missing_data_model(self):
         fails = check_multiple_attrs(test_dict, inference_data)
         assert not fails
 
+    def test_mv_missing_data_model(self):
+        data = ma.masked_values([[1, 2], [2, 2], [-1, 4], [2, -1], [-1, -1]], value=-1)
+
+        model = pm.Model()
+        with model:
+            mu = pm.Normal("mu", 0, 1, shape=2)
+            sd_dist = pm.HalfNormal.dist(1.0)
+            chol, *_ = pm.LKJCholeskyCov("chol_cov", n=2, eta=1, sd_dist=sd_dist, compute_corr=True)
+            pm.MvNormal("y", mu=mu, chol=chol, observed=data)
+            trace = pm.sample(100, chains=2)
+
+        # make sure that data is really missing
+        (y_missing,) = model.missing_values
+        assert y_missing.tag.test_value.shape == (4,)
+        inference_data = from_pymc3(trace=trace, model=model)
+        test_dict = {
+            "posterior": ["mu", "chol_cov"],
+            "observed_data": ["y"],
+            "log_likelihood": ["y"],
+        }
+        fails = check_multiple_attrs(test_dict, inference_data)
+        assert not fails
+
     @pytest.mark.parametrize("log_likelihood", [True, False, ["y1"]])
     def test_multiple_observed_rv(self, log_likelihood):
         y1_data = np.random.randn(10)