mhw_threshold.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "57a62ba4-d4e1-446d-bbe3-446bc36a57db",
   "metadata": {},
   "outputs": [],
   "source": [
    "import xarray as xr\n",
    "from dask.distributed import Client\n",
    "import dask.array as da\n",
    "from dask import delayed\n",
    "import numpy as np\n",
    "import os"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3f4864c4-9510-461b-be1f-69fe7c9d2f14",
   "metadata": {},
   "outputs": [],
   "source": [
    "client = Client()\n",
    "print(\"Dask dashboard available at:\", client.dashboard_link)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e6652c61-058f-4467-a0bb-7ba699204ddb",
   "metadata": {},
   "outputs": [],
   "source": [
    "coords_to_drop =['st_edges_ocean','nv','st_ocean']\n",
    "vars_to_drop =['Time_bounds','average_DT','average_T1','average_T2','st_ocean']\n",
    "\n",
    "# preprocesser to drop unwanted variables\n",
    "def drop_stuff(ds, coords_to_drop,vars_to_drop):\n",
    "    \"\"\"\n",
    "    Preprocessor function to drop specified coordinates and variables from a dataset loaded via xr.open_mfdataset\n",
    "\n",
    "    Parameters:\n",
    "        ds (xarray.Dataset): The dataset from which coordinates & variables are to be dropped.\n",
    "        coords_to_drop (list of str): List of coordinate names to drop.\n",
    "        vars_to_drop(list of str): List of variable names to drop\n",
    "\n",
    "    Returns:\n",
    "        xarray.Dataset: Dataset with specified coordinates and variables dropped.\n",
    "    \"\"\"\n",
    "    # Drop coordinates if they are in the dataset\n",
    "    ds = ds.drop_vars(coords_to_drop, errors='ignore')\n",
    "    ds = ds.drop_vars(vars_to_drop, errors='ignore')\n",
    "    return ds"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fb2b349c-3bd4-435f-90c9-8c4c4ac6891a",
   "metadata": {},
   "outputs": [],
   "source": [
    "def process_threshold(ds, time_dim, start, end, variable, period):\n",
    "    # Rechunk the data along the time dimension\n",
    "    data = getattr(ds, variable).sel(**{time_dim: slice(start, end)}).chunk({'Time': -1, 'xt_ocean': 50, 'yt_ocean': 50}).persist()\n",
    "\n",
    "    # Debug: Print the shape and chunking of the data\n",
    "    print(f\"Processing period {period}, data shape: {data.shape}, chunks: {data.chunks}\")\n",
    "\n",
    "    # Calculate the 90th percentile\n",
    "    percentile_90 = data.groupby('Time.month').quantile(0.9, dim='Time').compute()\n",
    "\n",
    "    # Convert the result back to an xarray DataArray\n",
    "    percentile_90_da = xr.DataArray(\n",
    "        percentile_90,\n",
    "         coords={\n",
    "            'xt_ocean': data.coords['xt_ocean'],\n",
    "            'yt_ocean': data.coords['yt_ocean'],\n",
    "            'month': np.arange(1, 13)\n",
    "        },\n",
    "        dims=['yt_ocean', 'xt_ocean', 'month']\n",
    "    )\n",
    "\n",
    "    return percentile_90_da"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "109d50f9-2665-4896-832b-9435ed3a9da0",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define the periods for processing\n",
    "GWL_periods = {\n",
    "    'current': ('1995-01-01', '2014-12-31'),\n",
    "    'GW1p2': ('2001-01-01', '2020-12-31'),\n",
    "    'GW1p5': ('2015-01-01', '2034-12-31'),\n",
    "    'GW2p0': ('2030-01-01', '2049-12-31'),\n",
    "    'GW3p0': ('2053-01-01', '2072-12-31'),\n",
    "    'GW4p0': ('2074-01-01', '2093-12-31')\n",
    "}\n",
    "\n",
    "# Directory paths for SST\n",
    "dir1_new = '/g/data/fp2/OFAM3/jra55_historical.1/surface/'\n",
    "dir2_new = '/g/data/fp2/OFAM3/jra55_rcp8p5/surface/'\n",
    "\n",
    "# Load datasets with chunking\n",
    "dsst1 = xr.open_mfdataset(dir1_new + 'ocean_temp_sfc_*', parallel=True, \n",
    "                                        preprocess = lambda x: drop_stuff(x, \n",
    "                                         coords_to_drop, \n",
    "                                         vars_to_drop)).squeeze()\n",
    "dsst2 = xr.open_mfdataset(dir2_new + 'ocean_temp_sfc*.nc', parallel=True, \n",
    "                                        preprocess = lambda x: drop_stuff(x, \n",
    "                                         coords_to_drop, \n",
    "                                         vars_to_drop)).squeeze()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6051a369-37c0-4d25-a3bc-16d72d2752d7",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Process SST\n",
    "for period, (start, end) in GWL_periods.items():\n",
    "    if period == 'current':\n",
    "        percentile_90th_da = process_threshold(dsst1, 'Time', start, end, 'temp', period)\n",
    "    else:\n",
    "        percentile_90th_da = process_threshold(dsst2, 'Time', start, end, 'temp', period)\n",
    "    file_path_90th = f'/g/data/ia39/ncra/ocean/peacey/sst_percentile_monthly__{period}.nc'\n",
    "    # Save 90th percentile as netCDF\n",
    "    percentile_90_da.to_netcdf(file_path_90th, compute=True)\n",
    "\n",
    "client.close()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python [conda env:analysis3]",
   "language": "python",
   "name": "conda-env-analysis3-py"
  },
  "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
}