From bc8d1520b66e4cb2cc290950801d4bae114628d2 Mon Sep 17 00:00:00 2001
From: Mattia Almansi <m.almansi@bopen.eu>
Date: Wed, 13 Mar 2024 10:21:25 +0100
Subject: [PATCH] fix future

---
 .../extreme_temperature_indices_future.ipynb  | 61 ++++++++-----------
 1 file changed, 27 insertions(+), 34 deletions(-)

diff --git a/notebooks/wp4/extreme_temperature_indices_future.ipynb b/notebooks/wp4/extreme_temperature_indices_future.ipynb
index fcf7792..b51d510 100644
--- a/notebooks/wp4/extreme_temperature_indices_future.ipynb
+++ b/notebooks/wp4/extreme_temperature_indices_future.ipynb
@@ -85,7 +85,7 @@
     "assert timeseries in (\"annual\", \"DJF\", \"MAM\", \"JJA\", \"SON\")\n",
     "\n",
     "# Choose CORDEX or CMIP6\n",
-    "collection_id = \"CMIP6\"\n",
+    "collection_id = \"CORDEX\"\n",
     "assert collection_id in (\"CORDEX\", \"CMIP6\")\n",
     "\n",
     "# Define region for analysis\n",
@@ -104,16 +104,6 @@
     "chunks = {\"year\": 1}"
    ]
   },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "if collection_id == \"CORDEX\":\n",
-    "    raise NotImplementedError(f\"{collection_id=}\")"
-   ]
-  },
   {
    "attachments": {},
    "cell_type": "markdown",
@@ -206,7 +196,7 @@
     "request_cordex = {\n",
     "    \"format\": \"zip\",\n",
     "    \"domain\": cordex_domain,\n",
-    "    \"experiment\": \"historical\",\n",
+    "    \"experiment\": \"rcp_8_5\",\n",
     "    \"horizontal_resolution\": \"0_11_degree_x_0_11_degree\",\n",
     "    \"temporal_resolution\": \"daily_mean\",\n",
     "    \"variable\": \"maximum_2m_temperature_in_the_last_24_hours\",\n",
@@ -233,8 +223,8 @@
     "def get_cordex_years(\n",
     "    year_start,\n",
     "    year_stop,\n",
-    "    start_years=[1971, 1976, 1981, 1986, 1991],\n",
-    "    end_years=[1975, 1980, 1985, 1990, 1995],\n",
+    "    start_years=list(range(2006, 2097, 5)),\n",
+    "    end_years=list(range(2010, 2101, 5)),\n",
     "):\n",
     "    start_year = []\n",
     "    end_year = []\n",
@@ -308,9 +298,8 @@
     "\n",
     "    ds = xr.open_mfdataset(paths)\n",
     "    in_files = f\"{tmpdir}/rechunked.zarr\"\n",
-    "    ds.chunk({dim: -1 if dim == \"time\" else \"auto\" for dim in ds.dims}).to_zarr(\n",
-    "        in_files\n",
-    "    )\n",
+    "    chunks = {dim: -1 if dim == \"time\" else \"auto\" for dim in ds.dims}\n",
+    "    ds.chunk(chunks).to_zarr(in_files)\n",
     "\n",
     "    datasets = [\n",
     "        icclim.index(\n",
@@ -418,7 +407,6 @@
     "    \"year_stop\": year_stop,\n",
     "    \"resample\": False,\n",
     "}\n",
-    "\n",
     "ds_regrid = download.download_and_transform(\n",
     "    requests=[request | {model_key: model_regrid} for request in request_sim[1]],\n",
     "    **kwargs,\n",
@@ -445,7 +433,10 @@
    },
    "outputs": [],
    "source": [
-    "grid_out = ds_regrid[[\"latitude\", \"longitude\"]].reset_coords(drop=True)\n",
+    "coords = [\"latitude\", \"longitude\"]\n",
+    "grid_out = ds_regrid[coords]\n",
+    "coords_to_drop = set(grid_out.coords) - set(coords) - set(grid_out.dims)\n",
+    "grid_out = ds_regrid[coords].reset_coords(coords_to_drop, drop=True)\n",
     "grid_out.attrs = {}\n",
     "\n",
     "interpolated_datasets = []\n",
@@ -453,23 +444,25 @@
     "for model in models:\n",
     "    print(f\"{model=}\")\n",
     "    # Original model\n",
-    "    model_datasets[model] = download.download_and_transform(\n",
-    "        requests=[request | {model_key: model} for request in request_sim[1]],\n",
-    "        **kwargs,\n",
-    "        transform_func_kwargs=transform_func_kwargs,\n",
-    "    )\n",
-    "\n",
-    "    # Interpolated model\n",
     "    ds = download.download_and_transform(\n",
     "        requests=[request | {model_key: model} for request in request_sim[1]],\n",
     "        **kwargs,\n",
-    "        transform_func_kwargs=transform_func_kwargs\n",
-    "        | {\n",
-    "            \"grid_out\": grid_out,\n",
-    "            \"method\": interpolation_method,\n",
-    "            \"skipna\": True,\n",
-    "        },\n",
+    "        transform_func_kwargs=transform_func_kwargs,\n",
     "    )\n",
+    "    model_datasets[model] = ds\n",
+    "\n",
+    "    if model != model_regrid:\n",
+    "        # Interpolated model\n",
+    "        ds = download.download_and_transform(\n",
+    "            requests=[request | {model_key: model} for request in request_sim[1]],\n",
+    "            **kwargs,\n",
+    "            transform_func_kwargs=transform_func_kwargs\n",
+    "            | {\n",
+    "                \"grid_out\": grid_out,\n",
+    "                \"method\": interpolation_method,\n",
+    "                \"skipna\": True,\n",
+    "            },\n",
+    "        )\n",
     "    interpolated_datasets.append(ds.expand_dims(model=[model]))\n",
     "\n",
     "ds_interpolated = xr.concat(interpolated_datasets, \"model\")"
@@ -499,7 +492,7 @@
     "}\n",
     "\n",
     "# Remove boundaries\n",
-    "isel = {coord: slice(5, -5) for coord in (\"latitude\", \"longitude\")}\n",
+    "isel = {dim: slice(5, -5) for dim in set(ds.dims) - {\"model\", \"index\"}}\n",
     "ds_interpolated = ds_interpolated.isel(isel)\n",
     "\n",
     "# Edit attributes\n",
@@ -828,7 +821,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.11.7"
+   "version": "3.11.8"
   }
  },
  "nbformat": 4,