Merge pull request #309 from devforfu/dict_param_fix

Make batch runner variable parameters more explicit and robust

mesa/batchrunner.py

@@ -7,6 +7,8 @@
 
 """
 from itertools import product
+import copy
+
 import pandas as pd
 from tqdm import tqdm
 
@@ -23,19 +25,21 @@ class BatchRunner:
     entire DataCollector object.
 
     """
-    def __init__(self, model_cls, parameter_values, iterations=1,
-                 max_steps=1000, model_reporters=None, agent_reporters=None,
-                 display_progress=True):
+    def __init__(self, model_cls, variable_parameters, fixed_parameters=None,
+                 iterations=1, max_steps=1000, model_reporters=None,
+                 agent_reporters=None, display_progress=True):
         """ Create a new BatchRunner for a given model with the given
         parameters.
 
         Args:
             model_cls: The class of model to batch-run.
-            parameter_values: Dictionary of parameters to their values or
+            variable_parameters: Dictionary of parameters to their values or
                 ranges of values. For example:
                     {"param_1": range(5),
                      "param_2": [1, 5, 10],
                       "const_param": 100}
+            fixed_parameters: Dictionary of parameters that stay same through
+                all batch runs.
             iterations: The total number of times to run the model for each
                 combination of parameters.
             max_steps: The upper limit of steps above which each run will be halted
@@ -52,7 +56,8 @@ def __init__(self, model_cls, parameter_values, iterations=1,
         """
         self.model_cls = model_cls
         self.parameter_values = {param: self.make_iterable(vals)
-                                 for param, vals in parameter_values.items()}
+                                 for param, vals in variable_parameters.items()}
+        self.fixed_values = fixed_parameters or {}
         self.iterations = iterations
         self.max_steps = max_steps
 
@@ -72,14 +77,14 @@ def run_all(self):
         params = self.parameter_values.keys()
         param_ranges = self.parameter_values.values()
         run_count = 0
-
         if self.display_progress:
             pbar = tqdm(total=len(list(product(*param_ranges))) * self.iterations)
 
         for param_values in list(product(*param_ranges)):
             kwargs = dict(zip(params, param_values))
+            model = self._try_to_init_model(kwargs)
+
             for _ in range(self.iterations):
-                model = self.model_cls(**kwargs)
                 self.run_model(model)
                 # Collect and store results:
                 if self.model_reporters:
@@ -88,7 +93,8 @@ def run_all(self):
                 if self.agent_reporters:
                     agent_vars = self.collect_agent_vars(model)
                     for agent_id, reports in agent_vars.items():
-                        key = tuple(list(param_values) + [run_count, agent_id])
+                        key = tuple(
+                            list(param_values) + [run_count, agent_id])
                         self.agent_vars[key] = reports
                 if self.display_progress:
                     pbar.update()
@@ -126,33 +132,37 @@ def collect_agent_vars(self, model):
         return agent_vars
 
     def get_model_vars_dataframe(self):
