New version of JAXified template synthesis!
This MR is a reimplementation of !11 (closed) , after rebasing the changes.
This version of the JAXified template synthesis is ready to be used and can be merged with the beta release.
A few major updates:
- only JIT compile the heavy functions, which are called outside (there is a wrapper in the class)
- make everything else that do not rely on it being dynamic to be in numpy arrays to make it static
- in particular remove dependency of ModelMeta since its not required to make everything PyTree readable, just important functions
- Include dependency with jax_radio_tools
- implemented filtering & resampling functionality with SlicedShower
TODOs:
- check and make the naming of attributes consistent with the numpy version (#17 (closed) )
- implement the CoREASHDF5 reader into the JAX version
- Implement some way to access slice-specific variables (#3 )
- Make the template saving consistent with numpy version (#17 (closed) )