🚀[FEA]: make zenith angle code device agnostic and hide warnings

[nbren12]

Is this a new feature, an improvement, or a change to existing functionality?

Improvement

How would you describe the priority of this feature request

Medium

Please provide a clear description of problem you would like to solve.

In earth2mip we have needed to make some hacks to get the zenith angle code working with jax and numpy. There are also some noisy NaN warnings in arccos. The out of domain error occurs for polar latitudes where the sun never sets/rises. We handle this case explicitly so there is no need to pass a warning to the user.

This came up when @akshaysubr was reviewing some code:

• https://github.com/NVIDIA/earth2mip/pull/153#discussion_r1433353233
• https://github.com/NVIDIA/earth2mip/pull/153#discussion_r1433353821

Describe any alternatives you have considered

No response