balance.sample_frame

SampleFrame: an explicit-role DataFrame container for the Balance library.

Stores covariates, weights, outcomes, outcomes_hat, and ignored columns with explicit role metadata, replacing the inference-by-exclusion pattern used in the legacy Sample class.

class balance.sample_frame.SampleFrame

A DataFrame container with explicit column-role metadata.

SampleFrame stores data as a single internal pd.DataFrame but with explicit metadata tracking which columns belong to which role: covars (X), weights (W), outcomes (Y), outcomes_hat (Y_hat), ignored.

Must be constructed via SampleFrame.from_frame() or SampleFrame.from_sample().

Mutability:

SampleFrame is mostly-immutable at the data level. The underlying DataFrame and column-role assignments are set at construction time and are not replaced afterwards. All data-access properties (e.g. df_covars, df_weights) return copies, so callers cannot mutate internal state through the returned objects.

Controlled mutation points (methods that intentionally modify the instance in-place):

• set_active_weight() — changes which weight column is active.

• add_weight_column() — appends a new weight column to the frame.

• set_weight_metadata() — updates weight provenance metadata.

These mutations are intentional and expected as part of normal usage (e.g. after calling BalanceFrame.adjust()). Outside of these methods the object behaves as immutable.

add_outcomes_hat_column(name: str, values: Series, metadata: dict[str, Any] | None = None) → None

Attach a predicted-outcome (Y_hat) column to the SampleFrame.

Mirrors add_weight_column(): the column is appended to the internal DataFrame in place, registered under the outcomes_hat role, and optionally associated with provenance metadata (stored in _prediction_metadata). Because the column takes the outcomes_hat role, it is not treated as a covariate.

By convention, predicted-outcome columns are named "<outcome>_hat" (e.g. "happiness_hat"). A logging.warning is emitted — but no error raised — when name does not end in _hat.

Parameters:

• name (str) – Name for the new outcomes_hat column.

• values (pd.Series) – Predicted values. Must match the DataFrame length, unless it is a shorter pd.Series — in which case values are aligned by index and missing rows are filled with NaN.

• metadata (dict, optional) – Provenance metadata for the new column (e.g. the fitting method or learner).

Raises:

ValueError – If name is already an outcomes_hat column, if name already exists in the DataFrame, or if values is longer than the DataFrame (or is a non-Series with a different length).

Examples

>>> import pandas as pd
>>> from balance.sample_frame import SampleFrame
>>> df = pd.DataFrame({"id": ["1", "2", "3", "4"],
...                    "age": [25, 30, 35, 40],
...                    "weight": [1.0, 1.0, 1.0, 1.0]})
>>> sf = SampleFrame.from_frame(df)
>>> sf.add_outcomes_hat_column("happiness_hat",
...                            pd.Series([52., 58., 68., 79.]))
>>> sf.df_outcomes_hat["happiness_hat"].tolist()
[52.0, 58.0, 68.0, 79.0]
>>> list(sf.df_covars.columns)
['age']

add_weight_column(name: str, values: Series, metadata: dict[str, Any] | None = None) → None

Add a new weight column to the SampleFrame.

The column is appended to the internal DataFrame and registered as a weight column. Optionally associates provenance metadata.

Parameters:

• name (str) – Name for the new weight column.

• values (pd.Series) – Weight values. Must match the DataFrame length, unless it is a shorter pd.Series — in which case values are aligned by index and missing rows are filled with NaN (this supports adjustment functions that drop rows internally, e.g., na_action="drop"). Note: this column is a history column, not the active weight — the active weight is set separately via set_weights().

• metadata (dict, optional) – Provenance metadata for the new column.

Raises:

ValueError – If name is already a registered weight column, if name already exists in the DataFrame as a non-weight column, or if values is longer than the DataFrame or is a non-Series with a different length.

Examples

