ENHANCEMENT: Autograd to jax

sandialabs / WecOptTool

WEC Design Optimization Toolbox

https://sandialabs.github.io/WecOptTool/

GNU General Public License v3.0

12 stars 20 forks source link

ENHANCEMENT: Autograd to jax #319

Open jorgeypcb opened 3 months ago

jorgeypcb commented 3 months ago

Description

This PR is a demonstration of the changes being made to the source code in order to transition from autograd to jax FIX #118 .

Type of PR

[ ] Other: Enhancement

Checklist for PR

[ ] Authors read the contribution guidelines
[ ] The pull request is from an issue branch (not main) on your fork, to the main branch in WecOptTool.
[ ] The authors have given the admins edit access
[ ] All changes adhere to the style guide including PEP8, Docstrings, and Type Hints.
[ ] Modified the documentation if applicable
[ ] Modified or added a new test
[ ] All tests didn't pass, still working on passing test_waves.py and test_integration.py
[ ] Reference or close any relevant issues

Additional details

cmichelenstrofer commented 3 months ago

Looking good! I added some minor comments.

cmichelenstrofer commented 2 months ago

@jorgeypcb jax.scipy.optimize.minimize currently only has BFGS implemented. But... could we use just scipy.optimize.minimize and use Jax manually for providing the gradients?

jorgeypcb commented 1 month ago

@jorgeypcb jax.scipy.optimize.minimize currently only has BFGS implemented. But... could we use just scipy.optimize.minimize and use Jax manually for providing the gradients?

We could try! I will be testing that today and if it works I will see how it compares to the other pull request I did where I used the original setup with some slight changes to get it to work with cyipopt minimize_ipopt, to see which option is better. I did that pull request with my RageTech account 👍