Evaluate joint probability density function of a Markov random field

Question

The following code evaluates probability mass function for all possible states of a model.

def PDF(size):
    b = np.random.randn(size)
    J = np.random.randn(size, size)

    density_func = np.zeros(2**size)
    states = dec2bin(size) 

    for i in range(2**size):
        density_func[i] = np.exp((np.dot(b, states[i,:]) + np.dot(np.dot(states[i],J),states[i])))

    Z = np.sum(density_func)
    density_func = density_func / Z 

    return density_func

utility functions

def bitfield(n,size):
    x = [int(x) for x in  bin(n) [ 2 :]]
    x = [0] * (size - len(x)) + x
    return x

def dec2bin(size):
    states = []
    for i in range(2**size):
        binary = bitfield(i, size)
        states.append(binary)
    return np.array(states)

The model is grid graph Markov random field. Each node of the graph can have two states {0, 1}, so the total number of possible states of the model is 2^{total number of nodes}. Each instance of the for loop is calculating the probability of that state. At end I want to calculate the joint probability distribution of all the nodes. The joint probability distribution of this model is as follows

$$ p(\mathbf{x}) = \tfrac{1}{z} exp(\mathbf{b}\cdot\mathbf{x} + \mathbf{x}\cdot\mathbf{J}\cdot\mathbf{x}) $$

\$b\$ and \$J\$ are model parameters and \$x\$ is a vector of present states of the system.

\$Z\$ is the normalizing constant which is calculated by summing up 'values' of each possible state of the system to convert them into probabilities.

When the variable size is greater than 25 the code takes long time to execute. Is there a way to vectorize this code and speed it up?

Answer 1

A first go at the problem:

Generating the states

Since you are generating all possible bit patterns with a given number of bits, it's wasteful to first create a binary representation as a string, and then convert it back to integers digit by digit. As often itertools can help you here. Especially itertools.product. Your use case is basically one of the examples given in its documentation.

from itertools import product

def generate_all_possible_states(size):
    states = np.empty((2**size, size))
    for i, value in enumerate(product((0, 1), repeat=size)):
        states[i, :] = value
    return states

The reimplementation also avoids the dynamically growing list and allocates a numpy array of appropriate size beforehand.

Moving repeated computations

You are calculating exp for every single scalar that you put in density_func. That is wasteful since it does not allow numpy to play its strengths, namely applying the same operation to a lot of values aka vectorization. Fortunately, there is an easy way out (ot the loop ;-)) and you can just compute np.exp on the whole density_func vector.

The first part of the sum can also be easily moved out of the loop.

density_func = np.empty(2**size)
for i in range(2**size):
    density_func[i] = np.dot(np.dot(states[i], J), states[i])

density_func += np.dot(states, b)
density_func = np.exp(density_func)
# maybe a teeny-tiny bit faster:
# density_func = np.exp(density_func + np.dot(states, b))

The order of operation matters here. Intializing density_func to np.dot(states, b) and then adding the results in the loop like density_func[i] += np.dot(np.dot(states[i], J), states[i]) is consistently slower.

---

Preliminary timing

As an example, for size = 18, your original implementation takes about \$5.5s\$ here on my machine. With just the two modifications above, you can get down to \$1.9s\$ (both with the same fixed seed for the random numbers).

• size = 20: \$23.4s\$ vs. \$8.0s\$
• size = 21: \$47.9s\$ vs. \$16.4s\$

Maybe I will have time to revisit this later on this week to update the answer with a proper analysis on how both of them scale for larger sizes. My strong intuition at this point would be that they both have a complexity of \$\mathcal{O}(2^n)\$ and the difference between them will actually remain a constant factor (\$\approx 3\times\$ faster).

---

Things that didn't work (so far)

Things that I tried, that didn't work for me:

1. Eliminating the magic of np.dot and explicitely using np.inner or np.matmul / @ (in Python 3). np.dot seems to have some tricks up its sleeves and is still faster. Maybe I will have to read the doc more carefully.

2. Trying to compute the second part of the sum in a vectorized manner. I'm a little bit too tired at the moment to come up with a clever solution to this.