>>> import pandas as pd
>>> from balance.sample_frame import SampleFrame
>>> df = pd.DataFrame({"id": ["1", "2"], "x": [10, 20],
...                    "weight": [1.0, 2.0]})
>>> sf = SampleFrame.from_frame(df)
>>> sf.add_weight_column("w_adj", pd.Series([1.5, 1.5]),
...                      metadata={"method": "rake"})
>>> sf._column_roles["weights"]
['weight', 'w_adj']

property covar_columns: list[str]

Names of the covariate columns.

Returns a copy so that callers cannot accidentally mutate the internal column-role registry.

Returns:

Covariate column names.

Return type:

list[str]

Examples

>>> import pandas as pd
>>> from balance.sample_frame import SampleFrame
>>> df = pd.DataFrame({"id": ["1", "2"], "age": [25, 30],
...                    "income": [50000, 60000], "weight": [1.0, 1.0]})
>>> sf = SampleFrame.from_frame(df)
>>> sf.covar_columns
['age', 'income']

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

Return a BalanceDFCovars for this SampleFrame.

Creates a covariate analysis view backed by this SampleFrame, inheriting any linked sources set via _links.

Parameters:

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

Returns:

Covariate view backed by this SampleFrame.

Return type:

BalanceDFCovars

Examples

>>> import pandas as pd
>>> from balance.sample_frame import SampleFrame
>>> sf = SampleFrame.from_frame(
...     pd.DataFrame({"id": [1, 2], "x": [10.0, 20.0],
...                   "weight": [1.0, 1.0]}))
>>> sf.covars().df.columns.tolist()
['x']

property df: DataFrame

Full DataFrame reconstruction.

property df_covars: DataFrame

Covariate columns as a DataFrame.

Returns a copy so that callers cannot accidentally mutate the internal data.

Returns:

A copy of the covariate columns.

Return type:

pd.DataFrame

Examples

>>> import pandas as pd
>>> from balance.sample_frame import SampleFrame
>>> df = pd.DataFrame({"id": ["1", "2"], "age": [25, 30],
...                    "income": [50000, 60000], "weight": [1.0, 1.0]})
>>> sf = SampleFrame.from_frame(df)
>>> covars = sf.df_covars
>>> covars["age"] = [999, 999]
>>> list(sf.df_covars["age"])  # internal data unchanged
[25.0, 30.0]

property df_ignored: DataFrame | None

Ignored columns, or None.

Returns a copy so that callers cannot accidentally mutate the internal data.

Returns:

A copy of ignored columns, or

None if no ignored columns are registered.

Return type:

pd.DataFrame | None

Examples

>>> import pandas as pd
>>> from balance.sample_frame import SampleFrame
>>> df = pd.DataFrame({"id": ["1", "2"], "x": [10, 20],
...                    "weight": [1.0, 1.0], "region": ["US", "UK"]})
>>> sf = SampleFrame.from_frame(df, ignored_columns=["region"])
>>> m = sf.df_ignored
>>> m["region"] = ["XX", "XX"]
>>> list(sf.df_ignored["region"])  # internal data unchanged
['US', 'UK']

property df_outcomes: DataFrame | None

Outcome columns, or None if no outcomes.

Returns a copy so that callers cannot accidentally mutate the internal data.

Returns:

A copy of outcome columns, or None if

no outcome columns are registered.

Return type:

pd.DataFrame | None

Examples

>>> import pandas as pd
>>> from balance.sample_frame import SampleFrame
>>> df = pd.DataFrame({"id": ["1", "2"], "x": [10, 20],
...                    "weight": [1.0, 1.0], "y": [5, 6]})
>>> sf = SampleFrame.from_frame(df, outcome_columns=["y"])
>>> out = sf.df_outcomes
>>> out["y"] = [999, 999]
>>> list(sf.df_outcomes["y"])  # internal data unchanged
[5.0, 6.0]

property df_outcomes_hat: DataFrame | None

Predicted-outcome (Y_hat) columns, or None if none.

Mirrors df_outcomes for the outcomes_hat role. Returns a copy so that callers cannot accidentally mutate the internal data.

Returns:

A copy of the outcomes_hat columns, or None

if no outcomes_hat columns are registered.

Return type:

pd.DataFrame | None

Examples

