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)
    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 2total 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?

  • 2
    \$\begingroup\$ Welcome to Code Review! I will help reviewers if you describe the problem you are trying to solve in more detail. At the moment its not immediately obvious why the system is modeled as it is in your code. \$\endgroup\$ – AlexV Jul 22 '19 at 13:43
  • \$\begingroup\$ Welcome to Code Review! The current question title, which states your concerns about the code, is too general to be useful here. Please edit to the site standard, which is for the title to simply state the task accomplished by the code. Please see How to get the best value out of Code Review: Asking Questions for guidance on writing good question titles. \$\endgroup\$ – Toby Speight Jul 22 '19 at 14:04
  • 3
    \$\begingroup\$ I have added more context to my question and changed the title. \$\endgroup\$ – papabiceps Jul 22 '19 at 14:17
  • 1
    \$\begingroup\$ You can use MathJax here on Code Review to typeset your math expressions directly in the question. \$\endgroup\$ – AlexV Jul 22 '19 at 14:36
  • \$\begingroup\$ I have edited the code. Apologies for the mixup. \$\endgroup\$ – papabiceps Jul 22 '19 at 22:15

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.

  • \$\begingroup\$ I have run your code and timed both of the ways using timeit. I'm also seeing a ~3x speedup. For n = 21, my code clocked 6.21s per loop and your code clocked 2.49s with numba.jit it clocked 0.548s per loop. That's a great a speedup. I'm eagerly for your further analysis. How do I train myself to write optimized code like this, any suggestions ? Thank you for the help :) \$\endgroup\$ – papabiceps Jul 25 '19 at 9:13
  • \$\begingroup\$ I just found out another way to do this with no for loop and all matrices which is in total 6x faster than my first attempt but I could only verify this for n = 10, anything more than that I'm running out of memory and for some n its slower than your method. This is the code to my method pastebin.com/2D64jQQr \$\endgroup\$ – papabiceps Jul 25 '19 at 14:25
  • \$\begingroup\$ Basically I do the operations as shown in the original equation in the question and do trace(result) to get the PDF. Maybe if I use np.einsum it would be faster because I'm doing a lot of wasteful calculations. \$\endgroup\$ – papabiceps Jul 25 '19 at 14:28
  • 1
    \$\begingroup\$ Regarding your other question: Python Data Science Handbook by Jake VanderPlas as well as his PyCon talks (2017, 2018) are a good start to get going with optimization potential. \$\endgroup\$ – AlexV Jul 25 '19 at 20:33
  • 1
    \$\begingroup\$ try this np.einsum('ij,jk,ik->i',states, J, states). It is giving me more speedup as the value of n is increases. Ref: stackoverflow.com/questions/57216521/… \$\endgroup\$ – papabiceps Jul 26 '19 at 9:47

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service, privacy policy and cookie policy

Not the answer you're looking for? Browse other questions tagged or ask your own question.