Add rake fit metadata and predict_weights reconstruction (#435) - #435
Closed
talgalili wants to merge 1 commit into
Closed
Add rake fit metadata and predict_weights reconstruction (#435)#435talgalili wants to merge 1 commit into
predict_weights reconstruction (#435)#435talgalili wants to merge 1 commit into
Conversation
Contributor
|
@talgalili has exported this pull request. If you are a Meta employee, you can view the originating Diff in D103532489. |
There was a problem hiding this comment.
Pull request overview
This PR extends the balance weighting workflow so rake fits can persist enough artifacts to reconstruct weights later via BalanceFrame.predict_weights(). It fits into the recent BalanceFrame.fit / replayable-model work by bringing rake closer to the same artifact-driven prediction flow already used for other weighting methods.
Changes:
- Add optional rake fit-metadata persistence, including stored contingency-table artifacts and training weights.
- Add rake support to
BalanceFrame.predict_weights()for in-place reconstruction anddata=...transfer scoring. - Add rake-focused tests and a changelog entry describing the new metadata/prediction behavior.
Reviewed changes
Copilot reviewed 5 out of 5 changed files in this pull request and generated 4 comments.
Show a summary per file
| File | Description |
|---|---|
tests/test_rake.py |
Adds unit tests for rake fit-metadata storage behavior. |
tests/test_balance_frame.py |
Adds end-to-end tests for rake predict_weights() reconstruction and transfer scenarios. |
balance/weighting_methods/rake.py |
Persists rake fit artifacts and documents the new reconstruction workflow. |
balance/balance_frame.py |
Wires rake into fit() defaults and predict_weights() dispatch/reconstruction logic. |
CHANGELOG.md |
Documents the new rake metadata and prediction support. |
Comment on lines
+403
to
+415
| if store_fit_metadata: | ||
| model.update( | ||
| { | ||
| "store_fit_metadata": True, | ||
| "variables": alphabetized_variables, | ||
| "variables_before_transformations": list(variables), | ||
| "categories": categories, | ||
| "m_fit": m_fit, | ||
| "m_sample": m_sample, | ||
| "na_action": na_action, | ||
| "transformations": transformations_to_apply, | ||
| "transformations_origin": transformations, | ||
| "training_sample_weights": sample_weights.copy(), |
Comment on lines
+2139
to
+2143
| 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 " |
Comment on lines
+2234
to
+2239
| ratio = np.divide( | ||
| m_fit, | ||
| m_sample, | ||
| out=np.zeros_like(m_fit, dtype=float), | ||
| where=m_sample != 0, | ||
| ) |
Comment on lines
+188
to
+189
| For this reason, balance always emits a warning on transferred scoring | ||
| (``predict_weights(data=...)``). |
meta-codesync
Bot
changed the title
predict_weights reconstruction (#433)predict_weights reconstruction (#435)
talgalili
pushed a commit
to talgalili/balance
that referenced
this pull request
May 3, 2026
…esearch#435) Summary: Differential Revision: D103532489 Pulled By: talgalili
talgalili
force-pushed
the
export-D103532489
branch
from
May 3, 2026 09:59
5b6d8c9 to
f2644c0
Compare
talgalili
pushed a commit
to talgalili/balance
that referenced
this pull request
May 3, 2026
…esearch#435) Summary: Differential Revision: D103532489 Pulled By: talgalili
talgalili
force-pushed
the
export-D103532489
branch
from
May 3, 2026 10:14
f2644c0 to
1eda581
Compare
Comment on lines
+2300
to
+2305
| 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() | ||
| ) |
Comment on lines
+2149
to
+2164
| if source is not None and isinstance(transformations_origin, dict): | ||
| # Explicit dicts can still contain data-dependent helpers | ||
| # (quantize, fct_lump, ...) whose bins/levels are recomputed from | ||
| # the scoring data. Stored cell ratios then no longer line up | ||
| # with the transformed scoring cells, so transferred weights | ||
| # would be silently incorrect. Reject the known cases up-front; | ||
| # users who really need this can pass deterministic wrappers. | ||
| 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 | ||
| } |
Comment on lines
+1121
to
+1122
| if resolved_method is built_in_rake: | ||
| kwargs.setdefault("store_fit_metadata", True) |
Comment on lines
+19
to
+21
| - `BalanceFrame.fit(method="rake")` now enables `store_fit_metadata=True` by | ||
| default so fitted rake models can be reused with | ||
| `BalanceFrame.predict_weights()` without refitting. |
talgalili
pushed a commit
to talgalili/balance
that referenced
this pull request
May 3, 2026
…esearch#435) Summary: Differential Revision: D103532489 Pulled By: talgalili
talgalili
force-pushed
the
export-D103532489
branch
from
May 3, 2026 11:31
1eda581 to
46264e1
Compare
talgalili
pushed a commit
to talgalili/balance
that referenced
this pull request
May 3, 2026
…esearch#435) Summary: Differential Revision: D103532489 Pulled By: talgalili
talgalili
pushed a commit
to talgalili/balance
that referenced
this pull request
May 3, 2026
…esearch#435) Summary: Differential Revision: D103532489 Pulled By: talgalili
talgalili
force-pushed
the
export-D103532489
branch
from
May 3, 2026 13:46
46264e1 to
a40a024
Compare
Comment on lines
+2159
to
+2184
| 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 resolved_method is built_in_poststratify: | ||
| kwargs.setdefault("store_fit_metadata", True) | ||
| if resolved_method is built_in_rake: | ||
| kwargs.setdefault("store_fit_metadata", True) |
Comment on lines
+2325
to
+2331
| 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() | ||
| ) |
Comment on lines
+12
to
+15
| silently breaking. **Migration**: convert any lambda/closure | ||
| transformations to top-level functions or | ||
| `functools.partial(...)` of top-level functions, or pass | ||
| `store_fit_metadata=False` to opt out of fit-metadata persistence |
…esearch#435) Summary: Differential Revision: D103532489 Pulled By: talgalili
talgalili
force-pushed
the
export-D103532489
branch
from
May 3, 2026 13:57
a40a024 to
3bf51fc
Compare
Contributor
|
@talgalili merged this pull request in 38142c6. |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary:
Differential Revision: D103532489
Pulled By: talgalili