[FEATURE] Change Sample and BalanceDF roles

[talgalili]

The current design saves data in Sample, e.g.:
https://github.com/facebookresearch/balance/blob/main/balance/sample_class.py#L188
And then set_target links to another sample:
https://github.com/facebookresearch/balance/blob/main/balance/sample_class.py#L554
Then balanceDF extracts information from Sample and does stuff on it (plot, summary, etc.)

This is an odd design.

A better design should have something like
Balancedf (what Sample is currently) - including X, W (many weights), Y (outcomes), Y_hat (estimated Y). Each are a df sharing the same index, but we want to know which is which for the different diagnostics and model fitting.
** For each X, we want to also keep which transformation should be applied to it (bucketing, NA filling, formula, etc).
The transformation output might be memoized.
We WANT the original X to also be kept, so that we could see the impact of weights on the original scale, not just the transformed one.
For each Y_hat we want to know how we estimated it.
For each W we want to know how we estimated it.

Use Balancedf for sample and target population. Combine them together with ‘Sample’ object.
PROBLEM: the current implementation stores sample and target as ‘Sample’, and you have ‘Sample’ with another one as ‘Target’ (which is a bad design). Instead, we should make sure that when we combine sample and target, both will be saved as balancedf
We want Sample to apply the same transformations for both sample and target

Going with something like this will give us a more consistent architecture which will enable the package to develop to more directions and use-cases.

This is closely related to:
#51

All the names and structure is a bit mixed.

Alternative object structure

A potential alternative names and hierarchy (all tentative):

1. SampleDF (maybe SampleDFs, or SampleFrame) = this is a class with 5 DataFrame(s), for X, W (including the baseline/unadjusted/design-weights, and the adjusted ones - at least one adjust option), Y, Y_hat, and Misc (columns that are kept, but not really needed for the core analysis). The name 'sample' is good, since we have responders and target, but it's fine that they are a subset of another population. This could also be a single DataFrame, just with the attributes of which column is what.
2. BalancedDF (or BalancedDFs, or BalancedFrame, or BalancedSample) = this is a class that includes two SampleDF, one for responders (will have Y, can also be called something like BiasedDF or something like this) and one for the target. This one implements core capabilities like getting a single df of everything, responders and target with all columns together - original index and overall index. Also things like download, and plot, and summary, etc. And from it, we can get a subset of classes. BalancedDFCovars, BalancedDFWeights, BalancedDFOutcomes - each with their own set of methods (plot, df, etc.)

Then it would be something like:

from balance import SampleFrame, BalancedFrame

sf_responders = SampleFrame.from_frame(df1, covars_columns, weights_columns, outcome_columns)
sf_target = SampleFrame.from_frame(df2, covars_columns)
bf = BalancedFrame.set_sample(sf_responders).set_target(sf_target)

# Then some form of 
bf.adjust()
bf.summary() #etc.
bf.covars().plot() #etc.

The challenge with the naming is that 'sample' has to do with taking a sample from 'something', and it can be a sample from the target population. So set_sample could have just as well been used for the target. So it's not ideal. But set_responders is too long. set_target could also be set_population or set_target_population but these are too long.

NOTE: Doing this transition before balance 1.0.0 is critical (since this will be the 'official' design).

[talgalili]

A potential breakdown of how to do it (thanks to claude)

Balance Redesign: Incremental Diff Plan
Overview
The goal is to transition from the current architecture:

Sample(df) → Sample.set_target(Sample) → Sample.adjust() → Sample.covars().plot()

To the new architecture:

SampleFrame(X, W, Y, Y_hat, Misc) → BalancedFrame(SampleFrame_responders, SampleFrame_target) → bf.adjust() → bf.covars().plot()

There are 12 diffs organized in 4 phases.

Phase 1: Introduce SampleFrame (the new data container)
Diff 1: Create SampleFrame class with core data storage
Goal: Introduce the SampleFrame class that holds 5 DataFrames (X, W, Y, Y_hat, Misc) sharing the same index. This is the "one idea" — a proper multi-DataFrame container with column-role awareness.

Files to create

Create balance/sample_frame.py
Details:

SampleFrame stores data as a single internal pd.DataFrame (like Sample._df today), but with explicit metadata tracking which columns belong to which role: covars (X), weights (W), outcomes (Y), predicted_outcomes (Y_hat), and misc.
Factory method SampleFrame.from_frame(df, covars_columns=None, weight_column=None, outcome_columns=None, predicted_outcome_columns=None, misc_columns=None, id_column=None, ...) — mimics the current Sample.from_frame() signature but adds the new column roles.
If covars_columns is not specified, infer them by exclusion (same logic as current Sample._covar_columns_names()).
If weight_column is not specified, guess it (same guess_id_column pattern).
Default weights to 1.0 if absent (same as current).
Properties: .covars_df, .weights_df, .outcomes_df, .predicted_outcomes_df, .misc_df, .id_column, .df (full reconstruction).
repr / str showing shape and column roles.
No adjustment logic, no linking, no plotting — just data storage and column-role metadata.
Register in init.py with a # noqa import so it's importable but not yet the default path.
Backward compatibility: Fully additive — no existing code is changed. Sample continues to work exactly as before.

