separate seasons

notebooks/wp4/clima_and_bias_cordex_cmip6_regionalised.ipynb

@@ -53,7 +53,7 @@
     "year_stop = 2000\n",
     "\n",
     "# Choose annual or seasonal timeseries\n",
-    "timeseries = \"JJA\"\n",
+    "timeseries = \"annual\"\n",
     "assert timeseries in (\"annual\", \"DJF\", \"MAM\", \"JJA\", \"SON\")\n",
     "\n",
     "# Variable\n",
@@ -260,15 +260,11 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "def rechunk(obj):\n",
-    "    return obj.chunk(**{dim: 1 for dim in set(obj.dims) & {\"season\", \"year\"}})\n",
-    "\n",
-    "\n",
     "def compute_regridded_timeseries(\n",
-    "    ds, year_start, year_stop, annual, grid_out=None, **kwargs\n",
+    "    ds, year_start, year_stop, timeseries, grid_out=None, **kwargs\n",
     "):\n",
     "    # Drop useless data\n",
-    "    if annual:\n",
+    "    if timeseries == \"annual\":\n",
     "        mask = (ds[\"time\"].dt.year >= year_start) & (ds[\"time\"].dt.year <= year_stop)\n",
     "    else:\n",
     "        # Select years (shift -1 to get D(year-1)J(year)F(year))\n",
@@ -290,11 +286,9 @@
     "    da = ds[var_name].rename(field)\n",
     "\n",
     "    # Create timeseries\n",
-    "    if annual:\n",
-    "        da = diagnostics.annual_weighted_mean(da)\n",
-    "    else:\n",
-    "        da[\"year\"] = da[\"year\"].astype(int)\n",
-    "        da = da.groupby(\"year\").map(diagnostics.seasonal_weighted_mean)\n",
+    "    da = diagnostics.annual_weighted_mean(da)\n",
+    "    if timeseries != \"annual\":\n",
+    "        da = da.assign_coords(season=timeseries)\n",
     "\n",
     "    # Regid\n",
     "    if grid_out is not None:\n",
@@ -311,7 +305,6 @@
     "            da = da.assign_coords(height=ds[\"height\"] if \"height\" in ds else None)\n",
     "        else:\n",
     "            raise ValueError\n",
-    "    da = rechunk(da)\n",
     "    da.encoding[\"chunksizes\"] = tuple(map(max, da.chunks))\n",
     "    return da.to_dataset()"
    ]
@@ -330,8 +323,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "annual = timeseries == \"annual\"\n",
-    "transform_chunks = annual\n",
+    "transform_chunks = timeseries == \"annual\"\n",
     "print(\"Downloading and processing ERA5\")\n",
     "ds_era = download.download_and_transform(\n",
     "    *request_era,\n",
@@ -341,10 +333,10 @@
     "    transform_func_kwargs={\n",
     "        \"year_start\": year_start,\n",
     "        \"year_stop\": year_stop,\n",
-    "        \"annual\": annual,\n",
+    "        \"timeseries\": timeseries,\n",
     "    },\n",
     ").expand_dims(model=[\"ERA5\"])\n",
-    "ds_era = rechunk(ds_era)\n",
+    "ds_era = ds_era.chunk(year=chunks.get(\"year\", -1))\n",
     "\n",
     "datasets = []\n",
     "for model in models:\n",
@@ -360,12 +352,12 @@
     "            \"grid_out\": ds_era[[\"longitude\", \"latitude\"]],\n",
     "            \"year_start\": year_start,\n",
     "            \"year_stop\": year_stop,\n",
-    "            \"annual\": annual,\n",
+    "            \"timeseries\": timeseries,\n",
     "            \"method\": \"bilinear\",\n",
     "            \"periodic\": periodic,\n",
     "        },\n",
     "    )\n",
-    "    datasets.append(rechunk(ds).expand_dims(model=[model]))\n",
+    "    datasets.append(ds.chunk(year=chunks.get(\"year\", -1)).expand_dims(model=[model]))\n",
     "ds_sim = xr.concat(datasets, \"model\")"
    ]
   },
@@ -391,8 +383,6 @@
     "    ],\n",
     "    \"model\",\n",
     ")\n",
-    "if timeseries != \"annual\":\n",
-    "    ds_timeseries = ds_timeseries.sel(season=timeseries)\n",
     "ds_timeseries = utils.regionalise(\n",
     "    ds_timeseries, lon_slice=lon_slice, lat_slice=lat_slice\n",
     ")\n",
@@ -405,7 +395,7 @@
     "        transform_func=utils.regionalise,\n",
     "        transform_func_kwargs={\"lon_slice\": lon_slice, \"lat_slice\": lat_slice},\n",
     "    )\n",
-    "    da_lsm = ds_lsm[\"lsm\"].squeeze(\"time\", drop=True).compute()\n",
+    "    da_lsm = ds_lsm[\"lsm\"].squeeze(\"time\", drop=True)\n",
     "    ds = ds.where(da_lsm > 0.5)\n",
     "ds = ds.compute()\n",
     "with xr.set_options(keep_attrs=True):\n",