Extract rake predict_weights logic into weighting_methods/rake.py

Summary:
Moves the bulk of `BalanceFrame._predict_weights_rake` (~210 lines) into a new `predict_weights_from_model` function in `balance/weighting_methods/rake.py`. The BalanceFrame method becomes a thin ~30-line wrapper.

Motivation: balance_frame.py was already ~2,300 lines and rake-specific predict logic was awkwardly far from rake's fit code. Co-locating the lifecycle (fit + replay + transfer) inside `weighting_methods/rake.py` makes it easier to reason about rake's contract — anything that depends on the rake-specific model dict structure (m_fit, m_sample, categories, training_*_weights, etc.) now lives next to the function that produces that structure. This also establishes a template for the same extraction on cbps, poststratify, and ipw in subsequent diffs.

What stays in `BalanceFrame._predict_weights_rake` (~30-line wrapper): resolve `self` vs `source`, validate that `data` has a target set, extract sample/target covariate DataFrames and design weights from the BalanceFrame, call into `predict_weights_from_model`, apply weight name and `_align_to_index` on the result.

What moves to `weighting_methods/rake.py::predict_weights_from_model`: all metadata validation (required keys, types, transformations_origin/data-dependent guards), covariate column existence checks, stored-transformations replay, NA handling (drop / add_indicator), category-code mapping against stored fit-time categories, per-cell ratio compute (m_fit / m_sample at the scored cells only), unconditional transfer warning, joint-support validation, target-sum determination across in-place vs transfer × na_action variants, trim_weights, and na_action="drop" full-index restoration.

Side benefit: `predict_weights_from_model` is callable directly from outside BalanceFrame — users holding a fitted rake model dict and plain DataFrames can replay/transfer weights without going through the BalanceFrame class. No public API change: `BalanceFrame.predict_weights()` behavior is unchanged.

Differential Revision: D103588867

Author: talgalili

balance/balance_frame.py

@@ -2109,246 +2109,25 @@ def _predict_weights_rake(
         model: dict[str, Any],
         source: BalanceFrame | None = None,
     ) -> pd.Series:
-        required = (
-            "variables",
-            "variables_before_transformations",
-            "categories",
-            "m_fit",
-            "m_sample",
-            "na_action",
-            "transformations",
-        )
-        if source is None:
-            required = required + ("training_sample_weights", "training_target_weights")
-        missing = [key for key in required if key not in model]
-        if missing:
-            raise ValueError(
-                "Rake model is missing fit-time metadata "
-                f"({missing}) for predict_weights(). "
-                "Call BalanceFrame.fit(method='rake') or run rake(..., "
-                "store_fit_metadata=True)."
-            )
-
-        variables = model.get("variables")
-        input_variables = model.get("variables_before_transformations")
-        categories = model.get("categories")
-        m_fit = model.get("m_fit")
-        m_sample = model.get("m_sample")
-        transformations_origin = model.get("transformations_origin")
-        if (
-            not isinstance(variables, list)
-            or not isinstance(input_variables, list)
-            or not isinstance(categories, list)
-        ):
-            raise ValueError("Rake model metadata is malformed for predict_weights().")
-        if not isinstance(m_fit, np.ndarray) or not isinstance(m_sample, np.ndarray):
-            raise ValueError("Rake model is missing stored contingency tables.")
+        from balance.weighting_methods.rake import predict_weights_from_model
 
         bf = source if source is not None else self
         if source is not None and bf._sf_target is None:
             raise ValueError(
                 "data must have a target set for rake predict_weights(data=...)."
             )
