From 1ee790ebf15bfd7c04b3b87e23837b0547a8464e Mon Sep 17 00:00:00 2001
From: Matt Davis <jiffyclub@gmail.com>
Date: Tue, 15 Jul 2014 13:48:26 -0700
Subject: [PATCH] Add simulation.run function for running models across years.

---
 urbansim/sim/simulation.py            | 39 +++++++++++++++++++++++----
 urbansim/sim/tests/test_simulation.py | 39 +++++++++++++++++++++++++++
 2 files changed, 73 insertions(+), 5 deletions(-)

diff --git a/urbansim/sim/simulation.py b/urbansim/sim/simulation.py
index 61702935..842c53c3 100644
--- a/urbansim/sim/simulation.py
+++ b/urbansim/sim/simulation.py
@@ -1,3 +1,5 @@
+from __future__ import print_function
+
 import inspect
 from collections import Callable
 
@@ -106,7 +108,7 @@ class _TableFuncWrapper(object):
     def __init__(self, name, func):
         self.name = name
         self._func = func
-        self._arg_list = inspect.getargspec(func).args
+        self._arg_list = set(inspect.getargspec(func).args)
         self._columns = []
 
     @property
@@ -157,7 +159,7 @@ def __init__(self, table_name, column_name, func):
         self.table_name = table_name
         self.name = column_name
         self._func = func
-        self._arg_list = inspect.getargspec(func).args
+        self._arg_list = set(inspect.getargspec(func).args)
 
     def __call__(self):
         kwargs = {t: get_table(t) for t in self._arg_list}
@@ -202,10 +204,12 @@ class _ModelFuncWrapper(object):
     def __init__(self, model_name, func):
         self.name = model_name
         self._func = func
-        self._arg_list = inspect.getargspec(func).args
+        self._arg_list = set(inspect.getargspec(func).args)
 
-    def __call__(self):
-        kwargs = {t: get_table(t) for t in self._arg_list}
+    def __call__(self, year=None):
+        kwargs = {t: get_table(t) for t in self._arg_list if t != 'year'}
+        if 'year' in self._arg_list:
+            kwargs['year'] = year
         return self._func(**kwargs)
 
 
@@ -355,6 +359,10 @@ def add_model(model_name, func):
     """
     Add a model function to the simulation.
 
+    Model argument names are used for injecting known tables of the same name.
+    The argument name "year" may be used to have the current simulation
+    year injected.
+
     Parameters
     ----------
     model_name : str
@@ -392,3 +400,24 @@ def get_model(model_name):
         return _MODELS[model_name]
     else:
         raise KeyError('no model named {}'.format(model_name))
+
+
+def run(models, years=None):
+    """
+    Run models in series, optionally repeatedly over some years.
+
+    Parameters
+    ----------
+    models : list of str
+        List of models to run identified by their name.
+
+    """
+    years = years or [None]
+
+    for year in years:
+        if year:
+            print('Running year {}'.format(year))
+        for model_name in models:
+            print('Running model {}'.format(model_name))
+            model = get_model(model_name)
+            model(year=year)
diff --git a/urbansim/sim/tests/test_simulation.py b/urbansim/sim/tests/test_simulation.py
index 85e500b2..bd89ce68 100644
--- a/urbansim/sim/tests/test_simulation.py
+++ b/urbansim/sim/tests/test_simulation.py
@@ -3,6 +3,7 @@
 from pandas.util import testing as pdt
 
 from .. import simulation as sim
+from ...utils.testing import assert_frames_equal
 
 
 @pytest.fixture
@@ -127,3 +128,41 @@ def test_model(test_table):
             {'a': [5, 7, 9],
              'b': [4, 5, 6]},
             index=['x', 'y', 'z']))
+
+
+def test_model_run(df, clear_sim):
+    sim.add_table('test_table', df)
+
+    @sim.table('table_func')
+    def table_func(test_table):
+        tt = test_table.to_frame()
+        tt['c'] = [7, 8, 9]
+        return tt
+
+    @sim.column('table_func', 'new_col')
+    def new_col(test_table, table_func):
+        tt = test_table.to_frame()
+        tf = table_func.to_frame(columns=['c'])
+        return tt['a'] + tt['b'] + tf['c']
+
+    @sim.model('test_model1')
+    def test_model1(year, test_table, table_func):
+        tf = table_func.to_frame(columns=['new_col'])
+        test_table[year] = tf['new_col'] + year
+
+    @sim.model('test_model2')
+    def test_model2(test_table):
+        tt = test_table.to_frame()
+        test_table['a'] = tt['a'] ** 2
+
+    sim.run(models=['test_model1', 'test_model2'], years=[2000, 3000])
+
+    test_table = sim.get_table('test_table')
+    assert_frames_equal(
+        test_table.to_frame(),
+        pd.DataFrame(
+            {'a': [1, 16, 81],
+             'b': [4, 5, 6],
+             2000: [2012, 2015, 2018],
+             3000: [3012, 3017, 3024]},
+            index=['x', 'y', 'z']))