balance.balance_frame

BalanceFrame: workflow orchestrator for survey/observational data reweighting.

Pairs a responder SampleFrame with a target SampleFrame and exposes an immutable adjust() method that returns a new, weight-augmented BalanceFrame.

class balance.balance_frame.BalanceFrame(sample: SampleFrame | None = None, target: SampleFrame | None = None)

A pairing of responder and target SampleFrames for survey weighting.

BalanceFrame holds two SampleFrame instances — responders (the sample to be reweighted) and target (the population benchmark) — and provides methods for adjusting responder weights and computing diagnostics.

BalanceFrame is immutable by convention: adjust() returns a new BalanceFrame rather than modifying the existing one. This makes it safe to keep a reference to the pre-adjustment state.

Must be constructed via the public constructor BalanceFrame(sample=..., target=...) which delegates to the internal _create() factory.

responders

The responder sample.

Type:

SampleFrame

target

The target population.

Type:

SampleFrame

is_adjusted

Whether adjust() has been called.

Type:

bool

Examples

>>> import pandas as pd
>>> from balance.sample_frame import SampleFrame
>>> from balance.balance_frame import BalanceFrame
>>> resp = SampleFrame.from_frame(
...     pd.DataFrame({"id": [1, 2], "x": [10.0, 20.0], "weight": [1.0, 1.0]}))
>>> tgt = SampleFrame.from_frame(
...     pd.DataFrame({"id": [3, 4], "x": [15.0, 25.0], "weight": [1.0, 1.0]}))
>>> bf = BalanceFrame(sample=resp, target=tgt)
>>> bf.is_adjusted
False
>>> adjusted = bf.adjust(method="ipw")
>>> adjusted.is_adjusted
True
>>> bf.is_adjusted  # original unchanged
False

adjust(target: BalanceFrame | None = None, method: str | Callable[..., Any] = 'ipw', *args: Any, **kwargs: Any) → Self

Adjust responder weights to match the target. Returns a NEW BalanceFrame.

The original BalanceFrame is not modified (immutable pattern). The returned BalanceFrame has is_adjusted == True and the pre-adjustment responders stored in unadjusted.

The active weight column always keeps its original name (e.g., "weight"). Its values are overwritten with the new adjusted weights. The full weight history is tracked via additional columns:

Weight columns after each adjustment

After	Weight columns in responders	Active ("weight")
Before adjust	weight	original design weights
1st adjust	weight, weight_pre_adjust, weight_adjusted_1	= weight_adjusted_1 values
2nd adjust	weight_adjusted_2	= weight_adjusted_2 values
3rd adjust	weight_adjusted_3	= weight_adjusted_3 values

Compound / sequential adjustments: adjust() can be called multiple times. Each call uses the current (previously adjusted) weights as design weights, so adjustments compound. For example, run IPW first to correct broad imbalances, then rake on a specific variable for fine-tuning:

adjusted_ipw = bf.adjust(method="ipw", max_de=2)
adjusted_final = adjusted_ipw.adjust(method="rake")

The original unadjusted baseline is always preserved:

• _sf_sample_pre_adjust always points to the original (pre-first-adjustment) SampleFrame.

• _links["unadjusted"] always points to the original unadjusted BalanceFrame, so 3-way comparisons (adjusted vs original vs target) and asmd_improvement() show total improvement across all adjustment steps.

• model stores only the latest adjustment’s model dict.

Parameters:

• target – Optional target BalanceFrame/Sample. If provided, calls set_target(target) first, then adjusts. If None, uses the already-set target.

• method – The weighting method to use. Built-in options: "ipw", "cbps", "rake", "poststratify", "null". A callable with the same signature as the built-in methods is also accepted.

• *args – Positional arguments (forwarded on recursive call only).

• **kwargs – Additional keyword arguments forwarded to the adjustment function (e.g. max_de, transformations).

Returns:

A new, adjusted BalanceFrame.

Raises:

ValueError – If method is a string that doesn’t match any registered adjustment method.

Examples

>>> import pandas as pd
>>> from balance.sample_frame import SampleFrame
>>> from balance.balance_frame import BalanceFrame
>>> resp = SampleFrame.from_frame(
...     pd.DataFrame({"id": [1, 2, 3], "x": [10.0, 20.0, 30.0],
...                   "weight": [1.0, 1.0, 1.0]}))
>>> tgt = SampleFrame.from_frame(
...     pd.DataFrame({"id": [4, 5, 6], "x": [15.0, 25.0, 35.0],
...                   "weight": [1.0, 1.0, 1.0]}))
>>> bf = BalanceFrame(sample=resp, target=tgt)
>>> adjusted = bf.adjust(method="ipw")
>>> adjusted.is_adjusted
True
>>> adjusted2 = adjusted.adjust(method="null")
>>> adjusted2.is_adjusted
True

property adjustment_history: list[dict[str, Any]]

Chronological adjustment model history.

Returns a best-effort read-only copy of each recorded adjustment step so callers can inspect compound reweighting workflows without mutating internal state. The model property remains the latest adjustment model for backward compatibility.

Deep-copy is attempted per entry; when a payload is not deepcopy-safe, the method falls back to copying the step dictionary and nested model mapping (when present).

Examples

>>> import pandas as pd
>>> from balance.sample_frame import SampleFrame
>>> from balance.balance_frame import BalanceFrame
>>> resp = SampleFrame.from_frame(
...     pd.DataFrame({"id": [1, 2], "x": [1.0, 2.0], "weight": [1.0, 1.0]}))
>>> tgt = SampleFrame.from_frame(
...     pd.DataFrame({"id": [3, 4], "x": [1.5, 2.5], "weight": [1.0, 1.0]}))
>>> bf = BalanceFrame(sample=resp, target=tgt).adjust(method="null")
>>> len(bf.adjustment_history)
1

aipw() → Series

Doubly-robust (AIPW) estimate μ̂_DR of the target-population mean.

Combines the fitted outcome model ĝ (from fit_outcome_model()) with the responder weights w to form the augmented / one-sample AIPW (doubly-robust) estimate of the target mean, per outcome column:

μ̂_DR = wmean(ĝ(X_T), w_T) + wmean(Y − ĝ(X_S), w)

(the augmentation runs over responders with an observed Y). It is consistent if either the outcome model or the weighting model is correct. Companion to the other estimate accessors: outcomes().mean() gives μ̂_IPW and outcomes_hat().mean() gives μ̂_OM; this gives μ̂_DR.

Requires a fitted outcome model and a target (set_target()). Uses whatever weight columns are present (any balance weights). With constant responder weights it warns and reduces to μ̂_OM.

Point estimate only – no confidence interval (an honest AIPW interval must jointly capture the weighting- and outcome-model uncertainty; see the TODOs in balance.outcome_models.aipw).

Returns:

μ̂_DR indexed by outcome column.

Return type:

pd.Series

Raises:

ValueError – If no outcome model has been fit, if no target is set, or if the responders carry no observed outcomes.

Examples

