pythonfor-loop

Increase speed of two for loops in Python


I have a double for loop in a Python code. I want to increase the execution speed, but the matrix is built using data from the previous row and column. How can I achieve this?

import numpy as np

num_rows, num_cols = 100, 5

matrix = np.zeros((num_rows, num_cols))

matrix[0, :] = np.random.rand(num_cols)
matrix[:, 0] = np.random.rand(num_rows)

coeff1 = np.random.rand(num_rows)
coeff2 = np.random.rand(num_rows)
coeff3 = np.random.rand(num_rows)

result = np.zeros_like(matrix)

for j in range(1, num_cols):
    for n in range(1, num_rows):
        
        term1 = coeff1[n] * matrix[n-1, j-1]
        term2 = coeff2[n] * matrix[n, j-1]
        term3 = coeff3[n] * matrix[n-1, j]
        
        result[n, j] = term1 + term2 + term3

Solution

  • Make it Fully vectorized :

    # Vectorized calculation for each j > 0
     # Prev row, prev col (term1)
    term1 = coeff1[1:, None] * matrix[:-1, :-1] 
    # Current row, prev col (term2)
    term2 = coeff2[1:, None] * matrix[1:, :-1]  
    # Prev row, current col (term3) 
    term3 = coeff3[1:, None] * matrix[:-1, 1:]  
    
    # Summing the terms to get the result matrix (vectorized)
    result[1:, 1:] = term1 + term2 + term3