Autograd -> Jax conversion

[jtgrasb]

Description

Convert autograd to jax.

Wavebot tutorial is working for the first optimization.

[jtgrasb]

Wavebot is working now when converted to jax. Working on updating the code so that only the necessary functions use jax.numpy while the rest use numpy.

[jtgrasb]

To do:

• Add in functions such as vmap, grad, and jit to increase code performance.
• Resolve MacOS issues on GitHub actions.

[jtgrasb]

I added just-in-time compilation to the optimization (objective function, constraints, and relative gradients) using jax.jit which should speed up the code. I also had to change the call to block_diag() in the mimo_transfer_mat() function to a revised function that is now jittable. There were a couple of assert statements that I commented out for now because assertions related to the dynamic value (some were able to be left in because they were only checking static properties, which is jittable).

Here is the computation time with and without jit for the AquaHarmonics parameter sweep cell:

Implementation	Time
Autograd	~ 38 s
Jax without jit	~130 s
Jax with jit	~20 s

Based on this issue thread, scipy does computation on numpy arrays which means the data type keeps getting converted back and forth between jax and numpy arrays. This is why the jax without jit has such a large computation time and the jax with jit computation time is not so much faster than autograd.

jax.vmap:

• Would need to write functions differently to vectorize. Would be good but a lot of work - save for future work.

To do:

• Resolve MacOS issues on GitHub actions.

[rgcoe]

@cmichelenstrofer - I'm trying to debug the CI on this PR for macOS. The key part of the failing CI log seems to be the following, which is either due to some mismatch of jax and jaxlib (see jax-ml/jax#14036) and/or something about arm64 (I think this seems likely based on my work below). I'd like your help thinking about this and deciding whether our current workflow of installing partially with mamba and partially with pip really makes sense.

E   AttributeError: partially initialized module 'jax' has no attribute 'version' (most likely due to a circular import)
=========================== short test summary info ============================
ERROR tests/test_core.py - RuntimeError: This version of jaxlib was built using AVX instructions, which your CPU and/or operating system do not support. This error is frequently encountered on macOS when running an x86 Python installation on ARM hardware. In this case, try installing an ARM build of Python. Otherwise, you may be able work around this issue by building jaxlib from source.
ERROR tests/test_integration.py - AttributeError: partially initialized module 'jax' has no attribute 'version' (most likely due to a circular import)
ERROR tests/test_pto.py - AttributeError: partially initialized module 'jax' has no attribute 'version' (most likely due to a circular import)
ERROR tests/test_utilities.py - AttributeError: partially initialized module 'jax' has no attribute 'version' (most likely due to a circular import)
ERROR tests/test_waves.py - AttributeError: partially initialized module 'jax' has no attribute 'version' (most likely due to a circular import)
!!!!!!!!!!!!!!!!!!! Interrupted: 5 errors during collection !!!!!!!!!!!!!!!!!!!!
========================= 5 errors in 99.80s (0:01:39) =========================

When I install locally, I do the following trying to replicate the install commands in pr.yml.

mamba create -n tmp_wot
mamba activate tmp_wot
mamba install python=3.12 capytaine wavespectra

This seems to work on the GitHub runner (see, e.g., https://github.com/sandialabs/WecOptTool/actions/runs/17650331026/job/50178598578), but my machine tells me there aren't capytaine and wavespectra binaries for arm64 on conda-forge (see #324 for our previous discussion on this). So I instead do the following, which works fine.

mamba create -n tmp_wot
mamba activate tmp_wot
mamba create -n tmp_wot
mamba install pip
cd wecopttool
pip install .
pip install gmsh pygmsh coveralls pytest # we could probably make a special set of optional dependencies in pyproject.toml to avoid this line
coverage run -m pytest

Based on all of this, I pushed an update to pr.yml with f1e6e51 to basically do the whole installation via pip. Note that this bypasses the installation of capytaine and wavespectra via mamba which I believe you added to allow for caching (see #35)¹ It seems to fix the problem on macOS, Ubuntu still works fine, and now Windows fails due a Segmentation fault 🤪 -- maybe this is due to these warnings which we have actually been receiving for a while about cygpath?

Footnotes

1. note that installation via pip on macOS and Ubuntu takes only 1-2 minutes (Windows is 8 minutes) ↩

[jtgrasb]

For some reason, fully removing mamba and pip installing in editable mode using -e worked to get the tests passing on my own fork, but didn't work for this PR and the windows test is still failing.