>>> import pandas as pd
>>> from sklearn.linear_model import LinearRegression
>>> from balance.sample_frame import SampleFrame
>>> from balance.balance_frame import BalanceFrame
>>> resp = SampleFrame.from_frame(
...     pd.DataFrame({"id": [1, 2, 3, 4],
...                   "x": [10.0, 20.0, 30.0, 40.0],
...                   "y": [1.0, 2.0, 3.0, 4.0],
...                   "weight": [1.0, 2.0, 1.0, 2.0]}),
...     outcome_columns=["y"])
>>> tgt = SampleFrame.from_frame(
...     pd.DataFrame({"id": [5, 6], "x": [15.0, 35.0],
...                   "weight": [1.0, 1.0]}))
>>> bf = BalanceFrame(sample=resp, target=tgt)
>>> _ = bf.fit_outcome_model(model=LinearRegression())
>>> bf.aipw().index.tolist()
['y']

covars(formula: str | list[str] | None = None) → Any

Return a BalanceDFCovars for the responders.

The returned object carries linked target (and unadjusted, if adjusted) views so that methods like .mean() and .asmd() automatically include comparisons across sources.

Parameters:

formula – Optional formula string (or list) for model matrix construction. Passed through to BalanceDFCovars.

Returns:

Covariate view with linked sources.

Return type:

BalanceDFCovars

Examples

>>> import pandas as pd
>>> from balance.sample_frame import SampleFrame
>>> from balance.balance_frame import BalanceFrame
>>> resp = SampleFrame.from_frame(
...     pd.DataFrame({"id": [1, 2], "x": [10.0, 20.0], "weight": [1.0, 1.0]}))
>>> tgt = SampleFrame.from_frame(
...     pd.DataFrame({"id": [3, 4], "x": [15.0, 25.0], "weight": [1.0, 1.0]}))
>>> bf = BalanceFrame(sample=resp, target=tgt)
>>> bf.covars().df.columns.tolist()
['x']

design_matrix(on: Literal['sample'], *, data: BalanceFrame | None = None) → DataFrame

design_matrix(on: Literal['target'], *, data: BalanceFrame | None = None) → DataFrame

design_matrix(on: Literal['both'] = 'both', *, data: BalanceFrame | None = None) → tuple[DataFrame, DataFrame]

Return the IPW model’s design matrices.

Returns the model matrices (feature matrices) built by the stored preprocessing pipeline — after formula expansion, one-hot encoding, NA indicator addition, scaling, and penalty weighting.

When data is provided, the stored preprocessing is applied to data’s covariates and the result is returned without caching. When data is None (default), stored/cached matrices for this object’s own data are returned (original behavior).

Parameters:

• on – Which population’s matrix to return. "sample" returns the respondent matrix, "target" returns the target matrix, and "both" returns (sample_matrix, target_matrix).

• data – An optional BalanceFrame whose covariates are transformed using this object’s stored preprocessing pipeline. The data BalanceFrame does not need to be adjusted — it just provides covariates. Must have matching covariate column names.

Returns:

A model-matrix DataFrame, or a tuple of two DataFrames when on="both".

Raises:

ValueError – If the object is not IPW-adjusted, if target is missing for on in {"target", "both"}, if recomputation of sample-side artifacts is required but no target is available, if on is invalid, or if data has mismatched covariate columns.

Notes

When data is None and stored fit artifacts are stale for the current rows (e.g., after set_fitted_model()), this method recomputes and caches refreshed matrices. That cache update is an intentional in-memory mutation. When data is provided, no caching occurs.

Examples

>>> import pandas as pd
>>> from balance.sample_frame import SampleFrame
>>> from balance.balance_frame import BalanceFrame
>>> resp = SampleFrame.from_frame(
...     pd.DataFrame({"id": [1, 2], "x": [0.0, 1.0], "weight": [1.0, 1.0]}))
>>> tgt = SampleFrame.from_frame(
...     pd.DataFrame({"id": [3, 4], "x": [0.2, 0.8], "weight": [1.0, 1.0]}))
>>> adjusted = BalanceFrame(sample=resp, target=tgt).fit(method="ipw")
>>> x_s, x_t = adjusted.design_matrix(on="both")
>>> x_s.shape[0], x_t.shape[0]
(2, 2)

property df: DataFrame

Flat user-facing DataFrame from the responders.

Returns the responder data with columns ordered as: id → covariates → outcomes → weight → ignored.

Returns:

Ordered copy of the responder’s data.

Return type:

pd.DataFrame

property df_all: DataFrame

Combined DataFrame with all samples, distinguished by a "source" column.

Concatenates the responder, target, and (if adjusted) unadjusted DataFrames vertically, adding a "source" column with values "self", "target", and "unadjusted" respectively.

Returns:

A DataFrame with all rows from responder, target,

and optionally unadjusted SampleFrames, plus a "source" column.

Return type:

pd.DataFrame

Examples

>>> import pandas as pd
>>> from balance.sample_frame import SampleFrame
>>> from balance.balance_frame import BalanceFrame
>>> resp = SampleFrame.from_frame(
...     pd.DataFrame({"id": [1, 2], "x": [10.0, 20.0], "weight": [1.0, 1.0]}))
>>> tgt = SampleFrame.from_frame(
...     pd.DataFrame({"id": [3, 4], "x": [15.0, 25.0], "weight": [1.0, 1.0]}))
>>> bf = BalanceFrame(sample=resp, target=tgt)
>>> bf.df_all["source"].unique().tolist()
['self', 'target']

property df_ignored: DataFrame | None

Ignored columns from the responder SampleFrame, or None.

property df_outcomes_hat: DataFrame | None

Predicted-outcome (Y_hat) columns from the responder, or None.

property df_responders: DataFrame

The responder data as a DataFrame.

property df_responders_unadjusted: DataFrame

The original (pre-adjustment) responder data as a DataFrame.

property df_target: DataFrame | None

The target data as a DataFrame, or None if not yet set.

diagnostics(weights_impact_on_outcome_method: str | None = 't_test', weights_impact_on_outcome_conf_level: float = 0.95) → DataFrame

Table of diagnostics about the adjusted BalanceFrame.

Produces a DataFrame with columns ["metric", "val", "var"] containing size information, weight diagnostics, model details, covariate ASMD, and optionally outcome-weight impact statistics. Delegates to _build_diagnostics().

Parameters:

• weights_impact_on_outcome_method – Method for computing outcome-weight impact. Pass None to skip. Defaults to "t_test".

• weights_impact_on_outcome_conf_level – Confidence level for outcome impact intervals. Defaults to 0.95.

Returns:

Diagnostics table with columns

["metric", "val", "var"].

Return type:

pd.DataFrame

Raises:

ValueError – If this BalanceFrame has not been adjusted.

Examples

>>> import pandas as pd
>>> from balance.sample_frame import SampleFrame
>>> from balance.balance_frame import BalanceFrame
>>> resp = SampleFrame.from_frame(
...     pd.DataFrame({"id": ["1", "2"], "x": [0, 1],
...                   "weight": [1.0, 2.0]}))
>>> tgt = SampleFrame.from_frame(
...     pd.DataFrame({"id": ["3", "4"], "x": [0, 1],
...                   "weight": [1.0, 1.0]}))
>>> bf = BalanceFrame(sample=resp, target=tgt)
>>> adjusted = bf.adjust(method="null")
>>> adjusted.diagnostics().columns.tolist()
['metric', 'val', 'var']