Diff 2: Add per-column transformation metadata to SampleFrame
Goal: Each covariate column in SampleFrame can have an associated transformation specification (bucketing, NA filling, formula, etc.), with memoized output.

Files to modify:

balance/sample_frame.py
Details:

Add a _transformations dict attribute: Dict[str, TransformSpec] mapping column names to transformation specifications.
TransformSpec is a small dataclass TransformSpec(func: Callable, kwargs: dict, description: str). Lives in balance/sample_frame.py or a new balance/transform_spec.py.
Methods:
set_transformations(transformations: Dict[str, Callable | TransformSpec]) — set per-column transformations. Validates column names exist in covars.
set_default_transformations() — delegates to balance.adjustment.default_transformations() to auto-detect quantize/fct_lump per column.
transformed_covars_df (property) — applies transformations and returns the result. Uses functools.lru_cache or a private _transformed_cache dict, invalidated when _transformations or the underlying data changes.
clear_transform_cache() — manual cache invalidation.
The original X is always preserved in _df. The transformed version is computed on demand and cached. This satisfies the requirement "we WANT the original X to also be kept."
The transformation dict is shallow-copied when SampleFrame is deep-copied, so adjustments don't mutate the original.
Backward compatibility: Fully additive.

Diff 3: Add weight and prediction provenance metadata to SampleFrame
Goal: For each W column, track how it was estimated. For each Y_hat column, track how it was estimated. This gives full provenance.

Files to modify:

balance/sample_frame.py
Details:

Add _weight_metadata: Dict[str, Dict[str, Any]] — maps weight column names to metadata dicts (e.g., {"method": "ipw", "model": {...}, "source_sample": "unadjusted"}).
Add _prediction_metadata: Dict[str, Dict[str, Any]] — maps Y_hat column names to metadata dicts (e.g., {"model_type": "logistic", "features_used": [...], "performance": {...}}).
Methods:
set_weight_metadata(column: str, metadata: dict) — stores provenance for a weight column.
set_prediction_metadata(column: str, metadata: dict) — stores provenance for a Y_hat column.
weight_metadata(column: str = None) -> dict — retrieves metadata. If column is None and there's only one weight column, returns that one.
prediction_metadata(column: str = None) -> dict — same pattern.
Support multiple weight columns: SampleFrame can hold baseline/design weights AND adjusted weights simultaneously. The "active" weight column is tracked via a _active_weight_column: str attribute, with a .set_active_weight(column_name) method.
.weights_df returns only the active weight column (for backward compat with single-weight consumers).
.all_weights_df returns all weight columns.
Backward compatibility: Fully additive. No existing code changes.

Phase 2: Introduce BalancedFrame (the new sample+target container)
Diff 4: Create BalancedFrame class combining two SampleFrames
Goal: Introduce BalancedFrame — the replacement for "Sample-with-target-linked." It holds two SampleFrame objects (responders + target) as first-class citizens, not via a hidden _links dict.

Files to create

Create balance/balanced_frame.py
Modify balance/init.py to export BalancedFrame
Details:

BalancedFrame holds:
_responders: SampleFrame — the sample being adjusted (has Y, W, potentially Y_hat).
_target: SampleFrame — the target population (typically has X and W, rarely Y).
_unadjusted: SampleFrame | None — snapshot of responders before adjustment (set by adjust()).
Construction via builder pattern:
bf = BalancedFrame(responders=sf_responders, target=sf_target)

Or via a fluent API:
bf = BalancedFrame.from_frames(responders=sf_responders, target=sf_target)

Properties:
.responders → SampleFrame
.target → SampleFrame
.unadjusted → SampleFrame | None
.is_adjusted → bool
Validation on construction:
Responders and target must share the same covariate column names (warning if not exact match, error if zero overlap).
IDs should not overlap (warning if they do).
Transformation consistency: BalancedFrame ensures the same transformations are applied to both responders and target covariates. Method: .set_transformations(...) sets transformations on both SampleFrames simultaneously.
repr / str showing both frames' shapes and adjustment status.
Backward compatibility: Fully additive. Sample unchanged.

Diff 5: Add adjust() to BalancedFrame
Goal: Implement the adjustment workflow on BalancedFrame, delegating to the existing weighting methods.

Files to modify:

balance/balanced_frame.py
Details:

