Faster matrix calculation in numpy

Question

Is there some faster variant of computing the following matrix (from this paper), given a nxn matrix M and a n-vector X: ?

I currently compute it as follows:

#M, X are given as numpy arrays
G = np.zeros((n,n))
for i in range(0,n):
    for j in range(i,n):
        xi = X[i]
        if i == j:
            G[i,j] = abs(xi)
        else:
            xi2 = xi*xi
            xj = X[j]
            xj2 = xj*xj
            mij = M[i,j]
            mid = (xi2 - xj2)/mij
            top =  mij*mij + mid*mid + 2*xi2 + 2*xj2
            G[i,j] = math.sqrt(top)/2

This is very slow, but I suspect there is a nicer "numpythonic" way of doing this instead of looping...

EDIT: While all answers work and are much faster than my naive implementation, I chose the one I benchmarked to be the fastest. Thanks!

Why didn't you use the formula for all `G[i, j]`? Seems like diagonal is simply assigned to X and lower triangle is not computed. — swag2198, Nov 26 '20 at 09:26
@swag2198 the formula is symmetric, so I just need the upper triangular part of the matrix — berndibus, Nov 26 '20 at 09:28
@Moritz yes, but I unfortunately can't use it in this setting. If I were free to use whatever I wanted, I would have done it in C++ anyways.. But the numpy version from the answers is fast enough! — berndibus, Nov 26 '20 at 09:57

swag2198 · Accepted Answer · 2020-11-26T10:00:44.387

Quite straightforward actually.

import math
import numpy as np

n = 5
M = np.random.rand(n, n)
X = np.random.rand(n)

Your code and result:

G = np.zeros((n,n))
for i in range(0,n):
    for j in range(i,n):
        xi = X[i]
        if i == j:
            G[i,j] = abs(xi)
        else:
            xi2 = xi*xi
            xj = X[j]
            xj2 = xj*xj
            mij = M[i,j]
            mid = (xi2 - xj2)/mij
            top =  mij*mij + mid*mid + 2*xi2 + 2*xj2
            G[i,j] = math.sqrt(top)/2

array([[0.77847813, 5.26334534, 0.8794082 , 0.7785694 , 0.95799072],
       [0.        , 0.15662266, 0.88085031, 0.47955479, 0.99219171],
       [0.        , 0.        , 0.87699707, 8.92340836, 1.50053712],
       [0.        , 0.        , 0.        , 0.45608367, 0.95902308],
       [0.        , 0.        , 0.        , 0.        , 0.95774452]])

Using broadcasting:

temp = M**2 + ((X[:, None]**2 - X[None, :]**2) / M)**2 + 2 * (X[:, None]**2) + 2 * (X[None, :]**2)
G = np.sqrt(temp) / 2

array([[0.8284724 , 5.26334534, 0.8794082 , 0.7785694 , 0.95799072],
       [0.89251217, 0.25682736, 0.88085031, 0.47955479, 0.99219171],
       [0.90047282, 1.10306597, 0.95176428, 8.92340836, 1.50053712],
       [0.85131766, 0.47379576, 0.87723514, 0.55013345, 0.95902308],
       [0.9879939 , 1.46462011, 0.99516443, 0.95774481, 1.02135642]])

Note that you did not use the formula directly for diagonal elements and only computed for upper triangular region of G. I simply implemented the formula to calculate all G[i, j].

Note: If diagonal elements of M don't matter and they contain some zeros, just add some offset to avoid the divide by zero error like:

M[np.arange(n), np.arange(n)] += 1e-5

# Do calculation to get G

# Assign diagonal to X
G[np.arange(n), np.arange(n)] = abs(X)

This looks great, and much faster! But it gives an error if a value of M is zero, which is the case for me for some of the diagonals (that's why I used the if-else clause, sorry for not writing that in the question...) — berndibus, Nov 26 '20 at 09:35
That is a good question; I believe it doesn't matter as these values will never be used. I just set them to zero in my current code, by using numpy's nan_to_num. — berndibus, Nov 26 '20 at 09:47
I added some note to my answer. You can offset the diagonal to do the same calculation, then assign the diagonal of `G` to `X` later if you wish so. This basically gives you the `if i == j` part of your original code. — swag2198, Nov 26 '20 at 09:56

score 2 · Answer 2 · answered Nov 26 '20 at 09:24

First, you function is not you equation. As this line

 mid = (xi2 - xj2)/mij

should be

 mid = (xi - xj)/mij

Second, I use numpy generate your equation.

Generate test data

test_m = np.array(
    [
        [1, 2, 3, 4, 5],
        [1, 2, 3, 4, 5],
        [1, 2, 3, 4, 5],
        [1, 2, 3, 4, 5],
        [1, 2, 3, 4, 5],
    ]
)
test_x = np.array([5, 6, 7, 8, 9])

build function

def solve(m, x):
    x_size = x.shape[0]
    x = x.reshape(1, -1)
    reshaped_x = x.reshape(-1, 1)
    result = np.sqrt(
        m ** 2
        + ((reshaped_x - x) / m) ** 2
        + 2 * np.repeat(reshaped_x, x_size, axis=1) ** 2
        + 2 * np.repeat(x, x_size, axis=0) ** 2
    ) / 2
    return result

run

print(solve(test_m, test_x))

In fact, the result part could be simpfy like this:

    result = np.sqrt(
        m ** 2
        + ((reshaped_x - x) / m) ** 2
        + 2 * reshaped_x ** 2
        + 2 * x ** 2
    ) / 2

score 1 · Answer 3 · answered Nov 26 '20 at 10:03

Tested with googles colab:

import numba import numpy as np import math

# your implementation

def bench_1(n):
    #M, X are given as numpy arrays
    G = np.zeros((n,n))
    M = np.random.rand(n, n)
    X = np.random.rand(n)
    for i in range(0,n):
        for j in range(i,n):
            xi = X[i]
            if i == j:
                G[i,j] = abs(xi)
            else:
                xi2 = xi*xi
                xj = X[j]
                xj2 = xj*xj
                mij = M[i,j]
                mid = (xi2 - xj2)/mij
                top =  mij*mij + mid*mid + 2*xi2 + 2*xj2
                G[i,j] = math.sqrt(top)/2
    return G

  %%timeit
  n = 1000
  bench_1(n)

  1 loop, best of 3: 1.61 s per loop

Using Numba to compile the function:

@numba.jit(nopython=True, parallel=True)
def bench_2(n):
  #M, X are given as numpy arrays
  G = np.zeros((n,n))
  M = np.random.rand(n, n)
  X = np.random.rand(n)
  for i in range(0,n):
      for j in range(i,n):
          xi = X[i]
          if i == j:
              G[i,j] = abs(xi)
          else:
              xi2 = xi*xi
              xj = X[j]
              xj2 = xj*xj
              mij = M[i,j]
              mid = (xi2 - xj2)/mij
              top =  mij*mij + mid*mid + 2*xi2 + 2*xj2
              G[i,j] = math.sqrt(top)/2
  return G

%%timeit
n = 1000
bench_2(n)

The slowest run took 88.13 times longer than the fastest. This could mean that an intermediate result is being cached.
1 loop, best of 3: 9.8 ms per loop

Faster matrix calculation in numpy

3 Answers3