Hi, thanks for your hints @jessegrabowski.
So far I have not pursued that route because of the way VAR is fitted in the notebook.
The predictions are made for one time-step only (probably better convergence properties) but that is not possible in my case: I have gappy tide-gauge records that only partially overlap, and never all together. That means I have to make a forward prediction from time 0 all the way through the time-series (though it could be possible in theory to use observations instead of previous time-step prediction whenever they are available). That means using scan, which I am not good at… (I might be doing something wrong, but I failed just to reproduce the AR(1) class with scan – it does not even start the sampling). In contrast my current approach, though less transparent, seems to work quite well, and uses half the parameters (because of the symmetry of the correlation matrix, reflected in the LKJCholesky triangular structure). Here is what I do:
ns = 14 # number of locations, also reflected in the `station` dimension
ny = 200 # number of years for the AR(1) process
# Parameters for the AR(1)
rho = pm.Uniform("oceandyn_ar1_rho", dims='station')
# ... I calculated sigma below so that the DIFF of the AR(1) process has variance = 1
sigma = at.sqrt((1+rho)/2)
# Here 14 individual AR(1) processes are generated (with 5 years throw-away warmup)
# and the diff is taken (I work on the rate, and the cumulative part is AR(1))
warmup = 5
ar1_slr = pm.AR(f"oceandyn_ar1_slr", rho=rho[None].T, sigma=sigma, ar_order=1, steps=ny+warmup,
shape=(ns, ny+warmup+1), init_dist=pm.Normal.dist(0, 1)) # size years + 1
uncorrelated_timeseries = at.diff(ar1_slr[:, warmup:], axis=1).T # years x stations
# Now I use the Cholesky matrix to introduce correlation
chol, _, _ = pm.LKJCholeskyCov("oceandyn_chol", n=ns, eta=1, sd_dist=pm.HalfNormal.dist(25, dims=ns) )
correlated_timeseries = at.dot(uncorrelated_timeseries, chol.T) # picks up the variance of chol
# And later, the likelihood function ..
sat_rate = ... # observations
sat_sigma = pm.HalfNormal("satellite_error", 5, dims="station")
pm.Normal(f"satellite_obs", mu=correlated_timeseries[1994-1900:2019-1900+1], sigma=sat_sigma, observed=sat_rate, dims=("sat_year", "station"))
Somewhat surprisingly, pyMC is able to derive the posterior of such a large guessing, solving in time and space. That is quite a convenient way not to have to split up the code into tuning and prediction. However I am still encountering convergence issues (which I think are due to other aspects), so the discussion is not closed yet.