experimentation notebooks for the convolution of full maps

experimentation/convolution-numpy.ipynb

@@ -0,0 +1,164 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "0",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import numpy as np\n",
+    "\n",
+    "import healpix_convolution as hc\n",
+    "from healpix_convolution.convolution import convolve\n",
+    "from healpix_convolution.plotting import plot_healpix"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "1",
+   "metadata": {},
+   "source": [
+    "## load data"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "2",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import xarray as xr\n",
+    "\n",
+    "url = \"/home/jmagin/work/data/destine/average_surface_temperature.zarr\"\n",
+    "ds = xr.open_dataset(url, engine=\"zarr\", chunks={}).isel(oceanModelLayer=0)\n",
+    "ds"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "3",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "computed = ds.compute()\n",
+    "data = computed[\"avg_thetao\"].data\n",
+    "cell_ids = computed[\"cell_ids\"].data\n",
+    "resolution = computed[\"cell_ids\"].attrs[\"resolution\"]\n",
+    "indexing_scheme = computed[\"cell_ids\"].attrs[\"indexing_scheme\"]"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "4",
+   "metadata": {},
+   "source": [
+    "## compute the kernel"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "5",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "%%time\n",
+    "kernel = hc.kernels.gaussian_kernel(\n",
+    "    cell_ids,\n",
+    "    resolution=resolution,\n",
+    "    indexing_scheme=indexing_scheme,\n",
+    "    sigma=0.001,\n",
+    "    truncate=3,\n",
+    ")\n",
+    "kernel"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "6",
+   "metadata": {},
+   "source": [
+    "## convolve the data"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "7",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "%%time\n",
+    "# `axis` does not work, the spatial dimension needs to be last\n",
+    "# use `numpy.moveaxis` if it is not already the last dimension\n",
+    "result = convolve(data, kernel)\n",
+    "result"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "8",
+   "metadata": {},
+   "source": [
+    "## plot the result"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "9",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import matplotlib.pyplot as plt\n",
+    "import cartopy.crs as ccrs\n",
+    "\n",
+    "fig, axes = plt.subplots(\n",
+    "    ncols=1, nrows=2, figsize=(16, 14), subplot_kw={\"projection\": ccrs.Mollweide()}\n",
+    ")\n",
+    "plot_healpix(\n",
+    "    data,\n",
+    "    cell_ids,\n",
+    "    ax=axes[0],\n",
+    "    resolution=resolution,\n",
+    "    cmap=\"plasma\",\n",
+    "    title=\"original SST\",\n",
+    ")\n",
+    "plot_healpix(\n",
+    "    result,\n",
+    "    cell_ids,\n",
+    "    ax=axes[1],\n",
+    "    resolution=resolution,\n",
+    "    cmap=\"plasma\",\n",
+    "    title=\"result of the convolution\",\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "10",
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  }
+ ],
+ "metadata": {
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}