From fba7f1998d11cdca619f884ba79fd01fa2ccc851 Mon Sep 17 00:00:00 2001
From: Mattia Almansi <m.almansi@bopen.eu>
Date: Thu, 4 Apr 2024 15:25:44 +0200
Subject: [PATCH] cleanup

---
 notebooks/wp5/ozone_merged_uv.ipynb | 74 ++++++++---------------------
 1 file changed, 21 insertions(+), 53 deletions(-)

diff --git a/notebooks/wp5/ozone_merged_uv.ipynb b/notebooks/wp5/ozone_merged_uv.ipynb
index f633b8e..b9f2823 100644
--- a/notebooks/wp5/ozone_merged_uv.ipynb
+++ b/notebooks/wp5/ozone_merged_uv.ipynb
@@ -49,30 +49,12 @@
     "\n",
     "# Regions\n",
     "region_slices = {\n",
-    "    \"global\": {\n",
-    "        \"lat_slice\": slice(-90, 90),\n",
-    "        \"lon_slice\": slice(0, 360),\n",
-    "    },\n",
-    "    \"tropics\": {\n",
-    "        \"lat_slice\": slice(-25, 25),\n",
-    "        \"lon_slice\": slice(0, 360),\n",
-    "    },\n",
-    "    \"NH mid-latitudes\": {\n",
-    "        \"lat_slice\": slice(30, 60),\n",
-    "        \"lon_slice\": slice(0, 360),\n",
-    "    },\n",
-    "    \"SH mid-latitudes\": {\n",
-    "        \"lat_slice\": slice(-30, -60),\n",
-    "        \"lon_slice\": slice(0, 360),\n",
-    "    },\n",
-    "    \"NH polar\": {\n",
-    "        \"lat_slice\": slice(60, 90),\n",
-    "        \"lon_slice\": slice(0, 360),\n",
-    "    },\n",
-    "    \"SH polar\": {\n",
-    "        \"lat_slice\": slice(-60, -90),\n",
-    "        \"lon_slice\": slice(0, 360),\n",
-    "    },\n",
+    "    \"global\": {\"lat_slice\": slice(-90, 90), \"lon_slice\": slice(0, 360)},\n",
+    "    \"tropics\": {\"lat_slice\": slice(-25, 25), \"lon_slice\": slice(0, 360)},\n",
+    "    \"NH mid-latitudes\": {\"lat_slice\": slice(30, 60), \"lon_slice\": slice(0, 360)},\n",
+    "    \"SH mid-latitudes\": {\"lat_slice\": slice(-30, -60), \"lon_slice\": slice(0, 360)},\n",
+    "    \"NH polar\": {\"lat_slice\": slice(60, 90), \"lon_slice\": slice(0, 360)},\n",
+    "    \"SH polar\": {\"lat_slice\": slice(-60, -90), \"lon_slice\": slice(0, 360)},\n",
     "}\n",
     "\n",
     "# Variable to show\n",
@@ -119,21 +101,14 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "def convert_source_to_time(ds):\n",
-    "    # Naming convention: YYYYMM-*.nc\n",
-    "    ds[\"source\"] = pd.to_datetime(ds[\"source\"].str.slice(None, 6), format=\"%Y%m\")\n",
-    "    return ds.rename(source=\"time\")\n",
+    "def add_time_dim(ds):\n",
+    "    time = pd.to_datetime(ds.attrs[\"time_coverage_start\"][:6], format=\"%Y%m\")\n",
+    "    return ds.expand_dims(time=[time])\n",
     "\n",
     "\n",
     "def spatial_weighted_mean(ds, lon_slice, lat_slice):\n",
-    "    ds = convert_source_to_time(ds)\n",
     "    ds = utils.regionalise(ds, lon_slice=lon_slice, lat_slice=lat_slice)\n",
-    "    return diagnostics.spatial_weighted_mean(ds)\n",
-    "\n",
-    "\n",
-    "def time_weighted_mean(ds):\n",
-    "    ds = convert_source_to_time(ds)\n",
-    "    return diagnostics.time_weighted_mean(ds)"
+    "    return diagnostics.spatial_weighted_mean(ds)"
    ]
   },
   {
@@ -149,7 +124,12 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "kwargs = {\"collection_id\": collection_id, \"requests\": requests, \"chunks\": {\"year\": 1}}\n",
+    "kwargs = {\n",
+    "    \"collection_id\": collection_id,\n",
+    "    \"requests\": requests,\n",
+    "    \"chunks\": {\"year\": 1},\n",
+    "    \"preprocess\": add_time_dim,\n",
+    "}\n",
     "\n",
     "# Timeseries\n",
     "datasets = []\n",
@@ -165,10 +145,7 @@
     "ds_timeseries = xr.concat(datasets, \"region\")\n",
     "\n",
     "# Original data\n",
-    "ds = download.download_and_transform(\n",
-    "    **kwargs,\n",
-    "    transform_func=convert_source_to_time,\n",
-    ")"
+    "ds = download.download_and_transform(**kwargs)"
    ]
   },
   {
@@ -186,20 +163,11 @@
    "source": [
     "da_obs = ds[f\"{varname}_number_of_observations\"]\n",
     "plot.projected_map(da_obs.sum(dim=\"time\", keep_attrs=True), show_stats=False)\n",
+    "plt.show()\n",
     "\n",
-    "fig, axs = plt.subplots(2, 1)\n",
-    "for ax, dim in zip(axs.flatten(), (\"longitude\", \"latitude\")):\n",
-    "    plot_obj = da_obs.sum(dim, keep_attrs=True).plot.contourf(\n",
-    "        x=\"time\",\n",
-    "        ax=ax,\n",
-    "        levels=range(0, 10_500 + 1, 1_500),\n",
-    "        extend=\"max\",\n",
-    "        add_colorbar=False,\n",
-    "    )\n",
-    "label = da_obs.attrs[\"long_name\"]\n",
-    "isplit = len(label) // 2\n",
-    "label = \"\\n\".join([label[:isplit], label[isplit:]])\n",
-    "_ = fig.colorbar(plot_obj, ax=axs.flatten(), label=label)"
+    "with xr.set_options(keep_attrs=True):\n",
+    "    da_sum = da_obs.sum(\"longitude\", keep_attrs=True) / da_obs.sizes[\"longitude\"]\n",
+    "da_sum.plot(x=\"time\")"
    ]
   },
   {