IAP GITLAB

Skip to content

Template generation in JAX version has large memory overhead

With the current setup of how origin showers are generated in the JAX version, there is a large memory overhead of > 30 GB when loading the origin shower. This is most likely due to how the arrays are stored, as the emission traces are stored as (Npol x Nant x Nsamples x Nslices), which JAX (probably) does some internal compilation.

While this can be (and is currently) mitigated by pre-generating templates on a workstation with enough CPU resources, this is not practical in the long run if we want e.g. perform parity checks with NumPy versions, observe dependence with origin showers, or for simple plotting purposes.

What I have observed so far is that the initialisation of the SlicedShower was the major overhead, and the actual template generation is not as limited via memory. So we would only need to think how to efficiently import the origin coreas hdf5 files.