BalancedFrame.adjust(method="ipw", **kwargs) -> BalancedFrame:
Validates that _target is set.
Gets transformation-applied covariates for both responders and target (using SampleFrame.transformed_covars_df if transformations are set, otherwise raw covars).
Calls balance.adjustment._find_adjustment_method(method) (reuses existing infrastructure).
Calls the adjustment function with the same contract: (sample_df, sample_weights, target_df, target_weights, **kwargs).
Creates a new BalancedFrame (immutable pattern):
Deep-copies responders SampleFrame.
Adds the adjusted weights as a new weight column (e.g., "weight_adjusted"), sets it as active.
Stores _unadjusted pointing to the original responders.
Stores the model in _adjustment_model and in the weight metadata.
Returns the new BalancedFrame.
Also support adjust(method=some_callable) for custom adjustment functions.
BalancedFrame.model() → returns _adjustment_model (same as current Sample.model()).
Backward compatibility: Fully additive. Sample.adjust() unchanged.

Diff 6: Add covars(), weights(), outcomes() accessors and BalanceDF integration to BalancedFrame
Goal: Wire BalancedFrame into the existing BalanceDFCovars, BalanceDFWeights, BalanceDFOutcomes system so plotting, summary, and diagnostics work.

Files to modify:

balance/balanced_frame.py
balance/balancedf_class.py (minor: make BalanceDF accept SampleFrame in addition to Sample)
Details:

BalancedFrame.covars() → returns a new BalancedFrameCovars wrapper (or reuses BalanceDFCovars if we adapt it).
Strategy: Create thin adapter classes (BalancedFrameCovars, BalancedFrameWeights, BalancedFrameOutcomes) that wrap the underlying BalanceDF* objects but add the linked-sample behavior using BalancedFrame's explicit _responders/_target/_unadjusted instead of the _links dict.
These adapters expose the same public API: .plot(), .mean(), .asmd(), .summary(), etc.
Internally, they construct BalanceDFCovars(sample_frame_as_sample) by creating a lightweight adapter that makes SampleFrame look like Sample to the existing BalanceDF code. This avoids rewriting all the stats/plotting logic.
BalancedFrame.weights() → similar pattern.
BalancedFrame.outcomes() → similar pattern, returns None if no outcomes.
BalancedFrame.summary() → delegates to diagnostics logic (reuses existing Sample.summary() internals).
BalancedFrame.diagnostics() → same.
Backward compatibility: The BalanceDF* classes gain a new code path but the old Sample-based path is untouched.

Phase 3: Bridge old and new APIs
Diff 7: Add SampleFrame.from_sample() and Sample.to_sample_frame() conversion methods
Goal: Enable seamless conversion between Sample and SampleFrame so users can migrate incrementally.

Files to modify:

balance/sample_frame.py
balance/sample_class.py
Details:

SampleFrame.from_sample(sample: Sample) -> SampleFrame:
Extracts sample._df, sample.id_column.name, sample.weight_column.name, outcome column names, ignored column names.
Constructs a SampleFrame preserving all data and metadata.
If sample.is_adjusted(), also copies the _adjustment_model into weight metadata.
Sample.to_sample_frame() -> SampleFrame:
Convenience method that calls SampleFrame.from_sample(self).
BalancedFrame.from_sample(sample: Sample) -> BalancedFrame:
If sample.has_target(): creates BalancedFrame(responders=SampleFrame.from_sample(sample), target=SampleFrame.from_sample(sample._links["target"])).
If sample.is_adjusted(): also sets _unadjusted from sample._links["unadjusted"].
Sample.to_balanced_frame() -> BalancedFrame:
Convenience method wrapping the above.
Backward compatibility: Additive methods on existing classes. No behavior changes.

Diff 8: Add BalancedFrame.to_sample() reverse conversion and deprecation warnings on Sample
Goal: Enable reverse conversion (new → old) and begin signaling deprecation of the old API.

Files to modify:

balance/balanced_frame.py
balance/sample_class.py
Details:

BalancedFrame.to_sample() -> Sample:
Constructs a Sample from the responders SampleFrame, sets target and unadjusted via _links.
This enables users who adopt BalancedFrame to still call APIs that expect Sample.
Add warnings.warn(..., DeprecationWarning) to:
Sample.set_target(): "Use BalancedFrame(responders=..., target=...) instead. Sample.set_target() will be removed in balance 1.0.0."
Sample.adjust(): "Use BalancedFrame.adjust() instead. Sample.adjust() will be removed in balance 1.0.0."
Sample.from_frame(): No deprecation yet — it's still the entry point until SampleFrame.from_frame() is fully proven.
The deprecation warnings are off by default — controlled by a module-level flag balance.SHOW_DEPRECATION_WARNINGS = False. Users opt in by setting it to True. This avoids breaking existing notebooks/scripts with noisy warnings before 1.0.
Backward compatibility: All old code continues to work. Warnings are opt-in.

Phase 4: Feature parity and migration
Diff 9: Port summary(), diagnostics(), and design_effect() to BalancedFrame
Goal: Achieve feature parity for the most important analytical methods.

Files to modify:

balance/balanced_frame.py
Details:

