Include phase corrections in JAX version

Description

The phase correction, which allows for fluctuations in the zenith angle, is currently not implemented in JAX. This is one large feature that is currently missing, and should be the first to be implemented towards v0.8.

Checklist

• Implement functionality
• Verify with same geometry
• Verify for varying geometries
• Verify with numpy version