Allow suspend/resume of simulations on Batch runs

src/scripts/dev/scenarios/suspend-resume-test.py

@@ -0,0 +1,54 @@
+"""
+This file defines a batch run of a large population for a long time with all disease modules and full use of HSIs
+It's used for calibrations (demographic patterns, health burdens and healthsystem usage)
+
+Run on the batch system using:
+```tlo batch-submit src/scripts/calibration_analyses/scenarios/long_run_all_diseases.py```
+
+or locally using:
+    ```tlo scenario-run src/scripts/calibration_analyses/scenarios/long_run_all_diseases.py```
+
+"""
+
+from tlo import Date, logging
+from tlo.analysis.utils import get_parameters_for_status_quo
+from tlo.methods.fullmodel import fullmodel
+from tlo.scenario import BaseScenario
+
+
+class SuspendResumeTest(BaseScenario):
+    def __init__(self):
+        super().__init__()
+        self.seed = 0
+        self.start_date = Date(2010, 1, 1)
+        self.end_date = Date(2012, 1, 1)  # The simulation will stop before reaching this date.
+        self.pop_size = 1_000
+        self.number_of_draws = 2
+        self.runs_per_draw = 2
+
+    def log_configuration(self):
+        return {
+            'filename': 'suspend_resume_test',  # <- (specified only for local running)
+            'directory': './outputs',  # <- (specified only for local running)
+            'custom_levels': {
+                '*': logging.WARNING,
+                'tlo.methods.demography': logging.INFO,
+                'tlo.methods.demography.detail': logging.WARNING,
+                'tlo.methods.healthburden': logging.INFO,
+                'tlo.methods.healthsystem': logging.INFO,
+                'tlo.methods.healthsystem.summary': logging.INFO,
+                "tlo.methods.contraception": logging.INFO,
+            }
+        }
+
+    def modules(self):
+        return fullmodel()
+
+    def draw_parameters(self, draw_number, rng):
+        return get_parameters_for_status_quo()
+
+
+if __name__ == '__main__':
+    from tlo.cli import scenario_run
+
+    scenario_run([__file__])

src/scripts/test_suspend_resume/analysis_extract_data.py

