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. I did this here and also updated the inputs to jax.random.normal and got the tests to pass now.

To do:

• Understand why editable mode fixes the issue then merge.

[coveralls]

Pull Request Test Coverage Report for Build 18686165478

Warning: This coverage report may be inaccurate.

This pull request's base commit is no longer the HEAD commit of its target branch. This means it includes changes from outside the original pull request, including, potentially, unrelated coverage changes.

• For more information on this, see Tracking coverage changes with pull request builds.
• To avoid this issue with future PRs, see these Recommended CI Configurations.
• For a quick fix, rebase this PR at GitHub. Your next report should be accurate.

Details

• 0 of 120 (0.0%) changed or added relevant lines in 4 files are covered.
• 1140 unchanged lines in 7 files lost coverage.
• Overall coverage decreased (-93.8%) to 0.0%

---

Changes Missing Coverage	Covered Lines	Changed/Added Lines	%
wecopttool/utilities.py	0	4	0.0%
wecopttool/controllers.py	0	14	0.0%
wecopttool/pto.py	0	29	0.0%
wecopttool/core.py	0	73	0.0%

Files with Coverage Reduction	New Missed Lines	%
wecopttool/init.py	26	0.0%
wecopttool/geom.py	32	0.0%
wecopttool/waves.py	75	0.0%
wecopttool/controllers.py	79	0.0%
wecopttool/utilities.py	160	0.0%
wecopttool/pto.py	176	0.0%
wecopttool/core.py	592	0.0%

Totals
Change from base Build 17107381189:	-93.8%
Covered Lines:	0
Relevant Lines:	1426

---

💛 - Coveralls

[jtgrasb]

After encountering additional GitHub CI errors with uploading to coveralls (errors also present on dev branch), I had to edit the workflow further. Here are the changes I made to the pr.yml and push.yml:

• removed caching - this removed a lot of steps that may save time when setting up the environment, but aren't that important because the majority of the time is taken by actually running the tests.
• Replacing coverage run -m pytest with pytest --cov=wecopttool --cov-report=xml resolved some issues that were occurring when uploading to coveralls. pytest-cov integrates more seamlessly with pytest so this workflow should be better moving forward.

This PR should now be ready for merge.