Autograd -> Jax conversion

.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",