Implement SFTtoAge in Jax and use vmap

[ufuk-cakir]

The conversion from stellar formation time to age in this line

rubix/rubix/utils.py

Line 105 in 9bdcd8a

# TODO maybe implement this in JAX?

should be implemented in Jax using the cosmology function

rubix/rubix/cosmology/base.py

Line 110 in 9bdcd8a

def lookback_to_z(self, redshift):

and then we can vmap over all the stars particles.

This would be nice because this conversion can take some time if we have a lot of stars in the galaxy, and implementing in jax and vmap will give a huge performance boost