monthly ocean color

bopen · Aug 11, 2023 · ced0fe7 · ced0fe7
1 parent 855bfc1
commit ced0fe7
Showing 1 changed file with 36 additions and 48 deletions.
diff --git a/notebooks/wp5/WIP_ocean_color_reflectance.ipynb b/notebooks/wp5/WIP_ocean_color_reflectance.ipynb
@@ -23,11 +23,9 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "import tempfile\n",
-    "\n",
-    "import cartopy.crs as ccrs\n",
     "import matplotlib.pyplot as plt\n",
     "import numpy as np\n",
+    "import tqdm\n",
     "import xarray as xr\n",
     "from c3s_eqc_automatic_quality_control import diagnostics, download, plot, utils\n",
     "\n",
@@ -55,7 +53,14 @@
     "\n",
     "# Variable to analyse\n",
     "wavelength = 443\n",
-    "assert wavelength in (412, 443, 490, 510, 560, 665)"
+    "assert wavelength in (412, 443, 490, 510, 560, 665)\n",
+    "\n",
+    "# Regions to plot\n",
+    "regions = (\n",
+    "    {\"lon_slice\": slice(-180, 180), \"lat_slice\": slice(90, -90)},\n",
+    "    {\"lon_slice\": slice(-55, -40), \"lat_slice\": slice(30, 15)},\n",
+    "    {\"lon_slice\": slice(-80, 0), \"lat_slice\": slice(50, 0)},\n",
+    ")"
    ]
   },
   {
@@ -109,24 +114,26 @@
    "outputs": [],
    "source": [
     "def rechunk(obj):\n",
-    "    chunks = {\"time\": -1, \"year\": 1, \"longitude\": 270, \"latitude\": 270}\n",
+    "    chunks = {\"year\": 1, \"longitude\": 270, \"latitude\": 270}\n",
     "    return obj.chunk(**{k: v for k, v in chunks.items() if k in obj.dims})\n",
     "\n",
     "\n",
-    "def rrs_annual_weighted_log_mean(ds, wavelength):\n",
+    "def rrs_monthly_weighted_log_reductions(ds, wavelength):\n",
     "    name = f\"Rrs_{wavelength}\"\n",
     "    da = rechunk(ds[name])\n",
-    "    with tempfile.TemporaryDirectory() as tmpdir:\n",
-    "        da.to_zarr(tmpdir)\n",
-    "        da = xr.open_dataarray(tmpdir, engine=\"zarr\", chunks=dict(da.chunksizes))\n",
-    "        weights = np.abs(np.cos(np.deg2rad(da[\"latitude\"])))\n",
-    "        with xr.set_options(keep_attrs=True):\n",
-    "            da = 10 ** diagnostics.annual_weighted_mean(\n",
-    "                np.log10(da * weights), weights=False\n",
-    "            )\n",
-    "        da = rechunk(da.compute())\n",
-    "        da.encoding[\"chunksizes\"] = tuple(map(max, da.chunks))\n",
-    "        return da.to_dataset(name=name)"
+    "    weights = np.abs(np.cos(np.deg2rad(da[\"latitude\"])))\n",
+    "    with xr.set_options(keep_attrs=True):\n",
+    "        da = np.log10(da)\n",
+    "        da = da.groupby(\"time.year\").map(\n",
+    "            diagnostics.monthly_weighted_mean, weights=False\n",
+    "        )\n",
+    "        da = 10 ** (da + np.log10(weights))\n",
+    "        da = da.persist()\n",
+    "        da_mean = da.mean(\"month\").expand_dims(reduction=[\"mean\"])\n",
+    "        da_std = (10 ** np.log10(da).std(\"month\")).expand_dims(reduction=[\"std\"])\n",
+    "    da = xr.concat([da_mean, da_std], \"reduction\")\n",
+    "    da.encoding[\"chunksizes\"] = tuple(map(max, da.chunks))\n",
+    "    return da.to_dataset(name=name)"
    ]
   },
   {
@@ -145,63 +152,44 @@
    "outputs": [],
    "source": [
     "datasets = []\n",
-    "for year in range(year_start, year_stop + 1):\n",
-    "    print(f\"{year=}\")\n",
+    "for year in tqdm.tqdm(range(year_start, year_stop + 1), desc=\"year\"):\n",
     "    requests = download.update_request_date(\n",
     "        request, start=f\"{year}-01\", stop=f\"{year}-12\", stringify_dates=True\n",
     "    )\n",
     "    ds = download.download_and_transform(\n",
     "        collection_id,\n",
     "        requests,\n",
     "        transform_chunks=False,\n",
-    "        transform_func=rrs_annual_weighted_log_mean,\n",
+    "        transform_func=rrs_monthly_weighted_log_reductions,\n",
     "        transform_func_kwargs={\"wavelength\": wavelength},\n",
     "        chunks={\"year\": 1, \"month\": 1},\n",
     "        **open_mfdataset_kwargs,\n",
     "    )\n",
     "    datasets.append(rechunk(ds))\n",
-    "ds = xr.concat(datasets, \"year\")"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "id": "94b4fe5c-2eae-4847-bb17-66fbd531b522",
-   "metadata": {},
-   "source": [
-    "## Plot global maps"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "id": "afd7484e-f23a-4558-83ad-779504e68c8d",
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "da = ds[f\"Rrs_{wavelength}\"]\n",
-    "plot_kwargs = {\"col\": \"year\", \"col_wrap\": 2, \"robust\": True}\n",
-    "\n",
-    "da = da.coarsen(latitude=5, longitude=5).mean()\n",
-    "facet = plot.projected_map(da, projection=ccrs.Robinson(), **plot_kwargs)"
+    "da = xr.concat(datasets, \"year\")[f\"Rrs_{wavelength}\"]"
    ]
   },
   {
    "cell_type": "markdown",
-   "id": "1a4a6b69-cca8-4b7f-84ba-a724909db848",
+   "id": "7aaf7e33-a440-4c6a-a3dd-0e08010e3eb1",
    "metadata": {},
    "source": [
-    "## Plot regional maps"
+    "## Plot maps"
    ]
   },
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "9a030fe6-6cf4-4f65-914a-fdb70ee1d943",
+   "id": "032853a4-e6fb-4d89-8593-76f5163da3f4",
    "metadata": {},
    "outputs": [],
    "source": [
-    "da_region = utils.regionalise(da, lon_slice=slice(-55, -40), lat_slice=slice(30, 15))\n",
-    "facet = plot.projected_map(da_region, projection=ccrs.PlateCarree(), **plot_kwargs)"
+    "for regionalise_kwargs in regions:\n",
+    "    da_region = utils.regionalise(da, **regionalise_kwargs)\n",
+    "    for reduction, da_reduction in da_region.groupby(\"reduction\"):\n",
+    "        plot.projected_map(da_reduction, col=\"year\", robust=True)\n",
+    "        plt.suptitle(f\"{reduction.title()} ({year_start}-{year_stop})\")\n",
+    "        plt.show()"
    ]
   }
  ],