-        """ Generate a pandas DataFrame from the model-level variables collected.
+        """ Generate a pandas DataFrame from the model-level variables
+        collected.
 
         """
-        index_col_names = list(self.parameter_values.keys())
-        index_col_names.append("Run")
-        records = []
-        for key, val in self.model_vars.items():
-            record = dict(zip(index_col_names, key))
-            for k, v in val.items():
-                record[k] = v
-            records.append(record)
-        return pd.DataFrame(records)
+        return self._prepare_report_table(self.model_vars)
 
     def get_agent_vars_dataframe(self):
         """ Generate a pandas DataFrame from the agent-level variables
         collected.
 
         """
-        index_col_names = list(self.parameter_values.keys())
-        index_col_names += ["Run", "AgentID"]
+        return self._prepare_report_table(self.agent_vars)
+
+    def _prepare_report_table(self, vars_dict):
+        """
+        Creates a dataframe from collected records and sorts it using 'Run'
+        column as a key.
+        """
+        index_cols = list(self.parameter_values.keys()) + ['Run']
+
         records = []
-        for key, val in self.agent_vars.items():
-            record = dict(zip(index_col_names, key))
-            for k, v in val.items():
-                record[k] = v
+        for k, v in vars_dict.items():
+            record = dict(zip(index_cols, k))
+            record.update(v)
             records.append(record)
-        return pd.DataFrame(records)
+
+        df = pd.DataFrame(records)
+        rest_cols = set(df.columns) - set(index_cols)
+        ordered = df[index_cols + list(sorted(rest_cols))]
+        ordered.sort_values(by='Run', inplace=True)
+        return ordered
 
     @staticmethod
     def make_iterable(val):
@@ -161,3 +171,33 @@ def make_iterable(val):
             return val
         else:
             return [val]
+
+    def _try_to_init_model(self, variable_params):
+        """
+        Attempts to instantiate a model with specific variable parameters set
+        and additional fixed parameters if any.
+
+        Args:
+            variable_params: A mapping of a specific set of variable parameters.
+        """
+        if not self.fixed_values:
+            return self.model_cls(**variable_params)
+
+        try:
+            kv = copy.deepcopy(variable_params)
+            kv.update(self.fixed_values)
+            return self.model_cls(**kv)
+
+        except TypeError:
+            import inspect
+            sig = inspect.signature(self.model_cls.__init__)
+            last_arg = list(sig.parameters.values())[-1]
+            valid_types = (last_arg.POSITIONAL_OR_KEYWORD,
+                           last_arg.VAR_POSITIONAL)
+            if last_arg.kind in valid_types:
+                variable_params[last_arg.name] = self.fixed_values
+                return self.model_cls(**variable_params)
+
+        msg = ('Cannot configure model with variable '
+               'params {} and fixed params {}')
+        raise ValueError(msg.format(variable_params, self.fixed_values))

mesa/time.py

@@ -139,7 +139,7 @@ class StagedActivation(BaseScheduler):
     shuffle_between_stages = False
     stage_time = 1
 
-    def __init__(self, model, stage_list=["step"], shuffle=False,
+    def __init__(self, model, stage_list=None, shuffle=False,
                  shuffle_between_stages=False):
         """ Create an empty Staged Activation schedule.
 