>>> import pandas as pd
>>> from balance.sample_frame import SampleFrame
>>> df = pd.DataFrame({"id": ["1", "2", "3", "4"],
...                    "age": [25, 30, 35, 40],
...                    "weight": [1.0, 1.0, 1.0, 1.0]})
>>> sf = SampleFrame.from_frame(df)
>>> sf.add_outcomes_hat_column("happiness_hat",
...                            pd.Series([52., 58., 68., 79.]))
>>> sf.df_outcomes_hat["happiness_hat"].tolist()
[52.0, 58.0, 68.0, 79.0]

property df_weights: DataFrame

Active weight column as a single-column DataFrame.

Returns a copy so that callers cannot accidentally mutate the internal data.

Returns:

A copy of the active weight column, or an empty

DataFrame if no active weight is set.

Return type:

pd.DataFrame

Examples

>>> import pandas as pd
>>> from balance.sample_frame import SampleFrame
>>> df = pd.DataFrame({"id": ["1", "2"], "x": [10, 20],
...                    "weight": [1.0, 2.0]})
>>> sf = SampleFrame.from_frame(df)
>>> w = sf.df_weights
>>> w["weight"] = [999.0, 999.0]
>>> list(sf.df_weights["weight"])  # internal data unchanged
[1.0, 2.0]

fit_outcome_model(*, model: OutcomeLearner = 'auto', outcome_columns: list[str] | str | None = None, variables: list[str] | str | None = None, formula: str | list[str] | None = None, transformations: str | dict[str, Any] | None = None, na_action: str = 'add_indicator', use_model_matrix: bool | str = 'auto', weighted: bool = False, calibrate: bool = False, inplace: bool = True) → Self

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

This is the outcome-modelling counterpart to fit() (which fits the weighting model). It fits a regressor (continuous outcome) or classifier (binary outcome) per resolved outcome column on this frame’s covariates and observed outcome(s), and stores the resulting model dict (fitted estimators + preprocessing) on outcome_model. Like sklearn’s fit(), it does not produce predictions — call predict_outcomes() (or fit_predict_outcomes()) to write the <outcome>_hat columns.

The fit is unweighted by default (weighted=False): the outcome model estimates E[Y|X] and is usually best left unbiased by the design weights. Pass weighted=True to fit with the frame’s currently-applied weight (weighted least squares / weighted boosting), aligned to the covariate index before it is passed through. Rows whose outcome Y is missing (NaN) are dropped before fitting — the covariates and weights are aligned to the retained rows, so the fit uses only complete outcome observations. Re-fitting drops any <outcome>_hat columns left by a previous predict_outcomes() so a stale prediction cannot linger against a newly-fit model.

Parameters:

• model – "auto" (a HistGradientBoosting regressor/classifier chosen by outcome type), a single sklearn estimator (cloned per outcome column — one estimator type only when the outcomes are mixed), a {"_discrete": clf, "_continuous": reg} type map, or a {outcome_column: estimator} column map.

• outcome_columns – Outcome column name(s) to model. Defaults to all of this frame’s outcome_columns; raises if none exist.

• variables – Covariate column name(s) to use as the model inputs X. Defaults to None (all of this frame’s covariates); pass a subset (validated against the covariate role) to restrict the model to those columns.

• formula – Optional patsy formula(s) forwarded to the one-hot path.

• transformations – Reserved for parity with IPW; must be None (the replay-safe default) — a non-None value raises.

• na_action – Missing-value handling for the design matrix; only "add_indicator" (default) is supported ("drop" raises).

• use_model_matrix – "auto" (default) picks the native-categorical path for tree/boosting learners on scikit-learn >= 1.4 and the one-hot + scaler path otherwise; pass True/False to force.

• weighted – When True, fit with the frame’s currently-applied weight (weighted least squares / weighted boosting). Defaults to False — the outcome model predicts E[Y|X] and is usually best left unweighted. A weighted fit whose estimator’s fit does not accept sample_weight raises TypeError.

• calibrate – When True, wrap each classifier in CalibratedClassifierCV (binary outcomes only).