fit(*, target: BalanceFrame | SampleFrame | None = None, method: str | Callable[..., Any] = 'ipw', inplace: bool = True, **kwargs: Any) → Self

Fit a weighting model and return the fitted BalanceFrame.

This is the sklearn-style entry point for survey weight adjustment. Like sklearn’s fit(), it learns model parameters, mutates self (by default), and returns self. In survey weighting, fitting the propensity model inherently produces adjusted weights (the two are inseparable), so the returned object contains both the fitted model and the adjusted weights — analogous to how KMeans.fit() stores labels_ on the fitted object.

Workflow — basic fitting (sklearn-style, inplace=True):

bf = BalanceFrame(sample=respondents, target=population)
bf.fit(method="ipw")       # mutates bf, returns bf
bf.weights().df            # the adjusted weights

Workflow — functional style (inplace=False):

adjusted = bf.fit(method="ipw", inplace=False)

Workflow — fit on subset, apply to holdout:

fitted = train_bf.fit(method="ipw")
scored = holdout_bf.set_fitted_model(fitted, inplace=False)
holdout_weights = scored.predict_weights()

Alternatively, design_matrix(), predict_proba(), and predict_weights() accept a data= argument so the holdout workflow becomes a single line: fitted.predict_weights(data=holdout_bf).

Parameters:

• target – Optional target population to set before fitting. If provided, this method calls set_target(target, inplace=False) first, preserving immutability.

• method – Adjustment method name ("ipw", "cbps", "rake", "poststratify", "null") or a custom callable with the weighting-method signature.

• inplace – If True (default), mutate this object with the fitted state and return self — matching sklearn’s fit() convention. If False, return a new adjusted BalanceFrame without modifying self.

• **kwargs – Keyword arguments forwarded to adjust().

Returns:

The fitted BalanceFrame — self when inplace=True, a new object when inplace=False.

Raises:

ValueError – If no target is available and none is provided, if method is invalid, or if na_action='drop' is combined with stored fit artifacts.

Examples

>>> import pandas as pd
>>> from balance.sample_frame import SampleFrame
>>> from balance.balance_frame import BalanceFrame
>>> resp = SampleFrame.from_frame(
...     pd.DataFrame({"id": [1, 2], "x": [0.0, 1.0], "weight": [1.0, 1.0]}))
>>> tgt = SampleFrame.from_frame(
...     pd.DataFrame({"id": [3, 4], "x": [0.2, 0.8], "weight": [1.0, 1.0]}))
>>> adjusted = BalanceFrame(sample=resp, target=tgt).fit(method="null")
>>> bool(adjusted.is_adjusted)
True

Notes

For the built-in IPW method, fit() enables store_fit_metadata=True and store_fit_matrices=True by default so design_matrix()/predict_proba()/predict_weights() can consume fit-time artifacts. This may increase memory usage for large inputs; pass these kwargs explicitly as False to opt out. For the built-in CBPS method, fit() enables store_fit_metadata=True by default so predict_weights() can reconstruct CBPS scoring artifacts. Pass store_fit_metadata=False to opt out. For the built-in poststratify method, fit() enables store_fit_metadata=True by default so predict_weights() can reconstruct poststratification cell-ratio artifacts, while direct adjust(method='poststratify') remains metadata-light unless store_fit_metadata=True is passed explicitly. For the built-in rake method, fit() enables store_fit_metadata=True by default so predict_weights() can replay/transfer fitted rake artifacts. This may increase memory usage because contingency tables and fit metadata are stored; pass store_fit_metadata=False to opt out. When metadata is stored, transformations must be pickleable (lambdas/closures are rejected at fit time so the resulting BalanceFrame remains serializable). Note: in-place predict_weights() works for any rake fit, including the default transformations='default'. Transfer scoring (predict_weights(data=...)), however, rejects models fitted with transformations='default' or with explicit dicts containing known data-dependent helpers (quantize, fct_lump) because those transformations recompute bins/levels from each scoring sample, which would silently invalidate the stored cell ratios. To use rake’s data=... transfer path, pass deterministic transformations at fit time or re-fit rake on the scoring data.

fit_outcome_model(*, target: BalanceFrame | SampleFrame | None = None, inplace: bool = True, **kwargs: Any) → Self

Fit an outcome model ĝ(X) ≈ E[Y|X] on the responders and store it.

This is the BalanceFrame entry point to the outcome-modelling axis (the counterpart to fit() for the weighting axis). It delegates to the responder (self._sf_sample): the fitted model is stored on _sf_sample._outcome_model — the single source of truth read back through outcome_model. Because the model lives on the responder SampleFrame, it rides along whenever _sf_sample is deep-copied — in particular it survives :meth:`adjust` (which deep-copies _sf_sample) and is preserved across set_target().

On a Sample, this BalanceFrame method takes MRO precedence over SampleFrame.fit_outcome_model (Sample is both classes), so the model correctly lands on _sf_sample rather than on the Sample’s own inherited attribute.

Like sklearn’s fit() (and SampleFrame.fit_outcome_model()), it does not produce predictions — call predict_outcomes() (or fit_predict_outcomes()) to write the <outcome>_hat columns.

Parameters:

• target – Optional target population to set before fitting (mirrors fit()). When provided, set_target(target, inplace=False) is applied first so immutability is preserved; the fit then proceeds on the returned object.

• inplace – If True (default), mutate this object (store the model on the responder) and return self; if False, return a new copy with the model stored and leave self untouched — matching fit() / SampleFrame.fit_outcome_model().

• **kwargs – Forwarded to SampleFrame.fit_outcome_model() (e.g. model, outcome_columns, variables, weighted, formula, na_action, use_model_matrix, calibrate).

Returns:

The BalanceFrame with the fitted outcome model stored on its responder (self when inplace, else a new copy).

Raises:

ValueError – If there are no outcome columns to fit (and none are passed), or for any underlying fit_outcome_model error.

Examples

>>> import pandas as pd
>>> from balance.sample_frame import SampleFrame
>>> from balance.balance_frame import BalanceFrame
>>> resp = SampleFrame.from_frame(
...     pd.DataFrame({"id": ["1", "2", "3", "4"],
...                   "age": [25.0, 30.0, 35.0, 40.0],
...                   "happiness": [50.0, 55.0, 65.0, 80.0],
...                   "weight": [1.0, 1.0, 1.0, 1.0]}),
...     outcome_columns=["happiness"])
>>> tgt = SampleFrame.from_frame(
...     pd.DataFrame({"id": ["5", "6"], "age": [28.0, 38.0],
...                   "weight": [1.0, 1.0]}))
>>> bf = BalanceFrame(sample=resp, target=tgt)
>>> _ = bf.fit_outcome_model()
>>> bf.outcome_model["method"]
'outcome_model'

End-to-end on the bundled load_data() dataset — the g-computation estimate μ̂_OM (the weighted mean of the predicted outcome on the target). No adjust() is required: the outcome model is fit on the responders independently of the weights.

