Fix single-variable rake fallback support metadata

Author: neuralsorcerer

balance/weighting_methods/rake.py

@@ -312,9 +312,12 @@ def rake(
         }
         if store_fit_metadata:
             cell_ratio = post["model"]["cell_weight_ratio"]
-            categories = [[str(x) for x in cell_ratio.index]]
+            supported_ratio = cell_ratio[np.isfinite(cell_ratio.to_numpy(dtype=float))]
+            categories = [[str(x) for x in supported_ratio.index]]
             m_sample = np.ones(len(categories[0]), dtype=np.float64)
-            m_fit = (cell_ratio.to_numpy(dtype=np.float64) * m_sample).reshape(-1)
+            m_fit = (supported_ratio.to_numpy(dtype=np.float64) * m_sample).reshape(-1)
+            sample_index = post["model"]["sample_index"]
+            target_index = post["model"]["target_index"]
             model.update(
                 {
                     "store_fit_metadata": True,
@@ -328,8 +331,8 @@ def rake(
                     "na_action": post["model"]["na_action"],
                     "transformations": post["model"]["transformations"],
                     "transformations_origin": transformations,
-                    "training_sample_weights": sample_weights.copy(),
-                    "training_target_weights": target_weights.copy(),
+                    "training_sample_weights": sample_weights.loc[sample_index].copy(),
+                    "training_target_weights": target_weights.loc[target_index].copy(),
                     "weight_trimming_mean_ratio": weight_trimming_mean_ratio,
                     "weight_trimming_percentile": weight_trimming_percentile,
                     "keep_sum_of_weights": keep_sum_of_weights,

tests/test_rake.py

@@ -340,6 +340,65 @@ def test_rake_single_variable_default_transform_with_na_predict_works(self) -> N
         self.assertEqual(predicted.shape[0], sample_df.shape[0])
         self.assertTrue(np.isfinite(predicted.fillna(0).to_numpy()).all())
 
+    def test_rake_single_variable_target_only_levels_rejected_in_transfer(self) -> None:
+        sample_df = pd.DataFrame({"x": ["a", "a", "b", "b"]})
+        target_df = pd.DataFrame({"x": ["a", "a", "b", "b", "c"]})
+        sample_w = pd.Series([1.0, 1.0, 1.0, 1.0])
+        target_w = pd.Series([1.0] * len(target_df))
+
+        fitted = rake(
+            sample_df=sample_df,
+            sample_weights=sample_w,
+            target_df=target_df,
+            target_weights=target_w,
+            variables=["x"],
+            transformations=None,
+            store_fit_metadata=True,
+        )
+        new_sample_df = pd.DataFrame({"x": ["a", "c"]})
+        new_sample_w = pd.Series([1.0, 1.0])
+        with self.assertRaisesRegex(
+            ValueError, "do not map to stored fit-time categories"
+        ):
+            _predict_weights_from_model(
+                model=fitted["model"],
+                sample_df=new_sample_df,
+                sample_weights_full=new_sample_w,
+                target_df=target_df,
+                target_weights=target_w,
+                is_transfer=True,
+            )
+
+    def test_rake_single_variable_drop_na_replay_uses_dropped_target_total(
+        self,
+    ) -> None:
+        sample_df = pd.DataFrame({"x": ["a", "b", "a"]})
+        target_df = pd.DataFrame({"x": ["a", "b", np.nan, "b"]})
+        sample_w = pd.Series([1.0, 1.0, 1.0])
+        target_w = pd.Series([1.0, 1.0, 10.0, 1.0])
+
+        fitted = rake(
+            sample_df=sample_df,
+            sample_weights=sample_w,
+            target_df=target_df,
+            target_weights=target_w,
+            variables=["x"],
+            transformations=None,
+            na_action="drop",
+            store_fit_metadata=True,
+        )
+        replay = _predict_weights_from_model(
+            model=fitted["model"],
+            sample_df=sample_df,
+            sample_weights_full=sample_w,
+            target_df=target_df,
+            target_weights=target_w,
+            is_transfer=False,
+        )
+        self.assertAlmostEqual(
+            float(replay.sum()), float(fitted["weight"].sum()), places=8
+        )
+
     def test_rake_weights(self) -> None:
         """
         Test basic rake weighting functionality with categorical data.