{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "0",
   "metadata": {},
   "source": [
    "# Run the n-layer energy balance model\n",
    "\n",
    "This notebook shows examples of extending the 3-layer energy balance model to general n.\n",
    "\n",
    "For the two and three layer cases we'll take the MLE estimates from Cummins et al. (2020) for HadGEM2-ES, and we'll use the GISS forcing. Where n > 3 the data is fake."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1",
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as pl\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import pooch\n",
    "\n",
    "from fair.energy_balance_model import EnergyBalanceModel"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2",
   "metadata": {},
   "outputs": [],
   "source": [
    "df_forcing = pd.read_csv('../tests/test_data/RFMIP_ERF_tier2_GISS-E2-1-G.csv')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3",
   "metadata": {},
   "outputs": [],
   "source": [
    "ebm3 = EnergyBalanceModel(\n",
    "    ocean_heat_capacity=[3.62, 9.47, 98.66],\n",
    "    ocean_heat_transfer=[0.54, 2.39, 0.63],\n",
    "    deep_ocean_efficacy=1.59,\n",
    "    gamma_autocorrelation=1.73,\n",
    "    sigma_xi=0.32,\n",
    "    sigma_eta=0.43,\n",
    "    forcing_4co2=6.35,\n",
    "    stochastic_run=True,\n",
    "    seed=16\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4",
   "metadata": {},
   "outputs": [],
   "source": [
    "ebm3.add_forcing(forcing = df_forcing['GISS-E2-1-G TOT'].values, timestep=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5",
   "metadata": {},
   "outputs": [],
   "source": [
    "ebm3.run()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6",
   "metadata": {},
   "outputs": [],
   "source": [
    "ebm3.temperature"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7",
   "metadata": {},
   "outputs": [],
   "source": [
    "time = np.arange(1850.5, 2101)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8",
   "metadata": {},
   "outputs": [],
   "source": [
    "pl.plot(time, ebm3.temperature[:,0], label='surface / top ocean layer')\n",
    "pl.plot(time, ebm3.temperature[:,1], label='second ocean layer')\n",
    "pl.plot(time, ebm3.temperature[:,2], label='deep ocean layer')\n",
    "pl.ylabel('K relative to 1850')\n",
    "pl.title('SSP2-4.5 temperature change')\n",
    "pl.legend()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9",
   "metadata": {},
   "outputs": [],
   "source": [
    "ebm3.emergent_parameters()\n",
    "ebm3.ecs, ebm3.tcr"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "10",
   "metadata": {},
   "outputs": [],
   "source": [
    "ebm2 = EnergyBalanceModel(\n",
    "    ocean_heat_capacity=[7.73, 89.29],\n",
    "    ocean_heat_transfer=[0.63, 0.52],\n",
    "    deep_ocean_efficacy=1.52,\n",
    "    gamma_autocorrelation=1.58,\n",
    "    sigma_xi=0.64,\n",
    "    sigma_eta=0.43,\n",
    "    stochastic_run=True,\n",
    "    forcing_4co2=6.86,\n",
    "    seed=16\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "11",
   "metadata": {},
   "outputs": [],
   "source": [
    "ebm2.add_forcing(forcing = df_forcing['GISS-E2-1-G TOT'].values, timestep=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "12",
   "metadata": {},
   "outputs": [],
   "source": [
    "ebm2.emergent_parameters()\n",
    "ebm2.ecs, ebm2.tcr"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "13",
   "metadata": {},
   "outputs": [],
   "source": [
    "ebm2.run()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "14",
   "metadata": {},
   "outputs": [],
   "source": [
    "pl.plot(time, ebm2.temperature[:,0], label='surface / top ocean layer')\n",
    "pl.plot(time, ebm2.temperature[:,1], label='deep ocean layer')\n",
    "pl.ylabel('K relative to 1850')\n",
    "pl.title('SSP2-4.5 temperature change')\n",
    "pl.legend()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "15",
   "metadata": {},
   "outputs": [],
   "source": [
    "# this is not based on a tuning to any existing CMIP6 model, but I try to get the TCR close to the \n",
    "# HadGEM2 2- and 3-layer cases.\n",
    "ebm4 = EnergyBalanceModel(\n",
    "    ocean_heat_capacity=[1.3, 9, 20, 80],\n",
    "    ocean_heat_transfer=[0.54, 3, 3, 0.63],\n",
    "    deep_ocean_efficacy=1.2,\n",
    "    gamma_autocorrelation=1.73,\n",
    "    sigma_xi=0.32,\n",
    "    sigma_eta=0.43,\n",
    "    forcing_4co2=6.35,\n",
    "    stochastic_run=True,\n",
    "    seed=16\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "16",
   "metadata": {},
   "outputs": [],
   "source": [
    "ebm4.emergent_parameters()\n",
    "ebm4.ecs, ebm4.tcr"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "17",
   "metadata": {},
   "outputs": [],
   "source": [
    "ebm4.add_forcing(forcing = df_forcing['GISS-E2-1-G TOT'].values, timestep=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "18",
   "metadata": {},
   "outputs": [],
   "source": [
    "ebm4.run()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "19",
   "metadata": {},
   "outputs": [],
   "source": [
    "pl.plot(time, ebm4.temperature[:,0], label='surface / top ocean layer')\n",
    "pl.plot(time, ebm4.temperature[:,1], label='second ocean layer')\n",
    "pl.plot(time, ebm4.temperature[:,2], label='third ocean layer')\n",
    "pl.plot(time, ebm4.temperature[:,3], label='deep ocean layer')\n",
    "pl.ylabel('K relative to 1850')\n",
    "pl.title('SSP2-4.5 temperature change')\n",
    "pl.legend()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "20",
   "metadata": {},
   "outputs": [],
   "source": [
    "# let's go totally crazy\n",
    "ebm10 = EnergyBalanceModel(\n",
    "    ocean_heat_capacity=[0.6, 1.3, 2, 5, 7, 10, 45, 70, 80, 130],\n",
    "    ocean_heat_transfer=[0.54, 4, 5, 5, 5, 5, 5, 5, 5, 0.63],\n",
    "    deep_ocean_efficacy=1.2,\n",
    "    gamma_autocorrelation=1.73,\n",
    "    sigma_xi=0.32,\n",
    "    sigma_eta=0.43,\n",
    "    forcing_4co2=6.35,\n",
    "    stochastic_run=True,\n",
    "    seed=16\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "21",
   "metadata": {},
   "outputs": [],
   "source": [
    "ebm10.emergent_parameters()\n",
    "ebm10.ecs, ebm10.tcr"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "22",
   "metadata": {},
   "outputs": [],
   "source": [
    "ebm10.add_forcing(forcing = df_forcing['GISS-E2-1-G TOT'].values, timestep=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "23",
   "metadata": {},
   "outputs": [],
   "source": [
    "ebm10.run()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "24",
   "metadata": {},
   "outputs": [],
   "source": [
    "pl.plot(time, ebm10.temperature[:,0], label='surface / top ocean layer')\n",
    "pl.plot(time, ebm10.temperature[:,1], label='second ocean layer')\n",
    "pl.plot(time, ebm10.temperature[:,2], label='third ocean layer')\n",
    "pl.plot(time, ebm10.temperature[:,3], label='fourth ocean layer')\n",
    "pl.plot(time, ebm10.temperature[:,4], label='fifth ocean layer')\n",
    "pl.plot(time, ebm10.temperature[:,5], label='sixth ocean layer')\n",
    "pl.plot(time, ebm10.temperature[:,6], label='seventh ocean layer')\n",
    "pl.plot(time, ebm10.temperature[:,7], label='eighth ocean layer')\n",
    "pl.plot(time, ebm10.temperature[:,8], label='ninth ocean layer')\n",
    "pl.plot(time, ebm10.temperature[:,9], label='deep ocean layer')\n",
    "pl.ylabel('K relative to 1850')\n",
    "pl.title('SSP2-4.5 temperature change')\n",
    "pl.legend()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "25",
   "metadata": {},
   "outputs": [],
   "source": [
    "pl.plot(time, ebm2.temperature[:,0], label='two layer model')\n",
    "pl.plot(time, ebm3.temperature[:,0], label='three layer model')\n",
    "pl.plot(time, ebm4.temperature[:,0], label='four layer model')\n",
    "pl.plot(time, ebm10.temperature[:,0], label='ten layer model')\n",
    "pl.ylabel('K relative to 1850')\n",
    "pl.title('SSP2-4.5 temperature change')\n",
    "pl.legend()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "26",
   "metadata": {},
   "outputs": [],
   "source": [
    "pl.plot(time, ebm2.toa_imbalance, label='two layer model')\n",
    "pl.plot(time, ebm3.toa_imbalance, label='three layer model')\n",
    "pl.plot(time, ebm4.toa_imbalance, label='four layer model')\n",
    "pl.plot(time, ebm10.toa_imbalance, label='ten layer model')\n",
    "pl.ylabel('W/m2 relative to 1850')\n",
    "pl.title('SSP2-4.5 TOA radiation change')\n",
    "pl.legend()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "27",
   "metadata": {},
   "source": [
    "## Repeat everything with stochastic forcing switched off"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "28",
   "metadata": {},
   "outputs": [],
   "source": [
    "ebm3 = EnergyBalanceModel(\n",
    "    ocean_heat_capacity=[3.62, 9.47, 98.66],\n",
    "    ocean_heat_transfer=[0.54, 2.39, 0.63],\n",
    "    deep_ocean_efficacy=1.59,\n",
    "    gamma_autocorrelation=1.73,\n",
    "    sigma_xi=0.32,\n",
    "    sigma_eta=0.43,\n",
    "    forcing_4co2=6.35,\n",
    "    stochastic_run=False,\n",
    "    seed=16\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "29",
   "metadata": {},
   "outputs": [],
   "source": [
    "ebm3.add_forcing(forcing = df_forcing['GISS-E2-1-G TOT'].values, timestep=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "30",
   "metadata": {},
   "outputs": [],
   "source": [
    "ebm3.run()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "31",
   "metadata": {},
   "outputs": [],
   "source": [
    "ebm3.temperature"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "32",
   "metadata": {},
   "outputs": [],
   "source": [
    "time = np.arange(1850.5, 2101)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "33",
   "metadata": {},
   "outputs": [],
   "source": [
    "pl.plot(time, ebm3.temperature[:,0], label='surface / top ocean layer')\n",
    "pl.plot(time, ebm3.temperature[:,1], label='second ocean layer')\n",
    "pl.plot(time, ebm3.temperature[:,2], label='deep ocean layer')\n",
    "pl.ylabel('K relative to 1850')\n",
    "pl.title('SSP2-4.5 temperature change')\n",
    "pl.legend()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "34",
   "metadata": {},
   "outputs": [],
   "source": [
    "ebm3.emergent_parameters()\n",
    "ebm3.ecs, ebm3.tcr"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "35",
   "metadata": {},
   "outputs": [],
   "source": [
    "ebm2 = EnergyBalanceModel(\n",
    "    ocean_heat_capacity=[7.73, 89.29],\n",
    "    ocean_heat_transfer=[0.63, 0.52],\n",
    "    deep_ocean_efficacy=1.52,\n",
    "    gamma_autocorrelation=1.58,\n",
    "    sigma_xi=0.64,\n",
    "    sigma_eta=0.43,\n",
    "    stochastic_run=False,\n",
    "    forcing_4co2=6.86,\n",
    "    seed=16\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "36",
   "metadata": {},
   "outputs": [],
   "source": [
    "ebm2.add_forcing(forcing = df_forcing['GISS-E2-1-G TOT'].values, timestep=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "37",
   "metadata": {},
   "outputs": [],
   "source": [
    "ebm2.emergent_parameters()\n",
    "ebm2.ecs, ebm2.tcr"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "38",
   "metadata": {},
   "outputs": [],
   "source": [
    "ebm2.run()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "39",
   "metadata": {},
   "outputs": [],
   "source": [
    "pl.plot(time, ebm2.temperature[:,0], label='surface / top ocean layer')\n",
    "pl.plot(time, ebm2.temperature[:,1], label='deep ocean layer')\n",
    "pl.ylabel('K relative to 1850')\n",
    "pl.title('SSP2-4.5 temperature change')\n",
    "pl.legend()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "40",
   "metadata": {},
   "outputs": [],
   "source": [
    "# this is not based on a tuning to any existing CMIP6 model, but I try to get the TCR close to the \n",
    "# HadGEM2 2- and 3-layer cases.\n",
    "ebm4 = EnergyBalanceModel(\n",
    "    ocean_heat_capacity=[1.3, 9, 20, 80],\n",
    "    ocean_heat_transfer=[0.54, 3, 3, 0.63],\n",
    "    deep_ocean_efficacy=1.2,\n",
    "    gamma_autocorrelation=1.73,\n",
    "    sigma_xi=0.32,\n",
    "    sigma_eta=0.43,\n",
    "    forcing_4co2=6.35,\n",
    "    stochastic_run=False,\n",
    "    seed=16\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "41",
   "metadata": {},
   "outputs": [],
   "source": [
    "ebm4.emergent_parameters()\n",
    "ebm4.ecs, ebm4.tcr"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "42",
   "metadata": {},
   "outputs": [],
   "source": [
    "ebm4.add_forcing(forcing = df_forcing['GISS-E2-1-G TOT'].values, timestep=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "43",
   "metadata": {},
   "outputs": [],
   "source": [
    "ebm4.run()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "44",
   "metadata": {},
   "outputs": [],
   "source": [
    "pl.plot(time, ebm4.temperature[:,0], label='surface / top ocean layer')\n",
    "pl.plot(time, ebm4.temperature[:,1], label='second ocean layer')\n",
    "pl.plot(time, ebm4.temperature[:,2], label='third ocean layer')\n",
    "pl.plot(time, ebm4.temperature[:,3], label='deep ocean layer')\n",
    "pl.ylabel('K relative to 1850')\n",
    "pl.title('SSP2-4.5 temperature change')\n",
    "pl.legend()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "45",
   "metadata": {},
   "outputs": [],
   "source": [
    "# let's go totally crazy\n",
    "ebm10 = EnergyBalanceModel(\n",
    "    ocean_heat_capacity=[0.6, 1.3, 2, 5, 7, 10, 45, 70, 80, 130],\n",
    "    ocean_heat_transfer=[0.54, 4, 5, 5, 5, 5, 5, 5, 5, 0.63],\n",
    "    deep_ocean_efficacy=1.2,\n",
    "    gamma_autocorrelation=1.73,\n",
    "    sigma_xi=0.32,\n",
    "    sigma_eta=0.43,\n",
    "    forcing_4co2=6.35,\n",
    "    stochastic_run=False,\n",
    "    seed=16\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "46",
   "metadata": {},
   "outputs": [],
   "source": [
    "ebm10.emergent_parameters()\n",
    "ebm10.ecs, ebm10.tcr"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "47",
   "metadata": {},
   "outputs": [],
   "source": [
    "ebm10.add_forcing(forcing = df_forcing['GISS-E2-1-G TOT'].values, timestep=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "48",
   "metadata": {},
   "outputs": [],
   "source": [
    "ebm10.run()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "49",
   "metadata": {},
   "outputs": [],
   "source": [
    "pl.plot(time, ebm10.temperature[:,0], label='surface / top ocean layer')\n",
    "pl.plot(time, ebm10.temperature[:,1], label='second ocean layer')\n",
    "pl.plot(time, ebm10.temperature[:,2], label='third ocean layer')\n",
    "pl.plot(time, ebm10.temperature[:,3], label='fourth ocean layer')\n",
    "pl.plot(time, ebm10.temperature[:,4], label='fifth ocean layer')\n",
    "pl.plot(time, ebm10.temperature[:,5], label='sixth ocean layer')\n",
    "pl.plot(time, ebm10.temperature[:,6], label='seventh ocean layer')\n",
    "pl.plot(time, ebm10.temperature[:,7], label='eighth ocean layer')\n",
    "pl.plot(time, ebm10.temperature[:,8], label='ninth ocean layer')\n",
    "pl.plot(time, ebm10.temperature[:,9], label='deep ocean layer')\n",
    "pl.ylabel('K relative to 1850')\n",
    "pl.title('SSP2-4.5 temperature change')\n",
    "pl.legend()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "50",
   "metadata": {},
   "outputs": [],
   "source": [
    "pl.plot(time, ebm2.temperature[:,0], label='two layer model')\n",
    "pl.plot(time, ebm3.temperature[:,0], label='three layer model')\n",
    "pl.plot(time, ebm4.temperature[:,0], label='four layer model')\n",
    "pl.plot(time, ebm10.temperature[:,0], label='ten layer model')\n",
    "pl.ylabel('K relative to 1850')\n",
    "pl.title('SSP2-4.5 temperature change')\n",
    "pl.legend()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "51",
   "metadata": {},
   "outputs": [],
   "source": [
    "pl.plot(time, ebm2.toa_imbalance, label='two layer model')\n",
    "pl.plot(time, ebm3.toa_imbalance, label='three layer model')\n",
    "pl.plot(time, ebm4.toa_imbalance, label='four layer model')\n",
    "pl.plot(time, ebm10.toa_imbalance, label='ten layer model')\n",
    "pl.ylabel('W/m2 relative to 1850')\n",
    "pl.title('SSP2-4.5 TOA radiation change')\n",
    "pl.legend()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "52",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "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.9.15"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}