Fig5_direction_sig_trend_mk.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "90c313a9-54a2-4f66-8465-31cee2fa875a",
   "metadata": {},
   "source": [
    "# Calculate Theil-Sen Trend (Mann Kendall Significance) for Annual Time Series for Observational Product and CORDEX Simulations"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "03bf2851-f5b4-4ae3-a777-d2353151e0aa",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Import modules\n",
    "%matplotlib inline\n",
    "%run /g/data/w40/ri9247/code/aus_precip_benchmarking/master_functions_bmf.ipynb\n",
    "import pymannkendall as mk\n",
    "from scipy import stats\n",
    "import fnmatch\n",
    "import xarray as xr\n",
    "import pandas as pd\n",
    "import matplotlib.pyplot as plt\n",
    "import matplotlib as mpl\n",
    "import matplotlib.ticker as mticker\n",
    "import matplotlib.gridspec as gridspec"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "06f8097a-5103-4b70-bb11-3f26e54159b0",
   "metadata": {},
   "source": [
    "Citation for using mann kendall trend:\n",
    "    Hussain et al., (2019). pyMannKendall: a python package for non parametric Mann Kendall family of trend tests. Journal of Open Source Software, 4(39), 1556, https://doi.org/10.21105/joss.01556"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f2da9e47-1dfe-4597-b096-78ee7c99be97",
   "metadata": {},
   "source": [
    "## Define Spatiotemporal Boundaries"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f0fc0b51-5e0a-4c23-836e-03bf34ad34dc",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Define region of interest (lat and lon boundaries for all of Australia)\n",
    "lat_slice = slice(-44.5,-10)\n",
    "lon_slice = slice(112,156.25)\n",
    "time_slice = slice(\"1976-01-01\", \"2005-12-30\")\n",
    "\n",
    "# Define temporal scale for SPI variable (1 for 3-months, 2 for 6-months, 3 for 12-months); use \"None\" if not using the SPI variable\n",
    "iscale = None\n",
    "\n",
    "# Define season as a list of month numbers. If not wanting a seasonal breakdown, use None.\n",
    "# Suggestions: Wet 6-month: [10,11,12,1,2,3]; Dry 6-month: [4,5,6,7,8,9]; Wet 3-month: [12,1,2]; Dry 3-month: [6,7,8]\n",
    "season = None \n",
    "season_name = 'Annual'\n",
    "\n",
    "# Path to mask datasets\n",
    "qc_mask_ds = xr.open_dataset('/g/data/w40/ri9247/CORDEX-Australasia/data/obs/AUS-44i_grid/no_indices/mask/AUS-44i_combined_quality_mask_no_oceans.nc')\n",
    "ipcc_mask_ds = None #xr.open_dataset('/g/data/w40/ri9247/CORDEX-Australasia/data/obs/AUS-44i_grid/no_indices/mask/AUS44i_IPCC_AUS_regions_mask.nc')\n",
    "\n",
    "# Extract masks over Australia\n",
    "qc_mask = qc_mask_ds.station_mask.sel(lat=lat_slice, lon=lon_slice)\n",
    "ipcc_mask = None #ipcc_mask_ds.mask"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "942eb0bd-2f3e-4e1d-98bf-c0c637df9031",
   "metadata": {},
   "source": [
    "## Define Keywords for Data and Data Paths"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "92cf5ae8-a297-414e-bf33-17ad8a4b9487",
   "metadata": {},
   "source": [
    "Keywords needed are bracketed {}. <br><br>\n",
    "**RCM** database is organized following: <br>\n",
    "**Climpact Indices:** <br>\n",
    "parent_directory/{time_period}/{index-keyword}/{variable}_{time_average} _dataset_file.nc \n",
    "<br><br>\n",
    "**Observations** database is organized as follows: <br>\n",
    "obs_parent_directory/grid_type}/climpact/{variable}/{variable} _{time_average}_agcd_historical_v1_1950-2020.nc <br><br>\n",
    "**Keyword Options:** <br>\n",
    "grid_type: 'AUS-44i_grid', 'native_grid', 'one_degree_grid' <br>\n",
    "time_period: 'historical', 'rcp85' <br>\n",
    "variable: See list of Climpact Indices at: https://climpact-sci.org/indices; must be all-lowercase <br>\n",
    "time_average: 'ANN', 'MON'"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0f7080e9-cb94-4018-98ff-b8aa1beaf8e1",
   "metadata": {},
   "source": [
    "### Using Climpact Indices"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "67a610c3-d9bf-4442-91c1-4f3ee973e37b",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Define keywords to be used with f-strings to build the path to datasets\n",
    "time_period = 'historical'\n",
    "variable = 'prcptot'\n",
    "time_average = 'ANN'\n",
    "grid_type = 'AUS-44i_grid'\n",
    "\n",
    "# Define paths to data\n",
    "model_master_path = '/g/data/ks32/CLEX_Data/CORDEX_Australasia_Indices/v1-0/'\n",
    "obs_master_path = '/g/data/w40/ri9247/CORDEX-Australasia/data/obs/'\n",
    "\n",
    "model_data_path = model_master_path + f'{time_period}/{variable}/'\n",
    "agcd_data_path = obs_master_path + f'{grid_type}/climpact/{variable}/{variable}_{time_average}_agcd_historical_v1_1950-2020.nc'"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "50df002e-6ff2-4278-a950-070a4187845c",
   "metadata": {},
   "source": [
    "## Get Data File Paths"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "60519bd6-4ab1-436e-bd86-24a7b97879b8",
   "metadata": {},
   "source": [
    "### Models"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7a5c8625-f7b1-4a19-b3ff-4f4a3bba2476",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Get all model paths in full ensemble (this will be stored in a Pandas DataFrame)\n",
    "model_paths = get_model_files(model_data_path, time_average)\n",
    "\n",
    "# Expand the column with to see the full paths stored in the Pandas Dataframe\n",
    "pd.set_option('display.max_colwidth', None)\n",
    "\n",
    "# Option to check that correct file paths are selected\n",
    "model_paths"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9adff227-7f5a-4262-800e-0faf76996cf9",
   "metadata": {
    "tags": []
   },
   "source": [
    "## Get Weighted Spatial Average at each Year for the time series and Calculate Mann-Kendall Metrics based on time series data"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3f5247f8-3f7a-4791-94da-8c140dc76e0d",
   "metadata": {},
   "source": [
    "### Obs - AGCD"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "63d91367-8760-4316-8690-4454ad310d0b",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Define Data Type (in case unit conversion needed)\n",
    "data_type = 'obs'\n",
    "\n",
    "# Get Time Series\n",
    "obs_ts = get_annual_average(agcd_data_path, variable, time_slice, lat_slice, lon_slice, data_type, iscale, qc_mask)\n",
    "\n",
    "# Calcualte Trend and Significance\n",
    "obs_ts_mk = mk.original_test(obs_ts)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5b0d6919-dd42-4b5b-8225-bc0a36f4ea2e",
   "metadata": {},
   "source": [
    "### Models"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a2c3bc76-d319-48f5-94e6-1bbc4589d420",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Define a data_type (daily and monthly model data need to be converted to proper units)\n",
    "data_type = 'model_climpact'\n",
    "\n",
    "# Intialize empty dictionaries to store model timeseries and Mann Kendall metrics \n",
    "model_ann_ts_dict = {}\n",
    "model_ts_mk_dict = {}\n",
    "\n",
    "# Initialize dataframe to store theil-sen trend value for sorting plots\n",
    "model_ts_master_df = pd.DataFrame(columns=[\"theil-sen-slope\", \"dataset_name\"])\n",
    "\n",
    "# Loop through RCM simulations; store in a dictionary\n",
    "for i, row in model_paths.iterrows():\n",
    "    \n",
    "    # Get weighted spatial annual average for time series\n",
    "    model_ts = get_annual_average(model_paths.iloc[i,2], variable, time_slice, lat_slice, lon_slice, data_type, iscale, qc_mask)\n",
    "    \n",
    "    # Calculate Mann Kendall Metrics based on time series\n",
    "    model_mk_ts = mk.original_test(model_ts)\n",
    "    \n",
    "    # Get slope value\n",
    "    ts_slope = model_mk_ts[7]\n",
    "    \n",
    "    # Create a dataframe to store the slope value and corresponding dataset\n",
    "    model_ts_df = pd.DataFrame({'theil-sen-slope': ts_slope, 'dataset_name': f'{model_paths.iloc[i,1]}   {model_paths.iloc[i,0]}'}, index=[0])   \n",
    "    \n",
    "    # Append slope dataframe to the master dataframe \n",
    "    model_ts_master_df = pd.concat([model_ts_master_df, model_ts_df])\n",
    "    \n",
    "    # Add to dictionaries\n",
    "    model_ann_ts_dict[f'{model_paths.iloc[i,1]}   {model_paths.iloc[i,0]}'] = model_ts\n",
    "    model_ts_mk_dict[f'{model_paths.iloc[i,1]}   {model_paths.iloc[i,0]}'] = model_mk_ts\n",
    "    \n",
    "# Sort dataframe from largest to smallest for plotting later    \n",
    "model_ts_master_df = model_ts_master_df.sort_values(\"theil-sen-slope\", ascending=False)\n",
    "print(model_ts_mk_dict.keys())"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7ae2d832-51e0-45b0-86fc-0aed7aff2551",
   "metadata": {},
   "source": [
    "## Calculate Weighted Spatial Average for Sorting the Plots"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "141c25d1-2db8-4edc-8faf-dcdf7c225855",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Intialize an empty dictionary to store model climatologies as DataArrays\n",
    "climatology_model_dict = {}\n",
    "\n",
    "# Loop through RCM simulations and calculate the model bias; store in a dictionary\n",
    "for i, row in model_paths.iterrows():\n",
    "    climatology_xr = get_climatology(model_paths.iloc[i,2], variable, time_slice, lat_slice, lon_slice, data_type, season, iscale, qc_mask)\n",
    "    climatology_model_dict[f'{model_paths.iloc[i,1]}   {model_paths.iloc[i,0]}'] = climatology_xr\n",
    "\n",
    "# Option to print dictionary keys\n",
    "print(climatology_model_dict.keys())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ea11be01-db11-427b-9df9-c9d21f49a008",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Set up empty data frame to store values\n",
    "model_weighted_avg_master_df = pd.DataFrame(columns=[\"weighted_average\", \"dataset_name\"])\n",
    "\n",
    "# Loop through Dictionary of climatologies and calculate spatial weighted average\n",
    "for model_name in climatology_model_dict.keys():\n",
    "    model_weighted_avg = get_weighted_spatial_average_from_data(climatology_model_dict[model_name]).round(2)\n",
    "        \n",
    "    # Convert the resulting xarray to a pandas df\n",
    "    model_weighted_avg_df = pd.DataFrame({'weighted_average': model_weighted_avg.item(0), 'dataset_name': f'{model_name}'}, index=[0])   \n",
    "    \n",
    "    # Append the master time series \n",
    "    model_weighted_avg_master_df = pd.concat([model_weighted_avg_master_df, model_weighted_avg_df])\n",
    "    \n",
    "model_weighted_avg_master_df = model_weighted_avg_master_df.sort_values(\"weighted_average\", ascending=False)\n",
    "\n",
    "# Option to print weighted spatial average values\n",
    "print(model_weighted_avg_master_df)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d87ac4fa-02a1-4c20-8a68-1f169306ae81",
   "metadata": {},
   "source": [
    "## Plot Time series and trend lines"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "32490bf8-0132-4a2e-9d48-e8ad0fc9b7d3",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Define function to plot trend line color based on Mann Kendall significance\n",
    "def get_trend_line_color(mk_trend):\n",
    "    if mk_trend == \"increasing\":\n",
    "        color = 'darkturquoise'\n",
    "    elif mk_trend == \"decreasing\":\n",
    "        color = 'darkgoldenrod'\n",
    "    elif mk_trend == \"no trend\":\n",
    "        color = 'blueviolet'\n",
    "    return color"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8646813a-f5e7-4b55-9365-2722b2199cb6",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Option to test that color output is correct\n",
    "color_test = get_trend_line_color(obs_ts_mk[0])\n",
    "color_test"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9a4d5dd2-c9f9-4ed7-904b-02f4ea19f47c",
   "metadata": {},
   "source": [
    "## Plot Time Series and Trend Lines"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ed2a1316-a94e-45a5-aaef-3d226441f45a",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Define figure size\n",
    "fig = plt.figure(figsize=(17,18))\n",
    "fig.suptitle(\"Annual Average Total Precipitation (1976-2005)\", fontsize=16, y=0.92)\n",
    "\n",
    "# Setup axes for all subplots\n",
    "gs = gridspec.GridSpec(7,4,width_ratios=[1,1,1,1], height_ratios=[1,1,1,1,1,1,1])\n",
    "\n",
    "row_max = 6\n",
    "col_max = 3\n",
    "row = 1\n",
    "col = 0\n",
    "\n",
    "axs = {}\n",
    "\n",
    "# Define AGCD as only dataset on the top row\n",
    "axs['AGCD'] = fig.add_subplot(gs[0,0])\n",
    "\n",
    "# Set up Axes labels (this loops through the sorted Pandas DF to assign axes positions based on wetness)\n",
    "for i, df_row_mod in model_weighted_avg_master_df.iterrows():\n",
    "    axs[f'{df_row_mod[1]}'] = fig.add_subplot(gs[row,col])\n",
    "    \n",
    "    if col == col_max:\n",
    "        row = row + 1\n",
    "        col = 0 \n",
    "    else:\n",
    "        col = col + 1\n",
    "        \n",
    "# Define plots where we want to plot x and/or y axis labels\n",
    "y_label = [0,1,5,9,13,17]\n",
    "x_label = [22,23,24]\n",
    "both_label = [21]\n",
    "\n",
    "# Add subplot titles\n",
    "for name, ax in axs.items():\n",
    "    ax.set_title(name, fontsize=13)\n",
    "    ax.set_xlim(1976, 2005)\n",
    "    ax.set_ylim(0,1200)\n",
    "    \n",
    "# Set axis ticks for only left and bottom figures    \n",
    "for i, ax in enumerate(fig.axes):\n",
    "    ax.set_xticklabels([])\n",
    "    ax.set_yticklabels([])\n",
    "    if i in y_label:\n",
    "        ax.set_yticklabels([0,200,400,600,800,1000,1200])\n",
    "        ax.set_ylabel(\"mm/year\", fontsize=12)\n",
    "    elif i in x_label:\n",
    "        ax.set_xticklabels([1975, 1980, 1985, 1990, 1995, 2000, 2005])\n",
    "        ax.set_xlabel(\"year\", fontsize=11)\n",
    "    elif i in both_label:\n",
    "        ax.set_yticklabels([0,200,400,600,800,1000,1200])\n",
    "        ax.set_xticklabels([1975, 1980, 1985, 1990, 1995, 2000, 2005])\n",
    "        ax.set_xlabel(\"year\", fontsize=11)\n",
    "        ax.set_ylabel(\"mm/year\", fontsize=11)\n",
    "    else:\n",
    "        continue\n",
    "\n",
    "axs['AGCD'].plot(obs_ts.year, obs_ts, color='black', linestyle='-')\n",
    "axs['AGCD'].plot(obs_ts.year, obs_ts_mk[7] * (obs_ts.year - 1976) + obs_ts_mk[8], color=get_trend_line_color(obs_ts_mk[0]), linestyle='-')\n",
    "axs['AGCD'].text(1978, 60, 'Slope: %.2f' % (obs_ts_mk[7]), fontsize=12, color='black')\n",
    "axs['AGCD'].text(1993, 60, f'{(obs_ts_mk[0])}', fontsize=12, color=get_trend_line_color(obs_ts_mk[0]))\n",
    "    \n",
    "# Models\n",
    "for model_name in model_ann_ts_dict.keys():\n",
    "    axs[model_name].plot(model_ann_ts_dict[model_name].year, model_ann_ts_dict[model_name], color='black', linestyle='-')\n",
    "    axs[model_name].plot(model_ann_ts_dict[model_name].year, model_ts_mk_dict[model_name][7] * (model_ann_ts_dict[model_name].year - 1976) + model_ts_mk_dict[model_name][8], \n",
    "                         color=get_trend_line_color(model_ts_mk_dict[model_name][0]), linestyle='-')\n",
    "    axs[model_name].text(1978, 60, 'Slope: %.2f' % (model_ts_mk_dict[model_name][7]), fontsize=12, color='black')\n",
    "    axs[model_name].text(1993, 60, f'{(model_ts_mk_dict[model_name][0])}', fontsize=12, color=get_trend_line_color(model_ts_mk_dict[model_name][0]))"
   ]
  }
 ],
 "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.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}