Autograd -> Jax conversion (sandialabs#433)

* Update CONTRIBUTING.md to indicate PRs should be to the new `dev` branch

* Update RELEASING.md to reflect new workflow with the `dev` branch

* update docstrings (sandialabs#326)

* damping naming and consistently change radiation damping (sandialabs#328)

* issue 321 fd_to_td() bug (sandialabs#329)

* bug bix : DC and Nyquist frequency should not be devided by two before ifft

* Changed td_to_fd to scale single sided frequency components rather than TD signal

* minor bug fix from issue332 sandialabs#332

* nodf -> ndof (sandialabs#334)

* add DOI for Daniel's paper (sandialabs#336)

* Lower tolerance for new test to fix CI failing occasionally

* hyperlinks no longer have formatting, plus other small adjustments (sandialabs#348)

* Merge to dev, not main (sandialabs#349)

* Dev version of documentation site (sandialabs#347)

* added initial file changes based on sphinx_multiversion docs and WEC-Sim implementation

* removed sphinx-multiversion since it is no longer supported and made manual multiversion

* now uses absolute paths, commented out linkcheck for debugging

* fixed docstring errors in utilities module

* updating files again that somehow got reverted

* fixing path in conf.py

* don't run tutorials (will revert later)

* handle file moves correctly, fixed if statement to make other versions appear

* fixed two bugs in versions template

* reverted temp changes, changes latest to main

* switched latest to main

* main branch now in root directory of pages

* fixed URLs with change from last commit

* make other branches visible before building

* switched main branch tag for more testing

* fixed typo

* switched dev branch to an existing branch

* renamed main to latest, changed version.html file name to avoid confusion

* added prints about moving files so Sphinx output isn't misleading

* fixed typo with quotations

* changed versions.html name back because that broke things I guess

* modified contributing documentation to reflect changes

* add logic to remove duplicate 'latest' branch

* Fixed pathing when already on latest

* remove typo

* Troubleshooting complete, switching back to correct branches for deployment

* Removed extra word in docstring

* removed redundant function

* fixed pathing so returns to same file (and fixes tutorial/API docs)

* changed latest branch for demonstration

* switched back latest branch for deployment

* updated with new Capytaine docs URL

* Add warnings when adding inertia and hydrostatic stiffness automatically (sandialabs#346)

* CI workflow cleanup (sandialabs#352)

* removed conda environment from workflows since newer capytaine/wavespectra work with Windows

* fixed unnecessary capitalization

* still create CI conda environment to fix Mac environment failures

* added conda env fully back in, push workflow deploys docs, split PR workflow

* conda environment activates again

* mambaforge instead of miniforge

* manual cache reset

* reset to older version of setup-miniconda to troubleshoot

* Updated workflows to newest Python version and changed references to supported versions (sandialabs#390)

Co-authored-by: jtgrasb <[email protected]>

* Revert to Python 3.12 (sandialabs#394)

* Try specifying subversion

* Test new cache

* revert to 3.12

* Revert comment back to normal

* use dev for docs and restrict sphinx (sandialabs#396)

* Remove Sphinx version requirement (sandialabs#409)

* v3.0.3

* v3.1

* Trying to convert tutorial 1

* Convert to jax progress

* post-processing

* clear outputs

* wavebot tutorial running

* wavebot tutorial running

* Update to jax and numpy

* Revert wavebot execution count

* update pyproject.toml

* Specify jax version for mac

* try jaxlib

* no jaxlib

* add jax to environment manually

* add jaxlib to env

* ad jaxlib to pyproject

* install jax manually for macos

* conda init

* install jax and jaxlib on macos

* try arm64

* remove arm

* make core optimization jittable

* try pinning jax version

* revert previous

* install entirely w. pip

* Add verbose outputs to testing

* editable mode

* no mamba

* editable mode

* update random inputs

* remove cache environment

* use pytest-cov

* remove cache clear

* Fix push.yml

---------

Co-authored-by: Carlos A. Michelén Ströfer <[email protected]>
Co-authored-by: Ryan Coe <[email protected]>
Co-authored-by: Daniel Gaebele <[email protected]>
Co-authored-by: Michael Devin <[email protected]>
Co-authored-by: mcdevin <[email protected]>

Author: 5 people

.github/workflows/pr.yml

@@ -23,52 +23,13 @@ jobs:
     # Checkout the WecOptTool repo
     - uses: actions/checkout@v3
 
-    # Cache the conda environment >>>
-    # The cache key includes the OS, the date, the hash of the pyproject.toml file, and the cache number.
-    # A new cache is created if:
-    #   - this is the first time today that this file is run (date changes)
-    #   - the content of pyproject.toml changes
-    #   - you manually change the value of the CACHE_NUMBER below
-    # Else the existing cache is used.
+    # Caching of environment has been removed for simplicity
     - name: Setup Miniforge
       uses: conda-incubator/setup-miniconda@v2
       with:
           miniforge-variant: Miniforge3
           miniforge-version: latest
           activate-environment: test-env
-          use-mamba: true
-
-    # save the date to include in the cache key
-    - name: Get Date
-      id: get-date
-      run: echo "DATE=$(/bin/date -u '+%Y%m%d')" >> $GITHUB_ENV
-      shell: bash
-
-    # create a conda yaml file
-    - name: Create environment.yml file
-      run: |
-        echo "name: test-env" >> environment.yml;
-        echo " " >> environment.yml
-        echo "dependencies:" >> environment.yml
-        echo "  - python=${{ matrix.python-version }}" >> environment.yml
-        echo "  - capytaine" >> environment.yml
-        echo "  - wavespectra" >> environment.yml
-        cat environment.yml
-
-    # use the cache if it exists
-    - uses: actions/cache@v3
-      env:
-        CACHE_NUMBER: 0  # increase to reset cache manually
-      with:
-        path: ${{ env.CONDA }}/envs
-        key: conda-${{ runner.os }}--${{ runner.arch }}--${{ env.DATE }}-${{ hashFiles('pyproject.toml') }}-${{ env.CACHE_NUMBER }}
-      id: cache
-
-    # if cache key has changed, create new cache
-    - name: Update environment
-      run: mamba env update -n test-env -f environment.yml
-      if: steps.cache.outputs.cache-hit != 'true'
-    # <<< Cache the conda environment
 
     # install libglu on ubuntu.
     - name: Install libglu
@@ -81,13 +42,13 @@ jobs:
       run: |
         conda activate test-env
         python3 -m pip install --upgrade pip
-        pip3 install gmsh pygmsh coveralls pytest
+        pip3 install gmsh pygmsh coveralls pytest pytest-cov
         pip3 install .
 
     # run all tests & coverage
     - name: Run Test
       shell: bash -l {0}
-      run: coverage run -m pytest
+      run: pytest --cov=wecopttool --cov-report=xml
 
     # upload coverage data
     - name: Upload coverage data to coveralls.io

.github/workflows/push.yml

@@ -23,52 +23,12 @@ jobs:
     # Checkout the WecOptTool repo
     - uses: actions/checkout@v3
 
-    # Cache the conda environment >>>
-    # The cache key includes the OS, the date, the hash of the pyproject.toml file, and the cache number.
-    # A new cache is created if:
-    #   - this is the first time today that this file is run (date changes)
-    #   - the content of pyproject.toml changes
-    #   - you manually change the value of the CACHE_NUMBER below
-    # Else the existing cache is used.
     - name: Setup Miniforge
       uses: conda-incubator/setup-miniconda@v2
       with:
           miniforge-variant: Miniforge3
           miniforge-version: latest
           activate-environment: test-env
-          use-mamba: true
-
-    # save the date to include in the cache key
-    - name: Get Date
-      id: get-date
-      run: echo "DATE=$(/bin/date -u '+%Y%m%d')" >> $GITHUB_ENV
-      shell: bash
-
-    # create a conda yaml file
-    - name: Create environment.yml file
-      run: |
-        echo "name: test-env" >> environment.yml;
-        echo " " >> environment.yml
-        echo "dependencies:" >> environment.yml
-        echo "  - python=${{ matrix.python-version }}" >> environment.yml
-        echo "  - capytaine" >> environment.yml
-        echo "  - wavespectra" >> environment.yml
-        cat environment.yml
-
-    # use the cache if it exists
-    - uses: actions/cache@v3
-      env:
-        CACHE_NUMBER: 0  # increase to reset cache manually
-      with:
-        path: ${{ env.CONDA }}/envs
-        key: conda-${{ runner.os }}--${{ runner.arch }}--${{ env.DATE }}-${{ hashFiles('pyproject.toml') }}-${{ env.CACHE_NUMBER }}
-      id: cache
-
-    # if cache key has changed, create new cache
-    - name: Update environment
-      run: mamba env update -n test-env -f environment.yml
-      if: steps.cache.outputs.cache-hit != 'true'
-    # <<< Cache the conda environment
 
     # install libglu on ubuntu.
     - name: Install libglu
@@ -81,13 +41,13 @@ jobs:
       run: |
         conda activate test-env
         python3 -m pip install --upgrade pip
-        pip3 install gmsh pygmsh coveralls pytest
+        pip3 install gmsh pygmsh coveralls pytest pytest-cov
         pip3 install .
 
     # run all tests & coverage
     - name: Run Test
       shell: bash -l {0}
-      run: coverage run -m pytest
+      run: pytest --cov=wecopttool --cov-report=xml
 
     # upload coverage data
     - name: Upload coverage data to coveralls.io

examples/tutorial_1_WaveBot.ipynb

@@ -30,7 +30,8 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "import autograd.numpy as np\n",
+    "import numpy as np\n",
+    "import jax.numpy as jnp\n",
     "import capytaine as cpy\n",
     "from capytaine.io.meshio import load_from_meshio\n",
     "import matplotlib.pyplot as plt\n",
@@ -372,7 +373,7 @@
     "\n",
     "def const_f_pto(wec, x_wec, x_opt, wave): # Format for scipy.optimize.minimize\n",
     "    f = pto.force(wec, x_wec, x_opt, wave, nsubsteps)\n",
-    "    return f_max - np.abs(f.flatten())\n",
+    "    return f_max - jnp.abs(f.flatten())\n",
     "\n",
     "ineq_cons = {'type': 'ineq',\n",
     "             'fun': const_f_pto,\n",
@@ -706,7 +707,7 @@
     "## Update PTO constraints and forcing\n",
     "def const_f_pto_2(wec, x_wec, x_opt, wave):\n",
     "    f = pto_2.force_on_wec(wec, x_wec, x_opt, wave, nsubsteps)\n",
-    "    return f_max - np.abs(f.flatten())\n",
+    "    return f_max - jnp.abs(f.flatten())\n",
     "ineq_cons_2 = {'type': 'ineq', 'fun': const_f_pto_2}\n",
     "constraints_2 = [ineq_cons_2]\n",
     "f_add_2 = {'PTO': pto_2.force_on_wec}"
@@ -989,7 +990,7 @@
     "\n",
     "    def const_f_pto(wec, x_wec, x_opt, wave):\n",
     "        f = pto.force(wec, x_wec, x_opt, wave, nsubsteps)\n",
-    "        return f_max - np.abs(f.flatten())\n",
+    "        return f_max - jnp.abs(f.flatten())\n",
     "\n",
     "    ineq_cons = {'type': 'ineq', 'fun': const_f_pto}\n",
     "    constraints = [ineq_cons]\n",

examples/tutorial_2_AquaHarmonics.ipynb

@@ -24,7 +24,8 @@
    "source": [
     "import capytaine as cpy\n",
     "from capytaine.io.meshio import load_from_meshio\n",
-    "import autograd.numpy as np\n",
+    "import numpy as np\n",
+    "import jax.numpy as jnp\n",
     "import matplotlib.pyplot as plt\n",
     "from matplotlib import cm\n",
     "from scipy.optimize import brute\n",
@@ -376,10 +377,10 @@
    "source": [
     "def f_buoyancy(wec, x_wec, x_opt, wave, nsubsteps=1):\n",
     "    \"\"\"Only the zeroth order component (doesn't include linear stiffness)\"\"\"\n",
-    "    return displacement * rho * g * np.ones([wec.ncomponents*nsubsteps, wec.ndof])\n",
+    "    return displacement * rho * g * jnp.ones([wec.ncomponents*nsubsteps, wec.ndof])\n",
     "\n",
     "def f_gravity(wec, x_wec, x_opt, wave, nsubsteps=1):\n",
-    "    return -1 * wec.inertia_matrix.item() * g * np.ones([wec.ncomponents*nsubsteps, wec.ndof])\n",
+    "    return -1 * wec.inertia_matrix.item() * g * jnp.ones([wec.ncomponents*nsubsteps, wec.ndof])\n",
     "\n",
     "def f_pretension_wec(wec, x_wec, x_opt, wave, nsubsteps=1):\n",
     "    \"\"\"Pretension force as it acts on the WEC\"\"\"\n",
@@ -389,18 +390,18 @@
     "\n",
     "def f_pto_passive(wec, x_wec, x_opt, wave, nsubsteps=1):\n",
     "    pos = wec.vec_to_dofmat(x_wec)\n",
-    "    vel = np.dot(wec.derivative_mat,pos)\n",
-    "    acc = np.dot(wec.derivative_mat, vel)\n",
+    "    vel = jnp.dot(wec.derivative_mat,pos)\n",
+    "    acc = jnp.dot(wec.derivative_mat, vel)\n",
     "    time_matrix = wec.time_mat_nsubsteps(nsubsteps)\n",
     "    spring = -(gear_ratios['spring']*airspring['gamma']*airspring['area']*\n",
     "              airspring['press_init']/airspring['vol_init']) * pos\n",
-    "    f_spring = np.dot(time_matrix,spring)\n",
+    "    f_spring = jnp.dot(time_matrix,spring)\n",
     "    fric = -(friction_pto  + \n",
     "                friction['Bpneumatic_spring_static1']*\n",
     "                gear_ratios['spring']) * vel\n",
-    "    f_fric = np.dot(time_matrix,fric)\n",
+    "    f_fric = jnp.dot(time_matrix,fric)\n",
     "    inertia = inertia_pto * acc\n",
-    "    f_inertia = np.dot(time_matrix,inertia)\n",
+    "    f_inertia = jnp.dot(time_matrix,inertia)\n",
     "    return f_spring + f_fric + f_inertia\n",
     "\n",
     "def f_pto_line(wec, x_wec, x_opt, wave, nsubsteps=1):\n",
@@ -452,18 +453,18 @@
     "\n",
     "def const_peak_torque_pto(wec, x_wec, x_opt, wave):\n",
     "    torque = pto.force(wec, x_wec, x_opt, wave, nsubsteps)\n",
-    "    return torque_peak_max - np.abs(torque.flatten())\n",
+    "    return torque_peak_max - jnp.abs(torque.flatten())\n",
     "\n",
     "def const_speed_pto(wec, x_wec, x_opt, wave):\n",
     "    rot_vel = pto.velocity(wec, x_wec, x_opt, wave, nsubsteps)\n",
-    "    return rot_speed_max - np.abs(rot_vel.flatten())\n",
+    "    return rot_speed_max - jnp.abs(rot_vel.flatten())\n",
     "\n",
     "def const_power_pto(wec, x_wec, x_opt, wave):\n",
     "    power_mech = (\n",
     "        pto.velocity(wec, x_wec, x_opt, wave, nsubsteps) *\n",
     "        pto.force(wec, x_wec, x_opt, wave, nsubsteps)\n",
     "    )\n",
-    "    return power_max - np.abs(power_mech.flatten())\n",
+    "    return power_max - jnp.abs(power_mech.flatten())\n",
     "\n",
     "def constrain_min_tension(wec, x_wec, x_opt, wave):\n",
     "    total_tension = -1*f_pto_line(wec, x_wec, x_opt, wave, nsubsteps)\n",

examples/tutorial_3_LUPA.ipynb

@@ -40,7 +40,8 @@
     "import gmsh, pygmsh\n",
     "import capytaine as cpy\n",
     "from capytaine.io.meshio import load_from_meshio\n",
-    "import autograd.numpy as np\n",
+    "import numpy as np\n",
+    "import jax.numpy as jnp\n",
     "import matplotlib.pyplot as plt\n",
     "import xarray as xr\n",
     "from scipy.optimize import brute\n",
@@ -766,8 +767,8 @@
     "# maximum stroke\n",
     "stroke_max = 0.5  # m\n",
     "def const_stroke_pto(wec, x_wec, x_opt, wave): \n",
-    "    pos = pto.position(wec, x_wec, x_opt, wave, nsubsteps)\n",
-    "    return stroke_max - np.abs(pos.flatten())\n",
+    "    pos = pto.position(wec, x_wec, x_opt, waves, nsubsteps)\n",
+    "    return stroke_max - jnp.abs(pos.flatten())\n",
     "\n",
     "## GENERATOR\n",
     "# peak torque\n",
@@ -776,20 +777,20 @@
     "    \"\"\"Instantaneous torque must not exceed max torque Tmax - |T| >=0 \n",
     "    \"\"\"\n",
     "    torque = pto.force(wec, x_wec, x_opt, wave, nsubsteps) / gear_ratio(radius)\n",
-    "    return generator['max_torque'] - np.abs(torque.flatten())\n",
+    "    return generator['max_torque'] - jnp.abs(torque.flatten())\n",
     "\n",
     "# continuous torque\n",
     "def const_torque_pto(wec, x_wec, x_opt, wave, radius=default_radius): \n",
     "    \"\"\"RMS torque must not exceed max continous torque \n",
     "        Tmax_conti - Trms >=0 \"\"\"\n",
     "    torque = pto.force(wec, x_wec, x_opt, wave, nsubsteps) / gear_ratio(radius)\n",
-    "    torque_rms = np.sqrt(np.mean(torque.flatten()**2))\n",
+    "    torque_rms = jnp.sqrt(jnp.mean(torque.flatten()**2))\n",
     "    return generator['continuous_torque'] - torque_rms\n",
     "\n",
     "# max speed\n",
     "def const_speed_pto(wec, x_wec, x_opt, wave, radius=default_radius): \n",
     "    rot_vel = pto.velocity(wec, x_wec, x_opt, wave, nsubsteps) * gear_ratio(radius)\n",
-    "    return generator['max_speed'] - np.abs(rot_vel.flatten())\n",
+    "    return generator['max_speed'] - jnp.abs(rot_vel.flatten())\n",
     "\n",
     "## Constraints\n",
     "constraints = [\n",
@@ -1263,7 +1264,7 @@
     "    stroke_max = 0.5  # m\n",
     "    def const_stroke_pto(wec, x_wec, x_opt, wave): \n",
     "        pos = pto.position(wec, x_wec, x_opt, wave, nsubsteps)\n",
-    "        return stroke_max - np.abs(pos.flatten())\n",
+    "        return stroke_max - jnp.abs(pos.flatten())\n",
     "\n",
     "    ## GENERATOR\n",
     "    # peak torque\n",
@@ -1272,20 +1273,20 @@
     "        \"\"\"Instantaneous torque must not exceed max torque Tmax - |T| >=0 \n",
     "        \"\"\"\n",
     "        torque = pto.force(wec, x_wec, x_opt, wave, nsubsteps) / gear_ratio(radius)\n",
-    "        return generator['max_torque'] - np.abs(torque.flatten())\n",
+    "        return generator['max_torque'] - jnp.abs(torque.flatten())\n",
     "\n",
     "    # continuous torque\n",
     "    def const_torque_pto(wec, x_wec, x_opt, wave): \n",
     "        \"\"\"RMS torque must not exceed max continous torque \n",
     "            Tmax_conti - Trms >=0 \"\"\"\n",
     "        torque = pto.force(wec, x_wec, x_opt, wave, nsubsteps) / gear_ratio(radius)\n",
-    "        torque_rms = np.sqrt(np.mean(torque.flatten()**2))\n",
+    "        torque_rms = jnp.sqrt(jnp.mean(torque.flatten()**2))\n",
     "        return generator['continuous_torque'] - torque_rms\n",
     "\n",
     "    # max speed\n",
     "    def const_speed_pto(wec, x_wec, x_opt, wave): \n",
     "        rot_vel = pto.velocity(wec, x_wec, x_opt, wave, nsubsteps) * gear_ratio(radius)\n",
-    "        return generator['max_speed'] - np.abs(rot_vel.flatten())\n",
+    "        return generator['max_speed'] - jnp.abs(rot_vel.flatten())\n",
     "\n",
     "    ## Constraints\n",
     "    constraints = [\n",