notebooks/06_bias_correction.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "78552b45-219e-49bd-80e2-62f680a860ff",
   "metadata": {},
   "source": [
    "# 06: Bias-correction\n",
    "*Develop a model that resolves differences between the climate model data and those from a more-detailed reference historical timeseries, and then use that model to ensure that future projections also reflect that level of detail.*"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "68b09f8c-391e-4854-acc4-c2cd88368df3",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "import fsspec\n",
    "import geopandas as gpd\n",
    "import numpy as np\n",
    "import xarray as xr\n",
    "from dask.distributed import Client\n",
    "from utils import gcm_list\n",
    "from xclim import sdba\n",
    "from xclim.sdba.adjustment import QuantileDeltaMapping"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "940d5a09-49c8-49c7-a471-35291cf3dda6",
   "metadata": {
    "tags": []
   },
   "source": [
    "Set up cluster to handle multiprocessing using a Dask client."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f37d494d-fdfd-4d18-b598-3594fede09ea",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "client = Client(n_workers=32)\n",
    "client"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7c930206-5144-4191-b895-84188e71f1dd",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "def load_projection(gcm, scenario):\n",
    "    \"\"\"\n",
    "    Load in a WBGT in the shade estimate produced by in `05_aggregate.ipynb`.\n",
    "    \"\"\"\n",
    "    ds = xr.open_zarr(\n",
    "        f\"s3://carbonplan-extreme-heat/temp/wbgt-shade-regions/{gcm}-{scenario}.zarr\"\n",
    "    )\n",
    "    ds[\"WBGT\"].attrs = {}\n",
    "    ds[\"WBGT\"].attrs[\"units\"] = \"degC\"\n",
    "    ds[\"processing_id\"] = ds[\"processing_id\"].astype(\"int\")\n",
    "    return ds"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "165dd8a7-d08e-4ef0-9de0-6c7c8b9086fc",
   "metadata": {
    "tags": []
   },
   "source": [
    "Load UHE-daily estimates developed in `05_aggregate.ipynb`. This data will be the reference for every bias-correction."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e15d692e-bc2c-4b4e-b473-a4b59f88a7ab",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "ref = xr.open_zarr(\n",
    "    \"s3://carbonplan-climate-impacts/extreme-heat/v1.0/inputs/wbgt-UHE-daily-historical.zarr\"\n",
    ")\n",
    "ref[\"WBGT\"].attrs[\"units\"] = \"degC\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a3c1769d-799b-4961-8369-d20af02e03f5",
   "metadata": {},
   "source": [
    "Load the region information."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c9ae188e-48ec-4358-9e0c-72d973d53484",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "with fsspec.open(\n",
    "    \"s3://carbonplan-climate-impacts/extreme-heat/v1.0/inputs/all_regions_and_cities.json\"\n",
    ") as file:\n",
    "    regions_df = gpd.read_file(file)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4af15747-d4bc-40a1-aa18-e57dbf094e09",
   "metadata": {},
   "outputs": [],
   "source": [
    "def train_bias_correction(ref_ts, model_ts, gcm):\n",
    "    \"\"\"\n",
    "    Prep timeseries for training and train the bias-correction model\n",
    "    \"\"\"\n",
    "\n",
    "    # convert all ts to the no-leap calendar and convert back to\n",
    "    # gregorian after prediction\n",
    "    ref_ts = ref_ts.convert_calendar(\"noleap\")\n",
    "    model_ts = model_ts.convert_calendar(\n",
    "        \"noleap\",\n",
    "        dim=\"time\",\n",
    "        align_on=\"year\",\n",
    "        missing=np.nan,\n",
    "        use_cftime=None,\n",
    "    )\n",
    "\n",
    "    # gap fill by linearly interpolating\n",
    "    model_ts = model_ts.interpolate_na(dim=\"time\", method=\"linear\").chunk({\"time\": -1})\n",
    "    group = sdba.Grouper(\"time.dayofyear\", window=31)\n",
    "    nquantiles = 100\n",
    "\n",
    "    # train the same model but using different groupers\n",
    "    trained_model = QuantileDeltaMapping.train(\n",
    "        ref_ts, model_ts, nquantiles=nquantiles, kind=\"+\", group=group\n",
    "    )\n",
    "\n",
    "    return trained_model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ffe5675f-6634-4fc4-8938-f4f2cc007656",
   "metadata": {},
   "outputs": [],
   "source": [
    "def apply_bias_correction(trained_model, ts, gcm, out_store):\n",
    "    \"\"\"\n",
    "    Apply trained bias-correction model to each model timeseries (whether historic\n",
    "    or future).\n",
    "    \"\"\"\n",
    "    bias_corrected = trained_model.adjust(ts)\n",
    "\n",
    "    # the rolling monthly bias-correction\n",
    "    # works with no-leap calendars so convert it back to gregorian\n",
    "    bias_corrected = (\n",
    "        bias_corrected.convert_calendar(\n",
    "            \"gregorian\",\n",
    "            align_on=\"year\",\n",
    "            missing=np.nan,\n",
    "            use_cftime=None,\n",
    "        )\n",
    "        .interpolate_na(dim=\"time\", method=\"linear\")\n",
    "        .chunk({\"time\": -1})\n",
    "    )\n",
    "\n",
    "    bias_corrected.to_zarr(out_store, mode=\"w\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c3d88f1b-e79d-4d31-bc60-f395f98611f7",
   "metadata": {},
   "source": [
    "Data isn't available for all regions. Only apply bias-correction where data is available in both the reference and the modelled datasets."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dc94cae2-36d9-4511-9223-acf6ddb5d8ae",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "ref_places = ref.processing_id.values\n",
    "modelled_places = load_projection(\"ACCESS-CM2\", \"historical\")[\n",
    "    \"WBGT\"\n",
    "].processing_id.values\n",
    "valid_ids = list(set(ref_places) & set(modelled_places))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "94ed2565-02ca-4309-a377-5ef89cdc531a",
   "metadata": {
    "tags": []
   },
   "source": [
    "Define the periods over which the bias-correction will occur."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ce8c4a8b-da59-4ce4-9119-0f1f986aab11",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "analysis_period_dict = {\n",
    "    \"historical\": slice(\"1985\", \"2014\"),\n",
    "    \"ssp245-2030\": slice(\"2020\", \"2039\"),\n",
    "    \"ssp245-2050\": slice(\"2040\", \"2059\"),\n",
    "}"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8a4cb78c-a62d-49f8-af7a-58a5626d7825",
   "metadata": {},
   "source": [
    "Subset the reference dataset to the historical time period used for training (1985-2014)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9b4ff787-4312-44a4-90fe-fd6c055da2da",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "ref = (\n",
    "    ref.sel(time=analysis_period_dict[\"historical\"])\n",
    "    .sel(processing_id=valid_ids)\n",
    "    .chunk({\"time\": -1, \"processing_id\": 850})\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1af916c5-21c2-4a5e-b8b4-bfc343c28392",
   "metadata": {},
   "source": [
    "Load in the different datasets into a dictionary, which, instead of an Xarray object, allows for the different calendars that different GCMs use. Then, for each GCM separately, train a bias-correction model and use it to bias-correct the historic and future projections."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a2165a09-d44b-4d80-82f8-c8996757780e",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "for gcm in gcm_list:\n",
    "    ts_dict = {}\n",
    "    ts_dict[\"reference\"] = ref[\"WBGT\"]\n",
    "    for scenario in [\"historical\", \"ssp245-2030\", \"ssp245-2050\"]:\n",
    "        ts_dict[scenario] = load_projection(gcm, scenario.split(\"-\")[0])[\"WBGT\"]\n",
    "        ts_dict[scenario] = (\n",
    "            ts_dict[scenario]\n",
    "            .sel(processing_id=valid_ids)\n",
    "            .chunk({\"time\": -1, \"processing_id\": 850})\n",
    "        )\n",
    "        ts_dict[scenario] = ts_dict[scenario].sel(time=analysis_period_dict[scenario])\n",
    "\n",
    "    trained_model = train_bias_correction(\n",
    "        ts_dict[\"reference\"], ts_dict[\"historical\"], gcm\n",
    "    )\n",
    "\n",
    "    for scenario in [\"historical\", \"ssp245-2030\", \"ssp245-2050\"]:\n",
    "        apply_bias_correction(\n",
    "            trained_model,\n",
    "            ts_dict[scenario],\n",
    "            gcm,\n",
    "            f\"s3://carbonplan-scratch/extreme-heat/wbgt-shade-regions/{gcm}-{scenario}-bc.zarr\",\n",
    "        )"
   ]
  }
 ],
 "metadata": {
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}