Updates to HTM Swtiching Method to Fix Issue #1411

src/scripts/comparison_of_horizontal_and_vertical_programs/analysis_hss_elements.py

@@ -0,0 +1,278 @@
+"""Produce plots to show the impact each the healthcare system (overall health impact) when running under different
+scenarios (scenario_impact_of_healthsystem.py)"""
+
+import argparse
+import textwrap
+from pathlib import Path
+from typing import Tuple
+
+import numpy as np
+import pandas as pd
+from matplotlib import pyplot as plt
+
+from tlo import Date
+from tlo.analysis.utils import (
+    CAUSE_OF_DEATH_OR_DALY_LABEL_TO_COLOR_MAP,
+    extract_results,
+    get_color_cause_of_death_or_daly_label,
+    make_age_grp_lookup,
+    order_of_cause_of_death_or_daly_label,
+    summarize,
+)
+
+
+def apply(results_folder: Path, output_folder: Path, resourcefilepath: Path = None):
+    """Produce standard set of plots describing the effect of each TREATMENT_ID.
+    - We estimate the epidemiological impact as the EXTRA deaths that would occur if that treatment did not occur.
+    - We estimate the draw on healthcare system resources as the FEWER appointments when that treatment does not occur.
+    """
+
+    TARGET_PERIOD = (Date(2020, 1, 1), Date(2030, 12, 31))
+
+    # Definitions of general helper functions
+    make_graph_file_name = lambda stub: output_folder / f"{stub.replace('*', '_star_')}.png"  # noqa: E731
+
+    _, age_grp_lookup = make_age_grp_lookup()
+
+    def target_period() -> str:
+        """Returns the target period as a string of the form YYYY-YYYY"""
+        return "-".join(str(t.year) for t in TARGET_PERIOD)
+
+    def get_parameter_names_from_scenario_file() -> Tuple[str]:
+        """Get the tuple of names of the scenarios from `Scenario` class used to create the results."""
+        from scripts.comparison_of_horizontal_and_vertical_programs.scenario_hss_elements import \
+            HSSElements
+        e = HSSElements()
+        return tuple(e._scenarios.keys())
+
+    def get_num_deaths(_df):
+        """Return total number of Deaths (total within the TARGET_PERIOD)"""
+        return pd.Series(data=len(_df.loc[pd.to_datetime(_df.date).between(*TARGET_PERIOD)]))
+
+    def get_num_dalys(_df):
+        """Return total number of DALYS (Stacked) by label (total within the TARGET_PERIOD).
+        Throw error if not a record for every year in the TARGET PERIOD (to guard against inadvertently using
+        results from runs that crashed mid-way through the simulation.
+        """
+        years_needed = [i.year for i in TARGET_PERIOD]
+        assert set(_df.year.unique()).issuperset(years_needed), "Some years are not recorded."
+        return pd.Series(
+            data=_df
+            .loc[_df.year.between(*years_needed)]
+            .drop(columns=['date', 'sex', 'age_range', 'year'])
+            .sum().sum()
+        )
+
+    def set_param_names_as_column_index_level_0(_df):
+        """Set the columns index (level 0) as the param_names."""
+        ordered_param_names_no_prefix = {i: x for i, x in enumerate(param_names)}
+        names_of_cols_level0 = [ordered_param_names_no_prefix.get(col) for col in _df.columns.levels[0]]
+        assert len(names_of_cols_level0) == len(_df.columns.levels[0])
+        _df.columns = _df.columns.set_levels(names_of_cols_level0, level=0)
+        return _df
+
+    def find_difference_relative_to_comparison(_ser: pd.Series,
+                                               comparison: str,
+                                               scaled: bool = False,
+                                               drop_comparison: bool = True,
+                                               ):
+        """Find the difference in the values in a pd.Series with a multi-index, between the draws (level 0)
+        within the runs (level 1), relative to where draw = `comparison`.
+        The comparison is `X - COMPARISON`."""
+        return _ser \
+            .unstack(level=0) \
+            .apply(lambda x: (x - x[comparison]) / (x[comparison] if scaled else 1.0), axis=1) \
+            .drop(columns=([comparison] if drop_comparison else [])) \
+            .stack()
+
+    def do_bar_plot_with_ci(_df, annotations=None, xticklabels_horizontal_and_wrapped=False, put_labels_in_legend=True):
+        """Make a vertical bar plot for each row of _df, using the columns to identify the height of the bar and the
+         extent of the error bar."""
+
+        substitute_labels = 'ABCDEFGHIJKLMNOPQRSTUVWXYZ'
+
+        yerr = np.array([
+            (_df['mean'] - _df['lower']).values,
+            (_df['upper'] - _df['mean']).values,
+        ])
+
+        xticks = {(i + 0.5): k for i, k in enumerate(_df.index)}
+
+        # Define colormap (used only with option `put_labels_in_legend=True`)
+        cmap = plt.get_cmap("tab20")
+        rescale = lambda y: (y - np.min(y)) / (np.max(y) - np.min(y))
+        colors = list(map(cmap, rescale(np.array(list(xticks.keys()))))) if put_labels_in_legend else None
+
+        fig, ax = plt.subplots(figsize=(10, 5))
+        ax.bar(
+            xticks.keys(),
+            _df['mean'].values,
+            yerr=yerr,
+            alpha=0.8,
+            ecolor='black',
+            color=colors,
+            capsize=10,
+            label=xticks.values()
+        )
+        if annotations:
+            for xpos, ypos, text in zip(xticks.keys(), _df['upper'].values, annotations):
+                ax.text(xpos, ypos*1.15, text, horizontalalignment='center', rotation='vertical', fontsize='x-small')
+        ax.set_xticks(list(xticks.keys()))
+
+        if put_labels_in_legend:
+            # Update xticks label with substitute labels
+            # Insert legend with updated labels that shows correspondence between substitute label and original label
+            xtick_values = [letter for letter, label in zip(substitute_labels, xticks.values())]
+            xtick_legend = [f'{letter}: {label}' for letter, label in zip(substitute_labels, xticks.values())]
+            h, l = ax.get_legend_handles_labels()
+            ax.legend(h, xtick_legend, loc='center left', fontsize='small', bbox_to_anchor=(1, 0.5))
+            ax.set_xticklabels(list(xtick_values))
+        else:
+            if not xticklabels_horizontal_and_wrapped:
+                # xticklabels will be vertical and not wrapped
+                ax.set_xticklabels(list(xticks.values()), rotation=90)
+            else:
+                wrapped_labs = ["\n".join(textwrap.wrap(_lab, 20)) for _lab in xticks.values()]
+                ax.set_xticklabels(wrapped_labs)
+
+        ax.grid(axis="y")
+        ax.spines['top'].set_visible(False)
+        ax.spines['right'].set_visible(False)
+        fig.tight_layout()
+
+        return fig, ax
+
+    # %% Define parameter names
+    param_names = get_parameter_names_from_scenario_file()
+
+    # %% Quantify the health gains associated with all interventions combined.
+
+    # Absolute Number of Deaths and DALYs
+    num_deaths = extract_results(
+        results_folder,
+        module='tlo.methods.demography',
+        key='death',
+        custom_generate_series=get_num_deaths,
+        do_scaling=True
+    ).pipe(set_param_names_as_column_index_level_0)
+
+    num_dalys = extract_results(
+        results_folder,
+        module='tlo.methods.healthburden',
+        key='dalys_stacked',
+        custom_generate_series=get_num_dalys,
+        do_scaling=True
+    ).pipe(set_param_names_as_column_index_level_0)
+
+    # %% Charts of total numbers of deaths / DALYS
+    num_dalys_summarized = summarize(num_dalys).loc[0].unstack().reindex(param_names)
+    num_deaths_summarized = summarize(num_deaths).loc[0].unstack().reindex(param_names)
+
+    name_of_plot = f'Deaths, {target_period()}'
+    fig, ax = do_bar_plot_with_ci(num_deaths_summarized / 1e6)
+    ax.set_title(name_of_plot)
+    ax.set_ylabel('(Millions)')
+    fig.tight_layout()
+    ax.axhline(num_deaths_summarized.loc['Baseline', 'mean']/1e6, color='black', alpha=0.5)
+    fig.savefig(make_graph_file_name(name_of_plot.replace(' ', '_').replace(',', '')))
+    fig.show()
+    plt.close(fig)
+
+    name_of_plot = f'All Scenarios: DALYs, {target_period()}'
+    fig, ax = do_bar_plot_with_ci(num_dalys_summarized / 1e6)
+    ax.set_title(name_of_plot)
+    ax.set_ylabel('(Millions)')
+    ax.axhline(num_dalys_summarized.loc['Baseline', 'mean']/1e6, color='black', alpha=0.5)
+    fig.tight_layout()
+    fig.savefig(make_graph_file_name(name_of_plot.replace(' ', '_').replace(',', '')))
+    fig.show()
+    plt.close(fig)
+
+
+    # %% Deaths and DALYS averted relative to Status Quo
+    num_deaths_averted = summarize(
+        -1.0 *
+        pd.DataFrame(
+            find_difference_relative_to_comparison(
+                num_deaths.loc[0],
+                comparison='Baseline')
+        ).T
+    ).iloc[0].unstack().reindex(param_names).drop(['Baseline'])
+
+    pc_deaths_averted = 100.0 * summarize(
+        -1.0 *
+        pd.DataFrame(
+            find_difference_relative_to_comparison(
+                num_deaths.loc[0],
+                comparison='Baseline',
+                scaled=True)
+        ).T
+    ).iloc[0].unstack().reindex(param_names).drop(['Baseline'])
+
+    num_dalys_averted = summarize(
+        -1.0 *
+        pd.DataFrame(
+            find_difference_relative_to_comparison(
+                num_dalys.loc[0],
+                comparison='Baseline')
+        ).T
+    ).iloc[0].unstack().reindex(param_names).drop(['Baseline'])
+
+    pc_dalys_averted = 100.0 * summarize(
+        -1.0 *
+        pd.DataFrame(
+            find_difference_relative_to_comparison(
+                num_dalys.loc[0],
+                comparison='Baseline',
+                scaled=True)
+        ).T
+    ).iloc[0].unstack().reindex(param_names).drop(['Baseline'])
+
+    # DEATHS
+    name_of_plot = f'Additional Deaths Averted vs Baseline, {target_period()}'
+    fig, ax = do_bar_plot_with_ci(
+        num_deaths_averted.clip(lower=0.0),
+        annotations=[
+            f"{round(row['mean'], 0)} ({round(row['lower'], 1)}-{round(row['upper'], 1)}) %"
+            for _, row in pc_deaths_averted.clip(lower=0.0).iterrows()
+        ]
+    )
+    ax.set_title(name_of_plot)
+    ax.set_ylabel('Additional Deaths Averted')
+    fig.tight_layout()
+    fig.savefig(make_graph_file_name(name_of_plot.replace(' ', '_').replace(',', '')))
+    fig.show()
+    plt.close(fig)
+
+    # DALYS
+    name_of_plot = f'Additional DALYs Averted vs Baseline, {target_period()}'
+    fig, ax = do_bar_plot_with_ci(
+        (num_dalys_averted / 1e6).clip(lower=0.0),
+        annotations=[
+            f"{round(row['mean'])} ({round(row['lower'], 1)}-{round(row['upper'], 1)}) %"
+            for _, row in pc_dalys_averted.clip(lower=0.0).iterrows()
+        ]
+    )
+    ax.set_title(name_of_plot)
+    ax.set_ylabel('Additional DALYS Averted \n(Millions)')
+    fig.tight_layout()
+    fig.savefig(make_graph_file_name(name_of_plot.replace(' ', '_').replace(',', '')))
+    fig.show()
+    plt.close(fig)
+
+    # todo: Neaten graphs
+    # todo: Graph showing difference broken down by disease (this can be cribbed from the calcs about wealth from the
+    #  third set of analyses in the overview paper).
+    # todo: other metrics of health
+    # todo: other graphs, broken down by age/sex (this can also be cribbed from overview paper stuff)
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument("results_folder", type=Path)  # outputs/horizontal_and_vertical_programs-2024-05-16
+    args = parser.parse_args()
+
+    apply(
+        results_folder=args.results_folder,
+        output_folder=args.results_folder,
+        resourcefilepath=Path('./resources')
+    )