From 7df6b9637c05f4d2c5a3bbef639e0a6040bfafd9 Mon Sep 17 00:00:00 2001 From: igerber Date: Sun, 28 Jun 2026 07:51:23 -0400 Subject: [PATCH] perf(within-transform): drop the redundant full-frame copy in within_transform within_transform -- the two-way fixed-effects demeaning helper shared by TwoWayFixedEffects, SunAbraham, WooldridgeDiD, and BaconDecomposition -- copied the entire input frame defensively before pd.concat-ing the demeaned columns onto it. But the demean is read-only and concat does not mutate its inputs, so that copy is pure overhead. Compute the demeaned columns first, then attach them as a single consolidated block via pd.concat with no defensive pre-copy (under pandas copy-on-write the original columns are shared, not copied). Peak RSS of a wide TwoWayFixedEffects(vcov_type="hc1") fit drops ~8% (964 -> 886 MB at 400k units x 6 covariates); the win scales with panel width. Bit-identical: proven at atol=0 for TWFE (incl. the survey replicate-weight path), SunAbraham, Wooldridge, and Bacon -- frame assembly only, the demean arithmetic is unchanged. The inplace/suffix params are decoupled (inplace = in-place attach vs concat; suffix = column naming), matching demean_by_group. New unit tests lock the modes + the no-fragmentation many-column path. PR-C of the memory-scaling work. The ImputationDiD/TwoStageDiD _iterative_demean vectorization is deferred (TODO) -- it is non-bit-identical on pandas 3.0 and analytical-path-only. Co-Authored-By: Claude Opus 4.8 (1M context) --- CHANGELOG.md | 10 +++ TODO.md | 1 + diff_diff/utils.py | 115 ++++++++++++++++--------------- tests/test_within_transform.py | 122 +++++++++++++++++++++++++++++++++ 4 files changed, 193 insertions(+), 55 deletions(-) create mode 100644 tests/test_within_transform.py diff --git a/CHANGELOG.md b/CHANGELOG.md index 65233905..10511c7d 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -70,6 +70,16 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 tests, and the full Python⇄Rust equivalence suite (`tests/test_rust_backend.py`). ### Performance +- **`within_transform` no longer takes a redundant full-frame copy.** The two-way within + (fixed-effects) demeaning helper — shared by `TwoWayFixedEffects`, `SunAbraham`, + `WooldridgeDiD`, and `BaconDecomposition` — copied the entire input frame defensively before + `pd.concat`-ing the demeaned columns onto it, even though the demean is read-only and + `concat` does not mutate its inputs. That copy is removed: the demeaned columns are attached + as a single consolidated block via `pd.concat` (under pandas copy-on-write the original + columns are shared, not copied). Peak RSS of a wide `TwoWayFixedEffects(vcov_type="hc1")` fit + drops ~8% (e.g. 964 → 886 MB at 400k units × 6 covariates); the win scales with panel width. + **Bit-identical** (proven at `atol=0` for TWFE incl. the replicate-weight path, SunAbraham, + Wooldridge, and Bacon) — frame assembly only, the demean arithmetic is unchanged. - **`ImputationDiD` conservative-variance: cache the untreated-projection factorization per fit.** The exact imputation projection `v_untreated = -A_0 (A_0'[W]A_0)^{-1} A_1'w` has a target-invariant design (`A_0`/`A_1`/factorization) and a target-specific RHS (`A_1'w`), but diff --git a/TODO.md b/TODO.md index 5eec4702..28b51fbe 100644 --- a/TODO.md +++ b/TODO.md @@ -65,6 +65,7 @@ generic sparse-FE, QR+SVD rank-detection redundancy, `check_finite` bypass — m | `HeterogeneousAdoptionDiD` Phase 3 Stute: Appendix-D vectorized form replaces the per-iteration OLS refit with a single precomputed `M = I - X(X'X)^{-1}X'` applied to `eps*eta` (~2× faster, functionally identical). Shipped the literal-refit form to match paper text. | `had_pretests.py::stute_test` | Phase 3 | Mid | Low | | Rust faer SVD ndarray-to-faer conversion overhead (minimal vs SVD cost). | `rust/src/linalg.rs:67` | #115 | Quick | Low | | Multiplier-bootstrap weight chunking (CallawaySantAnna, EfficientDiD, and HAD — all wired through `diff_diff/bootstrap_chunking.py`) covers the **unstratified** survey-PSU generation (the default unit-level bootstrap — `cluster=None`, equivalently `cluster="unit"` — the large-`n_units` OOM case). Remaining gap: the **stratified** survey-PSU generator (`generate_survey_multiplier_weights_batch`, per-stratum + lonely-PSU pooling + FPC) still materializes the full `(n_bootstrap × n_psu)` matrix (consumed via sliced blocks). Stratified designs have few PSUs so this rarely OOMs; tile per-stratum generation over draws (each stratum's draws are independent → contiguous draw-blocks reproduce the stream bit-identically) if a large-PSU stratified design hits memory. | `diff_diff/bootstrap_chunking.py::iter_survey_multiplier_weight_blocks` | follow-up | Mid | Low | +| `ImputationDiD._iterative_demean` (`imputation.py:1064`) and its near-identical twin `TwoStageDiD._iterative_demean` (`two_stage.py:2201`) rebuild a `pd.Series`+`groupby().transform()` every alternating-projection iteration. Precompute the unit/time group codes once and use `np.bincount` for the per-iteration group means. **Not bit-identical** on pandas 3.0 (`np.bincount` is naive accumulation; pandas' `group_mean` is Kahan-compensated → ~5.8e-11, the same order as the demean's `tol=1e-10`), so this needs a tolerance/REGISTRY-Note treatment + golden re-validation. Modest peak effect (the per-iteration Series are transient), analytical-path-only (the bootstrap reuses #562's cached projection). Optimize both twins together (cross-surface-twins). | `imputation.py`, `two_stage.py` | PR-C deferral | Mid | Low | ### Testing / docs diff --git a/diff_diff/utils.py b/diff_diff/utils.py index 8ad63d25..a7c84e75 100644 --- a/diff_diff/utils.py +++ b/diff_diff/utils.py @@ -2667,10 +2667,17 @@ def within_transform( time : str Column name for time period identifier. inplace : bool, default False - If True, modifies the original columns. If False, creates new columns - with the specified suffix. + Controls how the demeaned columns are attached. If False (default), they + are concatenated onto the input as a single block and a new frame is + returned; the input frame is not mutated (no defensive deep copy is taken + — the demean is read-only and ``concat`` does not mutate its inputs). If + True, they are written onto the passed frame in place (the caller must own + it) and that frame is returned. Independent of ``suffix``. suffix : str, default "_demeaned" - Suffix for new column names when inplace=False. + Column naming, independent of ``inplace``. A non-empty suffix writes the + demeaned values to new ``f"{var}{suffix}"`` columns (originals preserved); + ``suffix=""`` overwrites the source columns. Assigning to an existing + column name overwrites it rather than appending a duplicate label. weights : np.ndarray, optional Observation weights for weighted group means. max_iter : int, default 100 @@ -2699,8 +2706,12 @@ def within_transform( >>> df = within_transform(df, ['y', 'x'], 'unit_id', 'year') >>> # df now has 'y_demeaned' and 'x_demeaned' columns """ - if not inplace: - data = data.copy() + # Column naming (``suffix``) is independent of how the demeaned columns are + # attached (``inplace``): an empty suffix targets the source column (overwrite), + # a non-empty suffix a new ``f"{var}{suffix}"`` column. The demean below only + # reads ``data``, so no defensive copy is taken up front. + target_cols = [var if not suffix else f"{var}{suffix}" for var in variables] + demeaned_values: List[np.ndarray] = [] if weights is not None: # Weighted within-transformation via iterative alternating projections @@ -2717,37 +2728,19 @@ def _weighted_group_demean(x, groups, w, w_sum): return x - wx_sum / w_sum non_converged_vars: List[str] = [] - if inplace: - for var in variables: - x = data[var].values.astype(np.float64) - converged = False - for _iter in range(max_iter): - x_old = x.copy() - x = _weighted_group_demean(x, unit_groups, w, unit_w_sum) - x = _weighted_group_demean(x, time_groups, w, time_w_sum) - if np.max(np.abs(x - x_old)) < tol: - converged = True - break - if not converged: - non_converged_vars.append(var) - data[var] = x - else: - demeaned_data = {} - for var in variables: - x = data[var].values.astype(np.float64) - converged = False - for _iter in range(max_iter): - x_old = x.copy() - x = _weighted_group_demean(x, unit_groups, w, unit_w_sum) - x = _weighted_group_demean(x, time_groups, w, time_w_sum) - if np.max(np.abs(x - x_old)) < tol: - converged = True - break - if not converged: - non_converged_vars.append(var) - demeaned_data[f"{var}{suffix}"] = x - demeaned_df = pd.DataFrame(demeaned_data, index=data.index) - data = pd.concat([data, demeaned_df], axis=1) + for var in variables: + x = data[var].values.astype(np.float64) + converged = False + for _iter in range(max_iter): + x_old = x.copy() + x = _weighted_group_demean(x, unit_groups, w, unit_w_sum) + x = _weighted_group_demean(x, time_groups, w, time_w_sum) + if np.max(np.abs(x - x_old)) < tol: + converged = True + break + if not converged: + non_converged_vars.append(var) + demeaned_values.append(x) if non_converged_vars: warn_if_not_converged( False, @@ -2760,22 +2753,34 @@ def _weighted_group_demean(x, groups, w, w_sum): unit_grouper = data.groupby(unit, sort=False) time_grouper = data.groupby(time, sort=False) - if inplace: - for var in variables: - unit_means = unit_grouper[var].transform("mean") - time_means = time_grouper[var].transform("mean") - grand_mean = data[var].mean() - data[var] = data[var] - unit_means - time_means + grand_mean - else: - demeaned_data = {} - for var in variables: - unit_means = unit_grouper[var].transform("mean") - time_means = time_grouper[var].transform("mean") - grand_mean = data[var].mean() - demeaned_data[f"{var}{suffix}"] = ( - data[var] - unit_means - time_means + grand_mean - ).values - demeaned_df = pd.DataFrame(demeaned_data, index=data.index) - data = pd.concat([data, demeaned_df], axis=1) - - return data + for var in variables: + unit_means = unit_grouper[var].transform("mean") + time_means = time_grouper[var].transform("mean") + grand_mean = data[var].mean() + demeaned_values.append((data[var] - unit_means - time_means + grand_mean).values) + + if inplace: + # Write onto the passed frame (the caller must own it); an existing + # same-named column is overwritten. Used for the in-place overwrite + # callers (small column sets); large suffixed sets take the concat path + # below to avoid per-column block fragmentation. + for col, vals in zip(target_cols, demeaned_values): + data[col] = vals + return data + + # Default (non-inplace): attach the demeaned columns as a single consolidated + # block via ``pd.concat``. No defensive copy of ``data`` is taken — the demean + # above is read-only and concat does not mutate its inputs, so the caller's + # frame is preserved (and shared rather than copied under copy-on-write). This + # both drops the redundant copy the old code took before the concat and avoids + # the ``DataFrame is highly fragmented`` warning of N per-column inserts. + new_block = pd.DataFrame(dict(zip(target_cols, demeaned_values)), index=data.index) + # Honor the overwrite contract: if a target name already exists (``suffix=""``, + # or re-demeaning a frame that already carries the suffix), drop it first so the + # concat replaces it instead of producing a duplicate label. ``drop`` returns a + # new frame (the input is not mutated). The common case — fresh suffixed targets + # — has no collision and skips straight to the concat. + collisions = [c for c in target_cols if c in data.columns] + if collisions: + data = data.drop(columns=collisions) + return pd.concat([data, new_block], axis=1) diff --git a/tests/test_within_transform.py b/tests/test_within_transform.py new file mode 100644 index 00000000..9cc2949b --- /dev/null +++ b/tests/test_within_transform.py @@ -0,0 +1,122 @@ +"""Unit tests for the copy-avoiding ``within_transform`` refactor (PR-C). + +Locks the ``inplace``/``suffix`` decoupling, the batch-assignment overwrite +semantics (load-bearing for TWFE's replicate refit, which re-demeans a frame that +already carries the suffix), and the absence of a pandas fragmentation warning on +the many-column path (SunAbraham can demean 100+ interaction columns). +""" + +import warnings + +import numpy as np +import pandas as pd + +from diff_diff.utils import within_transform + + +def _panel(nu=30, nt=5, k=2, seed=0): + rng = np.random.default_rng(seed) + n = nu * nt + data = { + "unit": np.repeat(np.arange(nu), nt), + "time": np.tile(np.arange(nt), nu), + } + for j in range(k): + data[f"v{j}"] = rng.standard_normal(n) + return pd.DataFrame(data) + + +def _ref_demean(df, var): + """Reference unweighted two-way within transform: y - y_i. - y_.t + y_..""" + s = df[var] + return ( + s + - df.groupby("unit")[var].transform("mean") + - df.groupby("time")[var].transform("mean") + + s.mean() + ).values + + +class TestWithinTransformInplaceSuffix: + def test_inplace_false_leaves_input_untouched(self): + df = _panel() + orig = df.copy() + out = within_transform(df, ["v0", "v1"], "unit", "time") + pd.testing.assert_frame_equal(df, orig) # default inplace=False: input not mutated + assert out is not df + assert "v0_demeaned" in out.columns and "v0" in out.columns + np.testing.assert_array_equal(out["v0_demeaned"].values, _ref_demean(df, "v0")) + np.testing.assert_array_equal(out["v1_demeaned"].values, _ref_demean(df, "v1")) + + def test_inplace_true_suffix_mutates_same_object_keeps_originals(self): + df = _panel() + v0_orig = df["v0"].values.copy() + out = within_transform(df, ["v0", "v1"], "unit", "time", inplace=True) + assert out is df # same object, mutated in place (no copy) + np.testing.assert_array_equal(df["v0"].values, v0_orig) # original preserved + np.testing.assert_array_equal(df["v0_demeaned"].values, _ref_demean(df, "v0")) + + def test_inplace_true_empty_suffix_overwrites_source(self): + df = _panel() + ref = _ref_demean(df, "v0") + within_transform(df, ["v0"], "unit", "time", inplace=True, suffix="") + assert "v0_demeaned" not in df.columns + np.testing.assert_array_equal(df["v0"].values, ref) + + def test_redemean_existing_suffix_overwrites_no_duplicate(self): + # The TWFE replicate scenario: a frame that already carries the suffix is + # re-demeaned. The batch assignment must OVERWRITE the existing column + # (single label) rather than append a duplicate — a duplicate label would + # make ``df[col].values`` 2-D and break the downstream np.column_stack. + df = _panel() + within_transform(df, ["v0"], "unit", "time", inplace=True) # adds v0_demeaned + out = within_transform(df, ["v0"], "unit", "time", inplace=True) # re-demean + assert list(out.columns).count("v0_demeaned") == 1 + assert out["v0_demeaned"].values.ndim == 1 + + def test_non_inplace_empty_suffix_overwrites_no_duplicate(self): + # inplace=False + suffix="" targets the existing source column; the concat + # path must drop the original first so the result has ONE "v0" column + # (a duplicate label would make .values 2-D), while leaving the input frame + # unmutated. + df = _panel() + ref = _ref_demean(df, "v0") + out = within_transform(df, ["v0"], "unit", "time", suffix="") + assert list(out.columns).count("v0") == 1 + assert out["v0"].values.ndim == 1 + np.testing.assert_array_equal(out["v0"].values, ref) + np.testing.assert_array_equal(df["v0"].values, _panel()["v0"].values) # input intact + + def test_non_inplace_redemean_existing_suffix_no_duplicate(self): + # inplace=False re-demean of a frame that already carries the suffixed + # column overwrites it (single label), not a duplicate. + df = within_transform(_panel(), ["v0"], "unit", "time") # has v0_demeaned + out = within_transform(df, ["v0"], "unit", "time") # re-demean, inplace=False + assert list(out.columns).count("v0_demeaned") == 1 + assert out["v0_demeaned"].values.ndim == 1 + + def test_weighted_inplace_matches_non_inplace(self): + rng = np.random.default_rng(1) + df = _panel(seed=1) + w = rng.uniform(0.5, 2.0, len(df)) + with warnings.catch_warnings(): + warnings.simplefilter("ignore") + a = within_transform(df.copy(), ["v0"], "unit", "time", weights=w) + b = df.copy() + within_transform(b, ["v0"], "unit", "time", weights=w, inplace=True) + np.testing.assert_array_equal(a["v0_demeaned"].values, b["v0_demeaned"].values) + + +class TestWithinTransformManyColumns: + def test_many_columns_no_fragmentation_warning(self): + # SunAbraham can demean 100+ interaction columns; the default (concat) + # path must attach them as one consolidated block and NOT trigger pandas' + # "DataFrame is highly fragmented" PerformanceWarning that per-column + # inserts would. + df = _panel(nu=20, nt=4, k=150) + cols = [f"v{j}" for j in range(150)] + with warnings.catch_warnings(): + warnings.simplefilter("error", pd.errors.PerformanceWarning) + out = within_transform(df, cols, "unit", "time") # default inplace=False -> concat + assert all(f"{c}_demeaned" in out.columns for c in cols) + assert out is not df