Why SCICO?#

Advantages of JAX-based Design#

The vast majority of scientific computing packages in Python are based on NumPy and SciPy. SCICO, in contrast, is based on JAX, which provides most of the same features, but with the addition of automatic differentiation, GPU support, and just-in-time (JIT) compilation. (The availability of these features in SCICO is subject to some caveats.) SCICO users and developers are advised to become familiar with the differences between JAX and NumPy..

While recent advances in automatic differentiation have primarily been driven by its important role in deep learning, it is also invaluable in a functional minimization framework such as SCICO. The most obvious advantage is allowing the use of gradient-based minimization methods without the need for tedious mathematical derivation of an expression for the gradient. Equally valuable, though, is the ability to automatically compute the adjoint operator of a linear operator, the manual derivation of which is often time-consuming.

GPU support and JIT compilation both offer the potential for significant code acceleration, with the speed gains that can be obtained depending on the algorithm/function to be executed. In many cases, a speed improvement by an order of magnitude or more can be obtained by running the same code on a GPU rather than a CPU, and similar speed gains can sometimes also be obtained via JIT compilation.

The figure below shows timing results obtained on a compute server with an Intel Xeon Gold 6230 CPU and NVIDIA GeForce RTX 2080 Ti GPU. It is interesting to note that for FiniteDifference the GPU provides no acceleration, while JIT provides more than an order of magnitude of speed improvement on both CPU and GPU. For DFT and Convolve, significant JIT acceleration is limited to the GPU, which also provides significant acceleration over the CPU.

Timing results for SCICO operators on CPU and GPU with and without JIT