allow multiple variables

bopen · Oct 12, 2023 · 7f88aae · 7f88aae
1 parent d8f2fbb
commit 7f88aae
Showing 1 changed file with 34 additions and 29 deletions.
diff --git a/notebooks/wp5/ocean_color_timeseries.ipynb b/notebooks/wp5/ocean_color_timeseries.ipynb
@@ -54,16 +54,16 @@
     "stop = \"2023-03\"\n",
     "\n",
     "# Variable to analyse\n",
-    "variable = \"chlor_a\"\n",
-    "assert variable in (\n",
+    "variables = [\"chlor_a\", \"Rrs_443\", \"Rrs_560\"]\n",
+    "assert set(variables) <= {\n",
     "    \"chlor_a\",\n",
     "    \"Rrs_412\",\n",
     "    \"Rrs_443\",\n",
     "    \"Rrs_490\",\n",
     "    \"Rrs_510\",\n",
     "    \"Rrs_560\",\n",
     "    \"Rrs_665\",\n",
-    ")\n",
+    "}\n",
     "\n",
     "# Regions to plot\n",
     "regions = {\n",
@@ -96,17 +96,10 @@
     "collection_id = \"satellite-ocean-colour\"\n",
     "\n",
     "request = {\n",
-    "    \"variable\": \"remote_sensing_reflectance\"\n",
-    "    if variable.startswith(\"Rrs\")\n",
-    "    else \"mass_concentration_of_chlorophyll_a\",\n",
     "    \"projection\": \"regular_latitude_longitude_grid\",\n",
     "    \"version\": \"6_0\",\n",
     "    \"format\": \"zip\",\n",
-    "}\n",
-    "\n",
-    "requests = download.update_request_date(\n",
-    "    request, start=start, stop=stop, stringify_dates=True\n",
-    ")"
+    "}"
    ]
   },
   {
@@ -151,21 +144,32 @@
    "outputs": [],
    "source": [
     "datasets = []\n",
-    "for region, slices in regions.items():\n",
-    "    ds = download.download_and_transform(\n",
-    "        collection_id,\n",
-    "        requests,\n",
-    "        transform_func=regionalised_spatial_weighted_mean,\n",
-    "        transform_func_kwargs={\"variable\": variable} | slices,\n",
-    "        chunks={\"year\": 1, \"month\": 1},\n",
-    "    )\n",
-    "    datasets.append(ds.expand_dims(latitudes=[region]))\n",
-    "ds = xr.concat(datasets, \"latitudes\")\n",
+    "for variable in variables:\n",
+    "    for region, slices in regions.items():\n",
+    "        requests = download.update_request_date(\n",
+    "            request\n",
+    "            | {\n",
+    "                \"variable\": \"remote_sensing_reflectance\"\n",
+    "                if variable.startswith(\"Rrs\")\n",
+    "                else \"mass_concentration_of_chlorophyll_a\"\n",
+    "            },\n",
+    "            start=start,\n",
+    "            stop=stop,\n",
+    "            stringify_dates=True,\n",
+    "        )\n",
+    "        ds = download.download_and_transform(\n",
+    "            collection_id,\n",
+    "            requests,\n",
+    "            transform_func=regionalised_spatial_weighted_mean,\n",
+    "            transform_func_kwargs={\"variable\": variable} | slices,\n",
+    "            chunks={\"year\": 1, \"month\": 1, \"variable\": 1},\n",
+    "        )\n",
+    "        datasets.append(ds.expand_dims(latitudes=[region]))\n",
+    "ds = xr.merge(datasets)\n",
     "\n",
-    "# Extract DataArray\n",
-    "da = ds[variable]\n",
-    "da_global = da.sel(latitudes=[\"Global\"])\n",
-    "da_regional = da.drop_sel(latitudes=\"Global\")"
+    "# Extract global and regional\n",
+    "ds_global = ds.sel(latitudes=[\"Global\"])\n",
+    "ds_regional = ds.drop_sel(latitudes=\"Global\")"
    ]
   },
   {
@@ -235,10 +239,11 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "for da_to_plot in [da_global, da_regional]:\n",
-    "    for plot_func in plot_daily, plot_monthly:\n",
-    "        plot_func(da_to_plot)\n",
-    "        plt.show()"
+    "for variable in variables:\n",
+    "    for ds in [ds_global, ds_regional]:\n",
+    "        for plot_func in plot_daily, plot_monthly:\n",
+    "            plot_func(ds[variable])\n",
+    "            plt.show()"
    ]
   }
  ],