I have a Python program (available on github) that uses naive gradient descent to find approximate solutions to the Thomson Problem. It works surprisingly well, but I've been wondering if there's a way to speed up my gradient descent implementation.

The relevant code snippet is the one below: given a set of points in n-space (the rows of the array variable) it computes the gradient function, which for this problem is known analytically.

def grad_func(array, p=1):
    m = len(array)
    n = len(array[0])

    grad_mtx = np.zeros([m, n])
    dm = distance_matrix(array, array)

    for i in range(m):
        grad = 0.0
        # iterate over all points j where i != j
        for j in range(m):
            if i != j:
                new_vec = (array[i] - array[j])/dm[i, j]**(p+1)
                grad = grad + new_vec

        grad_mtx[i] = -2.0*p*grad

    return grad_mtx

Using the kernprof profiler I can see that a large percentage of the run time is spent on the lines

new_vec = (array[i] - array[j])/dm[i, j]**(p+1)
grad = grad + new_vec

For instance, with $n=3$ and $m=100$ the program takes about 21s on my machine. Of this, around 14s is spent on the repeated evaluation of these two lines.

It seems to me that there might be a way to get rid of least one of the for loops by a judicious use of the numpy einsum function, or some other approach. Does any one have any suggestions? For the purpose of this question I am solely interested in improving this implementation of naive gradient descent, not in a better algorithm.

  • $\begingroup$ Are you open to writing parts of this in Cython or C? $\endgroup$ Sep 25, 2023 at 2:02
  • 1
    $\begingroup$ Does grad_mtx = ((array[:,None,:]-array[None,:,:])/dm[:,:,None]).sum(axis=1) work and give the correct result? To avoid zero division, one would have to set dm.diag()=1. $\endgroup$ Sep 25, 2023 at 7:09
  • $\begingroup$ @LutzLehmann I will check $\endgroup$
    – Martin C.
    Sep 26, 2023 at 6:49
  • $\begingroup$ @DanielShapero preferably not, but I am not one to look a gift horse in the mouth $\endgroup$
    – Martin C.
    Sep 26, 2023 at 6:50
  • $\begingroup$ @LutzLehmann It took me a while to get around to this, but indeed, this increases the speed of the algorithm by two orders of magnitude! The code that took 21s to run now takes less than a second. Thank you - if you want to you can expand your comment to an answer, and I will mark it as accepted. $\endgroup$
    – Martin C.
    Sep 30, 2023 at 9:36

1 Answer 1


Instead of using einsum to get compiled loops, one could also first use the provided tools of numpy for elementary array operations.

grad_mtx[i] =  -2.0*p*sum((array[i] - array[j])/dm[i, j]**(p+1) for j in range(m) if j!=i)

still uses interpreted loops. (Try numba.jit, this might be similar to using einsum.)

To make it an operation over array slices using compiled numpy loops, all arrays need to have the same structure, automatic broadcasting gets it wrong here. Thus the difference becomes array[:,None]-array[None,:].

To get the contraction per sum right, one now has to specify axis=1 to select the correct index for summation.

There is still an up-to-now hidden vector dimension in third place. As dm is contracted over this dimension, it is safer in division to mark this likewise with an empty range.

On the diagonal, if it were not excluded, the quotient is of the type 0/0. To avoid getting errors or wrong values there, set the diagonal of dm to a non-zero value

dm.diag() = 1
grad_mtx = -2*p*((array[:,None,:]-array[None,:,:])/dm[:,:,None]**(p+1)).sum(axis=1)

Such use of numpy facilities usually speeds the code up by a factor 10 to 20. Using cython to directly implement this operation in C can result in an additional speed-up by a factor of 30 to 100. This progression can be found in discussions about the mandelbrot fractal (using masks with numpy, returning to pixel iteration loops in cython) or matrix operations


Your Answer

By clicking “Post Your Answer”, you agree to our terms of service and acknowledge you have read our privacy policy.

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