Update Level 1b consumable availability to reflect 1b-2 merge

src/scripts/data_file_processing/healthsystem/consumables/consumable_resource_analyses_with_lmis/consumables_availability_estimation.py

@@ -765,7 +765,33 @@ def get_inflow_to_outflow_ratio_by_item_and_facilitylevel(_df):
 mwanza_1b = sf.loc[(sf.district_std == 'Mwanza') & (sf.fac_type_tlo == '1a')].copy().assign(fac_type_tlo='1b')
 sf = pd.concat([sf, mwanza_1b], axis=0, ignore_index=True)
 
-# 4) Copy all the results to create a level 0 with an availability equal to half that in the respective 1a
+# 4) Update the availability Xpert (item_code = 187)
+# First add rows for Xpert at level 1b by cloning rows for level 2 -> only if not already present
+xpert_item = sf['item_code'].eq(187)
+level_2    = sf['fac_type_tlo'].eq('2')
+level_1b   = sf['fac_type_tlo'].eq('1b')
+
+# Clone rows from level 2
+base   = sf.loc[level_2 & xpert_item].copy()
+new_rows  = base.copy()
+new_rows['fac_type_tlo'] = '1b'
+
+# Add rows to main availability dataframe and drop duplicates, if any
+sf = pd.concat([sf, new_rows], ignore_index=True)
+id_cols = [c for c in sf.columns if c != 'available_prop']
+dupe_mask = sf.duplicated(subset=id_cols, keep=False)
+dupes = sf.loc[dupe_mask].sort_values(id_cols)
+sf = sf.drop_duplicates(subset=id_cols, keep='first').reset_index(drop=True)
+
+# Compute the average availability Sep–Dec (months >= 9) for level 2, item 187
+sep_to_dec = sf['month'].ge(9)
+new_xpert_availability = sf.loc[level_2 & xpert_item & sep_to_dec, 'available_prop'].mean()
+# Assign new availability to relevant facility levels
+levels_1b_2_or_3 = sf['fac_type_tlo'].isin(['1b', '2', '3'])
+xpert_item = sf['item_code'].eq(187)
+sf.loc[levels_1b_2_or_3 & xpert_item, 'available_prop'] = new_xpert_availability
+
+# 5) Copy all the results to create a level 0 with an availability equal to half that in the respective 1a
 all_1a = sf.loc[sf.fac_type_tlo == '1a']
 all_0 = sf.loc[sf.fac_type_tlo == '1a'].copy().assign(fac_type_tlo='0')
 all_0.available_prop *= 0.5
@@ -878,6 +904,151 @@ def interpolate_missing_with_mean(_ser):
 full_set_interpolated = full_set_interpolated.reset_index()
 #full_set_interpolated = full_set_interpolated.reset_index().merge(item_code_category_mapping, on = 'item_code', how = 'left', validate = 'm:1')
 
