# Possible ways to speed up spatio-temporal GP modeling

**URL:** <https://discourse.pymc.io/t/possible-ways-to-speed-up-spatio-temporal-gp-modeling/13127>\
**Category:** v5\
**Tags:** modeling\
**Created:** [October 18, 2023, 5:22pm UTC](https://discourse.pymc.io/t/possible-ways-to-speed-up-spatio-temporal-gp-modeling/13127 "2023-10-18T17:22:10Z")\
**Posts on this page:** 17\
**Page:** 1

<div class="post-metadata">

**Author:** ![seunjeong](https://avatars.discourse-cdn.com/v4/letter/s/ecae2f/32.png) [@seunjeong](https://discourse.pymc.io/u/seunjeong)\
**Post date:** [October 18, 2023, 5:22pm UTC](https://discourse.pymc.io/t/possible-ways-to-speed-up-spatio-temporal-gp-modeling/13127/1 "2023-10-18T17:22:10Z")

</div>

Hi All,

I’ve been working on adding more features to my PyMC model for a couple of months now. I feel that I’m making progress, especially with the advanced features available in PyMC.

I’m now modeling spatio-temporal satellite data, and the model is shown below. While the model seems to generate the expected posterior distribution, it runs very slowly.

Interestingly, without optimizing the length scale (i.e., when using a fixed value for the length scale), the JAX-enabled Gaussian process finishes in about 10 minutes. However, once I optimize the length scale parameters (for both space and time), the processing time shoots up to more than 1 day on a CPU and over 8 hours on Colab’s single T4 GPU with JAX enabled.

Given that my total number of observations is around 700, the long processing time isn’t entirely unexpected. However, the duration seems excessive, even for a JAX-based simulation. Using multiple GPUs only reduces the time to about 4 hours (i.e., half of the 8 hours) if I use 1500 samples (with 1000 warm-up samples) for each chain.

I would greatly appreciate any advice on improving the speed of the code.

Thank you.

* * *

```auto
def ST_model ():
    with pm.Model() as model:

        normal_dist = pm.Normal.dist(mu=mu_lambda, sigma=sigma_lambda)
        Lambda = pm.Truncated("Lambda", normal_dist, shape=m, lower=0, upper=10)

        #===============================================================================
        # Mean function
        #===============================================================================
        mu = LinearMean(K = K, Lambda = Lambda)

        #===============================================================================
        # Kernel variance
        #===============================================================================
        var = pm.Exponential("kernel_var", 2)

        #===============================================================================
        # Kernel length
        #===============================================================================
        length_T = pm.Exponential("kernel_length_T", 1/(7/(31*NUM_MONTHS)))
        length_S = pm.Gamma("kernel_length_S", 0.5, 0.2)
        
        #===============================================================================
        # Noise
        #===============================================================================
        noise = pm.HalfCauchy ("noise", 1, shape = 1)
        
        #===============================================================================
        # Covariance function
        # Note we are not using kronecker product
        #===============================================================================
        # Spatial covariance
        spatial_cov = pm.gp.cov.Matern52(input_dim=3, ls=[length_S, length_S], active_dims=[0,1])
        
        # Temporal covariance function, with kernel variance
        temporal_cov = var * pm.gp.cov.ExpQuad(input_dim=3, ls=length_T, active_dims=[2])

        # Combined spatio-temporal covariance. 
        cov_func = spatial_cov * temporal_cov

        #===============================================================================
        # specify the GP:
        #===============================================================================
        gp = pm.gp.Latent(mean_func=mu, cov_func=cov_func)
        f = gp.prior("f", X=Xs)

        #===============================================================================
        # Likelihood
        #===============================================================================
        Y_ = pm.Normal("Y_", mu=f, sigma = noise, observed=Y)
        
        ################################################################################
        # Sampling
        ################################################################################
        trace = pm.sample_prior_predictive()
        trace.extend (pm.sampling_jax.sample_numpyro_nuts (
                            NUM_SAMPLE, tune=TUNE,\
                            target_accept=0.9, random_seed=RANDOM_SEED,\
                            idata_kwargs = {'log_likelihood': True}, chains=NUM_CHAIN))

```

---

<div class="post-metadata">

**Author:** ![jessegrabowski](https://yyz2.discourse-cdn.com/flex036/user_avatar/discourse.pymc.io/jessegrabowski/32/5010_2.png) [@jessegrabowski](https://discourse.pymc.io/u/jessegrabowski)\
**Post date:** [October 18, 2023, 5:30pm UTC](https://discourse.pymc.io/t/possible-ways-to-speed-up-spatio-temporal-gp-modeling/13127/2 "2023-10-18T17:30:23Z")

</div>

You could try replacing `pm.gp.Latent` with [`pm.gp.HSGP`](https://www.pymc.io/projects/docs/en/stable/api/gp/generated/pymc.gp.HSGP.html). This is an approximate method that is intended to speed up models with GP latent variables. See [this talk](https://www.youtube.com/watch?v=ri5sJAdcYHk) by @bwengals for all the details. There are also good usage examples in the docs I linked.

Edit: Actually the nice examples I was thinking of were in the video, not the docs.

---

<div class="post-metadata">

**Author:** ![seunjeong](https://avatars.discourse-cdn.com/v4/letter/s/ecae2f/32.png) [@seunjeong](https://discourse.pymc.io/u/seunjeong)\
**Post date:** [October 18, 2023, 5:34pm UTC](https://discourse.pymc.io/t/possible-ways-to-speed-up-spatio-temporal-gp-modeling/13127/3 "2023-10-18T17:34:32Z")

</div>

Thank you so much for your prompt and valuable advice. I will look into it and report if I find anything interesting.

Again, thank you so much!

---

<div class="post-metadata">

**Author:** ![seunjeong](https://avatars.discourse-cdn.com/v4/letter/s/ecae2f/32.png) [@seunjeong](https://discourse.pymc.io/u/seunjeong)\
**Post date:** [October 18, 2023, 5:51pm UTC](https://discourse.pymc.io/t/possible-ways-to-speed-up-spatio-temporal-gp-modeling/13127/4 "2023-10-18T17:51:20Z")

</div>

Following the example, I quickly tested HSGP with:

```auto
 gp = pm.gp.HSGP(m=[35, 35, 35], c=4.0, mean_func=mu, cov_func=cov_func)
 f = gp.prior("f", X=Xs)

```

But, I had the following error message:

**“NotImplementedError: The power spectral density of products of covariance functions is not implemented.”**

I guess I am multiplying two covariances here, but other than that, I don’t have a clue. I would appreciate it if any advice is offered.

Thank you.

---

<div class="post-metadata">

**Author:** ![jessegrabowski](https://yyz2.discourse-cdn.com/flex036/user_avatar/discourse.pymc.io/jessegrabowski/32/5010_2.png) [@jessegrabowski](https://discourse.pymc.io/u/jessegrabowski)\
**Post date:** [October 18, 2023, 6:52pm UTC](https://discourse.pymc.io/t/possible-ways-to-speed-up-spatio-temporal-gp-modeling/13127/5 "2023-10-18T18:52:56Z")

</div>

With the caveat that I’m not a GP expert, here are some thoughts:

1. You could resort to a linear combination of the kernels? Perhaps model the log of your data instead of the levels? That’s the usual trick with multiplicative models. I’m not 100% sure if the logic holds in the kernel space.
2. You could make two separate HSGP approximations, one for the spatial and one for the temporal. Extract their HSGP representations (called `phi` and `sqrt_psd`, see 45:56 in the video), multiply each by a set of parameters, then multiply together the results. You’d essentially be doing GP(time) \* GP(space) “by hand”
3. Similar to 2, but with fewer parameters – you could extract those basis features and try to multiplying them together, then defining only a single set of parameters for the resulting product?
4. Nuclear option: see if anyone has come up with a closed-form for the Matern52 \* ExpQuad kernel, and use that directly. My understanding is that you’re _technically_ not allowed to multiply GP kernels, but that it’s almost correct so people do it anyway. Since both kernel functions you’re using are exponential (and popular), it seems at least conceivable that there’s a closed form result, and thus there might be research into a more correct kernel to use?

These are just spitballing. I hope a real expert can weigh in as well.

---

<div class="post-metadata">

**Author:** ![bwengals](https://yyz2.discourse-cdn.com/flex036/user_avatar/discourse.pymc.io/bwengals/32/1237_2.png) [@bwengals](https://discourse.pymc.io/u/bwengals)\
**Post date:** [October 18, 2023, 8:17pm UTC](https://discourse.pymc.io/t/possible-ways-to-speed-up-spatio-temporal-gp-modeling/13127/6 "2023-10-18T20:17:29Z")

</div>

To add to what @jessegrabowski said, the prior on lengthscale has a _huge_ effect on sampling time, like 10 minutes to 1 day huge isn’t unreasonable. But ~700 data points really should be doable IMO. The problem with the `Gamma(0.5, 0.2)` prior is that it puts a lot of mass at very small lengthscales, below the resolution of your data. Say your x data is on a rectangular grid, 1 “unit” apart, as in x[1] - x[0] = 1. Roughly speaking, lengthscale values smaller than 1 (though it’s not a strict boundary) are below the resolution of your data. When this happens the GP starts to turn into a model of IID Gaussian noise – which is already covered by your likelihood.

As a starting point, you could try using an inverse gamma prior. It has a nice property where it has low mass at small values. Using this code, set your `lower` value carefully, it should at least be the resolution of your x data.

```python
params_ls = pm.find_constrained_prior(
    distribution=pm.InverseGamma,
    lower=2,
    upper=10,
    init_guess={"alpha": 2, "beta": 10},
    mass=0.9,
)
d = pm.InverseGamma.dist(**params_ls)
s = pm.draw(d, 10_000)
plt.hist(s, 100);

```

I would double check the distribution at the end because every now and then `pm.find_constrained_prior` will return something that’s not quite right without erroring.

This all applies btw whether or not you’re using HSGP or regular GP. Something to watch out for with HSGPs on spatial / time series models is that it’s not amazing for super small lengthscales, relative to the overall width of your x data. Or in other words for modeling very local relationships between data. To do this you start needing `m` to get really high and you lose the speed gains.

---

<div class="post-metadata">

**Author:** ![seunjeong](https://avatars.discourse-cdn.com/v4/letter/s/ecae2f/32.png) [@seunjeong](https://discourse.pymc.io/u/seunjeong)\
**Post date:** [October 18, 2023, 8:20pm UTC](https://discourse.pymc.io/t/possible-ways-to-speed-up-spatio-temporal-gp-modeling/13127/7 "2023-10-18T20:20:22Z")

</div>

Thanks very much for your insightful suggestions, as always! I just finished watching the video and will try to get help there.

With respect to the multiplication of kernels, I was trying to implement the Kroncker product assuming separable time and space kernels. I could’ve used PyMC’s some Kron functions - I couldn’t find one that is applicable or I didn’t know enough.

The problem is that all observation locations of the satellite do not have data at a particular measurement time. In other words, depending on the measurement time, some locations have missing data. Because of this, I could not use a fixed input of lon, lat, and time, but I had to make a mesh of lon, lat, and time, excluding those missing locations and times. This lon, lat, time mesh (`mn x 3` array) allows me to make an `mn x mn` (big) matrix for both space and time. And then, I just need to do an element-wise multiplication instead of using the Kronecker product. Here, the `active_dims` parameter of PyMC was a game changer for me.

I guess my logic seems correct here, but I’d like to hear advice if I am wrong or not efficient.

By the way, I saw a multiplication example from this PyMC doc (though old): [Kronecker Structured Covariances — PyMC3 3.11.4 documentation](https://www.pymc.io/projects/docs/en/v3.11.4/pymc-examples/examples/gaussian_processes/GP-Kron.html).

Thank you so much!

---

<div class="post-metadata">

**Author:** ![bwengals](https://yyz2.discourse-cdn.com/flex036/user_avatar/discourse.pymc.io/bwengals/32/1237_2.png) [@bwengals](https://discourse.pymc.io/u/bwengals)\
**Post date:** [October 18, 2023, 8:25pm UTC](https://discourse.pymc.io/t/possible-ways-to-speed-up-spatio-temporal-gp-modeling/13127/8 "2023-10-18T20:25:16Z")

</div>

The Kronecker stuff sounds like an option then!

> The problem is that all observation locations of the satellite do not have data at a particular measurement time. In other words, depending on the measurement time, some locations have missing data.

How many are missing? If you use `LatentKron` (I think this will work for `MarginalKron` too but I’d have to try it) you can still pass a full grid of `X_spatial` and `X_time`. Then index the GP where you actually have observed data and feed that into the likelihood and then do NUTS sampling. You’ll get the result of the GP interpolating the missing spatial/time points “for free”.

---

<div class="post-metadata">

**Author:** ![jessegrabowski](https://yyz2.discourse-cdn.com/flex036/user_avatar/discourse.pymc.io/jessegrabowski/32/5010_2.png) [@jessegrabowski](https://discourse.pymc.io/u/jessegrabowski)\
**Post date:** [October 18, 2023, 8:50pm UTC](https://discourse.pymc.io/t/possible-ways-to-speed-up-spatio-temporal-gp-modeling/13127/9 "2023-10-18T20:50:23Z")

</div>

Totally orthogonal to the discussion of GPs, but if the number of raster cells in your data is smaller than the number of timesteps, you could also consider writing the model as a statespace system. I’m thinking basically a diagonal VAR(p) with spatially correlated errors. This handles missing data very naturally, and might be a more convenient way to represent things. The wrinkle is that the models we have in PyMC at the moment are just slow for very large state vectors (and yours probably would be – in your case n \times p, where n is the number of rasters and p is the number of lags in the VAR).

---

<div class="post-metadata">

**Author:** ![seunjeong](https://avatars.discourse-cdn.com/v4/letter/s/ecae2f/32.png) [@seunjeong](https://discourse.pymc.io/u/seunjeong)\
**Post date:** [October 18, 2023, 9:33pm UTC](https://discourse.pymc.io/t/possible-ways-to-speed-up-spatio-temporal-gp-modeling/13127/10 "2023-10-18T21:33:46Z")

</div>

That sounds amazing! Although I haven’t calculated, I think the number of missing data points for a given month is \> ~90% in the full space-time grid. That’s because the satellite returns to the same location only every few days, and even if it comes back, it passes over a section (i.e., swath) of the land. Then there is a quality issue, so we toss out lots of data points.

If there is a simple toy example for this full grid + indexing, that would be extremely helpful.

Thank you so much!

---

<div class="post-metadata">

**Author:** ![seunjeong](https://avatars.discourse-cdn.com/v4/letter/s/ecae2f/32.png) [@seunjeong](https://discourse.pymc.io/u/seunjeong)\
**Post date:** [October 18, 2023, 9:43pm UTC](https://discourse.pymc.io/t/possible-ways-to-speed-up-spatio-temporal-gp-modeling/13127/11 "2023-10-18T21:43:42Z")

</div>

Thanks for the suggestion! By the way, I saw the new state space (SS) model implementation. I’ve used state space models in R for a while, primarily simple random walk models. I was delighted to see your developing SS models within the PyMC framework.

I guess my raster cells are too many, although someone in my group is testing spatial aggregation before doing any GP or analytical multivariate normal modeling. But your idea of a diagonal VAR sounds very interesting.

I will keep paying attention to your SS models.

Thank you very much!

---

<div class="post-metadata">

**Author:** ![bwengals](https://yyz2.discourse-cdn.com/flex036/user_avatar/discourse.pymc.io/bwengals/32/1237_2.png) [@bwengals](https://discourse.pymc.io/u/bwengals)\
**Post date:** [October 18, 2023, 9:47pm UTC](https://discourse.pymc.io/t/possible-ways-to-speed-up-spatio-temporal-gp-modeling/13127/12 "2023-10-18T21:47:54Z")

</div>

> If there is a simple toy example for this full grid + indexing, that would be extremely helpful.

Sure. Slight modification of [this example](https://www.pymc.io/projects/examples/en/latest/gaussian_processes/GP-Kron.html#example-1). I removed all the comments except for what’s relevant for this discussion.

```python
import arviz as az
import matplotlib as mpl
import numpy as np
import pymc as pm
plt = mpl.pyplot

RANDOM_SEED = 12345
rng = np.random.default_rng(RANDOM_SEED)

n1, n2 = (50, 30)
x1 = np.linspace(0, 5, n1)
x2 = np.linspace(0, 3, n2)

X = pm.math.cartesian(x1[:, None], x2[:, None])

l1_true = 0.8
l2_true = 1.0
eta_true = 1.0

cov = (
    eta_true**2
    * pm.gp.cov.Matern52(2, l1_true, active_dims=[0])
    * pm.gp.cov.Cosine(2, ls=l2_true, active_dims=[1])
)

K = cov(X).eval()
f_true = rng.multivariate_normal(np.zeros(X.shape[0]), K, 1).flatten()

sigma_true = 0.25
y_true = f_true + sigma_true * rng.standard_normal(X.shape[0])

# indices of the y data we're keeping
keep_percentage = 0.5
n_keep = int(np.round(keep_percentage * len(y_true)))
ix = np.arange(len(y_true))
observed_ix = np.random.choice(ix, size=n_keep, replace=False)
y = y_true[observed_ix] # now your observed data is a selection of y_true, and you know the indices

with pm.Model() as model:
    ls1 = pm.Gamma("ls1", alpha=2, beta=2)
    ls2 = pm.Gamma("ls2", alpha=2, beta=2)
    eta = pm.HalfNormal("eta", sigma=2)

    cov_x1 = pm.gp.cov.Matern52(1, ls=ls1)
    cov_x2 = eta**2 * pm.gp.cov.Cosine(1, ls=ls2)

    sigma = pm.HalfNormal("sigma", sigma=2)

    gp = pm.gp.LatentKron(cov_funcs=[cov_x1, cov_x2])
    f = gp.prior("f", Xs=Xs)

    # This line is the only one that's different. Index the GP where you observed data
    f_obs = f[observed_ix]
    y_ = pm.Normal("y_", mu=f_obs, sigma=sigma, observed=y)

```

Now your trace of `f` will cover the full grid of your data. It will give you the same posterior of `f` as if you used the GP to “predict” what the `f` values would be there.

If you have 700 total data points, how big is your spatial grid, and how many points in time?

---

<div class="post-metadata">

**Author:** ![seunjeong](https://avatars.discourse-cdn.com/v4/letter/s/ecae2f/32.png) [@seunjeong](https://discourse.pymc.io/u/seunjeong)\
**Post date:** [October 18, 2023, 10:05pm UTC](https://discourse.pymc.io/t/possible-ways-to-speed-up-spatio-temporal-gp-modeling/13127/13 "2023-10-18T22:05:52Z")

</div>

Thanks very much for this piece of code! I had not realized the impact of the length scale on the speed (or maybe also on model performance). Yes, I will try this.

Thank you!

---

<div class="post-metadata">

**Author:** ![seunjeong](https://avatars.discourse-cdn.com/v4/letter/s/ecae2f/32.png) [@seunjeong](https://discourse.pymc.io/u/seunjeong)\
**Post date:** [October 18, 2023, 10:19pm UTC](https://discourse.pymc.io/t/possible-ways-to-speed-up-spatio-temporal-gp-modeling/13127/14 "2023-10-18T22:19:50Z")

</div>

Wow, thank you so much! I am getting a lot of help today. I really appreciate your time!

I love your statement:

> It will give you the same posterior of `f` as if you used the GP to “predict” what the `f` values would be there.

It is really great that we do get predictions while doing inference!

For the time, we are doing hourly modeling by month for the entire 2020. So, for a month, the full time index will be a size of 24\*30.

For the exact spatial domain boundary, I’d have to ask my postdoc, but it is roughly the Southern California area (LA County + suburban areas). The satellite resolution is ~ 2 km. My postdoc is using an analytical solution, and, in parallel, I am now into PyMC GP!

Thank you!

---

<div class="post-metadata">

**Author:** ![seunjeong](https://avatars.discourse-cdn.com/v4/letter/s/ecae2f/32.png) [@seunjeong](https://discourse.pymc.io/u/seunjeong)\
**Post date:** [October 19, 2023, 1:02am UTC](https://discourse.pymc.io/t/possible-ways-to-speed-up-spatio-temporal-gp-modeling/13127/15 "2023-10-19T01:02:03Z")

</div>

Hi Bill,

Following your example, I tried to use

```auto
gp = pm.gp.LatentKron(cov_funcs=[cov_x1, cov_x2])
f = gp.prior("f", Xs=Xs)

```

But I have difficulty in setting up `Xs.`

Following [Kronecker Structured Covariances — PyMC3 3.11.4 documentation](https://www.pymc.io/projects/docs/en/v3.11.4/pymc-examples/examples/gaussian_processes/GP-Kron.html), which seems to be the example you are using, I set up `Xs` as:

```auto
Xs = [lons[:, None], lats[:, None]], dt_full_norm[:, None]]

```

where lons and lats are unique longitudes and latitudes, which will make a full grid if I use `np.meshgrid`. For time, `dt_full_norm` is the entire time stamp.

Then I applied as:

```auto
spatial_cov = pm.gp.cov.Matern52(input_dim=3, ls=[length_S, length_S], active_dims=[0,1])
temporal_cov = var*pm.gp.cov.ExpQuad(input_dim=3, ls=length_T, active_dims=[2])

gp = pm.gp.LatentKron(cov_funcs=[spatial_cov, temporal_cov])
f = gp.prior("f", Xs=Xs)

```

Then I have the following error message:

> ValueError: Must provide a covariance function for each X

I think this happens because my `Xs` is a list instead of `n x 3` array. I used to use `n x 3` array, which was manually constructed by looking at only data points where observations are available. But `Xs = [lons[:, None], lats[:, None]], dt_full_norm[:, None]]` yields a list.

I think I am missing something here. Any suggestion?

Thank you.

---

<div class="post-metadata">

**Author:** ![bwengals](https://yyz2.discourse-cdn.com/flex036/user_avatar/discourse.pymc.io/bwengals/32/1237_2.png) [@bwengals](https://discourse.pymc.io/u/bwengals)\
**Post date:** [October 21, 2023, 7:20am UTC](https://discourse.pymc.io/t/possible-ways-to-speed-up-spatio-temporal-gp-modeling/13127/16 "2023-10-21T07:20:05Z")

</div>

> [@seunjeong](#):
>
> ```auto
> spatial_cov = pm.gp.cov.Matern52(input_dim=3, ls=[length_S, length_S], active_dims=[0,1])
> temporal_cov = var*pm.gp.cov.ExpQuad(input_dim=3, ls=length_T, active_dims=[2])
> 
> ```

I think you just need to change input\_dim and active\_dims. Since it’s separable, each cov\_func is just acting over each part of Xs.

Can you try it with this?

```python
spatial_cov = pm.gp.cov.Matern52(input_dim=2, ls=[length_S, length_S])
temporal_cov = var*pm.gp.cov.ExpQuad(input_dim=1, ls=length_T)

```

EDIT: changed input\_dim on temporal\_cov from 2 to 1.

---

<div class="post-metadata">

**Author:** ![seunjeong](https://avatars.discourse-cdn.com/v4/letter/s/ecae2f/32.png) [@seunjeong](https://discourse.pymc.io/u/seunjeong)\
**Post date:** [October 22, 2023, 12:12am UTC](https://discourse.pymc.io/t/possible-ways-to-speed-up-spatio-temporal-gp-modeling/13127/17 "2023-10-22T00:12:27Z")

</div>

Thank you so much! I quickly tried it, but it didn’t work. However, from your comment, I have an idea of how I should try in a few different ways. If none of my tries work, I will ask for your help!

Again, thanks so much!