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
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