{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Working with 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 `AtmoSim` (as in \"atmospheric simulation\"). \n",
    "The code below provides an example of basic usage of `AtmoSim`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "from pathlib import Path\n",
    "\n",
    "import iris\n",
    "from aeolus.core import AtmoSim"
   ]
  },
  {
   "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().parent / \"tests\" / \"data\" / \"test_data\" / \"pp\" / \"sample_t1e_2d_mean.pp\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To load the data, use `iris` functions."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "inp_data = iris.load(sample_file)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "iris.cube.CubeList"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "type(inp_data)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This cube list can be then used to instantiate `AtmoSim`. It is also 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": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "my_run = AtmoSim(\n",
    "    cubes=inp_data,\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": 6,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "aeolus.core.AtmoSim"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "type(my_run)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The loaded data are stored as a `CubeList` as the `_cubes` attribute."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0: convective_rainfall_flux / (kg m-2 s-1) (latitude: 90; longitude: 144)\n",
      "1: convective_snowfall_flux / (kg m-2 s-1) (latitude: 90; longitude: 144)\n",
      "2: high_type_cloud_area_fraction / (1) (latitude: 90; longitude: 144)\n",
      "3: low_type_cloud_area_fraction / (1)  (latitude: 90; longitude: 144)\n",
      "4: medium_type_cloud_area_fraction / (1) (latitude: 90; longitude: 144)\n",
      "5: stratiform_rainfall_flux / (kg m-2 s-1) (latitude: 90; longitude: 144)\n",
      "6: stratiform_snowfall_flux / (kg m-2 s-1) (latitude: 90; longitude: 144)\n"
     ]
    }
   ],
   "source": [
    "print(my_run._cubes)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Constants that have been used in the model:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Trap1eConstants(earth_day [s], stefan_boltzmann [W m-2 K-4], water_heat_vaporization [m2 s-2], water_molecular_weight [kg mol-1], molar_gas_constant [J K-1 mol-1], boltzmann [m^2 kg s^-2 K^-1], avogadro [mol-1], gravity [m s-2], radius [m], day [s], solar_constant [W m-2], reference_surface_pressure [Pa], 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], condensible_density [kg m-3], condensible_heat_vaporization [m2 s-2])"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "my_run.const"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Individual cubes can be accessed either by via the `_cubes` attribute and `iris` methods, or, more conveniently, as keys or attributes of `AtmoSim`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "convective_rainfall_flux / (kg m-2 s-1) (latitude: 90; longitude: 144)\n",
      "    Dimension coordinates:\n",
      "        latitude                                 x              -\n",
      "        longitude                                -              x\n",
      "    Scalar coordinates:\n",
      "        forecast_period                 36360.0 hours, bound=(36000.0, 36720.0) hours\n",
      "        forecast_reference_time         2004-05-13 09:00:00\n",
      "        time                            2008-07-06 09:00:00, bound=(2008-06-21 09:00:00, 2008-07-21 09:00:00)\n",
      "    Cell methods:\n",
      "        mean                            time (1 hour)\n",
      "    Attributes:\n",
      "        STASH                           m01s05i205\n",
      "        planet_conf                     Trap1eConstants(earth_day [s], stefan_boltzmann [W m-2 K-4], water_heat_vaporization...\n",
      "        source                          'Data from Met Office Unified Model'\n",
      "        um_version                      '11.3'\n"
     ]
    }
   ],
   "source": [
    "print(my_run.cv_rain)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "convective_rainfall_flux / (kg m-2 s-1) (latitude: 90; longitude: 144)\n",
      "    Dimension coordinates:\n",
      "        latitude                                 x              -\n",
      "        longitude                                -              x\n",
      "    Scalar coordinates:\n",
      "        forecast_period                 36360.0 hours, bound=(36000.0, 36720.0) hours\n",
      "        forecast_reference_time         2004-05-13 09:00:00\n",
      "        time                            2008-07-06 09:00:00, bound=(2008-06-21 09:00:00, 2008-07-21 09:00:00)\n",
      "    Cell methods:\n",
      "        mean                            time (1 hour)\n",
      "    Attributes:\n",
      "        STASH                           m01s05i205\n",
      "        planet_conf                     Trap1eConstants(earth_day [s], stefan_boltzmann [W m-2 K-4], water_heat_vaporization...\n",
      "        source                          'Data from Met Office Unified Model'\n",
      "        um_version                      '11.3'\n"
     ]
    }
   ],
   "source": [
    "print(my_run[\"cv_rain\"])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "which is equivalent to typing"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "convective_rainfall_flux / (kg m-2 s-1) (latitude: 90; longitude: 144)\n",
      "    Dimension coordinates:\n",
      "        latitude                                 x              -\n",
      "        longitude                                -              x\n",
      "    Scalar coordinates:\n",
      "        forecast_period                 36360.0 hours, bound=(36000.0, 36720.0) hours\n",
      "        forecast_reference_time         2004-05-13 09:00:00\n",
      "        time                            2008-07-06 09:00:00, bound=(2008-06-21 09:00:00, 2008-07-21 09:00:00)\n",
      "    Cell methods:\n",
      "        mean                            time (1 hour)\n",
      "    Attributes:\n",
      "        STASH                           m01s05i205\n",
      "        planet_conf                     Trap1eConstants(earth_day [s], stefan_boltzmann [W m-2 K-4], water_heat_vaporization...\n",
      "        source                          'Data from Met Office Unified Model'\n",
      "        um_version                      '11.3'\n"
     ]
    }
   ],
   "source": [
    "print(my_run._cubes.extract_cube(\"convective_rainfall_flux\"))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Short-cuts are defined using the `Model` container (see the \"Model field names\" example)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'convective_rainfall_flux'"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "my_run.model.cv_rain"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python [conda env:aeolus_test]",
   "language": "python",
   "name": "conda-env-aeolus_test-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.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}