BalancedFrame.summary() -> str:
Reuses the logic from Sample.summary() (lines 1552-1727 of sample_class.py), but reads from self._responders, self._target, self._unadjusted instead of self._links.
Refactor the summary logic into a standalone function _build_summary(responders, target, unadjusted, adjustment_model) in a new balance/summary.py (or within balanced_frame.py), callable by both Sample.summary() and BalancedFrame.summary().
BalancedFrame.diagnostics() -> pd.DataFrame:
Same approach — extract logic into shared function, call from both old and new classes.
BalancedFrame.design_effect(), .design_effect_prop(), .covar_means(), .outcome_sd_prop(), .outcome_variance_ratio():
All delegate to the shared extracted functions.
BalancedFrame.plot_weight_density() → delegates to .weights().plot().
Backward compatibility: Sample.summary() and Sample.diagnostics() now call the shared functions internally but produce identical output.

Diff 10: Port keep_only_some_rows_columns(), to_csv(), to_download(), and CLI support to BalancedFrame
Goal: Complete the feature parity for data manipulation and export methods.

Files to modify:

balance/balanced_frame.py
balance/cli.py
Details:

BalancedFrame.keep_only_some_rows_columns(row_condition=None, columns=None) -> BalancedFrame:
Filters both _responders and _target (and _unadjusted if present).
Returns a new BalancedFrame.
BalancedFrame.to_csv(path), BalancedFrame.to_download():
Delegates to responders' data, including adjustment info.
BalancedFrame.df property:
Returns a combined DataFrame with both responders and target rows, distinguished by a "source" column ("responders" / "target"). This is the "single df of everything" mentioned in the design.
CLI updates in balance/cli.py:
Add a --use-new-api flag (default False).
When set, the CLI constructs SampleFrame + BalancedFrame instead of Sample.
Both paths produce identical output (same diagnostics DataFrame).
Backward compatibility: CLI default behavior unchanged. New flag is opt-in.

Diff 11: Add SampleFrame.from_frame() convenience that auto-detects column roles, plus documentation
Goal: Make SampleFrame as easy to use as Sample.from_frame() for new users.

Files to modify:

balance/sample_frame.py
balance/init.py (update welcome message to mention new API)
Details:

Enhance SampleFrame.from_frame() with the same auto-detection logic as Sample.from_frame():
Auto-detect id column via guess_id_column().
Auto-detect weight column (look for "weight" / "weights").
Auto-create default weight of 1.0 if no weight column found.
Type standardization (Int64 → float64, etc.).
All the same validation (no null ids, no null weights, non-negative weights, unique ids).
Add a SampleFrame.from_csv(path, ...) convenience method.
Add comprehensive docstrings to all public methods of SampleFrame and BalancedFrame with usage examples.
Update balance/init.py:
Export SampleFrame and BalancedFrame prominently.
Update version to 0.16.0 or similar pre-1.0 version.
Update the welcome message to mention the new API is available.
Backward compatibility: Fully additive.

Diff 12: Wire Sample internals to use SampleFrame under the hood
Goal: Make Sample a thin wrapper around SampleFrame internally, reducing code duplication and preparing for eventual removal.

Files to modify:

balance/sample_class.py
Details:

Refactor Sample so that internally it holds a SampleFrame instead of raw _df, weight_column, _outcome_columns, _ignored_column_names:
Sample._sample_frame: SampleFrame replaces the individual attributes.
Sample.from_frame() now constructs a SampleFrame first, then wraps it.
All existing properties (df, covars(), weights(), outcomes(), ignored_columns()) delegate to self._sample_frame.
Sample.set_target() still returns a Sample (not a BalancedFrame), but internally stores the target as a SampleFrame.
Sample.adjust() still works identically but internally delegates to shared adjustment logic.
The _links dict is preserved for backward compat, but now stores SampleFrame objects that are wrapped in Sample on access.
All existing tests must pass with zero changes — this is a pure internal refactor.
Private attributes (_df, weight_column, id_column, _outcome_columns) are preserved as computed properties delegating to _sample_frame, so any code accessing them (including tests) continues to work.
Backward compatibility: 100% — identical public API, identical behavior. Only internal implementation changes. This is the riskiest diff and should have thorough test coverage verification.

Summary Table
Diff Phase Title Key Idea Risk
1 1 Create SampleFrame Multi-DataFrame container with column-role metadata Low (additive)
2 1 Transformation metadata Per-column transform specs with memoized output, original X preserved Low (additive)
3 1 Weight & prediction provenance Multiple weights, Y_hat tracking, provenance metadata Low (additive)
4 2 Create BalancedFrame Two SampleFrames (responders + target) as first-class citizens Low (additive)
5 2 BalancedFrame.adjust() Adjustment workflow on new class, reusing existing methods Low (additive)
6 2 BalanceDF integration Wire covars(), weights(), outcomes() into existing plotting/stats Medium (touches balancedf_class.py)
7 3 Old → New conversion Sample.to_sample_frame(), BalancedFrame.from_sample() Low (additive)
8 3 New → Old + deprecation BalancedFrame.to_sample(), opt-in deprecation warnings Low (warnings off by default)
9 4 Port analytics summary(), diagnostics(), design_effect() on BalancedFrame Medium (refactors shared logic)
10 4 Port export + CLI to_csv(), keep_only_some_rows_columns(), CLI flag Low-Medium
11 4 Convenience + docs from_frame() auto-detection, docstrings, version bump Low (additive)
12 4 Internal unification Sample becomes thin wrapper around SampleFrame High (internal refactor, must pass all tests)
Post-1.0 (not in this plan, but the natural next step)
Remove Sample.set_target() and Sample.adjust() (breaking change, hence 1.0).
Make BalancedFrame the primary API.
Potentially rename Sample to just re-export SampleFrame with a deprecation alias.