from balance import load_data, Sample

target_df, sample_df = load_data()
bf = Sample.from_frame(
    sample_df, outcome_columns=["happiness"]
).set_target(Sample.from_frame(target_df))
bf.fit_outcome_model(model="auto")   # fit ĝ(X) ≈ E[Y|X] on the responders
bf.predict_outcomes(on="target")     # score the target with the fitted model
bf.outcomes_hat().mean()             # target row = μ̂_OM (g-computation estimate)

fit_predict_outcomes(*, target: BalanceFrame | SampleFrame | None = None, on: Literal['sample', 'target', 'both'] | None = None, populate: bool = True, **fit_kwargs: Any) → DataFrame | tuple[DataFrame, DataFrame]

Fit an outcome model then replay it (on the target by default).

Convenience wrapper: calls fit_outcome_model() (in place, on the responder) with **fit_kwargs and then predict_outcomes() with the given on / populate.

Parameters:

• target – Optional target to set before fitting (see fit_outcome_model()).

• on – Which population(s) to predict on afterwards. None (the default) resolves to "target" when a target is set, else "sample" (see predict_outcomes()).

• populate – Whether to persist the <outcome>_hat columns.

• **fit_kwargs – Forwarded to fit_outcome_model().

Returns:

The predictions DataFrame (or a (sample, target) tuple when on="both"), from predict_outcomes().

Examples

>>> import pandas as pd
>>> from balance.sample_frame import SampleFrame
>>> from balance.balance_frame import BalanceFrame
>>> resp = SampleFrame.from_frame(
...     pd.DataFrame({"id": ["1", "2", "3", "4"],
...                   "age": [25.0, 30.0, 35.0, 40.0],
...                   "happiness": [50.0, 55.0, 65.0, 80.0],
...                   "weight": [1.0, 1.0, 1.0, 1.0]}),
...     outcome_columns=["happiness"])
>>> tgt = SampleFrame.from_frame(
...     pd.DataFrame({"id": ["5", "6"], "age": [28.0, 38.0],
...                   "weight": [1.0, 1.0]}))
>>> bf = BalanceFrame(sample=resp, target=tgt)
>>> preds = bf.fit_predict_outcomes(on="target")
>>> list(preds.columns)
['happiness_hat']

classmethod from_sample(sample: Any) → BalanceFrame

Convert a Sample to a BalanceFrame.

The Sample must have a target set (via Sample.set_target). If the Sample is adjusted, the adjustment state (unadjusted responders, model) is preserved.

Parameters:

sample – A Sample instance with a target.

Returns:

A new BalanceFrame mirroring the Sample’s data,

target, and adjustment state.

Return type:

BalanceFrame

Raises:

• TypeError – If sample is not a Sample instance.

• ValueError – If sample does not have a target set.

Examples

>>> import pandas as pd
>>> from balance.sample_class import Sample
>>> from balance.balance_frame import BalanceFrame
>>> s = Sample.from_frame(
...     pd.DataFrame({"id": [1, 2], "x": [10.0, 20.0], "weight": [1.0, 1.0]}))
>>> t = Sample.from_frame(
...     pd.DataFrame({"id": [3, 4], "x": [15.0, 25.0], "weight": [1.0, 1.0]}))
>>> bf = BalanceFrame.from_sample(s.set_target(t))
>>> bf.is_adjusted
False

property has_target: _CallableBool

Whether this BalanceFrame has a target population set.

Returns a dual-use _CallableBool: both bf.has_target and bf.has_target() work (the latter for backward compatibility).

Examples

>>> import pandas as pd
>>> from balance.sample_frame import SampleFrame
>>> from balance.balance_frame import BalanceFrame
>>> resp = SampleFrame.from_frame(
...     pd.DataFrame({"id": [1, 2], "x": [10.0, 20.0], "weight": [1.0, 1.0]}))
>>> bf = BalanceFrame(sample=resp)
>>> bf.has_target
False
>>> tgt = SampleFrame.from_frame(
...     pd.DataFrame({"id": [3, 4], "x": [15.0, 25.0], "weight": [1.0, 1.0]}))
>>> bf.set_target(tgt)
>>> bf.has_target
True

property id_column: str | None

The id column name, delegated to _sf_sample.

Changed in 0.20.0 to return the name (str) instead of data (pd.Series). Use id_series for data.

property id_series: Series | None

The id column as a Series, delegated to _sf_sample.

property is_adjusted: _CallableBool

Whether this BalanceFrame has been adjusted.

Returns a _CallableBool so both bf.is_adjusted (property) and bf.is_adjusted() (legacy call) work.

For compound adjustments (calling adjust() multiple times), is_adjusted is True after the first adjustment and remains True for all subsequent adjustments. The original unadjusted baseline is always preserved in _sf_sample_pre_adjust.

keep_only_some_rows_columns(rows_to_keep: str | None = None, columns_to_keep: list[str] | None = None) → BalanceFrame

Return a new BalanceFrame with filtered rows and/or columns.

Returns a deep copy with the requested subset applied to the responder, target, and (if adjusted) unadjusted SampleFrames. The original BalanceFrame is unchanged (immutable pattern).

Parameters:

• rows_to_keep – A boolean expression string evaluated via pd.DataFrame.eval to select rows. Applied to each SampleFrame’s underlying DataFrame. For example: 'x > 10' or 'gender == "Female"'. Defaults to None (all rows kept).

• columns_to_keep – Covariate column names to retain. Special columns (id, weight) are always kept. Defaults to None (all columns kept).

Returns:

A new BalanceFrame with the filters applied.

Return type:

BalanceFrame

Examples

>>> import pandas as pd
>>> from balance.sample_frame import SampleFrame
>>> from balance.balance_frame import BalanceFrame
>>> resp = SampleFrame.from_frame(
...     pd.DataFrame({"id": [1, 2, 3], "x": [10.0, 20.0, 30.0],
...                   "weight": [1.0, 1.0, 1.0]}))
>>> tgt = SampleFrame.from_frame(
...     pd.DataFrame({"id": [4, 5, 6], "x": [15.0, 25.0, 35.0],
...                   "weight": [1.0, 1.0, 1.0]}))
>>> bf = BalanceFrame(sample=resp, target=tgt)
>>> filtered = bf.keep_only_some_rows_columns(rows_to_keep="x > 15")
>>> len(filtered._sf_sample._df)
2

property model: dict[str, Any] | None

The adjustment model dictionary, or None if not adjusted.

Examples

>>> import pandas as pd
>>> from balance.sample_frame import SampleFrame
>>> from balance.balance_frame import BalanceFrame
>>> resp = SampleFrame.from_frame(
...     pd.DataFrame({"id": [1, 2], "x": [10.0, 20.0], "weight": [1.0, 1.0]}))
>>> tgt = SampleFrame.from_frame(
...     pd.DataFrame({"id": [3, 4], "x": [15.0, 25.0], "weight": [1.0, 1.0]}))
>>> bf = BalanceFrame(sample=resp, target=tgt)
>>> bf.model is None
True

