[feature request] GPLVMs with missing data

I’m interested in using the pyro.contrib.gp.models.GPLVM model in the case where some points in Y (the observed data) are missing at random.

GPy has an example here.

This functionality isn’t available in pyro at the moment (I think - couldn’t find examples for this). Is this something that can be added? I’d be happy to contribute some code to do this.

[InfProbSciX]

This worked! This is the class for reference:

class MaskedGaussian(gp.likelihoods.Gaussian):
    def forward(self, f_loc, f_var, y=None):
        y_var = f_var + self.variance
        y_dist = dist.Normal(f_loc, y_var.sqrt())

        if y is not None:
            if y.isnan().any():
                y_dist = dist.MaskedDistribution(y_dist, ~y.isnan())
                y = torch.masked_fill(y, y.isnan(), -999.)

            y_dist = y_dist.expand_by(y.shape[:-f_loc.dim()]).to_event(y.dim())
        return pyro.sample(self._pyro_get_fullname("y"), y_dist, obs=y)

I think I’m slightly in favour of changing the existing Gaussian likelihood class actually, and I can add an optional argument to the __init__ function like mask_missing_data=False as there’s a lot of repeated code here. What would you say?

Also, another quick comment, would it be possible to use this likelihood with the (Sparse)GPRegression classes? I ask because likelihood doesn’t seem to be an optional argument for those classes, so perhaps it’d be a bit more work to get them to work under the missing data regime.

[fehiepsi]

@InfProbSciX I think you will need to also do y = torch.masked_fill(y, y.isnan(), -999.) if y has NaN value. This note explains well why NaN happens during backward.

By the way, the MaskedGaussian interface looks great to me. 👍