LCOA R script + python hook

Author: sangeetabhatia03

src/scripts/lcoa_inputs_from_tlo_analyses/analysis_effect_of_treatment_ids.py

@@ -14,15 +14,12 @@
 import pickle
 from pathlib import Path
 import pandas as pd
-import numpy as np
-from matplotlib import pyplot as plt
+
+
 from tlo import Date
 from tlo.util import create_age_range_lookup
 
-from scripts.lcoa_inputs_from_tlo_analyses.fig_utils import (
-    do_bar_plot_with_ci,
-    plot_multiindex_dot_with_interval,
-)
+
 from scripts.lcoa_inputs_from_tlo_analyses.results_processing_utils import (
     get_counts_of_appts,
     get_counts_of_hsi_by_short_treatment_id,
@@ -63,7 +60,6 @@
     extract_results,
     get_color_short_treatment_id,
     make_age_grp_lookup,
-    squarify_neat,
     summarize,
 )
 # python src/scripts/lcoa_inputs_from_tlo_analyses/analysis_effect_of_treatment_ids.py outputs/s.bhatia@imperial.ac.uk/effect_of_each_treatment_id-2026-02-12T120859Z figs/ --target-start=2010-01-01 --target-end=2025-12-31
@@ -339,6 +335,10 @@ def apply(
         pipe(set_param_names_as_column_index_level_0, param_names=param_names)
     )
 
+    capacity_used_by_cadre = (
+        compute_summary_statistics(capacity_used_by_cadre, central_measure='median')
+    )
+
     results['dalys'] = dalys
     results['dalys_averted'] = dalys_averted if do_comparison else None
     results['pc_dalys_averted'] = pc_dalys_averted if do_comparison else None

src/scripts/lcoa_inputs_from_tlo_analyses/fig_utils.py

@@ -437,6 +437,65 @@ def do_label_barh_plot(_df: pd.DataFrame, _ax):
                 size=7,
             )
 
+def plot_cadre_time_by_draw_stacked(
+    _df: pd.DataFrame,
+    stat: str = "central",
+    figsize: tuple[float, float] | None = None,
+):
+    """Plot horizontal stacked bars of cadre time use by draw for one summary stat."""
+    if not isinstance(_df.columns, pd.MultiIndex) or _df.columns.nlevels != 2:
+        raise ValueError("_df columns must be a 2-level MultiIndex with levels for draw and stat.")
+
+    stat_level_name = "stat" if "stat" in _df.columns.names else _df.columns.names[1]
+    available_stats = pd.Index(_df.columns.get_level_values(stat_level_name).unique())
+    if stat not in available_stats:
+        raise ValueError(f"Statistic '{stat}' not found. Available stats: {available_stats.tolist()}")
+
+    _plot = _df.xs(stat, axis=1, level=stat_level_name).T.fillna(0.0)
+    if _plot.empty:
+        raise ValueError(f"No plottable data remain for stat '{stat}'.")
+
+    _plot = _plot.loc[_plot.sum(axis=1).sort_values(ascending=True).index]
+
+    if figsize is None:
+        fig_height = max(6, min(0.35 * len(_plot.index) + 3, 20))
+        figsize = (12, fig_height)
+    fig, ax = plt.subplots(figsize=figsize)
+
+    cadre_colors = list(plt.get_cmap("tab10").colors)
+    left = np.zeros(len(_plot.index), dtype=float)
+    y = np.arange(len(_plot.index))
+
+    for i, cadre in enumerate(_plot.columns):
+        values = _plot[cadre].to_numpy(dtype=float)
+        ax.barh(
+            y,
+            values,
+            left=left,
+            color=cadre_colors[i % len(cadre_colors)],
+            label=str(cadre),
+        )
+        left += values
+
+    ax.set_yticks(y)
+    ax.set_yticklabels([str(draw) for draw in _plot.index])
+    ax.set_xlabel("Time used")
+    ax.set_ylabel("Draw")
+    ax.grid(axis="x")
+    ax.spines["top"].set_visible(False)
+    ax.spines["right"].set_visible(False)
+    ax.legend(
+        loc="lower right",
+        fontsize=12,
+        handlelength=2.4,
+        handleheight=1.6,
+        borderpad=1.0,
+        labelspacing=0.8,
+        frameon=True,
+    )
+    fig.tight_layout()
+    return fig, ax
+
 def plot_hsi_counts_stacked_bar(_df: pd.DataFrame, plot_stat: str = "central"):
     """Plot horizontal stacked bars of HSI counts by draw for a selected summary statistic."""
     if not isinstance(_df.columns, pd.MultiIndex) or _df.columns.nlevels != 2:

