sync hit-rate

bopen · Mar 11, 2024 · 075e5f3 · 075e5f3
1 parent fb33335
commit 075e5f3
Showing 1 changed file with 30 additions and 11 deletions.
diff --git a/notebooks/wp3/hit_rate.ipynb b/notebooks/wp3/hit_rate.ipynb
@@ -94,17 +94,17 @@
     "}\n",
     "\n",
     "# Regions\n",
+    "defined_regions_module_name = \"ar6.land\"\n",
     "regions = [\n",
-    "    \"EAF\",\n",
     "    \"ENA\",\n",
     "    \"MED\",\n",
-    "    \"NEB\",\n",
+    "    \"NEAF\",\n",
+    "    \"NES\",\n",
+    "    \"NWS\",\n",
     "    \"SAS\",\n",
     "    \"SEA\",\n",
     "    \"WNA\",\n",
-    "    \"WSA\",\n",
-    "]\n",
-    "assert set(regions) <= set(regionmask.defined_regions.srex.abbrevs)"
+    "]"
    ]
   },
   {
@@ -122,7 +122,17 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "ax = regionmask.defined_regions.srex[regions].plot(\n",
+    "def get_defined_regions(defined_regions_module_name):\n",
+    "    defined_regions = regionmask.defined_regions\n",
+    "    for module in defined_regions_module_name.split(\".\"):\n",
+    "        defined_regions = getattr(defined_regions, module)\n",
+    "    return defined_regions\n",
+    "\n",
+    "\n",
+    "defined_regions = get_defined_regions(defined_regions_module_name)\n",
+    "assert set(regions) <= set(defined_regions.abbrevs)\n",
+    "\n",
+    "ax = defined_regions[regions].plot(\n",
     "    projection=ccrs.Robinson(),\n",
     "    label=\"abbrev\",\n",
     "    add_ocean=True,\n",
@@ -212,10 +222,11 @@
     "    return ds\n",
     "\n",
     "\n",
-    "def compute_tercile_occupation(ds, region, detrend):\n",
+    "def compute_tercile_occupation(ds, region, defined_regions_module_name, detrend):\n",
     "    # Mask region\n",
-    "    mask = regionmask.defined_regions.srex.mask(ds)\n",
-    "    index = regionmask.defined_regions.srex.map_keys(region)\n",
+    "    defined_regions = get_defined_regions(defined_regions_module_name)\n",
+    "    mask = defined_regions.mask(ds)\n",
+    "    index = defined_regions.map_keys(region)\n",
     "    ds = ds.where((mask == index).compute(), drop=True)\n",
     "    ds = ds.expand_dims(region=[region])\n",
     "\n",
@@ -293,7 +304,11 @@
     "            chunks=chunks,\n",
     "            transform_chunks=False,\n",
     "            transform_func=compute_tercile_occupation,\n",
-    "            transform_func_kwargs={\"region\": region, \"detrend\": detrend},\n",
+    "            transform_func_kwargs={\n",
+    "                \"region\": region,\n",
+    "                \"defined_regions_module_name\": defined_regions_module_name,\n",
+    "                \"detrend\": detrend,\n",
+    "            },\n",
     "        )\n",
     "        (da,) = ds.data_vars.values()\n",
     "        dataarrays.append(da.rename(variable))\n",
@@ -338,7 +353,11 @@
     "                    chunks=chunks,\n",
     "                    transform_chunks=False,\n",
     "                    transform_func=compute_tercile_occupation,\n",
-    "                    transform_func_kwargs={\"region\": region, \"detrend\": detrend},\n",
+    "                    transform_func_kwargs={\n",
+    "                        \"region\": region,\n",
+    "                        \"defined_regions_module_name\": defined_regions_module_name,\n",
+    "                        \"detrend\": detrend,\n",
+    "                    },\n",
     "                    backend_kwargs={\n",
     "                        \"time_dims\": (\n",
     "                            \"forecastMonth\",\n",