[talgalili]

Asking cluade to write the data models before and after

Now let me create the before/after diagram. Here's a comprehensive view of how the main objects behave before (current architecture) and after (proposed architecture):

---

BEFORE: Current Architecture

┌─────────────────────────────────────────────────────────────────────┐
│                        CURRENT FLOW                                 │
│                                                                     │
│  Sample.from_frame(df)                                              │
│  ┌──────────────────────────────────────┐                           │
│  │ Sample                               │                           │
│  │  _df: pd.DataFrame (ALL columns)     │                           │
│  │  id_column: pd.Series                │                           │
│  │  weight_column: pd.Series            │                           │
│  │  _outcome_columns: pd.DataFrame|None │                           │
│  │  _ignored_column_names: List[str]    │                           │
│  │  _links: {}  (empty)                 │                           │
│  │  _adjustment_model: None             │                           │
│  │                                      │                           │
│  │  Covariates = _df minus special cols │                           │
│  │  (computed on the fly, not stored)   │                           │
│  └──────────────┬───────────────────────┘                           │
│                 │                                                    │
│                 │ .set_target(target_sample)                         │
│                 │ (deepcopy self, store target in _links)            │
│                 ▼                                                    │
│  ┌──────────────────────────────────────┐                           │
│  │ Sample (deepcopy)                    │                           │
│  │  _df: same as above                  │                           │
│  │  _links: ──────────────────────────┐ │                           │
│  │    "target" ──► Sample (target)    │ │                           │
│  │  _adjustment_model: None           │ │                           │
│  └──────────────┬─────────────────────┘─┘                           │
│                 │                                                    │
│                 │ .adjust(method="ipw")                              │
│                 │ (deepcopy self, compute new weights,               │
│                 │  store model + links)                              │
│                 ▼                                                    │
│  ┌──────────────────────────────────────┐                           │
│  │ Sample (deepcopy, adjusted)          │                           │
│  │  _df: same columns, new weights      │                           │
│  │  weight_column: ADJUSTED weights     │                           │
│  │  _links: ──────────────────────────┐ │                           │
│  │    "target" ──────► Sample (target)│ │                           │
│  │    "unadjusted" ──► Sample (pre-   │ │                           │
│  │                      adjustment)   │ │                           │
│  │  _adjustment_model: {model dict}   │ │                           │
│  └──────────────┬─────────────────────┘─┘                           │
│                 │                                                    │
│                 │ .covars() / .weights() / .outcomes()               │
│                 ▼                                                    │
│  ┌──────────────────────────────────────────────────────────┐       │
│  │ BalanceDFCovars / BalanceDFWeights / BalanceDFOutcomes    │       │
│  │  __df: relevant DataFrame slice                          │       │
│  │  __sample: back-reference to Sample                      │       │
│  │  __name: "covars" | "weights" | "outcomes"               │       │
│  │                                                          │       │
│  │  .plot() / .mean() / .asmd() ──►                         │       │
│  │    _BalanceDF_child_from_linked_samples()                 │       │
│  │    walks __sample._links to build:                        │       │
│  │    {                                                      │       │
│  │      "self":        BalanceDFCovars(adjusted),            │       │
│  │      "target":      BalanceDFCovars(target),              │       │
│  │      "unadjusted":  BalanceDFCovars(original)             │       │
│  │    }                                                      │       │
│  │    → Comparative plots/stats across all three             │       │
│  └──────────────────────────────────────────────────────────┘       │
└─────────────────────────────────────────────────────────────────────┘

Key Issues in Current Design

Problems:
  1. Sample._df is a monolithic DataFrame — column roles (X, W, Y, Y_hat)
     are determined by exclusion logic, not explicit metadata
  2. _links is an opaque dict — relationships between samples are implicit
  3. No support for multiple weight columns (original vs adjusted)
  4. Transformations (bucketing, NA fill) happen inline during adjust(),
     losing the original covariate values
  5. No provenance tracking for weights or predictions
  6. deep_copy chains create hidden object graphs

---

AFTER: Proposed Architecture