src/scripts/lcoa_inputs_from_tlo_analyses/figures_effect_of_treatment_ids.py

@@ -17,6 +17,7 @@
     make_graph_file_name,
     do_barh_plot_with_ci,
     do_bar_plot_with_ci,
+    plot_cadre_time_by_draw_stacked,
     plot_deaths_by_period_for_cause,
     plot_deaths_by_period_for_draw,
     plot_hsi_counts_by_period_for_draw,
@@ -33,37 +34,45 @@
 def load_results_files(results_files: list[Path]) -> dict[Path, dict]:
     loaded = {}
     for results_file in results_files:
+        print(f"Loading results file: {results_file}")
         with open(results_file, "rb") as f:
             loaded[results_file] = pickle.load(f)
     return loaded
 
 
 def apply(results_files: list[Path], output_folder: Path, resourcefilepath: Path = None):
     """Produce standard plots describing effect of each TREATMENT_ID."""
+    print("Starting figure generation for treatment-ID effects.")
+    print(f"Output folder: {output_folder}")
 
     param_names = get_parameter_names_from_scenario_file()
+    print(f"Loaded parameter names: {len(param_names)}")
 
     all_results = load_results_files(results_files)
     primary_results = all_results[results_files[0]]
+    print(f"Using primary results from: {results_files[0]}")
 
     num_deaths_averted = primary_results.get('num_deaths_averted')
     pc_deaths_averted = primary_results.get('pc_deaths_averted')
-    num_dalys_averted = primary_results.get('num_dalys_averted')
+    dalys_averted = primary_results.get('dalys_averted')
     pc_dalys_averted = primary_results.get('pc_dalys_averted')
     icers = primary_results.get('icers_summarized')
     comparison_metrics_available = all(
         metric is not None
         for metric in (
             num_deaths_averted,
             pc_deaths_averted,
-            num_dalys_averted,
+            dalys_averted,
             pc_dalys_averted,
             icers,
         )
     )
+    print(f"Comparison metrics available: {comparison_metrics_available}")
 
     counts_of_hsi_in_implementation_period = primary_results['counts_of_hsi_by_period']
     counts_of_hsi_in_implementation_period = counts_of_hsi_in_implementation_period.drop(['2010-2041'], level=1)
