Numba interface

Hi Yoel,

Per our discussion, I have created a tentative wrapper around the functions in the fluids library and the chemicals library. Numba is never going to work perfectly for all cases, and optimizations to make numba perfect sometimes make CPython and especially PyPy slower. However I am very interested in it for what uses it can have. All I am looking to target with numba at this point is computation functions, not lookup functions. It seems anything with a dictionary or a string is probably going to get slower as well, often not possible to compile.

I believe you have more experience with it than I do, so I especially welcome your thoughts.

Here are some examples of the current wrapper. There are a few more tests.

import chemicals.numba
import chemicals.numba_vectorized

chemicals.numba.Antoine(100.0, 8.7687, 395.744, -6.469) # speed

chemicals.numba_vectorized.Psat_IAPWS(np.linspace(100.0,200.0, 3)) # vectorized computation speed

I find that a function has to be pretty slow to benefit from a single-point numba version of a function; but for functions with vector arguments, at 10 elements it is already showing significantly improved performance, compared to numpy also.

I think for some applications numba really wants a custom function as well, not to wrap the existing one. The optimizations of e.g.

def zs_to_ws(zs, MWs):
    cmps = range(len(zs))
    ws = [zs[i]*MWs[i] for i in cmps]
    Mavg = 0.0    # Cannot use sum and list comprehension with numba; otherwise Mavg = 1.0/sum(ws)
    for v in ws:
        Mavg += v
    Mavg = 1.0/Mavg
    for i in cmps:
        ws[i] *= Mavg
    return ws

Make the numba function return a list, but what we want there is a numpy array. It is not as elegant but I have no problem writing a few duplicate functions.

@numba.njit
def zs_to_ws(zs, MWs):
    ws = zs*MWs
    Mavg = 1.0/np.sum(ws)
    ws *= Mavg
    return ws 

It seems some functions need to be re-written in a more verbose way to work with numba also, for example https://github.com/CalebBell/fluids/blob/master/fluids/friction.py#L2917

I consider this all pretty preliminary.

[yoelcortes]

Sounds good! I’ll see how I add the new jitclasses over the weekend, I’ll probably try the macro to remove the slots and repr methods, hopefully have less code duplication. Thanks!

[CalebBell]

Hi Yoel, I have responded to your concerns in-line for clarity:

When we would like make changes for numba compatibility, we can just make this switch in the “IS_NUMBA” block. Correct We can adjust code with macros and try-except blocks to make them compatible with numba. Correct We can make new code in the numba.py module for any difficult cases. Correct No need to worry about the use of functions within functions as a completely new module is created with all __funcs as part of their namespace. Correct

Exceptions to the “one size fits all” can stack up (making the numba module too large)

Yep. I am very open to having more files and/or a numba_interface folder but for now let’s keep it in the same file and try to be clear about where the lines may be drawn in the future.

When there is enough junk for 3 or 4 numba-specific files I think that would be a good time to revisit the issue and make the change if needed. Right now it feels like it will just add overhead in my mind but I agree it may be needed.

The IS_NUMBA sections may become big in some modules. Also, when making improvements, we will need to worry about both the numba and non-numba sections as one whole.

Yep. It already sucks 😦. It sucks less than complete code duplication though.

The addition of try-except blocks will make both original functions and numba compilled functions slower… although only significant for small functions.

The macro can be used to avoid this problem; so far I haven’t needed to slow anything down, and the module level try/except IS_NUMBA block only fails once per block while loading. Future less-experienced collaborators may feel less motivated to contribute when they read code that has been patched up for numba.

Fair point; the macros and having a second implementation are things I hope to avoid as much as possible. Most of the time it seems possible to code to support PyPy, Numba, and CPython very well.

The high level code is not too explicit and possibly hard to debug in the future. I’m sure you could handle it easily tho.

It will get better in time; for now the only tests that matter to me are the functions running with numba and getting the same result. The only happy thing about numba for is is that it is fast, the rest of it sucks.

I hope to keep the numba and regular code the mostly same, and in one place. This means the documentation is in one place, and the function signature, and when editing the function you don’t have to switch between different files. But this will not always be possible and this is all very experimental 😃 Still, I am quite please with all the things I have already made work.