@@ -154,7 +154,7 @@ def __init__(self, model, stage_list=["step"], shuffle=False,
 
         """
         super().__init__(model)
-        self.stage_list = stage_list
+        self.stage_list = stage_list or ["step"]
         self.shuffle = shuffle
         self.shuffle_between_stages = shuffle_between_stages
         self.stage_time = 1 / len(self.stage_list)

tests/test_batchrunner.py

@@ -1,46 +1,63 @@
 """
 Test the BatchRunner
 """
+from functools import reduce
+from operator import mul
 import unittest
 
 from mesa import Agent, Model
-from mesa.batchrunner import BatchRunner
 from mesa.time import BaseScheduler
+from mesa.batchrunner import BatchRunner
+
 
 NUM_AGENTS = 7
 
 
 class MockAgent(Agent):
     """
-    Minimalistic model for testing purposes
+    Minimalistic agent implementation for testing purposes
     """
-    def __init__(self, unique_id, val):
+    def __init__(self, unique_id, model, val):
+        super().__init__(unique_id, model)
         self.unique_id = unique_id
         self.val = val
 
     def step(self):
-        """
-        increment val by 1
-        """
         self.val += 1
 
 
 class MockModel(Model):
     """
     Minimalistic model for testing purposes
     """
-    def __init__(self, model_param, agent_param):
-        """
-        Args:
-            model_param (any): parameter specific to the model
-            agent_param (int): parameter specific to the agent
-        """
-        self.schedule = BaseScheduler(None)
-        self.model_param = model_param
+    def __init__(self, variable_model_param, variable_agent_param,
+                 fixed_model_param=None, schedule=None, **kwargs):
+        super().__init__()
+        self.schedule = BaseScheduler(None) if schedule is None else schedule
+        self.variable_model_param = variable_model_param
+        self.variable_agent_param = variable_agent_param
+        self.fixed_model_param = fixed_model_param
+        self.n_agents = kwargs.get('n_agents', NUM_AGENTS)
         self.running = True
-        for i in range(NUM_AGENTS):
-            a = MockAgent(i, agent_param)
-            self.schedule.add(a)
+        self.init_agents()
+
+    def init_agents(self):
+        for i in range(self.n_agents):
+            self.schedule.add(MockAgent(i, self, self.variable_agent_param))
+
+    def step(self):
+        self.schedule.step()
+
+
+class MockDictionaryModel(Model):
+
+    def __init__(self, variable_param, fixed_params):
+        super().__init__()
+        self.variable_param = variable_param
+        self.fixed_name = fixed_params.get('fixed_name', None)
+        self.running = True
+        self.schedule = BaseScheduler(None)
+        self.schedule.add(MockAgent(1, self, 0))
 
     def step(self):
         self.schedule.step()
@@ -51,44 +68,94 @@ class TestBatchRunner(unittest.TestCase):
     Test that BatchRunner is running batches
     """
     def setUp(self):
-        """
-        Create the model and run it for some steps
-        """
-        self.model_reporter = {"model": lambda m: m.model_param}
-        self.agent_reporter = {
+        self.mock_model = MockModel
+        self.model_reporters = {
+            "reported_variable_value": lambda m: m.variable_model_param,
+            "reported_fixed_value": lambda m: m.fixed_model_param
+        }
+        self.agent_reporters = {
             "agent_id": lambda a: a.unique_id,
-            "agent_val": lambda a: a.val}
-        self.params = {
-            'model_param': range(3),
-            'agent_param': [1, 8],
+            "agent_val": lambda a: a.val
         }
+        self.variable_params = {
+            "variable_model_param": range(3),
+            "variable_agent_param": [1, 8]
+        }
+        self.fixed_params = None
         self.iterations = 17
-        self.batch = BatchRunner(
-            MockModel,
-            self.params,
+        self.max_steps = 3
+
+    def launch_batch_processing(self):
+        batch = BatchRunner(
+            self.mock_model,
+            variable_parameters=self.variable_params,
+            fixed_parameters=self.fixed_params,
             iterations=self.iterations,
-            max_steps=3,
-            model_reporters=self.model_reporter,
-            agent_reporters=self.agent_reporter)
-        self.batch.run_all()
+            max_steps=self.max_steps,
+            model_reporters=self.model_reporters,
+            agent_reporters=self.agent_reporters)
+        batch.run_all()
+        return batch
+
+    @property
+    def model_runs(self):
+        """
+        Returns total number of batch runner's iterations.
+        """
+        return (reduce(mul, map(len, self.variable_params.values())) *
+                self.iterations)
 
     def test_model_level_vars(self):
         """
         Test that model-level variable collection is of the correct size
         """
-        model_vars = self.batch.get_model_vars_dataframe()
-        rows = len(self.params['model_param']) * \
-            len(self.params['agent_param']) * \
-            self.iterations
-        assert model_vars.shape == (rows, 4)
+        batch = self.launch_batch_processing()
+        model_vars = batch.get_model_vars_dataframe()
+        expected_cols = (len(self.variable_params) +
+                         len(self.model_reporters) +
+                         1)  # extra column with run index
+
+        assert model_vars.shape == (self.model_runs, expected_cols)
 
     def test_agent_level_vars(self):
         """
         Test that agent-level variable collection is of the correct size
         """
-        agent_vars = self.batch.get_agent_vars_dataframe()
-        rows = NUM_AGENTS * \
-            len(self.params['agent_param']) * \
-            len(self.params['model_param']) * \
-            self.iterations
-        assert agent_vars.shape == (rows, 6)
+        batch = self.launch_batch_processing()
+        agent_vars = batch.get_agent_vars_dataframe()
+        expected_cols = (len(self.variable_params) +
+                         len(self.agent_reporters) +
+                         1)  # extra column with run index
+
+        assert agent_vars.shape == (self.model_runs * NUM_AGENTS, expected_cols)
+
+    def test_model_with_fixed_parameters_as_kwargs(self):
+        """
+        Test that model with fixed parameters passed like kwargs is
+        properly handled
+        """
+        self.fixed_params = {'fixed_model_param': 'Fixed', 'n_agents': 1}
+        batch = self.launch_batch_processing()
+        model_vars = batch.get_model_vars_dataframe()
+        agent_vars = batch.get_agent_vars_dataframe()
+
+        assert len(model_vars) == len(agent_vars)
+        assert len(model_vars) == self.model_runs
+        assert model_vars['reported_fixed_value'].unique() == ['Fixed']
+
+    def test_model_with_fixed_parameters_as_dict(self):
+        self.mock_model = MockDictionaryModel
+        self.model_reporters = {'reported_fixed_param': lambda m: m.fixed_name}
+        self.agent_reporters = {}
+        self.fixed_params = {'fixed_name': 'DictModel'}
+        self.variable_params = {'variable_param': [1, 2, 3]}
+
+        batch = self.launch_batch_processing()
+        model_vars = batch.get_model_vars_dataframe()
+        expected_cols = (len(self.variable_params) +
+                         len(self.model_reporters) +
+                         1)
+
+        assert model_vars.shape == (self.model_runs, expected_cols)
+        assert (model_vars['reported_fixed_param'].iloc[0] ==
+                self.fixed_params['fixed_name'])