Summarise function to summarise results -- with more flexibilitythan previous utility function

src/tlo/analysis/utils.py

@@ -10,7 +10,7 @@
 from collections.abc import Mapping
 from pathlib import Path
 from types import MappingProxyType
-from typing import Callable, Dict, Iterable, List, Optional, TextIO, Tuple, Union
+from typing import Callable, Dict, Iterable, List, Optional, TextIO, Tuple, Union, Literal
 
 import git
 import matplotlib.colors as mcolors
@@ -306,43 +306,90 @@ def generate_series(dataframe: pd.DataFrame) -> pd.Series:
     return _concat
 
 
-def summarize(results: pd.DataFrame, only_mean: bool = False, collapse_columns: bool = False) -> pd.DataFrame:
+def summarise(
+    results: pd.DataFrame,
+    central_measure: Literal["mean", "median"] = "median",
+    width_of_range: float = 0.95,
+    only_central: bool = False,
+    collapse_columns: bool = False,
+) -> pd.DataFrame:
     """Utility function to compute summary statistics
 
-    Finds mean value and 95% interval across the runs for each draw.
+    Finds a central value and a specified interval across the runs for each draw. By default, this uses a central
+     measure of the median and a 95% interval range.
+
+    :Param: results: The pd.DataFame of results.
+    :Param: central_measure: The name of the central measure to use - either 'mean' or 'median'.
+    :Param: width_of_range: The width of the range to compute the statistics (e.g. 0.95 for the 95% interval).
+    :Param: collapse_columns: Whether to simplify the columnar index if there is only one run (cannot be done otherwise)
+    :Param: only_central: Whether to only report the central value (dropping the range).
+
     """
+    stats = dict()
+
+    if central_measure == 'mean':
+        stats.update({'central': results.groupby(axis=1, by='draw', sort=False).mean()})
+    elif central_measure == 'median':
+        stats.update({'central': results.groupby(axis=1, by='draw', sort=False).median()})
+    else:
+        raise ValueError(f"Unknown stat: {central_measure}")
 
-    summary = pd.concat(
+    stats.update(
         {
-            'mean': results.groupby(axis=1, by='draw', sort=False).mean(),
-            'lower': results.groupby(axis=1, by='draw', sort=False).quantile(0.025),
-            'upper': results.groupby(axis=1, by='draw', sort=False).quantile(0.975),
-        },
-        axis=1
+            'lower': results.groupby(axis=1, by='draw', sort=False).quantile((1.-width_of_range)/2.),
+            'upper': results.groupby(axis=1, by='draw', sort=False).quantile(1.-(1.-width_of_range)/2.),
+        }
     )
+
+    summary = pd.concat(stats, axis=1)
     summary.columns = summary.columns.swaplevel(1, 0)
     summary.columns.names = ['draw', 'stat']
-    summary = summary.sort_index(axis=1)
+    summary = summary.sort_index(axis=1).reindex(columns=['lower', 'central', 'upper'], level=1)
 
-    if only_mean and (not collapse_columns):
+    if only_central and (not collapse_columns):
         # Remove other metrics and simplify if 'only_mean' across runs for each draw is required:
-        om: pd.DataFrame = summary.loc[:, (slice(None), "mean")]
+        om: pd.DataFrame = summary.loc[:, (slice(None), "central")]
         om.columns = [c[0] for c in om.columns.to_flat_index()]
         om.columns.name = 'draw'
         return om
 
     elif collapse_columns and (len(summary.columns.levels[0]) == 1):
         # With 'collapse_columns', if number of draws is 1, then collapse columns multi-index:
         summary_droppedlevel = summary.droplevel('draw', axis=1)
-        if only_mean:
-            return summary_droppedlevel['mean']
+        if only_central:
+            return summary_droppedlevel['central']
         else:
             return summary_droppedlevel
 
     else:
         return summary
 
 