@@ -0,0 +1,152 @@
+"""Retrieve total DALYs per year"""
+
+import argparse
+from pathlib import Path
+from typing import Tuple
+
+import pandas as pd
+
+from tlo import Date
+from tlo.analysis.utils import extract_results
+
+# Range of years considered
+min_year = 2014
+max_year = 2020
+
+
+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(min_year, 1, 1), Date(max_year, 1, 1))
+
+    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.test_suspend_resume.scenario_test_resume_with_change import (
+            ImpactOfHealthSystemMode,
+        )
+        e = ImpactOfHealthSystemMode()
+        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)"""
+        return pd.Series(
+            data=_df
+            .loc[_df.year.between(*[i.year for i in TARGET_PERIOD])]
+            .drop(columns=['date', 'sex', 'age_range', 'year'])
+            .sum().sum()
+        )
+
+    def get_num_dalys_by_cause(_df):
+        """Return number of DALYs by cause by label (total within the TARGET_PERIOD)"""
+        return pd.Series(
+            data=_df
+            .loc[_df.year.between(*[i.year for i in TARGET_PERIOD])]
+            .drop(columns=['date', 'sex', 'age_range', 'year'])
+            .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
+
+
+    # Obtain parameter names for this scenario file
+    param_names = get_parameter_names_from_scenario_file()
+    print(param_names)
+
+    # ================================================================================================
+    # TIME EVOLUTION OF TOTAL DALYs
+
+    year_target = 2023 # This global variable will be passed to custom function
+    def get_num_dalys_by_year(_df):
+        """Return total number of DALYs (Stacked) by label (total within the TARGET_PERIOD)"""
+        return pd.Series(
+            data=_df
+            .loc[_df.year == year_target]
+            .drop(columns=['date', 'sex', 'age_range', 'year'])
+            .sum().sum()
+        )
+
+    ALL = {}
+    # Plot time trend show year prior transition as well to emphasise that until that point DALYs incurred
+    # are consistent across different policies
+    this_min_year = 2010
+    for year in range(this_min_year, max_year+1):
+        year_target = year
+        num_dalys_by_year = extract_results(
+            results_folder,
+            module='tlo.methods.healthburden',
+            key='dalys_stacked',
+            custom_generate_series=get_num_dalys_by_year,
+            do_scaling=True
+        ).pipe(set_param_names_as_column_index_level_0)
+        ALL[year_target] = num_dalys_by_year
+        print(num_dalys_by_year.columns)
+    # Concatenate the DataFrames into a single DataFrame
+    concatenated_df = pd.concat(ALL.values(), keys=ALL.keys())
+    concatenated_df.index = concatenated_df.index.set_names(['date', 'index_original'])
+    concatenated_df = concatenated_df.reset_index(level='index_original',drop=True)
+    dalys_by_year = concatenated_df
+    dalys_by_year.to_csv('ConvertedOutputs/Total_DALYs_with_time.csv', index=True)
+
+if __name__ == "__main__":
+    rfp = Path('resources')
+
+    parser = argparse.ArgumentParser(
+        description="Produce plots to show the impact each set of treatments",
+        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
+    )
+    parser.add_argument(
+        "--output-path",
+        help=(
+            "Directory to write outputs to. If not specified (set to None) outputs "
+            "will be written to value of --results-path argument."
+        ),
+        type=Path,
+        default=None,
+        required=False,
+    )
+    parser.add_argument(
+        "--resources-path",
+        help="Directory containing resource files",
+        type=Path,
+        default=Path('resources'),
+        required=False,
+    )
+    parser.add_argument(
+        "--results-path",
+        type=Path,
+        help=(
+            "Directory containing results from running "
+            "src/scripts/test_suspend_resume/scenario_test_resume_with_change.py"
+        ),
+        default=None,
+        required=False
+    )
+    args = parser.parse_args()
+    assert args.results_path is not None
+    results_path = args.results_path
+
+    output_path = results_path if args.output_path is None else args.output_path
+
+    apply(
+        results_folder=results_path,
+        output_folder=output_path,
+        resourcefilepath=args.resources_path
+    )

src/scripts/test_suspend_resume/scenario_test_resume_with_change.py

@@ -0,0 +1,147 @@
+"""This Scenario file run the model under different assumptions for HR capabilities expansion in order to estimate the
+impact that is achieved under each.
+
+Run on the batch system using:
+```
+tlo batch-submit src/scripts/healthsystem/impact_of_policy/scenario_impact_of_const_capabilities_expansion.py
+```
+
+or locally using:
+```
+tlo scenario-run src/scripts/healthsystem/impact_of_policy/scenario_impact_of_const_capabilities_expansion.py
+```
+
+"""
+from pathlib import Path
+from typing import Dict
+
+import pandas as pd
+
+from tlo import Date, logging
+from tlo.analysis.utils import get_parameters_for_status_quo, mix_scenarios
+from tlo.methods.fullmodel import fullmodel
+from tlo.scenario import BaseScenario
+
+
+class ImpactOfHealthSystemMode(BaseScenario):
+    def __init__(self):
+        super().__init__()
+        self.seed = 0
+        self.start_date = Date(2010, 1, 1)
+        self.end_date = self.start_date + pd.DateOffset(years=8)
+        self.pop_size = 20_000
+        self._scenarios = self._get_scenarios()
+        self.number_of_draws = len(self._scenarios)
+        self.runs_per_draw = 3
+
+    def log_configuration(self):
+        return {
+            'filename': 'effect_of_capabilities_scaling',
+            'directory': Path('./outputs'),  # <- (specified only for local running)
+            'custom_levels': {
+                '*': logging.WARNING,
+                'tlo.methods.demography': logging.INFO,
+                'tlo.methods.demography.detail': logging.WARNING,
+                'tlo.methods.healthburden': logging.INFO,
+                'tlo.methods.healthsystem.summary': logging.INFO,
+            }
+        }
+
+    def modules(self):
+        return (
+            fullmodel()
+        )
+
+    def draw_parameters(self, draw_number, rng):
+        if draw_number < self.number_of_draws:
+            return list(self._scenarios.values())[draw_number]
+        else:
+            return
+
+    def _get_scenarios(self) -> Dict[str, Dict]:
+        """Return the Dict with values for the parameters that are changed, keyed by a name for the scenario.
+        """
+
+        self.YEAR_OF_CHANGE = 2014
+
+        return {
+
+            # =========== STATUS QUO ============
+            "Baseline":
+                mix_scenarios(
+                    self._baseline(),
+                ),
+
+            "Mode 2 no rescaling":
+                mix_scenarios(
+                    self._baseline(),
+                    {
+                     "HealthSystem": {
+                        "mode_appt_constraints_postSwitch": 2,      # <-- Mode 2 post-change to show effects of HRH
+                      },
+                    }
+                ),
+
+             "Mode 2 with rescaling":
+                mix_scenarios(
+                    self._baseline(),
+                    {
+                     "HealthSystem": {
+                        "mode_appt_constraints_postSwitch": 2,      # <-- Mode 2 post-change to show effects of HRH
+                        "scale_to_effective_capabilities": True,
+                      },
+                    }
+                ),
+
+            "Mode 2 with rescaling and funded plus":
+                mix_scenarios(
+                    self._baseline(),
+                    {
+                     "HealthSystem": {
+                        "mode_appt_constraints_postSwitch": 2,      # <-- Mode 2 post-change to show effects of HRH
+                        "scale_to_effective_capabilities": True,
+                        "use_funded_or_actual_staffing_postSwitch": "funded_plus",
+                      },
+                    }
+                ),
+
+            "Mode 1 perfect consumables":
+                mix_scenarios(
+                    self._baseline(),
+                    {
+                     "HealthSystem": {
+                        "cons_availability_postSwitch":"all",
+                      },
+                    }
+                ),
+
+        }
+
+    def _baseline(self) -> Dict:
+        """Return the Dict with values for the parameter changes that define the baseline scenario. """
+        return mix_scenarios(
+            get_parameters_for_status_quo(),
+            {
+                "HealthSystem": {
+                    "year_mode_switch":self.YEAR_OF_CHANGE,
+                    "year_cons_availability_switch":self.YEAR_OF_CHANGE,
+                    "year_use_funded_or_actual_staffing_switch":self.YEAR_OF_CHANGE,
+                    "mode_appt_constraints": 1,
+                    "cons_availability": "default",
+                    "use_funded_or_actual_staffing": "actual",
+                    "mode_appt_constraints_postSwitch": 1,
+                    "use_funded_or_actual_staffing_postSwitch":"actual",
+                    "cons_availability_postSwitch":"default",
+                    "policy_name": "Naive",
+                    "tclose_overwrite": 1,
+                    "tclose_days_offset_overwrite": 7,
+
+                }
+            },
+        )
+
+
+if __name__ == '__main__':
+    from tlo.cli import scenario_run
+
+    scenario_run([__file__])