Modeling longitudinal data among patients by coupling gaussian processes with random effects

Hi,

I’ve put the .conditional in the context, then I could indeed use the index as a random variable (I used pm.Categorical). And it works, though I’m still figuring out how to correctly interpret the outcome.

n_new = 10
n_vars = X_sc.shape[1]
X_new = np.random.normal(0, 1, (n_new, n_vars))

with pm.Model() as model:
    
    # Random effect
    mean_intercept = pm.Normal("mean_intercept", mu = 0, sigma = 5)
    std_intercept = pm.HalfNormal("std_intercept", 5)
    
    # intercept
    intercept = pm.Normal(
        "intercept",
        mu = mean_intercept,
        sigma = std_intercept,
        shape = n_idx
    )
    
    # GP
    length_scale = pm.Gamma("length_scale", alpha=2, beta=1)
    amplitude = pm.Gamma("amplitude", alpha=2, beta=1)

    cov = amplitude ** 2 * pm.gp.cov.Matern52(n_vars, length_scale)
    gp = pm.gp.Latent(cov_func = cov)
    gp_prior = gp.prior("gp_prior", X = X_sc)

    # index as a distribution
    idx_d = pm.Categorical('idx_d', p=np.repeat(1, n_idx)/n_idx, observed=idx)
    
    # Random GP
    σ = pm.HalfCauchy("σ", beta=5)
    gp_ri = pm.Deterministic("gp_ri", intercept[idx_d] + gp_prior)
    y_hat = pm.Normal("y_hat", mu = gp_ri, sigma = σ, observed=y_sc)

    # Sample
    trace = pm.sample(1000, tune=1000, chains=2, cores=2, return_inferencedata=True)

    # Predict
    gp_pred = gp.conditional("gp_pred", X_sc) # X_new)
    pred_samples = pm.sample_posterior_predictive(
        trace = trace.posterior,
        var_names = ["gp_pred", "y_hat"],
        samples = 500
    )

Thanks for the answer! I’ll try adding GPs for the next round.

1 Like