0d_prepare_diagnostics.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "38130fe6-9d8c-4edc-8b37-2abe0782b8f7",
   "metadata": {},
   "source": [
    "# Prepare atmospheric and ocean data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "1c233f26-1fbe-4e0a-91cc-19323aa33a71",
   "metadata": {},
   "outputs": [],
   "source": [
    "from dask.distributed import Client,LocalCluster\n",
    "from dask_jobqueue import PBSCluster"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "85975a6c-38e4-401e-a515-430e23698f03",
   "metadata": {},
   "outputs": [],
   "source": [
    "import xarray as xr\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "\n",
    "import os"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "8cfd113e-ac71-48b0-98fe-e04d75390608",
   "metadata": {},
   "outputs": [],
   "source": [
    "import functions as fn"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "e6852ea2-f6c0-4575-8e60-3accf788f1bf",
   "metadata": {},
   "outputs": [],
   "source": [
    "def anomalise_and_process(\n",
    "    load,\n",
    "    dataset,\n",
    "    filepath,\n",
    "    data_root_path,\n",
    "    var_name,\n",
    "    years,\n",
    "    subset_region,\n",
    "    subset_level,\n",
    "    preprocess_func,\n",
    "    mfdataset_chunks,\n",
    "    anom_freq,\n",
    "    new_var_name,\n",
    "    long_name,\n",
    "    units,\n",
    "    chunks\n",
    "):\n",
    "    \"\"\"\n",
    "    Load or process and write data.\n",
    "    \n",
    "    Data is processed by loading from the desired dataset a variable over a specified period.\n",
    "    The data are anomalised by removing the anom_freq means computed over the entire\n",
    "    period.\n",
    "    \n",
    "    load : Bool. Whether to load or process.\n",
    "    dataset : \"era5\" or \"hadisst\"\n",
    "    filepath : file to load saved data\n",
    "    data_root_path : path to raw era5 or hadisst data on Gadi.\n",
    "    var_name : string. name of era5 variable (ignored if other dataset).\n",
    "    years : range. Period to process.\n",
    "    subset_region : list. Region to subset over. See functions.open_era_data\n",
    "    subset_level : list. Pressure levels to select. See functions.open_era_data\n",
    "    preprocess_func : Preprocess using xr.open_mfdataset if desired. See functions.open_era_data\n",
    "    mfdataset_chunks : Chunks to use for preprocessing in xr.open_mfdataset. See functions.open_era_data\n",
    "    anom_freq : string. Over what frequency to anomalise, e.g. month or dayofyear\n",
    "    new_var_name : desired name of variable.\n",
    "    long_name : long name for xarray Dataset attributes.\n",
    "    units : units for xarray Dataset attributes.\n",
    "    chunks : dict. Specify chunks before writing to zarr.\n",
    "    \"\"\"\n",
    "    if load:\n",
    "        return xr.open_zarr(filepath, consolidated=True)\n",
    "    else:\n",
    "        print('Opening files...')\n",
    "        if dataset == 'era5':\n",
    "            ds = fn.open_era_data(\n",
    "                root_path=data_root_path,\n",
    "                variable=var_name,\n",
    "                years=years,\n",
    "                subset_region=subset_region,\n",
    "                subset_level=subset_level,\n",
    "                preprocess_func=preprocess_func,\n",
    "                mfdataset_chunks=mfdataset_chunks\n",
    "            )\n",
    "        elif dataset == 'hadisst':\n",
    "            ds = xr.open_zarr(data_root_path, consolidated=True)\n",
    "            ds = ds.sel(time=slice(str(years[0]), str(years[-1])))\n",
    "            ds['time'] = pd.date_range(str(years[0]), str(years[-1])+'-12-01', freq='1MS')\n",
    "            ds = xr.where(ds > -2, ds, np.nan)\n",
    "        else:\n",
    "            raise ValueError(\"Incorrect dataset.\")\n",
    "            \n",
    "        ds = ds.rename({\n",
    "            'latitude': 'lat',\n",
    "            'longitude': 'lon',\n",
    "            var_name: new_var_name\n",
    "        })\n",
    "        da = ds[new_var_name]\n",
    "        da = da.chunk(chunks)\n",
    "            \n",
    "        print('Anomalising...')\n",
    "        da_anoms = da.groupby('time.' + anom_freq).apply(lambda x: x - x.mean('time'))\n",
    "            \n",
    "        print('Attributes and writing...')\n",
    "        da_anoms = da_anoms.assign_attrs({\n",
    "            'long_name': long_name,\n",
    "            'short_name': new_var_name,\n",
    "            'units': units\n",
    "        })\n",
    "            \n",
    "        ds_anoms = da_anoms.to_dataset(name=new_var_name + '_anom')\n",
    "        encoding = {new_var_name + '_anom': {'dtype': 'float32'}}\n",
    "        ds_anoms.to_zarr(\n",
    "            filepath,\n",
    "            mode='w',\n",
    "            consolidated=True,\n",
    "            encoding=encoding\n",
    "        )\n",
    "        \n",
    "        return ds_anoms"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "94bd3c58-ac6e-46b7-941e-b22b89260c44",
   "metadata": {},
   "outputs": [],
   "source": [
    "data_fp = '/g/data/w42/dr6273/work/data/'\n",
    "years = range(1959, 2022)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "aaa293ee-5eb1-4402-8f31-c36025c3d9db",
   "metadata": {},
   "source": [
    "# Daily data on extended Australia region"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "31d46b5f-7785-440c-b923-8ce392bdc37b",
   "metadata": {},
   "source": [
    "### To process the hourly data we need significant resources - 3 full nodes."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "e25ead5e-0743-4b11-ba46-22d7eccdbb8c",
   "metadata": {},
   "outputs": [],
   "source": [
    "# One node on Gadi has 48 cores - try and use up a full node before going to multiple nodes (jobs)\n",
    "\n",
    "walltime = '00:30:00'\n",
    "cores = 48\n",
    "memory = str(4 * cores) + 'GB'\n",
    "\n",
    "cluster = PBSCluster(walltime=str(walltime), cores=cores, memory=str(memory), processes=cores,\n",
    "                     job_extra_directives=['-q normal',\n",
    "                                           '-P w42',\n",
    "                                           '-l ncpus='+str(cores),\n",
    "                                           '-l mem='+str(memory),\n",
    "                                           '-l storage=gdata/w42+gdata/rt52+gdata/xv83'],\n",
    "                     local_directory='$TMPDIR',\n",
    "                     job_directives_skip=[\"select\"])\n",
    "                     # python=os.environ[\"DASK_PYTHON\"])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "aef43f7a-0ce1-463e-a55e-b7ef3ac8d8f3",
   "metadata": {},
   "outputs": [],
   "source": [
    "cluster.scale(jobs=3)\n",
    "client = Client(cluster)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "991e01c8-eeeb-48c3-b146-6b679ffc9857",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "    <div style=\"width: 24px; height: 24px; background-color: #e1e1e1; border: 3px solid #9D9D9D; border-radius: 5px; position: absolute;\"> </div>\n",
       "    <div style=\"margin-left: 48px;\">\n",
       "        <h3 style=\"margin-bottom: 0px;\">Client</h3>\n",
       "        <p style=\"color: #9D9D9D; margin-bottom: 0px;\">Client-64aeac34-d28e-11ed-a628-0000018afe80</p>\n",
       "        <table style=\"width: 100%; text-align: left;\">\n",
       "\n",
       "        <tr>\n",
       "        \n",
       "            <td style=\"text-align: left;\"><strong>Connection method:</strong> Cluster object</td>\n",
       "            <td style=\"text-align: left;\"><strong>Cluster type:</strong> dask_jobqueue.PBSCluster</td>\n",
       "        \n",
       "        </tr>\n",
       "\n",
       "        \n",
       "            <tr>\n",
       "                <td style=\"text-align: left;\">\n",
       "                    <strong>Dashboard: </strong> <a href=\"http://10.6.44.33:8787/status\" target=\"_blank\">http://10.6.44.33:8787/status</a>\n",
       "                </td>\n",
       "                <td style=\"text-align: left;\"></td>\n",
       "            </tr>\n",
       "        \n",
       "\n",
       "        </table>\n",
       "\n",
       "        \n",
       "            <button style=\"margin-bottom: 12px;\" data-commandlinker-command=\"dask:populate-and-launch-layout\" data-commandlinker-args='{\"url\": \"http://10.6.44.33:8787/status\" }'>\n",
       "                Launch dashboard in JupyterLab\n",
       "            </button>\n",
       "        \n",
       "\n",
       "        \n",
       "            <details>\n",
       "            <summary style=\"margin-bottom: 20px;\"><h3 style=\"display: inline;\">Cluster Info</h3></summary>\n",
       "            <div class=\"jp-RenderedHTMLCommon jp-RenderedHTML jp-mod-trusted jp-OutputArea-output\">\n",
       "    <div style=\"width: 24px; height: 24px; background-color: #e1e1e1; border: 3px solid #9D9D9D; border-radius: 5px; position: absolute;\">\n",
       "    </div>\n",
       "    <div style=\"margin-left: 48px;\">\n",
       "        <h3 style=\"margin-bottom: 0px; margin-top: 0px;\">PBSCluster</h3>\n",
       "        <p style=\"color: #9D9D9D; margin-bottom: 0px;\">6669c5ca</p>\n",
       "        <table style=\"width: 100%; text-align: left;\">\n",
       "            <tr>\n",
       "                <td style=\"text-align: left;\">\n",
       "                    <strong>Dashboard:</strong> <a href=\"http://10.6.44.33:8787/status\" target=\"_blank\">http://10.6.44.33:8787/status</a>\n",
       "                </td>\n",
       "                <td style=\"text-align: left;\">\n",
       "                    <strong>Workers:</strong> 0\n",
       "                </td>\n",
       "            </tr>\n",
       "            <tr>\n",
       "                <td style=\"text-align: left;\">\n",
       "                    <strong>Total threads:</strong> 0\n",
       "                </td>\n",
       "                <td style=\"text-align: left;\">\n",
       "                    <strong>Total memory:</strong> 0 B\n",
       "                </td>\n",
       "            </tr>\n",
       "            \n",
       "        </table>\n",
       "\n",
       "        <details>\n",
       "            <summary style=\"margin-bottom: 20px;\">\n",
       "                <h3 style=\"display: inline;\">Scheduler Info</h3>\n",
       "            </summary>\n",
       "\n",
       "            <div style=\"\">\n",
       "    <div>\n",
       "        <div style=\"width: 24px; height: 24px; background-color: #FFF7E5; border: 3px solid #FF6132; border-radius: 5px; position: absolute;\"> </div>\n",
       "        <div style=\"margin-left: 48px;\">\n",
       "            <h3 style=\"margin-bottom: 0px;\">Scheduler</h3>\n",
       "            <p style=\"color: #9D9D9D; margin-bottom: 0px;\">Scheduler-e9fe845a-b778-450c-9d9d-a7b2b45285db</p>\n",
       "            <table style=\"width: 100%; text-align: left;\">\n",
       "                <tr>\n",
       "                    <td style=\"text-align: left;\">\n",
       "                        <strong>Comm:</strong> tcp://10.6.44.33:37627\n",
       "                    </td>\n",
       "                    <td style=\"text-align: left;\">\n",
       "                        <strong>Workers:</strong> 0\n",
       "                    </td>\n",
       "                </tr>\n",
       "                <tr>\n",
       "                    <td style=\"text-align: left;\">\n",
       "                        <strong>Dashboard:</strong> <a href=\"http://10.6.44.33:8787/status\" target=\"_blank\">http://10.6.44.33:8787/status</a>\n",
       "                    </td>\n",
       "                    <td style=\"text-align: left;\">\n",
       "                        <strong>Total threads:</strong> 0\n",
       "                    </td>\n",
       "                </tr>\n",
       "                <tr>\n",
       "                    <td style=\"text-align: left;\">\n",
       "                        <strong>Started:</strong> 1 minute ago\n",
       "                    </td>\n",
       "                    <td style=\"text-align: left;\">\n",
       "                        <strong>Total memory:</strong> 0 B\n",
       "                    </td>\n",
       "                </tr>\n",
       "            </table>\n",
       "        </div>\n",
       "    </div>\n",
       "\n",
       "    <details style=\"margin-left: 48px;\">\n",
       "        <summary style=\"margin-bottom: 20px;\">\n",
       "            <h3 style=\"display: inline;\">Workers</h3>\n",
       "        </summary>\n",
       "\n",
       "        \n",
       "\n",
       "    </details>\n",
       "</div>\n",
       "\n",
       "        </details>\n",
       "    </div>\n",
       "</div>\n",
       "            </details>\n",
       "        \n",
       "\n",
       "    </div>\n",
       "</div>"
      ],
      "text/plain": [
       "<Client: 'tcp://10.6.44.33:37627' processes=0 threads=0, memory=0 B>"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "client"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "64ceb0c6-7c01-4aa3-9d75-d54fda53c864",
   "metadata": {},
   "outputs": [],
   "source": [
    "extended_aus = fn.get_extended_AUS_boundary()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5e470f65-d86e-4c84-b5c7-913a180d6a6a",
   "metadata": {},
   "source": [
    "### Note:\n",
    "\n",
    "I couldn't process hourly data to daily using the function above.\n",
    "So, try loading hourly data beforehand, using preprocess to select subregion, level and calculate averages.\n",
    "Write these daily data to file, then load using this function to anomalise as before."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "678b8b8c-5008-46d8-ba35-558c433ea5c2",
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_files(file_path, var, years):\n",
    "    \"\"\"\n",
    "    Get list of files\n",
    "    \"\"\"\n",
    "    fp_list = []\n",
    "    for year in years:\n",
    "        fp_dir = file_path+var+'/'+str(year)+'/'\n",
    "        for fp in sorted(os.listdir(fp_dir)):\n",
    "            fp_list.append(fp_dir+fp)\n",
    "    return fp_list"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "e0b8cf82-6f14-4b67-9f9c-41b6db5a86f5",
   "metadata": {},
   "outputs": [],
   "source": [
    "def hourly_to_daily(\n",
    "    data_path,\n",
    "    variable,\n",
    "    years,\n",
    "    pre_process_func,\n",
    "    daily_func,\n",
    "    output_path,\n",
    "    new_var_name=None,\n",
    "    ):\n",
    "    \"\"\"\n",
    "    Process hourly data to daily (24 hour means starting at 1400UTC).\n",
    "    \n",
    "    data_path : directory for hourly data\n",
    "    variable : variable name\n",
    "    years : range of years to process\n",
    "    pre_process_fun : function to preprocess data in open_mfdataset\n",
    "    daily_func : function which processes hourly data (i.e. to 24 or 12 hours in this analysis)\n",
    "    output_path : path to write output\n",
    "    new_var_name : None, or string. Desired new variable name.\n",
    "    \"\"\"\n",
    "    files = get_files(data_path, variable, years) # Get files\n",
    "    \n",
    "    ds = xr.open_mfdataset( # Open dataset\n",
    "        files,\n",
    "        chunks={'time': 24, 'level': 1},\n",
    "        preprocess=pre_process_func,\n",
    "        compat='override',\n",
    "        coords='minimal',\n",
    "        engine='netcdf4'\n",
    "    )\n",
    "    \n",
    "    # ds = fn.daily_mean_1400(ds) # daily mean (24 hours from 1400UTC)\n",
    "    ds = daily_func(ds)\n",
    "    ds = ds.chunk({'time': 365})\n",
    "        \n",
    "    if isinstance(new_var_name, str):\n",
    "        print('Renaming variable')\n",
    "        ds = ds.rename({variable: new_var_name})\n",
    "        variable = new_var_name\n",
    "    \n",
    "    encoding = {variable: {'dtype': 'float32'}} # specify encoding\n",
    "    \n",
    "    ds.to_zarr( # Write\n",
    "        output_path,\n",
    "        mode='w',\n",
    "        consolidated=True,\n",
    "        encoding=encoding\n",
    "    )"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0706f76d-f5cc-41fe-ba68-c067c5841cec",
   "metadata": {},
   "source": [
    "### 500 hPa geopotential height\n",
    "\n",
    "- convert by dividing the geopotential by the Earth's gravitational acceleration, g (=9.80665 m s-2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "285e51d4-004c-41bc-b943-0faaf707f204",
   "metadata": {},
   "outputs": [],
   "source": [
    "def z_preprocess(ds):\n",
    "    \"\"\"\n",
    "    Select desired levels and subregion. Convert to\n",
    "    geopotential height by dividing the geopotential\n",
    "    by Earth's gravitational acceleration\n",
    "    g = 9.80665 m s-2\n",
    "    \n",
    "    Source: https://cds.climate.copernicus.eu/cdsapp#!/dataset/reanalysis-era5-pressure-levels-monthly-means\n",
    "    \"\"\"\n",
    "    ds = ds.rename({\n",
    "        'longitude': 'lon',\n",
    "        'latitude': 'lat'\n",
    "    })\n",
    "    ds = ds.sel(\n",
    "        level=levels,\n",
    "        lon=slice(extended_aus[0], extended_aus[1]),\n",
    "        lat=slice(extended_aus[2], extended_aus[3])\n",
    "    )\n",
    "    \n",
    "    # Convert from geopotential to geopotential height\n",
    "    g = 9.80665\n",
    "    ds['z'] = ds['z'] / g\n",
    "    \n",
    "    # New attributes\n",
    "    ds['z'] = ds['z'].assign_attrs({\n",
    "        'long_name': 'Geopotential height',\n",
    "        'standard_name': 'geopotential height',\n",
    "        'units': 'm'\n",
    "    })\n",
    "\n",
    "    return ds"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "cc31b591-325a-4a9b-afb9-7beb0be33fb1",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "/g/data/w42/dr6273/work/data/era5/z/z_era5_daily_1400UTC_pl_1959-2021_extended_Aus_region.zarr\n"
     ]
    }
   ],
   "source": [
    "dataset = 'era5'\n",
    "levels = [500]\n",
    "z_output_path = data_fp + dataset + '/z/z_' + dataset + '_daily_1400UTC_pl_' + str(years[0]) + '-' + str(years[-1]) + '_extended_Aus_region.zarr'\n",
    "print(z_output_path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "a2cac7c4-33ac-4688-8eb3-c265e960f5f3",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 14min 46s, sys: 37.2 s, total: 15min 23s\n",
      "Wall time: 19min 17s\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "hourly_to_daily(\n",
    "    data_path='/g/data/rt52/era5/pressure-levels/reanalysis/',\n",
    "    variable='z',\n",
    "    years=years,\n",
    "    pre_process_func=z_preprocess,\n",
    "    daily_func=fn.daily_mean_1400,\n",
    "    output_path=z_output_path\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9b1e48ec-7d30-49bc-b65e-224a4aff4c96",
   "metadata": {},
   "source": [
    "Open daily data, anomalise and write"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "c12f4b7a-9810-46b3-8072-265f73524265",
   "metadata": {},
   "outputs": [],
   "source": [
    "z = xr.open_zarr(z_output_path, consolidated=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "e97f7fc8-4017-4872-88a3-2cd652952b3d",
   "metadata": {},
   "outputs": [],
   "source": [
    "z_anoms = z.groupby('time.dayofyear').apply(lambda x: x - x.mean('time'))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "2066de67-fde5-4e78-8e0e-bf8acca35d11",
   "metadata": {},
   "outputs": [],
   "source": [
    "z_anoms = z_anoms.chunk({'time': 365})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "813e1599-eb8e-4018-b662-4592a9e286f2",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 1min 47s, sys: 2.26 s, total: 1min 49s\n",
      "Wall time: 1min 50s\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "<xarray.backends.zarr.ZarrStore at 0x14d2dff1be60>"
      ]
     },
     "execution_count": 20,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "%%time\n",
    "z_anoms.to_zarr( # Write\n",
    "    data_fp + dataset + '/z/z_anoms_' + dataset + '_daily_1400UTC_pl_' + str(years[0]) + '-' + str(years[-1]) + '_extended_Aus_region.zarr',\n",
    "    mode='w',\n",
    "    consolidated=True,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6fc0d727-8ecc-40d6-9eae-6d2f55addd0d",
   "metadata": {},
   "source": [
    "### MSLP"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "752246c0-9166-4d9c-b95a-931b50090bb1",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Using preprocess in open_mfdataset to select desired levels improves performance\n",
    "#  versus doing a .sel() afterwards\n",
    "def mslp_preprocess(ds):\n",
    "    \"\"\"\n",
    "    Select desired levels and subregion\n",
    "    \"\"\"\n",
    "    ds = ds.rename({\n",
    "        'longitude': 'lon',\n",
    "        'latitude': 'lat'\n",
    "    })\n",
    "    ds = ds.sel(\n",
    "        lon=slice(extended_aus[0], extended_aus[1]),\n",
    "        lat=slice(extended_aus[2], extended_aus[3])\n",
    "    )\n",
    "    return ds"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "5c13f84e-33ca-4c1c-b47b-99620cca83a9",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "/g/data/w42/dr6273/work/data/era5/mslp/mslp_era5_daily_1400UTC_1959-2021_extended_Aus_region.zarr\n"
     ]
    }
   ],
   "source": [
    "dataset = 'era5'\n",
    "mslp_output_path = data_fp + dataset + '/mslp/mslp_' + dataset + '_daily_1400UTC_' + str(years[0]) + '-' + str(years[-1]) + '_extended_Aus_region.zarr'\n",
    "print(mslp_output_path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "28839909-3473-48bc-a883-574857475c44",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Renaming variable\n",
      "CPU times: user 9min 44s, sys: 17.7 s, total: 10min 2s\n",
      "Wall time: 10min 58s\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "hourly_to_daily(\n",
    "    data_path='/g/data/rt52/era5/single-levels/reanalysis/',\n",
    "    variable='msl',\n",
    "    years=years,\n",
    "    pre_process_func=mslp_preprocess,\n",
    "    daily_func=fn.daily_mean_1400,\n",
    "    output_path=mslp_output_path,\n",
    "    new_var_name='mslp'\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "beec3a85-2a6f-4483-9dbb-126da905af3b",
   "metadata": {},
   "source": [
    "Open daily data, anomalise and write"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "8f644ee2-aeef-424b-a077-a9919410bf51",
   "metadata": {},
   "outputs": [],
   "source": [
    "mslp = xr.open_zarr(mslp_output_path, consolidated=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "e4b15a7b-cef8-40e8-b989-4e27184e4c56",
   "metadata": {},
   "outputs": [],
   "source": [
    "mslp_anoms = mslp.groupby('time.dayofyear').apply(lambda x: x - x.mean('time'))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "id": "123d99aa-d8d2-4f1c-a180-765ad141a308",
   "metadata": {},
   "outputs": [],
   "source": [
    "mslp_anoms = mslp_anoms.chunk({'time': 365})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "id": "1b37f5e5-23cd-49a3-aeac-fd89fd87b12c",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 2min 19s, sys: 2.63 s, total: 2min 22s\n",
      "Wall time: 2min 22s\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "<xarray.backends.zarr.ZarrStore at 0x154f80908b30>"
      ]
     },
     "execution_count": 23,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "%%time\n",
    "mslp_anoms.to_zarr( # Write\n",
    "    data_fp + dataset + '/mslp/mslp_anoms_' + dataset + '_daily_1400UTC_' + str(years[0]) + '-' + str(years[-1]) + '_extended_Aus_region.zarr',\n",
    "    mode='w',\n",
    "    consolidated=True\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c909063e-baa9-481d-8123-888356f6f364",
   "metadata": {},
   "source": [
    "### Fraction of cloud cover (tcc). Just do this for daytime hours (~2000 - 0800 UTC)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "133b373f-a481-487e-9e08-ce5e9edb8d5b",
   "metadata": {},
   "outputs": [],
   "source": [
    "def daytime_mean(ds):\n",
    "    \"\"\"\n",
    "    Daytime average of ds. Computed for 12 hour periods\n",
    "    starting at 2000UTC (so an eastern Aus 6am, roughly).\n",
    "    Currently hard-coded as its easier with open_mfdataset preprocess/\n",
    "    \"\"\"\n",
    "    ds_12 = ds.rolling(time=12).mean()\n",
    "    ds_daily = ds_12.isel(time=ds_12.time.dt.hour == 8)\n",
    "    return ds_daily"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "3317721b-3fc1-431e-bdb5-aae18721c37c",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Using preprocess in open_mfdataset to select desired levels improves performance\n",
    "#  versus doing a .sel() afterwards\n",
    "def tcc_preprocess(ds):\n",
    "    \"\"\"\n",
    "    Select desired levels and subregion\n",
    "    \"\"\"\n",
    "    ds = ds.rename({\n",
    "        'longitude': 'lon',\n",
    "        'latitude': 'lat'\n",
    "    })\n",
    "    ds = ds.sel(\n",
    "        lon=slice(extended_aus[0], extended_aus[1]),\n",
    "        lat=slice(extended_aus[2], extended_aus[3])\n",
    "    )\n",
    "    return ds"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "d9f29d8b-0604-4609-b010-d4837bab2d65",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "/g/data/w42/dr6273/work/data/era5/tcc/tcc_era5_12hr_2000-0800UTC_1959-2021_extended_Aus_region.zarr\n"
     ]
    }
   ],
   "source": [
    "dataset = 'era5'\n",
    "tcc_output_path = data_fp + dataset + '/tcc/tcc_' + dataset + '_12hr_2000-0800UTC_' + str(years[0]) + '-' + str(years[-1]) + '_extended_Aus_region.zarr'\n",
    "print(tcc_output_path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "e877e366-9753-48e9-81db-05a4b6eb6f0b",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 8min 59s, sys: 20.3 s, total: 9min 20s\n",
      "Wall time: 11min 2s\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "hourly_to_daily(\n",
    "    data_path='/g/data/rt52/era5/single-levels/reanalysis/',\n",
    "    variable='tcc',\n",
    "    years=years,\n",
    "    pre_process_func=tcc_preprocess,\n",
    "    daily_func=daytime_mean,\n",
    "    output_path=tcc_output_path\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b9212d1e-eb06-4e46-9429-9d99f5738271",
   "metadata": {},
   "source": [
    "Open daily data, anomalise and write"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "id": "011a7f1f-9d98-4021-8c16-6d1cc70a0009",
   "metadata": {},
   "outputs": [],
   "source": [
    "tcc = xr.open_zarr(tcc_output_path, consolidated=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "id": "11f1f8bd-06a4-4963-bd0f-34e0593b9fab",
   "metadata": {},
   "outputs": [],
   "source": [
    "tcc_anoms = tcc.groupby('time.dayofyear').apply(lambda x: x - x.mean('time'))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "id": "3fa4bfab-ddf1-499a-b611-370b36453642",
   "metadata": {},
   "outputs": [],
   "source": [
    "tcc_anoms = tcc_anoms.chunk({'time': 365})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "id": "ba39137d-e233-4d5d-951d-2bc34294829c",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 1min 28s, sys: 2.88 s, total: 1min 31s\n",
      "Wall time: 1min 36s\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "<xarray.backends.zarr.ZarrStore at 0x14c600c70660>"
      ]
     },
     "execution_count": 30,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "%%time\n",
    "tcc_anoms.to_zarr( # Write\n",
    "    data_fp + dataset + '/tcc/tcc_anoms_' + dataset + '_12hr_2000-0800UTC_' + str(years[0]) + '-' + str(years[-1]) + '_extended_Aus_region.zarr',\n",
    "    mode='w',\n",
    "    consolidated=True\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1d66db36-37be-4879-9de3-e876689a7c6c",
   "metadata": {},
   "source": [
    "### Close cluster"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 51,
   "id": "f2bb44c9-e80e-45d0-87ed-72d130df4929",
   "metadata": {},
   "outputs": [],
   "source": [
    "client.close()\n",
    "cluster.close()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e7211568-3f3c-41f4-9af7-5c565e961153",
   "metadata": {},
   "source": [
    "# Monthly data, global"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1e3b1cf6-0d13-4733-ba04-f674abb74706",
   "metadata": {},
   "source": [
    "### To process monthly data we need fewer resources"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 52,
   "id": "6523804c-2b00-4cf9-b5c6-6457529d233a",
   "metadata": {},
   "outputs": [],
   "source": [
    "# One node on Gadi has 48 cores - try and use up a full node before going to multiple nodes (jobs)\n",
    "\n",
    "walltime = '00:20:00'\n",
    "cores = 8\n",
    "memory = str(4 * cores) + 'GB'\n",
    "\n",
    "cluster = PBSCluster(walltime=str(walltime), cores=cores, memory=str(memory), processes=cores,\n",
    "                     job_extra_directives=['-q normal',\n",
    "                                           '-P w42',\n",
    "                                           '-l ncpus='+str(cores),\n",
    "                                           '-l mem='+str(memory),\n",
    "                                           '-l storage=gdata/w42+gdata/rt52+gdata/xv83'],\n",
    "                     local_directory='$TMPDIR',\n",
    "                     job_directives_skip=[\"select\"])\n",
    "                     # python=os.environ[\"DASK_PYTHON\"])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 53,
   "id": "810a169e-010e-4ace-8dfc-f8084ca5a416",
   "metadata": {},
   "outputs": [],
   "source": [
    "cluster.scale(jobs=1)\n",
    "client = Client(cluster)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 54,
   "id": "2257b3ef-be6e-447a-82bb-0f7490cce5c7",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "    <div style=\"width: 24px; height: 24px; background-color: #e1e1e1; border: 3px solid #9D9D9D; border-radius: 5px; position: absolute;\"> </div>\n",
       "    <div style=\"margin-left: 48px;\">\n",
       "        <h3 style=\"margin-bottom: 0px;\">Client</h3>\n",
       "        <p style=\"color: #9D9D9D; margin-bottom: 0px;\">Client-979ccdda-d8c8-11ed-a371-00000081fe80</p>\n",
       "        <table style=\"width: 100%; text-align: left;\">\n",
       "\n",
       "        <tr>\n",
       "        \n",
       "            <td style=\"text-align: left;\"><strong>Connection method:</strong> Cluster object</td>\n",
       "            <td style=\"text-align: left;\"><strong>Cluster type:</strong> dask_jobqueue.PBSCluster</td>\n",
       "        \n",
       "        </tr>\n",
       "\n",
       "        \n",
       "            <tr>\n",
       "                <td style=\"text-align: left;\">\n",
       "                    <strong>Dashboard: </strong> <a href=\"http://10.6.1.34:8787/status\" target=\"_blank\">http://10.6.1.34:8787/status</a>\n",
       "                </td>\n",
       "                <td style=\"text-align: left;\"></td>\n",
       "            </tr>\n",
       "        \n",
       "\n",
       "        </table>\n",
       "\n",
       "        \n",
       "            <button style=\"margin-bottom: 12px;\" data-commandlinker-command=\"dask:populate-and-launch-layout\" data-commandlinker-args='{\"url\": \"http://10.6.1.34:8787/status\" }'>\n",
       "                Launch dashboard in JupyterLab\n",
       "            </button>\n",
       "        \n",
       "\n",
       "        \n",
       "            <details>\n",
       "            <summary style=\"margin-bottom: 20px;\"><h3 style=\"display: inline;\">Cluster Info</h3></summary>\n",
       "            <div class=\"jp-RenderedHTMLCommon jp-RenderedHTML jp-mod-trusted jp-OutputArea-output\">\n",
       "    <div style=\"width: 24px; height: 24px; background-color: #e1e1e1; border: 3px solid #9D9D9D; border-radius: 5px; position: absolute;\">\n",
       "    </div>\n",
       "    <div style=\"margin-left: 48px;\">\n",
       "        <h3 style=\"margin-bottom: 0px; margin-top: 0px;\">PBSCluster</h3>\n",
       "        <p style=\"color: #9D9D9D; margin-bottom: 0px;\">dff7e70e</p>\n",
       "        <table style=\"width: 100%; text-align: left;\">\n",
       "            <tr>\n",
       "                <td style=\"text-align: left;\">\n",
       "                    <strong>Dashboard:</strong> <a href=\"http://10.6.1.34:8787/status\" target=\"_blank\">http://10.6.1.34:8787/status</a>\n",
       "                </td>\n",
       "                <td style=\"text-align: left;\">\n",
       "                    <strong>Workers:</strong> 0\n",
       "                </td>\n",
       "            </tr>\n",
       "            <tr>\n",
       "                <td style=\"text-align: left;\">\n",
       "                    <strong>Total threads:</strong> 0\n",
       "                </td>\n",
       "                <td style=\"text-align: left;\">\n",
       "                    <strong>Total memory:</strong> 0 B\n",
       "                </td>\n",
       "            </tr>\n",
       "            \n",
       "        </table>\n",
       "\n",
       "        <details>\n",
       "            <summary style=\"margin-bottom: 20px;\">\n",
       "                <h3 style=\"display: inline;\">Scheduler Info</h3>\n",
       "            </summary>\n",
       "\n",
       "            <div style=\"\">\n",
       "    <div>\n",
       "        <div style=\"width: 24px; height: 24px; background-color: #FFF7E5; border: 3px solid #FF6132; border-radius: 5px; position: absolute;\"> </div>\n",
       "        <div style=\"margin-left: 48px;\">\n",
       "            <h3 style=\"margin-bottom: 0px;\">Scheduler</h3>\n",
       "            <p style=\"color: #9D9D9D; margin-bottom: 0px;\">Scheduler-b0ea1966-8678-4b94-955e-f01754a50b2f</p>\n",
       "            <table style=\"width: 100%; text-align: left;\">\n",
       "                <tr>\n",
       "                    <td style=\"text-align: left;\">\n",
       "                        <strong>Comm:</strong> tcp://10.6.1.34:38595\n",
       "                    </td>\n",
       "                    <td style=\"text-align: left;\">\n",
       "                        <strong>Workers:</strong> 0\n",
       "                    </td>\n",
       "                </tr>\n",
       "                <tr>\n",
       "                    <td style=\"text-align: left;\">\n",
       "                        <strong>Dashboard:</strong> <a href=\"http://10.6.1.34:8787/status\" target=\"_blank\">http://10.6.1.34:8787/status</a>\n",
       "                    </td>\n",
       "                    <td style=\"text-align: left;\">\n",
       "                        <strong>Total threads:</strong> 0\n",
       "                    </td>\n",
       "                </tr>\n",
       "                <tr>\n",
       "                    <td style=\"text-align: left;\">\n",
       "                        <strong>Started:</strong> Just now\n",
       "                    </td>\n",
       "                    <td style=\"text-align: left;\">\n",
       "                        <strong>Total memory:</strong> 0 B\n",
       "                    </td>\n",
       "                </tr>\n",
       "            </table>\n",
       "        </div>\n",
       "    </div>\n",
       "\n",
       "    <details style=\"margin-left: 48px;\">\n",
       "        <summary style=\"margin-bottom: 20px;\">\n",
       "            <h3 style=\"display: inline;\">Workers</h3>\n",
       "        </summary>\n",
       "\n",
       "        \n",
       "\n",
       "    </details>\n",
       "</div>\n",
       "\n",
       "        </details>\n",
       "    </div>\n",
       "</div>\n",
       "            </details>\n",
       "        \n",
       "\n",
       "    </div>\n",
       "</div>"
      ],
      "text/plain": [
       "<Client: 'tcp://10.6.1.34:38595' processes=0 threads=0, memory=0 B>"
      ]
     },
     "execution_count": 54,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "client"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bcb0e12e-ce17-45e8-95e2-af5465c285f5",
   "metadata": {},
   "source": [
    "## SST data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 55,
   "id": "a3a4cc8a-c7cd-46a2-927a-9fb1ec4cef97",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "/g/data/w42/dr6273/work/data/hadisst/sst/sst_anom_hadisst_moda_sfc_1959-2021.zarr\n"
     ]
    }
   ],
   "source": [
    "dataset = 'hadisst'\n",
    "sst_fp = data_fp + dataset + '/sst/sst_anom_'+dataset+'_moda_sfc_'+str(years[0])+'-'+str(years[-1])+'.zarr'\n",
    "print(sst_fp)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 57,
   "id": "af7464ec-e3bd-476a-b515-d52313323279",
   "metadata": {},
   "outputs": [],
   "source": [
    "sst_anoms = anomalise_and_process(\n",
    "    load=True,\n",
    "    dataset='hadisst',\n",
    "    filepath=sst_fp,\n",
    "    data_root_path='/g/data/w42/dr6273/work/data/hadisst/sst/sst_hadisst_187001-202301.zarr',\n",
    "    var_name='sst',\n",
    "    years=years,\n",
    "    subset_region=None,\n",
    "    subset_level=None,\n",
    "    preprocess_func=None,\n",
    "    mfdataset_chunks=None,\n",
    "    anom_freq='month',\n",
    "    new_var_name='sst',\n",
    "    long_name='Sea surface temperature',\n",
    "    units='K',\n",
    "    chunks={'time': 12, 'lat': -1, 'lon': -1}\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5246d095-cdc0-4dbc-9dbc-be70e0c48736",
   "metadata": {},
   "source": [
    "### Nino 3.4"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 58,
   "id": "049e920a-f65f-4ab6-b9d9-834fc4867b3a",
   "metadata": {},
   "outputs": [],
   "source": [
    "nino34 = sst_anoms.sst_anom.sel(\n",
    "    lat=slice(5, -5),\n",
    "    lon=slice(-170, -120)\n",
    ").mean(['lat','lon']).to_dataset(name='nino34')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 62,
   "id": "bcce4396-a9e8-4293-a89f-2c9c9c1c556d",
   "metadata": {},
   "outputs": [],
   "source": [
    "nino34_dt = fn.detrend_dim(nino34.nino34, 'time').to_dataset(name='nino34_detrended')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 63,
   "id": "1c413519-b77f-4ffa-921e-f8ae3353091f",
   "metadata": {},
   "outputs": [],
   "source": [
    "nino34 = nino34.merge(nino34_dt)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 64,
   "id": "19924ccc-5123-497c-9573-ac482b9ffc3f",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<xarray.backends.zarr.ZarrStore at 0x14ae3cdd69d0>"
      ]
     },
     "execution_count": 64,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "nino34_fp = data_fp + dataset + '/climate_modes/'+dataset+'_nino34_'+str(years[0])+'-'+str(years[-1])+'.zarr'\n",
    "nino34.to_zarr(nino34_fp, mode='w', consolidated=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "317dcbec-9492-49ed-adaa-fe85f01bc207",
   "metadata": {},
   "source": [
    "### Dipole mode index"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 65,
   "id": "0d09bddd-24ab-458c-b7f6-8af674d9e615",
   "metadata": {},
   "outputs": [],
   "source": [
    "def calc_dmi(da):\n",
    "    \"\"\"\n",
    "    Calculate Dipole Mode Index\n",
    "    \"\"\"    \n",
    "    da_W = da.sel(lat=slice(10, -10), lon=slice(50, 70)).mean(['lat', 'lon'])\n",
    "    da_E = da.sel(lat=slice(0, -10), lon=slice(90, 110)).mean(['lat', 'lon'])\n",
    "    \n",
    "    return (da_W - da_E)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 66,
   "id": "6ca9f879-23e1-4bee-992b-e63e0ae8e33e",
   "metadata": {},
   "outputs": [],
   "source": [
    "dmi = calc_dmi(sst_anoms.sst_anom).to_dataset(name='dmi')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 67,
   "id": "bde8187f-36ea-4d69-8017-9036892beed1",
   "metadata": {},
   "outputs": [],
   "source": [
    "dmi_dt = fn.detrend_dim(dmi.dmi, 'time').to_dataset(name='dmi_detrended')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 68,
   "id": "11817727-ff3b-4dfb-bb5c-23176288562a",
   "metadata": {},
   "outputs": [],
   "source": [
    "dmi = dmi.merge(dmi_dt)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 69,
   "id": "810362cc-6909-45dc-b6c9-96f7dd4da4aa",
   "metadata": {},
   "outputs": [],
   "source": [
    "dmi_fp = data_fp + dataset + '/climate_modes/' + dataset + '_dmi_' + str(years[0]) + '-' + str(years[-1]) + '.zarr'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 70,
   "id": "947dea34-97dd-4103-a6f6-b315067a8805",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<xarray.backends.zarr.ZarrStore at 0x14ae31cb66c0>"
      ]
     },
     "execution_count": 70,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "dmi.to_zarr(dmi_fp, mode='w', consolidated=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "48788be9-d88b-4220-8351-dac6f350de2a",
   "metadata": {},
   "source": [
    "## MSLP data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "121b6b4a-d58d-45fc-b7e2-19327a830295",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "/g/data/w42/dr6273/work/data/era5/mslp/mslp_anom_era5_moda_sfc_1959-2021.zarr\n"
     ]
    }
   ],
   "source": [
    "dataset = 'era5'\n",
    "mslp_fp = data_fp + dataset + '/mslp/mslp_anom_'+dataset+'_moda_sfc_'+str(years[0])+'-'+str(years[-1])+'.zarr'\n",
    "print(mslp_fp)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "c8ffdb1d-717d-4550-9957-3364d2c0b49e",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Opening files...\n",
      "Anomalising...\n",
      "Attributes and writing...\n"
     ]
    }
   ],
   "source": [
    "mslp_anoms = anomalise_and_process(\n",
    "    load=True,\n",
    "    dataset='era5',\n",
    "    filepath=mslp_fp,\n",
    "    data_root_path='/g/data/rt52/era5/single-levels/monthly-averaged/',\n",
    "    var_name='msl',\n",
    "    years=years,\n",
    "    subset_region=None,\n",
    "    subset_level=None,\n",
    "    preprocess_func=None,\n",
    "    mfdataset_chunks=None,\n",
    "    anom_freq='month',\n",
    "    new_var_name='mslp',\n",
    "    long_name='Mean sea level pressure',\n",
    "    units='Pa',\n",
    "    chunks={'time': 12, 'lat': -1, 'lon': -1}\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8132f010-d73d-4f7a-bb4d-37de4e153468",
   "metadata": {},
   "source": [
    "### Southern Annular Mode index"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "a871d7d6-3e27-4861-91d9-2a6d7a43fa7c",
   "metadata": {},
   "outputs": [],
   "source": [
    "def calc_sam(mslp_anoms, lat_name='lat', lon_name='lon', time_name='time'):\n",
    "    \"\"\"\n",
    "    Calculate Southern Annular Mode index from MSLP anomalies. The SAM index is\n",
    "    defined as the difference in MSLP anomalies between 40 and 65 degrees South.\n",
    "    MSLP anomalies are first normalised by dividing by their standard deviation\n",
    "    (calculated as a function of calendar month).\n",
    "    \"\"\"\n",
    "    \n",
    "    mslp_40 = mslp_anoms.interp({lat_name: -40}).mean(lon_name)\n",
    "    mslp_65 = mslp_anoms.interp({lat_name: -65}).mean(lon_name)\n",
    "    \n",
    "    norm_40 = mslp_40.groupby(time_name+'.month').apply(lambda x: x / x.std(time_name))\n",
    "    norm_65 = mslp_65.groupby(time_name+'.month').apply(lambda x: x / x.std(time_name))\n",
    "    \n",
    "    return norm_40 - norm_65"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "5b7de44d-9610-4b6e-8bd5-1c180c8ed2ff",
   "metadata": {},
   "outputs": [],
   "source": [
    "sam = calc_sam(mslp_anoms)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "8de16708-569e-4247-929c-fd0a132a0050",
   "metadata": {},
   "outputs": [],
   "source": [
    "sam = sam.rename({'mslp_anom': 'sam'})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "5fca1bed-434a-4cd8-9c15-630e0ea8dc36",
   "metadata": {},
   "outputs": [],
   "source": [
    "sam_dt = fn.detrend_dim(sam.sam, 'time').to_dataset(name='sam_detrended')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "2d6db17d-4a20-4b62-b6e6-355035feb5f0",
   "metadata": {},
   "outputs": [],
   "source": [
    "sam = sam.merge(sam_dt)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "id": "27cf60fc-e64a-4fef-a800-e69fb0876b37",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<xarray.backends.zarr.ZarrStore at 0x1539ae3e6ce0>"
      ]
     },
     "execution_count": 22,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "sam_fp = data_fp + dataset + '/climate_modes/'+dataset+'_sam_'+str(years[0])+'-'+str(years[-1])+'.zarr'\n",
    "sam.to_zarr(sam_fp, mode='w', consolidated=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "38f21d89-34f0-48aa-a72d-f29aa3e81d0d",
   "metadata": {},
   "source": [
    "# Close cluster"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 71,
   "id": "dca99be2-eba1-4457-aac3-5a4552044885",
   "metadata": {},
   "outputs": [],
   "source": [
    "client.close()\n",
    "cluster.close()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e520de80-b857-4da8-ab24-65a903886be0",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python (pangeo)",
   "language": "python",
   "name": "pangeo"
  },
  "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.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}