8_Correlations.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Calculate pearson correlations between climate variables and fluxes\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import sys\n",
    "import xarray as xr\n",
    "import numpy as np\n",
    "from scipy import stats\n",
    "import scipy as sp\n",
    "import geopandas as gpd\n",
    "from odc.geo.xr import assign_crs\n",
    "import pandas as pd\n",
    "from odc.algo import xr_reproject\n",
    "# from odc.geo.geobox import zoom_out\n",
    "from matplotlib import pyplot as plt\n",
    "from datacube.utils.dask import start_local_dask\n",
    "\n",
    "sys.path.append('/g/data/os22/chad_tmp/dea-notebooks/Tools/')\n",
    "from dea_tools.spatial import xr_rasterize\n",
    "\n",
    "sys.path.append('/g/data/os22/chad_tmp/NEE_modelling/')\n",
    "from _collect_prediction_data import round_coords, collect_prediction_data "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "client = start_local_dask(mem_safety_margin='2Gb')\n",
    "client"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "var = 'GPP'\n",
    "suffix='20230320'\n",
    "# results_name = var+'_2003_2022_1km_quantiles_'+suffix+'.nc'\n",
    "results_name ='AusEFlux_'+var+'_2003_2022_1km_quantiles_v1.1.nc'\n",
    "chunks_data = {'x':1100,'y':1100, 'time':1}\n",
    "chunks = {'longitude':1100,'latitude':1100, 'time':1}\n",
    "# mask_path = '/g/data/os22/chad_tmp/NEE_modelling/results/prediction_data/mask_5km.nc'\n",
    "# data_path = '/g/data/os22/chad_tmp/NEE_modelling/results/prediction_data/data_5km.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=chunks)[var+'_median']\n",
    "ds"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# grid = zoom_out(ds.odc.geobox, 2)\n",
    "# ds = xr_reproject(ds, geobox=grid.compat, resampling='average').compute()\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Open predictor data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "data = collect_prediction_data(time_start='2003',\n",
    "                             time_end='2022',\n",
    "                             verbose=False,\n",
    "                             export=False,\n",
    "                             chunks=chunks_data\n",
    "                             )\n",
    "data = data.rename({'x':'longitude', 'y':'latitude'}).chunk(chunks)\n",
    "data"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Calculate climatology and anomaly\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ds_clim_mean = ds.groupby('time.month').mean().compute()\n",
    "# ds_anom = (ds.groupby('time.month') - ds_clim_mean).compute()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Create table of correlations per bioclimatic region"
   ]
  },
  {
   "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": {},
   "outputs": [],
   "source": [
    "clim_vars = ['rain_anom', 'rain_cml3_anom', 'rain_cml6_anom','rain_cml12_anom', 'tavg_anom', 'srad_anom', 'kNDVI_anom']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ds_anom = assign_crs(ds_anom, crs='EPSG:4326')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "outer = {}\n",
    "for index, row in gdf.iterrows():\n",
    "    print(row['region_name'])\n",
    "    mask = xr_rasterize(gdf.iloc[[index]], ds_anom.isel(time=1))\n",
    "    mask = round_coords(mask)\n",
    "    # mask = mask.rename({'latitude':'y', 'longitude':'x'})\n",
    "    inner = {}\n",
    "    for v in clim_vars:\n",
    "        var_anom_region = data[v].where(mask).compute()\n",
    "        ds_anom_region = ds_anom.where(mask)\n",
    "        r2 = xr.corr(ds_anom_region.chunk(chunks),\n",
    "                     var_anom_region.chunk(chunks),\n",
    "                     dim='time').compute()\n",
    "        r2 = r2.mean(['latitude', 'longitude'])\n",
    "        print('  ', v, r2.values)\n",
    "        inner[v] = r2.values\n",
    "    outer[row['region_name']] = inner"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "df = pd.DataFrame(outer)\n",
    "df"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "df.to_csv('/g/data/os22/chad_tmp/NEE_modelling/results/'+var+'_anomaly_bioregion_correlations.csv')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Over all of Aus"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# r2 = xr.corr(ds_anom_region, var_anom_region, dim='time').compute()\n",
    "# r2 = r2.mean(['x', 'y'])\n",
    "# print(r2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# ax2_ylim = -100,100\n",
    "# ax_ylim = -1,1\n",
    "\n",
    "# fig,ax=plt.subplots(1,2, figsize=(18,5), gridspec_kw={'width_ratios': [3, 1]})\n",
    "# ax2 = ax[0].twinx()\n",
    "# var_anom.mean(['x','y']).rolling(time=3).mean().plot(ax=ax2, label=data_var, c='orange')\n",
    "# ds_anom.sum(['x','y']).rolling(time=3).mean().plot(ax=ax[0], label=var)\n",
    "# ax[0].legend(loc=(0.80,0.925))\n",
    "# ax2.legend(loc=(0.80,0.85))\n",
    "# ax2.set_ylabel(data_var+' Anomaly', fontsize=15)\n",
    "# ax[0].set_xlabel('')\n",
    "# ax2.set_ylim(ax2_ylim)\n",
    "# ax[0].set_ylim(ax_ylim)\n",
    "# ax[0].text(.05, .90, 'r={:.2f}'.format(r2[0]),\n",
    "#             transform=ax[0].transAxes, fontsize=15)\n",
    "# ax[0].set_ylabel(var+' Anomalies (PgC y⁻¹)', fontsize=15)\n",
    "# ax[0].tick_params(axis='x', labelsize=14)\n",
    "# ax[0].tick_params(axis='y', labelsize=14)\n",
    "# ax2.tick_params(axis='y', labelsize=14)\n",
    "# ax[1].tick_params(axis='x', labelsize=14)\n",
    "# ax[1].tick_params(axis='y', labelsize=14)\n",
    "\n",
    "# ax[0].axhline(0, c='grey', linestyle='--')\n",
    "\n",
    "# ax3 = ax[1].twinx()\n",
    "\n",
    "# var_clim_mean.mean(['x','y']).plot(ax=ax3, label=data_var, c='orange')\n",
    "# ds_clim_mean.mean(['x','y']).plot(ax=ax[1], label='NEE')\n",
    "# ax3.set_ylabel(data_var)\n",
    "# ax3.set_ylabel(data_var, fontsize=15)\n",
    "# ax[1].set_ylabel(var+' (PgC y⁻¹)', fontsize=15)\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\"]) \n",
    "# ax[1].set_xlabel('')\n",
    "# ax3.tick_params(axis='y', labelsize=14)\n",
    "# ax[0].set_title(None)\n",
    "# ax[1].set_title(None)\n",
    "# ax2.set_title(None)\n",
    "# ax3.set_title(None)\n",
    "# plt.tight_layout();\n",
    "# plt.savefig('/g/data/os22/chad_tmp/NEE_modelling/results/figs/'+var+'_Aus_'+data_var+'_correlations.png')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Per-pixel correlations"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Correlations with climatology"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "c = {'latitude':1100,'longitude':1100}\n",
    "\n",
    "precip_clim_corr = xr.corr(ds_clim_mean.chunk(c), data['rain'].groupby('time.month').mean(), dim='month').compute()\n",
    "# srad_clim_corr = xr.corr(ds_clim_mean.chunk(c), data['srad'].groupby('time.month').mean(), dim='month').compute()\n",
    "# tavg_clim_corr = xr.corr(ds_clim_mean.chunk(c), data['tavg'].groupby('time.month').mean(), dim='month').compute()\n",
    "# kNDVI_clim_corr = xr.corr(ds_clim_mean.chunk(c), data['kNDVI'].groupby('time.month').mean(), dim='month').compute()\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Plot correlations with climatology"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig,ax = plt.subplots(2,2, figsize=(12,12), sharey=True, sharex=True)\n",
    "precip_clim_corr.plot.imshow(vmin=-0.8, vmax=0.8, cmap='RdBu_r', ax=ax[0,0], add_colorbar=False)\n",
    "ax[0,0].set_title(var+' Climatology & Rain Climatology',  fontsize=18);\n",
    "ax[0,0].set_yticklabels([])\n",
    "ax[0,0].set_ylabel('')\n",
    "ax[0,0].set_xlabel('')\n",
    "ax[0,0].set_xticklabels([])\n",
    "\n",
    "# tavg_clim_corr.plot.imshow(vmin=-0.8, vmax=0.8, cmap='RdBu_r', ax=ax[0,1], add_colorbar=False)\n",
    "# ax[0,1].set_title(var+' Climatology & TAVG Climatology',  fontsize=18);\n",
    "# ax[0,1].set_yticklabels([])\n",
    "# ax[0,1].set_ylabel('')\n",
    "# ax[0,1].set_xlabel('')\n",
    "# ax[0,1].set_xticklabels([])\n",
    "\n",
    "# srad_clim_corr.plot.imshow(vmin=-0.8, vmax=0.8, cmap='RdBu_r', ax=ax[1,0], add_colorbar=False)\n",
    "# ax[1,0].set_title(var+' Climatology & SRAD Climatology',  fontsize=18);\n",
    "# ax[1,0].set_yticklabels([])\n",
    "# ax[1,0].set_ylabel('')\n",
    "# ax[1,0].set_xlabel('')\n",
    "# ax[1,0].set_xticklabels([])\n",
    "\n",
    "# im = kNDVI_clim_corr.plot.imshow(vmin=-0.8, vmax=0.8, cmap='RdBu_r', ax=ax[1,1], add_colorbar=False)\n",
    "# ax[1,1].set_title(var+' Climatology & kNDVI Climatology', fontsize=18)\n",
    "# ax[1,1].set_yticklabels([])\n",
    "# ax[1,1].set_ylabel('')\n",
    "# ax[1,1].set_xlabel('')\n",
    "# ax[1,1].set_xticklabels([])\n",
    "\n",
    "plt.tight_layout();\n",
    "fig.savefig('/g/data/os22/chad_tmp/NEE_modelling/results/figs/'+var+'_climatology_perpixel_climate_correlations.png',\n",
    "            bbox_inches='tight', dpi=300)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Correlations with anomalies"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "kNDVI_anom_corr = xr.corr(ds_anom.chunk(chunks), data['kNDVI_anom'], dim='time').compute()\n",
    "rain_anom_corr = xr.corr(ds_anom.chunk(chunks), data['rain_anom'], dim='time').compute()\n",
    "# vpd_anom_corr = xr.corr(ds_anom.chunk(chunks), data['vpd'], dim='time').compute()\n",
    "srad_anom_corr = xr.corr(ds_anom.chunk(chunks), data['srad_anom'], dim='time').compute()\n",
    "tavg_anom_corr = xr.corr(ds_anom.chunk(chunks), data['tavg_anom'], dim='time').compute()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Plot correlations with anomalies"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "anom_data = [rain_anom_corr,  tavg_anom_corr, srad_anom_corr, kNDVI_anom_corr]\n",
    "clim_vars = ['Rainfall', 'Air Temp.','Solar Rad.', 'kNDVI']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig,axes = plt.subplots(1,4, figsize=(24,7), sharey=True, sharex=True)\n",
    "\n",
    "for ax, ds, clim in zip(axes.ravel(), anom_data, clim_vars):\n",
    "\n",
    "    im = ds.plot.imshow(vmin=-0.8, vmax=0.8, cmap='RdBu_r', ax=ax, add_colorbar=False)\n",
    "    ax.set_title(var+' Anomalies & '+clim+' Anomalies',  fontsize=20);\n",
    "    ax.set_yticklabels([])\n",
    "    ax.set_ylabel('')\n",
    "    ax.set_xlabel('')\n",
    "    ax.set_xticklabels([])\n",
    "\n",
    "plt.tight_layout();\n",
    "fig.savefig('/g/data/os22/chad_tmp/NEE_modelling/results/figs/'+var+'_anomalies_perpixel_climate_correlations.png', bbox_inches='tight')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# get handles and labels for reuse\n",
    "# label_params = ax[1,1].get_legend_handles_labels() \n",
    "cbar = fig.colorbar(im, spacing='uniform', ax=ax, orientation='horizontal', shrink=0.4);\n",
    "\n",
    "# cbar\n",
    "figl, axl = plt.subplots(figsize=(11,4))\n",
    "axl.axis(False)\n",
    "cbar = plt.colorbar(im, spacing='uniform', ax=axl, orientation='horizontal')\n",
    "cbar.ax.tick_params(labelsize=20)\n",
    "cbar.set_label(\"Pearson's Correlation\",size=20)\n",
    "figl.savefig('/g/data/os22/chad_tmp/NEE_modelling/results/figs/correlation_legend.png', bbox_inches='tight', dpi=300)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Correlations with climate vars"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# kndvi_corr = xr.corr(ds.chunk(chunks), data['kNDVI'], dim='time').compute()\n",
    "# rain_corr = xr.corr(ds.chunk(chunks), data['rain'], dim='time').compute()\n",
    "# vpd_corr = xr.corr(ds.chunk(chunks), data['vpd'], dim='time').compute()\n",
    "# srad_corr = xr.corr(ds.chunk(chunks), data['srad'], dim='time').compute()\n",
    "# tavg_corr = xr.corr(ds.chunk(chunks), data['tavg'], dim='time').compute()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Plots by bioregion"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ax2_ylim = -0.11,0.11"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Dictionary to save results \n",
    "results = {}\n",
    "for index, row in gdf.iterrows():\n",
    "\n",
    "    # Generate a polygon mask to keep only data within the polygon\n",
    "    mask = xr_rasterize(gdf.iloc[[index]], var_anom.isel(time=1))\n",
    "    mask['latitude'] = mask.latitude.astype('float32')\n",
    "    mask['longitude'] = mask.longitude.astype('float32')  \n",
    "    mask['latitude'] = np.array([round(i,4) for i in mask.latitude.values])\n",
    "    mask['longitude'] = np.array([round(i,4) for i in mask.longitude.values])\n",
    "    mask = mask.rename({'latitude':'y', 'longitude':'x'})\n",
    "    \n",
    "    # Mask dataset to set pixels outside the polygon to `NaN`\n",
    "    var_anom_region = var_anom.where(mask)\n",
    "    ds_anom_region = ds_anom.where(mask)\n",
    "    \n",
    "    var_clim_mean_region = var_clim_mean.where(mask)\n",
    "    ds_clim_mean_region = ds_clim_mean.where(mask)\n",
    "    \n",
    "    r2 = xr.corr(ds_anom_region, var_anom_region, dim='time').compute()\n",
    "    r2 = r2.mean(['x', 'y'])\n",
    "    print(row['region_name'], r2)\n",
    "\n",
    "    fig,ax=plt.subplots(1,2, figsize=(18,5), gridspec_kw={'width_ratios': [3, 1]})\n",
    "    ax2 = ax[0].twinx()\n",
    "    var_anom_region.mean(['x','y']).rolling(time=3).mean().plot(ax=ax2, label=data_var, c='orange')\n",
    "    ds_anom_region.mean(['x','y']).rolling(time=3).mean().plot(ax=ax[0], label=var)\n",
    "    \n",
    "    ax[0].legend(loc=(0.80,0.925))\n",
    "    ax2.legend(loc=(0.80,0.85))\n",
    "    ax2.set_ylabel(data_var+' Anomaly', fontsize=15)\n",
    "    ax[0].set_xlabel('')\n",
    "    ax2.set_ylim(ax2_ylim)\n",
    "    ax[0].set_ylim(ax_ylim)\n",
    "    ax[0].text(.05, .90, 'r={:.2f}'.format(r2[0]),\n",
    "                transform=ax[0].transAxes, fontsize=15)\n",
    "    ax[0].set_ylabel(var+' Anomalies (gC m\\N{SUPERSCRIPT TWO} m⁻¹)', fontsize=15)\n",
    "    ax[0].tick_params(axis='x', labelsize=14)\n",
    "    ax[0].tick_params(axis='y', labelsize=14)\n",
    "    ax2.tick_params(axis='y', labelsize=14)\n",
    "    ax[1].tick_params(axis='x', labelsize=14)\n",
    "    ax[1].tick_params(axis='y', labelsize=14)\n",
    "\n",
    "    ax[0].axhline(0, c='grey', linestyle='--')\n",
    "\n",
    "    ax3 = ax[1].twinx()\n",
    "\n",
    "    var_clim_mean_region.mean(['x','y']).plot(ax=ax3, label=data_var, c='orange')\n",
    "    ds_clim_mean_region.mean(['x','y']).plot(ax=ax[1], label='NEE')\n",
    "    \n",
    "    ax3.set_ylabel(data_var)\n",
    "    ax3.set_ylabel(data_var, fontsize=15)\n",
    "    ax[1].set_ylabel(var+' (gC m\\N{SUPERSCRIPT TWO} m⁻¹)', fontsize=15)\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\"]) \n",
    "    ax[1].set_xlabel('')\n",
    "    ax3.tick_params(axis='y', labelsize=14)\n",
    "    ax[0].set_title(None)\n",
    "    ax[1].set_title(None)\n",
    "    ax2.set_title(None)\n",
    "    ax3.set_title(None)\n",
    "    \n",
    "    plt.suptitle(row['region_name'], fontsize=18)\n",
    "    plt.tight_layout();\n",
    "    plt.savefig('/g/data/os22/chad_tmp/NEE_modelling/results/figs/'+var+'_'+row['region_name']+'_'+data_var+'_correlations.png')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Variable with highest correlation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "corrs = xr.merge([\n",
    "    np.abs(rain_corr.rename('rain')),\n",
    "    np.abs(vpd_corr.rename('vpd')),\n",
    "    np.abs(srad_corr.rename('srad')),\n",
    "    np.abs(tavg_corr.rename('tavg'))\n",
    "])\n",
    "\n",
    "max_corrs = corrs.to_array(\"variable\").idxmax(\"variable\")\n",
    "max_corrs = xr.where(max_corrs == 'rain', 1, max_corrs)\n",
    "max_corrs = xr.where(max_corrs == 'vpd', 2, max_corrs)\n",
    "max_corrs = xr.where(max_corrs == 'srad', 3, max_corrs)\n",
    "max_corrs = xr.where(max_corrs == 'tavg', 4, max_corrs)\n",
    "\n",
    "max_corrs = max_corrs.astype(np.float32)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "rain_corrs = xr.merge([\n",
    "    np.abs(precip_corr.rename('precip_anom')),\n",
    "    np.abs(precip_3_corr.rename('precip_3_anom')),\n",
    "    np.abs(precip_6_corr.rename('precip_6_anom')),\n",
    "    np.abs(precip_12_corr.rename('precip_12_anom'))\n",
    "])\n",
    "\n",
    "rain_max_corrs = rain_corrs.to_array(\"variable\").idxmax(\"variable\")\n",
    "\n",
    "rain_max_corrs = xr.where(rain_max_corrs == 'precip_anom', 1, rain_max_corrs)\n",
    "rain_max_corrs = xr.where(rain_max_corrs == 'precip_3_anom', 2, rain_max_corrs)\n",
    "rain_max_corrs = xr.where(rain_max_corrs == 'precip_6_anom', 3, rain_max_corrs)\n",
    "rain_max_corrs = xr.where(rain_max_corrs == 'precip_12_anom', 4, rain_max_corrs)\n",
    "\n",
    "rain_max_corrs = rain_max_corrs.astype(np.float32)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Plots"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig,ax = plt.subplots(1,1, figsize=(10,7))\n",
    "im = max_corrs.plot.imshow(add_colorbar=False, ax=ax)\n",
    "cbar = fig.colorbar(im, spacing='uniform', ax=ax, orientation='vertical', shrink=0.4)\n",
    "cbar.set_ticks([1,2,3,4])\n",
    "cbar.set_ticklabels(['Rain', 'VPD', 'SRAD', 'TAVG'], fontsize=10)\n",
    "plt.title('Climate Variable with Maximum Absolute Correlation with '+var);"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig,ax = plt.subplots(1,1, figsize=(10,7))\n",
    "im = rain_max_corrs.plot.imshow(add_colorbar=False, ax=ax)\n",
    "cbar = fig.colorbar(im, spacing='uniform', ax=ax, orientation='vertical', shrink=0.4)\n",
    "cbar.set_ticks([1,2,3,4])\n",
    "cbar.set_ticklabels(['Rain', 'Rain-3', 'Rain-6', 'Rain-12'], fontsize=10)\n",
    "plt.title('Rainfall Variable with Maximum Absolute Correlation with '+var);"
   ]
  },
  {
   "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
}