{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Using `climlab` to solve the 1D diffusive Energy Balance Model"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## An interactive tutorial for ENV / ATM 415: Climate Laboratory\n",
    "\n",
    "### Spring 2016\n",
    "\n",
    "____________________________________"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## The model\n",
    "\n",
    "The equation for the model is\n",
    "\n",
    "$$ C \\frac{\\partial T}{\\partial t} = ASR - OLR + \\frac{D}{\\cos\\phi} \\frac{\\partial}{\\partial \\phi} \\left( \\cos\\phi \\frac{\\partial T}{\\partial \\phi} \\right)$$\n",
    "\n",
    "Where, following our standard notation:\n",
    "- $T$ is the surface temperature (in this case, the **zonal average** temperature)\n",
    "- $C$ is the heat capacity in J m$^{-2}$ K$^{-1}$\n",
    "- $ASR$ is the absorbed solar radiation\n",
    "- $OLR$ is the outgoing longwave radiation\n",
    "- $D$ is the thermal diffusivity in units of W m$^{-2}$ K$^{-1}$\n",
    "- $\\phi$ is latitude"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Longwave parameterization\n",
    "\n",
    "We will use a very simple parameterization linking surface temperature to OLR:\n",
    "\n",
    "$$ OLR = A + B~T $$"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Shortwave parameterization\n",
    "\n",
    "$$ ASR = (1-\\alpha) Q $$\n",
    "\n",
    "and we will use the **annual average insolation** (ignore seasons).\n",
    "\n",
    "The albedo $\\alpha$ (at least for now) will be a simple prescribed function of latitude (larger at the poles than at the equator) that mimics the observed annual mean albedo."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## The EBM is already coded up in `climlab`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "#  We start with the usual import statements\n",
    "\n",
    "%matplotlib inline\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "import climlab"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "#  create a new model with all default parameters (except the grid size)\n",
    "\n",
    "mymodel = climlab.EBM_annual(num_lat = 30)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "#  What did we just do?\n",
    "\n",
    "print mymodel"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## What is a `climlab Process`?\n",
    "\n",
    "``climlab`` is a flexible Python engine for process-oriented climate modeling.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Every model in `climlab` consists of some **state variables**, and **processes** that operate on those state variables.\n",
    "\n",
    "The **state variables** are the climate variables that we will step forward in time. For our EBM, it is just the surface temperature: "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "mymodel.Ts"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We have an array of temperatures in degrees Celsius.  Let's see how big this array is:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "mymodel.Ts.shape"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Every state variable exists on a spatial grid. In this case, the grid is an array of latitude points:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "mymodel.lat"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can, for example, **plot the current temperature versus latitude**:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "plt.plot(mymodel.lat, mymodel.Ts)\n",
    "plt.xlabel('Latitude')\n",
    "plt.ylabel('Temperature (deg C)')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "It is based on a very general concept of a model as a collection of individual, interacting processes. \n",
    "\n",
    "Every model in ``climlab`` is a collection of individual, interacting physical processes. \n",
    "\n",
    "In this case, the ``EBM_annual`` object has a list of 4 sub-processes. Each one is basically one of the terms in our energy budget equation above:\n",
    "\n",
    "- diffusion\n",
    "- Longwave radiation ($A+BT$)\n",
    "- Insolation (annual mean)\n",
    "- Albedo (in this case, a fixed spatial variation)\n",
    "\n",
    "To see this explicitly:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false,
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "#  The dictionary of sub-processes:\n",
    "mymodel.subprocess"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## So what does it do?\n",
    "\n",
    "Well for one thing, we can **step the model forward in time**. This is basically just like we did manually with our zero-dimensional model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "#  Make a copy of the original temperature array\n",
    "initial = mymodel.Ts.copy()\n",
    "#  Take a single timestep forward!\n",
    "mymodel.step_forward()\n",
    "#  Check out the difference\n",
    "print mymodel.Ts - initial"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Looks like the temperature got a bit colder near the equator and warmer near the poles"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "#  How long is a single timestep?\n",
    "mymodel.timestep"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This value is in seconds. It is actually 1/90th of a year (so, to step forward one year, we need 90 individual steps). This is a default value -- we could change it if we wanted to."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Actually we can see a bunch of important parameters for this model all together in a dictionary called `param`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "mymodel.param"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Accessing the model diagnostics\n",
    "\n",
    "Each model can have lots of different **diagnostic quantities**. This means anything that can be calculated from the current model state -- in our case, from the temperature.\n",
    "\n",
    "For the EBM, this includes (among other things) the OLR and the ASR."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "#  Each climlab model has a dictionary called diagnostics.\n",
    "#  Let's look at the names of the fields in this dicionary:\n",
    "mymodel.diagnostics.keys()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "#  We can access individual fields either through standard dictionary notation:\n",
    "mymodel.diagnostics['ASR']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "#  Or using the more interactive 'dot' notation:\n",
    "mymodel.ASR"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "#  Let's use the diagnostics to make a plot of the current state of the radiation\n",
    "\n",
    "plt.plot(mymodel.lat, mymodel.ASR, label='ASR')\n",
    "plt.plot(mymodel.lat, mymodel.OLR, label='OLR')\n",
    "plt.xlabel('Latitude')\n",
    "plt.ylabel('W/m2')\n",
    "plt.legend()\n",
    "plt.grid()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This plot shows that $ASR > OLR$ (system is gaining extra energy) across the tropics, and $ASR < OLR$ (system is losing energy) near the poles.\n",
    "\n",
    "It's interesting to ask how close this model is to **global energy balance**.\n",
    "\n",
    "For that, we have to take an **area-weighted global average**."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "climlab.global_mean(mymodel.net_radiation)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "climlab.global_mean(mymodel.Ts)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Running the model out to equilibrium\n",
    "\n",
    "We write a loop to integrate the model through many timesteps, and then check again for energy balance:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "#  Loop 90 times for 1 year of simulation\n",
    "for n in range(90):\n",
    "    mymodel.step_forward()\n",
    "\n",
    "print 'Global mean temperature is %0.1f degrees C.' %climlab.global_mean(mymodel.Ts)\n",
    "print 'Global mean energy imbalance is %0.1f W/m2.' %climlab.global_mean(mymodel.net_radiation)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Since there is still a significant imbalance, we are not yet at equilibrium. We should step forward again.\n",
    "\n",
    "This time, let's use a convenient built-in function instead of writing our own loop:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "mymodel.integrate_years(1.)\n",
    "\n",
    "print 'Global mean temperature is %0.1f degrees C.' %climlab.global_mean(mymodel.Ts)\n",
    "print 'Global mean energy imbalance is %0.1f W/m2.' %climlab.global_mean(mymodel.net_radiation)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We are now quite close to equilibrium. Let's make some plots"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "plt.plot(mymodel.lat, mymodel.Ts)\n",
    "plt.xlabel('Latitude')\n",
    "plt.ylabel('Temperature (deg C)')\n",
    "plt.grid()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "plt.plot(mymodel.lat_bounds, mymodel.heat_transport())\n",
    "plt.xlabel('Latitude')\n",
    "plt.ylabel('PW')\n",
    "plt.grid()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's make some nice plots of all the terms in the energy budget.\n",
    "\n",
    "We'll define a reusable function to make the plots."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def ebm_plot(e, return_fig=False):    \n",
    "    templimits = -20,32\n",
    "    radlimits = -340, 340\n",
    "    htlimits = -6,6\n",
    "    latlimits = -90,90\n",
    "    lat_ticks = np.arange(-90,90,30)\n",
    "    \n",
    "    fig = plt.figure(figsize=(8,12))\n",
    "\n",
    "    ax1 = fig.add_subplot(3,1,1)\n",
    "    ax1.plot(e.lat, e.Ts)\n",
    "    ax1.set_ylim(templimits)\n",
    "    ax1.set_ylabel('Temperature (deg C)')\n",
    "    \n",
    "    ax2 = fig.add_subplot(3,1,2)\n",
    "    ax2.plot(e.lat, e.ASR, 'k--', label='SW' )\n",
    "    ax2.plot(e.lat, -e.OLR, 'r--', label='LW' )\n",
    "    ax2.plot(e.lat, e.net_radiation, 'c-', label='net rad' )\n",
    "    ax2.plot(e.lat, e.heat_transport_convergence(), 'g--', label='dyn' )\n",
    "    ax2.plot(e.lat, e.net_radiation.squeeze() + e.heat_transport_convergence(), 'b-', label='total' )\n",
    "    ax2.set_ylim(radlimits)\n",
    "    ax2.set_ylabel('Energy budget (W m$^{-2}$)')\n",
    "    ax2.legend()\n",
    "    \n",
    "    ax3 = fig.add_subplot(3,1,3)\n",
    "    ax3.plot(e.lat_bounds, e.heat_transport() )\n",
    "    ax3.set_ylim(htlimits)\n",
    "    ax3.set_ylabel('Heat transport (PW)')\n",
    "    \n",
    "    for ax in [ax1, ax2, ax3]:\n",
    "        ax.set_xlabel('Latitude')\n",
    "        ax.set_xlim(latlimits)\n",
    "        ax.set_xticks(lat_ticks)\n",
    "        ax.grid()\n",
    "    \n",
    "    if return_fig:\n",
    "        return fig"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "ebm_plot(mymodel)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## What if we start from a very different initial temperature?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "model2 = climlab.EBM_annual(num_lat=30)\n",
    "print model2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "#  The default initial temperature\n",
    "print model2.Ts"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "#  Now let's change it to be 15 degrees everywhere\n",
    "model2.Ts[:] = 15.\n",
    "\n",
    "print model2.Ts"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "model2.compute_diagnostics()\n",
    "\n",
    "ebm_plot(model2)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Why is the heat transport zero everywhere?"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now let's look at an animation showing how all the terms in the energy budget evolve from this initial condition toward equilibrium.\n",
    "\n",
    "The code below generates a series of snapshots, that we can stitch together into an animation with a graphics program."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "#  Code to generate individual frames for the animation\n",
    "\n",
    "#  You need to specify a valid path for your computer, or else this won't work\n",
    "\n",
    "#  Uncomment the code below to re-create the frames from the animation\n",
    "\n",
    "#nsteps = 90\n",
    "#for n in range(nsteps):\n",
    "#    fig = ebm_plot(model2, return_fig=True)\n",
    "#    \n",
    "#    filepath = '/Users/br546577/Desktop/temp/ebm_animation_frames/'\n",
    "#    filename = filepath + 'ebm_animation' + str(n).zfill(4) + '.png'\n",
    "#    fig.savefig(filename)\n",
    "#    plt.close()\n",
    "#    \n",
    "#    model2.step_forward()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We will look at the animation together in class."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 2",
   "language": "python",
   "name": "python2"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 2
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython2",
   "version": "2.7.11"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}