Quick way to divide matrix entries K_ij by K_ii*K_jj in Python

Viewed 61

In Python, I have a matrix K of dimensions (N x N). I want to normalize K by dividing every entry K_ij by sqrt(K_(i,i)*K_(j,j)). What is a fast way to achieve this in Python without iterating through every entry?

My current solution is:

import numpy as np
K = np.random.rand(3,3)
diag = np.diag(K)
for i in range(np.shape(K)[0]):
    for j in range(np.shape(K)[1]):
        K[i,j] = K[i,j]/np.sqrt(diag[i]*diag[j])
1 Answers

Of course you have to iterate through every entry, at least internally. For square matrices:

K / np.sqrt(np.einsum('ii,jj->ij', K, K))

If the matrix is not square, you first have to define what should replace the "missing" values K[i,i] where i > j etc.

Alternative: use numba to leave your loop as is, get free speedup, and even avoid intermediate allocation:

@njit
def normalize(K):
    M = np.empty_like(K)
    m, n = K.shape
    for i in range(m):
        Kii = K[i,i]
        for j in range(n):
            Kjj = K[j,j]
            M[i,j] = K[i,j] / np.sqrt(Kii * Kjj)
    return M
Related