+def summarize(
+    results: pd.DataFrame,
+    only_mean: bool = False,
+    collapse_columns: bool = False
+):
+    """Utility function to compute summary statistics
+
+    Finds mean value and 95% interval across the runs for each draw.
+
+    NOTE: This provides the legacy functionality of `summarize` that is hard-wired to use `means` (the kwarg is
+     `only_mean` and the name of the column in the output is `mean`). Please move to using the new and more flexible
+     version of `summarize` that allows the use of medians and is flexible to allow other forms of summary measure in
+     the future.
+    """
+    output = summarise(
+        results=results,
+        central_measure='mean',
+        only_central=only_mean,
+        collapse_columns=collapse_columns,
+    )
+    if output.columns.nlevels > 1:
+        output = output.rename(columns={'central': 'mean'}, level=1)  # rename 'central' to 'mean'
+    return output
+
+
 def get_grid(params: pd.DataFrame, res: pd.Series):
     """Utility function to create the arrays needed to plot a heatmap.
 
@@ -1129,7 +1176,7 @@ def get_parameters_for_status_quo() -> Dict:
             "equip_availability": "all",  # <--- NB. Existing calibration is assuming all equipment is available
         },
     }
-    
+
 def get_parameters_for_standard_mode2_runs() -> Dict:
     """
     Returns a dictionary of parameters and their updated values to indicate

tests/test_analysis.py

@@ -22,6 +22,7 @@
     order_of_coarse_appt,
     order_of_short_treatment_ids,
     parse_log_file,
+    summarise,
     summarize,
     unflatten_flattened_multi_index_in_logging,
 )
@@ -571,7 +572,7 @@ def check_parameters(self) -> None:
     sim.simulate(end_date=Date(year_of_change + 2, 1, 1))
 
 
-def test_summarize():
+def test_summarise():
     """Check that the summarize utility function works as expected."""
 
     results_multiple_draws = pd.DataFrame(
@@ -602,10 +603,10 @@ def test_summarize():
             columns=pd.MultiIndex.from_tuples(
                 [
                     ("DrawA", "lower"),
-                    ("DrawA", "mean"),
+                    ("DrawA", "central"),
                     ("DrawA", "upper"),
                     ("DrawB", "lower"),
-                    ("DrawB", "mean"),
+                    ("DrawB", "central"),
                     ("DrawB", "upper"),
                 ],
                 names=("draw", "stat"),
@@ -618,7 +619,7 @@ def test_summarize():
                 ]
             ),
         ),
-        summarize(results_multiple_draws),
+        summarise(results_multiple_draws, central_measure='mean'),
     )
 
     # Without collapsing and only mean
@@ -628,19 +629,32 @@ def test_summarize():
             index=["TimePoint0", "TimePoint1"],
             data=np.array([[10.0, 1500.0], [10.0, 1500.0]]),
         ),
-        summarize(results_multiple_draws, only_mean=True),
+        summarise(results_multiple_draws, central_measure='mean', only_central=True),
     )
 
     # With collapsing (as only one draw)
     pd.testing.assert_frame_equal(
         pd.DataFrame(
-            columns=pd.Index(["lower", "mean", "upper"], name="stat"),
+            columns=pd.Index(["lower", "central", "upper"], name="stat"),
             index=["TimePoint0", "TimePoint1"],
             data=np.array([[0.5, 10.0, 19.5], [0.5, 10.0, 19.5], ]),
         ),
-        summarize(results_one_draw, collapse_columns=True),
+        summarise(results_one_draw, central_measure='mean', collapse_columns=True),
     )
 
+    # Check that summarize() produces legacy behaviour:
+    pd.testing.assert_frame_equal(
+        summarise(results_multiple_draws, central_measure='mean').rename(columns={'central': 'mean'}, level=1),
+        summarize(results_multiple_draws)
+    )
+    pd.testing.assert_frame_equal(
+        summarise(results_multiple_draws, central_measure='mean', only_central=True),
+        summarize(results_multiple_draws, only_mean=True)
+    )
+    pd.testing.assert_frame_equal(
+        summarise(results_one_draw, central_measure='mean', collapse_columns=True),
+        summarize(results_one_draw, collapse_columns=True)
+    )
 
 def test_control_loggers_from_same_module_independently(seed, tmpdir):
     """Check that detailed/summary loggers in the same module can configured independently."""