-        if source is not None and transformations_origin == "default":
-            raise ValueError(
-                "Rake predict_weights(data=...) is unsupported for models fitted "
-                "with transformations='default' because those transformations are "
-                "data-dependent and not replayable across new samples. Re-fit on "
-                "the scoring data or fit with explicit deterministic transformations."
-            )
-        if source is not None and isinstance(transformations_origin, dict):
-            # Best-effort guard: reject explicit dicts that directly
-            # reference balance's known data-dependent helpers
-            # (quantize, fct_lump). These recompute bins/levels from the
-            # scoring data, so stored cell ratios no longer line up with
-            # the transformed scoring cells and transfer would silently
-            # return incorrect weights.
-            #
-            # This guard does NOT catch indirect uses such as
-            # ``functools.partial(fct_lump, prop=0.1)``, top-level wrapper
-            # functions, or user-defined data-dependent transformations.
-            # The general invariant is: any callable whose output for a
-            # row depends on other rows in the input is unsafe to replay
-            # on a different sample. Users supplying such transformations
-            # are responsible for either (a) wrapping them as
-            # deterministic functions of stored fit-time parameters or
-            # (b) re-fitting rake on the scoring data.
-            from balance.utils.data_transformation import fct_lump, quantize
-
-            data_dependent_helpers = {quantize, fct_lump}
-            offenders = sorted(
-                {
-                    getattr(fn, "__name__", repr(fn))
-                    for fn in transformations_origin.values()
-                    if fn in data_dependent_helpers
-                }
-            )
-            if offenders:
-                raise ValueError(
-                    "Rake predict_weights(data=...) is unsupported for models "
-                    f"fitted with data-dependent transformations ({', '.join(offenders)}). "
-                    "These recompute bins/levels from the scoring data, so "
-                    "stored cell ratios no longer line up with the transformed "
-                    "cells. Re-fit on the scoring data or fit with deterministic "
-                    "transformations."
-                )
-        sample_covars = bf._sf_sample.df_covars
-        target_covars = _assert_type(bf._sf_target).df_covars
-        for column in input_variables:
-            if (
-                column not in sample_covars.columns
-                or column not in target_covars.columns
-            ):
-                raise ValueError(
-                    "Rake predict_weights() cannot find required covariate "
-                    f"'{column}' in both sample and target."
-                )
-        sample_df = sample_covars.loc[:, input_variables]
-        target_df = target_covars.loc[:, input_variables]
-        sample_weights_full = bf._sf_sample.df_weights.iloc[:, 0]
-        training_sample_weights = model.get("training_sample_weights")
-        sample_weights = sample_weights_full
-        na_action = cast(str, model.get("na_action", "add_indicator"))
 
