Performance issues with double for loops, EM-algorithm

The alpha that I create is not created in a crazily efficient way I know, but the main bottleneck in this code is in what I call the "Maximization step" (see comments in code). There are double for-loops twice (with list comprehensions inside), and the reason for this is that I need to use the computed vector from the first double for-loop in the second double for-loop. Here's some proof of my claim, I tested how long the different parts of the code took;

Finishing the vector alphas took 0.1451436 s.

Finishing the first bottleneck for loops took 13.7402911 s

Finishing the second bottleneck for loops took 13.1357222 s

A single complete loop in the while step took 26.9219609 s

There should be a more numpy-focused way to update the three vectors mu, pi, sigma_squared, I just can't seem to figure out how. Any help to make this quicker would be greatly appreciated.

import numpy as np
import random as rnd
from scipy.stats import norm

no_distributions = 2 # range of L
range_of_i = 255
range_of_m = 100

iterations = 100
mu = [5.7e-04, 5.9e-04]
sigma_squared = [9.5e-07, 9.5e-07]
weights = [0.49, 0.51]
alphas = []
for i in range(range_of_i):
alpha = []
for m in range(range_of_m):
sub_populations = [rnd.gauss(mu[i], np.sqrt(sigma_squared[i])) for i in range(no_distributions)]

alphas.append(alpha)

threshold = 10
thresh_limit = 1e-4
initial_parameters = np.array([[-1, 1], [0.2, 0.2], [0.5, 0.5]])
# we iterate between the expectation and maximization steps until convergence
while threshold > thresh_limit:

# mu, pi and sigma_squared need to be numpy arrays because we want to add vectors together the numpy way
mu_initial = initial_parameters[0, :]
sigma_squared_initial = initial_parameters[1, :]
pi_initial = initial_parameters[2, :]
mu_updated = np.zeros(no_distributions)
pi_updated = np.zeros(no_distributions)
sigma_squared_updated = np.zeros(no_distributions)
indicator_normalized = []

# Expectation step
for i in range(range_of_i):
indicator = [pi_initial[l] * norm.pdf(alphas[i], mu_initial[l], np.sqrt(sigma_squared_initial[l])) for l
in range(no_distributions)]
indicator_normalized.append([l / sum(indicator) for l in indicator])

indicator_normalized = np.array(indicator_normalized)
indicator_normalized = indicator_normalized.transpose(0, 2, 1)

# Maximization step
for i in range(range_of_i):
for j in range(np.array(alphas).shape[1]):
mu_updated += np.array(
[indicator_normalized[i][j][l] * alphas[i][j] / sum(sum(indicator_normalized))[l] for l in
range(no_distributions)])
pi_updated += np.array(
[indicator_normalized[i][j][l] / (range_of_i * range_of_m)
for l in range(no_distributions)])

# same for loop again needed because we want to use the complete mu_vector to calculate sigma_squared
for i in range(range_of_i):
for j in range(np.array(alphas).shape[1]):
sigma_squared_updated += np.array([indicator_normalized[i][j][l] * (
alphas[i][j] - mu_updated[l]) ** 2 / sum(sum(indicator_normalized))[l] for l in
range(no_distributions)])

# update parameters and calculate the new threshold to see if convergence has been met
parameters_updated = np.array([mu_updated, sigma_squared_updated, pi_updated])
threshold = sum(sum(np.abs(parameters_updated - initial_parameters)))
initial_parameters = parameters_updated

print (parameters_updated)


The code basically implements an EM-algorithm where the purpose is to find better estimates (better yet, the MLE) of the parameters $$\\mu, \sigma, \pi\$$, which are part of a gaussian mixture model. If you want to understand where the calculations come from, I think the easiest way is just to show you a short description of the algorithm and the formulas below.

• Suppose the initial parameter vector is $$\\hat\theta = (\{\hat\pi_l\}_{l=1}^L, \{\hat\mu_l\}_{l=1}^L, \{\hat\sigma_l^2\}_{l=1}^L)\$$

• Expectation step: Compute the expected value of the indicator variable that indicates which population (e.g., the population of skilled or unskilled managers) $$\\alpha_{ij}\$$ is drawn from:
\\begin{align*}\hat{p}_{ijl} & = \hat{Pr}(\alpha_{ij}\text{ comes from Group }l) \\ & = \frac{\hat\pi_l\phi(\alpha_{ij};\hat\mu_l,\hat\sigma^2_l)}{\sum\limits_{l=1}^{L}\hat\pi_l\phi(\alpha_{ij};\hat\mu_l,\hat\sigma^2_l)}, i = 1, \ldots , N; j = 1, \ldots, M; l = 1, \ldots, L \end{align*}\
where $$\\phi(\cdot;\mu,\sigma^2)\$$ is the density of the normal distribution $$\\mathcal{N}(\mu, \sigma^2)\$$

• Maximization Step: Compute the weighted means and variances with weights obtiained from the Expectation Step:
$$\\tilde\mu_l=\frac{\sum_{ij}\hat{p}_{ijl}~\alpha_{ij}}{\sum_{ij}\hat{p}_{ijl}}, \tilde\sigma^2_l = \frac{\sum_{ij}\hat{p}_{ijl}~(\alpha_{ij} - \tilde\mu_l)^2}{\sum_{ij}\hat{p}_{ijl}}, \tilde\pi_l = \frac{\sum_{ij}\hat{p}_{ijl}}{MN}, l = 1, \ldots, L.\$$

• Iterate between the Expectation Step and the Maximization Step until convergence.

Congratulations to you for finding a major bottleneck yourself. There are a few more to get rid of.

