search

renecotyfanboy / jaxspec

jaxspec is an X-ray spectra Bayesian analysis package, relying on JAX to enable just in time compilation
https://jaxspec.readthedocs.io/en/latest/
MIT License
23 stars 1 forks source link

Transition to Equinox as a backend #73

Open renecotyfanboy opened 1 year ago

renecotyfanboy commented 1 year ago

haiku won't be improved in the future, but will be maintained to be compatible with newer JAX releases. It would be great to consider switching to equinox as a backend for JAXspec, since it is much lighter and would lead to less spaghetti code in the engine.

renecotyfanboy commented 1 year ago

After some tryouts, I figured out a nice code architecture which doesn't rely on weird metaclass as it is done right now. Equinox modules are great and can be instantiated unlike haiku modules. However I can't get it to work right now since I'm not familiar yet with this library.

UML diagram

Here, the natural model building SpectralModel = Phabs() * ( Powerlaw() + Blackbody()) is handled by the compositions defined in the Spectral model class. Individual components are built with a single node graph in the ModelComponent class constructor. I'll dive into this later!

patrick-kidger commented 1 month ago

I recently came across this library. I'd love to see Equinox used as the backend :)

renecotyfanboy commented 1 month ago

Wow, hi Patrick, I am very honoured to see you there. This project has been in the pipeline for a while, I've been playing around with equinox some time ago and it clearly seems feasible to transpose the internal logic of model building based on it. I just need to practise with it a bit more to get used to it, and that should be no problem!