Merge pull request #103 from vadmbertr/docs-typos

Fix some typos in the doc

Author: vadmbertr

docs/examples/duacs.ipynb

@@ -26,7 +26,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 2,
+   "execution_count": null,
    "id": "9b8f05cd-b082-4d59-b552-6ee23a22e779",
    "metadata": {
     "ExecuteTime": {
@@ -37,16 +37,7 @@
      "hide-input"
     ]
    },
-   "outputs": [
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "/Users/bertrava/miniforge3/envs/jaxparrow/lib/python3.11/site-packages/tqdm/auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n",
-      "  from .autonotebook import tqdm as notebook_tqdm\n"
-     ]
-    }
-   ],
+   "outputs": [],
    "source": [
     "import cartopy.crs as ccrs\n",
     "import cmocean.cm as cmo\n",
@@ -98,24 +89,18 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 3,
+   "execution_count": null,
    "id": "58e0299b-feab-40d7-9927-39dc05a93d01",
    "metadata": {
     "ExecuteTime": {
      "end_time": "2024-06-10T13:46:14.436118Z",
      "start_time": "2024-06-10T13:45:55.061171Z"
-    }
+    },
+    "tags": [
+     "hide-input"
+    ]
    },
-   "outputs": [
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "INFO - 2026-03-06T16:31:36Z - Selected dataset version: \"202411\"\n",
-      "INFO - 2026-03-06T16:31:36Z - Selected dataset part: \"default\"\n"
-     ]
-    }
-   ],
+   "outputs": [],
    "source": [
     "dataset_options = {\n",
     "    \"dataset_id\": \"cmems_obs-sl_glo_phy-ssh_my_allsat-l4-duacs-0.125deg_P1D\",\n",
@@ -179,29 +164,6 @@
     "## Reconstructing cyclogeostrophic currents using `jaxparrow`"
    ]
   },
-  {
-   "cell_type": "markdown",
-   "id": "d86fcbd4-a3f9-4c9e-aeae-b7737a351ba2",
-   "metadata": {},
-   "source": [
-    "`jaxparrow` uses C-grids, following NEMO convention. U, V, and F points are automatically derived from the T points."
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 5,
-   "id": "c34b9f78-3c33-4e66-8051-0933da97f16f",
-   "metadata": {
-    "ExecuteTime": {
-     "end_time": "2024-06-10T13:46:20.842589Z",
-     "start_time": "2024-06-10T13:46:20.622322Z"
-    }
-   },
-   "outputs": [],
-   "source": [
-    "lon_t, lat_t = jnp.meshgrid(duacs_ds.longitude.values, duacs_ds.latitude.values)"
-   ]
-  },
   {
    "cell_type": "markdown",
    "id": "1983b1f8-a7f9-4559-b589-8787c4525061",
@@ -225,7 +187,10 @@
     "ExecuteTime": {
      "end_time": "2024-06-10T13:46:21.798714Z",
      "start_time": "2024-06-10T13:46:20.873348Z"
-    }
+    },
+    "tags": [
+     "hide-input"
+    ]
    },
    "outputs": [],
    "source": [
@@ -259,6 +224,8 @@
     "\n",
     "BATCH_SIZE = 10\n",
     "\n",
+    "lon_t, lat_t = jnp.meshgrid(duacs_ds.longitude.values, duacs_ds.latitude.values)\n",
+    "\n",
     "duacs_ds = duacs_ds.where(np.abs(lat_t) >= 5)\n",
     "duacs_ds = duacs_ds.chunk(chunks={\"time\": BATCH_SIZE, \"latitude\": -1, \"longitude\": -1})\n",
     "\n",

docs/examples/gradient_wind.ipynb

@@ -70,35 +70,11 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 4,
-   "metadata": {
-    "tags": [
-     "hide-input"
-    ]
-   },
-   "outputs": [
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "W0313 11:46:02.754300 1377963 cpp_gen_intrinsics.cc:74] Empty bitcode string provided for eigen. Optimizations relying on this IR will be disabled.\n"
-     ]
-    }
-   ],
-   "source": [
-    "lat, lon, ssh, ug, vg, ucg, vcg, land_mask = gaussian_eddy.simulate_gaussian_eddy(R0=50e3, eta0=-.2)"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 5,
+   "execution_count": null,
    "metadata": {},
    "outputs": [],
    "source": [
-    "uvg = kinematics.magnitude(ug, vg)\n",
-    "uvcg = kinematics.magnitude(ucg, vcg)\n",
-    "\n",
-    "max_diff = jnp.nanmax(jnp.abs(uvcg - uvg))"
+    "lat, lon, ssh, ug, vg, ucg, vcg, land_mask = gaussian_eddy.simulate_gaussian_eddy(R0=50e3, eta0=-.2)"
    ]
   },
   {
@@ -153,6 +129,18 @@
     "$v_g = -2x \\frac{g \\eta_0}{f R_0^2} \\exp^{-(r/R_0)^2} = -2x \\frac{g \\eta}{f R_0^2}$"
    ]
   },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "uvg = kinematics.magnitude(ug, vg)\n",
