Eliminate or fix numba jit instances without nopython=True

[SeanMatthewNolan]

Numba jit compiler cannot compile certain Python types or functions. If attempted without care, jit compiling can actually slow the code by requiring the code to exit the compiled code into Python. The nopython=True keyword argument will cause numba to throw an error instead of slowing the code.

[SeanMatthewNolan]

When I get a chance, I'll incorporate numba well. Numba jit used naively will usually result in worse runtime because one has to make sure that the functions are not repeatedly compiled and that they actually compile. I was successful in incorporating numba jit with a custom shooting solver and the speed-ups are great when careful so it is definitely worthwhile.

[msparapa]

Do you have any more info on this? From trying to implement custom functions, I understand this a bit better now. I believe all jit instances where jit(nopython=True) are equivalent to njit(). Do we want our solver to fall back on the nopython=False?

[SeanMatthewNolan]

We want to avoid falling back because on python mode because numba in python mode is significantly slower than just python. I think a try except statement would be better.

[msparapa]

Alright so maybe a try: njit, then if that doesn't work just use the function spit out by eval()? The eval turns a lambda str into a function, but then we use jit() on it. Skip jit() altogether if nopython=True fails?

[SeanMatthewNolan]

I think so, but it may not actually fix the underlying problem of jit not working. Honestly, last I checked, I think jit is doomed to fail every time. I will get to fixing/diagnosing it by the end of week if you are willing to punt it to me. It's near the top of my to-do list.

[SeanMatthewNolan]

As I work on this and go in-depth to the Shooting algorithm, I suggest moving the derv_func, quad_func, bc_func arguments out of the solve method into the arguments for initializing the "shooter" object (is equally applicable to any like algorithm). From experience, it makes more sense because these don't generally change and the "shooter" object's solve method may be called repeatedly. This allows processes (#17) and jit compiled code to remain with the class, and overhead to generate these only is required once. The "solve" method should only take the arguments for the current guess and the values needed to define the BCs. @msparapa I can probably explain this better in person with showing how I did it in my own solver.

[msparapa]

Do you mean placing them as args in __new__? So `if args then solve; else return shooting? I think I can make this work more intuitively as a map. See #91. Also look at my documentation for Commutator() [here](https://beluga.readthedocs.io/en/latest/modules/liepack/liepack.html). In one of the examples, I "preload" the commutator's first argument with a lie algebra element. Then I can use that object as a function by inserting the second argument whenever. You can see this code in action inbeluga/liepack/liepack.py` for the dexpinv() function (line 212). I make adg by not giving Commutator all the required arguments, but then I use it repeatedly later on line 218 with a different input every time, but the first input is saved and always the same. I think its computationally efficient, looks good, and mathematically "correct" when posed like this.

[SeanMatthewNolan]

Hmm, kinda? I think I am thinking of something simpler. The arguments would go in init. It would look something like this:

Called once: tpbvp_solver = bvp_algo(deriv_func, quad_func, bc_func, *args)

Called at every continuation step: sol = tpbvp_solver.solve(solinit, cont_values)

[msparapa]

See the morphism class I just pushed in branch I91. I think my solution is more complicated at the core, but to a developer it should be simpler. Check out how I changed Commutator() from what it is in the master branch to what it is now in branch I91 (there's still a minor bug with repeated usage that I'll fix later). So I'm imaging your block of code will now look like:

tpbvp_solver = algo(deriv_func, quad_func, bc_func, **any_other_kwargs)

some loop:
    sol = tpbvp_solver(solinit, cont_values, **more_kwargs)

but, you can also call it like

sol = bvp_algo(deriv_func, quad_func, bc_func, solinit, cont_values, **all_the_kwargs)

[msparapa]

I bumped the close and comment button midway through my thought 😞

So the bvp_algorithm behaves both as a function and a class that can store information, depending on how the developer calls it. This format might seem a tiny bit strange when calling ode45 and other engineering-style written functions, but for other functions written in a more mathematical-style, its much more similar to what the equations represent.

If you look at Commutator() in LiePack, I think I have the interface set up so that it should be extremely easy to build off of. I did include 1 thing thats redundant. We can use either class variables, like the

inputs = (LieAlg, LieAlg)
outputs = (LieAlg,)

or we can use annotations, which is the

def map(self, g: LieAlg, h: LieAlg) -> LieAlg:

In Python, the second example where I'm using annotations is primarily just for reading. Code is not necessarily supposed to be dictated by it. Which do you prefer?

[SeanMatthewNolan]

Maybe I am missing the benefit, but in my opinion, that just obfuscates what is happening in the case of the tpbvp solver. I prefer the simplicity of having a separate method on the object. It seems to be the clearest way and most in line with OOP to have objects with properties with methods that "do things". It's explicit that you are making a solver object with those values are held in that object. If you need to call everything in one line, you can use sol = bvp_algo(deriv_func, quad_func, bc_func, *args).solve(solinit, cont_values).

[msparapa]

@SeanMatthewNolan Can you check the latest commit I pushed for this issue? I'm having the solver try jit(nopython=True) for functions and if that fails, it will not use numba, but rather just a normal Python function. Is this what you were looking for?

[SeanMatthewNolan]

@msparapa Sorry for not responding earlier. Yeah, that's probably the best we can do because the function needs to be called to be compiled. We need to also rewrite deriv_func to be compiled too. From looking at the code, it looks like that isn't happening.

[msparapa]

That's a longer term goal of mine. Before, we were using a combination of pystache and Python code which caused issues when saving data, and I'm assuming parallelization as well but I haven't tried that yet. It looks like a mess right now, but uses less external packages and handles all the functions in a similar manner. I'd like it to eventually handle all functions the exact same way without even knowing if we're using icrm, traditional, or even direct methods.

For this task I think it's just sufficient to try and run the njit() code and if it fails, to stick with pure Python.