Mutable data used in .prior(X) of a GP

Questions v5
1

BLUF: is there a reason that gp.prior() should not take a mutable data tensor as input for X ?

There is not much documentation on using gp.prior(), but it seems like you could use a mutable data object as input to the X values. However, I was getting a NotConstantValueError in pytensorf.py constant_fold. I changed cov.py _slice to disable this check.

def _slice(self, X, Xs=None):
xdims = X.shape[-1]
if isinstance(xdims, Variable):
[xdims] = constant_fold([xdims], raise_not_constant=False)
if self.input_dim != xdims:
warnings.warn(
f"Only {self.input_dim} column(s) out of {xdims} are"
" being used to compute the covariance function. If this"
" is not intended, increase ‘input_dim’ parameter to"
" the number of columns to use. Ignore otherwise.",
UserWarning,
)
X = pt.as_tensor_variable(X[:, self.active_dims])
if Xs is not None:
Xs = pt.as_tensor_variable(Xs[:, self.active_dims])
return X, Xs

With HSGP it sure looks like I am able to update the X values with pm.set_dim() & pm.set_data() and get a good update (after accounting for aligning the means of the X)

2

I think the check is there just to provide that helper warning message. If you’re not changing the last dimension of X you could pass that info with x = pm.MutableData(..., shape=(None, 12)), for example, where with 12 your restrict yourself to always have that size for the last dimension and None allows the other dim to vary.

Then the constant_fold would work?

CC @bwengals, not sure if there’s any reason to require static shape?

3

Yup that’s right. @wmbelk are you having the problem described in this issue (that I forgot about…)?

For HSGP you can use a X = pm.MutableData, but as described in the gp.prior_linearized docstring. For gp.prior and other GP implementations, I guess you maybe “could” mutate the input data (I can’t think of a reason to), but if you’re doing it to get predictions from the GP it won’t give you the right answer.

4

Is the wrong answer just because of the mean alignment?
Admittedly, I have only worked with your cherry blossom example, but when I input new year values extended out and just far enough back in time to maintain the mean year… I get (by the eyeball) very similar results.
from your git hub (using: phi, sqrt_psd = gp.prior_linearized(X) )

engals_extension
engals_extension1433×813 117 KB

After alteration:

belks-extension
belks-extension1433×813 113 KB

It could very well be quirk that this happened to work out… just curious about what would lead to your intuition about getting the wrong answer from a projection using pm.set_data on a HSGP.
What I did instead:

x = df[“year”].values

with pm.Model() as model:

model.add_coord(‘year’, df.year.tolist(),mutable=True)
eta = pm.Exponential(“eta”, lam=0.25, shape=1)
ell_params = pm.find_constrained_prior(
pm.InverseGamma,
lower=20, upper=200,
init_guess={“alpha”: 2, “beta”: 20},
mass=0.95,
)
ell = pm.InverseGamma(“ell”, **ell_params)
# Specify covariance function and GP
cov = eta[0]**2 * pm.gp.cov.Matern52(1, ls=ell)
intercept = pm.Normal(“intercept”, mu=df[“temp”].mean(), sigma=2 * df[“temp”].std())
gp = pm.gp.HSGP(m=[400], L=[1565.0], cov_func=cov, mean_func=pm.gp.mean.Constant(intercept))
x_mu = pm.ConstantData(‘x_mu’,np.mean(x))
X = pm.MutableData(f"X_mu_{x_mu}", x[:, None])
X = X
f = gp.prior(“f”, X=X, dims=“year”)
mu = pm.Deterministic(“mu”, f, dims=“year”)
sigma = pm.HalfNormal(“sigma”, sigma=2 * df[“temp”].std())
pm.Normal(“y”, mu= f, sigma=sigma, observed=df[“temp”].values, shape=X.shape[0 ])

Then:

def extend_x_keep_mean(x_extension=50, target_mean=np.mean(x), x=x, make_int=True):

maintain_mean_of_x_halfwidth = x_extension + np.max(x) - target_mean
x_new = np.arange((target_mean - maintain_mean_of_x_halfwidth),
(target_mean + maintain_mean_of_x_halfwidth))
if make_int:
x_new = x_new.astype(int)
print(f’{x_new.mean()-target_mean=}')
return x_new

x_new = extend_x_keep_mean()

Then:

with model:

model.set_dim(“year”, new_length=len(x_new),coord_values=x_new.tolist())
model.set_data(name=f"X_mu_{x_mu}",values=x_new[:, None], coords={“year”: x_new.tolist()})
ppc = pm.sample_posterior_predictive(idata_gp, var_names=[ ‘f’])

1 Like
5

Yeah, it looks like what you’re doing gave the right answer! The reason it’s not recommended is because you definitely can’t do this with any other gp.prior non-HSGP implementation. In the code of HSGP.prior it calls HSGP.prior_linearized, so as long as you handle that scaling of X correctly it’ll work out. Nice!

1 Like