• inplace – If True (default), mutate this frame (store the model) and return self; if False, return a new copy with the model stored and leave self untouched — mirroring fit().

Returns:

The frame with the fitted outcome model stored (self when inplace, else a new copy).

Raises:

• ValueError – If there are no outcome columns to fit (and none are passed), if a requested outcome/variable column is not registered, or for the underlying fit_outcome_model errors (na_action="drop", non-None transformations, a model column map missing an outcome column, a single estimator for mixed-type outcomes, etc.).

• TypeError – If weighted=True but the resolved estimator’s fit does not accept sample_weight.

Examples

>>> import pandas as pd
>>> from balance.sample_frame import SampleFrame
>>> df = 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]})
>>> sf = SampleFrame.from_frame(df, outcome_columns=["happiness"])
>>> sf.fit_outcome_model()
<balance.sample_frame.SampleFrame object at ...>
>>> sf.outcome_model["method"]
'outcome_model'
>>> sf.df_outcomes_hat is None  # fit does NOT persist Ŷ
True

fit_predict_outcomes(*, populate: bool = True, **fit_kwargs: Any) → DataFrame

Fit an outcome model then predict on this frame, in one call.

Convenience wrapper mirroring sklearn’s fit_predict: it calls fit_outcome_model() (in place) with **fit_kwargs and then predict_outcomes() on this frame, persisting <outcome>_hat when populate=True (the default).

Parameters:

• populate – When True (default), persist the predictions onto this frame as <outcome>_hat columns.

• **fit_kwargs – Keyword arguments forwarded to fit_outcome_model() (e.g. model, outcome_columns, variables, weighted). inplace is always True here.

Returns:

The predictions, one "<outcome>_hat" column per

fitted outcome.

Return type:

pd.DataFrame

Examples

>>> import pandas as pd
>>> from balance.sample_frame import SampleFrame
>>> df = 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]})
>>> sf = SampleFrame.from_frame(df, outcome_columns=["happiness"])
>>> preds = sf.fit_predict_outcomes()
>>> list(preds.columns)
['happiness_hat']
>>> sf.outcome_model["method"]
'outcome_model'

classmethod from_frame(df: DataFrame, id_column: str | None = None, covar_columns: list[str] | None = None, weight_column: str | None = None, outcome_columns: list[str] | tuple[str, ...] | str | None = None, outcomes_hat_columns: list[str] | tuple[str, ...] | str | None = None, ignored_columns: list[str] | tuple[str, ...] | str | None = None, check_id_uniqueness: bool = True, standardize_types: bool = True, use_deepcopy: bool = True, id_column_candidates: list[str] | tuple[str, ...] | str | None = None) → SampleFrame

Create a SampleFrame from a pandas DataFrame with auto-detection.

Infers id, weight, and covariate columns from column names when not explicitly provided. Validates the data (e.g., unique IDs, non-negative weights) and standardizes dtypes (Int64 -> float64, pd.NA -> np.nan).

Parameters:

• df (pd.DataFrame) – The input DataFrame containing survey or observational data.

• id_column (str, optional) – Name of the column to use as row identifier. If None, guessed from common names ("id", "ID", etc.).

• covar_columns (list of str, optional) – Explicit list of covariate column names. If None, inferred by exclusion (all columns minus id, weight, outcome, outcomes_hat, and ignored columns).

• weight_column (str, optional) – Name of the column containing sampling weights. If None, guesses "weight"/"weights" or creates one filled with 1.0.

• outcome_columns (list of str or str, optional) – Column names to treat as outcome variables.

• outcomes_hat_columns (list of str or str, optional) – Column names to treat as predicted-outcome (Y_hat) variables.

• ignored_columns (list of str or str, optional) – Column names to ignore (excluded from covariates).

• check_id_uniqueness (bool) – Whether to verify id uniqueness. Defaults to True.

• standardize_types (bool) – Whether to standardize dtypes. Defaults to True.

• use_deepcopy (bool) – Whether to deep-copy the input DataFrame. Defaults to True.