┌─────────────────────────────────────────────────────────────────────┐
│                        NEW FLOW                                     │
│                                                                     │
│  SampleFrame.from_frame(df, covars_columns=..., weight_column=...) │
│  ┌──────────────────────────────────────────────────────────────┐   │
│  │ SampleFrame                                                  │   │
│  │                                                              │   │
│  │  _df: pd.DataFrame (ALL data, same as before)                │   │
│  │                                                              │   │
│  │  Column-role metadata (EXPLICIT, not by exclusion):          │   │
│  │  ┌────────────────────────────────────────────┐              │   │
│  │  │ _column_roles:                             │              │   │
│  │  │   covars (X):     ["age","gender","income"]│              │   │
│  │  │   weights (W):    ["weight"]               │              │   │
│  │  │   outcomes (Y):   ["converted"]            │              │   │
│  │  │   predicted (Ŷ):  ["p_converted"]          │              │   │
│  │  │   misc:           ["source","region"]      │              │   │
│  │  │   id:             "respondent_id"          │              │   │
│  │  └────────────────────────────────────────────┘              │   │
│  │                                                              │   │
│  │  Transformation layer (Diff 2):                              │   │
│  │  ┌────────────────────────────────────────────┐              │   │
│  │  │ _transformations: Dict[str, TransformSpec] │              │   │
│  │  │   "age" → TransformSpec(quantize, k=10)    │              │   │
│  │  │   "income" → TransformSpec(bucketize, ...) │              │   │
│  │  │                                            │              │   │
│  │  │ .transformed_covars_df (cached, on-demand) │              │   │
│  │  │ Original X always preserved in _df         │              │   │
│  │  └────────────────────────────────────────────┘              │   │
│  │                                                              │   │
│  │  Weight provenance (Diff 3):                                 │   │
│  │  ┌────────────────────────────────────────────┐              │   │
│  │  │ _weight_metadata:                          │              │   │
│  │  │   "weight" → {method: "design"}            │              │   │
│  │  │   "weight_adjusted" → {method: "ipw",      │              │   │
│  │  │                        model: {...}}        │              │   │
│  │  │ _active_weight_column: "weight_adjusted"   │              │   │
│  │  │                                            │              │   │
│  │  │ .weights_df → active weight only           │              │   │
│  │  │ .all_weights_df → all weight columns       │              │   │
│  │  └────────────────────────────────────────────┘              │   │
│  │                                                              │   │
│  │  Prediction provenance (Diff 3):                             │   │
│  │  ┌────────────────────────────────────────────┐              │   │
│  │  │ _prediction_metadata:                      │              │   │
│  │  │   "p_converted" → {model: "logistic",      │              │   │
│  │  │                    features: [...],         │              │   │
│  │  │                    performance: {...}}      │              │   │
│  │  └────────────────────────────────────────────┘              │   │
│  │                                                              │   │
│  │  Properties:                                                 │   │
│  │    .covars_df, .weights_df, .outcomes_df,                    │   │
│  │    .predicted_outcomes_df, .misc_df, .df                     │   │
│  └──────────────────────────────────────────────────────────────┘   │
│                                                                     │
│                                                                     │
│  BalancedFrame(responders=sf_resp, target=sf_target)                │
│  ┌──────────────────────────────────────────────────────────────┐   │
│  │ BalancedFrame                                                │   │
│  │                                                              │   │
│  │  Explicit, first-class references (NOT a _links dict):       │   │
│  │  ┌─────────────────────────────────────────────────────┐     │   │
│  │  │                                                     │     │   │
│  │  │  _responders: SampleFrame ◄── the sample to adjust  │     │   │
│  │  │       ┌──────────────────────────┐                  │     │   │
│  │  │       │ covars, weights, Y, Ŷ,  │                  │     │   │
│  │  │       │ transformations,         │                  │     │   │
│  │  │       │ weight provenance        │                  │     │   │
│  │  │       └──────────────────────────┘                  │     │   │
│  │  │                                                     │     │   │
│  │  │  _target: SampleFrame ◄── the target population     │     │   │
│  │  │       ┌──────────────────────────┐                  │     │   │
│  │  │       │ covars, weights          │                  │     │   │
│  │  │       │ (typically no Y)         │                  │     │   │
│  │  │       └──────────────────────────┘                  │     │   │
│  │  │                                                     │     │   │
│  │  │  _unadjusted: SampleFrame | None                    │     │   │
│  │  │       (snapshot before adjustment, set by adjust())  │     │   │
│  │  │                                                     │     │   │
│  │  └─────────────────────────────────────────────────────┘     │   │
│  │                                                              │   │
│  │  .is_adjusted → bool                                        │   │
│  │  _adjustment_model: dict | None                              │   │
│  │                                                              │   │
│  └──────────────┬───────────────────────────────────────────────┘   │
│                 │                                                    │
│                 │ .adjust(method="ipw")                              │
│                 ▼                                                    │
│  ┌──────────────────────────────────────────────────────────────┐   │
│  │ BalancedFrame (NEW instance, immutable pattern)              │   │
│  │                                                              │   │
│  │  _responders: SampleFrame (deepcopy)                         │   │
│  │    → new weight col "weight_adjusted" added                  │   │
│  │    → _active_weight_column = "weight_adjusted"               │   │
│  │    → weight_metadata["weight_adjusted"] = {method, model}    │   │
│  │    → original weights PRESERVED in "weight" column           │   │
│  │                                                              │   │
│  │  _target: SampleFrame (same reference or copy)               │   │
│  │                                                              │   │
│  │  _unadjusted: SampleFrame (original responders snapshot)     │   │
│  │                                                              │   │
│  │  _adjustment_model: {model dict}                             │   │
│  └──────────────┬───────────────────────────────────────────────┘   │
│                 │                                                    │
│                 │ .covars() / .weights() / .outcomes()               │
│                 ▼                                                    │
│  ┌──────────────────────────────────────────────────────────────┐   │
│  │ BalancedFrameCovars / Weights / Outcomes (adapters)           │   │
│  │                                                              │   │
│  │  Wraps existing BalanceDF* classes, but uses explicit         │   │
│  │  _responders / _target / _unadjusted instead of _links dict  │   │
│  │                                                              │   │
│  │  .plot() / .mean() / .asmd() directly builds:                │   │
│  │  {                                                           │   │
│  │    "self":       covars from adjusted _responders,            │   │
│  │    "target":     covars from _target,                         │   │
│  │    "unadjusted": covars from _unadjusted                     │   │
│  │  }                                                           │   │
│  │  No _links traversal needed — references are explicit.        │   │
│  └──────────────────────────────────────────────────────────────┘   │
└─────────────────────────────────────────────────────────────────────┘

