Writing tests for the log probability of the sum to zero ICAR prior via `pm.Potential`

Development
,
1

Hi there,

I am trying to implement spatial models that use an ICAR prior. There is an issue related to the ICAR prior currently open, and I think the issue is still open because it was hard to test the log probability function because the multivariate distribution itself is a singular Gaussian.

I am trying to use the pm.Potential to test the log probability for various adjacency matrices. I don’t know how to use pm.Potential, but I believe it should be able to be used to replicate a Stan model such as:
image
from this Stan case study. I have never had a parameter vector within a PyMC model before that I haven’t been able to define by calling a PyMC distribution to specify its prior. It feels like I need to call some Aesara function to define the phi vector as a set of empty? nodes on the graph (but I don’t know what I am talking about :slight_smile:).

If anyone has any thoughts, knowledge, or tips on using pm.Potential within a model, that would be appreciated.

Also, I would love to learn more about Aesara, and if you have any pointers to valuable parts of the Aesara documentation, etc., for this type of thing, that would be awesome.

Thanks,
Conor

1 Like
2

pm.Potential adds an arbitrary term to the model logp. The simplest case, say you wanted to add a unit gaussian logpdf prior term but there was no pm.Normal implemented:

def normal_logpdf(value, mu, sigma):
  ...

with pm.Model() as m:
  x = pm.Flat("x")
  pm.Potential("gaussian_term", normal_logpdf(x, 0, 1))

Which would be equivalent to:

with pm.Model() as m:
  x = pm.Normal("x", 0, 1)
1 Like
3

Thank you, Ricardo; that is just what I was looking for! Rather than write the tests, I couldn’t help but try to write a model :slight_smile:

with pm.Model(coords={"num_areas": np.arange(N)}) as bym_model:
    # precision priors 
    tau_theta = pm.Gamma("tau_theta", alpha=3.2761, beta=1.81)
    tau_phi = pm.Gamma("tau_phi", alpha=1, beta=1)
    # transform from precision to standard deviation 
    sigma_theta = pm.Deterministic("sigma_theta", 1/at.sqrt(tau_theta))
    sigma_phi = pm.Deterministic("sigma_phi", 1/at.sqrt(tau_phi))

    # independent random effect prior, constraining to mean 0
    theta = pm.Normal("theta", mu=0, sigma=1, dims="num_areas")
    theta_cons = pm.Deterministic("theta_cons", theta-at.mean(theta), dims="num_areas")
    # spatial ICAR random effect prior 
    phi = pm.Flat("phi", dims="num_areas")
    pm.Potential("spatial_diff", pairwise_diff(phi, node1, node2)) # `pairwise_diff` denotes to equivalent sigma=1 prior
    phi_cons = pm.Deterministic("phi_cons", phi-at.mean(phi), dims="num_areas")
    
    # regression coefficient priors 
    beta0 = pm.Normal("beta0", mu=0, sigma=5)
    beta1 = pm.Normal("beta1", mu=0, sigma=5)
    
    # linear predictor 
    eta = pm.Deterministic("eta", logE + beta0 + beta1*x + phi_cons*sigma_phi + theta_cons*sigma_theta, dims="num_areas") 

    # likelihood
    obs = pm.Poisson("obs", at.exp(eta), observed=y, dims="num_areas")

The sampling isn’t great, with some divergences, tree depth issues, and low acceptance probabilities. However, this is probably a combination of needing to fiddle with sampler hyperparameters, reparameterizing the model or dealing with constraints differently.

Onto attempting to write some tests for the ICAR distribution!

5 Likes
4

Sorry I saw this late, but I want to throw in my 2c here for reference. Instead of explicitly subtracting the mean of theta, theta - at.mean(theta), you can use another potential to get the sum to zero or mean zero constraint,

  1. “soft” sum to zero constraint, as implemented in the Stan case study
logp_func = pm.Normal.dist(mu=0.0, sigma=np.sqrt(0.001))
pm.Potential("zero_sum", pm.logp(logp_func, pm.math.sum(theta)))

Or, there is also
2. “hard” sum to zero constraint (credit to @aseyboldt)

import aesara.tensor as at


def ZeroSumNormal(name, *, sigma=None, size, model=None):
    model = pm.modelcontext(model=model)

    def extend_axis(value, axis):
        n_out = value.shape[axis] + 1
        sum_vals = value.sum(axis, keepdims=True)
        norm = sum_vals / (at.sqrt(n_out) + n_out)
        fill_val = norm - sum_vals / at.sqrt(n_out)
        out = at.concatenate([value, fill_val], axis=axis)
        return out - norm
    
    raw = pm.Normal(f"{name}_raw", sigma=sigma, size=size - 1)
    raw = extend_axis(raw, 0)
    return pm.Deterministic(name, raw)

with pm.Model() as model:
    theta = ZeroSumNormal("theta", sigma=1.0, size=10)
    tr = pm.sample(return_inferencedata=False)

tr["theta"].sum(axis=1)
3 Likes