Make batch_run pytestable by adding main() functions

examples/bank_reserves/batch_run.py

@@ -185,38 +185,16 @@ def run_model(self):
     "reserve_percent": 5,
 }
 
-if __name__ == "__main__":
+
+def main():
+    # The existing batch run logic here
     data = mesa.batch_run(
         BankReservesModel,
         br_params,
     )
     br_df = pd.DataFrame(data)
     br_df.to_csv("BankReservesModel_Data.csv")
 
-    # The commented out code below is the equivalent code as above, but done
-    # via the legacy BatchRunner class. This is a good example to look at if
-    # you want to migrate your code to use `batch_run()` from `BatchRunner`.
-    # Things to note:
-    # - You have to set "reserve_percent" in br_params to `[5]`, because the
-    #   legacy BatchRunner doesn't auto-detect that it is single-valued.
-    # - The model reporters need to be explicitly specified in the legacy
-    #   BatchRunner
-    """
-    from mesa.batchrunner import BatchRunnerMP
-    br = BatchRunnerMP(
-        BankReservesModel,
-        nr_processes=2,
-        variable_parameters=br_params,
-        iterations=2,
-        max_steps=1000,
-        model_reporters={"Data Collector": lambda m: m.datacollector},
-    )
-    br.run_all()
-    br_df = br.get_model_vars_dataframe()
-    br_step_data = pd.DataFrame()
-    for i in range(len(br_df["Data Collector"])):
-        if isinstance(br_df["Data Collector"][i], DataCollector):
-            i_run_data = br_df["Data Collector"][i].get_model_vars_dataframe()
-            br_step_data = br_step_data.append(i_run_data, ignore_index=True)
-    br_step_data.to_csv("BankReservesModel_Step_Data.csv")
-    """
+
+if __name__ == "__main__":
+    main()

examples/sugarscape_g1mt/run.py

@@ -61,47 +61,45 @@ def assess_results(results, single_agent):
 
 
 # Run the model
+def main():
+    args = sys.argv[1:]
+
+    if len(args) == 0:
+        server.launch()
+
+    elif args[0] == "-s":
+        print("Running Single Model")
+        model = SugarscapeG1mt()
+        model.run_model()
+        model_results = model.datacollector.get_model_vars_dataframe()
+        model_results["Step"] = model_results.index
+        agent_results = model.datacollector.get_agent_vars_dataframe()
+        agent_results = agent_results.reset_index()
+        assess_results(model_results, agent_results)
+
+    elif args[0] == "-b":
+        print("Conducting a Batch Run")
+        params = {
+            "width": 50,
+            "height": 50,
+            "vision_min": range(1, 4),
+            "metabolism_max": [2, 3, 4, 5],
+        }
+
+        results_batch = mesa.batch_run(
+            SugarscapeG1mt,
+            parameters=params,
+            iterations=1,
+            number_processes=1,
+            data_collection_period=1,
+            display_progress=True,
+        )
+
+        assess_results(results_batch, None)
 
-args = sys.argv[1:]
-
-if len(args) == 0:
-    server.launch()
-
-elif args[0] == "-s":
-    print("Running Single Model")
-    # instantiate the model
-    model = SugarscapeG1mt()
-    # run the model
-    model.run_model()
-    # Get results
-    model_results = model.datacollector.get_model_vars_dataframe()
-    # Convert to make similar to batch_run_results
-    model_results["Step"] = model_results.index
-    agent_results = model.datacollector.get_agent_vars_dataframe()
-    agent_results = agent_results.reset_index()
-    # assess the results
-    assess_results(model_results, agent_results)
-
-elif args[0] == "-b":
-    print("Conducting a Batch Run")
-    # Batch Run
-    params = {
-        "width": 50,
-        "height": 50,
-        "vision_min": range(1, 4),
-        "metabolism_max": [2, 3, 4, 5],
-    }
-
-    results_batch = mesa.batch_run(
-        SugarscapeG1mt,
-        parameters=params,
-        iterations=1,
-        number_processes=1,
-        data_collection_period=1,
-        display_progress=True,
-    )
-
-    assess_results(results_batch, None)
-
-else:
-    raise Exception("Option not found")
+    else:
+        raise Exception("Option not found")
+
+
+if __name__ == "__main__":
+    main()