Spatiotemporal variational demo model always emits mean 0

[sethaxen]

Describe the bug

With SpatioTemporalSparseCVI, after training, space_time_predict_f emits only zeros for the mean in the demo.

To reproduce

1. Run the complete spatiotemporal variational example
2. Call the following:

>>> mu_f, var_f = model.space_time_predict_f(X_grid)
>>> mu_f
<tf.Tensor: shape=(2500, 1), dtype=float64, numpy=
array([[0.],
       [0.],
       [0.],
       ...,
       [0.],
       [0.],
       [0.]])>
>>> 
>>> np.max(np.abs(mu_f.numpy()))
0.0

Expected behaviour

It is surprising that the predicted mean is zero everywhere. This is visible in the plots in the notebook, where the prediction looks nothing like the data even after training: https://secondmind-labs.github.io/markovflow/notebooks/spatio_temporal.html?highlight=spatiotemporalsparsecvi:

I expected the mean to look like it had learned something from the data.

System information

• OS: Pop!_OS 22.04 LTS
• Python version: 3.7
• Markovflow version: develop branch
• TensorFlow version: 2.4.3
• GPflow version: 2.2.1

Additional context
If I modify the demo to instead use SpatioTemporalSparseVariational, space_time_predict_f does not predict all zeros. After fitting for many more iterations I get a reasonable-looking result:

It appears that 7323200 may have been the commit that broke SpatioTemporalSparseCVI If we check out the previous commit (fc82c0a), this is the plot at the end of the notebook:

I suspect the culprit is these lines:

markovflow/markovflow/models/spatio_temporal_variational.py

Lines 433 to 437 in 67e80c8

	self._posterior = ConditionalProcess(
	posterior_dist=self.dist_q,
	kernel=self.kernel,
	conditioning_time_points=self.inducing_time,
	)

Since dist_q is a property that is generated each time it is called, perhaps declaring _posterior in the constructor makes it out-of-sync with the trained parameters?

If we just convert these lines to

    @property
    def posterior(self) -> PosteriorProcess:
        return ConditionalProcess(
            posterior_dist=self.dist_q,
            kernel=self.kernel,
            conditioning_time_points=self.inducing_time,
        )

then we again get a non-zero mean with predictions that look like the data:

[sethaxen]

@vincentadam87 I can open a PR with the proposed fix and an integration test that after training the predicted mean is non-zero, but I would not know where to begin to write unit tests for the spatiotemporal models. How do you recommend proceeding?

[vincentadam87]

I assume most objects / functions you use here are already tested so an integration test would be enough.

One test that would be cool:
spatiotemp vs gpflow.gpr for a product kernel

If you have a factorial 2d data set
coordinate 1 (time): t1, t2
coordinate 2: x1, x2
and you have 4 data points in a factorial grid : (t1, x1) (t1, x2) (t2, x1) (t2, x2)

For the sparse spatio temp cvi, gaussian noise, choosing z_s = [x1,x2] and z_t = [t1,t2]
after inference, you should get the same posterior and loss (elbo=likelihood)

most of our tests are similar : in a simple setting, CVI boils down to GPR (classic GP regression)
It would be awesome if you could do that, but no worries otherwise :)