• id_column_candidates (list of str, optional) – Candidate id column names to try when id_column is None.

Returns:

A validated SampleFrame with standardized dtypes.

Return type:

SampleFrame

Raises:

ValueError – If the id column contains nulls or duplicates, if the weight column contains nulls or negative values, or if specified outcome/outcomes_hat/ignore columns are missing from the DataFrame.

Examples

>>> import pandas as pd
>>> df = pd.DataFrame({"id": [1, 2, 3], "weight": [1.0, 2.0, 1.5],
...                    "age": [25, 30, 35], "income": [50000, 60000, 70000]})
>>> sf = SampleFrame.from_frame(df)
>>> list(sf.df_covars.columns)
['age', 'income']

classmethod from_sample(sample: Any) → SampleFrame

Convert a Sample to a SampleFrame.

Preserves the Sample’s tabular data and column role assignments: id column, weight column, outcome columns, and ignored columns. Covariate columns are inferred by exclusion, matching the Sample’s own logic.

The internal DataFrame is deep-copied so that the resulting SampleFrame is fully independent of the original Sample.

Warning

Data not preserved in the conversion

The following Sample attributes are not carried over:

• _adjustment_model — the fitted model dictionary stored by adjust().

• _links — references to target, unadjusted, and other linked Samples (used by BalanceDF for comparative display).

• Column ordering may differ after a round-trip (Sample → SampleFrame → Sample), since SampleFrame stores columns grouped by role rather than preserving the original DataFrame column order.

Parameters:

sample – A Sample instance.

Returns:

A new SampleFrame mirroring the Sample’s data and

column roles.

Return type:

SampleFrame

Raises:

TypeError – If sample is not a Sample instance.

Examples

>>> import pandas as pd
>>> from balance.sample_class import Sample
>>> from balance.sample_frame import SampleFrame
>>> s = Sample.from_frame(
...     pd.DataFrame({"id": [1, 2], "x": [10.0, 20.0], "weight": [1.0, 2.0]}))
>>> sf = SampleFrame.from_sample(s)
>>> list(sf.df_covars.columns)
['x']

property id_column: str

Name of the ID column.

Note

In balance 0.20.0, id_column was changed from returning ID data (pd.Series) to returning the column name (str), for consistency with weight_column. If you need ID data, use id_series.

Returns:

The ID column name.

Return type:

str

Examples

>>> import pandas as pd
>>> from balance.sample_frame import SampleFrame
>>> df = pd.DataFrame({"id": ["1", "2"], "x": [10, 20],
...                    "weight": [1.0, 1.0]})
>>> sf = SampleFrame.from_frame(df)
>>> sf.id_column
'id'

property id_series: Series

The ID column as a Series.

Returns a copy so that callers cannot accidentally mutate the internal data.

Returns:

A copy of the ID column.

Return type:

pd.Series

Examples

>>> import pandas as pd
>>> from balance.sample_frame import SampleFrame
>>> df = pd.DataFrame({"id": ["1", "2"], "x": [10, 20],
...                    "weight": [1.0, 1.0]})
>>> sf = SampleFrame.from_frame(df)
>>> ids = sf.id_series
>>> ids.iloc[0] = "MUTATED"
>>> sf.id_series.iloc[0]  # internal data unchanged
'1'

property ignored_columns: list[str]

Names of the ignored columns.

Returns a copy so that callers cannot accidentally mutate the internal column-role registry.

Returns:

Ignored column names (empty list if none).

Return type:

list[str]

Examples

>>> import pandas as pd
>>> from balance.sample_frame import SampleFrame
>>> df = pd.DataFrame({"id": ["1", "2"], "x": [10, 20],
...                    "weight": [1.0, 1.0], "region": ["US", "UK"]})
>>> sf = SampleFrame.from_frame(df, ignored_columns=["region"])
>>> sf.ignored_columns
['region']

property outcome_columns: list[str]

Names of the outcome columns.

Returns a copy so that callers cannot accidentally mutate the internal column-role registry.

Returns:

Outcome column names (empty list if none).

Return type:

list[str]

Examples

