[WIP] jit/core redesign

[s-m-e]

Objective

This PR serves as foundation for both orbit and state arrays. It focusses on functionalities provided by the core module and their invocation. All relevant core functions are designed to work equally on CPUs and GPUs as either universal functions or generalized universal functions. As a "side-effect", all relevant core functions allow parallel operation with full broadcasting semantics.

Summary

All dependencies of Orbit and state classes towards core are refactored as follows:

• Functions decorated by either numba.vectorize and numba.guvectorize serve as the only interface between regular uncompiled Python code and core
• Functions decorated by numba.vectorize and numba.guvectorize only call functions decorated by numba.jit/numba.cuda.jit
• Functions decorated by numba.jit/numba.cuda.jit can only call each other.
• Functions decorated by numba.vectorize, numba.guvectorize and numba.jit/numba.cuda.jit:
  • are only allowed to depend on Python's standard library's math module, but not numpy - except for certain details like enforcing floating point precision
  • are fully typed, loosely following numba semantics via shortcuts

The above mentioned "hierarchy" of decorators is imposed by CUDA-compatibility. While functions decorated by numba.jit (targets cpu and parallel) can be called from uncompiled Python code, functions decorated by numba.cuda.jit (target cuda) are considered "device functions" and can not be called by uncompiled Python code directly. They are supposed to be called by CUDA-kernels (or other device functions) only (slightly simplifying the actual situation as implemented by numba). If the target is set to cuda, functions decorated by numba.vectorize and numba.guvectorize become CUDA kernels.

Eliminating numpy as a dependency serves two purposes. While it also contributes to CUDA-compatiblity, it additionally makes the code significantly faster on CPUs.

New decorators are introduced, wrapping numba.jit, numba.cuda.jit, numba.vectorize and numba.guvectorize, centralizing compiler options and target switching (cpu, parallel or cuda) as well as simplifying typing:
- vjit: Wraps numba.vectorize. Functions decorated by it carry the suffix _vf.
- gjit: Wraps numba.guvectorize. Functions decorated by it carry the suffix _gf.
- hjit: Wraps numba.jit or numba.cuda.jit, depending on compiler target. Functions decorated by it carry the suffix _hf.
- djit: Variation of hjit with fixed function signature for user-provided functions used by Cowell

All mentioned wrappers are found in core.jit.

As a result of name suffixes, a number of core module functions have been renamed making the package intentionally backwards-incompatible. Functions not yet using the new infrastructure can be recognized based on lack of suffix. core functions dynamically generating (and compiling) other functions carry _hb, _vb and _gb suffixes.

Math

The former _math module has become a first-class citizen as core.math, fully compiled by the above mentioned infrastructure.

All compiled code now enforces a single floating point precision level, which can be configured by users. The default is FP64 / double precision. For simplicity, the type shortcut is f. Additional infrastructure can be found in core.math.ieee754.

core.math contains a number of replacements for numpy operations, mostly found in core.math.linalg. All of those functions do not allocate memory and are free of side-effects including a lack of changes to their parameters. 3D vectors are expressed as tuples (type shortcut V, replacing Tuple([f,f,f])). Matrices are expressed as tuples of tuples (type shortcut M, replacing Tuple([V,V,V])).

core.math also replaces (some) required scipy functions:
- scipy.interpolate.interp1d is replaced by core.math.interpolate.interp_hb. It custom-compiles 1D linear interpolators, embedding data statically into the compiled functions.
- scipy.integrate.solve_ivp, scipy.integrate.DOP853 and scipy.optimize.brentq are replaced by core.math.ivp.

Style

core modules now explicitly export APIs via __all__.

Settings

This PR introduces a new settings module. It currently serves to control compiler options, e.g. the compiler target (cpu, parallel and cuda). Settings can be switched by either setting environment variables or importing the settings module before any other (sub-) module is imported.

Logging

This PR introduces basic logging functionality. Logger name equals package name. The logger can also be imported from the new debug module. At the moment, it only logs compiler status messages and issues.

Blocking numba issues for CUDA-compatibility

• math.nextafter: nextafter (via both math and numpy) missing for CUDA numba/numba#9435
• Signatures for function references: CUDA: Providing a signature for a function reference when eagerly compiling does not work numba/numba#9398
• Return type mismatch: On target cuda, return type of function may not match type given in signature string numba/numba#8400
• Multiple outputs for guvectorize: Allow multiple outputs for guvectorize on CUDA target numba/numba#8303
• Missing __name__ attribute for 'CUDAUFuncDispatcher': 'CUDAUFuncDispatcher' object has no attribute '__name__' numba/numba#8272

Non-blocking numba issues for CUDA-compatibility with workaround present

• Tuple of functions: Tuple of functions (with identical signature) does not compile for CUDA numba/numba#9420
• Keyword arguments: Add support of keyword argument to Dispatchers from numba.cuda.jit numba/numba#7870

Non-blocking numba issues unrelated to CUDA with workaround present

• New size variable i.e. () -> (n) in output argument: @guvectorize not accepting new size variable (i.e. () -> (n)) in output argument numba/numba#2797
• Building the cache in parallel: Cache causes Segmentation Faults when generated in parallel numba/numba#4807

TODO

• Code-style, doc strings and comments in core.math.ivp
• Compiled non-linear interpolators
• Documentation

---

📚 Documentation preview 📚: https://hapsira--7.org.readthedocs.build/en/7/