The figure from the last post was quite slickly generate with a little Python pattern that I want to share. It is all based on my new Python knowledge of the “with” statement, and my new mainstay of pythonic data manipulation: pandas.
For easy reference, the compartmental model from that last post looks like this:
The “with” command was a little hard for me to read up on, but I figured out a simple way to use it, the context manager:
from contextlib import contextmanager @contextmanager def columns_in(df): col_names = df.columns col_list = [df[col] for col in col_names] try: yield tuple(col_list) finally: for i, col in enumerate(col_names): df[col] = col_list[i]
With “with”, the python code for the simulation is almost exactly the mathematical description of the simulation:
with columns_in(state) as (C,D,q,r,s): for t in range(T-1): C[t+1] = C[t] + q[t] - r[t] if C[t+1] < 0: C[t+1] = 0 D[t+1] = D[t] + r[t] - s[t] q[t+1] = s[t+1-7] if t % 7 == 6 else 0 # special case, first pickup, need to deliver something if t == 6: q[t+1] = 70. r[t+1] = 8. s[t+1] = D[t] if t % 7 == 6 else 0 # special case, forget to take the diapers out if t == 34: s[t+1] = 0
Pretty cool, huh? Here is a notebook full of it (code).