>>> import pandas as pd
>>> from balance.sample_frame import SampleFrame
>>> df = pd.DataFrame({"id": ["1", "2"], "x": [10, 20],
...                    "weight": [1.0, 1.0], "y": [5, 6]})
>>> sf = SampleFrame.from_frame(df, outcome_columns=["y"])
>>> sf.outcome_columns
['y']

property outcome_model: dict[str, Any] | None

The fitted outcome-model dictionary, or None if not fit.

Mirrors model (which holds the weighting model) for the outcome-modelling axis. The dict is produced by fit_outcome_model() and consumed by predict_outcomes(); its keys are documented on balance.outcome_models.fit_outcome_model() ("method", "fit", "X_matrix_columns", "perf", …).

Returns:

The stored _outcome_model dict, or

None when no outcome model has been fit.

Return type:

dict[str, Any] | None

Examples

>>> import pandas as pd
>>> from balance.sample_frame import SampleFrame
>>> df = 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]})
>>> sf = SampleFrame.from_frame(df, outcome_columns=["happiness"])
>>> sf.outcome_model is None
True
>>> sf.fit_outcome_model()
<balance.sample_frame.SampleFrame object at ...>
>>> sf.outcome_model["method"]
'outcome_model'

outcomes() → Any | None

Return a BalanceDFOutcomes, or None.

Returns None if this SampleFrame has no outcome columns.

Returns:

Outcome view backed by this SampleFrame,

or None if no outcomes are defined.

Return type:

BalanceDFOutcomes or None

Examples

>>> import pandas as pd
>>> from balance.sample_frame import SampleFrame
>>> sf = 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"])
>>> sf.outcomes().df.columns.tolist()
['y']

outcomes_hat() → Any | None

Return a BalanceDFOutcomesHat, or None.

Returns None if this SampleFrame has no predicted-outcome (outcomes_hat / Y_hat) columns. When present, the returned view exposes the weighted mean / CI machinery over the Y_hat columns (one column per predicted outcome), mirroring outcomes().

Returns:

Predicted-outcome view backed by this

SampleFrame, 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
>>> df = pd.DataFrame({"id": ["1", "2", "3", "4"],
...                    "age": [25, 30, 35, 40],
...                    "weight": [1.0, 1.0, 1.0, 1.0]})
>>> sf = SampleFrame.from_frame(df)
>>> sf.add_outcomes_hat_column("happiness_hat",
...                            pd.Series([52., 58., 68., 79.]))
>>> sf.outcomes_hat().df.columns.tolist()
['happiness_hat']
>>> SampleFrame.from_frame(df).outcomes_hat() is None
True

property outcomes_hat_columns: list[str]

Names of the predicted-outcome (Y_hat) columns.

Returns a copy so that callers cannot accidentally mutate the internal column-role registry.

Returns:

outcomes_hat column names (empty list if none).

Return type:

list[str]

Examples

>>> import pandas as pd
>>> from balance.sample_frame import SampleFrame
>>> df = pd.DataFrame({"id": ["1", "2"], "x": [10, 20],
...                    "weight": [1.0, 1.0], "p_y": [0.3, 0.7]})
>>> sf = SampleFrame.from_frame(df, outcomes_hat_columns=["p_y"])
>>> sf.outcomes_hat_columns
['p_y']

predict_outcomes(*, data: SampleFrame | None = None, populate: bool | None = None) → DataFrame

Predict outcomes_hat from the stored model, optionally persisting.

Requires a model fit by fit_outcome_model(). Predictions are produced on this frame’s covariates (or on data’s covariates when a SampleFrame is passed via data=) by replaying the model’s stored preprocessing. The returned DataFrame has one column per fitted outcome, named "<outcome>_hat".

When populate is True, the predicted columns are written onto this frame via add_outcomes_hat_column() (a same-named Ŷ column is dropped and re-added).

Parameters:

• data – Optional SampleFrame whose covariates to score. Defaults to None, meaning predict on this frame’s own covariates. When scoring a different frame, prefer populate=False: persisting predictions row-indexed by data onto this frame would row-misalign (they are reindexed to this frame, NaN-padding or dropping non-matching rows).

