Changes

Clubb JAX (#1351)

* Testing new porting shortcut by creating a jax version of advance_xm_wpxp.py - it wen't well. Also updating the porting shortcut with more tips.

* Small update to porting shortcut

* Iterative improvements to porting shortcut

* Using porting script to port advance_xp2_xpyp in one go - still some issues that seems to require rerunning to fix issues, which is sad but maybe the best we can do with such a simple method

* Little bug fixes

* Updates to make the ports create better jax in the sense that it is more jit-ready

* Adding a more general plan to also use CLUBB-JAX and using it on advance_xm_wpxp

* Cleanup

* More properly jax-izing the current jax ports

* More formatting of ported jax code

* More formatting of ported jax code

* Updates in pursuit of full jit-compatibility

* Adding jit compile test

* Fully jitting advance_xp2_xpyp

* Copying parameter_indices from CLUBB-JAX but subtracting 1 to avoid -1 in code

* Little update to port shortcut

* Jitting and formatting CLUBB-JAX

* Mostly formatting

* Updating jax code

* Improving stats outputting

* Compile mono flux range loops with JAX

* Improvement to jax mixing length port

* Organizing stats code

* Making jax work on GPU

* Improvements to prescribe forcings code

* Mainly prescribe forcings formatting work

* Small improvements, getting rid of unneccesary weird jit statement

* Radiation porting completed. Bugsrad and silhs radiation were
intentionally left out for now.

* Removing stale code and adding a proposal.md file detailing a potential future feature to make the python_api work with current jax code

* Proper port of model flags and parameters and constants and such - those there were all kinda put into one file before

* new jenkinsfile added

* Tiny cleanup and adding back some stats outputs for microphysics that we need for some reason

* limited to only one process. Find out built time

* made jenkins jax tests faster by running them not so extensively

* Mainly working on README, also some other tiny things: changing default test stats list to standard instead of all_stats, and renaming run_jax.sh to run_jax_wrapper.sh

---------

Co-authored-by: Daniel Ebel <debel@uwm.edu>