Python efficiency - nested loop structure with matrices

Question

I am new to python, but have been coding in MATLAB for years. I've been working to code a specific problem, which requires nested for loops for multiple matrices to solve a series of equations. I knew how I would code this in MATLAB and then used chatGPT to help convert into python syntax. I am using Spyder interface and python 3.9. Below shows the related code described here, and I would like to understand specific ways to make this run more efficiently. I've read previous answers for similar questions using itertools, but when applied, they don't seem to work.

import numpy as np

R1init=[]
R2init=[]
L1init=[]
L2init=[]
p1init=[]
p2init=[]
m1init=[]
m2init=[]
dVrinit=[]
dVlinit=[]

R1 = np.arange(50, 200.001, 2)
R2 = R1  # Since R1 and R2 are identical
L1 = -1*R1
L2 = np.arange(-50,-300.001,-10)

#converted values to be thousands, assuming this would help with matrix sizes.
dVl = 194329/1000  
dVr = 51936/1000  
dVg = 188384/1000  
DR = 0. 
DB = 0.

m1 = np.abs(dVl / R1) 
m2 = np.abs(dVr / L2)


j1 = 0
j2 = 0

for i in R1:
    for j in R2:
        for k in L1:
            for m in L2:
                for n in m1:
                    for q in m2:
                        p1 = ((j2*(1+q)-q)*m+j+dVr)/i
                        p2 = 1-j2*(1+q)+q-(i/m)*(1-j1*(1+n)+n-p1)+dVg/m
                        dVrchk = (q-(j2*q)-q)*m+(p1*i)-j+DR+DB
                        dVlchk =(j1-n+(j1*n))*i+k-(p2*m)
                        dVgchk = (1-j1-p1+n-j1*n)*i-(1-j2-p2+q-j2*q)*m
                        if 0<p2<1.05 and 0<p1<1.05 and dVl-100<dVlchk<dVl+100 and dVr-  100<dVrchk<dVr+100:
                            R1init.append(i)
                            R2init.append(j)
                            L1init.append(k)
                            L2init.append(m)
                            p1init.append(p1)
                            p2init.append(p2)
                            m1init.append(n)
                            m2init.append(q)
                            dVrinit.append(dVrchk)
                            dVlinit.append(dVlchk)

Answer 1

I would suggest using Numba to speed this up, and altering the order that you check the conditions in. That is sufficient to solve this in under a second.

For example, your condition on p1 does not involve the variables n or k, so you don't need to loop over all combinations of these to check them.

Similarly, p2 doesn't depend on k, so it can be done outside the k loop.

This does have the disadvantage that it changes the order of the solutions, but I'm guessing you don't strongly care what order they're returned in.

These optimizations allow this to be done in under a second.

import numba as nb
from numba.typed import List


@nb.njit()
def search_inner(R1, R2, L1, L2, m1, m2):
    dVl = 194329/1000
    dVr = 51936/1000
    dVg = 188384/1000
    DR = 0. 
    DB = 0.

    R1init = List()
    R2init = List()
    L1init = List()
    L2init = List()
    p1init = List()
    p2init = List()
    m1init = List()
    m2init = List()
    dVrinit = List()
    dVlinit = List()
    j1 = 0
    j2 = 0

    for i in R1:
        for j in R2:
            for q in m2:
                for m in L2:
                    # i j q m
                    p1 = ((j2*(1+q)-q)*m+j+dVr)/i
                    if not (0 < p1 < 1.05):
                        continue
                    for n in m1:
                        # q i m n p1
                        p2 = 1-j2*(1+q)+q-(i/m)*(1-j1*(1+n)+n-p1)+dVg/m
                        if not (0 < p2 < 1.05):
                            continue
                        for k in L1:
                            # i j q m p1
                            dVrchk = (q-(j2*q)-q)*m+(p1*i)-j+DR+DB
                            if not (dVr - 100 < dVrchk < dVr + 100):
                                continue
                            # n i k m p2
                            dVlchk =(j1-n+(j1*n))*i+k-(p2*m)
                            if not (dVl - 100 < dVlchk < dVl + 100):
                                continue
                            dVgchk = (1-j1-p1+n-j1*n)*i-(1-j2-p2+q-j2*q)*m
                            # dVgchk is ignored here - Bug?
                            R1init.append(i)
                            R2init.append(j)
                            L1init.append(k)
                            L2init.append(m)
                            p1init.append(p1)
                            p2init.append(p2)
                            m1init.append(n)
                            m2init.append(q)
                            dVrinit.append(dVrchk)
                            dVlinit.append(dVlchk)

    ret = {
        'R1init': R1init,
        'R2init': R2init,
        'L1init': L1init,
        'L2init': L2init,
        'p1init': p1init,
        'p2init': p2init,
        'm1init': m1init,
        'm2init': m2init,
        'dVrinit': dVrinit,
        'dVlinit': dVlinit,
    }
    return ret
def search():
    #converted values to be thousands, assuming this would help with matrix sizes.
    dVl = 194329/1000
    dVr = 51936/1000

    R1 = np.arange(50, 200.001, 2)
    R2 = R1  # Since R1 and R2 are identical
    L1 = -1*R1
    L2 = np.arange(-50,-300.001,-10)

    m1 = np.abs(dVl / R1)
    m2 = np.abs(dVr / L2)

    ret = search_inner(R1, R2, L1, L2, m1, m2)
    # Unwrap typed lists into arrays
    ret = {k: np.array(v, dtype='float64') for k, v in ret.items()}
    return ret