0c_regional_averages.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "7f5990f0-2587-4772-b2a8-fdc40b19e1a7",
   "metadata": {},
   "source": [
    "# Calculate data for REZ and GCCSA regions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "c4c27199-6aba-41e7-9b25-c2ca19fa2c13",
   "metadata": {},
   "outputs": [],
   "source": [
    "from dask.distributed import Client,LocalCluster\n",
    "from dask_jobqueue import PBSCluster"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "bde72e92-135b-4a88-acfe-89f3021b4983",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/g/data/w42/dr6273/apps/conda/envs/pangeo/lib/python3.10/site-packages/distributed/node.py:182: UserWarning: Port 8787 is already in use.\n",
      "Perhaps you already have a cluster running?\n",
      "Hosting the HTTP server on port 42185 instead\n",
      "  warnings.warn(\n"
     ]
    }
   ],
   "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:15: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'],\n",
    "                     local_directory='$TMPDIR',\n",
    "                     job_directives_skip=[\"select\"])\n",
    "                     # python=os.environ[\"DASK_PYTHON\"])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "53272303-22b2-422c-a6a5-ee522e9c1210",
   "metadata": {},
   "outputs": [],
   "source": [
    "cluster.scale(jobs=1)\n",
    "client = Client(cluster)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "35b39eb8-53ae-4075-910f-236cca4ed27e",
   "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-9da2e779-d825-11ed-89a9-00000196fe80</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.71:42185/status\" target=\"_blank\">http://10.6.1.71:42185/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.71:42185/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;\">08a33cdc</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.71:42185/status\" target=\"_blank\">http://10.6.1.71:42185/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-5948a7f5-abaf-4d21-b157-3b7b6c4c59e1</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.71:46207\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.71:42185/status\" target=\"_blank\">http://10.6.1.71:42185/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.71:46207' processes=0 threads=0, memory=0 B>"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "client"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "3c591d9a-36a8-446c-96df-7895b5fe6198",
   "metadata": {},
   "outputs": [],
   "source": [
    "# client.close()\n",
    "# cluster.close()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "548e4b0b-fc16-4e1a-ab41-416e380001a5",
   "metadata": {},
   "outputs": [],
   "source": [
    "# %load_ext autoreload\n",
    "# %autoreload 2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "73f2c505-11be-4383-911b-14e4f2b17cd5",
   "metadata": {},
   "outputs": [],
   "source": [
    "import xarray as xr\n",
    "import numpy as np\n",
    "\n",
    "%matplotlib inline\n",
    "import matplotlib\n",
    "import matplotlib.pyplot as plt\n",
    "from mpl_toolkits.axes_grid1.inset_locator import inset_axes\n",
    "\n",
    "import cartopy.crs as ccrs\n",
    "import cartopy\n",
    "cartopy.config['pre_existing_data_dir'] = '/g/data/w42/dr6273/work/data/cartopy-data/'\n",
    "cartopy.config['data_dir'] = '/g/data/w42/dr6273/work/data/cartopy-data/'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "d999452a-8dad-4018-8d33-a49ea58f1a45",
   "metadata": {},
   "outputs": [],
   "source": [
    "import functions as fn\n",
    "\n",
    "plt_params = fn.get_plot_params()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f3233818-50d8-4b98-b4e5-bd37866d2f94",
   "metadata": {},
   "source": [
    "# Load masks"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "6ff78f87-200f-4cf6-9920-1044006f10cd",
   "metadata": {},
   "outputs": [],
   "source": [
    "rez_mask = fn.get_rez_mask()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "5fd006a3-adf2-4134-aed0-03322ee9fce0",
   "metadata": {},
   "outputs": [],
   "source": [
    "gccsa_mask = fn.get_gccsa_mask()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "97d83fe0-1b44-42d4-9caf-a956fa3fff36",
   "metadata": {},
   "source": [
    "# Load gridded data"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "01e7ca08-bb0e-4d19-a048-8136c664f0b7",
   "metadata": {},
   "source": [
    "First look at daily (1400). Have already looked at hourly and daily (0000)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "8e9c9b64-d53f-4d8b-823c-d04ee083ee12",
   "metadata": {},
   "outputs": [],
   "source": [
    "years = range(1959, 2022)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7972251b-f7f9-4457-8a6f-e6c131eab522",
   "metadata": {},
   "source": [
    "Mean solar shortwave radiation downwards"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "8aadf70d-cbe1-432f-a712-a99a547981da",
   "metadata": {},
   "outputs": [],
   "source": [
    "mssrd = xr.open_zarr(\n",
    "    '/g/data/w42/dr6273/work/data/era5/msdwswrf/msdwswrf_era5_daily_1400UTC_sfc_'+str(years[0])+'-'+str(years[-1])+'_AUS_region.zarr',\n",
    "    consolidated=True\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7c80fcf7-73dc-4d3d-81d1-e0d2da2dda4b",
   "metadata": {},
   "source": [
    "100m wind speed"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "d15366ac-9f9f-4d50-9f01-35add419c467",
   "metadata": {},
   "outputs": [],
   "source": [
    "w100 = xr.open_zarr(\n",
    "    '/g/data/w42/dr6273/work/data/era5/100w/100w_era5_daily_1400UTC_sfc_'+str(years[0])+'-'+str(years[-1])+'_AUS_region.zarr',\n",
    "    consolidated=True\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3ad7d122-175b-4a27-a437-a3578defd0d4",
   "metadata": {},
   "source": [
    "Runoff"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "205d502c-65e0-4923-b67f-997c660305bd",
   "metadata": {},
   "outputs": [],
   "source": [
    "ro = xr.open_zarr(\n",
    "    '/g/data/w42/dr6273/work/data/era5/ro/ro_era5_monthly-averaged_sfc_'+str(years[0])+'-'+str(years[-1])+'_REZ_region.zarr',\n",
    "    consolidated=True\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "93a21017-0618-4f03-b060-2355cc8894b2",
   "metadata": {},
   "source": [
    "2m temperature"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "d2471d3d-392c-4adc-9cab-bc75ff20f697",
   "metadata": {},
   "outputs": [],
   "source": [
    "t = xr.open_zarr(\n",
    "    '/g/data/w42/dr6273/work/data/era5/2t/2t_era5_daily_1400UTC_sfc_'+str(years[0])+'-'+str(years[-1])+'_extended_AUS_region.zarr',\n",
    "    consolidated=True\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d96ed336-52af-42e0-a7c8-64d86ca3334f",
   "metadata": {},
   "source": [
    "# Calculate regional averages"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "904d185f-caf7-466e-9808-cf9edc45dabe",
   "metadata": {},
   "outputs": [],
   "source": [
    "def region_spatial_mean(da, region, mask):\n",
    "    \"\"\"\n",
    "    Calculate spatial mean for a single region.\n",
    "    \"\"\"\n",
    "    region_mask = mask.sel(region=region)\n",
    "    return da.where(region_mask, drop=True).mean(['lat', 'lon'], skipna=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "ca7e15cd-5137-4ced-92c9-27ee2a07112f",
   "metadata": {},
   "outputs": [],
   "source": [
    "def calculate_spatial_means(da, mask):\n",
    "    \"\"\"\n",
    "    Calculate spatial means over all regions in da.\n",
    "    \"\"\"\n",
    "    da_list = []\n",
    "    for region in mask['region'].values:\n",
    "        region_mean = region_spatial_mean(da, region, mask)\n",
    "        da_list.append(region_mean)\n",
    "        \n",
    "    da_concat = xr.concat(\n",
    "        da_list,\n",
    "        dim='region'\n",
    "    )\n",
    "    return da_concat"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "846ebaa0-3bd3-4c2a-ace9-16a7042abcc1",
   "metadata": {},
   "outputs": [],
   "source": [
    "def to_single_chunk(da):\n",
    "    \"\"\"\n",
    "    Rechunk to a single chunk.\n",
    "    \"\"\"\n",
    "    return da.chunk({\n",
    "        'region': -1,\n",
    "        'time': -1\n",
    "    })"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "8ed3c122-528b-46a0-8ea2-7acef26401ed",
   "metadata": {},
   "outputs": [],
   "source": [
    "def calculate_and_write(da, mask, mask_name, var_name, time_freq_name):\n",
    "    \"\"\"\n",
    "    Calculate regional means, chunk and write.\n",
    "    \"\"\"\n",
    "    da = calculate_spatial_means(da, mask=mask)\n",
    "    da = to_single_chunk(da)\n",
    "    ds = da.to_dataset(name=var_name)\n",
    "    ds.to_zarr(\n",
    "        '/g/data/w42/dr6273/work/projects/Aus_energy/data/'+var_name+'_era5_'+time_freq_name+'_sfc_'+str(years[0])+'-'+str(years[-1])+'_'+mask_name+'_region_mean.zarr',\n",
    "        mode='w',\n",
    "        consolidated=True\n",
    "    )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "id": "ec7a259c-ff87-487b-91cd-5d4cc20bdc11",
   "metadata": {},
   "outputs": [],
   "source": [
    "calculate_and_write(\n",
    "    da=w100['w100'],\n",
    "    mask=rez_mask,\n",
    "    mask_name='REZ',\n",
    "    var_name='100w',\n",
    "    time_freq_name='daily_1400UTC'\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "id": "413ffc33-d323-4816-b80a-2a2b665b6137",
   "metadata": {},
   "outputs": [],
   "source": [
    "calculate_and_write(\n",
    "    da=mssrd['msdwswrf'],\n",
    "    mask=rez_mask,\n",
    "    mask_name='REZ',\n",
    "    var_name='mssrd',\n",
    "    time_freq_name='daily_1400UTC'\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "eaf91060-88f3-45f5-b472-a24303ca9228",
   "metadata": {},
   "outputs": [],
   "source": [
    "calculate_and_write(\n",
    "    da=t['t2m'],\n",
    "    mask=gccsa_mask,\n",
    "    mask_name='GCCSA',\n",
    "    var_name='t2m',\n",
    "    time_freq_name='daily_1400UTC'\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "id": "db5a4add-4a7e-462a-87f0-726d110d950b",
   "metadata": {},
   "outputs": [],
   "source": [
    "calculate_and_write(\n",
    "    da=ro['ro'],\n",
    "    mask=rez_mask,\n",
    "    mask_name='REZ',\n",
    "    var_name='ro',\n",
    "    time_freq_name='monthly-averaged'\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c50e8c0d-04cf-4910-9354-4dc8cd6ef47e",
   "metadata": {},
   "source": [
    "# Close cluster"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "id": "b8f19820-b88b-41de-8c1f-4b0c2289da79",
   "metadata": {},
   "outputs": [],
   "source": [
    "client.close()\n",
    "cluster.close()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9554ddf9-d0d6-43e0-a9db-e895a8d5ce2c",
   "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
}