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>

Pushing loop into advance_microphys

Little more loop push

More basic loop pushing

More loop pushing, also found what seems to be a bug with additional columns, theres code that uses only the first one. Im leaving this in for now to maintain BFBness, but labeled the problem areas with TODOs

Potentially final batch of loop pushes