+    "uvcg = kinematics.magnitude(ucg, vcg)\n",
+    "\n",
+    "max_diff = jnp.nanmax(jnp.abs(uvcg - uvg))"
+   ]
+  },
   {
    "cell_type": "code",
    "execution_count": 7,
@@ -406,11 +394,11 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 13,
+   "execution_count": null,
    "metadata": {},
    "outputs": [],
    "source": [
-    "mb_res = minimization_based(lat_t=lat, lon_t=lon, ssh_t=ssh,)\n",
+    "mb_res = minimization_based(lat_t=lat, lon_t=lon, ssh_t=ssh)\n",
     "\n",
     "umb_est = mb_res.ucg\n",
     "vmb_est = mb_res.vcg\n",

docs/examples/swot-enatl60.ipynb

@@ -283,24 +283,13 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 7,
+   "execution_count": null,
    "metadata": {
     "tags": [
      "hide-input"
     ]
    },
-   "outputs": [
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "/var/folders/xc/bksmt58x2nq8jshz2jbf9b_m0000gn/T/ipykernel_61290/1634369195.py:46: FutureWarning: The return type of `Dataset.dims` will be changed to return a set of dimension names in future, in order to be more consistent with `DataArray.dims`. To access a mapping from dimension names to lengths, please use `Dataset.sizes`.\n",
-      "  n_pixels_003 = swot_003_ds.dims[\"num_pixels\"]\n",
-      "/var/folders/xc/bksmt58x2nq8jshz2jbf9b_m0000gn/T/ipykernel_61290/1634369195.py:47: FutureWarning: The return type of `Dataset.dims` will be changed to return a set of dimension names in future, in order to be more consistent with `DataArray.dims`. To access a mapping from dimension names to lengths, please use `Dataset.sizes`.\n",
-      "  n_pixels_016 = swot_016_ds.dims[\"num_pixels\"]\n"
-     ]
-    }
-   ],
+   "outputs": [],
    "source": [
     "t_003_regridder = get_regridder(enatl60_t_ds, swot_003_ds)\n",
     "u_003_regridder = get_regridder(enatl60_u_ds, swot_003_ds)\n",
@@ -434,7 +423,11 @@
   {
    "cell_type": "code",
    "execution_count": 9,
-   "metadata": {},
+   "metadata": {
+    "tags": [
+     "hide-input"
+    ]
+   },
    "outputs": [],
    "source": [
     "optim = optax.chain(optax.clip(1), optax.sgd(learning_rate=5e-3))\n",
@@ -553,7 +546,11 @@
   {
    "cell_type": "code",
    "execution_count": 10,
-   "metadata": {},
+   "metadata": {
+    "tags": [
+     "hide-input"
+    ]
+   },
    "outputs": [
     {
      "data": {

docs/faq.md

@@ -4,23 +4,29 @@
 
 `Dask` is probably your best bet to handle large datasets that do not fit into (CPU or GPU) memory.
 
-On a CPU backend, one can use a chunk size of 1 along the time dimension and map the call to [`cyclogeostrophy`](api.md#jaxparrow.cyclogeostrophy.cyclogeostrophy) onto the dataset:
+On a CPU backend, one can use a chunk size of 1 along the time dimension and map the call to [`minimization_based`](api.md#jaxparrow.cyclogeostrophy.minimization_based) onto the dataset:
 
 ```python
 import dask
 import jax.numpy as jnp
 import numpy as np
 import xarray as xr
 
-from jaxparrow.cyclogeostrophy import cyclogeostrophy
+from jaxparrow import minimization_based
 
 
 def do_one_block(in_block):
-    ucg, vcg, ug, vg = cyclogeostrophy(
-        jnp.asarray(in_block.ssh.values), jnp.asarray(in_block.lat.values), jnp.asarray(in_block.lon.values), 
-        return_geos=True, return_grids=False
+    mb_result = minimization_based(
+        lat_t=jnp.asarray(in_block.lat.values), lon_tjnp.asarray(in_block.lon.values), 
+        ssh_t=jnp.asarray(in_block.ssh.values), 
+        return_geos=True
     )
 
+    ucg = mb_result.ucg
+    vcg = mb_result.vcg
+    ug = mb_result.ug
+    vg = mb_result.vg
+
     out_block = xr.Dataset(
         {
             "ucg": (in_block.ssh.dims, np.asarray(ucg)[None, :, :]),
@@ -70,16 +76,21 @@ from jaxparrow.cyclogeostrophy import cyclogeostrophy
 
 
 vmap_cyclogeostrophy = jax.vmap(
-    lambda *args: cyclogeostrophy(*args, return_geos=True, return_grids=False),
+    lambda ssh, lat, lon: minimization_based(lat_t=lat, lon_t=lon, ssh_t=ssh, return_geos=True),
     in_axes=(0, None, None)
 )
 
 
 def do_one_block_vmap(in_block: xr.Dataset):
-    ucg_3d, vcg_3d, ug_3d, vg_3d = vmap_cyclogeostrophy(
+    mb_result = vmap_cyclogeostrophy(
         jnp.asarray(in_block.ssh.values), jnp.asarray(in_block.lat.values), jnp.asarray(in_block.lon.values)
     )
 
+    ucg_3d = mb_result.uvg
+    vcg_3d = mb_result.vcg
+    ug_3d = mb_result.ug
+    vg_3d = mb_result.vg
+
     out_block = xr.Dataset(
         {
             "ucg": (in_block.ssh.dims, np.asarray(ucg_3d)),
@@ -135,12 +146,12 @@ optimizer = optax.chain(
 )
 ```
 
-And then pass the `optimizer` object as the `optim` argument of the [`cyclogeostrophy`](api.md#jaxparrow.cyclogeostrophy.cyclogeostrophy) function.
+And then pass the `optimizer` object as the `optim` argument of the [`minimization_based`](api.md#jaxparrow.cyclogeostrophy.minimization_based) function.
 This is employed in the [Pseudo-SWOT observations from eNATL60 model data](examples/swot-enatl60.ipynb) example.
 
 We also recommend using `JAX` floating point types with sufficient precision, e.g., `float64`:
 
 ```python
 import jax
 jax.config.update("jax_enable_x64", True)
-```
+```

jaxparrow/cyclogeostrophy/_core.py

@@ -386,10 +386,10 @@ def cyclogeostrophic_imbalance(
 
     # compute grid metrics once
     dx_e, dx_n, dy_e, dy_n, J = geometry.grid_metrics(lat_t, lon_t)
+    f = geometry.coriolis_factor(lat_t)
 
     return _cyclogeostrophic_imbalance(
-        ug_u, vg_v, ucg_u, vcg_v,
-        dx_u, dx_v, dy_u, dy_v, coriolis_factor_u, coriolis_factor_v, mask
+        ug, vg, ucg, vcg, dx_e, dx_n, dy_e, dy_n, J, f, land_mask
     )