Add tutorial for models with missing data

[fehiepsi]

We should able to run MCMC for models with missing data in NumPyro because NumPyro supports improper priors through param primitive. This is also a good chance to illustrate that NumPyro supports improper priors.

Related issues in other repos:

• Pyro: Support masked sample statement for batched imputation pyro#1106, add tutorial/example with missing data pyro#1676
• brmp: Support for model-based imputation of partially unobserved response variables brmp#43, Support for fully unobserved response variables brmp#44

[neerajprad]

I think the issue in Pyro refers to allowing a way to support models with sample statements of the kind pyro.sample('obs', dist(..), obs=[some partially observed tensor]), where we condition on the observed data and impute the missing entries. Could you clarify how having param statements in MCMC help this use case?

[fehiepsi]

I will use the most naive way to impute those missing data. For example,

def model(x, y):
    loc = sample('loc', Normal(0, 1))
    scale = sample('scale', LogNormal(0, 1))
    isnan = onp.isnan(x)
    x_impute = param('x_impute', np.zeros(isnan.sum()))
    x = ops.index_update(x, onp.nonzero(isnan), x_impute)  # update x
    sample('x', dist.Normal(loc, scale), obs=x)
    ...

I am not sure how to make this possible with pyro.sample... Does the above script do what pyro/brmp issues mentioned?

[neerajprad]

I see - I think your method will give us the ML estimate for the missing values.

I am not sure how to make this possible with pyro.sample.

I am not quite sure, but maybe sampling from the dist.Normal(..) might work:

    x_impute = sample('x_impute', dist.Normal(loc, scale))
    x = ops.index_update(x, onp.nonzero(isnan), x_impute)  # update x
    sample('x', dist.Normal(loc, scale), obs=x)

In any case, I think this will make for a good tutorial.

[fehiepsi]

the ML estimate for the missing values

Actually, those imputed values will have priors dist.Normal(loc, scale) because the statement

sample('x', dist.Normal(loc, scale), obs=x)

acts on the "merged" x (x_notnan & x_impute). Here I am using the fact that in NumPyro (not Pyro) MCMC

x = param('x', ...)
sample('x', dist.Normal(), obs=x)

and

x = sample('x', dist.Normal())

are equivalent (even for priors with constrained supports). This way also works for multivariate priors IIUC.

The script in your comment will double log_prob of x_impute, so I think that it is not quite correct. Maybe I am overlooking something?

[neerajprad]

Interesting. No, you are absolutely right about the double counting. What I wrote above is incorrect unless we are doing masked observes. Actually, your original solution is quite clever! This will make for an interesting tutorial, look forward to it.