---

Side-by-Side Comparison: Object Lifecycle

═══════════════════════════════════════════════════════════════════════
                    BEFORE (Sample)                    
═══════════════════════════════════════════════════════════════════════

Step 1: Create
    s = Sample.from_frame(df, outcome_columns=["y"])
    ┌─────────────────────┐
    │ Sample              │
    │  _df: full df       │  ← monolithic, roles by exclusion
    │  _links: {}         │
    │  covars = _df \ {id, weight, y}  (computed)
    └─────────────────────┘

Step 2: Link target
    s_t = s.set_target(target_sample)
    ┌─────────────────────┐     ┌──────────────┐
    │ Sample (deepcopy)   │────►│ Sample       │
    │  _links["target"]   │     │ (target)     │
    └─────────────────────┘     └──────────────┘
    NOTE: target is just a Sample too, same type,
          no structural distinction

Step 3: Adjust
    adj = s_t.adjust(method="ipw")
    ┌─────────────────────┐     ┌──────────────┐
    │ Sample (deepcopy)   │──┬─►│ Sample       │
    │  weight_column:     │  │  │ (target)     │
    │    OVERWRITTEN with  │  │  └──────────────┘
    │    adjusted weights  │  │  ┌──────────────┐
    │  _links["target"]   │  └─►│ Sample       │
    │  _links["unadj."]   │     │ (unadjusted) │
    │  _adjustment_model   │     │ = pre-adj self│
    └─────────────────────┘     └──────────────┘
    PROBLEM: original weights are LOST (overwritten)
    PROBLEM: transformations applied inline, original X lost

Step 4: Analyze
    adj.covars().plot()
    BalanceDFCovars walks _links dict → builds comparison dict


═══════════════════════════════════════════════════════════════════════
                    AFTER (SampleFrame + BalancedFrame)
═══════════════════════════════════════════════════════════════════════

Step 1: Create
    sf = SampleFrame.from_frame(df,
           covars_columns=["age","gender","income"],
           outcome_columns=["y"])
    ┌─────────────────────────────────────┐
    │ SampleFrame                         │
    │  _df: full df                       │
    │  _column_roles: {                   │  ← EXPLICIT role metadata
    │    covars: ["age","gender","income"],│
    │    weights: ["weight"],             │
    │    outcomes: ["y"],                 │
    │    predicted: [],                   │
    │    misc: []                         │
    │  }                                  │
    │  _transformations: {}               │  ← separate layer
    │  _weight_metadata: {}               │  ← provenance tracking
    │  _active_weight_column: "weight"    │
    └─────────────────────────────────────┘

Step 2: Link target  (no more set_target mutation!)
    bf = BalancedFrame(responders=sf_resp, target=sf_target)
    ┌──────────────────────────┐
    │ BalancedFrame            │
    │                          │
    │  _responders: ──────────►│ SampleFrame (responders)
    │  _target: ──────────────►│ SampleFrame (target)
    │  _unadjusted: None       │
    │                          │
    │  Both are explicitly     │  ← explicit, typed references
    │  typed SampleFrame       │     not an opaque dict
    │  objects, not generic    │
    │  "links"                 │
    └──────────────────────────┘

