6_Plot_predictions.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Plotting results of predictions\n",
    "\n",
    "Here, the outputs from `5_combine_ensembles.ipynb` are plotted. There are various types of plots, some that are used in the manuscript and some that are not\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import folium\n",
    "import odc.geo\n",
    "import numpy as np\n",
    "import xarray as xr\n",
    "import geopandas as gpd\n",
    "import contextily as ctx\n",
    "import matplotlib.pyplot as plt\n",
    "import matplotlib.dates as mdates\n",
    "from odc.geo.xr import assign_crs\n",
    "from odc.algo import xr_reproject\n",
    "from datacube.utils.cog import write_cog\n",
    "from matplotlib.colors import ListedColormap\n",
    "\n",
    "import sys\n",
    "sys.path.append('/g/data/os22/chad_tmp/NEE_modelling/')\n",
    "from _collect_prediction_data import allNaN_arg\n",
    "\n",
    "sys.path.append('/g/data/os22/chad_tmp/dea-notebooks/Tools/')\n",
    "from dea_tools.spatial import xr_rasterize\n",
    "\n",
    "import warnings\n",
    "warnings.filterwarnings(\"ignore\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from datacube.utils.dask import start_local_dask\n",
    "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 = 'NEE'\n",
    "suffix='20230320'\n",
    "# results_name = var+'_2003_2022_1km_quantiles_'+suffix+'.nc'\n",
    "# results_name = var+'_2003_2022_1km_quantiles.nc'\n",
    "results_name ='AusEFlux_'+var+'_2003_2022_1km_quantiles_v1.1.nc'\n",
    "# mask_path = '/g/data/os22/chad_tmp/NEE_modelling/results/prediction_data/mask_5km.nc'\n",
    "\n",
    "# upper = var+'_2003_2021_5km_LGBM_quantiles_'+suffix+'.nc'"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Open predictions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ds = xr.open_dataset('/g/data/os22/chad_tmp/NEE_modelling/results/predictions/'+results_name,\n",
    "                      chunks = dict(x=1100,y=1100, time=1)\n",
    "                      )\n",
    "ds = assign_crs(ds, crs='EPSG:4326')\n",
    "ds"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Convert to PgC/year"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "grid = ds.odc.geobox.to_crs('EPSG:3577')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ds = ds.odc.reproject(how=grid, resampling='bilinear')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "ds = ds.compute()\n",
    "# ds = ds.sel(quantile=0.5).compute()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "area_per_pixel = ds.geobox.resolution[1]**2\n",
    "\n",
    "if var != 'ET':\n",
    "    ds = ds * area_per_pixel * 1e-15 * 12 # (pgC/year)\n",
    "else:\n",
    "    ds = ds * 12 # (mm/year)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Fire emissions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fire = xr.open_dataarray('/g/data/os22/chad_tmp/NEE_modelling/data/FireEmissions_10km_monthly_2003_2022.nc')\n",
    "fire = fire.sum(['x', 'y'])\n",
    "fire = fire*12"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "\n",
    "## Interannual variability and annual mean"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "annual_ds = ds.sel(time=slice('2003', '2021')).resample(time='1Y').sum()\n",
    "annual_fire = fire.sel(time=slice('2003', '2021')).resample(time='1Y').sum()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "upper_percentile = (annual_ds[var+'_75th_percentile'].sum(['x','y']) / 12).mean().values\n",
    "lower_percentile = (annual_ds[var+'_25th_percentile'].sum(['x','y']) / 12).mean().values\n",
    "print('This Study Annual Mean (median model): ',(annual_ds[var+'_median'].sum(['x','y']) / 12).mean().values)\n",
    "print('This Study Annual Mean (median model)+FIRE: ', ((annual_ds[var+'_median'].sum(['x','y'])+annual_fire)/12).mean().values)\n",
    "print('This Study annual mean 75th percentile: ',upper_percentile)\n",
    "print('This Study annual mean 25th percentile: ',lower_percentile)\n",
    "print('IQR = ',upper_percentile-lower_percentile)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print('This Study IAV (median model): ',(annual_ds[var+'_median'].sum(['x','y']) / 12).std().values)\n",
    "print('This Study IAV (median model) + FIRE: ',((annual_ds[var+'_median'].sum(['x','y'])+annual_fire)/12).std().values)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ds_clim_mean = ds[var+'_median'].groupby('time.month').mean()\n",
    "print('Seasonal range of '+var, ds_clim_mean.sum(['x','y']).max().values - ds_clim_mean.sum(['x','y']).min().values)\n",
    "\n",
    "if var=='NEE':\n",
    "    ds_1D = ds[var+'_median'].sum(['x','y'])\n",
    "    ds_fire_1d = ds_1D+fire\n",
    "    ds_clim_mean_fire = ds_fire_1d.groupby('time.month').mean()\n",
    "    print('Seasonal range with fire of '+var, ds_clim_mean_fire.max().values - ds_clim_mean_fire.min().values)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Anomalies"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# ds_clim_mean = ds.groupby('time.month').mean()\n",
    "# ds_anom = (ds.groupby('time.month') - ds_clim_mean)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# ds_1d = ds.sum(['x', 'y'])\n",
    "\n",
    "# ds_fire_1d = ds_1d+fire\n",
    "# ds_fire_clim_mean = ds_fire_1d.groupby('time.month').mean()\n",
    "# ds_fire_anom = (ds_fire_1d.groupby('time.month') - ds_fire_clim_mean)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# ds_fire_anom.sel(quantile=0.5).rolling(time=12, min_periods=12).mean().plot(figsize=(11,5))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Zonal timeseries\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# ds_1D = ds.sum(['x','y'])\n",
    "\n",
    "ds_1D = ds[var+'_median'].rename(var).sum(['x','y'])\n",
    "upper_1D = ds[var+'_75th_percentile'].sum(['x','y'])\n",
    "lower_1D = ds[var+'_25th_percentile'].sum(['x','y'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ds_fire_1d.drop(['spatial_ref']).to_dataframe(name='NEE+fire')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "df_1d = ds_1D.drop(['spatial_ref']).to_dataframe()\n",
    "if var == 'NEE':\n",
    "    df_1d['NEE+fire'] = ds_fire_1d.drop(['spatial_ref']).to_dataframe(name='NEE+fire')\n",
    "\n",
    "fig,ax = plt.subplots(1,2, figsize=(15,5), sharey=True, gridspec_kw={'width_ratios': [4, 1]})\n",
    "ds_1D.plot(linestyle='-', ax=ax[0], label=var)\n",
    "ds_1D.rolling(time=12, min_periods=12).mean().plot(ax=ax[0], label=var+' 12-month rolling mean')\n",
    "ax[0].fill_between(ds_1D.time, lower_1D, upper_1D, alpha=0.2)\n",
    "df_1d.boxplot(ax=ax[1], rot=0, grid=False, fontsize=14, showmeans=True, color={'medians': 'blue'}, meanprops={'markerfacecolor': 'blue', 'markeredgecolor':'blue'})\n",
    "\n",
    "if var == 'NEE':\n",
    "    ds_fire_1d.plot(linestyle='--', ax=ax[0], label='NEE+Fire', c='purple', alpha=0.5)\n",
    "    ds_fire_1d.rolling(time=12, min_periods=12).mean().plot(ax=ax[0], label='NEE+fire 12-month rolling mean')\n",
    "    ax[0].axhline(0, c='grey', linestyle='--')\n",
    "    ax[1].axhline(0, c='grey', linestyle='--')\n",
    "\n",
    "# if var!='NEE':\n",
    "    # ax[0].set_ylim(1.0, 8.0)\n",
    "    \n",
    "if var != 'ET':\n",
    "    ax[0].set_ylabel(var+' PgC/year', fontsize=14)\n",
    "else:\n",
    "    ax[0].set_ylabel(var+' mm/year', fontsize=14)\n",
    "ax[0].xaxis.set_major_locator(mdates.YearLocator(2))\n",
    "if var =='ER':\n",
    "    ax[0].set_ylim(1.0, 8.0)\n",
    "ax[0].tick_params(axis='x', labelsize=14)\n",
    "ax[0].tick_params(axis='y', labelsize=14)\n",
    "ax[0].set_yticklabels(['{:3.1f}'.format(x) for x in ax[0].get_yticks()])\n",
    "ax[0].set_xlabel('')\n",
    "ax[0].set_title('')\n",
    "ax[0].grid('off', which='major', axis='both', linestyle='--', linewidth=0.75)\n",
    "ax[1].grid('off', which='major', axis='y', linestyle='--', linewidth=0.75)\n",
    "ax[0].legend(loc='best', fontsize=14)\n",
    "plt.tight_layout()\n",
    "plt.savefig('/g/data/os22/chad_tmp/NEE_modelling/results/figs/'+var+'_aus_timeseries_with_boxplot_'+suffix+'.png',\n",
    "           bbox_inches='tight', dpi=300)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Plots per bioregion"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "gdf = gpd.read_file('/g/data/os22/chad_tmp/NEE_modelling/data/bioclimatic_regions.geojson')\n",
    "gdf.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Dictionary to save results \n",
    "results = {}\n",
    "for index, row in gdf.iterrows():\n",
    "    print(row['region_name'])\n",
    "\n",
    "    # Generate a polygon mask to keep only data within the polygon\n",
    "    mask = xr_rasterize(gdf.iloc[[index]], ds)\n",
    "    \n",
    "    # Mask dataset to set pixels outside the polygon to `NaN`\n",
    "    dss = ds.where(mask)\n",
    "    \n",
    "    # Append results to a dictionary using the attribute\n",
    "    # column as an key\n",
    "    results.update({row['region_name']: dss})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, axs = plt.subplots(2,3, figsize=(15,10), sharex=True)\n",
    "\n",
    "for ax, k in zip(axs.ravel(), results.keys()):\n",
    "    \n",
    "    _1D = results[k].sel(quantile=0.5).sum(['x','y'])\n",
    "    \n",
    "    _1D.plot(linestyle='--', ax=ax, label=var, )\n",
    "    ax.text(.015, .95, f'Mean: {_1D.mean().values:.3f}, std: {_1D.std().values:.3f} PgC/yr',\n",
    "        transform=ax.transAxes, fontdict={'fontsize': 14})\n",
    "    if var=='NEE':\n",
    "        ax.axhline(0, c='grey', linestyle='--')\n",
    "    ax.set_title(k, fontdict={'fontsize': 14})\n",
    "    ax.set_ylabel('')\n",
    "    ax.set_xlabel('')\n",
    "\n",
    "fig.supylabel(var+' PgC/year', fontsize=14)\n",
    "plt.tight_layout()\n",
    "# plt.savefig('/g/data/os22/chad_tmp/NEE_modelling/results/figs/'+var+'_bioregion'+suffix+'.png')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Per-pixel summaries"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "annual_mean = ds.sel(time=slice('2003', '2021'))[var+'_median'].resample(time='1Y').sum().mean('time')\n",
    "annual_std = ds.sel(time=slice('2003', '2021'))[var+'_median'].resample(time='1Y').sum().std('time')\n",
    "\n",
    "scaled_variance = annual_std / annual_mean"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "monthly_mean = ds[var+'_median'].groupby('time.month').mean() #.sel(time=slice('2003','2015'))\n",
    "\n",
    "if var == 'NEE':\n",
    "    max_mean_month = allNaN_arg(monthly_mean.where(monthly_mean!=0), 'month', 'min')\n",
    "    # max_mean_season = allNaN_arg(monthly_season, 'season', 'min')\n",
    "else:\n",
    "    max_mean_month = allNaN_arg(monthly_mean.where(monthly_mean!=0), 'month', 'max')\n",
    "    # max_mean_season = allNaN_arg(monthly_season, 'season', 'max')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig,ax = plt.subplots(1,3, figsize=(18,7), sharey=True)\n",
    "font=20\n",
    "vmax = 0.000025\n",
    "vmin = 0\n",
    "vmax_sd_nee=0.0000015\n",
    "vmin_sd_nee=0\n",
    "vmax_sd=0.00000275\n",
    "vmin_sd=0.0000\n",
    "\n",
    "if var == 'NEE':\n",
    "    im = annual_mean.rename(var+' PgC/year').where(annual_mean!=0).plot.imshow(ax=ax[0], cmap='Spectral_r', vmax=3e-6, vmin=-3e-6,\n",
    "                                                                               add_colorbar=False, add_labels=False)\n",
    "    cbar = plt.colorbar(im, location='bottom', pad=0.025, ax=ax[0])\n",
    "    cbar.set_label(label=var+' PgC/year', size=15, weight='bold')\n",
    "    cbar.ax.tick_params(labelsize=15)\n",
    "    cbar.formatter.set_powerlimits((0, 0))\n",
    "else:\n",
    "    im = annual_mean.rename(var+' PgC/year').where(annual_mean!=0).plot.imshow(ax=ax[0], cmap='viridis', vmax=vmax, vmin=vmin,\n",
    "                                                                               add_colorbar=False, add_labels=False)\n",
    "    cbar = plt.colorbar(im, location='bottom', pad=0.025, ax=ax[0])\n",
    "    cbar.set_label(label=var+' PgC/year', size=15, weight='bold')\n",
    "    cbar.ax.tick_params(labelsize=15)\n",
    "\n",
    "if var=='NEE':\n",
    "    im1 = annual_std.rename('').where(annual_mean!=0).plot.imshow(ax=ax[1], cmap='magma',vmax=vmax_sd_nee,vmin=vmin_sd_nee,\n",
    "                                                                  add_colorbar=False, add_labels=False)\n",
    "else:\n",
    "    im1 = annual_std.rename('').where(annual_mean!=0).plot.imshow(ax=ax[1], cmap='magma',vmax=vmax_sd,vmin=vmin_sd,\n",
    "                                                                  add_colorbar=False, add_labels=False)\n",
    "\n",
    "cbar1 = plt.colorbar(im1, location='bottom', pad=0.025, ax=ax[1])\n",
    "cbar1.set_label(label=var+' Std. Dev. PgC/yr', size=15, weight='bold')\n",
    "cbar1.ax.tick_params(labelsize=15)\n",
    "\n",
    "im2 = max_mean_month.plot.imshow(vmin=1, vmax=12, ax=ax[2], cmap='twilight', add_colorbar=False, add_labels=False)\n",
    "cbar2 = plt.colorbar(im2, spacing='uniform', ax=ax[2], location='bottom', pad=0.025)\n",
    "cbar2.set_ticks([1,2,3,4,5,6,7,8,9,10,11,12])\n",
    "cbar2.set_ticklabels([\"J\",\"F\",\"M\",\"A\",\"M\",\"J\",\"J\",\"A\",\"S\",\"O\",\"N\",\"D\"], fontsize=15)\n",
    "cbar2.set_label(label=var+' Month of Max.', size=15, weight='bold')\n",
    "\n",
    "ax[0].set_yticklabels([])\n",
    "ax[0].set_ylabel('')\n",
    "ax[0].set_xlabel('')\n",
    "ax[0].set_xticklabels([])\n",
    "ax[1].set_yticklabels([])\n",
    "ax[1].set_xticklabels([])\n",
    "ax[2].set_yticklabels([])\n",
    "ax[2].set_xticklabels([])\n",
    "\n",
    "ax[0].set_title('Annual Mean of '+var,fontsize=font)\n",
    "ax[1].set_title('Std. Dev. in Annual Mean '+var, fontsize=font)\n",
    "\n",
    "if var=='NEE':\n",
    "    ax[2].set_title('Month of Maximum Carbon Uptake', fontsize=font);\n",
    "else:\n",
    "    ax[2].set_title('Month of Maximum '+var, fontsize=font);\n",
    "    \n",
    "plt.tight_layout();\n",
    "plt.savefig('/g/data/os22/chad_tmp/NEE_modelling/results/figs/'+var+'_perpixel_summaries_'+suffix+'.png',\n",
    "            bbox_inches='tight', dpi=300)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "y = monthly_mean.sum(['x','y']).values\n",
    "x = monthly_mean.month.values\n",
    "# std = monthly_std.mean(['x','y']).values\n",
    "\n",
    "fig,ax = plt.subplots(1,1, figsize=(7,4))\n",
    "ax.plot(x, y, marker='o', linestyle='--')\n",
    "if var=='NEE':\n",
    "    ax.axhline(0, c='grey')\n",
    "# ax.fill_between(x, y-std, y+std, alpha=0.5)\n",
    "ax.set_title('Monthly Mean '+var+' averaged over Australia')\n",
    "ax.set_xticks(range(1,13))\n",
    "ax.set_xticklabels([\"J\",\"F\",\"M\",\"A\",\"M\",\"J\",\"J\",\"A\",\"S\",\"O\",\"N\",\"D\"]) \n",
    "ax.set_xlabel('Month');\n",
    "ax.set_ylabel('PgC/year');\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Per bioregion climatology"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Dictionary to save results \n",
    "results = {}\n",
    "for index, row in gdf.iterrows():\n",
    "    print(row['region_name'])\n",
    "\n",
    "    # Generate a polygon mask to keep only data within the polygon\n",
    "    mask = xr_rasterize(gdf.iloc[[index]], ds)\n",
    "    \n",
    "    # Mask dataset to set pixels outside the polygon to `NaN`\n",
    "    dss = monthly_mean.where(mask)\n",
    "    \n",
    "    # Append results to a dictionary using the attribute\n",
    "    # column as an key\n",
    "    results.update({row['region_name']: dss})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, axs = plt.subplots(3,2, figsize=(12,15), sharex=True)\n",
    "\n",
    "for ax, k in zip(axs.ravel(), results.keys()):\n",
    "    \n",
    "    _1D = results[k].sum(['x','y'])\n",
    "    \n",
    "    _1D.plot(linestyle='--', ax=ax, label=var, )\n",
    "    ax.text(.015, .05, f'Mean: {_1D.mean().values:.3f}, std: {_1D.std().values:.3f} PgC/yr',\n",
    "        transform=ax.transAxes, fontdict={'fontsize': 14})\n",
    "    if var=='NEE':\n",
    "        ax.axhline(0, c='grey', linestyle='--')\n",
    "    ax.set_title(k, fontdict={'fontsize': 14})\n",
    "    ax.set_ylabel('')\n",
    "    ax.set_xlabel('')\n",
    "    ax.set_xticks(range(1,13))\n",
    "    ax.set_xticklabels([\"J\",\"F\",\"M\",\"A\",\"M\",\"J\",\"J\",\"A\",\"S\",\"O\",\"N\",\"D\"]) \n",
    "fig.supylabel(var+' PgC/year', fontsize=14)\n",
    "plt.tight_layout()\n",
    "# plt.savefig('/g/data/os22/chad_tmp/NEE_modelling/results/figs/'+var+'_bioregion_climatology_'+suffix+'.png')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Interactive map"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# annual_ds = ds.sel(quantile=0.5).resample(time='1Y').sum()\n",
    "# annual_mean = annual_ds.mean('time')\n",
    "# annual_std = annual_ds.std('time')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create folium Map (ipyleaflet is also supported)\n",
    "m = folium.Map(control_scale = True)\n",
    "\n",
    "# Plot each sample image with different colormap\n",
    "# max_mean_month.plot.imshow(vmin=1, vmax=12, ax=ax[0], cmap='twilight', add_colorbar=False)\n",
    "# np.absolute(scaled_variance).where(~mask).odc.add_to(m, opacity=1.0, cmap='magma', vmax=2)\n",
    "# max_mean_month.odc.add_to(m, opacity=1.0, cmap='twilight', vmin=1, vmax=12)\n",
    "# annual_std.where(annual_std>0).odc.add_to(m,  cmap='magma', robust=True, opacity=1.0) #opacity=1.0,\n",
    "annual_mean.where(annual_mean!=0).odc.add_to(m, cmap='RdBu_r', vmin=-300, vmax=300)\n",
    "\n",
    "\n",
    "# Zoom map to Australia\n",
    "m.fit_bounds(annual_mean.odc.map_bounds())\n",
    "\n",
    "tile = folium.TileLayer(\n",
    "        tiles = 'https://server.arcgisonline.com/ArcGIS/rest/services/World_Imagery/MapServer/tile/{z}/{y}/{x}',\n",
    "        attr = 'Esri',\n",
    "        name = 'Esri Satellite',\n",
    "        overlay = True,\n",
    "        control = True\n",
    "       ).add_to(m)\n",
    "\n",
    "folium.LayerControl().add_to(m)\n",
    "display(m)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Zoomed in maps"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "annual_ds = ds['NEE_median'].resample(time='1Y').sum().compute()\n",
    "annual_mean = annual_ds.mean('time').compute()\n",
    "annual_std = annual_ds.std('time').compute()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# #queensland\n",
    "x_slice=slice(2700,4200)\n",
    "y_slice=slice(850,1800)\n",
    "name='queensland'\n",
    "\n",
    "#SW WA\n",
    "# x_slice=slice(290,900)\n",
    "# y_slice=slice(2150,2510)\n",
    "# name='westernaus'\n",
    "\n",
    "#Cape York\n",
    "# x_slice=slice(2600,3600)\n",
    "# y_slice=slice(20,850)\n",
    "# name='capeyork'\n",
    "\n",
    "# # melbourne\n",
    "# x_slice=slice(2600,3800)\n",
    "# y_slice=slice(2400,2900)\n",
    "# name='melbourne'\n",
    "\n",
    "# NW WA\n",
    "# x_slice=slice(900,2000)\n",
    "# y_slice=slice(300,1000)\n",
    "# name='NW_WA'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig,ax = plt.subplots(1,3, figsize=(21,7), sharey=True) # figsize=(21,6),\n",
    "font=20\n",
    "\n",
    "if var == 'NEE':\n",
    "    im = annual_mean.isel(x=x_slice,\n",
    "                          y=y_slice).rename(var+' gC m\\N{SUPERSCRIPT TWO} y⁻¹').where(annual_mean!=0).plot.imshow(ax=ax[0],\n",
    "                                                                                                                 cmap='RdBu_r',\n",
    "                                                                                                                 # robust=True,\n",
    "                                                                                                                  vmin=-250, vmax=250,\n",
    "                                                                                                                 add_colorbar=False,\n",
    "                                                                                                                 add_labels=False)\n",
    "    cbar = plt.colorbar(im, location='bottom', pad=0.05, ax=ax[0])\n",
    "    cbar.set_label(label=var+' gC m\\N{SUPERSCRIPT TWO} year⁻¹', size=15, weight='bold')\n",
    "    cbar.ax.tick_params(labelsize=15)\n",
    "    # cbar.formatter.set_powerlimits((0, 0))\n",
    "else:\n",
    "    im = annual_mean.isel(x=x_slice, y=y_slice).rename(var+' gC m\\N{SUPERSCRIPT TWO} y⁻¹').where(annual_mean!=0).plot.imshow(ax=ax[0], cmap='viridis', robust=True, add_colorbar=False, add_labels=False)\n",
    "    cbar = plt.colorbar(im, location='bottom', pad=0.05, ax=ax[0])\n",
    "    cbar.set_label(label=var+' gC m\\N{SUPERSCRIPT TWO} y⁻¹', size=15, weight='bold')\n",
    "    cbar.ax.tick_params(labelsize=15)\n",
    "\n",
    "im2 = annual_mean.isel(x=x_slice, y=y_slice).where(annual_mean!=0).plot.imshow(ax=ax[2], cmap='viridis', robust=True, add_colorbar=False, add_labels=False)\n",
    "#ctx.add_basemap(ax[2], source=ctx.providers.Esri.WorldImagery, crs='EPSG:4326', attribution='Esri World Imagery', attribution_size=15)\n",
    "\n",
    "ax[2].set_aspect('auto')\n",
    "if var=='NEE':\n",
    "    im1 = annual_std.isel(x=x_slice, y=y_slice).rename('').where(annual_mean!=0).plot.imshow(ax=ax[1],\n",
    "                                                                                             cmap='magma',\n",
    "                                                                                             robust=True,\n",
    "                                                                                             # vmin=0, vmax=160,\n",
    "                                                                                             add_colorbar=False,\n",
    "                                                                                             add_labels=False)\n",
    "else:\n",
    "    im1 = annual_std.isel(x=x_slice, y=y_slice).rename('').where(annual_mean!=0).plot.imshow(ax=ax[1], cmap='magma',robust=True, add_colorbar=False, add_labels=False)\n",
    "\n",
    "cbar1 = plt.colorbar(im1, location='bottom', pad=0.05, ax=ax[1])\n",
    "cbar1.set_label(label=var+' Std. Dev.', size=15, weight='bold')\n",
    "cbar1.ax.tick_params(labelsize=15)\n",
    "\n",
    "cbar2 = plt.colorbar(im2, spacing='uniform', ax=ax[2], location='bottom', pad=0.05)\n",
    "\n",
    "\n",
    "ax[0].set_ylabel('')\n",
    "ax[0].set_xlabel('')\n",
    "# ax[0].set_yticklabels([])\n",
    "# ax[0].set_xticklabels([])\n",
    "# ax[1].set_yticklabels([])\n",
    "# ax[1].set_xticklabels([])\n",
    "# ax[2].set_yticklabels([])\n",
    "# ax[2].set_xticklabels([])\n",
    "\n",
    "ax[0].set_title('Annual Mean of '+var,fontsize=font)\n",
    "ax[1].set_title('Std.Dev. in Annual '+var, fontsize=font)\n",
    "\n",
    "ax[2].set_title('True Colour Image', fontsize=font);\n",
    "    \n",
    "plt.tight_layout()\n",
    "plt.savefig('/g/data/os22/chad_tmp/NEE_modelling/results/figs/'+var+'_zoomedin_perpixel_'+name+'.png', \n",
    "            bbox_inches='tight', dpi=300)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Summarizing all climatology fluxes"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import sys\n",
    "import numpy as np\n",
    "import xarray as xr\n",
    "import geopandas as gpd\n",
    "import matplotlib.pyplot as plt\n",
    "from odc.geo.xr import assign_crs\n",
    "\n",
    "sys.path.append('/g/data/os22/chad_tmp/dea-notebooks/Tools/')\n",
    "from dea_tools.spatial import xr_rasterize\n",
    "\n",
    "import warnings\n",
    "warnings.filterwarnings(\"ignore\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from datacube.utils.dask import start_local_dask\n",
    "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": [
    "results_name_er ='AusEFlux_ER_2003_2022_1km_quantiles_v1.1.nc'\n",
    "results_name_gpp ='AusEFlux_GPP_2003_2022_1km_quantiles_v1.1.nc'"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Open datasets"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "er = xr.open_dataset('/g/data/os22/chad_tmp/NEE_modelling/results/predictions/'+results_name_er,\n",
    "                      chunks=dict(latitude=1100, longitude=1100, time=1))['ER_median']\n",
    "er.attrs['nodata'] = np.nan\n",
    "er = assign_crs(er, crs='EPSG:4326')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "gpp = xr.open_dataset('/g/data/os22/chad_tmp/NEE_modelling/results/predictions/'+results_name_gpp,\n",
    "                      chunks=dict(latitude=1100, longitude=1100, time=1))['GPP_median']\n",
    "gpp.attrs['nodata'] = np.nan\n",
    "gpp = assign_crs(gpp, crs='EPSG:4326')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fire = xr.open_dataarray('/g/data/os22/chad_tmp/NEE_modelling/data/FireEmissions_10km_monthly_2003_2022.nc')\n",
    "# fire = fire.sum(['x', 'y'])\n",
    "fire = fire*12"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "er = er.odc.reproject(how=er.odc.geobox.to_crs('EPSG:3577'),\n",
    "                      resampling='bilinear').compute()\n",
    "\n",
    "gpp = gpp.odc.reproject(how=gpp.odc.geobox.to_crs('EPSG:3577'),\n",
    "                      resampling='bilinear').compute()\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Convert to total C in (petagrams) rather than gC/m2\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "er = er * er.geobox.resolution[1]**2 * 1e-15 *12\n",
    "gpp = gpp * gpp.geobox.resolution[1]**2 * 1e-15 *12\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Climatology"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import warnings\n",
    "warnings.filterwarnings(\"ignore\")\n",
    "\n",
    "er_clim_mean = er.groupby('time.month').mean()\n",
    "gpp_clim_mean = gpp.groupby('time.month').mean()\n",
    "fire_clim_mean = fire.groupby('time.month').mean()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Per bioregion"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "bio = assign_crs(xr.open_dataarray('/g/data/os22/chad_tmp/NEE_modelling/data/bioclimatic_regions.nc'), crs='EPSG:4326')\n",
    "gdf = gpd.read_file('/g/data/os22/chad_tmp/NEE_modelling/data/bioclimatic_regions.geojson')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def make_mask(xx, index):\n",
    "    mask = bio.odc.reproject(how=xx.odc.geobox, resampling='nearest')\n",
    "    mask = mask.astype(int)\n",
    "    mask['x'] = np.array([round(i,6) for i in mask.x.values])\n",
    "    mask['y'] = np.array([round(i,6) for i in mask.y.values])\n",
    "    mask = xr.where(mask==index, 1, 0)\n",
    "    return mask\n",
    "\n",
    "def round_coords(ds):\n",
    "    ds['x'] = np.array([round(i,6) for i in ds.x.values])\n",
    "    ds['y'] = np.array([round(i,6) for i in ds.y.values])\n",
    "    return ds"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Dictionary to save results \n",
    "er_results = {}\n",
    "for index, row in gdf.iterrows():\n",
    "    index=index+1\n",
    "    print(row['region_name'])\n",
    "    datasets = {}\n",
    "    \n",
    "    ds_mask = make_mask(er, index)\n",
    "    \n",
    "    dss = round_coords(er_clim_mean).where(ds_mask)\n",
    "        \n",
    "    datasets.update({\n",
    "                'This Study ER':dss,\n",
    "                }\n",
    "               )\n",
    "    # Append results to a dictionary using the attribute\n",
    "    # column as an key\n",
    "    er_results.update({row['region_name']: datasets})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Dictionary to save results \n",
    "gpp_results = {}\n",
    "for index, row in gdf.iterrows():\n",
    "    index=index+1\n",
    "    print(row['region_name'])\n",
    "    datasets = {}\n",
    "    \n",
    "    ds_mask = make_mask(gpp, index)\n",
    "    \n",
    "    dss = round_coords(gpp_clim_mean).where(ds_mask)\n",
    "        \n",
    "    datasets.update({\n",
    "                'This Study GPP':dss,\n",
    "                }\n",
    "               )\n",
    "    # Append results to a dictionary using the attribute\n",
    "    # column as an key\n",
    "    gpp_results.update({row['region_name']: datasets})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Dictionary to save results \n",
    "fire_results = {}\n",
    "for index, row in gdf.iterrows():\n",
    "    index=index+1\n",
    "    print(row['region_name'])\n",
    "    datasets = {}\n",
    "    \n",
    "    ds_mask = make_mask(fire, index)\n",
    "    \n",
    "    dss = round_coords(fire_clim_mean).where(ds_mask)\n",
    "        \n",
    "    datasets.update({\n",
    "                'GFAS Fire':dss,\n",
    "                }\n",
    "               )\n",
    "    # Append results to a dictionary using the attribute\n",
    "    # column as an key\n",
    "    fire_results.update({row['region_name']: datasets})"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Plot"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "regions = ['Tropics',\n",
    "'Savanna',\n",
    "'Warm Temperate',\n",
    "'Cool Temperate',\n",
    "'Mediterranean',\n",
    "'Desert']\n",
    "\n",
    "fig, ax = plt.subplots(1,2, figsize=(16,7))\n",
    "for i in regions:\n",
    "    \n",
    "    subtract =  er_results[i]['This Study ER'].sum(['x','y']) - gpp_results[i]['This Study GPP'].sum(['x','y'])\n",
    "    subtract.plot(label=i, ax=ax[1])\n",
    "    \n",
    "    if i=='Savanna':\n",
    "        fire_1d = fire_results[i]['GFAS Fire'].sum(['x','y'])\n",
    "        (subtract+fire_1d).plot(label=i+'+Fire', ax=ax[1], linestyle='--', color='orange')\n",
    "    # if i=='Desert':\n",
    "    #     fire_1d = fire_results[i]['GFAS Fire'].sum(['x','y'])\n",
    "    #     (subtract+fire_1d).plot(label=i+'+Fire', ax=ax[1], linestyle='--', color='brown')\n",
    "        \n",
    "    ax[1].legend(fontsize=13)\n",
    "    ax[1].set_title('NEE (ER-GPP) Climatologies', fontsize=20)\n",
    "    ax[1].set_ylabel('')\n",
    "    ax[1].tick_params(axis='y', labelsize=17)\n",
    "    ax[1].set_xlabel('')\n",
    "    ax[1].set_xticks(range(1,13))\n",
    "    ax[1].set_xticklabels([\"J\",\"F\",\"M\",\"A\",\"M\",\"J\",\"J\",\"A\",\"S\",\"O\",\"N\",\"D\"], fontsize=17)\n",
    "    ax[1].axhline(0, c='grey', linestyle='--')\n",
    "    \n",
    "er_results['Savanna']['This Study ER'].sum(['x','y']).plot(ax=ax[0], label='ER-Savanna', c='coral')\n",
    "gpp_results['Savanna']['This Study GPP'].sum(['x','y']).plot(ax=ax[0], label='GPP-Savanna', c='coral', linestyle='--')\n",
    "\n",
    "er_results['Desert']['This Study ER'].sum(['x','y']).plot(ax=ax[0], label='ER-Desert',c='royalblue' )\n",
    "gpp_results['Desert']['This Study GPP'].sum(['x','y']).plot(ax=ax[0], label='GPP-Desert',  c='royalblue', linestyle='--')\n",
    "\n",
    "ax[0].legend(fontsize=13)\n",
    "ax[0].set_title('Desert & Savanna Flux Climatologies', fontsize=20)\n",
    "ax[0].set_ylabel('PgC/yr', fontsize=20)\n",
    "ax[0].tick_params(axis='y', labelsize=17)\n",
    "ax[0].set_xlabel('')\n",
    "ax[0].set_xticks(range(1,13))\n",
    "ax[0].set_xticklabels([\"J\",\"F\",\"M\",\"A\",\"M\",\"J\",\"J\",\"A\",\"S\",\"O\",\"N\",\"D\"], fontsize=17)\n",
    "    \n",
    "plt.tight_layout()\n",
    "\n",
    "plt.savefig('/g/data/os22/chad_tmp/NEE_modelling/results/figs/desert_savanna_climatology.png',\n",
    "             bbox_inches='tight', dpi=300);     "
   ]
  },
  {
   "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.10.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}