Key Differences with NumPy version

Here we highlight the key differences between the JAX and NumPy versions of the template synthesis package, as there are some important differences in the implementation that one needs to consider when using either module.

  1. The JAX version does not have a SliceSynthesis class.

To fully utilize the potential of the JAX framework, we opt to store all slices in a single array. While this allows for better performance when mapping the template, this does mean that the TemplateSynthesis class is less flexible than the NumPy version. For example, it is not possible to access the amplitude and phase spectra of a single slice. Instead, one can access this through slice indexing. A more user-friendly method is in development.

  1. The JAX version can lead to memory overloads when using large arrays.

As JAX internally uses a just-in-time compilation (JIT) approach, it can lead to some memory issues when using large arrays ( i.e. origin showers with many antennas) . This is due to the fact that JAX will try to compile the function with the largest array size, which can lead to memory issues. To mitigate this, we recommend users to use the origin shower library, which stores origin showers with minimal memory constraints.

  1. The JAX version relies on the following external modules:

    • jax: The JAX library is used for automatic differentiation and GPU/TPU acceleration. This is the core of the JAX version which allows for the performance improvements. In particular, many functionalities are written in jax.numpy, which is the numpy API for JAX.

    • jax_radio_tools : This is a library that provides tools for radio signal processing in JAX, similar to the already widely used radio_tools.

    • interpax : An interpolation package for JAX. More details are given here.

    These libraries are necessary for the JAX version to function properly, and are not required for the NumPy version.

  2. Syntactic differences

There are several naming conventions that differ between the JAX and NumPy versions. While most syntactic differences have been resolved, there may be minor differences in the accessing of variables. Please raise an issue on GitLab if you encounter any issues regarding this.