-        sample_df, target_df = balance_adjustment.apply_transformations(
-            (sample_df, target_df), transformations=model.get("transformations")
+        target_frame = _assert_type(bf._sf_target)
+        predicted = predict_weights_from_model(
+            model=model,
+            sample_df=bf._sf_sample.df_covars,
+            sample_weights_full=bf._sf_sample.df_weights.iloc[:, 0],
+            target_df=target_frame.df_covars,
+            target_weights=target_frame.df_weights.iloc[:, 0],
+            is_transfer=source is not None,
         )
-        for column in variables:
-            if column not in sample_df.columns:
-                raise ValueError(
-                    "Rake transform output is missing stored variable "
-                    f"'{column}' required for predict_weights()."
-                )
-        sample_df = sample_df.loc[:, variables]
-        target_df = target_df.loc[:, variables]
 
-        if source is None:
-            if isinstance(training_sample_weights, pd.Series):
-                if na_action == "drop":
-                    sample_weights = training_sample_weights
-                elif training_sample_weights.index.equals(sample_weights_full.index):
-                    sample_weights = training_sample_weights
-                else:
-                    raise ValueError(
-                        "Rake predict_weights() requires compatible fit-time sample design "
-                        "weights for in-place replay. This can happen because "
-                        "store_fit_metadata is missing/incompatible, or because you're "
-                        "scoring a different sample; use predict_weights(data=...) "
-                        "for different samples."
-                    )
-            else:
-                raise ValueError(
-                    "Rake predict_weights() requires compatible fit-time sample design "
-                    "weights for in-place replay. This can happen because "
-                    "store_fit_metadata is missing/incompatible, or because you're "
-                    "scoring a different sample; use predict_weights(data=...) "
-                    "for different samples."
-                )
-
-        dropped_target_weights: pd.Series | None = None
-        if na_action == "drop":
-            sample_df, sample_weights = balance_util.drop_na_rows(
-                sample_df, sample_weights, "sample"
-            )
-            target_df, dropped_target_weights = balance_util.drop_na_rows(
-                target_df,
-                _assert_type(bf._sf_target).df_weights.iloc[:, 0],
-                "target",
-            )
-        elif na_action == "add_indicator":
-            sample_df = pd.DataFrame(_safe_fillna_and_infer(sample_df, "__NaN__"))
-            target_df = pd.DataFrame(_safe_fillna_and_infer(target_df, "__NaN__"))
-        else:
-            raise ValueError(
-                f"Rake model has invalid na_action metadata '{na_action}' for predict_weights()."
-            )
-        sample_df = sample_df.astype(str)
-        if m_fit.shape != m_sample.shape:
-            raise ValueError(
-                "Rake model metadata has incompatible fitted and sample table shapes."
-            )
-
-        category_maps = [{cat: i for i, cat in enumerate(cats)} for cats in categories]
-        code_columns = []
-        for column, cat_map in zip(variables, category_maps):
-            codes = sample_df[column].map(cat_map)
-            if bool(codes.isna().any()):
-                raise ValueError(
-                    "Rake predict_weights() found rows that do not map to stored fit-time "
-                    "categories. Re-fit with compatible covariates."
-                )
-            code_columns.append(codes.astype(int).to_numpy())
-        code_index = tuple(code_columns)
-
-        if source is not None:
-            logger.warning(
-                "Rake predict_weights(data=...): replaying fitted rake "
-                "artifacts on a different sample is a transfer operation, "
-                "not an exact fit. The stored cell-ratio surface "
-                "(m_fit / m_sample) encodes the *training* target's "
-                "marginal distribution; applying it to a new sample only "
-                "produces weights calibrated to that *training* target, "
-                "rescaled to the new target's total weight. The new "
-                "target's marginals are NOT re-balanced. Predictions are "
-                "therefore only valid when (a) the scoring sample's joint "
-                "distribution over the rake variables is similar to the "
-                "training sample's, AND (b) the scoring target's marginal "
-                "distribution is similar to the training target's. "
-                "Re-fit rake on the scoring sample/target for exact "
-                "marginal matching."
-            )
-        m_sample_at_cells = m_sample[code_index]
-        if bool((m_sample_at_cells <= 0).any()):
-            raise ValueError(
-                "Rake predict_weights() encountered sample rows in joint cells with "
-                "zero fit-time sample mass (m_sample==0). Re-fit rake on data with "
-                "compatible joint support."
-            )
-        # Compute ratios only for the scored cells. Avoids materializing a
-        # dense array the size of the full contingency table on every call
-        # (relevant for high-cardinality rakes).
-        cell_ratios = m_fit[code_index] / m_sample_at_cells
-        raw = pd.Series(
-            sample_weights.to_numpy() * cell_ratios,
-            index=sample_weights.index,
-            dtype=float,
-        )
-        target_weights = model.get("training_target_weights")
-        if source is None and not isinstance(target_weights, pd.Series):
-            raise ValueError(
-                "Rake predict_weights() requires compatible fit-time target design "
-                "weights for in-place replay. This can happen because "
-                "store_fit_metadata is missing/incompatible, or because you're "
-                "scoring a different sample; use predict_weights(data=...) "
-                "for different samples."
-            )
-        if source is not None:
-            if na_action == "drop" and isinstance(dropped_target_weights, pd.Series):
-                target_sum = float(dropped_target_weights.sum())
-            else:
-                target_sum = float(
-                    _assert_type(bf._sf_target).df_weights.iloc[:, 0].sum()
-                )
-        elif isinstance(target_weights, pd.Series):
-            target_sum = float(target_weights.sum())
-        elif na_action == "drop" and isinstance(dropped_target_weights, pd.Series):
-            target_sum = float(dropped_target_weights.sum())
-        else:
-            target_sum = float(_assert_type(bf._sf_target).df_weights.iloc[:, 0].sum())
-        predicted = balance_adjustment.trim_weights(
-            raw,
-            target_sum_weights=target_sum,
-            weight_trimming_mean_ratio=model.get("weight_trimming_mean_ratio"),
-            weight_trimming_percentile=model.get("weight_trimming_percentile"),
-            keep_sum_of_weights=bool(model.get("keep_sum_of_weights", True)),
-        )
         weight_name = getattr(_assert_type(bf.weight_series), "name", None)
-        if na_action == "drop":
-            predicted_full = pd.Series(
-                np.nan, index=sample_weights_full.index, dtype=float
-            ).rename(predicted.name)
-            predicted_full.loc[predicted.index] = predicted.to_numpy()
-            predicted = predicted_full
         return cast(
             pd.Series,
             self._align_to_index(