move setting of default stratifications (#297)

    Category: refactor
    JIRA issue: MIC-3890

We'd like to set default stratifications in the setup of the results manager than than the results stratifier.

Testing
Ran CIFF SAM for a few time steps, tests are passing.

docs/source/tutorials/exploration.rst

@@ -163,6 +163,9 @@ configuration by simply printing it.
             component_configs: True
         extrapolate:
             component_configs: True
+    stratification:
+        default:
+            component_configs: []
 
 
 What do we see here?  The configuration is *hierarchical*.  There are a set of

src/vivarium/framework/results/manager.py

@@ -29,6 +29,12 @@ class ResultsManager:
     `collect_metrics`).
     """
 
+    configuration_defaults = {
+        "stratification": {
+            "default": [],
+        }
+    }
+
     def __init__(self):
         self._metrics = Counter()
         self._results_context = ResultsContext()
@@ -59,6 +65,8 @@ def setup(self, builder: "Builder"):
 
         self.get_value = builder.value.get_value
 
+        self.set_default_stratifications(builder)
+
         builder.value.register_value_modifier("metrics", self.get_results)
 
     def on_time_step_prepare(self, event: Event):
@@ -78,7 +86,8 @@ def gather_results(self, event_name: str, event: Event):
         for results_group in self._results_context.gather_results(population, event_name):
             self._metrics.update(results_group)
 
-    def set_default_stratifications(self, default_stratifications: List[str]):
+    def set_default_stratifications(self, builder):
+        default_stratifications = builder.configuration.stratification.default
         self._results_context.set_default_stratifications(default_stratifications)
 
     def register_stratification(

tests/framework/results/test_interface.py

@@ -72,6 +72,7 @@ def test_register_observation(
     mgr = ResultsManager()
     interface = ResultsInterface(mgr)
     builder = mocker.Mock()
+    builder.configuration.stratification.default = []
     # Set up mock builder with mocked get_value call for Pipelines
     mocker.patch.object(builder, "value.get_value")
     builder.value.get_value = MethodType(mock_get_value, builder)

tests/framework/results/test_manager.py

@@ -224,3 +224,17 @@ def test_add_observation_nop_stratifications(
     )
     for m in match:
         assert m in caplog.text
+
+
+def test_setting_default_stratifications_at_setup(mocker):
+    """Test that set default stratifications happens at setup"""
+    mgr = ResultsManager()
+    builder = mocker.Mock()
+    mgr._results_context.set_default_stratifications = mocker.Mock()
+    mgr._results_context.set_default_stratifications.assert_not_called()
+
+    mgr.setup(builder)
+
+    mgr._results_context.set_default_stratifications.assert_called_once_with(
+        builder.configuration.stratification.default
+    )

tests/framework/results/test_stratification.py

@@ -2,6 +2,7 @@
 import pandas as pd
 import pytest
 
+from vivarium.framework.results.manager import ResultsManager
 from vivarium.framework.results.stratification import Stratification
 
 from .mocks import (
@@ -140,3 +141,15 @@ def test_stratification_call_raises(
     my_stratification = Stratification(name, sources, categories, mapper, is_vectorized)
     with pytest.raises(expected_exception):
         raise my_stratification(STUDENT_TABLE)
+
+
+@pytest.mark.parametrize("default_stratifications", [["age", "sex"], ["age"], []])
+def test_setting_default_stratifications(default_stratifications, mocker):
+    """Test that default stratifications are set as expected."""
+    mgr = ResultsManager()
+    builder = mocker.Mock()
+    builder.configuration.stratification.default = default_stratifications
+
+    mgr.setup(builder)
+
+    assert mgr._results_context.default_stratifications == default_stratifications