{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "7b826cdf",
   "metadata": {},
   "source": [
    "# Resetting time to zero after cloning a climlab process"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6510ab40",
   "metadata": {},
   "source": [
    "Brian Rose, 2/15/2022\n",
    "\n",
    "Here are some notes on how to reset a model's internal clock to zero after cloning a process with `climlab.process_like()`\n",
    "\n",
    "_These notes may become out of date after the next revision of climlab, because the calendar object that climlab uses will likely get replaced with something more robust in the future._\n",
    "\n",
    "For posterity, this is the version of climlab we're using in this example:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "72dfc3c3",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'0.7.13'"
      ]
     },
     "execution_count": 1,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import numpy as np\n",
    "import climlab\n",
    "climlab.__version__"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "acaad8bc",
   "metadata": {},
   "source": [
    "## The climlab time dictionary"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "182d4cdb",
   "metadata": {},
   "source": [
    "Every process object contains a `time` attribute, which is just a dictionary with various counters and information about timesteps.\n",
    "\n",
    "Here we create a single-column radiation model `m1` with a timestep of 1 day, and inspect its `time` dictionary:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "ad551015",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'timestep': 86400.0,\n",
       " 'num_steps_per_year': 365.2422,\n",
       " 'day_of_year_index': 0,\n",
       " 'steps': 0,\n",
       " 'days_elapsed': 0,\n",
       " 'years_elapsed': 0,\n",
       " 'days_of_year': array([  0.,   1.,   2.,   3.,   4.,   5.,   6.,   7.,   8.,   9.,  10.,\n",
       "         11.,  12.,  13.,  14.,  15.,  16.,  17.,  18.,  19.,  20.,  21.,\n",
       "         22.,  23.,  24.,  25.,  26.,  27.,  28.,  29.,  30.,  31.,  32.,\n",
       "         33.,  34.,  35.,  36.,  37.,  38.,  39.,  40.,  41.,  42.,  43.,\n",
       "         44.,  45.,  46.,  47.,  48.,  49.,  50.,  51.,  52.,  53.,  54.,\n",
       "         55.,  56.,  57.,  58.,  59.,  60.,  61.,  62.,  63.,  64.,  65.,\n",
       "         66.,  67.,  68.,  69.,  70.,  71.,  72.,  73.,  74.,  75.,  76.,\n",
       "         77.,  78.,  79.,  80.,  81.,  82.,  83.,  84.,  85.,  86.,  87.,\n",
       "         88.,  89.,  90.,  91.,  92.,  93.,  94.,  95.,  96.,  97.,  98.,\n",
       "         99., 100., 101., 102., 103., 104., 105., 106., 107., 108., 109.,\n",
       "        110., 111., 112., 113., 114., 115., 116., 117., 118., 119., 120.,\n",
       "        121., 122., 123., 124., 125., 126., 127., 128., 129., 130., 131.,\n",
       "        132., 133., 134., 135., 136., 137., 138., 139., 140., 141., 142.,\n",
       "        143., 144., 145., 146., 147., 148., 149., 150., 151., 152., 153.,\n",
       "        154., 155., 156., 157., 158., 159., 160., 161., 162., 163., 164.,\n",
       "        165., 166., 167., 168., 169., 170., 171., 172., 173., 174., 175.,\n",
       "        176., 177., 178., 179., 180., 181., 182., 183., 184., 185., 186.,\n",
       "        187., 188., 189., 190., 191., 192., 193., 194., 195., 196., 197.,\n",
       "        198., 199., 200., 201., 202., 203., 204., 205., 206., 207., 208.,\n",
       "        209., 210., 211., 212., 213., 214., 215., 216., 217., 218., 219.,\n",
       "        220., 221., 222., 223., 224., 225., 226., 227., 228., 229., 230.,\n",
       "        231., 232., 233., 234., 235., 236., 237., 238., 239., 240., 241.,\n",
       "        242., 243., 244., 245., 246., 247., 248., 249., 250., 251., 252.,\n",
       "        253., 254., 255., 256., 257., 258., 259., 260., 261., 262., 263.,\n",
       "        264., 265., 266., 267., 268., 269., 270., 271., 272., 273., 274.,\n",
       "        275., 276., 277., 278., 279., 280., 281., 282., 283., 284., 285.,\n",
       "        286., 287., 288., 289., 290., 291., 292., 293., 294., 295., 296.,\n",
       "        297., 298., 299., 300., 301., 302., 303., 304., 305., 306., 307.,\n",
       "        308., 309., 310., 311., 312., 313., 314., 315., 316., 317., 318.,\n",
       "        319., 320., 321., 322., 323., 324., 325., 326., 327., 328., 329.,\n",
       "        330., 331., 332., 333., 334., 335., 336., 337., 338., 339., 340.,\n",
       "        341., 342., 343., 344., 345., 346., 347., 348., 349., 350., 351.,\n",
       "        352., 353., 354., 355., 356., 357., 358., 359., 360., 361., 362.,\n",
       "        363., 364., 365.]),\n",
       " 'active_now': True}"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "mystate = climlab.column_state()\n",
    "m1 = climlab.radiation.RRTMG(state=mystate, timestep=climlab.utils.constants.seconds_per_day)\n",
    "m1.time"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d7793007",
   "metadata": {},
   "source": [
    "If we take a single time step forward, some elements in this dictionary get updated:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "3d4d6024",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'timestep': 86400.0,\n",
       " 'num_steps_per_year': 365.2422,\n",
       " 'day_of_year_index': 1,\n",
       " 'steps': 1,\n",
       " 'days_elapsed': 1.0,\n",
       " 'years_elapsed': 0,\n",
       " 'days_of_year': array([  0.,   1.,   2.,   3.,   4.,   5.,   6.,   7.,   8.,   9.,  10.,\n",
       "         11.,  12.,  13.,  14.,  15.,  16.,  17.,  18.,  19.,  20.,  21.,\n",
       "         22.,  23.,  24.,  25.,  26.,  27.,  28.,  29.,  30.,  31.,  32.,\n",
       "         33.,  34.,  35.,  36.,  37.,  38.,  39.,  40.,  41.,  42.,  43.,\n",
       "         44.,  45.,  46.,  47.,  48.,  49.,  50.,  51.,  52.,  53.,  54.,\n",
       "         55.,  56.,  57.,  58.,  59.,  60.,  61.,  62.,  63.,  64.,  65.,\n",
       "         66.,  67.,  68.,  69.,  70.,  71.,  72.,  73.,  74.,  75.,  76.,\n",
       "         77.,  78.,  79.,  80.,  81.,  82.,  83.,  84.,  85.,  86.,  87.,\n",
       "         88.,  89.,  90.,  91.,  92.,  93.,  94.,  95.,  96.,  97.,  98.,\n",
       "         99., 100., 101., 102., 103., 104., 105., 106., 107., 108., 109.,\n",
       "        110., 111., 112., 113., 114., 115., 116., 117., 118., 119., 120.,\n",
       "        121., 122., 123., 124., 125., 126., 127., 128., 129., 130., 131.,\n",
       "        132., 133., 134., 135., 136., 137., 138., 139., 140., 141., 142.,\n",
       "        143., 144., 145., 146., 147., 148., 149., 150., 151., 152., 153.,\n",
       "        154., 155., 156., 157., 158., 159., 160., 161., 162., 163., 164.,\n",
       "        165., 166., 167., 168., 169., 170., 171., 172., 173., 174., 175.,\n",
       "        176., 177., 178., 179., 180., 181., 182., 183., 184., 185., 186.,\n",
       "        187., 188., 189., 190., 191., 192., 193., 194., 195., 196., 197.,\n",
       "        198., 199., 200., 201., 202., 203., 204., 205., 206., 207., 208.,\n",
       "        209., 210., 211., 212., 213., 214., 215., 216., 217., 218., 219.,\n",
       "        220., 221., 222., 223., 224., 225., 226., 227., 228., 229., 230.,\n",
       "        231., 232., 233., 234., 235., 236., 237., 238., 239., 240., 241.,\n",
       "        242., 243., 244., 245., 246., 247., 248., 249., 250., 251., 252.,\n",
       "        253., 254., 255., 256., 257., 258., 259., 260., 261., 262., 263.,\n",
       "        264., 265., 266., 267., 268., 269., 270., 271., 272., 273., 274.,\n",
       "        275., 276., 277., 278., 279., 280., 281., 282., 283., 284., 285.,\n",
       "        286., 287., 288., 289., 290., 291., 292., 293., 294., 295., 296.,\n",
       "        297., 298., 299., 300., 301., 302., 303., 304., 305., 306., 307.,\n",
       "        308., 309., 310., 311., 312., 313., 314., 315., 316., 317., 318.,\n",
       "        319., 320., 321., 322., 323., 324., 325., 326., 327., 328., 329.,\n",
       "        330., 331., 332., 333., 334., 335., 336., 337., 338., 339., 340.,\n",
       "        341., 342., 343., 344., 345., 346., 347., 348., 349., 350., 351.,\n",
       "        352., 353., 354., 355., 356., 357., 358., 359., 360., 361., 362.,\n",
       "        363., 364., 365.]),\n",
       " 'active_now': True}"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "m1.step_forward()\n",
    "m1.time"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4b0e1b04",
   "metadata": {},
   "source": [
    "In particular, `steps` has increased by 1, and `days_elapsed` is now 1.0 (using a timestep of 1 day).\n",
    "\n",
    "Let's now clone this model. Both the state and the calendar are cloned, so our new model has the same date as `m1`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "53f0c754",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'timestep': 86400.0,\n",
       " 'num_steps_per_year': 365.2422,\n",
       " 'day_of_year_index': 1,\n",
       " 'steps': 1,\n",
       " 'days_elapsed': 1.0,\n",
       " 'years_elapsed': 0,\n",
       " 'days_of_year': array([  0.,   1.,   2.,   3.,   4.,   5.,   6.,   7.,   8.,   9.,  10.,\n",
       "         11.,  12.,  13.,  14.,  15.,  16.,  17.,  18.,  19.,  20.,  21.,\n",
       "         22.,  23.,  24.,  25.,  26.,  27.,  28.,  29.,  30.,  31.,  32.,\n",
       "         33.,  34.,  35.,  36.,  37.,  38.,  39.,  40.,  41.,  42.,  43.,\n",
       "         44.,  45.,  46.,  47.,  48.,  49.,  50.,  51.,  52.,  53.,  54.,\n",
       "         55.,  56.,  57.,  58.,  59.,  60.,  61.,  62.,  63.,  64.,  65.,\n",
       "         66.,  67.,  68.,  69.,  70.,  71.,  72.,  73.,  74.,  75.,  76.,\n",
       "         77.,  78.,  79.,  80.,  81.,  82.,  83.,  84.,  85.,  86.,  87.,\n",
       "         88.,  89.,  90.,  91.,  92.,  93.,  94.,  95.,  96.,  97.,  98.,\n",
       "         99., 100., 101., 102., 103., 104., 105., 106., 107., 108., 109.,\n",
       "        110., 111., 112., 113., 114., 115., 116., 117., 118., 119., 120.,\n",
       "        121., 122., 123., 124., 125., 126., 127., 128., 129., 130., 131.,\n",
       "        132., 133., 134., 135., 136., 137., 138., 139., 140., 141., 142.,\n",
       "        143., 144., 145., 146., 147., 148., 149., 150., 151., 152., 153.,\n",
       "        154., 155., 156., 157., 158., 159., 160., 161., 162., 163., 164.,\n",
       "        165., 166., 167., 168., 169., 170., 171., 172., 173., 174., 175.,\n",
       "        176., 177., 178., 179., 180., 181., 182., 183., 184., 185., 186.,\n",
       "        187., 188., 189., 190., 191., 192., 193., 194., 195., 196., 197.,\n",
       "        198., 199., 200., 201., 202., 203., 204., 205., 206., 207., 208.,\n",
       "        209., 210., 211., 212., 213., 214., 215., 216., 217., 218., 219.,\n",
       "        220., 221., 222., 223., 224., 225., 226., 227., 228., 229., 230.,\n",
       "        231., 232., 233., 234., 235., 236., 237., 238., 239., 240., 241.,\n",
       "        242., 243., 244., 245., 246., 247., 248., 249., 250., 251., 252.,\n",
       "        253., 254., 255., 256., 257., 258., 259., 260., 261., 262., 263.,\n",
       "        264., 265., 266., 267., 268., 269., 270., 271., 272., 273., 274.,\n",
       "        275., 276., 277., 278., 279., 280., 281., 282., 283., 284., 285.,\n",
       "        286., 287., 288., 289., 290., 291., 292., 293., 294., 295., 296.,\n",
       "        297., 298., 299., 300., 301., 302., 303., 304., 305., 306., 307.,\n",
       "        308., 309., 310., 311., 312., 313., 314., 315., 316., 317., 318.,\n",
       "        319., 320., 321., 322., 323., 324., 325., 326., 327., 328., 329.,\n",
       "        330., 331., 332., 333., 334., 335., 336., 337., 338., 339., 340.,\n",
       "        341., 342., 343., 344., 345., 346., 347., 348., 349., 350., 351.,\n",
       "        352., 353., 354., 355., 356., 357., 358., 359., 360., 361., 362.,\n",
       "        363., 364., 365.]),\n",
       " 'active_now': True}"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "m2 = climlab.process_like(m1)\n",
    "m2.time"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "30f49e63",
   "metadata": {},
   "source": [
    "## What if we want to clone the state, but reset the calendar to zero?\n",
    "\n",
    "One simple hack is just to keep a copy of the initial `time` dictionary prior to taking any steps forward. \n",
    "\n",
    "We can do this with the `dict`'s built-in `copy()` method."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "bf60d314",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "After cloning, m3 has taken 1 steps, and m4 has taken 1 steps.\n",
      "After replacing the time dict, m3 has taken 1 steps, and m4 has taken 0 steps.\n"
     ]
    }
   ],
   "source": [
    "mystate2 = climlab.column_state()\n",
    "m3 = climlab.radiation.RRTMG(state=mystate2, timestep=climlab.utils.constants.seconds_per_day)\n",
    "zero_time = m3.time.copy()\n",
    "m3.step_forward()\n",
    "m4 = climlab.process_like(m3)\n",
    "# Now both m3 and m4 have the same state:\n",
    "assert m3.Ts == m4.Ts\n",
    "assert np.all(m3.Tatm == m4.Tatm)\n",
    "# And they also have the same calendar:\n",
    "print('After cloning, m3 has taken {} steps, and m4 has taken {} steps.'.format(m3.time['steps'], m4.time['steps']))\n",
    "# But we can reset the calendar for m4 as if it had never taken a step forward:\n",
    "m4.time = zero_time\n",
    "print('After replacing the time dict, m3 has taken {} steps, and m4 has taken {} steps.'.format(m3.time['steps'], m4.time['steps']))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "55ab44e2",
   "metadata": {},
   "source": [
    "Since we haven't changed any model parameters, they should both evolve exactly the same way on their next timestep so the states remain the same:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "285e0851",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "After one more step, m3 has taken 2 steps, and m4 has taken 1 steps.\n"
     ]
    }
   ],
   "source": [
    "for model in [m3, m4]:\n",
    "    model.step_forward()\n",
    "assert m3.Ts == m4.Ts\n",
    "assert np.all(m3.Tatm == m4.Tatm)\n",
    "print('After one more step, m3 has taken {} steps, and m4 has taken {} steps.'.format(m3.time['steps'], m4.time['steps']))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2a0e8455",
   "metadata": {},
   "source": [
    "But if I now change a parameter in `m4`, their states will begin to differ:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "8fa3732a",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "One step after changing S0 in m4, m3 has taken 3 steps, and m4 has taken 2 steps.\n",
      "\n",
      "Now checking to see if the states are still the same:\n"
     ]
    },
    {
     "ename": "AssertionError",
     "evalue": "",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mAssertionError\u001b[0m                            Traceback (most recent call last)",
      "Input \u001b[0;32mIn [7]\u001b[0m, in \u001b[0;36m<module>\u001b[0;34m\u001b[0m\n\u001b[1;32m      5\u001b[0m \u001b[38;5;28mprint\u001b[39m(\u001b[38;5;124m'\u001b[39m\u001b[38;5;124m'\u001b[39m)\n\u001b[1;32m      6\u001b[0m \u001b[38;5;28mprint\u001b[39m(\u001b[38;5;124m'\u001b[39m\u001b[38;5;124mNow checking to see if the states are still the same:\u001b[39m\u001b[38;5;124m'\u001b[39m)\n\u001b[0;32m----> 7\u001b[0m \u001b[38;5;28;01massert\u001b[39;00m m3\u001b[38;5;241m.\u001b[39mTs \u001b[38;5;241m==\u001b[39m m4\u001b[38;5;241m.\u001b[39mTs\n",
      "\u001b[0;31mAssertionError\u001b[0m: "
     ]
    }
   ],
   "source": [
    "m4.subprocess['SW'].S0 += 10.\n",
    "for model in [m3, m4]:\n",
    "    model.step_forward()\n",
    "print('One step after changing S0 in m4, m3 has taken {} steps, and m4 has taken {} steps.'.format(m3.time['steps'], m4.time['steps']))\n",
    "print('')\n",
    "print('Now checking to see if the states are still the same:')\n",
    "assert m3.Ts == m4.Ts"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "80fe6861",
   "metadata": {},
   "source": [
    "The assertion fails because the surface temperatures of the two states have diverged, as expected."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2af8c883",
   "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.8.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}