model_matrix() → DataFrame

Return the model matrix of the responder covariates.

Constructs a model matrix using balance.util.model_matrix(), adding NA indicators for null values.

Returns:

The model matrix.

Return type:

pd.DataFrame

property outcome_model: dict[str, Any] | None

The fitted outcome-model dict, delegated to _sf_sample.

Single source of truth: the outcome model lives on the wrapped responder SampleFrame (_sf_sample._outcome_model), so it rides along whenever _sf_sample is deep-copied (e.g. by adjust()) and a Sample — which is both a BalanceFrame and a SampleFrame — never reports a stale inherited copy. Mirrors how _outcome_columns / df_outcomes_hat read through to _sf_sample, and parallels model (the weighting model) for the outcome-modelling axis.

Returns:

The stored _outcome_model dict from the responder, or None if no outcome model has been fit.

Return type:

dict[str, Any] | None

Examples

>>> import pandas as pd
>>> from balance.sample_frame import SampleFrame
>>> from balance.balance_frame import BalanceFrame
>>> resp = SampleFrame.from_frame(
...     pd.DataFrame({"id": [1, 2], "x": [10.0, 20.0],
...                   "y": [1.0, 0.0], "weight": [1.0, 1.0]}),
...     outcome_columns=["y"])
>>> tgt = SampleFrame.from_frame(
...     pd.DataFrame({"id": [3, 4], "x": [15.0, 25.0], "weight": [1.0, 1.0]}))
>>> bf = BalanceFrame(sample=resp, target=tgt)
>>> bf.outcome_model is None
True

outcomes() → Any | None

Return a BalanceDFOutcomes, or None.

Returns None if the responder SampleFrame has no outcome columns.

Returns:

Outcome view with linked sources,

or None if no outcomes are defined.

Return type:

BalanceDFOutcomes or None

Examples

>>> import pandas as pd
>>> from balance.sample_frame import SampleFrame
>>> from balance.balance_frame import BalanceFrame
>>> resp = SampleFrame.from_frame(
...     pd.DataFrame({"id": [1, 2], "x": [10.0, 20.0],
...                   "y": [1.0, 0.0], "weight": [1.0, 1.0]}),
...     outcome_columns=["y"])
>>> tgt = SampleFrame.from_frame(
...     pd.DataFrame({"id": [3, 4], "x": [15.0, 25.0], "weight": [1.0, 1.0]}))
>>> bf = BalanceFrame(sample=resp, target=tgt)
>>> bf.outcomes().df.columns.tolist()
['y']

outcomes_hat() → Any | None

Return a BalanceDFOutcomesHat, or None.

Returns None if the responder SampleFrame has no predicted-outcome (outcomes_hat / Y_hat) columns. The returned view carries linked target (and unadjusted, if adjusted) sources so that .mean() and .mean_with_ci() automatically expand across the sources that actually carry predicted outcomes – mirroring outcomes(). In particular, the weighted mean of the target’s Y_hat is the g-computation / outcome-model estimate μ̂_OM.

Returns:

Predicted-outcome view with linked

sources, or None if no outcomes_hat columns are defined.

Return type:

BalanceDFOutcomesHat or None

Examples

>>> import pandas as pd
>>> from balance.sample_frame import SampleFrame
>>> from balance.balance_frame import BalanceFrame
>>> resp = SampleFrame.from_frame(
...     pd.DataFrame({"id": [1, 2], "x": [10.0, 20.0], "weight": [1.0, 1.0]}))
>>> tgt = SampleFrame.from_frame(
...     pd.DataFrame({"id": [3, 4], "x": [15.0, 25.0], "weight": [1.0, 1.0]}))
>>> bf = BalanceFrame(sample=resp, target=tgt)
>>> bf.outcomes_hat() is None
True
>>> resp.add_outcomes_hat_column("y_hat", pd.Series([0.3, 0.7]))
>>> bf2 = BalanceFrame(sample=resp, target=tgt)
>>> bf2.outcomes_hat().df.columns.tolist()
['y_hat']

predict_outcomes(*, on: Literal['sample', 'target'] | None = None, data: BalanceFrame | None = None, populate: bool = True) → DataFrame

predict_outcomes(*, on: Literal['both'], data: BalanceFrame | None = None, populate: bool = True) → tuple[DataFrame, DataFrame]

Replay the fitted outcome model onto the sample and/or the target.

Requires a model fit on the responder (via fit_outcome_model()). The responder’s stored preprocessing is replayed on the requested population’s covariates and the <outcome>_hat columns are populated.

• on="target" (the default when a target is set): replay on the

  target covariates and populate the target’s outcomes_hat columns. The target frame is deep-copied before it is written (and the populated copy is stored back on self._sf_target), so a caller’s target object is never mutated in place. The weighted mean of the target’s outcomes_hat is the g-computation estimate μ̂_OM, read via outcomes_hat().mean().

• on="sample": populate the responder (_sf_sample) in place —

  equivalent to _sf_sample.predict_outcomes().

• on="both": do both and return ``(sample_predictions,

  target_predictions)`` — mirroring the on= return convention of predict_weights() / design_matrix().

Parameters:

• on – Which population(s) to score/populate. None (the default) resolves to "target" when a target is set (the estimate workflow) and to "sample" otherwise (so a target-less frame behaves like SampleFrame.predict_outcomes).

• data – Reserved for parity with predict_weights(); not yet supported here — pass None (foreign-frame transfer is set_fitted_outcome_model(), a later change).

• populate – When True (default), persist the predictions as <outcome>_hat columns on the scored population.

Returns:

The predictions DataFrame (one <outcome>_hat column per fitted outcome), or a (sample, target) tuple when on="both".

Raises:

ValueError – If no outcome model has been fit, if on="target" / "both" but no target is set, or if data is passed.

Examples

>>> import pandas as pd
>>> from balance.sample_frame import SampleFrame
>>> from balance.balance_frame import BalanceFrame
>>> resp = SampleFrame.from_frame(
...     pd.DataFrame({"id": ["1", "2", "3", "4"],
...                   "age": [25.0, 30.0, 35.0, 40.0],
...                   "happiness": [50.0, 55.0, 65.0, 80.0],
...                   "weight": [1.0, 1.0, 1.0, 1.0]}),
...     outcome_columns=["happiness"])
>>> tgt = SampleFrame.from_frame(
...     pd.DataFrame({"id": ["5", "6"], "age": [28.0, 38.0],
...                   "weight": [1.0, 1.0]}))
>>> bf = BalanceFrame(sample=resp, target=tgt)
>>> _ = bf.fit_outcome_model()
>>> preds = bf.predict_outcomes(on="target")
>>> list(preds.columns)
['happiness_hat']
>>> bf.outcomes_hat().df.columns.tolist()   # target now populated
['happiness_hat']

predict_proba(on: Literal['sample'], output: Literal['probability', 'link'] = 'probability', *, data: BalanceFrame | None = None) → Series

predict_proba(on: Literal['target'], output: Literal['probability', 'link'] = 'probability', *, data: BalanceFrame | None = None) → Series

