BedDays availability switch now correctly updates capacities

src/tlo/methods/bed_days.py

@@ -5,12 +5,12 @@
 
 """
 from collections import defaultdict
-from typing import Dict, Tuple
+from typing import Dict, Literal, Tuple
 
 import numpy as np
 import pandas as pd
 
-from tlo import Property, Types, logging
+from tlo import Date, Property, Types, logging
 
 # ---------------------------------------------------------------------------------------------------------
 #   CLASS DEFINITIONS
@@ -145,6 +145,40 @@ def initialise_beddays_tracker(self, model_to_data_popsize_ratio=1.0):
             assert not df.isna().any().any()
             self.bed_tracker[bed_type] = df
 
+    def switch_beddays_availability(
+        self,
+        new_availability: Literal["all", "none", "default"],
+        effective_on_and_from: Date,
+        model_to_data_popsize_ratio: float = 1.0,
+    ) -> None:
+        """
+        Action to be taken if the beddays availability changes in the middle
+        of the simulation.
+
+        If bed capacities are reduced below the currently scheduled occupancy,
+        inpatients are not evicted from beds and are allowed to remain in the
+        bed until they are scheduled to leave. Obviously, no new patients will
+        be admitted if there is no room in the new capacities.
+
+        :param new_availability: The new bed availability. See __init__ for details.
+        :param effective_on_and_from: First day from which the new capacities will be imposed.
+        :param model_to_data_popsize_ratio: As in initialise_population.
+        """
+        # Store new bed availability
+        self.availability = new_availability
+        # Before we update the bed capacity, we need to store its old values
+        # This is because we will need to update the trackers to reflect the new#
+        # maximum capacities for each bed type.
+        old_max_capacities: pd.DataFrame = self._scaled_capacity.copy()
+        # Set the new capacity for beds
+        self.set_scaled_capacity(model_to_data_popsize_ratio)
+        # Compute the difference between the new max capacities and the old max capacities
+        difference_in_max = self._scaled_capacity - old_max_capacities
+        # For each tracker, after the effective date, impose the difference on the max
+        # number of beds
+        for bed_type, tracker in self.bed_tracker.items():
+            tracker.loc[effective_on_and_from:] += difference_in_max[bed_type]
+
     def on_start_of_day(self):
         """Things to do at the start of each new day:
         * Refresh inpatient status

src/tlo/methods/healthsystem.py

@@ -2776,7 +2776,11 @@ def apply(self, population):
             self.module.consumables.on_start_of_day(self.module.sim.date)
 
         if 'beds_availability' in self._parameters:
-            self.module.bed_days.availability = self._parameters['beds_availability']
+            self.module.bed_days.switch_beddays_availability(
+                new_availability=self._parameters["beds_availability"],
+                effective_on_and_from=self.sim.date,
+                model_to_data_popsize_ratio=self.sim.modules["Demography"].initial_model_to_data_popsize_ratio
+            )
 
         if 'equip_availability' in self._parameters:
             self.module.equipment.availability = self._parameters['equip_availability']

tests/test_beddays.py

@@ -973,3 +973,74 @@ def apply(self, person_id, squeeze_factor):
     # Check that the facility_id is included for each entry in the `HSI_Events` log, including HSI Events for
     # in-patient appointments.
     assert not (log_hsi['Facility_ID'] == -99).any()
+
+def test_beddays_availability_switch(seed):
+    """
+    Test that calling bed_days.switch_beddays_availability correctly updates the
+    bed capacities and adjusts the existing trackers to reflect the new capacities.
+    """
+    sim = Simulation(start_date=start_date, seed=seed)
+    sim.register(
+        demography.Demography(resourcefilepath=resourcefilepath),
+        healthsystem.HealthSystem(resourcefilepath=resourcefilepath),
+    )
+
+    # get shortcut to HealthSystem Module
+    hs: healthsystem.HealthSystem = sim.modules["HealthSystem"]
+
+    # Create a simple bed capacity dataframe with capacity designated for two regions
+    hs.parameters["BedCapacity"] = pd.DataFrame(
+        data={
+            "Facility_ID": [
+                128, #<-- patient 0 is admitted here
+                129,
+            ],
+            "bedtype1": 5,
+            "bedtype2": 10,
+        }
+    )
+    sim.make_initial_population(n=100)
+    sim.simulate(end_date=start_date)
+
+    day_2 = start_date + pd.DateOffset(days=1)
+    day_3 = start_date + pd.DateOffset(days=2)
+    day_4 = start_date + pd.DateOffset(days=3)
+
+    bed_days = hs.bed_days
+    # Reset the bed occupancies
+    bed_days.initialise_beddays_tracker()
+    # Have a patient occupy a bed at the start of the simulation
+    bed_days.impose_beddays_footprint(person_id=0, footprint={"bedtype1": 3, "bedtype2": 0})
+
+    # Have the bed_days availability switch to "none" on the 2nd simulation day
+    bed_days.switch_beddays_availability("none", effective_on_and_from=day_2)
+
+    # We should now see that the scaled capacities are all zero
+    assert (
+        not bed_days._scaled_capacity.any().any()
+    ), "At least one bed capacity was not set to 0"
+    # We should also see that bedtype1 should have -1 beds available for days 2 and 3 of the simulation,
+    # due to the existing occupancy and the new capacity of 0.
+    # It should have 4 beds available on the first day (since the original capacity was 5 and the availability
+    # switch happens day 2).
+    # It should then have 0 beds available after (not including) day 3
+    bedtype1: pd.DataFrame = bed_days.bed_tracker["bedtype1"]
+    bedtype2: pd.DataFrame = bed_days.bed_tracker["bedtype2"]
+
+    assert (
+        bedtype1.loc[start_date, 128] == 4 and bedtype1.loc[start_date, 129] == 5
+    ), "Day 1 capacities were incorrectly affected"
+    assert (bedtype1.loc[day_2:day_3, 128] == -1).all() and (
+        bedtype1.loc[day_2:day_3, 129] == 0
+    ).all(), "Day 2 & 3 capacities were not updated correctly"
+    assert (
+        (bedtype1.loc[day_4:, :] == 0).all().all()
+    ), "Day 4 onwards did not have correct capacity"
+
+    # Bedtype 2 should have also have been updated, but there is no funny business here.
+    assert (
+        (bedtype2.loc[day_2:, :] == 0).all().all()
+    ), "Bedtype 2 was not updated correctly"
+    assert (
+        (bedtype2.loc[start_date, :] == 10).all().all()
+    ), "Bedtype 2 had capacity updated on the incorrect dates"