diff --git a/notebooks/gaussian_eddy.ipynb b/notebooks/gaussian_eddy.ipynb
index c56d227..98f1c52 100644
--- a/notebooks/gaussian_eddy.ipynb
+++ b/notebooks/gaussian_eddy.ipynb
@@ -7,8 +7,8 @@
    "metadata": {
     "collapsed": false,
     "ExecuteTime": {
-     "end_time": "2024-03-27T09:48:18.803393Z",
-     "start_time": "2024-03-27T09:48:17.490282Z"
+     "end_time": "2024-04-16T06:06:20.795629Z",
+     "start_time": "2024-04-16T06:06:19.516749Z"
     }
    },
    "outputs": [],
@@ -18,8 +18,9 @@
     "\n",
     "import matplotlib.pyplot as plt\n",
     "import numpy as np\n",
+    "import optax\n",
     "\n",
-    "from jaxparrow.cyclogeostrophy import _iterative, _variational\n",
+    "from jaxparrow.cyclogeostrophy import _iterative, _variational, LR_VAR\n",
     "from jaxparrow.geostrophy import _geostrophy\n",
     "from jaxparrow.tools.kinematics import magnitude\n",
     "from jaxparrow.tools.operators import interpolation\n",
@@ -54,8 +55,8 @@
    "metadata": {
     "collapsed": false,
     "ExecuteTime": {
-     "end_time": "2024-03-27T09:48:18.816049Z",
-     "start_time": "2024-03-27T09:48:18.803541Z"
+     "end_time": "2024-04-16T06:06:20.809311Z",
+     "start_time": "2024-04-16T06:06:20.796162Z"
     }
    },
    "outputs": [],
@@ -85,8 +86,8 @@
    "metadata": {
     "collapsed": false,
     "ExecuteTime": {
-     "end_time": "2024-03-27T09:48:19.037726Z",
-     "start_time": "2024-03-27T09:48:18.816067Z"
+     "end_time": "2024-04-16T06:06:21.038434Z",
+     "start_time": "2024-04-16T06:06:20.809370Z"
     }
    },
    "outputs": [],