predict_proba(on: Literal['both'] = 'both', output: Literal['probability', 'link'] = 'probability', *, data: BalanceFrame | None = None) → tuple[Series, Series]

Return IPW propensity scores.

Returns the propensity scores (predicted probabilities of being in the sample vs target) from the fitted IPW model. A target row with high propensity is well-represented in the sample; a low score indicates underrepresentation.

When data is provided, the stored model is applied to data’s covariates and fresh predictions are returned without caching. When data is None (default), stored/cached predictions for this object’s own data are returned (original behavior).

Parameters:

• on – Which population to predict on ("sample", "target", or "both").

• output – Output scale. "probability" returns class-1 propensity probabilities. "link" returns logit-transformed values.

• data – An optional BalanceFrame whose covariates are scored using this object’s stored model. Must have matching covariate column names. The data BalanceFrame needs a target for on="target" or on="both".

Returns:

A prediction Series, or a tuple of two Series when on="both".

Raises:

ValueError – If the object is not IPW-adjusted, if target is missing for on in {"target", "both"}, if recomputation of sample-side predictions is required but no target is available, if on is invalid, or if data has mismatched covariate columns.

Notes

When data is None and stored fit-time predictions are stale for the current rows, this method may recompute and cache refreshed probabilities/links. When data is provided, no caching occurs.

Examples

>>> import pandas as pd
>>> from balance.sample_frame import SampleFrame
>>> from balance.balance_frame import BalanceFrame
>>> resp = SampleFrame.from_frame(
...     pd.DataFrame({"id": [1, 2], "x": [0.0, 1.0], "weight": [1.0, 1.0]}))
>>> tgt = SampleFrame.from_frame(
...     pd.DataFrame({"id": [3, 4], "x": [0.2, 0.8], "weight": [1.0, 1.0]}))
>>> adjusted = BalanceFrame(sample=resp, target=tgt).fit(method="ipw")
>>> p = adjusted.predict_proba(on="target", output="probability")
>>> int(p.shape[0])
2

predict_weights(*, data: BalanceFrame | None = None) → Series

Predict responder weights from the fitted model’s artifacts.

Reconstructs adjusted survey weights from stored fit-time artifacts (propensity links, design weights, class balancing, trimming parameters). On the fitted object itself, the result is numerically equivalent to self.weights().df (within floating-point tolerance) and serves as a validation that the stored artifacts are sufficient to reproduce the adjustment.

When data is provided, computes weights for data’s sample using the stored model, without caching. This is the one-liner alternative to the set_fitted_model workflow:

fitted.predict_weights(data=holdout_bf)

When data is None (default), uses this object’s own data (original behavior).

Dispatches by the adjustment method stored in the model dict:

• IPW: uses stored fit-time metadata (links, class balancing, trimming, and design weights) to reproduce fitted responder weights.

• CBPS: rebuilds the CBPS scoring artifacts from stored metadata and supports both in-place and data=... holdout scoring.

• Poststratify: replays fitted cell-ratio metadata in-place and supports data=... transfer scoring. Models fitted with transformations='default' or direct data-dependent helpers such as balance.utils.data_transformation.quantize() / balance.utils.data_transformation.fct_lump() are rejected for transfer; pass deterministic transformations explicitly or re-fit poststratify on the scoring data.

• Rake: replays fitted cell-ratio artifacts in-place and supports data=... transfer scoring. See balance.weighting_methods.rake.rake() Notes for validity constraints and interpretation.

• Other methods: not yet supported — will raise with guidance.

Parameters:

data – An optional BalanceFrame whose sample covariates are scored using this object’s stored model. Must have matching covariate column names and a target set. Supported for IPW, CBPS, rake, and poststratify.

Returns:

A Series of predicted responder weights.

Raises:

ValueError – If no fitted model is available, if the method is unsupported, if required target data is missing, or if data has mismatched covariate columns.

Examples

>>> import pandas as pd
>>> from balance.sample_frame import SampleFrame
>>> from balance.balance_frame import BalanceFrame
>>> resp = SampleFrame.from_frame(
...     pd.DataFrame({"id": [1, 2], "x": [0.0, 1.0], "weight": [1.0, 1.0]}))
>>> tgt = SampleFrame.from_frame(
...     pd.DataFrame({"id": [3, 4], "x": [0.2, 0.8], "weight": [1.0, 1.0]}))
>>> adjusted = BalanceFrame(sample=resp, target=tgt).fit(method="ipw")
>>> w = adjusted.predict_weights()
>>> int(w.shape[0])
2

property responders: SampleFrame

Alias for _sf_sample (backward compat, will be removed).

set_as_pre_adjust(*, inplace: bool = False) → Self

Set the current responder state as the new pre-adjust baseline.

This “locks in” the current responder weights (which may already be adjusted and/or trimmed) as the baseline for future diagnostics and subsequent adjustments.

Parameters:

inplace – If True, mutate this object and return it. If False (default), return a new object with a deep-copied responder frame and reset baseline.

Returns:

BalanceFrame with _sf_sample_pre_adjust reset to the current responder SampleFrame state. In copy mode (inplace=False), only the responder frame is deep-copied and used to construct a new object (the full _links graph is not deep-copied). In inplace mode, the baseline is set to the existing responder frame object so baseline/current share identity, matching unadjusted-object semantics elsewhere in the API. Any current adjustment model is cleared because the object is no longer considered adjusted after this operation.

Examples

>>> import pandas as pd
>>> from balance.sample_frame import SampleFrame
>>> from balance.balance_frame import BalanceFrame
>>> resp = SampleFrame.from_frame(
...     pd.DataFrame({"id": [1, 2], "x": [10.0, 20.0], "weight": [1.0, 1.0]}))
>>> tgt = SampleFrame.from_frame(
...     pd.DataFrame({"id": [3, 4], "x": [15.0, 25.0], "weight": [1.0, 1.0]}))
>>> adjusted = BalanceFrame(sample=resp, target=tgt).adjust(method="null")
>>> baseline_locked = adjusted.set_as_pre_adjust()  # copy mode
>>> baseline_locked.is_adjusted
False
>>> _ = adjusted.set_as_pre_adjust(inplace=True)  # inplace mode

set_fitted_model(fitted: BalanceFrame, *, inplace: bool = True) → Self

Apply a fitted model from another BalanceFrame, producing a fully adjusted result.

This enables fit-then-apply workflows: fit on one BalanceFrame (e.g., a 20k subset) and apply the fitted model to another BalanceFrame (e.g., the remaining 980k) with the same covariate schema. The returned object is fully adjusted (is_adjusted is True, model is set, summary() works with 3-way comparison).

Workflow (inplace=False — returns new adjusted object):

fitted = train_bf.fit(method="ipw")
scored = holdout_bf.set_fitted_model(fitted, inplace=False)
scored.summary()  # full diagnostics on holdout

Workflow (inplace=True, default — mutates self):

holdout_bf.set_fitted_model(fitted)
holdout_bf.summary()

Currently set_fitted_model applies fitted IPW models directly. Other methods may have fit artifacts but are routed through their dedicated predict_weights(data=...) workflows instead of this method-specific application path.