Generally speaking it's quite a performance killer to convert data between Python and numpy repeatedly. And you do that a lot. You even do it to determine the number of iterations for some loops, e.g. in for j in range(np.array(alphas).shape[1]):. Since alphas is created beforehand and not expanded again, the first step would be to convert it into a numpy array once. But there is no real need to query its shape, since you know exactly how the shape of alphas is, it's exactly (range_of_i, range_of_m). So you convert $$\255 \times 100\$$ elements to numpy just to get range_of_m as an answer. Not really efficient, isn't it?

In addition, Python is also often not really fast at loops. There is a great talk from Jake VanderPlas from PyCon 2015 titled "Losing your Loops - Fast Numerical Computing with NumPy" (to be found on YouTube) on that topic. I highly recommend watching it!

After the general remarks, let's start from the beginning:

When you generate alphas, you draw samples from no_distributions distributions, and then discard all but one value if I'm not terribly mistaken. So why not choose the distribution beforehand, and sample that one?

alphas = np.empty((range_of_i, range_of_m))
dist_idxs = list(range(no_distributions))
sigma = np.sqrt(sigma_squared)
for i in range(range_of_i):
idxs = rnd.choices(dist_idxs, weights=weights, k=range_of_m)
alphas[i, :] = [rnd.gauss(mu[idx], sigma[idx]) for idx in idxs]


This code snippet also includes the idea to have alphas as numpy array directly. You could also keep the original approach to have a list of lists and then convert that, but this helps you to save some variables. I also opted to remove the inner loop in favor of a list comprehension. Maybe there is even a clever way to get rid of the outer loop (without using a list comprehension!), but I'll leave that as an exercise to you.

Next up on the list is the expectation step. Again, funny mix-up of Python and numpy/scipy code here. And although I have the strong feeling that there has to be a better solution, the following is the best I could come up with at the moment:

sigma_initial = np.sqrt(initial_parameters[1, :])

# ... other code ...

# Expectation step
indicator_normalized = np.empty((range_of_i, no_distributions, range_of_m))
for i in range(range_of_i):
indicator = np.array([
pi_initial[l] * norm.pdf(alphas[i], mu_initial[l], sigma_initial[l])
for l in range(no_distributions)
])
indicator_normalized[i, ...] = indicator / indicator.sum(axis=0)

indicator_normalized = indicator_normalized.transpose(0, 2, 1)
# summing over multiple axis needs numpy >= 1.7
indicator_sum = indicator_normalized.sum(axis=(0, 1))


Next up on the list: maximization.

Let's get rid of some loops here. A first step in that direction might look as follows:

# Maximization step
for i in range(range_of_i):
for j in range(range_of_m):
mu_updated += indicator_normalized[i, j, :] * alphas[i, j] / indicator_sum
pi_updated += indicator_normalized[i, j, :] / (range_of_i * range_of_m)

# same for loop again needed because we want to use the complete mu_vector to calculate sigma_squared
for i in range(range_of_i):
for j in range(range_of_m):
sigma_squared_updated += \
indicator_normalized[i, j, :] * (alphas[i, j] - mu_updated)**2 \
/ indicator_sum


But we can do better! Enter numpy broadcasting and slicing!

# Maximization step
mu_updated = np.sum(
indicator_normalized * alphas[..., np.newaxis] / indicator_sum,
axis=(0, 1)
)
pi_updated = np.sum(
indicator_normalized / (range_of_i * range_of_m), axis=(0, 1)
)

# same for loop again needed because we want to use the complete mu_vector to calculate sigma_squared
sigma_squared_updated = np.sum(
indicator_normalized * (alphas[..., np.newaxis] - mu_updated)**2
/ indicator_sum,
axis=(0, 1)
)


Explaining the full beauty of this would be beyond the scope of what I can do here. I highly recommend reading the chapter "Fancy Indexing" from the Python Data Science Handbook by Jake VanderPlas if you want to learn more. What this does in essence is to move those pesky loops into the numpy's C back end. You can check that those results really match by pasting both pieces of code into the same file and check with np.allclose(...).

Results

So what did we win be jumping through all those hoops and losing all those loops?

Your proposed solution that moved the sum(sum(...)) out of the loop takes around $$\5s\$$ to converge here on my machine. With the code above in its most loop-reduced form the algorithm finishes in about $$\0.2s\$$ to $$\0.25 s\$$ with the same result (not exactly the same cause the random numbers also have a little impact here). That's a speed-up of 20-25x. I can live with that for the moment ;-).
(Note: The timing results only cover the content of the while loop, not the initialization.)

I'm actually not a 100% sure the implementation is fully correct, but the new implementation is at least as correct as the original one. If you find an error here in my code, it already existed in the original solution because I tested for agreement between your and my solution ;-)

• I had to study this answer for a few hours before understanding it, I learned so much! Can't thank you enough, I'm gonna try using numpy for everything from here on out. – armara Jul 12 '19 at 10:27
• Welcome to the world of numpy ;-) I also still learn new things while working on different projects. A word of warning though: numpy is not a magic bullet! There are cases where numpy can slow your code down. But whenever you can use it to vectorize your code to get rid of nested loops as presented above, the performance will almost always be significantly better. – AlexV Jul 12 '19 at 11:57

I searched around this forum a bit more and found some advice to use the package cProfile. This helped me locate the problem which was the sum(sum(indicator_normalized)) which took more than 95% of the total time for the complete algorithm. I just moved out the sum and the total time for my algorithm went down to 1/19 of the total time it took before.

So the solution was to write this before the maximization step, and then use this variable inside the for loops;

indicator_sum = sum(sum(indicator_normalized))
# Maximization step

• Be that as it may, I would still leave your question open for others to suggest improvements to the code. – Reinderien Jul 10 '19 at 16:01