+    capacity_used_by_cadre = primary_results.get("capacity_used_by_cadre")
+
 
     result_df_by_period = pd.DataFrame([
         {'treatment_id_included': draw, 'nonzero_hsis': treatment_id, 'period': period}
@@ -82,18 +91,48 @@ def apply(results_files: list[Path], output_folder: Path, resourcefilepath: Path
         if param == "Nothing":
             continue
         draw = format_scenario_name(param)
+        print(f"Plotting yearly HSI counts for draw: {draw}")
         name_of_plot = f"Yearly HSI counts for {draw}"
+        # Since all HSIs will be delivered before the service availability switch
+        # retain only the treatment id of interest in this period to avoid plot
+        # clutter.
+        pre_switch_periods = (
+            ['2010-2010', '2011-2011', '2012-2012', '2013-2013',
+             '2014-2014', '2015-2015', '2016-2016', '2017-2017',
+             '2018-2018', '2019-2019', '2020-2020', '2021-2021',
+             '2022-2022', '2023-2023', '2024-2024', '2025-2025']
+        )
+        mask_other_periods = (
+            ~counts_of_hsi_in_implementation_period.
+            index.
+            get_level_values("period").
+            isin(pre_switch_periods)
+        )
+        mask_early_periods = (
+            counts_of_hsi_in_implementation_period.index.get_level_values("period").isin(pre_switch_periods) &
+            (counts_of_hsi_in_implementation_period.index.get_level_values("appt_type") == draw.replace("_*", ""))
+        )
+        plot_this = counts_of_hsi_in_implementation_period[mask_other_periods | mask_early_periods]
         fig, ax = plot_hsi_counts_by_period_for_draw(
-            counts_of_hsi_in_implementation_period,
+            plot_this,
             draw,
         )
         ax.set_title(name_of_plot)
         outfile = os.path.join(output_folder, make_graph_file_name(name_of_plot))
         fig.savefig(outfile)
         plt.close(fig)
 
+    print("Plotting capacity used by cadres across draws.")
+    fig, ax = plot_cadre_time_by_draw_stacked(capacity_used_by_cadre, stat="central")
+    name_of_plot = "Capacity Used by Cadres (2026-2040)"
+    ax.set_title(name_of_plot)
+    outfile = os.path.join(output_folder, make_graph_file_name(name_of_plot))
+    fig.savefig(outfile)
+    plt.close(fig)
+
     # Plot population growth
     total_population_in_implementation = primary_results['total_population_by_year']
+    print("Plotting population size by year.")
     fig, ax = plot_population_by_year(total_population_in_implementation / 1e6)
     name_of_plot = "Population size by year"
     ax.set_title(name_of_plot)
@@ -102,14 +141,14 @@ def apply(results_files: list[Path], output_folder: Path, resourcefilepath: Path
     plt.close(fig)
 
     # Plot number of deaths and DALYS by cause for each parameter, with confidence intervals, for the target period
-
-
-    num_dalys_by_cause_label_implementation = primary_results['num_dalys'].drop(['2010-2041'], level=1)
+    num_dalys_by_cause_label_implementation = primary_results['dalys'].drop(['2010-2041'], level=1)
 
     num_deaths_by_cause_label_implementation = primary_results['num_deaths'].drop(['2010-2041'], level=1)
+    print("Prepared deaths and DALYs by cause for plotting.")
 
     for param in param_names:
         draw = format_scenario_name(param)
+        print(f"Plotting deaths over time by cause for draw: {draw}")
         fig, ax = plot_deaths_by_period_for_draw(
             num_deaths_by_cause_label_implementation / 1e3,
             draw,
@@ -123,6 +162,7 @@ def apply(results_files: list[Path], output_folder: Path, resourcefilepath: Path
 
     cause_labels = num_deaths_by_cause_label_implementation.index.get_level_values("label").unique()
     for cause_label in cause_labels:
+        print(f"Plotting cause-specific time series for: {cause_label}")
         fig, ax = plot_deaths_by_period_for_cause(
             num_deaths_by_cause_label_implementation / 1e3,
             cause_label=cause_label,
@@ -146,6 +186,7 @@ def apply(results_files: list[Path], output_folder: Path, resourcefilepath: Path
         plt.close(fig)
 
     if comparison_metrics_available:
+        print("Plotting comparison metrics: deaths/DALYs averted, percentages, and ICERs.")
         deaths_averted_sorted = (num_deaths_averted.sort_values(by="central", ascending=True) / 1e3)
         fig_height = max(6, min(0.28 * len(deaths_averted_sorted.index) + 4, 18))
         fig, ax = plt.subplots(figsize=(10, fig_height))
@@ -160,8 +201,9 @@ def apply(results_files: list[Path], output_folder: Path, resourcefilepath: Path
         fig.tight_layout()
         fig.savefig(outfile)
         plt.close(fig)
+        print("Saved: Deaths Averted by Each Treatment ID")
 
-        dalys_averted_sorted = (num_dalys_averted.sort_values(by="central", ascending=True) / 1e3)
+        dalys_averted_sorted = (dalys_averted.sort_values(by="central", ascending=True) / 1e3)
         fig_height = max(6, min(0.28 * len(dalys_averted_sorted.index) + 4, 18))
         fig, ax = plt.subplots(figsize=(10, fig_height))
         name_of_plot = "DALYS Averted by Each Treatment ID"
@@ -175,6 +217,7 @@ def apply(results_files: list[Path], output_folder: Path, resourcefilepath: Path
         fig.tight_layout()
         fig.savefig(outfile)
         plt.close(fig)
+        print("Saved: DALYS Averted by Each Treatment ID")
 
         pc_deaths_averted_sorted = (pc_deaths_averted.sort_values(by="central", ascending=True))
         fig_height = max(6, min(0.28 * len(pc_deaths_averted_sorted.index) + 4, 18))
@@ -190,6 +233,7 @@ def apply(results_files: list[Path], output_folder: Path, resourcefilepath: Path
         fig.tight_layout()
         fig.savefig(outfile)
         plt.close(fig)
+        print("Saved: Percentage Deaths Averted by Each Treatment ID")
 
         pc_dalys_averted_sorted = (pc_dalys_averted.sort_values(by="central", ascending=True))
         fig_height = max(6, min(0.28 * len(pc_dalys_averted_sorted.index) + 4, 18))
@@ -205,6 +249,7 @@ def apply(results_files: list[Path], output_folder: Path, resourcefilepath: Path
         fig.tight_layout()
         fig.savefig(outfile)
         plt.close(fig)
+        print("Saved: Percentage DALYs Averted by Each Treatment ID")
 
         icers_sorted = icers.sort_values(by="central", ascending=True)
         # Do not plot treatment ids with very wide uncertainty
@@ -227,6 +272,9 @@ def apply(results_files: list[Path], output_folder: Path, resourcefilepath: Path
         fig.tight_layout()
         fig.savefig(outfile)
         plt.close(fig)
+        print("Saved: ICERs for Each Treatment ID")
+
+    print("Finished generating figures.")
 
 if __name__ == "__main__":
     parser = argparse.ArgumentParser()