Parameters:

• fitted – A BalanceFrame already adjusted with a supported method. Its fitted model is used to compute holdout weights.

• inplace – If True (default), mutate this object and return self. If False, return a new BalanceFrame with computed weights, leaving self unchanged.

Returns:

A fully adjusted BalanceFrame with holdout weights applied. self when inplace=True, a new object when inplace=False.

Raises:

ValueError – If fitted has no stored model, if the model method is not yet supported, or if covariate column names differ between self and fitted.

Examples

>>> import pandas as pd
>>> from balance.sample_frame import SampleFrame
>>> train_resp = SampleFrame.from_frame(
...     pd.DataFrame({"id": [1, 2], "x": [0.0, 1.0], "weight": [1.0, 1.0]}))
>>> train_tgt = SampleFrame.from_frame(
...     pd.DataFrame({"id": [3, 4], "x": [0.2, 0.8], "weight": [1.0, 1.0]}))
>>> holdout_resp = SampleFrame.from_frame(
...     pd.DataFrame({"id": [5, 6], "x": [0.1, 0.9], "weight": [1.0, 1.0]}))
>>> holdout_tgt = SampleFrame.from_frame(
...     pd.DataFrame({"id": [7, 8], "x": [0.3, 0.7], "weight": [1.0, 1.0]}))
>>> train_bf = BalanceFrame(sample=train_resp, target=train_tgt)
>>> holdout_bf = BalanceFrame(sample=holdout_resp, target=holdout_tgt)
>>> fitted = train_bf.fit(method="ipw")
>>> scored = holdout_bf.set_fitted_model(fitted, inplace=False)
>>> scored.is_adjusted
True
>>> scored.model is not None
True

set_fitted_outcome_model(fitted: BalanceFrame | SampleFrame, *, inplace: bool = True) → Self

Apply a fitted outcome model from another frame (train/holdout transfer).

This is the outcome-modelling counterpart to set_fitted_model() (the IPW weighting-axis train/holdout transfer): fit an outcome model ĝ(X) ≈ E[Y|X] on one frame (the train frame) and apply the fitted model to a different frame (the holdout / scoring frame) that shares the same covariate schema. Nothing is re-fit — the already-fitted model (including its stored preprocessing) is copied onto self sharing the fitted estimator objects by identity (exactly like set_fitted_model()), so self.predict_outcomes(on="target") then replays the transferred model on self’s own target and self.outcomes_hat().mean() is the g-computation estimate μ̂_OM on the holdout target.

On a Sample, this BalanceFrame method is reached via the MRO; the model lands on self._sf_sample (the single source of truth read back through outcome_model), so it rides along adjust() and is preserved across set_target().

Parameters:

• fitted – Another BalanceFrame (or SampleFrame / Sample) that already carries a fitted outcome model (via fit_outcome_model()). Its covariate columns must match self’s.

• inplace – If True (default), mutate this object (store the model on the responder) and return self; if False, return a new copy with the model stored and leave self unchanged — matching set_fitted_model() / fit_outcome_model().

Returns:

The BalanceFrame with the transferred outcome model stored on its responder (self when inplace, else a new copy).

Raises:

ValueError – If fitted has no stored outcome model; if the covariate column names differ between self and fitted; or if the fitted model used a non-deterministic transformations (e.g. quantize / fct_lump) or na_action="drop" that cannot be replayed deterministically on self.

Examples

>>> import pandas as pd
>>> from balance.sample_frame import SampleFrame
>>> from balance.balance_frame import BalanceFrame
>>> train_resp = SampleFrame.from_frame(
...     pd.DataFrame({"id": ["1", "2", "3", "4"],
...                   "age": [25.0, 30.0, 35.0, 40.0],
...                   "happiness": [50.0, 55.0, 65.0, 80.0],
...                   "weight": [1.0, 1.0, 1.0, 1.0]}),
...     outcome_columns=["happiness"])
>>> holdout_resp = SampleFrame.from_frame(
...     pd.DataFrame({"id": ["5", "6", "7", "8"],
...                   "age": [26.0, 31.0, 36.0, 41.0],
...                   "happiness": [51.0, 56.0, 66.0, 81.0],
...                   "weight": [1.0, 1.0, 1.0, 1.0]}),
...     outcome_columns=["happiness"])
>>> holdout_tgt = SampleFrame.from_frame(
...     pd.DataFrame({"id": ["9", "10"], "age": [28.0, 38.0],
...                   "weight": [1.0, 1.0]}))
>>> train_bf = BalanceFrame(sample=train_resp)
>>> _ = train_bf.fit_outcome_model()
>>> holdout_bf = BalanceFrame(sample=holdout_resp, target=holdout_tgt)
>>> scored = holdout_bf.set_fitted_outcome_model(train_bf, inplace=False)
>>> _ = scored.predict_outcomes(on="target")
>>> scored.outcome_model["fit"]["happiness"] is train_bf.outcome_model["fit"]["happiness"]
True

set_target(target: BalanceFrame | SampleFrame, inplace: bool | None = None) → Self

Set or replace the target population.

When target is a BalanceFrame (or subclass such as Sample), a deep copy of self is returned with the target set (immutable pattern). When target is a raw SampleFrame, the behaviour depends on inplace: True mutates self, False returns a new BalanceFrame.

Parameters:

• target – The target population — a BalanceFrame/Sample or a SampleFrame.

• inplace – If True, mutates self (only valid for SampleFrame targets). If False, returns a new copy. Defaults to None which auto-selects: copy for BalanceFrame targets, inplace for SampleFrame targets.

Returns:

BalanceFrame with the new target set.

Raises:

TypeError / ValueError – If target is not a BalanceFrame or SampleFrame, or if they share no covariate columns.

Examples

>>> import pandas as pd
>>> from balance.sample_frame import SampleFrame
>>> from balance.balance_frame import BalanceFrame
>>> resp = SampleFrame.from_frame(
...     pd.DataFrame({"id": [1, 2], "x": [10.0, 20.0], "weight": [1.0, 1.0]}))
>>> tgt = SampleFrame.from_frame(
...     pd.DataFrame({"id": [3, 4], "x": [15.0, 25.0], "weight": [1.0, 1.0]}))
>>> bf = BalanceFrame(sample=resp)
>>> bf.set_target(tgt)
>>> bf.has_target()
True

set_unadjusted(second: BalanceFrame) → Self

Set the unadjusted link for comparative analysis.

Returns a deep copy with _sf_sample_pre_adjust pointing at second’s responder SampleFrame, and _links["unadjusted"] pointing at second.

Parameters:

second – A BalanceFrame (or subclass) whose responder data becomes the unadjusted baseline.

Returns:

A new BalanceFrame with the unadjusted link set.

Raises:

TypeError – If second is not a BalanceFrame.

set_weights(weights: Series | float | None, *, use_index: bool = False) → None

Set or replace the responder weights.

Delegates to the underlying SampleFrame’s set_weights.

When called on an unadjusted BalanceFrame (is_adjusted is False), _sf_sample and _sf_sample_pre_adjust share the same DataFrame, so the change is visible to both automatically — changing base weights is not an adjustment.

