Add single-variable rake fallback to poststratify

Author: neuralsorcerer

CHANGELOG.md

@@ -49,6 +49,12 @@
     `fct_lump`). To enable transfer scoring, pass deterministic
     transformations at fit time (e.g. wrappers built around stored
     fit-time bin edges) or re-fit rake on the scoring data.
+- **`rake()` now auto-falls back to `poststratify()` for single-variable inputs.**
+  - Calls that pass exactly one adjustment variable no longer fail with a
+    variable-count assertion.
+  - The function logs a warning and delegates to `poststratify()` with the
+    same transformation/NA-handling/trimming options, so users can keep
+    `method="rake"` in generic workflows without branching on variable count.
 
 ## Code Quality & Refactoring
 

balance/weighting_methods/rake.py

@@ -20,6 +20,7 @@
 import pandas as pd
 from balance import adjustment as balance_adjustment, util as balance_util
 from balance.util import _safe_fillna_and_infer
+from balance.weighting_methods.poststratify import poststratify
 
 logger: logging.Logger = logging.getLogger(__package__)
 
@@ -261,12 +262,29 @@ def rake(
     sample_df = sample_df.loc[:, variables]
     target_df = target_df.loc[:, variables]
 
-    # TODO: When len(variables) == 1, fall back to poststratify instead of
-    # raising, so users can call adjust(method="rake") without worrying about
-    # the variable count.
+    if len(variables) == 1:
+        logger.warning(
+            "rake() received a single variable (%s); falling back to "
+            "poststratify() on that variable.",
+            variables[0],
+        )
+        return poststratify(
+            sample_df=sample_df,
+            sample_weights=sample_weights,
+            target_df=target_df,
+            target_weights=target_weights,
+            variables=variables,
+            transformations=transformations,
+            na_action=na_action,
+            weight_trimming_mean_ratio=weight_trimming_mean_ratio,
+            weight_trimming_percentile=weight_trimming_percentile,
+            keep_sum_of_weights=keep_sum_of_weights,
+            store_fit_metadata=store_fit_metadata,
+        )
+
     assert len(variables) > 1, (
-        "Must weight on at least two variables for raking. "
-        f"Currently have variables={variables} only"
+        "Must weight on at least one variable. "
+        "Received no common variables between sample and target."
     )
 
     transformations_to_apply = transformations

tests/test_rake.py

@@ -21,6 +21,7 @@
 from balance import adjustment as balance_adjustment
 from balance.sample_class import Sample
 from balance.util import _assert_type
+from balance.weighting_methods.poststratify import poststratify
 from balance.weighting_methods.rake import (
     _find_lcm_of_array_lengths,
     _hare_niemeyer_allocation,
@@ -124,14 +125,21 @@ def test_rake_input_assertions(self) -> None:
             pd.Series((1,) * n_rows),
         )
 
-        # Must pass more than one variable
-        self._assert_rake_raises_with_message(
-            "Must weight on at least two variables",
-            sample[["a"]],
-            pd.Series((1,) * n_rows),
-            target[["a"]],
-            pd.Series((1,) * n_rows),
+        # A single variable falls back to poststratify.
+        with self.assertLogs("balance.weighting_methods", level="WARNING") as cm:
+            single_var_result = rake(
+                sample_df=sample[["a"]],
+                sample_weights=pd.Series((1,) * n_rows),
+                target_df=target[["a"]],
+                target_weights=pd.Series((1,) * n_rows),
+            )
+        self.assertTrue(
+            any(
+                "falling back to poststratify" in message.lower()
+                for message in cm.output
+            )
         )
+        self.assertIn("weight", single_var_result)
 
         # Must pass weights for sample
         self._assert_rake_raises_with_message(
@@ -234,6 +242,46 @@ def test_rake_fails_when_all_na(self) -> None:
             transformations=None,
         )
 
+    def test_rake_single_variable_matches_poststratify(self) -> None:
+        sample_df = pd.DataFrame(
+            {"x": ["a", "a", "b", "b", "b"], "noise": [1, 2, 3, 4, 5]}
+        )
+        target_df = pd.DataFrame({"x": ["a", "a", "a", "b", "b", "b", "b"]})
+        sample_w = pd.Series([1.0, 2.0, 3.0, 4.0, 5.0])
+        target_w = pd.Series([1.0] * len(target_df))
+
+        with self.assertLogs("balance.weighting_methods", level="WARNING"):
+            rake_res = rake(
+                sample_df=sample_df,
+                sample_weights=sample_w,
+                target_df=target_df,
+                target_weights=target_w,
+                variables=["x"],
+                transformations=None,
+                na_action="add_indicator",
+                weight_trimming_mean_ratio=10.0,
+                keep_sum_of_weights=True,
+                store_fit_metadata=True,
+            )
+
+        post_res = poststratify(
+            sample_df=sample_df[["x"]],
+            sample_weights=sample_w,
+            target_df=target_df[["x"]],
+            target_weights=target_w,
+            variables=["x"],
+            transformations=None,
+            na_action="add_indicator",
+            weight_trimming_mean_ratio=10.0,
+            keep_sum_of_weights=True,
+            store_fit_metadata=True,
+        )
+
+        pd.testing.assert_series_equal(rake_res["weight"], post_res["weight"])
+        self.assertEqual(rake_res["model"]["method"], "poststratify")
+        self.assertTrue(rake_res["model"]["store_fit_metadata"])
+        self.assertIn("cell_weight_ratio", rake_res["model"])
+
     def test_rake_weights(self) -> None:
         """
         Test basic rake weighting functionality with categorical data.