• populate – Whether to persist the predictions onto this frame as <outcome>_hat columns. Defaults to True when scoring this frame (data=None) and False when data= is given, since a different frame’s row index would misalign onto this one.

Returns:

The predictions, one "<outcome>_hat" column per

fitted outcome, indexed by the scored covariate rows.

Return type:

pd.DataFrame

Raises:

ValueError – If no outcome model has been fit on this frame.

Examples

>>> import pandas as pd
>>> from balance.sample_frame import SampleFrame
>>> df = 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]})
>>> sf = SampleFrame.from_frame(df, outcome_columns=["happiness"])
>>> _ = sf.fit_outcome_model()
>>> preds = sf.predict_outcomes()
>>> list(preds.columns)
['happiness_hat']
>>> "happiness_hat" in sf.outcomes_hat_columns
True

rename_weight_column(old_name: str, new_name: str) → None

Rename a weight column in-place.

Renames the column in the DataFrame, updates the column roles list, active weight pointer, and weight metadata.

Parameters:

• old_name – Current name of the weight column.

• new_name – New name for the weight column.

Raises:

ValueError – If old_name is not a registered weight column, or if new_name already exists in the DataFrame.

set_active_weight(column_name: str) → None

Set which weight column is the active one.

The active weight column is the one returned by df_weights.

Parameters:

column_name (str) – Must be a registered weight column.

Raises:

ValueError – If column_name is not a weight column.

Examples

>>> import pandas as pd
>>> from balance.sample_frame import SampleFrame
>>> sf = SampleFrame._create(
...     df=pd.DataFrame({"id": [1], "x": [10], "w1": [1.0], "w2": [2.0]}),
...     id_column="id", covar_columns=["x"],
...     weight_columns=["w1", "w2"])
>>> sf.set_active_weight("w2")
>>> list(sf.df_weights.columns)
['w2']

set_weight_metadata(column: str, metadata: dict[str, Any]) → None

Store provenance metadata for a weight column.

Metadata is an arbitrary dict that can track adjustment method, hyperparameters, timestamps, or any other provenance information relevant to how the weight column was computed.

Parameters:

• column (str) – Name of the weight column.

• metadata (dict) – Arbitrary metadata dict (e.g. method name, hyperparameters, timestamp).

Raises:

ValueError – If column is not a registered weight column.

Examples

>>> import pandas as pd
>>> from balance.sample_frame import SampleFrame
>>> df = pd.DataFrame({"id": ["1", "2"], "x": [10, 20],
...                    "weight": [1.0, 2.0]})
>>> sf = SampleFrame.from_frame(df)
>>> sf.set_weight_metadata("weight", {"method": "ipw"})
>>> sf.weight_metadata()
{'method': 'ipw'}

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

Replace the active weight column values.

This is the canonical weight-update method for balance objects. Both SampleFrame and BalanceFrame use this implementation (BalanceFrame delegates here). It also satisfies the BalanceDFSource protocol and is used by BalanceDFWeights.trim() to update weight values after trimming.

If weights is a float, all rows are set to that value. If None, all rows are set to 1.0. If a Series, behavior depends on use_index:

• use_index=False (default): the Series must have the same length as the DataFrame; values are assigned positionally.

• use_index=True: values are aligned by index. Rows whose index is missing from weights are set to NaN (pandas index-alignment semantics), and a warning is emitted.

All weight values are cast to float64.

Parameters:

• weights – New weight values — a Series, scalar, or None.

• use_index – If True, align a Series by index instead of requiring an exact length match.

Raises:

ValueError – If no active weight column is set, or if use_index=False and a Series has a different length than the DataFrame.

Examples

