Fix single-variable rake validation

Author: neuralsorcerer

balance/weighting_methods/rake.py

@@ -267,10 +267,11 @@ def rake(
             (sample_df, target_df)
         )
 
-    assert len(variables) > 0, (
-        "Must weight on at least one variable. "
-        "Received no common variables between sample and target."
-    )
+    if len(variables) == 0:
+        raise ValueError(
+            "Must weight on at least one variable. "
+            "Received no common variables between sample and target."
+        )
 
     if store_fit_metadata:
         # Fail fast: persisting non-pickleable callables (e.g. lambdas,
@@ -317,9 +318,6 @@ def rake(
         f"Alphabetized variable order is as follows: {alphabetized_variables}."
     )
 
-    target_df = target_df.assign(weight=target_weights)
-    sample_df = sample_df.assign(weight=sample_weights)
-
     # Cast all data types as string to be explicit about each unique value
     # being its own group and to handle that `fillna()` above creates
     # series of type Object, which won't work for the ipfn script
@@ -339,8 +337,9 @@ def rake(
             )
         if len(target_over_set):
             if len(alphabetized_variables) == 1:
-                missing_level_target_weight = target_df.loc[
-                    target_df[variable].isin(target_over_set), "weight"
+                missing_mask = target_df[variable].isin(target_over_set)
+                missing_level_target_weight = target_weights.loc[
+                    target_df.index[missing_mask]
                 ].sum()
                 if missing_level_target_weight > 0:
                     raise ValueError(
@@ -359,6 +358,9 @@ def rake(
         f"Final covariates and levels that will be used in raking: {dict(zip(alphabetized_variables, categories))}."
     )
 
+    target_df = target_df.assign(weight=target_weights)
+    sample_df = sample_df.assign(weight=sample_weights)
+
     sample_sum_weights = sample_df["weight"].sum()
     target_sum_weights = target_df["weight"].sum()
 

tests/test_rake.py

@@ -134,9 +134,7 @@ def test_rake_input_assertions(self) -> None:
         )
         self.assertIn("weight", single_var_result)
 
-        with self.assertRaisesRegex(
-            AssertionError, "Must weight on at least one variable"
-        ):
+        with self.assertRaisesRegex(ValueError, "Must weight on at least one variable"):
             rake(
                 sample_df=pd.DataFrame({"a": [1, 2]}),
                 sample_weights=pd.Series([1.0, 1.0]),
@@ -339,7 +337,7 @@ def test_rake_single_variable_default_transform_with_na_predict_works(self) -> N
             atol=1e-8,
         )
 
-    def test_rake_single_variable_target_only_levels_rejected_in_transfer(self) -> None:
+    def test_rake_single_variable_target_only_levels_rejected_at_fit(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])