Noise strength map

Smooth the noisy current to find the noise strength map; some things are
a bit broken, like plotting and the axis directions, but its ok for
now...

Author: richard-lane

src/current_denoising/generation/applying_noise.py

@@ -0,0 +1,28 @@
+"""
+Utilities for applying noise to data
+
+"""
+
+import numpy as np
+from scipy.ndimage import gaussian_filter
+
+
+def noise_strength_map(current_grid: np.ndarray, *, filter_size: int) -> np.ndarray:
+    """
+    Get the expected noise strength map from a gridded field of current
+    velocities.
+
+    :param current_grid: gridded field showing current strength; land should be indicated
+                         with np.nan.
+    :param filter_size: size of Gaussian filter to apply, in grid points
+    :returns: a grid of the same shape as current_grid, showing the expected noise strength
+              at that grid point. Land values are set to np.nan.
+    """
+    land_mask = np.isnan(current_grid)
+
+    current_grid = np.where(land_mask, 0.0, current_grid)
+
+    filtered = gaussian_filter(current_grid, sigma=filter_size)
+    filtered[land_mask] = np.nan
+
+    return filtered

src/notebooks/plot_tile_rms.ipynb

@@ -10,6 +10,17 @@
     "Plot the RMS of the $8\\times8\\degree$ MDT current tiles, to visualise how the energy in the tiles varies - we want to exclude the regions near coastlines when picking our training data."
    ]
   },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "002e5fcb",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "%load_ext autoreload\n",
+    "%autoreload 2"
+   ]
+  },
   {
    "cell_type": "code",
    "execution_count": null,
@@ -141,6 +152,92 @@
     "\n",
     "fig.colorbar(im, ax=axes[\"A\"])"
    ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "3a5a91c1",
+   "metadata": {},
+   "source": [
+    "# Noise application\n",
+    "\n",
+    "We estimate noise strength by creating a noise strength map - visualise this here"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "2a99491a",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from current_denoising.generation import applying_noise\n",
+    "\n",
+    "filter_size = 5\n",
+    "\n",
+    "filtered = applying_noise.noise_strength_map(data, filter_size=filter_size)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "40fc172c",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "fig, axis = plt.subplots(1, 1, figsize=(12, 6))\n",
+    "\n",
+    "im = axis.imshow(filtered, **imshow_kw)\n",
+    "\n",
+    "earth_radius_km = 6370\n",
+    "axis.set_title(\n",
+    "    f\"{filter_size=} grid point\\n\"\n",
+    "    r\"$\\approx$\"\n",
+    "    f\"{earth_radius_km / 12 * filter_size:.0f}km at equator\"\n",
+    ")\n",
+    "fig.colorbar(im)\n",
+    "\n",
+    "lat, long = maps.lat_long_grid(data.shape)\n",
+    "extent = [long[0], long[-1], lat[0], lat[-1]]\n",
+    "im.set_extent(extent)\n",
+    "\n",
+    "fig.tight_layout()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "0b754829",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "\"\"\"Plot a tile near Australia to reproduce Fig 9 in Laura's paper\"\"\"\n",
+    "\n",
+    "import numpy as np\n",
+    "from current_denoising.generation import ioutils\n",
+    "\n",
+    "# NB the latitude is backwards here, so it is positive instead of negative\n",
+    "lat, long = -38, -70\n",
+    "size = 128\n",
+    "\n",
+    "# Get the index of the nearest grid point to each of these\n",
+    "indices = []\n",
+    "for a, b in zip([lat, long], maps.lat_long_grid(data.shape)):\n",
+    "    indices.append(int(np.argmin(np.abs(b - a))))\n",
+    "\n",
+    "tile = ioutils._tile(filtered, indices, size)\n",
+    "\n",
+    "fig, axis = plt.subplots()\n",
+    "axis.imshow(tile, **imshow_kw)\n",
+    "axis.set_axis_off()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "63bbc617",
+   "metadata": {},
+   "outputs": [],
+   "source": []
   }
  ],
  "metadata": {

src/tests/ut/test_applying_noise.py

@@ -0,0 +1,25 @@
+"""
+Tests for the noise application utils
+"""
+
+import numpy as np
+
+from current_denoising.generation import applying_noise
+
+
+def test_noise_strength_land_mask():
+    """
+    Check that the land is in the right place in our noise strength map
+    """
+    map = np.arange(100, dtype=np.float32).reshape(10, 10)
+    # Set some values to indicate that they're land
+    map[(map % 3).astype(bool)] = np.nan
+
+    assert (
+        np.isnan(map) == np.isnan(applying_noise.noise_strength_map(map, filter_size=5))
+    ).all()
+
+    assert (
+        np.isnan(map)
+        == np.isnan(applying_noise.noise_strength_map(map, filter_size=50))
+    ).all()