perf(within-transform): drop the redundant full-frame copy in within_transform

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

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_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)

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