🔒 `numba` for `scipy.optimize.root_scalar()`

[crisp cypress]

I am attempting to utilize numba.njit()/numba.jit() on my code to speed up my mathematical operations. I want to NJIT/JIT compile the following type of function, which I call repeatedly.

class MeshPoint():
    # @numba.njit()
    def internal(p1data, p2data, p3data):
        # p1 is below, p2 is above
        x1, x2 = p1data[0], p2data[0]
        y1, y2 = p1data[1], p2data[1]
        M1, M2 = p1data[2], p2data[2]
        theta1, theta2 = p1data[3], p2data[3]

        mu1, mu2 = AeroFunc.M2mu(M1), AeroFunc.M2mu(M2)
        
        nu1, nu2 = AeroFunc.M2nu(M1), AeroFunc.M2nu(M2)
        nu3 = (nu1 + nu2)/2 - (theta1 - theta2)/2
        theta3 = (theta1 + theta2)/2 - (nu1 - nu2)/2
        
        M3 = AeroFunc.nu2M(nu3)
        mu3 = AeroFunc.M2mu(M3)

        phi1 = (theta1 + theta3 + mu1 + mu3)/2
        phi2 = (theta2 + theta3 - mu2 - mu3)/2

        x3 = x1 * math.tan(phi1) - x2 * math.tan(phi2) + y2 - y1
        x3 /= math.tan(phi1) - math.tan(phi2)

        y3 = y1 + math.tan(phi1) * (x3 - x1)
        
        p3data[0] = x3
        p3data[1] = y3
        p3data[2] = M3
        p3data[3] = theta3
        
        return p3data

However, within this function, I am utilizing other aerodynamic functions, which are contained within AeroFunc():

class AeroFunc():
    @numba.njit()
    def M2nu(M): # Prandtl-Meyer Function
        # https://en.wikipedia.org/wiki/Prandtl-Meyer_function
        A = np.sqrt((gamma + 1)/(gamma - 1))
        B = np.sqrt(M**2 - 1)
        return A * np.arctan(B/A) - np.arctan(B)

    @numba.jit()
    def nu2M(nu): # Inverse Prandtl-Meyer Function
        # https://en.wikipedia.org/wiki/Prandtl-Meyer_function
        return root_scalar(lambda M: nu - AeroFunc.M2nu(M), x0 = M_inl, x1 = M_exi).root
    
    @numba.njit()
    def M2mu(M): # Mach Angle Function
        # https://en.wikipedia.org/wiki/Mach_wave
        return np.arcsin(1/M)

[velvet quartz]

@crisp cypress

Python help channel opened

Remember to:

• Ask your Python question, not if you can ask or if there's an expert who can help.
• Show a code sample as text (rather than a screenshot) and the error message, if you've got one.
• Explain what you expect to happen and what actually happens.

:warning: Do not pip install anything that isn't related to your question, especially if asked to over DMs.

Closes after a period of inactivity, or when you send !close.

[crisp cypress]

I have no issue with compiling AeroFunc.M2nu() or AeroFunc.M2mu(), but due to the use of scipy.optimize.root_scalar() within AeroFunc.nu2M(), I get the following traceback:

Traceback (most recent call last):
  File "[filepath]/diverging_nozzle_convergence.py", line 189, in <module>
    wall_x, wall_y, wall_series = expansion_section(y_inl, r_exp, M_inl, M_exi)
                                  ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "[filepath]/diverging_nozzle_convergence.py", line 135, in expansion_section
    wall_series[n, 2] = AeroFunc.nu2M(nu)
                        ^^^^^^^^^^^^^^^^^
  File "/opt/homebrew/lib/python3.12/site-packages/numba/core/dispatcher.py", line 468, in _compile_for_args
    error_rewrite(e, 'typing')
  File "/opt/homebrew/lib/python3.12/site-packages/numba/core/dispatcher.py", line 409, in error_rewrite
    raise e.with_traceback(None)
numba.core.errors.TypingError: Failed in nopython mode pipeline (step: convert make_function into JIT functions)
Cannot capture the non-constant value associated with variable 'nu' in a function that may escape.

File "Diverging Nozzle/diverging_nozzle_convergence.py", line 25:
    def nu2M(nu): # Inverse Prandtl-Meyer Function
        <source elided>
        # https://en.wikipedia.org/wiki/Prandtl-Meyer_function
        return root_scalar(lambda M: nu - AeroFunc.M2nu(M), x0 = M_inl, x1 = M_exi).root

@wind totem

[wind totem]

well you can't use njit with that stuff was kind of my point

I thought that numba would still compile what it can even without njit though

[nimble oyster]

@wind totem how did you get so good

[crisp cypress]

Wait, what do you mean?

[crisp cypress]

Which is why my intuitive response is to just create a self-contained scalar root finder

[wind totem]

you can't use njit and call a function like scipy.optimize.root_scalar I believe

[crisp cypress]

Unless there is a package that extends the scope of numba, like rocket-fft does for numpy.fft-functions.

[wind totem]

I know the numpy linalg functions just work with numba njit

Here's an analysis I did on the performance of different implementations Dijkstra's algorithm in a randomized 3d graph using numba vs pure python

[crisp cypress]

Assuming these are log plots, that's extremely significant

[wind totem]

the final implementation used np functions

so I don't know why you need extensions for fft

[crisp cypress]

Oh I used the fft a lot in another program

[velvet quartz]

Python help channel closed

This help channel has been closed and it's no longer possible to send messages here. If your question wasn't answered, feel free to create a new post in #1035199133436354600. To maximize your chances of getting a response, check out this guide on asking good questions.