Fix S-learner estimate_ate DataFrame handling and return shape (#1026)

[su-jin1425]

Proposed changes

Fixes #1026.

This PR fixes two issues in BaseSLearner.estimate_ate:

• Fixes a KeyError: (0, 0) when LRSRegressor.estimate_ate() is called with a pandas DataFrame.
• Makes BaseSLearner.estimate_ate() consistently return (ate, lb, ub), matching the other meta-learners.

The statsmodels parameters and confidence intervals are converted to NumPy arrays before positional indexing.

Tests were updated to reflect the consistent S-learner return signature, and a regression test was added for DataFrame input.

Before & After

Test

function T($c,$n) {
    git checkout --force $c | Out-Null
    Write-Host "`n========== $n =========="

    @'
import numpy as np
import pandas as pd
from causalml.inference.meta import LRSRegressor, BaseSRegressor
from sklearn.linear_model import LinearRegression

rng = np.random.RandomState(42)
X = pd.DataFrame(rng.normal(size=(200, 5)), columns=list("abcde"))
t = rng.binomial(1, 0.5, 200)
y = 2*t + X["a"].to_numpy() + 0.5*X["b"].to_numpy() + rng.normal(size=200)

try:
    result = LRSRegressor().estimate_ate(X=X, treatment=t, y=y)
    print("DataFrame: PASS ->", result)
except Exception as e:
    print("DataFrame: FAIL ->", type(e).__name__, e)

X = rng.normal(size=(200, 5))
y = 2*t + X[:, 0] + 0.5*X[:, 1] + rng.normal(size=200)

result = BaseSRegressor(learner=LinearRegression()).estimate_ate(
    X=X, treatment=t, y=y
)

print("S-learner:", result)
print("Return type:", type(result).__name__)
print("Return length:", len(result))
'@ | python
}

T "35f6e734a478eea6fdf1cfaf1b85e0f04ae99b83" "BEFORE #1026 — BUG"
T "ef98264f8d7b8cce18666464cabb66b7abc0e402" "AFTER #1026 — FIXED"

Types of changes

• Bugfix (non-breaking change which fixes an issue)
• New feature (non-breaking change which adds functionality)
• Breaking change (fix or feature that would cause existing functionality to not work as expected)
• Documentation Update (if none of the other choices apply)

Checklist

• I have read the CONTRIBUTING doc
• I have signed the CLA
• Lint and unit tests pass locally with my changes
• I have added tests that prove my fix is effective or that my feature works
• I have added necessary documentation (if appropriate)
• Any dependent changes have been merged and published in downstream modules

Further comments

The return_ci argument for BaseSLearner.estimate_ate() was removed so that S-learners follow the same estimate_ate() return convention as the other meta-learners.

The original DataFrame failure was reproduced against the parent commit and verified as fixed. CI passes on Python 3.11 and 3.12, and lint checks pass.

[Copilot]

Copilot was unable to review this pull request because the user who requested the review has reached their quota limit.

[jeongyoonlee]

docs/examples/meta_learners_with_synthetic_data_multiple_treatment.ipynb still passes the removed return_ci to BaseSRegressor.estimate_ate(), so six cells now raise TypeError. The S-learner ATE section appears twice in the notebook, so each of these is duplicated:

• ate_s = learner_s.estimate_ate(..., return_ci=False, bootstrap_ci=False) — lines 153 / 2430
• ate_s, ate_s_lb, ate_s_ub = learner_s.estimate_ate(..., return_ci=True, bootstrap_ci=False) — 215 / 2524
• ate_s_b, ate_s_lb_b, ate_s_ub_b = learner_s.estimate_ate(..., return_ci=True, bootstrap_ci=True, ...) — 280 / 2596

The last two only need return_ci dropped. The first is a point-estimate-only demo with no equivalent after this change — the following cell displays ate_s, which is now a 3-tuple — so please either drop that cell pair or index the ATE out. Re-run the notebook afterwards so the stored outputs match.

Minor: docs/examples/feature_interpretations_example.ipynb ends a cell with slearner.estimate_ate(X=X, treatment=w_multi, y=y) on a BaseSRegressor; its stored output is now stale (array → 3-tuple).

[su-jin1425]

I’ve updated both affected notebooks to remove the deprecated return_ci usage and handle the new (ate, lb, ub) return consistently. The changes are in commit 4cae40d. Full notebook execution wasn’t possible locally due to the missing Cython _criterion extension, but the affected cells were updated surgically without unrelated output changes.

[su-jin1425]

Use #1030 → #1033 → #1032 → #1031 this order to avoid merge conflicts

[jeongyoonlee]

LGTM. Thanks!