For loop is slowing down performance. Any alternatives?

Question

I am trying to improve the performance of the function below

import numpy as np
import time 

r_0 = 0.1

drt_measurement = 9.999999999999991269e+04  3.305791191233514031e-02
9.083278409243831993e+04    6.686534998595229651e-02
1.569368719222600794e-01    7.259131510582693403e-02
1.425501300345730638e-01    1.423914344911824392e-01
1.294822518377939102e-01    2.296111516784170581e-01
1.176123342498509028e-01    3.236305655613548882e-01
1.068305576352440167e-01    4.110383626909880905e-01
9.703716976156935570e-02    4.774860553090806148e-01
8.814156289893920748e-02    5.096622216542943118e-01
8.006143552369422711e-02    4.972726694760475352e-01
7.272203087054407433e-02    4.349122174992410828e-01
6.605544528827768380e-02    3.237240623831713071e-01
5.999999999999999778e-02    1.727623254500962879e-01


def backcalculate_impedance_from_drt_freq(drt_measurement, r_0):
    
    # Extract frequency and DRT from drt_measurement matrix
    drt_frequency = drt_measurement[:,0]
    drt = drt_measurement[:,1]
    
    # Frequency to angular frequency, ω, and scaling factor
    ω_drt = -np.log(2*np.pi*drt_frequency)
    ds = ω_drt[2] - ω_drt[1] # ds is a constant
    
    # Unscaling the DRT
    unscaled_drt = drt*ds
    
    # Initializing Zreal and ZImag vectors
    ZReal = np.empty((len(drt_frequency),1))
    ZImag = np.empty((len(drt_frequency),1))
    
    print("\nBackImpedance For loop")
    tic = time.time()
    # Solve for ZReal and ZImag
    for i in range(len(drt_frequency)):
        K1 = 1 / (1 + np.exp(2 * (-ω_drt[i] + ω_drt))) * unscaled_drt
        K2 = -np.exp(-ω_drt[i] + ω_drt) * K1
        ZReal[i,0] = sum(K1) + r_0
        ZImag[i,0] = sum(K2)
    
    toc = time.time()
    diff = toc - tic
    print(f'Executed in {diff} seconds or {diff*1000} milliseconds')  
    
    # Save Backcalculated Impedance into a (x,3) matrix
    drt_frequency = drt_frequency[:,None]
    backcalculated_impedance = np.hstack((drt_frequency,ZReal,ZImag))
    
    return backcalculated_impedance

Currently I am getting about ~10ms each time this function is called with the for loop accounting for the majority of the time which really slows down the overall execution. Is there a better alternative to the for loop or a way to vectorize its operations?

Answer 1

The function can be rewritten without the for loop by using vectorization:

def backcalculate_impedance_from_vectorized(drt_measurement, r_0):
    drt_frequency, drt = drt_measurement[:, 0], drt_measurement[:, 1]
    ω_drt = -np.log(2 * np.pi * drt_frequency)
    ds = ω_drt[2] - ω_drt[1]
    unscaled_drt = drt * ds
    ω_drt_expanded = ω_drt[:, np.newaxis]
    K1 = 1 / (1 + np.exp(2 * (-ω_drt_expanded + ω_drt))) * unscaled_drt
    ZReal = K1.sum(axis=1) + r_0
    ZImag = (-np.exp(-ω_drt_expanded + ω_drt) * K1).sum(axis=1)
    return np.hstack((drt_frequency[:, np.newaxis], ZReal[:, np.newaxis], ZImag[:, np.newaxis]))