From f4166a098c432fc1cc34ba45beeb863d36d26f1c Mon Sep 17 00:00:00 2001
From: Mattia Almansi <m.almansi@bopen.eu>
Date: Fri, 17 May 2024 12:23:44 +0200
Subject: [PATCH] fix ocean colour

---
 notebooks/wp5/ocean_color_timeseries.ipynb | 46 ++++++++++++++++------
 1 file changed, 33 insertions(+), 13 deletions(-)

diff --git a/notebooks/wp5/ocean_color_timeseries.ipynb b/notebooks/wp5/ocean_color_timeseries.ipynb
index f3fdefe..cf78dc5 100644
--- a/notebooks/wp5/ocean_color_timeseries.ipynb
+++ b/notebooks/wp5/ocean_color_timeseries.ipynb
@@ -66,13 +66,22 @@
     "\n",
     "# Regions to plot\n",
     "regions = {\n",
-    "    \"Global Ocean\": {\"lon_slice\": slice(-180, 180), \"lat_slice\": slice(90, -90)},\n",
+    "    \"Global Ocean\": {\n",
+    "        \"lon_slice\": slice(-180, 180),\n",
+    "        \"lat_slice\": slice(90, -90),\n",
+    "    },\n",
     "    \"Latitudes 50$^o$S-50$^o$N\": {\n",
     "        \"lon_slice\": slice(-180, 180),\n",
     "        \"lat_slice\": slice(50, -50),\n",
     "    },\n",
-    "    \"NASTG\": {\"lon_slice\": slice(-80, 0), \"lat_slice\": slice(50, 0)},\n",
-    "    \"SPSTG\": {\"lon_slice\": slice(-120, -110), \"lat_slice\": slice(-20, -30)},\n",
+    "    \"NASTG\": {\n",
+    "        \"lon_slice\": slice(-80, 0),\n",
+    "        \"lat_slice\": slice(50, 0),\n",
+    "    },\n",
+    "    \"SPSTG\": {\n",
+    "        \"lon_slice\": slice(-120, -110),\n",
+    "        \"lat_slice\": slice(-20, -30),\n",
+    "    },\n",
     "}\n",
     "for region, slices in regions.items():\n",
     "    # Enforce sorting as original data\n",
@@ -119,18 +128,23 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "def regionalised_spatial_weighted_mean(ds, variable, lon_slice, lat_slice):\n",
+    "def regionalised_spatial_weighted_mean(\n",
+    "    ds, variable, lon_slice, lat_slice, log_transformed, vmin, vmax\n",
+    "):\n",
     "    da = ds[variable]\n",
     "    da = utils.regionalise(da, lon_slice=lon_slice, lat_slice=lat_slice)\n",
-    "    with xr.set_options(keep_attrs=True):\n",
-    "        if variable == \"chlor_a\":\n",
-    "            da = da.where((da > 0.01) & (da < 1.0e2))\n",
+    "    if vmin is not None:\n",
+    "        da = da.where(da > vmin)\n",
+    "    if vmax is not None:\n",
+    "        da = da.where(da < vmax)\n",
+    "    if log_transformed:\n",
+    "        with xr.set_options(keep_attrs=True):\n",
     "            da = 10 ** diagnostics.spatial_weighted_mean(np.log10(da))\n",
-    "            da.attrs[\"long_name\"] = da.attrs[\"long_name\"].replace(\n",
-    "                \" (not log-transformed)\", \"\"\n",
-    "            )\n",
-    "        else:\n",
-    "            da = diagnostics.spatial_weighted_mean(da)\n",
+    "        da.attrs[\"long_name\"] = da.attrs[\"long_name\"].replace(\n",
+    "            \" (not log-transformed)\", \"\"\n",
+    "        )\n",
+    "    else:\n",
+    "        da = diagnostics.spatial_weighted_mean(da)\n",
     "    return da.to_dataset(name=variable)"
    ]
   },
@@ -168,7 +182,13 @@
     "            collection_id,\n",
     "            requests,\n",
     "            transform_func=regionalised_spatial_weighted_mean,\n",
-    "            transform_func_kwargs={\"variable\": variable} | slices,\n",
+    "            transform_func_kwargs=slices\n",
+    "            | {\n",
+    "                \"variable\": variable,\n",
+    "                \"vmin\": 1.0e-2 if variable == \"chlor_a\" else None,\n",
+    "                \"vmax\": 1.0e2 if variable == \"chlor_a\" else None,\n",
+    "                \"log_transformed\": variable == \"chlor_a\",\n",
+    "            },\n",
     "            chunks={\"year\": 1, \"month\": 1, \"variable\": 1},\n",
     "        )\n",
     "        datasets.append(ds.expand_dims(region=[region]))\n",