7_Compare_products.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Comparing fluxes with other products\n",
    "\n",
    "There are nine other products in total, including: LSMs, FLUXCOM, obs-GPP, and atmos. invs. \n",
    "\n",
    "This notebook produces various intercomparison plots, both gridded plots and zonal time-series plots"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import xarray as xr\n",
    "import rioxarray as rxr\n",
    "import geopandas as gpd\n",
    "import pandas as pd\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "import matplotlib as mpl\n",
    "from matplotlib.colors import ListedColormap\n",
    "from matplotlib.ticker import FormatStrFormatter\n",
    "\n",
    "from odc.geo.geobox import zoom_out\n",
    "from odc.geo.xr import assign_crs\n",
    "from odc.algo import xr_reproject\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' #ER #NEE #GPP\n",
    "var_fc = 'NEE' #TER #NEE #GPP\n",
    "var_cable = 'nbp' #nbp #er #gpp\n",
    "suffix='20230320'\n",
    "trim=True\n",
    "trim_date = '2003','2022'\n",
    "\n",
    "# results_name = var+'_2003_2022_1km_quantiles_'+suffix+'.nc'\n",
    "results_name ='AusEFlux_'+var+'_2003_2022_1km_quantiles_v1.1.nc'\n",
    "# upper = 'NEE_2003_2021_5km_LGBM_quantiles.nc'"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Open various datasets"
   ]
  },
  {
   "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(latitude=1100, longitude=1100, time=1))\n",
    "ds.attrs['nodata'] = np.nan\n",
    "ds = assign_crs(ds, crs='EPSG:4326')\n",
    "ds"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Fluxcom"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ds_fc = assign_crs(xr.open_dataset('/g/data/os22/chad_tmp/NEE_modelling/data/FLUXCOM/'+var_fc+'_rs.nc'), crs='epsg:4326').sel(time=slice('2003', '2015'))\n",
    "\n",
    "ds_fc.attrs['nodata'] = np.nan\n",
    "ds_fc['latitude'] = ds_fc.latitude.astype('float32')\n",
    "ds_fc['longitude'] = ds_fc.longitude.astype('float32')\n",
    "\n",
    "ds_fc_met = assign_crs(xr.open_dataset('/g/data/os22/chad_tmp/NEE_modelling/data/FLUXCOM/'+var_fc+'_rs_meteo_era5.nc').sel(time=slice('2003', '2018')), crs='epsg:4326')\n",
    "ds_fc_met = ds_fc_met.rename({'lat':'latitude', 'lon':'longitude'})\n",
    "ds_fc_met.attrs['nodata'] = np.nan\n",
    "ds_fc_met['latitude'] = ds_fc_met.latitude.astype('float32')\n",
    "ds_fc_met['longitude'] = ds_fc_met.longitude.astype('float32')\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ds_fc['time'] = ds.sel(time=slice('2003', '2015')).time\n",
    "ds_fc_met['time'] = ds.sel(time=slice('2003', '2018')).time"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### CABLE\n",
    "\n",
    "<!-- '/g/data/os22/chad_tmp/NEE_modelling/data/CABLE/CABLE_ER_10km_monthly_2000_2019.nc' -->"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cable = xr.open_dataarray('/g/data/os22/chad_tmp/NEE_modelling/data/CABLE/CABLE-BIOS/CABLE_BIOS_'+var_cable+'_25km_monthly_2003_2019.nc')\n",
    "cable_pop = xr.open_dataset('/g/data/os22/chad_tmp/NEE_modelling/data/CABLE/CABLE-POP_v10/CABLE-POP_'+var_cable+'_100km_monthly_2003_2020.nc')\n",
    "cable_pop = cable_pop[var_cable]\n",
    "# cable_nee = cable_nee.sel(time=slice('2003', '2015'))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### GOSAT Inversion\n",
    "\n",
    "Metz etal 2023"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "gosat = pd.read_csv('/g/data/os22/chad_tmp/NEE_modelling/data/metz_2023/aus_transcom_gosat_inversion_PgCyear.csv',\n",
    "                   parse_dates=True).set_index('time')\n",
    "gosat = gosat.to_xarray()\n",
    "gosat=gosat['GOSAT_INVERSION']\n",
    "gosat['time'] = ds.sel(time=slice('2009-04', '2018-12')).time"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Fire emissions GFASv12"
   ]
  },
  {
   "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": [
    "### Obs GPP"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if var == 'GPP':\n",
    "    modis_gpp = xr.open_dataarray('/g/data/os22/chad_tmp/NEE_modelling/data/1km/MODIS_GPP_1km_monthly_2002_2021.nc')\n",
    "    modis_gpp.attrs['nodata'] = np.nan\n",
    "\n",
    "    gosif_gpp = xr.open_dataarray('/g/data/os22/chad_tmp/NEE_modelling/data/5km/GOSIF_GPP_5km_monthly_2002_2021.nc')\n",
    "    gosif_gpp.attrs['nodata'] = np.nan\n",
    "    \n",
    "    diffuse = xr.open_dataarray('/g/data/os22/chad_tmp/NEE_modelling/data/1km/DIFFUSE_GPP_1km_monthly_2003_2021.nc')\n",
    "    diffuse.attrs['nodata'] = np.nan"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Villalobos et al \n",
    "\n",
    "Needed to fix the dataset before it could be opened. Used the R script in data/villalobos/"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fn = '/g/data/os22/chad_tmp/NEE_modelling/data/villalobos/post.tif'\n",
    "post = rxr.open_rasterio(fn).rename({'band':'month'})\n",
    "post['month']=ds.sel(time=slice('2015', '2019')).time.values\n",
    "post = post.rename({'month':'time'})\n",
    "#post = xr_reproject(post, geobox=post.odc.geobox.to_crs('epsg:3577').compat, resampling='nearest')\n",
    "# post = assign_crs(post, crs='epsg:3577')\n",
    "post.attrs['nodata'] = np.nan\n",
    "post.attrs['units'] = 'PgC/year'\n",
    "\n",
    "fn = '/g/data/os22/chad_tmp/NEE_modelling/data/villalobos/land.tif'\n",
    "cbl = rxr.open_rasterio(fn).rename({'band':'month'})\n",
    "cbl['month']=ds.sel(time=slice('2015', '2019')).time.values\n",
    "cbl = cbl.rename({'month':'time'})\n",
    "#cable = xr_reproject(cable, geobox=cable.odc.geobox.to_crs('epsg:3577').compat, resampling='nearest')\n",
    "# post = assign_crs(post, crs='epsg:3577')\n",
    "cbl.attrs['nodata'] = np.nan\n",
    "cbl.attrs['units'] = 'PgC/year'\n",
    "\n",
    "fn = '/g/data/os22/chad_tmp/NEE_modelling/data/villalobos/fossil.tif'\n",
    "ff = rxr.open_rasterio(fn).rename({'band':'month'})\n",
    "ff['month']=ds.sel(time=slice('2015', '2019')).time.values\n",
    "ff = ff.rename({'month':'time'})\n",
    "#ff = xr_reproject(ff, geobox=ff.odc.geobox.to_crs('epsg:3577').compat, resampling='nearest')\n",
    "# post = assign_crs(post, crs='epsg:3577')\n",
    "ff.attrs['nodata'] = np.nan\n",
    "ff.attrs['units'] = 'PgC/year'\n",
    "\n",
    "#subtract fossile fuels from inversion\n",
    "post = post - ff\n",
    "\n",
    "# create a mask for this janky dataset\n",
    "mask_inv = cbl.where(cbl==0).isel(time=1)\n",
    "mask_inv = mask_inv.astype(bool)\n",
    "\n",
    "post = post.where(mask_inv)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Reproject to Albers Equal Area"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# grid = zoom_out(ds.odc.geobox, 20)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ds = ds.odc.reproject(how=ds.odc.geobox.to_crs('EPSG:3577'),\n",
    "                      resampling='bilinear').compute()\n",
    "\n",
    "ds_upper = ds[[var+'_25th_percentile', var+'_75th_percentile']]\n",
    "ds = ds[var+'_median']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ds_fc = ds_fc.odc.reproject(how=ds_fc.odc.geobox.to_crs('EPSG:3577'),\n",
    "                      resampling='bilinear').compute()\n",
    "ds_fc = ds_fc[var_fc] #.where(~mask.astype(bool))\n",
    "ds_fc = ds_fc*30 #/day to /month\n",
    "\n",
    "ds_fc_met = ds_fc_met.odc.reproject(how=ds_fc_met.odc.geobox.to_crs('EPSG:3577'),\n",
    "                      resampling='bilinear').compute()\n",
    "ds_fc_met = ds_fc_met[var_fc] #.where(~mask.astype(bool))\n",
    "ds_fc_met = ds_fc_met*30 #/day to /month"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cable = assign_crs(cable, crs='epsg:4326')\n",
    "cable = cable.odc.reproject(how=cable.odc.geobox.to_crs('EPSG:3577'),\n",
    "                      resampling='bilinear').compute()\n",
    "\n",
    "cable_pop = assign_crs(cable_pop, crs='epsg:4326')\n",
    "cable_pop = cable_pop.odc.reproject(how=cable_pop.odc.geobox.to_crs('EPSG:3577'),\n",
    "                      resampling='bilinear').compute()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if var == 'GPP':\n",
    "    # grid = zoom_out(ds.odc.geobox, 5)\n",
    "    # grid = grid.to_crs('EPSG:3577')\n",
    "    modis_gpp = modis_gpp.chunk(dict(latitude=1100,longitude=1100)).odc.reproject(how=ds.odc.geobox.to_crs('EPSG:3577'),\n",
    "                      resampling='average').compute()\n",
    "    \n",
    "    gosif_gpp = gosif_gpp.odc.reproject(how=gosif_gpp.odc.geobox.to_crs('EPSG:3577'),\n",
    "                      resampling='average').compute()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if var == 'GPP':\n",
    "    diffuse = diffuse.chunk(dict(latitude=1100,longitude=1100)).odc.reproject(how=diffuse.odc.geobox.to_crs('EPSG:3577'),\n",
    "                      resampling='bilinear').compute()\n",
    "    diffuse = diffuse*30 #/day to /month"
   ]
  },
  {
   "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": [
    "ds = ds * ds.geobox.resolution[1]**2 * 1e-15 *12\n",
    "ds_upper = ds_upper * ds_upper.geobox.resolution[1]**2 * 1e-15 *12\n",
    "\n",
    "ds_fc = ds_fc * ds_fc.geobox.resolution[1]**2 * 1e-15 * 12\n",
    "ds_fc_met = ds_fc_met * ds_fc_met.geobox.resolution[1]**2 * 1e-15 * 12\n",
    "\n",
    "cable = cable * cable.geobox.resolution[1]**2 * 1e-15 *12\n",
    "cable_pop = cable_pop * cable_pop.geobox.resolution[1]**2 * 1e-15 *12\n",
    "\n",
    "if var == 'GPP':\n",
    "    modis_gpp = modis_gpp * modis_gpp.geobox.resolution[1]**2 * 1e-12\n",
    "    modis_gpp = modis_gpp*12\n",
    "    gosif_gpp = gosif_gpp * gosif_gpp.geobox.resolution[1]**2 * 1e-15\n",
    "    gosif_gpp = gosif_gpp*12\n",
    "    diffuse = diffuse * diffuse.geobox.resolution[1]**2 * 1e-15\n",
    "    diffuse = diffuse*12"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Plot time-series"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if trim:\n",
    "    ds = ds.sel(time=slice(trim_date[0], trim_date[1]))\n",
    "    ds_upper = ds_upper.sel(time=slice(trim_date[0], trim_date[1]))\n",
    "    ds_fc = ds_fc.sel(time=slice(trim_date[0], trim_date[1]))\n",
    "    ds_fc_met = ds_fc_met.sel(time=slice(trim_date[0], trim_date[1]))\n",
    "    cable = cable.sel(time=slice(trim_date[0], trim_date[1]))\n",
    "    cable_pop = cable_pop.sel(time=slice(trim_date[0], trim_date[1]))\n",
    "    \n",
    "    if var=='GPP':\n",
    "        modis_gpp = modis_gpp.sel(time=slice(trim_date[0], trim_date[1]))\n",
    "        gosif_gpp = gosif_gpp.sel(time=slice(trim_date[0], trim_date[1]))\n",
    "        diffuse = diffuse.sel(time=slice(trim_date[0], trim_date[1]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ds_1d = ds.sum(['x', 'y'])\n",
    "ds_1d_upper = ds_upper.sum(['x', 'y'])\n",
    "\n",
    "ds_fc_1d = ds_fc.sum(['x', 'y'])\n",
    "ds_fc_met_1d = ds_fc_met.sum(['x', 'y'])\n",
    "\n",
    "post_1d = post.sum(['x', 'y'])\n",
    "cable_1d = cable.sum(['x', 'y'])\n",
    "cable_pop_1d = cable_pop.sum(['x', 'y'])\n",
    "\n",
    "if var=='GPP':\n",
    "    gosif_gpp_1d = gosif_gpp.sum(['x', 'y'])\n",
    "    modis_gpp_1d = modis_gpp.sum(['x', 'y'])\n",
    "    diffuse_1d = diffuse.sum(['x', 'y'])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### add fire to NEE estimates"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if var=='NEE':\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": "markdown",
   "metadata": {},
   "source": [
    "### Plot time-series w/o rolling mean"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots(1,1, figsize=(15,7))\n",
    "ax.plot(ds_1d.time.values, ds_1d,  label='This study')\n",
    "\n",
    "ax.fill_between(ds_1d.time, ds_1d_upper[var+'_25th_percentile'].rolling(time=3, min_periods=1).mean(),\n",
    "                ds_1d_upper[var+'_75th_percentile'].rolling(time=3, min_periods=1).mean(), alpha=0.2)\n",
    "\n",
    "ax.plot(ds_fc_1d.time.values, ds_fc_1d, linestyle='--', label='Fluxcom-RS', alpha=0.8)\n",
    "ax.plot(ds_fc_met_1d.time.values, ds_fc_met_1d, linestyle='--', label='Fluxcom-Met', alpha=0.8)\n",
    "ax.plot(cable_1d.time.values, cable_1d, linestyle='--', label='CABLE-BIOS', alpha=0.8)\n",
    "ax.plot(cable_pop_1d.time.values, cable_pop_1d, linestyle='--', label='CABLE-POP', alpha=0.8)\n",
    "\n",
    "if var=='NEE':\n",
    "    ax.plot(post_1d.time.values, post_1d, linestyle='--', label='Inversion', alpha=0.8)\n",
    "    ax.plot(ds_1d.time.values, ds_fire_1d,c='blue', linestyle='--', label='This study+Fire')\n",
    "    plt.axhline(0, c='grey')\n",
    "\n",
    "if var == 'GPP':\n",
    "    ax.plot(modis_gpp_1d.time.values, modis_gpp_1d, linestyle='--', label='MODIS', alpha=0.75)\n",
    "    ax.plot(gosif_gpp_1d.time.values, gosif_gpp_1d, linestyle='--', label='GOSIF', alpha=0.75)\n",
    "    ax.plot(diffuse_1d.time.values, diffuse_1d, linestyle='--', label='DIFFUSE', alpha=0.75)\n",
    "\n",
    "fontsize=16\n",
    "ax.set_ylabel(var+'PgC/year')\n",
    "\n",
    "if var=='NEE':\n",
    "    ax.text(.015, .30, f'Fluxcom-RS: {ds_fc_1d.mean().values:.2f} PgC/yr',\n",
    "            transform=ax.transAxes, fontsize=fontsize)\n",
    "    ax.text(.015, .25, f'Fluxcom-Met: {ds_fc_met_1d.mean().values:.2f} PgC/yr',\n",
    "            transform=ax.transAxes, fontsize=fontsize)\n",
    "    ax.text(.015, .20, f'This study: {ds_1d.mean().values:.2f} PgC/yr',\n",
    "            transform=ax.transAxes, fontsize=fontsize)\n",
    "    ax.text(.015, .15, f'Cable-BIOS: {cable_1d.mean().values:.2f} PgC/yr',\n",
    "            transform=ax.transAxes, fontsize=fontsize)\n",
    "    ax.text(.015, .10, f'Cable-POP: {cable_pop_1d.mean().values:.2f} PgC/yr',\n",
    "            transform=ax.transAxes,fontsize=fontsize)\n",
    "    ax.text(.015, .05, f'Inversion: {post_1d.mean().values:.2f} PgC/yr',\n",
    "            transform=ax.transAxes, fontsize=fontsize)\n",
    "    ax.legend(loc='lower right', fontsize=fontsize);\n",
    "\n",
    "if var == 'GPP':\n",
    "    ax.text(.015, .95, f'Fluxcom-RS: {ds_fc_1d.mean().values:.2f} PgC/yr',\n",
    "            transform=ax.transAxes, fontsize=fontsize)\n",
    "    ax.text(.015, .90, f'Fluxcom-Met: {ds_fc_met_1d.mean().values:.2f} PgC/yr',\n",
    "            transform=ax.transAxes, fontsize=fontsize)\n",
    "    ax.text(.015, .85, f'This study: {ds_1d.mean().values:.2f} PgC/yr',\n",
    "            transform=ax.transAxes, fontsize=fontsize)\n",
    "    ax.text(.015, .80, f'GOSIF: {gosif_gpp_1d.mean().values:.2f} PgC/yr',\n",
    "            transform=ax.transAxes, fontsize=fontsize)\n",
    "    ax.text(.015, .75, f'MODIS: {modis_gpp_1d.mean().values:.2f} PgC/yr',\n",
    "            transform=ax.transAxes, fontsize=fontsize)\n",
    "    ax.text(.015, .70, f'Cable-BIOS: {cable_1d.mean().values:.2f} PgC/yr',\n",
    "            transform=ax.transAxes, fontsize=fontsize)\n",
    "    ax.text(.015, .65, f'Cable-POP: {cable_pop_1d.mean().values:.2f} PgC/yr',\n",
    "            transform=ax.transAxes, fontsize=fontsize)\n",
    "    ax.text(.015, .60, f'DIFFUSE: {diffuse_1d.mean().values:.2f} PgC/yr',\n",
    "            transform=ax.transAxes, fontsize=fontsize)\n",
    "    ax.legend(loc='upper right', fontsize=fontsize)\n",
    "\n",
    "if var == 'ER':\n",
    "    ax.text(.015, .95, f'Fluxcom-RS: {ds_fc_1d.mean().values:.2f} PgC/yr',\n",
    "            transform=ax.transAxes, fontsize=fontsize)\n",
    "    ax.text(.015, .90, f'Fluxcom-Met: {ds_fc_met_1d.mean().values:.2f} PgC/yr',\n",
    "            transform=ax.transAxes, fontsize=fontsize)\n",
    "    ax.text(.015, .85, f'This study: {ds_1d.mean().values:.2f} PgC/yr',\n",
    "            transform=ax.transAxes, fontsize=fontsize)\n",
    "    ax.text(.015, .80, f'Cable-BIOS: {cable_1d.mean().values:.2f} PgC/yr',\n",
    "            transform=ax.transAxes, fontsize=fontsize)\n",
    "    ax.text(.015, .75, f'Cable-POP: {cable_pop_1d.mean().values:.2f} PgC/yr',\n",
    "            transform=ax.transAxes, fontsize=fontsize)\n",
    "    ax.legend(loc='upper right', fontsize=fontsize)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Plot timeseries with rolling mean"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "roll_len=12"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fontsize=20\n",
    "line=2\n",
    "\n",
    "fig, ax = plt.subplots(1,1, figsize=(15,7))\n",
    "ax.plot(ds_1d.time.values, ds_1d.rolling(time=roll_len, min_periods=roll_len).mean(),  label='This study', linewidth=line)\n",
    "ax.fill_between(ds_1d.time, ds_1d_upper[var+'_25th_percentile'].rolling(time=roll_len, min_periods=roll_len).mean(),\n",
    "                ds_1d_upper[var+'_75th_percentile'].rolling(time=roll_len, min_periods=roll_len).mean(), alpha=0.2)\n",
    "\n",
    "ax.plot(ds_fc_1d.time.values, ds_fc_1d.rolling(time=roll_len, min_periods=roll_len).mean(), label='Fluxcom-RS', linewidth=line)\n",
    "ax.plot(ds_fc_met_1d.time.values, ds_fc_met_1d.rolling(time=roll_len, min_periods=roll_len).mean(),  label='Fluxcom-Met', linewidth=line)\n",
    "ax.plot(cable_1d.time.values, cable_1d.rolling(time=roll_len, min_periods=roll_len).mean(), label='CABLE-BIOS', linewidth=line)\n",
    "ax.plot(cable_pop_1d.time.values, cable_pop_1d.rolling(time=roll_len, min_periods=roll_len).mean(),  label='CABLE-POP', linewidth=line)\n",
    "\n",
    "if var=='NEE':\n",
    "    ax.plot(post_1d.time.values, post_1d.rolling(time=roll_len, min_periods=roll_len).mean(), label='Inversion', color='darkkhaki', linewidth=line)\n",
    "    ax.plot(gosat.time.values, gosat.rolling(time=roll_len, min_periods=roll_len).mean(), label='GOSAT-Inversion', color='gold', linewidth=line)\n",
    "    ax.plot(ds_1d.time.values, ds_fire_1d.rolling(time=roll_len, min_periods=roll_len).mean(),  color='black', label='This study+Fire', linewidth=line)\n",
    "    plt.axhline(0, c='grey')\n",
    "\n",
    "if var == 'GPP':\n",
    "    ax.plot(modis_gpp_1d.time.values, modis_gpp_1d.rolling(time=roll_len, min_periods=roll_len).mean(), label='MODIS', linewidth=line)\n",
    "    ax.plot(gosif_gpp_1d.time.values, gosif_gpp_1d.rolling(time=roll_len, min_periods=roll_len).mean(), label='GOSIF', linewidth=line)\n",
    "    ax.plot(diffuse_1d.time.values, diffuse_1d.rolling(time=roll_len, min_periods=roll_len).mean(), label='DIFFUSE', linewidth=line)\n",
    "\n",
    "if var=='NEE':\n",
    "    ax.text(.015, .375, f'Fluxcom-RS: {ds_fc_1d.mean().values:.2f} PgC/yr',\n",
    "            transform=ax.transAxes, fontsize=fontsize)\n",
    "    ax.text(.015, .325, f'Fluxcom-Met: {ds_fc_met_1d.mean().values:.2f} PgC/yr',\n",
    "            transform=ax.transAxes, fontsize=fontsize)\n",
    "    ax.text(.015, .275, f'This study: -0.58 PgC/yr',\n",
    "            transform=ax.transAxes, fontsize=fontsize)\n",
    "    ax.text(.015, .225, f'This study+Fire: -0.44 PgC/yr',\n",
    "            transform=ax.transAxes, fontsize=fontsize)\n",
    "    ax.text(.015, .175, f'Cable-BIOS: {cable_1d.mean().values:.2f} PgC/yr',\n",
    "            transform=ax.transAxes, fontsize=fontsize)\n",
    "    ax.text(.015, .125, f'Cable-POP: {cable_pop_1d.mean().values:.2f} PgC/yr',\n",
    "            transform=ax.transAxes,fontsize=fontsize)\n",
    "    ax.text(.015, .075, f'Inversion: {post_1d.mean().values:.2f} PgC/yr',\n",
    "            transform=ax.transAxes, fontsize=fontsize)\n",
    "    ax.text(.015, .025, f'GOSAT-Inversion: {gosat.mean().values:.2f} PgC/yr',\n",
    "            transform=ax.transAxes, fontsize=fontsize)\n",
    "\n",
    "if var == 'GPP':\n",
    "    ax.text(.015, .95, f'Fluxcom-RS: {ds_fc_1d.mean().values:.2f} PgC/yr',\n",
    "            transform=ax.transAxes, fontsize=fontsize)\n",
    "    ax.text(.015, .90, f'Fluxcom-Met: {ds_fc_met_1d.mean().values:.2f} PgC/yr',\n",
    "            transform=ax.transAxes, fontsize=fontsize)\n",
    "    ax.text(.015, .85, f'This study: 4.25 PgC/yr',\n",
    "            transform=ax.transAxes, fontsize=fontsize)\n",
    "    ax.text(.015, .80, f'GOSIF: {gosif_gpp_1d.mean().values:.2f} PgC/yr',\n",
    "            transform=ax.transAxes, fontsize=fontsize)\n",
    "    ax.text(.015, .75, f'MODIS: {modis_gpp_1d.mean().values:.2f} PgC/yr',\n",
    "            transform=ax.transAxes, fontsize=fontsize)\n",
    "    ax.text(.015, .70, f'Cable-BIOS: {cable_1d.mean().values:.2f} PgC/yr',\n",
    "            transform=ax.transAxes, fontsize=fontsize)\n",
    "    ax.text(.015, .65, f'Cable-POP: {cable_pop_1d.mean().values:.2f} PgC/yr',\n",
    "            transform=ax.transAxes, fontsize=fontsize)\n",
    "    ax.text(.015, .60, f'DIFFUSE: {diffuse_1d.mean().values:.2f} PgC/yr',\n",
    "            transform=ax.transAxes, fontsize=fontsize)\n",
    "\n",
    "if var == 'ER':\n",
    "    ax.text(.015, .95, f'Fluxcom-RS: {ds_fc_1d.mean().values:.2f} PgC/yr',\n",
    "            transform=ax.transAxes, fontsize=fontsize)\n",
    "    ax.text(.015, .90, f'Fluxcom-Met: {ds_fc_met_1d.mean().values:.2f} PgC/yr',\n",
    "            transform=ax.transAxes, fontsize=fontsize)\n",
    "    ax.text(.015, .85, f'This study: 3.64 PgC/yr',\n",
    "            transform=ax.transAxes, fontsize=fontsize)\n",
    "    ax.text(.015, .80, f'Cable-BIOS: {cable_1d.mean().values:.2f} PgC/yr',\n",
    "            transform=ax.transAxes, fontsize=fontsize)\n",
    "    ax.text(.015, .75, f'Cable-POP: {cable_pop_1d.mean().values:.2f} PgC/yr',\n",
    "            transform=ax.transAxes, fontsize=fontsize)\n",
    "\n",
    "ax.grid('off', which='major', axis='both', linestyle='--', linewidth=0.75)\n",
    "ax.grid('off', which='major', axis='y', linestyle='--', linewidth=0.75)\n",
    "ax.set_ylabel(var+' PgC/year', fontsize=fontsize);\n",
    "ax.tick_params(axis='x', labelsize=fontsize)\n",
    "ax.tick_params(axis='y', labelsize=fontsize)\n",
    "ax.set_yticklabels(['{:3.1f}'.format(x) for x in ax.get_yticks()])\n",
    "plt.tight_layout()\n",
    "fig.savefig('/g/data/os22/chad_tmp/NEE_modelling/results/figs/annual_mean_timeseries_'+var+'.png',\n",
    "            bbox_inches='tight', dpi=300)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Save legend of plot above"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# get handles and labels for reuse\n",
    "label_params = ax.get_legend_handles_labels() \n",
    "\n",
    "figl, axl = plt.subplots()\n",
    "axl.axis(False)\n",
    "axl.legend(*label_params, loc=\"best\", fontsize=20, markerscale=4.0, ncol=6)\n",
    "figl.savefig('/g/data/os22/chad_tmp/NEE_modelling/results/figs/legend_annual_mean_timeseries_'+var+'.png',\n",
    "             bbox_inches='tight', dpi=300)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Convert to anomalies"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Climatology"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import warnings\n",
    "warnings.filterwarnings(\"ignore\")\n",
    "\n",
    "ds_clim_mean = ds.groupby('time.month').mean()\n",
    "ds_upper_clim_mean = ds_upper.groupby('time.month').mean()\n",
    "\n",
    "ds_fc_clim_mean = ds_fc.groupby('time.month').mean()\n",
    "ds_fc_met_clim_mean = ds_fc_met.groupby('time.month').mean()\n",
    "post_clim_mean = post.groupby('time.month').mean()\n",
    "cable_clim_mean = cable.groupby('time.month').mean()\n",
    "cable_pop_clim_mean = cable_pop.groupby('time.month').mean()\n",
    "gosat_clim_mean = gosat.groupby('time.month').mean()\n",
    "\n",
    "if var=='GPP':\n",
    "    modis_gpp_clim_mean = modis_gpp.groupby('time.month').mean()\n",
    "    gosif_gpp_clim_mean = gosif_gpp.groupby('time.month').mean()\n",
    "    diffuse_clim_mean = diffuse.groupby('time.month').mean()\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Anomalies"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ds_anom = (ds.groupby('time.month') - ds_clim_mean)\n",
    "# ds_upper_anom = (ds_upper.groupby('time.month') - ds_upper_clim_mean) \n",
    "ds_fc_anom = (ds_fc.groupby('time.month') - ds_fc_clim_mean)\n",
    "ds_fc_met_anom = (ds_fc_met.groupby('time.month') - ds_fc_met_clim_mean)\n",
    "post_anom = (post.groupby('time.month') - post_clim_mean)\n",
    "cable_anom = (cable.groupby('time.month') - cable_clim_mean)\n",
    "cable_pop_anom = (cable_pop.groupby('time.month') - cable_pop_clim_mean)\n",
    "gosat_anom = (gosat.groupby('time.month') - gosat_clim_mean)\n",
    "\n",
    "if var=='GPP':\n",
    "    modis_gpp_anom = (modis_gpp.groupby('time.month') - modis_gpp_clim_mean)\n",
    "    gosif_gpp_anom = (gosif_gpp.groupby('time.month') - gosif_gpp_clim_mean)\n",
    "    diffuse_anom = (diffuse.groupby('time.month') - diffuse_clim_mean)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Interannual variability"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "annual_ds = ds.sel(time=slice('2003', '2021')).resample(time='1Y').sum()\n",
    "annual_ds_upper = ds_upper.sel(time=slice('2003', '2021')).resample(time='1Y').sum()\n",
    "annual_post = post.resample(time='1Y').sum()\n",
    "annual_gosat = gosat.resample(time='1Y').sum()\n",
    "annual_cable = cable.resample(time='1Y').sum()\n",
    "annual_cable_pop = cable_pop.resample(time='1Y').sum()\n",
    "annual_ds_fc = ds_fc.resample(time='1Y').sum()\n",
    "annual_ds_fc_met = ds_fc_met.resample(time='1Y').sum()\n",
    "\n",
    "if var=='GPP':\n",
    "    annual_modis = modis_gpp.resample(time='1Y').sum()\n",
    "    annual_gosif = gosif_gpp.resample(time='1Y').sum()\n",
    "    annual_diffuse = diffuse.resample(time='1Y').sum()\n",
    "\n",
    "if var=='NEE':\n",
    "    annual_fire = fire.resample(time='1Y').sum()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print('fluxcom-rs: ', (annual_ds_fc.sum(['x','y']) / 12).std().values)\n",
    "print('fluxcom-met: ',(annual_ds_fc_met.sum(['x','y']) / 12).std().values)\n",
    "print('CABLE-BIOS: ',(annual_cable.sum(['x','y']) / 12).std().values)\n",
    "print('CABLE-POP: ',(annual_cable_pop.sum(['x','y']) / 12).std().values)\n",
    "print('Inversion: ',(annual_post.sum(['x','y']) / 12).std().values)\n",
    "print('This Study: ',(annual_ds.sum(['x','y']) / 12).std().values)\n",
    "print('GOSAT-Inversion: ',(annual_gosat/ 12).std().values)\n",
    "\n",
    "if var=='GPP':\n",
    "    print('MODIS: ',(annual_modis.sum(['x','y']) / 12).std().values)\n",
    "    print('GOSIF: ',(annual_gosif.sum(['x','y']) / 12).std().values)\n",
    "    print('DIFFUSE: ',(annual_diffuse.sum(['x','y']) / 12).std().values)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Zonal summaries"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ds_1D = ds_anom.sum(['x','y'])\n",
    "# ds_upper_1D = ds_upper_anom.sum(['x','y'])\n",
    "\n",
    "ds_fc_1D = ds_fc_anom.sum(['x','y'])\n",
    "ds_fc_met_1D = ds_fc_met_anom.sum(['x','y'])\n",
    "\n",
    "post_1D = post_anom.sum(['x','y'])\n",
    "cable_1D = cable_anom.sum(['x','y'])\n",
    "cable_pop_1D = cable_pop_anom.sum(['x','y'])\n",
    "\n",
    "if var=='GPP':\n",
    "    modis_gpp_1D = modis_gpp_anom.sum(['x','y'])\n",
    "    gosif_gpp_1D = gosif_gpp_anom.sum(['x','y'])\n",
    "    diffuse_1D = diffuse_anom.sum(['x','y'])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Plot timeseries"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "roll_len=12"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fontsize=20\n",
    "\n",
    "fig, ax = plt.subplots(1,1, figsize=(15,7))\n",
    "ax.plot(ds_1D.time.values, ds_1D.rolling(time=roll_len, min_periods=roll_len).mean(),  label='This study', linewidth=line)\n",
    "ax.plot(ds_fc_1D.time.values, ds_fc_1D.rolling(time=roll_len, min_periods=roll_len).mean(),  label='Fluxcom-RS', linewidth=line)\n",
    "ax.plot(ds_fc_met_1D.time.values, ds_fc_met_1D.rolling(time=roll_len, min_periods=roll_len).mean(),  label='Fluxcom-Met', linewidth=line)\n",
    "ax.plot(cable_1D.time.values, cable_1D.rolling(time=roll_len, min_periods=roll_len).mean(), label='CABLE-BIOS', linewidth=line)\n",
    "ax.plot(cable_pop_1D.time.values, cable_pop_1D.rolling(time=roll_len, min_periods=roll_len).mean(), label='CABLE-POP', linewidth=line)\n",
    "\n",
    "\n",
    "if var=='NEE':\n",
    "    ax.plot(post_1D.time.values, post_1D.rolling(time=roll_len, min_periods=roll_len).mean(), label='Inversion',color='darkkhaki', linewidth=line)\n",
    "    ax.plot(gosat_anom.time.values, gosat_anom.rolling(time=roll_len, min_periods=roll_len).mean(), label='GOSAT-Inversion', color='gold', linewidth=line)\n",
    "    ax.plot(ds_1D.time.values, ds_fire_anom.rolling(time=roll_len, min_periods=roll_len).mean(),  c='black', label='This study+Fire', linewidth=line)\n",
    "\n",
    "if var == 'GPP':\n",
    "    ax.plot(modis_gpp_1D.time.values, modis_gpp_1D.rolling(time=roll_len, min_periods=roll_len).mean(),  label='MODIS', linewidth=line)\n",
    "    ax.plot(gosif_gpp_1D.time.values, gosif_gpp_1D.rolling(time=roll_len, min_periods=roll_len).mean(), label='GOSIF', linewidth=line)\n",
    "    ax.plot(diffuse_1D.time.values, diffuse_1D.rolling(time=roll_len, min_periods=roll_len).mean(), label='DIFFUSE', linewidth=line)\n",
    "\n",
    "ax.set_ylabel(var+' anomaly PgC/year', fontsize=fontsize)\n",
    "plt.axhline(0, c='grey');\n",
    "\n",
    "ax.grid('off', which='major', axis='both', linestyle='--', linewidth=0.75)\n",
    "ax.grid('off', which='major', axis='y', linestyle='--', linewidth=0.75)\n",
    "ax.tick_params(axis='x', labelsize=fontsize)\n",
    "ax.tick_params(axis='y', labelsize=fontsize)\n",
    "ax.set_yticklabels(['{:3.1f}'.format(x) for x in ax.get_yticks()])\n",
    "plt.tight_layout()\n",
    "fig.savefig('/g/data/os22/chad_tmp/NEE_modelling/results/figs/annual_anomaly_timeseries_'+var+'.png',\n",
    "            bbox_inches='tight' , dpi=300)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Climatology timeseries"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Aus wide "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print('Seasonal range of '+var, ds_clim_mean.sum(['x','y']).max().values - ds_clim_mean.sum(['x','y']).min().values)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fontsize=20\n",
    "line=2\n",
    "\n",
    "fig, ax = plt.subplots(1,1, figsize=(9,7))\n",
    "\n",
    "y = ds_clim_mean.sum(['x','y']).values\n",
    "x = ds_clim_mean.month.values\n",
    "\n",
    "if var=='NEE':\n",
    "    y_fire = ds_fire_clim_mean.values\n",
    "    x_fire = ds_fire_clim_mean.month.values\n",
    "\n",
    "y_fc = ds_fc_clim_mean.sum(['x','y']).values\n",
    "x_fc = ds_fc_clim_mean.month.values\n",
    "\n",
    "y_fc_met = ds_fc_met_clim_mean.sum(['x','y']).values\n",
    "x_fc_met = ds_fc_met_clim_mean.month.values\n",
    "\n",
    "y_post = post_clim_mean.sum(['x','y']).values\n",
    "x_post = post_clim_mean.month.values\n",
    "\n",
    "y_gosat = gosat_clim_mean.values\n",
    "x_gosat = gosat_clim_mean.month.values\n",
    "\n",
    "y_cable = cable_clim_mean.sum(['x','y']).values\n",
    "x_cable = cable_clim_mean.month.values\n",
    "\n",
    "y_cable_pop = cable_pop_clim_mean.sum(['x','y']).values\n",
    "x_cable_pop = cable_pop_clim_mean.month.values\n",
    "\n",
    "if var =='GPP':\n",
    "    y_modis_gpp = modis_gpp_clim_mean.sum(['x','y']).values\n",
    "    x_modis_gpp = modis_gpp_clim_mean.month.values\n",
    "\n",
    "    y_gosif_gpp = gosif_gpp_clim_mean.sum(['x','y']).values\n",
    "    x_gosif_gpp = gosif_gpp_clim_mean.month.values\n",
    "    \n",
    "    y_diffuse = diffuse_clim_mean.sum(['x','y']).values\n",
    "    x_diffuse = diffuse_clim_mean.month.values\n",
    "\n",
    "\n",
    "ax.plot(x, y, label='This Study', linewidth=line)\n",
    "if var=='NEE':\n",
    "    ax.plot(x_fire, y_fire, label='This Study+Fire', c='black', linewidth=line)\n",
    "ax.plot(x_fc, y_fc, label='Fluxcom-RS', linewidth=line)\n",
    "ax.plot(x_fc_met, y_fc_met, label='Fluxcom-Met', linewidth=line)\n",
    "ax.plot(x_cable, y_cable, label='Cable-BIOS', linewidth=line)\n",
    "ax.plot(x_cable_pop, y_cable_pop, label='Cable-POP', linewidth=line)\n",
    "\n",
    "if var=='NEE':\n",
    "    ax.plot(x_post, y_post, label='Inversion', color='darkkhaki', linewidth=line)\n",
    "    ax.plot(x_gosat, y_gosat, label='GOSAT-Inversion', color='gold', linewidth=line)\n",
    "\n",
    "if var =='GPP':\n",
    "    ax.plot(x_modis_gpp, y_modis_gpp, label='MODIS', linewidth=line)\n",
    "    ax.plot(x_gosif_gpp, y_gosif_gpp, label='GOSIF', linewidth=line)\n",
    "    ax.plot(x_diffuse, y_diffuse, label='DIFFUSE', linewidth=line)\n",
    "\n",
    "\n",
    "if var=='NEE':\n",
    "    ax.axhline(0, c='grey')\n",
    "# ax.legend(loc='best', fontsize=15)\n",
    "if var!='NEE':\n",
    "    ax.set_ylim(1.75,6.75)\n",
    "\n",
    "ax.set_yticklabels(['{:3.1f}'.format(x) for x in ax.get_yticks()])\n",
    "ax.set_xticks(range(1,13))\n",
    "ax.set_xticklabels([\"J\",\"F\",\"M\",\"A\",\"M\",\"J\",\"J\",\"A\",\"S\",\"O\",\"N\",\"D\"], fontsize=15)    \n",
    "ax.tick_params(axis='x', labelsize=20)\n",
    "ax.tick_params(axis='y', labelsize=20)\n",
    "ax.set_ylabel(var+' PgC/year', fontsize=25);\n",
    "plt.savefig('/g/data/os22/chad_tmp/NEE_modelling/results/figs/'+var+'_Aus_climatology_timeseries.png',\n",
    "            bbox_inches='tight', dpi=300)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Save the legend as file"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# get handles and labels for reuse\n",
    "label_params = ax.get_legend_handles_labels() \n",
    "\n",
    "figl, axl = plt.subplots()\n",
    "axl.axis(False)\n",
    "axl.legend(*label_params, loc=\"best\", fontsize=20, markerscale=2.0, ncol=4)\n",
    "figl.savefig('/g/data/os22/chad_tmp/NEE_modelling/results/figs/climatology_legend_'+var+'.png', bbox_inches='tight', dpi=300)"
   ]
  },
  {
   "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 = xr_reproject(bio, geobox=xx.odc.geobox.compat, 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",
    "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(ds, index)\n",
    "    ds_fc_mask = make_mask(ds_fc, index)\n",
    "    ds_fc_met_mask = make_mask(ds_fc_met, index)\n",
    "    cable_mask = make_mask(cable, index)\n",
    "    cable_pop_mask = make_mask(cable_pop, index)\n",
    "    post_mask = make_mask(post, index)\n",
    "    \n",
    "    dss = round_coords(ds_clim_mean).where(ds_mask)\n",
    "    dss_fc = round_coords(ds_fc_clim_mean).where(ds_fc_mask)\n",
    "    dss_fc_met = round_coords(ds_fc_met_clim_mean).where(ds_fc_met_mask)\n",
    "    dss_cable= round_coords(cable_clim_mean).where(cable_mask)\n",
    "    dss_cable_pop = round_coords(cable_pop_clim_mean).where(cable_pop_mask)\n",
    "    dss_post = round_coords(post_clim_mean).where(post_mask)\n",
    "    \n",
    "    if var=='GPP':\n",
    "        modis_gpp_mask = make_mask(modis_gpp, index)\n",
    "        gosif_gpp_mask = make_mask(gosif_gpp, index)\n",
    "        diffuse_mask = make_mask(diffuse, index)\n",
    "        dss_modis = round_coords(modis_gpp_clim_mean).where(modis_gpp_mask)\n",
    "        dss_gosif = round_coords(gosif_gpp_clim_mean).where(gosif_gpp_mask)\n",
    "        dss_diffuse = round_coords(diffuse_clim_mean).where(diffuse_mask)\n",
    "        \n",
    "        datasets.update({\n",
    "                    'This Study':dss,\n",
    "                    'Fluxcom-RS':dss_fc,\n",
    "                    'Fluxcom-Met':dss_fc_met,\n",
    "                    'CABLE-BIOS':dss_cable,\n",
    "                    'CABLE-POP':dss_cable_pop,\n",
    "                    'MODIS-GPP':dss_modis,\n",
    "                    'GOSIF-GPP':dss_gosif,\n",
    "                    'DIFFUSE-GPP':dss_diffuse,\n",
    "                    }\n",
    "                   )\n",
    "    if var=='NEE':\n",
    "        datasets.update({\n",
    "                        'This Study':dss,\n",
    "                        'Fluxcom-RS':dss_fc,\n",
    "                        'Fluxcom-Met':dss_fc_met,\n",
    "                        'CABLE-BIOS':dss_cable,\n",
    "                        'CABLE-POP':dss_cable_pop,\n",
    "                        'Inversion':dss_post\n",
    "                        }\n",
    "                       )\n",
    "    \n",
    "    if var=='ER':\n",
    "        datasets.update({\n",
    "                        'This Study':dss,\n",
    "                        'Fluxcom-RS':dss_fc,\n",
    "                        'Fluxcom-Met':dss_fc_met,\n",
    "                        'CABLE-BIOS':dss_cable,\n",
    "                        'CABLE-POP':dss_cable_pop,\n",
    "                        }\n",
    "                       )\n",
    "    \n",
    "    # Append results to a dictionary using the attribute\n",
    "    # column as an key\n",
    "    results.update({row['region_name']: datasets})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "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",
    "    for source, da in results[k].items():\n",
    "        _1D = da.sum(['x','y'])\n",
    "        if source=='Inversion':\n",
    "            _1D.plot(ax=ax, label=source, color='darkkhaki')\n",
    "        else:    \n",
    "            _1D.plot(ax=ax, label=source)\n",
    "    if var=='NEE':\n",
    "        ax.axhline(0, c='grey', linestyle='--')\n",
    "        \n",
    "    ax.yaxis.set_major_formatter(FormatStrFormatter('%.1f'))\n",
    "    # ax.set_yticklabels(['{:3.1f}'.format(y) for y in ax.get_yticks()])\n",
    "    ax.set_xticks(range(1,13))\n",
    "    ax.set_xticklabels([\"J\",\"F\",\"M\",\"A\",\"M\",\"J\",\"J\",\"A\",\"S\",\"O\",\"N\",\"D\"], fontsize=17)\n",
    "    ax.tick_params(axis='x', labelsize=20)\n",
    "    ax.tick_params(axis='y', labelsize=20)\n",
    "    ax.set_title(k, fontdict={'fontsize': 22})\n",
    "   \n",
    "    ax.set_ylabel('')\n",
    "    ax.set_xlabel('')\n",
    "\n",
    "#axs[0,0].legend(fontsize=15)\n",
    "fig.supylabel(var+' PgC/year', fontsize=25)\n",
    "plt.tight_layout()\n",
    "plt.savefig('/g/data/os22/chad_tmp/NEE_modelling/results/figs/'+var+'_bioregion_climatology_timeseries.png',\n",
    "            bbox_inches='tight', dpi=300);"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Climatology per-pixel"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if var == 'NEE':\n",
    "    max_mean_month = allNaN_arg(ds_clim_mean, 'month', 'min')\n",
    "    max_mean_month_fc = allNaN_arg(ds_fc_clim_mean, 'month', 'min')\n",
    "    max_mean_month_fc_met = allNaN_arg(ds_fc_met_clim_mean, 'month', 'min')\n",
    "    max_mean_month_post = allNaN_arg(post_clim_mean, 'month', 'min')\n",
    "    max_mean_month_cable = allNaN_arg(cable_clim_mean, 'month', 'min')\n",
    "    max_mean_month_cable_pop = allNaN_arg(cable_pop_clim_mean, 'month', 'min')\n",
    "\n",
    "else:\n",
    "    if var=='GPP':\n",
    "        max_mean_month = allNaN_arg(ds_clim_mean, 'month', 'max')\n",
    "        max_mean_month_fc = allNaN_arg(ds_fc_clim_mean, 'month', 'max')\n",
    "        max_mean_month_fc_met = allNaN_arg(ds_fc_met_clim_mean, 'month', 'max')\n",
    "        max_mean_month_modis = allNaN_arg(modis_gpp_clim_mean, 'month', 'max')\n",
    "        max_mean_month_gosif = allNaN_arg(gosif_gpp_clim_mean, 'month', 'max')\n",
    "        max_mean_month_diffuse = allNaN_arg(diffuse_clim_mean, 'month', 'max')\n",
    "        max_mean_month_cable = allNaN_arg(cable_clim_mean, 'month', 'max')\n",
    "        max_mean_month_cable_pop = allNaN_arg(cable_pop_clim_mean, 'month', 'max')\n",
    "    if var=='ER':\n",
    "        max_mean_month = allNaN_arg(ds_clim_mean, 'month', 'max')\n",
    "        max_mean_month_fc = allNaN_arg(ds_fc_clim_mean, 'month', 'max')\n",
    "        max_mean_month_fc_met = allNaN_arg(ds_fc_met_clim_mean, 'month', 'max')\n",
    "        max_mean_month_cable = allNaN_arg(cable_clim_mean, 'month', 'max')\n",
    "        max_mean_month_cable_pop = allNaN_arg(cable_pop_clim_mean, 'month', 'max')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "y_range = ' (2003-21)'\n",
    "font=15\n",
    "if var=='GPP':\n",
    "    mask = xr.where(ds_clim_mean>0, 1, 0).isel(month=0).drop('month')\n",
    "    fig,ax = plt.subplots(2,4, figsize=(22,10))\n",
    "    im = max_mean_month.where(mask).plot.imshow(vmin=1, vmax=12, ax=ax[0,0], cmap='twilight', add_colorbar=False, add_labels=False)\n",
    "    ax[0,0].set_title('This Study '+var+y_range, fontsize=font);\n",
    "\n",
    "    im = max_mean_month_fc.plot.imshow(vmin=1, vmax=12, ax=ax[0,1], cmap='twilight', add_colorbar=False, add_labels=False)\n",
    "    ax[0,1].set_title('Fluxcom-RS '+var+ ' (2003-2015)', fontsize=font);\n",
    "\n",
    "    im = max_mean_month_fc_met.plot.imshow(vmin=1, vmax=12, ax=ax[0,2], cmap='twilight', add_colorbar=False, add_labels=False)\n",
    "    ax[0,2].set_title('Fluxcom-Met '+var+' (2003-2018)', fontsize=font);\n",
    "\n",
    "    im = max_mean_month_cable.plot.imshow(vmin=1, vmax=12, ax=ax[0,3], cmap='twilight', add_colorbar=False, add_labels=False)\n",
    "    ax[0,3].set_title('Cable-BIOS '+var+y_range, fontsize=font);\n",
    "\n",
    "    im = max_mean_month_cable_pop.plot.imshow(vmin=1, vmax=12, ax=ax[1,0], cmap='twilight', add_colorbar=False, add_labels=False)\n",
    "    ax[1,0].set_title('Cable-POP '+var+y_range, fontsize=font);\n",
    "\n",
    "    im = max_mean_month_modis.where(mask).plot.imshow(vmin=1, vmax=12, ax=ax[1,1], cmap='twilight', add_colorbar=False, add_labels=False)\n",
    "    ax[1,1].set_title('MODIS '+var+y_range, fontsize=font);\n",
    "\n",
    "    im = max_mean_month_gosif.plot.imshow(vmin=1, vmax=12, ax=ax[1,2], cmap='twilight', add_colorbar=False, add_labels=False)\n",
    "    ax[1,2].set_title('GOSIF '+var+y_range, fontsize=font);\n",
    "    \n",
    "    im = max_mean_month_diffuse.where(~np.isnan(diffuse_clim_mean.isel(month=1))).plot.imshow(vmin=1, vmax=12, ax=ax[1,3], cmap='twilight', add_colorbar=False, add_labels=False)\n",
    "    ax[1,3].set_title('DIFFUSE '+var+y_range, fontsize=font);\n",
    "\n",
    "if var=='NEE':\n",
    "    mask = ds_clim_mean.mean('month')\n",
    "    mask = xr.where(mask==0, 0, 1)\n",
    "    fig,ax = plt.subplots(2,3, figsize=(18,10))\n",
    "    im = max_mean_month.where(mask).plot.imshow(vmin=1, vmax=12, ax=ax[0,0], cmap='twilight', add_colorbar=False, add_labels=False)\n",
    "    ax[0,0].set_title('This Study '+var+y_range, fontsize=font);\n",
    "\n",
    "    im = max_mean_month_fc.plot.imshow(vmin=1, vmax=12, ax=ax[0,1], cmap='twilight', add_colorbar=False, add_labels=False)\n",
    "    ax[0,1].set_title('Fluxcom-RS '+var+ ' (2003-2015)', fontsize=font);\n",
    "\n",
    "    im = max_mean_month_fc_met.plot.imshow(vmin=1, vmax=12, ax=ax[0,2], cmap='twilight', add_colorbar=False, add_labels=False)\n",
    "    ax[0,2].set_title('Fluxcom-Met '+var+' (2003-2018)', fontsize=font);\n",
    "\n",
    "    im = max_mean_month_cable.plot.imshow(vmin=1, vmax=12, ax=ax[1,0], cmap='twilight', add_colorbar=False, add_labels=False)\n",
    "    ax[1,0].set_title('Cable-BIOS '+var+y_range, fontsize=font);\n",
    "\n",
    "    im = max_mean_month_cable_pop.plot.imshow(vmin=1, vmax=12, ax=ax[1,1], cmap='twilight', add_colorbar=False, add_labels=False)\n",
    "    ax[1,1].set_title('Cable-POP '+var+y_range, fontsize=font);\n",
    "\n",
    "    im = max_mean_month_post.isel(x=range(12, 68), y=range(12, 68)).plot.imshow(vmin=1, vmax=12, ax=ax[1,2], cmap='twilight', add_colorbar=False, add_labels=False)\n",
    "    ax[1,2].set_title('Inversion '+var+' (2015-19)', fontsize=font);\n",
    "\n",
    "if var=='ER':\n",
    "    mask = ds_clim_mean.mean('month')\n",
    "    mask = xr.where(mask==0, 0, 1)\n",
    "    fig,ax = plt.subplots(2,3, figsize=(18,10))\n",
    "    im = max_mean_month.where(mask).plot.imshow(vmin=1, vmax=12, ax=ax[0,0], cmap='twilight', add_colorbar=False, add_labels=False)\n",
    "    ax[0,0].set_title('This Study '+var+y_range, fontsize=font);\n",
    "\n",
    "    im = max_mean_month_fc.plot.imshow(vmin=1, vmax=12, ax=ax[0,1], cmap='twilight', add_colorbar=False, add_labels=False)\n",
    "    ax[0,1].set_title('Fluxcom-RS '+var+ ' (2003-2015)', fontsize=font);\n",
    "\n",
    "    im = max_mean_month_fc_met.plot.imshow(vmin=1, vmax=12, ax=ax[0,2], cmap='twilight', add_colorbar=False, add_labels=False)\n",
    "    ax[0,2].set_title('Fluxcom-Met '+var+' (2003-2018)', fontsize=font);\n",
    "\n",
    "    im = max_mean_month_cable.plot.imshow(vmin=1, vmax=12, ax=ax[1,0], cmap='twilight', add_colorbar=False, add_labels=False)\n",
    "    ax[1,0].set_title('Cable-BIOS '+var+y_range, fontsize=font);\n",
    "\n",
    "    im = max_mean_month_cable_pop.plot.imshow(vmin=1, vmax=12, ax=ax[1,1], cmap='twilight', add_colorbar=False, add_labels=False)\n",
    "    ax[1,1].set_title('Cable-POP '+var+y_range, fontsize=font);\n",
    "    \n",
    "\n",
    "ax[0,0].set_yticks([])\n",
    "ax[0,1].set_yticks([])\n",
    "ax[0,2].set_yticks([])\n",
    "ax[1,0].set_yticks([])\n",
    "ax[1,1].set_yticks([])\n",
    "ax[1,2].set_yticks([])\n",
    "ax[0,0].set_xticks([])\n",
    "ax[0,1].set_xticks([])\n",
    "ax[0,2].set_xticks([])\n",
    "ax[1,0].set_xticks([])\n",
    "ax[1,1].set_xticks([])\n",
    "ax[1,2].set_xticks([])\n",
    "\n",
    "if var=='GPP':\n",
    "    ax[0,3].set_yticks([])\n",
    "    ax[1,3].set_yticks([])\n",
    "    ax[0,3].set_xticks([])\n",
    "    ax[1,3].set_xticks([])\n",
    "\n",
    "plt.tight_layout()\n",
    "plt.savefig('/g/data/os22/chad_tmp/NEE_modelling/results/figs/month_of_max_'+var+'.png',\n",
    "            bbox_inches='tight', dpi=300)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Save the colorbar"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig = plt.figure()\n",
    "ax = fig.add_axes([0.05, 0.80, 0.9, 0.1])\n",
    "# cbar = fig.colorbar(im, )\n",
    "cbar = mpl.colorbar.ColorbarBase(orientation='horizontal',\n",
    "                                 ax=ax,\n",
    "                                 norm = mpl.colors.Normalize(vmin=1, vmax=12),\n",
    "                                 ticks=[1,2,3,4,5,6,7,8,9,10,11,12],\n",
    "                                 cmap='twilight')\n",
    "\n",
    "# cbar.set_ticks([1,2,3,4,5,6,7,8,9,10,11,12])\n",
    "cbar.set_ticklabels([\"J\",\"F\",\"M\",\"A\",\"M\",\"J\",\"J\",\"A\",\"S\",\"O\",\"N\",\"D\"], fontsize=font)\n",
    "\n",
    "plt.savefig('/g/data/os22/chad_tmp/NEE_modelling/results/figs/monthofmax_colorbar.png', \n",
    "            bbox_inches='tight', dpi=300)\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
}