[Python] Nested for-loop alternative

I have a python code that calculates a function F(X), where both F and X are arrays of the same shape. F(X) uses another function called from a package that only accepts a scalar as an argument, but I need to calculate function for all values in X. Using nested for-loops is very time expensive. Is there a way i could optimize this code, or else other alternatives to looping over the dimensions of X?

import numpy as np

def F(X):
"""
Calculates a matrix F of the same shape as X, using nested for-loops
Args:
X (np.ndarray): Input array, likely representing angles or similar.
Returns:

F (np.ndarray): An array of values with the same shape as X.
"""
matrix = np.zeros(np.shape(X))

for i in range(np.shape(X)[0]):
for j in range(np.shape(X)[1]):

x_ij = X[i,j]

P = external_function(x_ij) # This is my external function that accepts a scalar and returns a 4x4 matrix
f = - P[0, 1] / P[0, 0]
matrix[i,j] = f

return matrix

def external_function(x):
P = np.zeros((4,4))
for i in range(4):
for j in range(4):
for k in range(0,10):
P[i,j] = np.cos((x + i + j +k))
return P

X = np.random.rand(100, 100)
F(X)


I tried an using np.vectorize:

def F_opt(X):
"""
Calculates a matrix F of the same shape as X, by vectorizing the
scalar-only 'function'.

Args:
X (np.ndarray): Input array, m x n

Returns:
F (np.ndarray): An array of values with the same shape as X.
"""

# Define the scalar operation that needs to be applied to each element
def scalar_operation(x_ij):

P = external_function(x_ij) # This is my external function that accepts a scalar and returns a 4x4 matrix

f = - P[0, 1] / P[0, 0]
return f

vectorized_scalar_operation = np.vectorize(scalar_operation, otypes=[float])

F = vectorized_scalar_operation(X)

return F



but i arrived at a very similar performance:

import time

X_large = X = np.random.rand(500, 500)

start_time = time.time()
F_large_optimized = F_opt(X_large)
end_time = time.time()
print(f"Optimized F_optimized took: {end_time - start_time:.4f} seconds")


start_time = time.time()
F_large_original = F(X_large)
end_time = time.time()
print(f"Original F_original took: {end_time - start_time:.4f} seconds")


yielding : Optimized F_optimized took: 57.4816 seconds Original F_original took: 64.0786 seconds

Continue reading...