Fix to read in Legacy versus Updated Consumables RF (#1791)

* fix script to read in legacy or updated consumable availability resource file

* fix tests and preseve module contract

---------

Co-authored-by: Tim Hallett <39991060+tbhallett@users.noreply.github.com>

Author: sakshimohan

src/tlo/methods/healthsystem.py

@@ -207,7 +207,8 @@ class HealthSystem(Module):
                           "consumables in the merged 1b/2 facility level."
         ),
         "availability_estimates": Parameter(
-            Types.DATA_FRAME, "Estimated availability of consumables in the LMIS dataset."
+            Types.DICT, "Estimated availability of consumables in the LMIS dataset. Dict contains all the databases "
+                        "that might be selected using `data_source_for_cons_availability_estimates`"
         ),
         "cons_availability": Parameter(
             Types.STRING,
@@ -646,18 +647,12 @@ def read_parameters(self, resourcefilepath: Optional[Path] = None):
             dtype={"Item_Code": int, "is_diagnostic": bool, "is_medicine": bool, "is_other": bool},
         ).set_index("Item_Code")
 
-        # Choose to read-in the updated availabilty estimates or the legacy availability estimates
-        if self.parameters["data_source_for_cons_availability_estimates"] == 'original':
-            filename_for_cons_availability_estimates = "ResourceFile_Consumables_availability_small_original.csv"
-        elif self.parameters["data_source_for_cons_availability_estimates"] == 'updated':
-            filename_for_cons_availability_estimates = "ResourceFile_Consumables_availability_small.csv"
-        else:
-            raise ValueError("data_source_for_cons_availability_estimates should be either 'original' or 'updated'")
-
-        self.parameters["availability_estimates"] = pd.read_csv(
-            path_to_resourcefiles_for_healthsystem / "consumables" / filename_for_cons_availability_estimates
-        )
-
+        def read_consumables(filename):
+            return pd.read_csv(path_to_resourcefiles_for_healthsystem / "consumables" / filename)
+        self.parameters["availability_estimates"] = {
+            "original": read_consumables("ResourceFile_Consumables_availability_small_original.csv"),
+            "updated": read_consumables("ResourceFile_Consumables_availability_small.csv"),
+        }
 
         # Data on the number of beds available of each type by facility_id
         self.parameters["BedCapacity"] = pd.read_csv(
@@ -797,11 +792,14 @@ def pre_initialise_population(self):
         self.bed_days = BedDays(hs_module=self, availability=self.get_beds_availability())
         self.bed_days.pre_initialise_population()
 
+        # Confirm availability data for consumables
+        _availability_data = self.update_consumables_availability_to_represent_merging_of_levels_1b_and_2(
+                self.parameters["availability_estimates"][self.parameters["data_source_for_cons_availability_estimates"]]
+            )
+
         # Initialise the Consumables class
         self.consumables = Consumables(
-            availability_data=self.update_consumables_availability_to_represent_merging_of_levels_1b_and_2(
-                self.parameters["availability_estimates"]
-            ),
+            availability_data=_availability_data,
             item_code_designations=self.parameters["consumables_item_designations"],
             rng=rng_for_consumables,
             availability=self.get_cons_availability(),

tests/test_consumables.py

@@ -354,7 +354,8 @@ def initialise_simulation(self, sim):
     )
 
     if data is not None:
-        sim.modules['HealthSystem'].parameters['availability_estimates'] = data
+        sim.modules['HealthSystem'].parameters['data_source_for_cons_availability_estimates'] = 'data'
+        sim.modules['HealthSystem'].parameters['availability_estimates'] = {'data' :data}
 
     sim.make_initial_population(n=100)
 
@@ -707,3 +708,61 @@ def test_consumables_availability_modes_that_depend_on_designations(seed):
             assert isinstance(comparison, bool), 'Comparison went wrong: {availability=}'
             assert not consumables._prob_item_codes_available.equals(default_consumables_availability),  \
                 f"No change in actual avaialbility when: {availability=}"
+
+
+def test_switch_between_different_cons_availability_databases(seed):
+    """Check that option to change consumable availability data source works as expected"""
+
+    # Import the 'raw' datasets for consumables availability
+    path_to_files = resourcefilepath / 'healthsystem' / 'consumables'
+    options_for_availability = {
+        'original': pd.read_csv(path_to_files / 'ResourceFile_Consumables_availability_small_original.csv'),
+        'updated': pd.read_csv(path_to_files / 'ResourceFile_Consumables_availability_small.csv'),
+    }
+
+    # Sample availability of the raw datasets, where we know they are different:
+    _facility_id = 2
+    _month = 1
+    def get_sample_availability(df):
+        return df.loc[
+            (df.Facility_ID == _facility_id) & (df.month == _month),
+            ['item_code', 'available_prop']
+        ].set_index('item_code')['available_prop'].to_dict()
+    sample_availability_in_raw_data = {k : get_sample_availability(v) for k, v in options_for_availability.items()}
+
+    # Confirm that the samples from these two raw datasets are different
+    assert sample_availability_in_raw_data['original'] != sample_availability_in_raw_data['updated']
+
+    # Use different options (swithced after module registration) and confirm that the live data being used in the
+    # module matches the intended raw data
+    for which_option in options_for_availability.keys():
+        sim = Simulation(
+            start_date=Date(2010, 1, 1),
+            seed=seed,
+            resourcefilepath=resourcefilepath
+        )
+
+
+        # Register the core modules
+        sim.register(
+            demography.Demography(),
+            healthsystem.HealthSystem(),
+        )
+
+        # Make the decision to switch following module registration
+        sim.modules['HealthSystem'].parameters['cons_availability'] = 'default'
+        sim.modules['HealthSystem'].parameters['data_source_for_cons_availability_estimates'] = which_option
+
+        # Initialise the simulation and capture the 'live' data being used in the Consumables module
+        sim.make_initial_population(n=100)
+        sim.simulate(end_date=sim.start_date)
+        hs = sim.modules['HealthSystem']
+        consumables_live_data = hs.consumables._availability_data
+
+        # check availability actually held by consumables class during simulation matches intended data source
+        sample_availability_in_live_data = consumables_live_data.loc[
+            (consumables_live_data.Facility_ID == _facility_id) & (consumables_live_data.month == _month),
+            ['item_code', 'available_prop']
+        ].set_index('item_code')['available_prop'].to_dict()
+
+        assert sample_availability_in_live_data == sample_availability_in_raw_data[which_option]