Step 3: Adjust  (returns new BalancedFrame, immutable)
    adj_bf = bf.adjust(method="ipw")
    ┌──────────────────────────┐
    │ BalancedFrame (NEW)      │
    │                          │
    │  _responders: ──────────►│ SampleFrame (deepcopy)
    │    │                     │   weight cols: ["weight",
    │    │                     │                 "weight_adjusted"]
    │    │                     │   _active: "weight_adjusted"
    │    │                     │   weight_metadata:
    │    │                     │     "weight" → {method: "design"}
    │    │                     │     "weight_adjusted" → {method: "ipw",
    │    │                     │                          model: {...}}
    │    │                     │   ORIGINAL weights PRESERVED ✓
    │    │                     │   ORIGINAL covars PRESERVED ✓
    │    │                     │   transformed_covars_df cached ✓
    │  _target: ──────────────►│ SampleFrame (target, same)
    │  _unadjusted: ──────────►│ SampleFrame (pre-adj snapshot)
    │  _adjustment_model: {}   │
    └──────────────────────────┘

Step 4: Analyze
    adj_bf.covars().plot()
    BalancedFrameCovars reads _responders, _target, _unadjusted
    directly — no _links dict traversal

---

Data Flow Comparison

BEFORE:
  pd.DataFrame ──► Sample ──set_target──► Sample ──adjust──► Sample ──covars──► BalanceDFCovars
       │              │                      │                  │                     │
       │         (monolithic df,        (deepcopy +        (deepcopy +          (walks _links
       │          roles by              _links["target"])   _links["unadj"])      dict to find
       │          exclusion)                                 weights              related
       │                                                    overwritten)         samples)

AFTER:
  pd.DataFrame ──► SampleFrame ──┐
                                  ├──► BalancedFrame ──adjust──► BalancedFrame ──covars──► Adapter
  pd.DataFrame ──► SampleFrame ──┘         │                        │                      │
       │              │                     │                        │                (reads explicit
       │         (explicit roles,     (explicit               (new instance,       _responders,
       │          transformations,    _responders +            orig weights         _target,
       │          provenance)        _target refs)             preserved,           _unadjusted)
       │                                                      provenance
       │                                                      tracked)

---

Migration Bridge (Phase 3)

┌──────────────┐                    ┌──────────────────┐
│   Sample     │◄──.to_sample()─────│  BalancedFrame   │
│   (old)      │                    │  (new)           │
│              │──.to_balanced_     │                  │
│              │   frame()─────────►│                  │
│              │                    │                  │
│              │──.to_sample_       │                  │
│              │   frame()──┐       │                  │
└──────────────┘            │       └──────────────────┘
                            ▼                 ▲
                   ┌──────────────┐           │
                   │ SampleFrame  │───────────┘
                   │ (new)        │  (used to construct
                   │              │   BalancedFrame)
                   └──────────────┘

Phase 4 / Diff 12: Sample becomes a thin wrapper
┌──────────────────────────────────────┐
│ Sample (thin wrapper)                │
│                                      │
│  _sample_frame: SampleFrame  ◄──── internal storage
│                                      │
│  .df ──► self._sample_frame.df       │  all properties
│  .covars() ──► delegates             │  delegate to
│  .weights() ──► delegates            │  SampleFrame
│  .set_target() ──► still returns     │
│                    Sample (compat)    │
│  .adjust() ──► still returns         │
│                Sample (compat)       │
└──────────────────────────────────────┘

---

Summary of Key Behavioral Differences

Aspect	Before (Sample)	After (SampleFrame + BalancedFrame)
Column roles	By exclusion (_df minus special cols)	Explicit metadata dict
Sample-target link	_links["target"] (opaque dict)	_target: SampleFrame (typed attribute)
Unadjusted ref	_links["unadjusted"] (opaque dict)	_unadjusted: SampleFrame (typed attribute)
Original weights	Overwritten on adjust()	Preserved; new weight column added
Transformations	Applied inline, originals lost	Separate layer; originals always in _df
Weight provenance	None	_weight_metadata per weight column
Prediction provenance	None	_prediction_metadata per Y_hat column
Multiple weights	Not supported	Multiple cols + _active_weight_column
Immutability	set_target() returns deepcopy; adjust() mutates copy	adjust() returns new BalancedFrame
Comparative analysis	_BalanceDF_child_from_linked_samples() walks _links	Adapter reads explicit _responders/_target/_unadjusted

[talgalili]

excalidraw image

[talgalili]

And a transition diagram:

[talgalili]

Another version

[talgalili]

nit: I'm debating between BalancedFrame and BalanceFrame. I think I prefer to go with BalanceFrame...

[talgalili]

@neuralsorcerer
FYI: what I wrote in this issue is my vision for the next stage for balance. It's what's required to get it to include outcome modeling (and doubly robust, etc.).
However, before we get there, I need to get internal alignment within Meta that there are people interested in this work.
This might take me 1-3 weeks to get.
Once done, I'd love to support this work.

Would you be interested in doing some of the diffs outlined?

[neuralsorcerer]

@talgalili I'd love to work on this :)

[talgalili]

Amazing, I'll let you know once I know 👍 🥇