+def update_level1b_availability(
+    full_set_interpolated: pd.DataFrame,
+    facilities_by_level: dict,
+    resourcefilepath: Path,
+    district_to_city_dict: dict,
+    weighting: str = "district_1b_to_2_ratio"
+) -> pd.DataFrame:
+    """
+    Updates the availability of Level 1b facilities to be the weighted average
+    of availability at Level 1b and 2 facilities, since these levels are merged
+    together in simulations.
+
+    weighting : {'level2', 'national_1b_to_2_ratio', 'district_1b_to_2_ratio'}, default 'district_1b_to_2_ratio'
+        Weighting strategy:
+            - 'level2': Replace 1b availability entirely with level 2 values.
+            - 'national_1b_to_2_ratio': Apply a single national 1b:2 ratio to all districts.
+            - 'district_1b_to_2_ratio': (default) Use district-specific 1b:2 ratios.
+    """
+    # Load and prepare base weights (facility counts)
+    # ---------------------------------------------------------------------
+    weight = (
+        pd.read_csv(resourcefilepath / 'healthsystem' / 'organisation' / 'ResourceFile_Master_Facilities_List.csv')
+        [["District", "Facility_Level", "Facility_ID", "Facility_Count"]]
+    )
+
+    # Keep only Level 1b and 2 facilities
+    lvl1b2_weights = weight[weight["Facility_Level"].isin(["1b", "2"])].copy()
+
+    # Compute weights depending on strategy
+    # ---------------------------------------------------------------------
+    if weighting == "level2":
+        # Force all weight on level 2
+        lvl1b2_weights = lvl1b2_weights[~lvl1b2_weights.District.str.contains("City")]
+        lvl1b2_weights["weight"] = (lvl1b2_weights["Facility_Level"] == "2").astype(float)
+        lvl1b2_weights = lvl1b2_weights.drop(columns = 'Facility_ID')
+
+    elif weighting == "national_1b_to_2_ratio":
+        lvl1b2_weights = lvl1b2_weights[~lvl1b2_weights.District.str.contains("City")]
+        # National total counts
+        national_counts = (
+            lvl1b2_weights.groupby("Facility_Level")["Facility_Count"].sum().to_dict()
+        )
+        total_fac = national_counts.get("1b", 0) + national_counts.get("2", 0)
+        if total_fac == 0:
+            raise ValueError("No facilities found at levels 1b or 2.")
+        lvl1b2_weights["weight"] = lvl1b2_weights["Facility_Level"].map(
+            {lvl: cnt / total_fac for lvl, cnt in national_counts.items()}
+        )
+        lvl1b2_weights = lvl1b2_weights.drop(columns='Facility_ID')
+
+    elif weighting == "district_1b_to_2_ratio":
+        # Replace city names with their parent districts (temporarily for grouping)
+        city_to_district_dict = {v: k for k, v in district_to_city_dict.items()}
+        lvl1b2_weights["District"] = lvl1b2_weights["District"].replace(city_to_district_dict)
+
+        # District-level weighting (default)
+        lvl1b2_weights = (
+            lvl1b2_weights
+            .groupby(["District", "Facility_Level"], as_index=False)["Facility_Count"]
+            .sum()
+        )
+
+        lvl1b2_weights["total_facilities"] = lvl1b2_weights.groupby("District")["Facility_Count"].transform("sum")
+        lvl1b2_weights["weight"] = lvl1b2_weights["Facility_Count"] / lvl1b2_weights["total_facilities"]
+
+    else:
+        raise ValueError(
+            f"Invalid weighting '{weighting}'. Choose from "
+            "'level2', 'national_1b_to_2_ratio', or 'district_1b_to_2_ratio'."
+        )
+
+    # Add back city districts (reverse mapping)
+    for source, destination in copy_source_to_destination.items():
+        new_rows = lvl1b2_weights.loc[lvl1b2_weights.District == source].copy()
+        new_rows.District = destination
+        lvl1b2_weights = pd.concat([lvl1b2_weights, new_rows], axis=0, ignore_index=True)
+
+    # Merge Facility_ID back
+    lvl1b2_weights = lvl1b2_weights.merge(
+        weight.loc[weight["Facility_Level"].isin(["1b", "2"]), ["District", "Facility_Level", "Facility_ID"]],
+        on=["District", "Facility_Level"],
+        how="left",
+        validate="1:1"
+    )
+
+    # Subset Level 1b and 2 facilities and apply weights
+    # ---------------------------------------------------------------------
+    lvl1b2_ids = list(facilities_by_level.get("1b", [])) + list(facilities_by_level.get("2", []))
+    full_set_interpolated_levels1b2 = full_set_interpolated[
+        full_set_interpolated["Facility_ID"].isin(lvl1b2_ids)
+    ].copy()
+
+    full_set_interpolated_levels1b2 = full_set_interpolated_levels1b2.merge(
+        lvl1b2_weights[["District", "Facility_Level", "Facility_ID", "weight"]],
+        on="Facility_ID",
+        how="left",
+        validate="m:1"
+    )
+
+    # Apply weighting
+    full_set_interpolated_levels1b2["available_prop"] *= full_set_interpolated_levels1b2["weight"]
+
+    # Aggregate to district-month-item level
+    full_set_interpolated_levels1b2 = (
+        full_set_interpolated_levels1b2
+        .groupby(["District", "month", "item_code"], as_index=False)["available_prop"]
+        .sum()
+    )
+    full_set_interpolated_levels1b2["Facility_Level"] = "1b"
+
+    # Reattach Facility_IDs for level 1b
+    full_set_interpolated_levels1b2 = full_set_interpolated_levels1b2.merge(
+        lvl1b2_weights.query("Facility_Level == '1b'")[["District", "Facility_Level", "Facility_ID", "weight"]],
+        on=["District", "Facility_Level"],
+        how="left",
+        validate="m:1"
+    )
+
+    # Replace old level 1b facilities and recompute weighted availability
+    # ---------------------------------------------------------------------
+    # Drop old Level 1b facilities
+    full_set_interpolated = full_set_interpolated[
+        ~full_set_interpolated["Facility_ID"].isin(facilities_by_level.get("1b", []))
+    ]
+
+    # Append new 1b facility data
+    full_set_interpolated = pd.concat(
+        [
+            full_set_interpolated,
+            full_set_interpolated_levels1b2[["Facility_ID", "month", "item_code", "available_prop"]]
+        ],
+        axis=0,
+        ignore_index=True
+    )
+
+    return full_set_interpolated
+
+full_set_interpolated = update_level1b_availability(
+    full_set_interpolated=full_set_interpolated,
+    facilities_by_level=facilities_by_level,
+    resourcefilepath=resourcefilepath,
+    district_to_city_dict=copy_source_to_destination,
+    weighting = 'district_1b_to_2_ratio',
+)
+
 # --- Check that the exported file has the properties required of it by the model code. --- #
 check_format_of_consumables_file(df=full_set_interpolated, fac_ids=fac_ids)