ENH: Enable vectorizaton / Jax

[hmgaudecker]

Is your feature request related to a problem?

We merged #425 before Jax is fully usable in GETTSIM. This issue is to remind ourselves of what needs to be done still.

Describe the solution you would like to see

Core thing is to change src._gettsim.functions_loader.vectorize_func to something like:

def _vectorize_func(func):
    backend = "jax" if USE_JAX else "numpy"
    return make_vectorizable(func, backend)

Doing so currently leads to failures because some elements cannot be vectorized yet. Examples found so far:

• calls of float(x) do not work on arrays
• piecewise_polynomial assumes scalar input for x (even though numpy.searchsorted is more general)
• The entire piecewise_functions module is using numpy instead of numpy_or_jax for now, which is unlikely to work with the jax backend.

Once these things are fixed, the above version of _vectorize_function should work and it will likely be faster than the current implementation already.

If, in addition, we knew which function arguments are to be vectorized over and which are not (some ideas for a decorator in GEP 6), we can call jax.vmap and make it even faster. This is more of a reminder, probably to be tackled in a separate PR.

[hmgaudecker]

One thing to think about would be to leave the numpy backend as is, i.e., write on scalars and call np.vectorize. Seems like this is easier on people not speaking Python as their lingua franca.

[ChristianZimpelmann]

One thing to think about would be to leave the numpy backend as is, i.e., write on scalars and call np.vectorize. Seems like this is easier on people not speaking Python as their lingua franca.

I thought this is the plan. What would be the alternative?

[hmgaudecker]

Yes, sorry, that was misleading and only comprehensible based on understanding the evolution of my line of thinking during the day.

The way @timmens has implemented make_vectorizable is that it rewrites the source code so that you can call all functions with vectors natively. Users will not notice they are not writing them for scalars, unless they use constructs that are not supported by the functions.

In all likelihood, the code will run orders of magnitude faster with that setting.

So we'll need to balance two goals

• making it easy on people developing GETTSIM, but who are not Python experts
• making GETTSIM run reasonably fast out of the box

The solution will probably be to allow for both kinds of backends, leaving the current state of things the default and making the numpy-vectorized-out-of-the-box option an opt-in, just as the Jax backend. We'll just need to develop some useful terminology.

[hmgaudecker]

Once that is done, remove renewed casting to numpy.datetime64 in demographic_vars.py / alter_monate. Currently required because np.vectorize() will cast to object:

> /mambaforge/envs/gettsim/lib/python3.11/site-packages/numpy/lib/function_base.py(2412)

2400         def _vectorize_call(self, func, args):                                                                                                                                                                                                                                  
[...]
2409                 # Convert args to object arrays first                                                                                                                                                                                                                           
2410                 inputs = [asanyarray(a, dtype=object) for a in args]