Error when using pm.Potential with pm.Data + coordinates

Questions
1

Hi everyone!

I am trying to use pm.Data and the coords keyword but I’m getting an error when sampling. Here’s the model:

coords = {
    'LoT': np.arange(len(L)),
    'cat': np.arange(L.shape[1]),
    'obs': np.arange(len(category_i))
}

with pm.Model(coords=coords) as model:
    
    L_data = pm.Data('L', L)#, dims=('LoT', 'cat'))
    cat_i_data = pm.Data('cat_i', category_i)#, dims='obs')
    out_i_data = pm.Data('outcome_i', outcome_i)#, dims='obs')
        
    L_dims = L.shape
    n_lots, n_cats = L.shape

    sigma = pm.HalfNormal('sigma', sigma=2, dims='cat')
    a_0 = pm.Normal('a_0', mu=0, sigma=2)
    a_1 = pm.Normal('a_1', mu=0, sigma=2)

    z_logp = tt.zeros((n_lots,), dtype='float')
    for z_i in range(n_lots):
        zi_logp = 0.
        for cat in range(n_cats):
            obs_idx = np.where(cat_i_data==cat)[0]
            muZ = a_0 + a_1 * L[z_i, cat]
            zi_logp = zi_logp + tt.sum(
                pm.Normal
                .dist(mu=muZ, sd=sigma[cat])
                .logp(out_i_data[obs_idx])
            )
        z_logp = tt.set_subtensor(z_logp[z_i], zi_logp)
    lp3 = pm.Deterministic('z_logp', z_logp, dims='LoT')
    tot_logp = pm.math.logsumexp(z_logp)
    pot = pm.Potential('e', tot_logp)

The model compiles fine, but when I try to sample:

image
image1282×503 72.5 KB

And in the trace zlog_p is always 0, so clearly the sampling isn’t working.

However, when I use the arrays directly instead of using their pm.Data version, everything works fine, and when I comment out

    lp3 = pm.Deterministic('z_logp', z_logp, dims='LoT')
    tot_logp = pm.math.logsumexp(z_logp)
    pot = pm.Potential('e', tot_logp)

everything also works fine (well, except that it’s a different model, but I mean that it runs fine). Also, when I comment out just pot = pm.Potential('e', tot_logp) the problem is still there. So my guess is that it has something to do with the Deterministic z_logp, but I’m not sure what about it is problematic. Thanks in advance for any ideas!

1 Like
2

I am not surprised the deterministic does not handle dims. It is not really supposed to do anything other than store the result of deterministic operations.

2 Likes
3

Is the issue solved by using dims in pm.Data but not in pm.Deterministic?

I thought dims in deterministics were only stored and passed downstream to ArviZ at conversion but had no effect within PyMC3

4

Unfortunately not! Even when changing the line to

    lp3 = pm.Deterministic('z_logp', z_logp)#, dims='LoT')

It gives the same error.

To be specific: if I use arrays directly instead of the pm.Data objects, everything works fine. So I don’t think the problem is about Deterministic and dims. It must have something to do with the combination of Deterministic, dims, and pm.Data.

(To generate fake data:)

sigma = 0.1
a_0, a_1 = 0, 1
z = 3 
L = np.array([
        [1, 2, 1],
        [3, 2, 1],
        [1, 4, 5],
        [5, 6, 8]
    ])

# category of the ith observation
category_i = np.repeat(np.arange(L.shape[1]), 1000)
mu_i = a_0 + a_1 * L[z,category_i]
outcome_i = np.random.normal(
    loc=mu_i,
    scale=[sigma]*len(mu_i)
)