Rewriting likelihood function with for loops using pytensor

1

Setup

I have a Potts model with spins s_i \in {1, …, q}.:

P(s_1, ..., s_n) \sim \exp \sum_{i < j} J_{ij}(s_i, s_j)),

where P denotes the probability of finding the configuration (s_1, ... s_n). I am given fixed values J_{ij}(s_i, s_j) for each pair (s_i, s_j) (so I am not doing Bayesian inference to learn J_{ij}).

What I would like to do

Sample from distribution P.

Challenge

Given a spin configuration, compute the J_{ij}(s_i, s_j) sum using PyTensor. I can do this in numpy, but since I want to sample using pymc, I need to write this using PyTensor.

Example

Suppose I define a spin configuration as follows, where I have 3 positions and 2 possible spins for each position with

J_{i<j}(s_i, s_j) = \begin{pmatrix} 2 & 1 \\ 1 & 2 \end{pmatrix}

for all 0 \leq i < j \leq 2 and J_{i = j}(s_i, s_j) = \begin{pmatrix} 0 & 0 \\ 0 & 0 \end{pmatrix} .

For the likelihood term for the potential, I need to calculate \exp(\sum_{i < j} J_{ij}(s_i, s_j)). For example, given the sequence \{0, 0, 1\}, I should get J_{01}(0, 0) = 2, J_{02}(0, 1) = 1, J_{12}(0, 1) = 1, which gives \exp \{2 + 1 + 1\} = \exp 4.

This is easy in numpy:

j_couplings = np.zeros([2, 2, 3, 3])
for k in range(3):
    for m in range(k + 1, 3):
        j_couplings[0, 0, k, m] = 2
        j_couplings[1, 0, k, m] = 1
        j_couplings[0, 1, k, m] = 1
        j_couplings[1, 1, k, m] = 2

def spin_config_to_energy(spin_config, j_couplings):
   # NEED TO REWRITE THIS USING PYTENSOR OPERATIONS
    energy = 0
    for i in range(len(spin_config)):
        for j in range(i + 1, len(spin_config)):
            energy += j_couplings[spin_config[i], spin_config[j], i, j]
    return energy

with pm.Model() as model:
    spin_config = pm.DiscreteUniform("spins", lower=0, 
                                  upper=1, 
                                  shape=3)
    pm.Potential('energy_function', spin_config_to_energy(spin_config, j_couplings))
2

Hi, I spent a little bit of time on this and thought this might be a helpful start. It seems like you can generate the j_couplings just fine without pytensor. But this might be helpful to get a sense of how pytensor thinks.

j_couplings = pt.zeros((2,2,3,3))

k = pt.as_tensor([0,0,1])
m = pt.as_tensor([1,2,2])

j_couplings = pt.set_subtensor(j_couplings[0,0,k,m],2)
j_couplings = pt.set_subtensor(j_couplings[1,0,k,m],1)
j_couplings = pt.set_subtensor(j_couplings[0,1,k,m],1)
j_couplings = pt.set_subtensor(j_couplings[1,1,k,m],2)

The goal is to rewrite this function using vectorized indices instead of a for loop. So if you could generate a vector that tracks the values of k and m, you can just pass that vector to set_subtensor. set_subtensor is pytensor’s version of M[0,0,1,0] = 2.

Similarly, we can rewrite the spin_config_to_energy function as

a = spin_config[k]
b = spin_config[m]

energy = pt.sum(j_couplings[a,b,k,m])

this shows how you might get rid of the for loop using the same principle as above.

Both of these approaches rely on the idea that you can use static k and m vectors. It seems like that’s your intention because you set the shape of discrete uniform to 3.

4 Likes