{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Loading Unified Model output"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "For convenience, `aeolus` provides a way of keeping loaded and processed data within one object along with extra metadata.\n",
    "The object is called `Run` (as in \"simulation run\"). \n",
    "The code below provides an example of basic usage of `Run`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "from pathlib import Path\n",
    "\n",
    "import iris\n",
    "\n",
    "from aeolus.core import Run"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Note that you can use either a single filename or a list of filenames, each of which is either a `str` or (recommended) `pathlib.Path` object."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "sample_file = Path.cwd() / \"sample_data\" / \"sample_t1e_2d_mean.pp\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "While instantiating `Run`, it is possible to add a short name, a long description of the experiment; and to specify a planet configuration with relevant constants (see \"Physical constants\" example for more info)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "my_run = Run(\n",
    "    files=sample_file,\n",
    "    name=\"t1e_example\",\n",
    "    description=\"This is some sample data from a UM simulation of tidally-locked Trappist-1e planet.\",\n",
    "    planet=\"trap1e\"  # this reads constants from a JSON file\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<aeolus.core.Run at 0x7f3c39042400>"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "my_run"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Constants that have been used in the model:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Trap1eConstants(earth_day [s], stefan_boltzmann [W m-2 K-4], molar_gas_constant [J K-1 mol-1], water_heat_vaporization [m2 s-2], water_molecular_weight [kg mol-1], gravity [m s-2], radius [m], day [s], solar_constant [W m-2], semi_major_axis [au], eccentricity [1], obliquity [degree], dry_air_spec_heat_press [m2 s-2 K-1], dry_air_molecular_weight [kg mol-1])"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "my_run.const"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The loaded data are stored as a `CubeList` under `raw` attribute."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[<iris 'Cube' of convective_rainfall_flux / (kg m-2 s-1) (latitude: 90; longitude: 144)>,\n",
       "<iris 'Cube' of convective_snowfall_flux / (kg m-2 s-1) (latitude: 90; longitude: 144)>,\n",
       "<iris 'Cube' of high_type_cloud_area_fraction / (1) (latitude: 90; longitude: 144)>,\n",
       "<iris 'Cube' of low_type_cloud_area_fraction / (1) (latitude: 90; longitude: 144)>,\n",
       "<iris 'Cube' of medium_type_cloud_area_fraction / (1) (latitude: 90; longitude: 144)>,\n",
       "<iris 'Cube' of stratiform_rainfall_flux / (kg m-2 s-1) (latitude: 90; longitude: 144)>,\n",
       "<iris 'Cube' of stratiform_snowfall_flux / (kg m-2 s-1) (latitude: 90; longitude: 144)>]"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "my_run.raw"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Minimal processing of loaded data"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The next step might involve some clean-up and post-processing of raw data.\n",
    "\n",
    "This is done by creating a function that takes `iris.cube.CubeList` as its 1st argument and returns another `iris.cube.CubeList` as output.\n",
    "The function is then passed to `Run.proc_data()` and its output stored as `Run.proc` attribute."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "For example, the function `roll_cube_pm180()` imported below takes a `Cube` and rolls its longitudes from 0...360 to -180...180."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "from aeolus.grid import roll_cube_pm180, ensure_bounds"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "It can then be applied to all `raw` cubes."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "def _prepare_cubes(cubelist):\n",
    "    \"\"\"Post-process data for easier analysis.\"\"\"\n",
    "    # Roll cubes\n",
    "    r_cubes = iris.cube.CubeList()\n",
    "    for cube in cubelist:\n",
    "        r_c = roll_cube_pm180(cube)\n",
    "        ensure_bounds(r_c)  # also, ensure that longitudes and latitudes have bounds\n",
    "        r_cubes.append(r_c)\n",
    "\n",
    "    return r_cubes"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "my_run.proc_data(_prepare_cubes)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[<iris 'Cube' of convective_rainfall_flux / (kg m-2 s-1) (latitude: 90; longitude: 144)>,\n",
       "<iris 'Cube' of convective_snowfall_flux / (kg m-2 s-1) (latitude: 90; longitude: 144)>,\n",
       "<iris 'Cube' of high_type_cloud_area_fraction / (1) (latitude: 90; longitude: 144)>,\n",
       "<iris 'Cube' of low_type_cloud_area_fraction / (1) (latitude: 90; longitude: 144)>,\n",
       "<iris 'Cube' of medium_type_cloud_area_fraction / (1) (latitude: 90; longitude: 144)>,\n",
       "<iris 'Cube' of stratiform_rainfall_flux / (kg m-2 s-1) (latitude: 90; longitude: 144)>,\n",
       "<iris 'Cube' of stratiform_snowfall_flux / (kg m-2 s-1) (latitude: 90; longitude: 144)>]"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "my_run.proc"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Check that it did what expected."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "DimCoord(array([  1.25,   3.75,   6.25,   8.75,  11.25,  13.75,  16.25,  18.75,\n",
      "        21.25,  23.75,  26.25,  28.75,  31.25,  33.75,  36.25,  38.75,\n",
      "        41.25,  43.75,  46.25,  48.75,  51.25,  53.75,  56.25,  58.75,\n",
      "        61.25,  63.75,  66.25,  68.75,  71.25,  73.75,  76.25,  78.75,\n",
      "        81.25,  83.75,  86.25,  88.75,  91.25,  93.75,  96.25,  98.75,\n",
      "       101.25, 103.75, 106.25, 108.75, 111.25, 113.75, 116.25, 118.75,\n",
      "       121.25, 123.75, 126.25, 128.75, 131.25, 133.75, 136.25, 138.75,\n",
      "       141.25, 143.75, 146.25, 148.75, 151.25, 153.75, 156.25, 158.75,\n",
      "       161.25, 163.75, 166.25, 168.75, 171.25, 173.75, 176.25, 178.75,\n",
      "       181.25, 183.75, 186.25, 188.75, 191.25, 193.75, 196.25, 198.75,\n",
      "       201.25, 203.75, 206.25, 208.75, 211.25, 213.75, 216.25, 218.75,\n",
      "       221.25, 223.75, 226.25, 228.75, 231.25, 233.75, 236.25, 238.75,\n",
      "       241.25, 243.75, 246.25, 248.75, 251.25, 253.75, 256.25, 258.75,\n",
      "       261.25, 263.75, 266.25, 268.75, 271.25, 273.75, 276.25, 278.75,\n",
      "       281.25, 283.75, 286.25, 288.75, 291.25, 293.75, 296.25, 298.75,\n",
      "       301.25, 303.75, 306.25, 308.75, 311.25, 313.75, 316.25, 318.75,\n",
      "       321.25, 323.75, 326.25, 328.75, 331.25, 333.75, 336.25, 338.75,\n",
      "       341.25, 343.75, 346.25, 348.75, 351.25, 353.75, 356.25, 358.75],\n",
      "      dtype=float32), standard_name='longitude', units=Unit('degrees'), coord_system=GeogCS(6371229.0), circular=True)\n"
     ]
    }
   ],
   "source": [
    "print(my_run.raw.extract_strict(\"convective_rainfall_flux\").coord(\"longitude\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "DimCoord(array([-178.75, -176.25, -173.75, -171.25, -168.75, -166.25, -163.75,\n",
      "       -161.25, -158.75, -156.25, -153.75, -151.25, -148.75, -146.25,\n",
      "       -143.75, -141.25, -138.75, -136.25, -133.75, -131.25, -128.75,\n",
      "       -126.25, -123.75, -121.25, -118.75, -116.25, -113.75, -111.25,\n",
      "       -108.75, -106.25, -103.75, -101.25,  -98.75,  -96.25,  -93.75,\n",
      "        -91.25,  -88.75,  -86.25,  -83.75,  -81.25,  -78.75,  -76.25,\n",
      "        -73.75,  -71.25,  -68.75,  -66.25,  -63.75,  -61.25,  -58.75,\n",
      "        -56.25,  -53.75,  -51.25,  -48.75,  -46.25,  -43.75,  -41.25,\n",
      "        -38.75,  -36.25,  -33.75,  -31.25,  -28.75,  -26.25,  -23.75,\n",
      "        -21.25,  -18.75,  -16.25,  -13.75,  -11.25,   -8.75,   -6.25,\n",
      "         -3.75,   -1.25,    1.25,    3.75,    6.25,    8.75,   11.25,\n",
      "         13.75,   16.25,   18.75,   21.25,   23.75,   26.25,   28.75,\n",
      "         31.25,   33.75,   36.25,   38.75,   41.25,   43.75,   46.25,\n",
      "         48.75,   51.25,   53.75,   56.25,   58.75,   61.25,   63.75,\n",
      "         66.25,   68.75,   71.25,   73.75,   76.25,   78.75,   81.25,\n",
      "         83.75,   86.25,   88.75,   91.25,   93.75,   96.25,   98.75,\n",
      "        101.25,  103.75,  106.25,  108.75,  111.25,  113.75,  116.25,\n",
      "        118.75,  121.25,  123.75,  126.25,  128.75,  131.25,  133.75,\n",
      "        136.25,  138.75,  141.25,  143.75,  146.25,  148.75,  151.25,\n",
      "        153.75,  156.25,  158.75,  161.25,  163.75,  166.25,  168.75,\n",
      "        171.25,  173.75,  176.25,  178.75]), bounds=array([[-180. , -177.5],\n",
      "       [-177.5, -175. ],\n",
      "       [-175. , -172.5],\n",
      "       [-172.5, -170. ],\n",
      "       [-170. , -167.5],\n",
      "       [-167.5, -165. ],\n",
      "       [-165. , -162.5],\n",
      "       [-162.5, -160. ],\n",
      "       [-160. , -157.5],\n",
      "       [-157.5, -155. ],\n",
      "       [-155. , -152.5],\n",
      "       [-152.5, -150. ],\n",
      "       [-150. , -147.5],\n",
      "       [-147.5, -145. ],\n",
      "       [-145. , -142.5],\n",
      "       [-142.5, -140. ],\n",
      "       [-140. , -137.5],\n",
      "       [-137.5, -135. ],\n",
      "       [-135. , -132.5],\n",
      "       [-132.5, -130. ],\n",
      "       [-130. , -127.5],\n",
      "       [-127.5, -125. ],\n",
      "       [-125. , -122.5],\n",
      "       [-122.5, -120. ],\n",
      "       [-120. , -117.5],\n",
      "       [-117.5, -115. ],\n",
      "       [-115. , -112.5],\n",
      "       [-112.5, -110. ],\n",
      "       [-110. , -107.5],\n",
      "       [-107.5, -105. ],\n",
      "       [-105. , -102.5],\n",
      "       [-102.5, -100. ],\n",
      "       [-100. ,  -97.5],\n",
      "       [ -97.5,  -95. ],\n",
      "       [ -95. ,  -92.5],\n",
      "       [ -92.5,  -90. ],\n",
      "       [ -90. ,  -87.5],\n",
      "       [ -87.5,  -85. ],\n",
      "       [ -85. ,  -82.5],\n",
      "       [ -82.5,  -80. ],\n",
      "       [ -80. ,  -77.5],\n",
      "       [ -77.5,  -75. ],\n",
      "       [ -75. ,  -72.5],\n",
      "       [ -72.5,  -70. ],\n",
      "       [ -70. ,  -67.5],\n",
      "       [ -67.5,  -65. ],\n",
      "       [ -65. ,  -62.5],\n",
      "       [ -62.5,  -60. ],\n",
      "       [ -60. ,  -57.5],\n",
      "       [ -57.5,  -55. ],\n",
      "       [ -55. ,  -52.5],\n",
      "       [ -52.5,  -50. ],\n",
      "       [ -50. ,  -47.5],\n",
      "       [ -47.5,  -45. ],\n",
      "       [ -45. ,  -42.5],\n",
      "       [ -42.5,  -40. ],\n",
      "       [ -40. ,  -37.5],\n",
      "       [ -37.5,  -35. ],\n",
      "       [ -35. ,  -32.5],\n",
      "       [ -32.5,  -30. ],\n",
      "       [ -30. ,  -27.5],\n",
      "       [ -27.5,  -25. ],\n",
      "       [ -25. ,  -22.5],\n",
      "       [ -22.5,  -20. ],\n",
      "       [ -20. ,  -17.5],\n",
      "       [ -17.5,  -15. ],\n",
      "       [ -15. ,  -12.5],\n",
      "       [ -12.5,  -10. ],\n",
      "       [ -10. ,   -7.5],\n",
      "       [  -7.5,   -5. ],\n",
      "       [  -5. ,   -2.5],\n",
      "       [  -2.5,    0. ],\n",
      "       [   0. ,    2.5],\n",
      "       [   2.5,    5. ],\n",
      "       [   5. ,    7.5],\n",
      "       [   7.5,   10. ],\n",
      "       [  10. ,   12.5],\n",
      "       [  12.5,   15. ],\n",
      "       [  15. ,   17.5],\n",
      "       [  17.5,   20. ],\n",
      "       [  20. ,   22.5],\n",
      "       [  22.5,   25. ],\n",
      "       [  25. ,   27.5],\n",
      "       [  27.5,   30. ],\n",
      "       [  30. ,   32.5],\n",
      "       [  32.5,   35. ],\n",
      "       [  35. ,   37.5],\n",
      "       [  37.5,   40. ],\n",
      "       [  40. ,   42.5],\n",
      "       [  42.5,   45. ],\n",
      "       [  45. ,   47.5],\n",
      "       [  47.5,   50. ],\n",
      "       [  50. ,   52.5],\n",
      "       [  52.5,   55. ],\n",
      "       [  55. ,   57.5],\n",
      "       [  57.5,   60. ],\n",
      "       [  60. ,   62.5],\n",
      "       [  62.5,   65. ],\n",
      "       [  65. ,   67.5],\n",
      "       [  67.5,   70. ],\n",
      "       [  70. ,   72.5],\n",
      "       [  72.5,   75. ],\n",
      "       [  75. ,   77.5],\n",
      "       [  77.5,   80. ],\n",
      "       [  80. ,   82.5],\n",
      "       [  82.5,   85. ],\n",
      "       [  85. ,   87.5],\n",
      "       [  87.5,   90. ],\n",
      "       [  90. ,   92.5],\n",
      "       [  92.5,   95. ],\n",
      "       [  95. ,   97.5],\n",
      "       [  97.5,  100. ],\n",
      "       [ 100. ,  102.5],\n",
      "       [ 102.5,  105. ],\n",
      "       [ 105. ,  107.5],\n",
      "       [ 107.5,  110. ],\n",
      "       [ 110. ,  112.5],\n",
      "       [ 112.5,  115. ],\n",
      "       [ 115. ,  117.5],\n",
      "       [ 117.5,  120. ],\n",
      "       [ 120. ,  122.5],\n",
      "       [ 122.5,  125. ],\n",
      "       [ 125. ,  127.5],\n",
      "       [ 127.5,  130. ],\n",
      "       [ 130. ,  132.5],\n",
      "       [ 132.5,  135. ],\n",
      "       [ 135. ,  137.5],\n",
      "       [ 137.5,  140. ],\n",
      "       [ 140. ,  142.5],\n",
      "       [ 142.5,  145. ],\n",
      "       [ 145. ,  147.5],\n",
      "       [ 147.5,  150. ],\n",
      "       [ 150. ,  152.5],\n",
      "       [ 152.5,  155. ],\n",
      "       [ 155. ,  157.5],\n",
      "       [ 157.5,  160. ],\n",
      "       [ 160. ,  162.5],\n",
      "       [ 162.5,  165. ],\n",
      "       [ 165. ,  167.5],\n",
      "       [ 167.5,  170. ],\n",
      "       [ 170. ,  172.5],\n",
      "       [ 172.5,  175. ],\n",
      "       [ 175. ,  177.5],\n",
      "       [ 177.5,  180. ]]), standard_name='longitude', units=Unit('degrees'), coord_system=GeogCS(5804071.0), circular=True)\n"
     ]
    }
   ],
   "source": [
    "print(my_run.proc.extract_strict(\"convective_rainfall_flux\").coord(\"longitude\"))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Note that the longitude coordinate is not only shifted by 180 degrees, but has automatically calculated `bounds`.\n",
    "\n",
    "In addition, note that `coord_system` in `proc` cubes is different, because it used `Run.const` attribute to redefine the planet radius correctly.\n",
    "By default the loaded raw data has Earth radius in its `coord_system`, so certain calculations (e.g. grid cell areas) might be incorrect."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python [conda env:aeolus_py37]",
   "language": "python",
   "name": "conda-env-aeolus_py37-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.7.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
