Merge pull request #497 from The-Strategy-Unit/alter_inequalities_params

alter way inequalities factors are loaded

Author: tomjemmett

params-schema.json

@@ -208,11 +208,6 @@
                     }
                 }
             },
-            "required": [
-                "level_up",
-                "zero_sum",
-                "level_down"
-            ],
             "additionalProperties": false
         },
         "activity_avoidance": {

src/nhp/model/model.py

@@ -191,17 +191,28 @@ def _load_inequalities_factors(self, data_loader: Data) -> None:
         * | `self.inequalities_factors`: a pandas.DataFrame containing the inequalities
           | factors to be used for the model iterations, given model parameters
         """
-        inequalities_df = data_loader.get_inequalities().set_index("sushrg_trimmed")
-        inequalities_factors = pd.concat(
+        inequalities_params = pd.DataFrame(
             [
-                inequalities_df.loc[hrgs][["icb", "imd_quintile", inequality_method]]
-                .reset_index()
-                .rename(columns={inequality_method: "factor"})
-                for inequality_method, hrgs in self.params["inequalities"].items()
-            ]
-        ).reset_index(drop=True)
-
-        self.inequalities_factors = inequalities_factors
+                {"inequality_type": k, "sushrg_trimmed": v}
+                for k, hrgs in self.params["inequalities"].items()
+                for v in hrgs
+            ],
+            columns=["inequality_type", "sushrg_trimmed"],
+        )
+
+        inequalities_df = (
+            data_loader.get_inequalities()
+            .drop(columns=["activity_rate", "fitted_line"])
+            .melt(
+                id_vars=["icb", "sushrg_trimmed", "imd_quintile"],
+                var_name="inequality_type",
+                value_name="factor",
+            )
+        )
+
+        self.inequalities_factors = inequalities_df.merge(
+            inequalities_params, on=["sushrg_trimmed", "inequality_type"], how="inner"
+        ).drop(columns=["inequality_type"])
 
     @staticmethod
     def generate_run_params(params):

tests/unit/nhp/model/test_model.py

@@ -808,15 +808,17 @@ def test_load_inequalities_factors(mock_model):
             "icb": ["ICB"] * 10,
             "sushrg_trimmed": ["HRG1"] * 5 + ["HRG2"] * 5,
             "imd_quintile": list(range(1, 6)) * 2,
+            "activity_rate": [0.1] * 10,
+            "fitted_line": [0.2] * 10,
             "level_up": [1] * 10,
             "level_down": [2] * 10,
             "zero_sum": [3] * 10,
         }
     )
     expected = pd.DataFrame(
         {
-            "sushrg_trimmed": ["HRG1"] * 5 + ["HRG2"] * 5,
             "icb": ["ICB"] * 10,
+            "sushrg_trimmed": ["HRG1"] * 5 + ["HRG2"] * 5,
             "imd_quintile": list(range(1, 6)) * 2,
             "factor": [1] * 5 + [2] * 5,
         }
@@ -827,6 +829,30 @@ def test_load_inequalities_factors(mock_model):
     pd.testing.assert_frame_equal(mdl.inequalities_factors, expected)
 
 
+def test_load_inequalities_factors_when_empty(mock_model):
+    # arrange
+    mdl = mock_model
+    mdl.params["inequalities"] = {}
+    data_loader = Mock()
+    data_loader.get_inequalities.return_value = pd.DataFrame(
+        {
+            "icb": ["ICB"] * 10,
+            "sushrg_trimmed": ["HRG1"] * 5 + ["HRG2"] * 5,
+            "imd_quintile": list(range(1, 6)) * 2,
+            "activity_rate": [0.1] * 10,
+            "fitted_line": [0.2] * 10,
+            "level_up": [1] * 10,
+            "level_down": [2] * 10,
+            "zero_sum": [3] * 10,
+        }
+    )
+    expected = pd.DataFrame([], columns=["icb", "sushrg_trimmed", "imd_quintile", "factor"])
+    # act
+    mdl._load_inequalities_factors(data_loader)
+    # assert
+    assert len(mdl.inequalities_factors) == 0
+
+
 def test_process_results(mock_model):
     # arrange
     # act & assert