Warning

If this BalanceFrame has already been fitted (i.e., adjust() has been called), calling set_weights() changes the design weights but does not invalidate the stored fit artifacts (_adjustment_model). The link values in those artifacts were computed using the old weights, so predict_weights() will use new current_sample_weights with stale links, producing a mathematical inconsistency. Users should re-fit (call adjust() again) after changing weights on an already-fitted BalanceFrame.

Parameters:

• weights – New weights. A Series, a scalar (broadcast to all rows), or None (sets all to 1.0).

• use_index – If True, align weights by index instead of requiring matching length. See SampleFrame.set_weights().

summary() → str

Consolidated summary of covariate balance, weight health, and outcomes.

Produces a multi-line summary combining covariate ASMD / KLD diagnostics, weight design effect, and outcome means. Delegates to _build_summary() after computing the necessary intermediate values.

When no target is set, returns a minimal summary with weight diagnostics and outcome means only.

Returns:

A human-readable multi-line summary string.

Return type:

str

Examples

>>> import pandas as pd
>>> from balance.sample_frame import SampleFrame
>>> from balance.balance_frame import BalanceFrame
>>> resp = SampleFrame.from_frame(
...     pd.DataFrame({"id": [1, 2, 3, 4], "x": [0, 1, 1, 0],
...                   "weight": [1.0, 2.0, 1.0, 1.0]}))
>>> tgt = SampleFrame.from_frame(
...     pd.DataFrame({"id": [5, 6, 7, 8], "x": [0, 0, 1, 1],
...                   "weight": [1.0, 1.0, 1.0, 1.0]}))
>>> bf = BalanceFrame(sample=resp, target=tgt)
>>> adjusted = bf.adjust(method="null")
>>> "Covariate diagnostics:" in adjusted.summary()
True

property target: SampleFrame | None

Alias for _sf_target (backward compat, will be removed).

to_csv(path_or_buf: str | Path | IO | None = None, **kwargs: Any) → str | None

Write the combined DataFrame to CSV.

Writes the output of df (responder + target + unadjusted rows with a "source" column) to a CSV file or string. Delegates to to_csv_with_defaults().

Parameters:

• path_or_buf – Destination. If None, returns the CSV as a string.

• **kwargs – Additional keyword arguments passed to pd.DataFrame.to_csv().

Returns:

CSV string if path_or_buf is None, else None.

Return type:

str or None

Examples

>>> import pandas as pd
>>> from balance.sample_frame import SampleFrame
>>> from balance.balance_frame import BalanceFrame
>>> resp = SampleFrame.from_frame(
...     pd.DataFrame({"id": [1, 2], "x": [10.0, 20.0], "weight": [1.0, 1.0]}))
>>> tgt = SampleFrame.from_frame(
...     pd.DataFrame({"id": [3, 4], "x": [15.0, 25.0], "weight": [1.0, 1.0]}))
>>> bf = BalanceFrame(sample=resp, target=tgt)
>>> "id" in bf.to_csv()
True

to_download(tempdir: str | None = None) → Any

Create a downloadable file link of the combined DataFrame.

Writes df to a temporary CSV file and returns an IPython FileLink for interactive download.

Parameters:

tempdir – Directory for the temp file. If None, uses tempfile.gettempdir().

Returns:

An IPython file link for downloading the CSV.

Return type:

FileLink

Examples

>>> import tempfile
>>> import pandas as pd
>>> from balance.sample_frame import SampleFrame
>>> from balance.balance_frame import BalanceFrame
>>> resp = SampleFrame.from_frame(
...     pd.DataFrame({"id": [1, 2], "x": [10.0, 20.0], "weight": [1.0, 1.0]}))
>>> tgt = SampleFrame.from_frame(
...     pd.DataFrame({"id": [3, 4], "x": [15.0, 25.0], "weight": [1.0, 1.0]}))
>>> bf = BalanceFrame(sample=resp, target=tgt)
>>> link = bf.to_download(tempdir=tempfile.gettempdir())

to_sample() → Any

Convert this BalanceFrame back to a Sample.

Reconstructs a Sample with the responder data and target set. If this BalanceFrame is adjusted, the returned Sample will also be adjusted — is_adjusted() returns True, has_target() returns True, and the original (unadjusted) weights are preserved via the "unadjusted" link.

Returns:

A Sample mirroring this BalanceFrame’s data, target,

and adjustment state.

Return type:

Sample

Examples

>>> import pandas as pd
>>> from balance.sample_frame import SampleFrame
>>> from balance.balance_frame import BalanceFrame
>>> resp = SampleFrame.from_frame(
...     pd.DataFrame({"id": [1, 2, 3], "x": [10.0, 20.0, 30.0],
...                   "weight": [1.0, 1.0, 1.0]}))
>>> tgt = SampleFrame.from_frame(
...     pd.DataFrame({"id": [4, 5, 6], "x": [15.0, 25.0, 35.0],
...                   "weight": [1.0, 1.0, 1.0]}))
>>> bf = BalanceFrame(sample=resp, target=tgt)
>>> s = bf.to_sample()
>>> s.has_target()
True

trim(ratio: float | int | None = None, percentile: float | tuple[float, float] | None = None, keep_sum_of_weights: bool = True, target_sum_weights: float | int | np.floating | None = None, *, inplace: bool = False) → Self

Trim extreme weights using mean-ratio clipping or percentile winsorization.

Delegates to SampleFrame.trim() for computation and weight history tracking, then wraps the result in a new BalanceFrame (preserving target, pre-adjust baseline, and links).

Parameters:

• ratio – Mean-ratio upper bound. Mutually exclusive with percentile.

• percentile – Percentile(s) for winsorization. Mutually exclusive with ratio.

• keep_sum_of_weights – Whether to rescale after trimming to preserve the original sum of weights.

• target_sum_weights – If provided, rescale trimmed weights so their sum equals this target.

• inplace – If True, mutate this BalanceFrame’s weights and return it. If False (default), return a new BalanceFrame.

Returns:

The BalanceFrame with trimmed weights (self if inplace, else a new instance).

property unadjusted: SampleFrame | None

Alias for _sf_sample_pre_adjust if adjusted, else None (backward compat).

property weight_series: Series | None

The active weight as a Series, delegated to _sf_sample.

weights() → Any

Return a BalanceDFWeights for the responders.

The returned object carries linked target (and unadjusted, if adjusted) views for comparative weight analysis.

Returns:

Weight view with linked sources.

Return type:

BalanceDFWeights

Examples

>>> import pandas as pd
>>> from balance.sample_frame import SampleFrame
>>> from balance.balance_frame import BalanceFrame
>>> resp = SampleFrame.from_frame(
...     pd.DataFrame({"id": [1, 2], "x": [10.0, 20.0], "weight": [1.0, 2.0]}))
>>> tgt = SampleFrame.from_frame(
...     pd.DataFrame({"id": [3, 4], "x": [15.0, 25.0], "weight": [1.0, 1.0]}))
>>> bf = BalanceFrame(sample=resp, target=tgt)
>>> bf.weights().df.columns.tolist()
['weight']