4_Predict_ensemble.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Generate an ensemble of gridded predictions\n",
    "\n",
    "Using the 30 models produced in `3_Generate_ensemble_of_models.ipynb`, we will generate an ensemble of 30 predictions. From this ensemble we will produce an uncertainty envelope, and a median prediction.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import sys\n",
    "import os\n",
    "import warnings\n",
    "import xarray as xr\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "from joblib import load\n",
    "from odc.geo.geobox import zoom_out\n",
    "from odc.algo import xr_reproject\n",
    "from datacube.utils.dask import start_local_dask\n",
    "from odc.geo.xr import assign_crs\n",
    "import odc.geo.xr\n",
    "# from dask.distributed import Client,Scheduler\n",
    "# from dask_jobqueue import SLURMCluster\n",
    "\n",
    "sys.path.append('/g/data/os22/chad_tmp/NEE_modelling/')\n",
    "from _collect_prediction_data import round_coords, collect_prediction_data \n",
    "\n",
    "sys.path.append('/g/data/os22/chad_tmp/dea-notebooks/Tools/')\n",
    "from dea_tools.classification import predict_xr, HiddenPrints"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "client = start_local_dask(mem_safety_margin='2Gb')\n",
    "client"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Analysis Parameters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "var = 'ET'\n",
    "suffix='20230320'\n",
    "results_path = '/g/data/os22/chad_tmp/NEE_modelling/results/predictions_uncertainty/'+var+'/'\n",
    "models_folder = '/g/data/os22/chad_tmp/NEE_modelling/results/models_uncertainty/'+var+'/'\n",
    "features_list = '/g/data/os22/chad_tmp/NEE_modelling/results/variables_'+suffix+'.txt'\n",
    "\n",
    "t1, t2='2003','2022'\n",
    "rescale=False"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Get paths to models"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model_list = [file for file in os.listdir(models_folder) if file.endswith(\".joblib\")]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Open predictor data\n",
    "\n",
    "At 1 km resolution, we need to pull the gridded feature layers in as dask arrays and compute on each time-step individually as the total memory requirements are very large. At 5 km resolution, its better to load the entire feature layer data into memory as it speeds up predictions."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# data = xr.open_dataset('/g/data/os22/chad_tmp/NEE_modelling/results/prediction_data/data_5km.nc')\n",
    "# mask = xr.open_dataarray('/g/data/os22/chad_tmp/NEE_modelling/results/prediction_data/mask_5km.nc')\n",
    "\n",
    "## open data\n",
    "data = collect_prediction_data(time_start=t1,\n",
    "                             time_end=t2,\n",
    "                             verbose=False,\n",
    "                             export=False,\n",
    "                             chunks=dict(latitude=680, longitude=1050, time=1) #chunks optimised\n",
    "                             )\n",
    "\n",
    "#precomputed the mask to save a little time\n",
    "mask = xr.open_dataarray('/g/data/os22/chad_tmp/NEE_modelling/data/1km/mask_1km_monthly_2003_2022.nc')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Check training and prediction variable order"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_vars = list(pd.read_csv(features_list))[0:-1]\n",
    "train_vars=[i[:-3] for i in train_vars]\n",
    "\n",
    "data = data[train_vars]\n",
    "\n",
    "if train_vars == list(data.data_vars):\n",
    "    print('Variables match, n: ', len(data.data_vars))\n",
    "else:\n",
    "    raise ValueError('Variables dont match')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Predict\n",
    "\n",
    "Loop through each model, and each time-step.  Mask the output with the urban mask.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#mask urban (5km res)\n",
    "# mask1 = xr.open_dataarray('/g/data/os22/chad_tmp/NEE_modelling/data/urban_mask_1km.nc')\n",
    "# mask1 = xr_reproject(mask1, geobox=data.odc.geobox.compat, resampling='mode')\n",
    "# mask1=round_coords(mask1)\n",
    "# mask1 = mask1.rename({'latitude':'y', 'longitude':'x'})\n",
    "\n",
    "#mask urban (1km res)\n",
    "mask1 = xr.open_dataarray('/g/data/os22/chad_tmp/NEE_modelling/data/urban_mask_1km.nc')\n",
    "mask1 = mask1.rename({'latitude':'y', 'longitude':'x'})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "\n",
    "# Loop through the 30 models\n",
    "for m in model_list:\n",
    "    name = m.split('.')[0]\n",
    "    \n",
    "    if os.path.exists(results_path+name+'.nc'):\n",
    "        print('skipping model '+name)\n",
    "        continue\n",
    "    \n",
    "    print('Model: ', name)\n",
    "    \n",
    "    warnings.filterwarnings(\"ignore\")\n",
    "    model = load(models_folder+m).set_params(n_jobs=1)\n",
    "    \n",
    "    results = []\n",
    "    i=0\n",
    "    #loop through the time-steps\n",
    "    for i in range(0, len(data.time)): \n",
    "        print(\"  {:03}/{:03}\\r\".format(i + 1, len(range(0, len(data.time)))), end=\"\")\n",
    "\n",
    "        with HiddenPrints():\n",
    "            warnings.filterwarnings(\"ignore\")\n",
    "            predicted = predict_xr(model,\n",
    "                                data.isel(time=i),\n",
    "                                proba=False,\n",
    "                                clean=True,\n",
    "                                chunk_size=875000, #this number is optimized to maximise pred speed.\n",
    "                                  ).compute()\n",
    "\n",
    "        predicted = predicted.Predictions.where(~mask.isel(time=i).compute())\n",
    "        predicted['time'] = data.isel(time=i).time.values\n",
    "        results.append(predicted.astype('float32'))\n",
    "        i+=1 \n",
    "    \n",
    "    ds = xr.concat(results, dim='time').sortby('time').rename(var).astype('float32')\n",
    "    \n",
    "    #mask urban\n",
    "    ds = ds.where(mask!=1).astype('float32')\n",
    "\n",
    "    #save results\n",
    "    ds.to_netcdf(results_path+name+'.nc')\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "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.9.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}