From c7590e23712a4d454006f72318ddf6325541bf0c Mon Sep 17 00:00:00 2001
From: Mattia Almansi <m.almansi@bopen.eu>
Date: Tue, 6 Feb 2024 10:30:00 +0100
Subject: [PATCH] fix anomaly

---
 notebooks/wp3/WIP-hit_rate.ipynb | 23 +++++++++++------------
 1 file changed, 11 insertions(+), 12 deletions(-)

diff --git a/notebooks/wp3/WIP-hit_rate.ipynb b/notebooks/wp3/WIP-hit_rate.ipynb
index f7b0d9f..3b1116e 100644
--- a/notebooks/wp3/WIP-hit_rate.ipynb
+++ b/notebooks/wp3/WIP-hit_rate.ipynb
@@ -183,6 +183,9 @@
    "outputs": [],
    "source": [
     "def compute_tercile_occupation(ds, region):\n",
+    "    # Anomaly\n",
+    "    ds = ds - diagnostics.time_weighted_mean(ds)\n",
+    "\n",
     "    # Reindex using year/month\n",
     "    time = ds[\"forecast_reference_time\"]\n",
     "    ds = ds.assign_coords(\n",
@@ -191,25 +194,21 @@
     "    )\n",
     "    ds = ds.set_index({time.name: (\"year\", \"month\")}).unstack(time.name)\n",
     "\n",
-    "    # Anomaly\n",
-    "    (da,) = ds.data_vars.values()\n",
-    "    da = da - diagnostics.spatial_weighted_mean(da, weights=False)\n",
-    "\n",
     "    # Mask region\n",
-    "    mask = regionmask.defined_regions.srex.mask(da)\n",
+    "    mask = regionmask.defined_regions.srex.mask(ds)\n",
     "    index = regionmask.defined_regions.srex.map_keys(region)\n",
-    "    da = da.where((mask == index).compute(), drop=True)\n",
+    "    ds = ds.where((mask == index).compute(), drop=True)\n",
     "\n",
     "    # Spatial mean\n",
-    "    da = diagnostics.spatial_weighted_mean(da, weights=False)\n",
+    "    ds = diagnostics.spatial_weighted_mean(ds, weights=False)\n",
     "\n",
     "    # Get quantiles\n",
-    "    quantiles = da.chunk(year=-1).quantile([1 / 3, 2 / 3], \"year\")\n",
-    "    mask = xr.zeros_like(da, None)\n",
-    "    mask = xr.where(da < quantiles.sel(quantile=1 / 3), -1, mask)\n",
-    "    mask = xr.where(da > quantiles.sel(quantile=2 / 3), 1, mask)\n",
+    "    quantiles = ds.chunk(year=-1).quantile([1 / 3, 2 / 3], \"year\")\n",
+    "    mask = xr.zeros_like(ds, None)\n",
+    "    mask = xr.where(ds < quantiles.sel(quantile=1 / 3), -1, mask)\n",
+    "    mask = xr.where(ds > quantiles.sel(quantile=2 / 3), 1, mask)\n",
     "\n",
-    "    return mask.to_dataset()"
+    "    return mask"
    ]
   },
   {