>>> import pandas as pd
>>> from balance.sample_frame import SampleFrame
>>> df = pd.DataFrame({"id": [1, 2], "x": [10, 20],
...                    "weight": [1.0, 2.0]})
>>> sf = SampleFrame.from_frame(df)
>>> sf.set_weights(pd.Series([3.0, 4.0]))
>>> sf.weight_series.tolist()
[3.0, 4.0]

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 trim_weights() for the computation, then writes the result back via set_weights(). A weight history column (weight_trimmed_N) is added so the pre-trim values are preserved.

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 numeric target value. (This is a general-purpose rescaling parameter — not related to the “target population” concept in BalanceFrame.)

• inplace – If True, mutate this SampleFrame and return it. If False (default), return a new SampleFrame with trimmed weights and the original left untouched.

Returns:

The SampleFrame with trimmed weights (self if inplace, else a new copy).

Examples

>>> import pandas as pd
>>> from balance.sample_frame import SampleFrame
>>> sf = SampleFrame.from_frame(
...     pd.DataFrame({"id": [1, 2, 3], "weight": [1.0, 2.0, 100.0]}))
>>> sf2 = sf.trim(ratio=2)
>>> sf2.weight_series.max() < 100.0
True
>>> "weight_trimmed_1" in sf2._df.columns
True

property weight_column: str | None

Name of the currently active weight column, or None.

Note

In balance 0.19.0, weight_column was changed from returning weight data (pd.Series) to returning the column name (str). If you need weight data, use weight_series.

Returns:

The active weight column name.

Return type:

str | None

Examples

>>> import pandas as pd
>>> from balance.sample_frame import SampleFrame
>>> df = pd.DataFrame({"id": ["1", "2"], "x": [10, 20],
...                    "weight": [1.0, 2.0]})
>>> sf = SampleFrame.from_frame(df)
>>> sf.weight_column
'weight'

property weight_columns_all: list[str]

Names of all registered weight columns.

Returns a copy so that callers cannot accidentally mutate the internal column-role registry.

Returns:

Weight column names.

Return type:

list[str]

Examples

>>> import pandas as pd
>>> from balance.sample_frame import SampleFrame
>>> sf = SampleFrame._create(
...     df=pd.DataFrame({"id": [1], "x": [10], "w1": [1.0], "w2": [2.0]}),
...     id_column="id", covar_columns=["x"],
...     weight_columns=["w1", "w2"])
>>> sf.weight_columns_all
['w1', 'w2']

weight_metadata(column: str | None = None) → dict[str, Any]

Retrieve metadata for a weight column.

Parameters:

column (str, optional) – Weight column name. Defaults to the active weight column.

Returns:

The metadata dict, or an empty dict if none was set.

Return type:

dict

Examples

>>> import pandas as pd
>>> from balance.sample_frame import SampleFrame
>>> df = pd.DataFrame({"id": ["1", "2"], "x": [10, 20],
...                    "weight": [1.0, 2.0]})
>>> sf = SampleFrame.from_frame(df)
>>> sf.weight_metadata()
{}

property weight_series: Series

Active weight column as a Series (BalanceDFSource protocol).

Returns the active weight column values as a pd.Series. This is the thin protocol-level accessor used by BalanceDF and its subclasses. Unlike df_weights which returns a single-column DataFrame, this returns a plain Series.

Returns:

The active weight column values.

Return type:

pd.Series

Raises:

ValueError – If no active weight column is set.

Examples

>>> import pandas as pd
>>> from balance.sample_frame import SampleFrame
>>> df = pd.DataFrame({"id": [1, 2], "x": [10, 20],
...                    "weight": [1.0, 2.0]})
>>> sf = SampleFrame.from_frame(df)
>>> sf.weight_series.tolist()
[1.0, 2.0]

weights() → Any

Return a BalanceDFWeights for this SampleFrame.

Creates a weight analysis view backed by this SampleFrame, inheriting any linked sources set via _links.

Returns:

Weight view backed by this SampleFrame.

Return type:

BalanceDFWeights

Examples

>>> import pandas as pd
>>> from balance.sample_frame import SampleFrame
>>> sf = SampleFrame.from_frame(
...     pd.DataFrame({"id": [1, 2], "x": [10.0, 20.0],
...                   "weight": [1.0, 2.0]}))
>>> sf.weights().df.columns.tolist()
['weight']