Fix over-allocation of Bed Days in #1399 (#1437)

* 1st pass writing combination method

* Write test that checks failure case in issue

* Expand on docstring for combining footprints method

* Force-cast to int when returning footprint since pd.datetime.timedelta doesn't know how to handle np.int64's

* Catch bug when determining priority on each day, write test to cover this case with a 3-bed types resolution

---------

Co-authored-by: Emmanuel Mnjowe <[email protected]>

src/tlo/methods/bed_days.py

@@ -318,6 +318,60 @@ def issue_bed_days_according_to_availability(self, facility_id: int, footprint:
 
         return available_footprint
 
+    def combine_footprints_for_same_patient(
+        self, fp1: Dict[str, int], fp2: Dict[str, int]
+    ) -> Dict[str, int]:
+        """
+        Given two footprints that are due to start on the same day, combine the two footprints by
+        overlaying the higher-priority bed over the lower-priority beds.
+
+        As an example, given the footprints,
+        fp1 = {"bedtype1": 2, "bedtype2": 0}
+        fp2 = {"bedtype1": 1, "bedtype2": 6}
+
+        where bedtype1 is higher priority than bedtype2, we expect the combined allocation to be
+        {"bedtype1": 2, "bedtype2": 5}.
+
+        This is because footprints are assumed to run in the order of the bedtypes priority; so
+        fp2's second day of being allocated to bedtype2 is overwritten by the higher-priority
+        allocation to bedtype1 from fp1. The remaining 5 days are allocated to bedtype2 since
+        fp1 does not require a bed after the first 2 days, but fp2 does.
+
+        :param fp1: Footprint, to be combined with the other argument.
+        :param pf2: Footprint, to be combined with the other argument.
+        """
+        fp1_length = sum(days for days in fp1.values())
+        fp2_length = sum(days for days in fp2.values())
+        max_length = max(fp1_length, fp2_length)
+
+        # np arrays where each entry is the priority of bed allocated by the footprint
+        # on that day. fp_priority[i] = priority of the bed allocated by the footprint on
+        # day i (where the current day is day 0).
+        # By default, fill with priority equal to the lowest bed priority; though all
+        # the values will have been explicitly overwritten after the next loop completes.
+        fp1_priority = np.ones((max_length,), dtype=int) * (len(self.bed_types) - 1)
+        fp2_priority = fp1_priority.copy()
+
+        fp1_at = 0
+        fp2_at = 0
+        for priority, bed_type in enumerate(self.bed_types):
+            # Bed type priority is dictated by list order, so it is safe to loop here.
+            # We will start with the highest-priority bed type and work to the lowest
+            fp1_priority[fp1_at:fp1_at + fp1[bed_type]] = priority
+            fp1_at += fp1[bed_type]
+            fp2_priority[fp2_at:fp2_at + fp2[bed_type]] = priority
+            fp2_at += fp2[bed_type]
+
+        # Element-wise minimum of the two priority arrays is then the bed to assign
+        final_priorities = np.minimum(fp1_priority, fp2_priority)
+        # Final footprint is then formed by converting the priorities into blocks of days
+        return {
+            # Cast to int here since pd.datetime.timedelta doesn't know what to do with
+            # np.int64 types
+            bed_type: int(sum(final_priorities == priority))
+            for priority, bed_type in enumerate(self.bed_types)
+        }
+
     def impose_beddays_footprint(self, person_id, footprint):
         """This is called to reflect that a new occupancy of bed-days should be recorded:
         * Cause to be reflected in the bed_tracker that an hsi_event is being run that will cause bed to be
@@ -345,9 +399,7 @@ def impose_beddays_footprint(self, person_id, footprint):
             remaining_footprint = self.get_remaining_footprint(person_id)
 
             # combine the remaining footprint with the new footprint, with days in each bed-type running concurrently:
-            combo_footprint = {bed_type: max(footprint[bed_type], remaining_footprint[bed_type])
-                               for bed_type in self.bed_types
-                               }
+            combo_footprint = self.combine_footprints_for_same_patient(footprint, remaining_footprint)
 
             # remove the old footprint and apply the combined footprint
             self.remove_beddays_footprint(person_id)

tests/test_beddays.py

@@ -2,6 +2,7 @@
 import copy
 import os
 from pathlib import Path
+from typing import Dict
 
 import pandas as pd
 import pytest
@@ -83,6 +84,88 @@ def test_beddays_in_isolation(tmpdir, seed):
     assert ([cap_bedtype1] * days_sim == tracker.values).all()
 
 
+def test_beddays_allocation_resolution(tmpdir, seed):
+    sim = Simulation(start_date=start_date, seed=seed)
+    sim.register(
+        demography.Demography(resourcefilepath=resourcefilepath),
+        healthsystem.HealthSystem(resourcefilepath=resourcefilepath),
+    )
+
+    # Update BedCapacity data with a simple table:
+    level2_facility_ids = [128, 129, 130]  # <-- the level 2 facilities for each region
+    # This ensures over-allocations have to be properly resolved
+    cap_bedtype1 = 10
+    cap_bedtype2 = 10
+    cap_bedtype3 = 10
+
+    # create a simple bed capacity dataframe
+    hs = sim.modules["HealthSystem"]
+    hs.parameters["BedCapacity"] = pd.DataFrame(
+        data={
+            "Facility_ID": level2_facility_ids,
+            "bedtype1": cap_bedtype1,
+            "bedtype2": cap_bedtype2,
+            "bedtype3": cap_bedtype3,
+        }
+    )
+
+    sim.make_initial_population(n=100)
+    sim.simulate(end_date=start_date)
+
+    # reset bed days tracker to the start_date of the simulation
+    hs.bed_days.initialise_beddays_tracker()
+
+    def assert_footprint_matches_expected(
+        footprint: Dict[str, int], expected_footprint: Dict[str, int]
+    ):
+        """
+        Asserts that two footprints are identical.
+        The footprint provided as the 2nd argument is assumed to be the footprint
+        that we want to match, and the 1st as the result of the program attempting
+        to resolve over-allocations.
+        """
+        assert len(footprint) == len(
+            expected_footprint
+        ), "Bed type footprints did not return same allocations."
+        for bed_type, expected_days in expected_footprint.items():
+            allocated_days = footprint[bed_type]
+            assert expected_days == allocated_days, (
+                f"Bed type {bed_type} was allocated {allocated_days} upon combining, "
+                f"but expected it to get {expected_days}."
+            )
+
+    # Check that combining footprints for a person returns the expected output
+
+    # SIMPLE 2-bed days case
+    # Test uses example fail case given in https://github.com/UCL/TLOmodel/issues/1399
+    # Person p has: bedtyp1 for 2 days, bedtype2 for 0 days.
+    # Person p then assigned: bedtype1 for 1 days, bedtype2 for 6 days.
+    # EXPECT: p's footprints are combined into bedtype1 for 2 days, bedtype2 for 5 days.
+    existing_footprint = {"bedtype1": 2, "bedtype2": 0, "bedtype3": 0}
+    incoming_footprint = {"bedtype1": 1, "bedtype2": 6, "bedtype3": 0}
+    expected_resolution = {"bedtype1": 2, "bedtype2": 5, "bedtype3": 0}
+    allocated_footprint = hs.bed_days.combine_footprints_for_same_patient(
+        existing_footprint, incoming_footprint
+    )
+    assert_footprint_matches_expected(allocated_footprint, expected_resolution)
+
+    # TEST case involve 3 different bed-types.
+    # Person p has: bedtype1 for 2 days, then bedtype3 for 4 days.
+    # p is assigned: bedtype1 for 1 day, bedtype2 for 3 days, and bedtype3 for 1 day.
+    # EXPECT: p spends 2 days in each bedtype;
+    # - Day 1 needs bedtype1 for both footprints
+    # - Day 2 existing footprint at bedtype1 overwrites incoming at bedtype2
+    # - Day 3 & 4 incoming footprint at bedtype2 overwrites existing allocation to bedtype3
+    # - Day 5 both footprints want bedtype3
+    # - Day 6 existing footprint needs bedtype3, whilst incoming footprint is over.s
+    existing_footprint = {"bedtype1": 2, "bedtype2": 0, "bedtype3": 4}
+    incoming_footprint = {"bedtype1": 1, "bedtype2": 3, "bedtype3": 1}
+    expected_resolution = {"bedtype1": 2, "bedtype2": 2, "bedtype3": 2}
+    allocated_footprint = hs.bed_days.combine_footprints_for_same_patient(
+        existing_footprint, incoming_footprint
+    )
+    assert_footprint_matches_expected(allocated_footprint, expected_resolution)
+
 def check_dtypes(simulation):
     # check types of columns
     df = simulation.population.props