Add tests for partial observed node dynamics

[maedoc]

As part of work on improving inference for the MPR model @mhashemi0873 noticed that latent V was an issue. We can use the new embedding functions to construct a proxy V,

# grab the rate variable from our events
r = rv[..., 0]
n_event, n_time = r.shape

# make autoregressive basis on r
n_lag = 3
r_ar = vbjax.embed_autoregress(r, n_lag=n_lag, lag=5) #.reshape((4, -1))
n_ar_time = r_ar.shape[2]

# subset our known voltage, for the span to be predicted from autoregressive basis
V_ar = rv[:,-n_ar_time:,1]

# reshaping (will be improved)
r_ar = r_ar.reshape((n_lag, -1))
V_ar = V_ar.reshape((1, -1))

# embed in polynomial basis w/ coefficients
basis, coef = vbjax.embed_polynomial(r_ar, V_ar, max_order=3)
print('basis shape', basis.shape)

# try to truncate coef
coef_keep = np.abs(coef) > np.percentile(np.abs(coef), [25.0])
print(coef_keep[:,0].sum(), coef.size)
# coef = coef * coef_keep
# skip since it is not essential

Vh = (coef.T @ basis).reshape((9, -1))

for i in range(9):
    subplot(4, 3, i+1)
    plot(Vh[i], 'k')
    plot(V_ar.reshape((9, -1))[i],'r')
    ylim([-5, 2.5])

subplot(4, 1, 4); plot(coef)

We should add at least one test per model with a simulation and reconstruction of partially observed states, testing that loss is under a certain tolerance, because this ensures that even for cases of real data, where we can't observe the state variables, that we can still estimate (up to some error) the relevant dynamics.