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).

### Like this:

Like Loading...

*Related*