@@ -116,8 +117,8 @@
    "metadata": {
     "collapsed": false,
     "ExecuteTime": {
-     "end_time": "2024-03-27T09:48:19.357171Z",
-     "start_time": "2024-03-27T09:48:19.039226Z"
+     "end_time": "2024-04-16T06:06:21.344013Z",
+     "start_time": "2024-04-16T06:06:21.039655Z"
     }
    },
    "outputs": [
@@ -148,8 +149,8 @@
    "metadata": {
     "collapsed": false,
     "ExecuteTime": {
-     "end_time": "2024-03-27T09:48:19.632375Z",
-     "start_time": "2024-03-27T09:48:19.355709Z"
+     "end_time": "2024-04-16T06:06:21.560665Z",
+     "start_time": "2024-04-16T06:06:21.336989Z"
     }
    },
    "outputs": [
@@ -204,8 +205,8 @@
    "metadata": {
     "collapsed": false,
     "ExecuteTime": {
-     "end_time": "2024-03-27T09:48:19.729186Z",
-     "start_time": "2024-03-27T09:48:19.629885Z"
+     "end_time": "2024-04-16T06:06:21.569698Z",
+     "start_time": "2024-04-16T06:06:21.551351Z"
     }
    },
    "outputs": [],
@@ -220,8 +221,8 @@
    "metadata": {
     "collapsed": false,
     "ExecuteTime": {
-     "end_time": "2024-03-27T09:48:20.076605Z",
-     "start_time": "2024-03-27T09:48:19.655143Z"
+     "end_time": "2024-04-16T06:06:21.910989Z",
+     "start_time": "2024-04-16T06:06:21.570352Z"
     }
    },
    "outputs": [
@@ -256,8 +257,8 @@
    "metadata": {
     "collapsed": false,
     "ExecuteTime": {
-     "end_time": "2024-03-27T09:48:20.387878Z",
-     "start_time": "2024-03-27T09:48:20.075007Z"
+     "end_time": "2024-04-16T06:06:22.161062Z",
+     "start_time": "2024-04-16T06:06:21.909967Z"
     }
    },
    "outputs": [
@@ -300,8 +301,8 @@
    "metadata": {
     "collapsed": false,
     "ExecuteTime": {
-     "end_time": "2024-03-27T09:48:20.607746Z",
-     "start_time": "2024-03-27T09:48:20.388923Z"
+     "end_time": "2024-04-16T06:06:22.449759Z",
+     "start_time": "2024-04-16T06:06:22.156730Z"
     }
    },
    "outputs": [],
@@ -321,8 +322,8 @@
    "metadata": {
     "collapsed": false,
     "ExecuteTime": {
-     "end_time": "2024-03-27T09:48:20.880799Z",
-     "start_time": "2024-03-27T09:48:20.611195Z"
+     "end_time": "2024-04-16T06:06:22.744903Z",
+     "start_time": "2024-04-16T06:06:22.452556Z"
     }
    },
    "outputs": [
@@ -357,8 +358,8 @@
    "metadata": {
     "collapsed": false,
     "ExecuteTime": {
-     "end_time": "2024-03-27T09:48:20.975547Z",
-     "start_time": "2024-03-27T09:48:20.882108Z"
+     "end_time": "2024-04-16T06:06:22.849715Z",
+     "start_time": "2024-04-16T06:06:22.739457Z"
     }
    },
    "outputs": [
@@ -389,8 +390,8 @@
    "metadata": {
     "collapsed": false,
     "ExecuteTime": {
-     "end_time": "2024-03-27T09:48:21.015802Z",
-     "start_time": "2024-03-27T09:48:20.974942Z"
+     "end_time": "2024-04-16T06:06:22.917070Z",
+     "start_time": "2024-04-16T06:06:22.849692Z"
     }
    },
    "outputs": [
@@ -439,8 +440,8 @@
    "metadata": {
     "collapsed": false,
     "ExecuteTime": {
-     "end_time": "2024-03-27T09:48:21.031821Z",
-     "start_time": "2024-03-27T09:48:21.017074Z"
+     "end_time": "2024-04-16T06:06:22.963314Z",
+     "start_time": "2024-04-16T06:06:22.905593Z"
     }
    },
    "outputs": [],
@@ -455,8 +456,8 @@
    "metadata": {
     "collapsed": false,
     "ExecuteTime": {
-     "end_time": "2024-03-27T09:48:21.319274Z",
-     "start_time": "2024-03-27T09:48:21.034151Z"
+     "end_time": "2024-04-16T06:06:23.318Z",
+     "start_time": "2024-04-16T06:06:22.929184Z"
     }
    },
    "outputs": [
@@ -491,8 +492,8 @@
    "metadata": {
     "collapsed": false,
     "ExecuteTime": {
-     "end_time": "2024-03-27T09:48:21.470194Z",
-     "start_time": "2024-03-27T09:48:21.320931Z"
+     "end_time": "2024-04-16T06:06:23.426827Z",
+     "start_time": "2024-04-16T06:06:23.314041Z"
     }
    },
    "outputs": [
@@ -535,8 +536,8 @@
    "metadata": {
     "collapsed": false,
     "ExecuteTime": {
-     "end_time": "2024-03-27T09:48:21.529543Z",
-     "start_time": "2024-03-27T09:48:21.417223Z"
+     "end_time": "2024-04-16T06:06:23.539343Z",
+     "start_time": "2024-04-16T06:06:23.425423Z"
     }
    },
    "outputs": [],
@@ -563,15 +564,16 @@
    "metadata": {
     "collapsed": false,
     "ExecuteTime": {
-     "end_time": "2024-03-27T09:48:22.586107Z",
-     "start_time": "2024-03-27T09:48:21.529896Z"
+     "end_time": "2024-04-16T06:06:24.603106Z",
+     "start_time": "2024-04-16T06:06:23.540016Z"
     }
    },
    "outputs": [],
    "source": [
+    "optim = optax.sgd(learning_rate=LR_VAR)\n",
     "u_cyclo_est, v_cyclo_est, _ = _variational(u_geos_u, v_geos_v, dXY, dXY, dXY, dXY,\n",
     "                                           coriolis_factor, coriolis_factor, mask,\n",
-    "                                           n_it=20, optim=\"sgd\", optim_kwargs=None,\n",
+    "                                           n_it=20, optim=optim,\n",
     "                                           return_losses=False)\n",
     "\n",
     "u_cyclo_est_t = interpolation(u_cyclo_est, axis=1, padding=\"left\")\n",
@@ -587,8 +589,8 @@
    "metadata": {
     "collapsed": false,
     "ExecuteTime": {
-     "end_time": "2024-03-27T09:48:22.914879Z",
-     "start_time": "2024-03-27T09:48:22.587554Z"
+     "end_time": "2024-04-16T06:06:24.874994Z",
+     "start_time": "2024-04-16T06:06:24.604585Z"
     }
    },
    "outputs": [
@@ -623,8 +625,8 @@
    "metadata": {
     "collapsed": false,
     "ExecuteTime": {
-     "end_time": "2024-03-27T09:48:23.006573Z",
-     "start_time": "2024-03-27T09:48:22.915936Z"
+     "end_time": "2024-04-16T06:06:24.968533Z",
+     "start_time": "2024-04-16T06:06:24.873642Z"
     }
    },
    "outputs": [
@@ -655,8 +657,8 @@
    "metadata": {
     "collapsed": false,
     "ExecuteTime": {
-     "end_time": "2024-03-27T09:48:23.024525Z",
-     "start_time": "2024-03-27T09:48:23.006882Z"
+     "end_time": "2024-04-16T06:06:24.986644Z",
+     "start_time": "2024-04-16T06:06:24.967457Z"
     }
    },
    "outputs": [
@@ -704,15 +706,15 @@
    "metadata": {
     "collapsed": false,
     "ExecuteTime": {
-     "end_time": "2024-03-27T09:48:23.378719Z",
-     "start_time": "2024-03-27T09:48:23.026195Z"
+     "end_time": "2024-04-16T06:06:25.444250Z",
+     "start_time": "2024-04-16T06:06:24.987409Z"
     }
    },
    "outputs": [],
    "source": [
     "u_cyclo_est, v_cyclo_est, _ = _iterative(u_geos_u, v_geos_v, dXY, dXY, dXY, dXY,\n",
     "                                         coriolis_factor, coriolis_factor, mask,\n",
-    "                                         n_it=20, res_eps=0.01, res_init=\"same\", \n",
+    "                                         n_it=20, res_eps=0.01,\n",
     "                                         use_res_filter=True, res_filter_size=3, \n",
     "                                         return_losses=False)\n",
     "\n",
@@ -729,8 +731,8 @@
    "metadata": {
     "collapsed": false,
     "ExecuteTime": {
-     "end_time": "2024-03-27T09:48:23.650239Z",
-     "start_time": "2024-03-27T09:48:23.380245Z"
+     "end_time": "2024-04-16T06:06:25.717737Z",
+     "start_time": "2024-04-16T06:06:25.445894Z"
     }
    },
    "outputs": [
@@ -765,8 +767,8 @@
    "metadata": {
     "collapsed": false,
     "ExecuteTime": {
-     "end_time": "2024-03-27T09:48:23.794821Z",
-     "start_time": "2024-03-27T09:48:23.646712Z"
+     "end_time": "2024-04-16T06:06:25.811113Z",
+     "start_time": "2024-04-16T06:06:25.716342Z"
     }
    },
    "outputs": [
@@ -797,8 +799,8 @@
    "metadata": {
     "collapsed": false,
     "ExecuteTime": {
-     "end_time": "2024-03-27T09:48:23.795617Z",
-     "start_time": "2024-03-27T09:48:23.743304Z"
+     "end_time": "2024-04-16T06:06:25.828628Z",
+     "start_time": "2024-04-16T06:06:25.810917Z"
     }
    },
    "outputs": [
@@ -834,15 +836,15 @@
    "metadata": {
     "collapsed": false,
     "ExecuteTime": {
-     "end_time": "2024-03-27T09:48:24.088571Z",
-     "start_time": "2024-03-27T09:48:23.759301Z"
+     "end_time": "2024-04-16T06:06:26.166281Z",
+     "start_time": "2024-04-16T06:06:25.829283Z"
     }
    },
    "outputs": [],
    "source": [
     "u_cyclo_est, v_cyclo_est, _ = _iterative(u_geos_u, v_geos_v, dXY, dXY, dXY, dXY,\n",
     "                                         coriolis_factor, coriolis_factor, mask,\n",
-    "                                         n_it=20, res_eps=0.01, res_init=\"same\", \n",
+    "                                         n_it=20, res_eps=0.01, \n",
     "                                         use_res_filter=False, res_filter_size=1, \n",
     "                                         return_losses=False)\n",
     "\n",
@@ -859,8 +861,8 @@
    "metadata": {
     "collapsed": false,
     "ExecuteTime": {
-     "end_time": "2024-03-27T09:48:24.400204Z",
-     "start_time": "2024-03-27T09:48:24.090075Z"
+     "end_time": "2024-04-16T06:06:26.436757Z",
+     "start_time": "2024-04-16T06:06:26.167302Z"
     }
    },
    "outputs": [
@@ -895,8 +897,8 @@
    "metadata": {
     "collapsed": false,
     "ExecuteTime": {
-     "end_time": "2024-03-27T09:48:24.450133Z",
-     "start_time": "2024-03-27T09:48:24.360068Z"
+     "end_time": "2024-04-16T06:06:26.527336Z",
+     "start_time": "2024-04-16T06:06:26.435532Z"
     }
    },
    "outputs": [
@@ -927,8 +929,8 @@
    "metadata": {
     "collapsed": false,
     "ExecuteTime": {
-     "end_time": "2024-03-27T09:48:24.476603Z",
-     "start_time": "2024-03-27T09:48:24.457198Z"
+     "end_time": "2024-04-16T06:06:26.557442Z",
+     "start_time": "2024-04-16T06:06:26.532119Z"
     }
    },
    "outputs": [
diff --git a/notebooks/gaussian_eddy/gaussian_eddy.md b/notebooks/gaussian_eddy/gaussian_eddy.md
index 4ed01bb..09e4dfc 100644
--- a/notebooks/gaussian_eddy/gaussian_eddy.md
+++ b/notebooks/gaussian_eddy/gaussian_eddy.md
@@ -4,8 +4,9 @@ import sys
 
 import matplotlib.pyplot as plt
 import numpy as np
+import optax
 
-from jaxparrow.cyclogeostrophy import _iterative, _variational
+from jaxparrow.cyclogeostrophy import _iterative, _variational, LR_VAR
 from jaxparrow.geostrophy import _geostrophy
 from jaxparrow.tools.kinematics import magnitude
 from jaxparrow.tools.operators import interpolation
@@ -273,9 +274,10 @@ mask = init_mask(u_geos_t)
 
 
 ```python
+optim = optax.sgd(learning_rate=LR_VAR)
 u_cyclo_est, v_cyclo_est, _ = _variational(u_geos_u, v_geos_v, dXY, dXY, dXY, dXY,
                                            coriolis_factor, coriolis_factor, mask,
-                                           n_it=20, optim="sgd", optim_kwargs=None,
+                                           n_it=20, optim=optim,
                                            return_losses=False)
 
 u_cyclo_est_t = interpolation(u_cyclo_est, axis=1, padding="left")
@@ -348,7 +350,7 @@ Use of a convolution filter when computing the residuals.
 ```python
 u_cyclo_est, v_cyclo_est, _ = _iterative(u_geos_u, v_geos_v, dXY, dXY, dXY, dXY,
                                          coriolis_factor, coriolis_factor, mask,
-                                         n_it=20, res_eps=0.01, res_init="same", 
+                                         n_it=20, res_eps=0.01, 
                                          use_res_filter=True, res_filter_size=3, 
                                          return_losses=False)
 
@@ -418,7 +420,7 @@ No convolution filter, original approach.
 ```python
 u_cyclo_est, v_cyclo_est, _ = _iterative(u_geos_u, v_geos_v, dXY, dXY, dXY, dXY,
                                          coriolis_factor, coriolis_factor, mask,
-                                         n_it=20, res_eps=0.01, res_init="same", 
+                                         n_it=20, res_eps=0.01, 
                                          use_res_filter=False, res_filter_size=1, 
                                          return_losses=False)