fix anomaly

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"
    ]
   },
   {