Skip to content

Commit

Permalink
update notebooks accordingly
Browse files Browse the repository at this point in the history
  • Loading branch information
@vadmbertr
vadmbertr committed Apr 16, 2024
1 parent e905552 commit a6768a9
Show file tree
Hide file tree
Showing 2 changed files with 68 additions and 64 deletions.
122 changes: 62 additions & 60 deletions notebooks/gaussian_eddy.ipynb
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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": [],
Expand All @@ -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",
Expand Down Expand Up @@ -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": [],
Expand Down Expand Up @@ -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": [],
Expand Down Expand Up @@ -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": [
Expand Down Expand Up @@ -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": [
Expand Down Expand Up @@ -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": [],
Expand All @@ -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": [
Expand Down Expand Up @@ -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": [
Expand Down Expand Up @@ -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": [],
Expand All @@ -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": [
Expand Down Expand Up @@ -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": [
Expand Down Expand Up @@ -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": [
Expand Down Expand Up @@ -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": [],
Expand All @@ -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": [
Expand Down Expand Up @@ -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": [
Expand Down Expand Up @@ -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": [],
Expand All @@ -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",
Expand All @@ -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": [
Expand Down Expand Up @@ -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": [
Expand Down Expand Up @@ -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": [
Expand Down Expand Up @@ -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",
Expand All @@ -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": [
Expand Down Expand Up @@ -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": [
Expand Down Expand Up @@ -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": [
Expand Down Expand Up @@ -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",
Expand All @@ -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": [
Expand Down Expand Up @@ -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": [
Expand Down Expand Up @@ -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": [
Expand Down
10 changes: 6 additions & 4 deletions notebooks/gaussian_eddy/gaussian_eddy.md
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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")
Expand Down Expand Up @@ -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)

Expand Down Expand Up @@ -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)

Expand Down

0 comments on commit a6768a9

Please sign in to comment.