CODEWITHSUNDEEP
pythonnumpyloopslinear-algebra
Timo-m
Timo-m

Reputation: 3

How can I improve the computational efficiency of looping over numpy linear solve

I have a system of equations that I need to solve for a large number of variables. The code is more complicated than below, but the problem is as follow:

I have an iterator:

iterator= np.arange(2000)   # array!

And arrays dependent on a function of that iterator:

A_11 = function_A11(iterator)
A_12 = function_A12(iterator)
A_21 = function_A21(iterator)
A_22 = function_A22(iterator)

B_1 = function_B1(iterator)
B_2 = function_B2(iterator)

X = np.zeros(2, 2000)

for i, (A11, A12, A21, A22, B1, B2) in enumerate(zip(A_11, A_12, A_21, A_22, B_1, B_2):
    A = np.array([[A11, A12],[A21,A22]]) 
    B = np.array([B1, B2])
    X[:,i] = np.linalg.solve(A,B)

The method works but it is computationally expensive and I have the feeling that I should be able to optimize this, e.g. by working with 3D arrays. Does anyone have any suggestions?

Thanks!

Tim

Upvotes: 0

Views: 59

Answers (1)

hpaulj
hpaulj

Reputation: 231385

arr= np.arange(2000)   # array!

A_11 = function_A11(arr)
A_12 = function_A12(arr)
A_21 = function_A21(arr)
A_22 = function_A22(arr)

B_1 = function_B1(arr)
B_2 = function_B2(arr)

combine

AA = np.array([[A_11,A_12],[A_21, A_22]])

I expect AA.shape to be (2,2,2000).

A = AA.transpose(2,0,1)

to get a (2000,2,2) that solve_ivp can use.

I haven't tested this.

Upvotes: 1

Related Questions