hapc 0.2.1__tar.gz → 2.1.0__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (55) hide show
  1. {hapc-0.2.1 → hapc-2.1.0}/CMakeLists.txt +13 -5
  2. {hapc-0.2.1/python/hapc.egg-info → hapc-2.1.0}/PKG-INFO +2 -3
  3. {hapc-0.2.1 → hapc-2.1.0}/pyproject.toml +2 -3
  4. hapc-2.1.0/python/hapc/__init__.py +99 -0
  5. hapc-2.1.0/python/hapc/ate.py +425 -0
  6. hapc-2.1.0/python/hapc/core.py +318 -0
  7. hapc-2.1.0/python/hapc/cv.py +545 -0
  8. hapc-2.1.0/python/hapc/single.py +665 -0
  9. {hapc-0.2.1 → hapc-2.1.0/python/hapc.egg-info}/PKG-INFO +2 -3
  10. {hapc-0.2.1 → hapc-2.1.0}/python/hapc.egg-info/SOURCES.txt +5 -3
  11. {hapc-0.2.1 → hapc-2.1.0}/python/hapc.egg-info/requires.txt +1 -2
  12. {hapc-0.2.1 → hapc-2.1.0}/setup.py +52 -7
  13. {hapc-0.2.1 → hapc-2.1.0}/src/bindings.cpp +52 -15
  14. hapc-2.1.0/src/cv_classi.cpp +113 -0
  15. {hapc-0.2.1 → hapc-2.1.0}/src/cv_fast_pchal.cpp +0 -6
  16. {hapc-0.2.1 → hapc-2.1.0}/src/fast_pchal.cpp +10 -4
  17. {hapc-0.2.1 → hapc-2.1.0}/src/hapc_core.hpp +32 -2
  18. {hapc-0.2.1 → hapc-2.1.0}/src/logistic_call.cpp +4 -8
  19. {hapc-0.2.1 → hapc-2.1.0}/src/pcghal_call.cpp +43 -8
  20. {hapc-0.2.1 → hapc-2.1.0}/src/pcghal_classi_call.cpp +43 -8
  21. {hapc-0.2.1 → hapc-2.1.0}/src/pcghal_cv.cpp +23 -4
  22. hapc-2.1.0/src/pcghal_cv_classi_cpp.cpp +355 -0
  23. {hapc-0.2.1 → hapc-2.1.0}/src/pcghal_cv_cpp.cpp +56 -19
  24. {hapc-0.2.1 → hapc-2.1.0}/src/r_bindings.cpp +128 -17
  25. {hapc-0.2.1 → hapc-2.1.0}/src/single_pcghal_cpp.cpp +23 -8
  26. {hapc-0.2.1 → hapc-2.1.0}/src/single_pchar.cpp +0 -11
  27. hapc-2.1.0/tests/test_api.py +76 -0
  28. hapc-2.1.0/tests/test_ate.py +162 -0
  29. hapc-2.1.0/tests/test_ate_hapc_diagnostics_example.py +184 -0
  30. hapc-2.1.0/tests/test_core.py +126 -0
  31. hapc-2.1.0/tests/test_logistic_regression.py +137 -0
  32. hapc-2.1.0/tests/test_r_vs_python_alpha.py +146 -0
  33. hapc-0.2.1/python/demo_single.py +0 -39
  34. hapc-0.2.1/python/hapc/__init__.py +0 -33
  35. hapc-0.2.1/python/hapc/core.py +0 -112
  36. hapc-0.2.1/python/hapc/cv.py +0 -340
  37. hapc-0.2.1/python/hapc/single.py +0 -259
  38. hapc-0.2.1/python/test_install.py +0 -65
  39. hapc-0.2.1/src/cv_classi.cpp +0 -314
  40. hapc-0.2.1/src/single_pcghal.cpp +0 -204
  41. hapc-0.2.1/tests/test_api.py +0 -68
  42. hapc-0.2.1/tests/test_core.py +0 -140
  43. hapc-0.2.1/tests/test_r_vs_python_alpha.py +0 -363
  44. {hapc-0.2.1 → hapc-2.1.0}/LICENSE +0 -0
  45. {hapc-0.2.1 → hapc-2.1.0}/MANIFEST.in +0 -0
  46. {hapc-0.2.1 → hapc-2.1.0}/README.md +0 -0
  47. {hapc-0.2.1 → hapc-2.1.0}/python/hapc.egg-info/dependency_links.txt +0 -0
  48. {hapc-0.2.1 → hapc-2.1.0}/python/hapc.egg-info/not-zip-safe +0 -0
  49. {hapc-0.2.1 → hapc-2.1.0}/python/hapc.egg-info/top_level.txt +0 -0
  50. {hapc-0.2.1 → hapc-2.1.0}/setup.cfg +0 -0
  51. {hapc-0.2.1 → hapc-2.1.0}/src/cross_kernel.cpp +0 -0
  52. {hapc-0.2.1 → hapc-2.1.0}/src/cv_fast_pchal_python.cpp +0 -0
  53. {hapc-0.2.1 → hapc-2.1.0}/src/mkernel.cpp +0 -0
  54. {hapc-0.2.1 → hapc-2.1.0}/src/pchal_design.cpp +0 -0
  55. {hapc-0.2.1 → hapc-2.1.0}/src/ridge_wrappers.cpp +0 -0
@@ -11,22 +11,29 @@ if(WIN32 AND EXISTS "C:/vcpkg")
11
11
  list(APPEND CMAKE_PREFIX_PATH "C:/vcpkg/installed/x64-windows")
12
12
  endif()
13
13
 
14
- # Find Python
14
+ # Find Python. We intentionally use the modern FindPython3 module and pass
15
+ # Python3_EXECUTABLE from setup.py so the build always targets the *same*
16
+ # interpreter that pip is using. Without this CMake may discover a newer/
17
+ # older system Python and produce a .so tagged for the wrong ABI.
15
18
  find_package(Python3 COMPONENTS Interpreter Development REQUIRED)
19
+ message(STATUS "Python3_EXECUTABLE: ${Python3_EXECUTABLE}")
20
+ message(STATUS "Python3_VERSION: ${Python3_VERSION}")
16
21
 
17
- # Find Eigen3 with fallback to FetchContent for Windows
22
+ # Find Eigen3 with fallback to FetchContent for Windows.
23
+ # Disabling Eigen's tests/docs avoids long Windows configuration hangs and
24
+ # matches the working CI configuration in v0.2.x wheels.
18
25
  find_package(Eigen3 QUIET NO_MODULE)
19
26
  if(NOT Eigen3_FOUND)
20
27
  message(STATUS "Eigen3 not found, downloading via FetchContent...")
28
+ set(EIGEN_BUILD_TESTING OFF CACHE BOOL "" FORCE)
29
+ set(EIGEN_BUILD_DOC OFF CACHE BOOL "" FORCE)
30
+ set(BUILD_TESTING OFF CACHE BOOL "" FORCE)
21
31
  FetchContent_Declare(
22
32
  Eigen3
23
33
  URL https://gitlab.com/libeigen/eigen/-/archive/3.4.0/eigen-3.4.0.tar.gz
24
34
  URL_HASH SHA256=8586084f71f9bde545ee7fa6d00288b264a2b7ac3607b974e54d13e7162c1c72
25
35
  DOWNLOAD_EXTRACT_TIMESTAMP TRUE
26
36
  )
27
- # Disable Eigen tests and demos to speed up configuration
28
- set(EIGEN_BUILD_TESTING OFF CACHE BOOL "" FORCE)
29
- set(EIGEN_BUILD_DOC OFF CACHE BOOL "" FORCE)
30
37
  FetchContent_MakeAvailable(Eigen3)
31
38
  endif()
32
39
 
@@ -52,6 +59,7 @@ set(HAPC_CORE_SOURCES
52
59
  src/cv_fast_pchal_python.cpp
53
60
  src/single_pcghal_cpp.cpp
54
61
  src/pcghal_cv_cpp.cpp
62
+ src/pcghal_cv_classi_cpp.cpp
55
63
  )
56
64
 
57
65
  # Python module
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: hapc
3
- Version: 0.2.1
3
+ Version: 2.1.0
4
4
  Summary: Highly Adaptive Principal Components
5
5
  Home-page: https://github.com/meixide/hapc
6
6
  Author: Carlos García Meixide
@@ -21,8 +21,7 @@ Requires-Python: >=3.8
21
21
  Description-Content-Type: text/markdown
22
22
  License-File: LICENSE
23
23
  Requires-Dist: numpy<2.3,>=1.24
24
- Requires-Dist: scipy>=1.7
25
- Requires-Dist: scikit-learn>=0.24
24
+ Requires-Dist: scikit-learn>=1.0
26
25
  Provides-Extra: dev
27
26
  Requires-Dist: pytest; extra == "dev"
28
27
  Requires-Dist: pytest-cov; extra == "dev"
@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
4
4
 
5
5
  [project]
6
6
  name = "hapc"
7
- version = "0.2.1"
7
+ version = "2.1.0"
8
8
  description = "Highly Adaptive Principal Components"
9
9
  readme = "README.md"
10
10
  requires-python = ">=3.8"
@@ -23,8 +23,7 @@ classifiers = [
23
23
  ]
24
24
  dependencies = [
25
25
  "numpy>=1.24,<2.3",
26
- "scipy>=1.7",
27
- "scikit-learn>=0.24",
26
+ "scikit-learn>=1.0",
28
27
  ]
29
28
 
30
29
  [project.optional-dependencies]
@@ -0,0 +1,99 @@
1
+ """HAPC: Highly Adaptive Principal Components.
2
+
3
+ Public API
4
+ ----------
5
+ High-level entry points (mirror the R package):
6
+
7
+ - :func:`hapc` — single-λ fit (gaussian / binomial; norm in {"sv","1","2"}).
8
+ - :func:`cv_hapc` — k-fold cross-validated fit.
9
+
10
+ Lower-level building blocks:
11
+
12
+ - :func:`design_hapc`, :func:`kernel_hapc`, :func:`cross_kernel_hapc`
13
+ - :func:`ridge_regression`, :func:`fast_pchal`
14
+ - :func:`pcghal`, :func:`pcghal_classification`, :func:`pc_hal_classi`
15
+ - :func:`single_pcghal`, :func:`single_lambda_fit`,
16
+ :func:`single_pcghal_classification`,
17
+ :func:`single_pcghal_classification_ridge_only`
18
+ - :func:`pcghal_cv`, :func:`pcghal_cv_classi`, :func:`fasthal_cv`
19
+ - :func:`ate_hapc` — ATE estimate + Wald CI via HAPC + outcome undersmoothing.
20
+ """
21
+
22
+ __version__ = "2.1.0"
23
+
24
+ from .core import (
25
+ DesignOutput,
26
+ OptimizerOutput,
27
+ cross_kernel_hapc,
28
+ design_hapc,
29
+ fast_pchal,
30
+ kernel_cross, # alias for cross_kernel_hapc (backward compat)
31
+ kernel_hapc,
32
+ mkernel, # alias for kernel_hapc (backward compat)
33
+ pchal_design, # alias for design_hapc (backward compat)
34
+ pcghal,
35
+ pcghal_classification,
36
+ pc_hal_classi,
37
+ ridge_regression,
38
+ )
39
+ from .single import (
40
+ SingleLambdaResult,
41
+ SinglePcghalClassificationResult,
42
+ SinglePcghalResult,
43
+ hapc,
44
+ single_lambda_fit,
45
+ single_pcghal,
46
+ single_pcghal_classification,
47
+ single_pcghal_classification_lasso,
48
+ single_pcghal_classification_ridge_only,
49
+ )
50
+ from .cv import (
51
+ CVResult,
52
+ cv_hapc,
53
+ fasthal_cv,
54
+ pcghal_cv,
55
+ pcghal_cv_classi,
56
+ pcghal_cv_classi_lasso,
57
+ )
58
+ from .ate import ATEResult, ate_hapc
59
+
60
+ __all__ = [
61
+ "__version__",
62
+ # high level
63
+ "hapc",
64
+ "cv_hapc",
65
+ "ate_hapc",
66
+ # design & kernels
67
+ "design_hapc",
68
+ "kernel_hapc",
69
+ "cross_kernel_hapc",
70
+ # solvers
71
+ "ridge_regression",
72
+ "fast_pchal",
73
+ "pcghal",
74
+ "pcghal_classification",
75
+ "pc_hal_classi",
76
+ # single-λ
77
+ "single_pcghal",
78
+ "single_lambda_fit",
79
+ "single_pcghal_classification",
80
+ "single_pcghal_classification_ridge_only",
81
+ "single_pcghal_classification_lasso",
82
+ # CV
83
+ "pcghal_cv",
84
+ "pcghal_cv_classi",
85
+ "pcghal_cv_classi_lasso",
86
+ "fasthal_cv",
87
+ # result types
88
+ "ATEResult",
89
+ "CVResult",
90
+ "DesignOutput",
91
+ "OptimizerOutput",
92
+ "SingleLambdaResult",
93
+ "SinglePcghalResult",
94
+ "SinglePcghalClassificationResult",
95
+ # backward-compat aliases
96
+ "pchal_design",
97
+ "mkernel",
98
+ "kernel_cross",
99
+ ]
@@ -0,0 +1,425 @@
1
+ """Average Treatment Effect estimation with HAPC + undersmoothing.
2
+
3
+ Provides :func:`ate_hapc`, a high-level convenience wrapper that:
4
+
5
+ 1. Cross-validates the **propensity** model (binomial, ``A ~ W``) on a
6
+ log-spaced grid
7
+ ``(log_lambda_prop_min, log_lambda_prop_max, grid_length_prop)``
8
+ and the **outcome** model (gaussian, ``Y ~ (A, W)``) on a separate grid
9
+ ``(log_lambda_out_min, log_lambda_out_max, grid_length_out)``
10
+ (each built like :func:`hapc.cv_hapc`).
11
+ 2. Fixes the propensity score at its CV-best λ.
12
+ 3. Computes σ = std of the ATE efficient influence function (EIF) at the
13
+ CV configuration ``(π̂_CV, μ̂_CV)``.
14
+ 4. Sweeps the **outcome** λ grid in **decreasing**
15
+ order (most smoothing → least smoothing) and stops at the first λ for
16
+ which ``|mean(EIF)| ≤ σ / (√n · log n)``. This is the **undersmoothed**
17
+ outcome model. If no λ in the grid meets the threshold, the smallest λ
18
+ is used.
19
+ 5. Returns a **doubly robust** ATE point estimate at the undersmoothed outcome
20
+ model and a ``(1 - alpha)`` Wald confidence interval from the EIF evaluated
21
+ at that estimate (see Notes).
22
+
23
+ The function does not implement sample splitting / cross-fitting:
24
+ nuisances are fit on the full sample and the EIF is evaluated on the same
25
+ sample. Bias control is provided by the undersmoothing step instead.
26
+ """
27
+
28
+ from typing import NamedTuple, Optional
29
+
30
+ import numpy as np
31
+
32
+ try:
33
+ from scipy.stats import norm as _normal
34
+ except ImportError as _e: # pragma: no cover
35
+ raise ImportError(
36
+ "scipy is required for ate_hapc (used for normal quantiles). "
37
+ "It ships transitively with scikit-learn; run `pip install scipy`."
38
+ ) from _e
39
+
40
+ from .cv import CVResult, cv_hapc
41
+ from .single import hapc as _hapc
42
+
43
+
44
+ class ATEResult(NamedTuple):
45
+ """Output of :func:`ate_hapc`.
46
+
47
+ Attributes
48
+ ----------
49
+ estimate : float
50
+ Doubly robust (AIPW-style) ATE at the undersmoothed outcome model:
51
+ ``mean(A/π̂·(Y-μ̂₁)+μ̂₁ - (1-A)/(1-π̂)·(Y-μ̂₀) - μ̂₀)``, matching the
52
+ efficient influence function used for the Wald interval (see Notes).
53
+ lower : float
54
+ Lower endpoint of the ``(1 - alpha)`` Wald confidence interval.
55
+ upper : float
56
+ Upper endpoint of the ``(1 - alpha)`` Wald confidence interval.
57
+ """
58
+
59
+ estimate: float
60
+ lower: float
61
+ upper: float
62
+
63
+
64
+ def _plot_ate_diagnostics(
65
+ cv_prop: CVResult,
66
+ cv_out: CVResult,
67
+ traj_lambdas: np.ndarray,
68
+ traj_abs_mean_eif: np.ndarray,
69
+ lam_prop_cv: float,
70
+ lam_out_cv: float,
71
+ lam_undersmooth: float,
72
+ threshold: float,
73
+ ) -> None:
74
+ """Raise ImportError if matplotlib is missing; otherwise show diagnostic figures."""
75
+ try:
76
+ import matplotlib.pyplot as plt
77
+ except ImportError as e: # pragma: no cover
78
+ raise ImportError(
79
+ "plot_diagnostics=True requires matplotlib. "
80
+ "Install with: pip install matplotlib"
81
+ ) from e
82
+
83
+ fig = plt.figure(figsize=(11.0, 7.5))
84
+ gs = fig.add_gridspec(2, 2, height_ratios=[1.0, 1.1], hspace=0.35, wspace=0.3)
85
+ ax_prop = fig.add_subplot(gs[0, 0])
86
+ ax_out = fig.add_subplot(gs[0, 1])
87
+ ax_traj = fig.add_subplot(gs[1, :])
88
+
89
+ lp = np.asarray(cv_prop.lambdas, dtype=float)
90
+ sp = np.asarray(cv_prop.mses, dtype=float)
91
+ lo = np.asarray(cv_out.lambdas, dtype=float)
92
+ so = np.asarray(cv_out.mses, dtype=float)
93
+
94
+ ax_prop.semilogx(lp, sp, "o-", color="C1", lw=1.5, ms=5)
95
+ ax_prop.axvline(lam_prop_cv, color="C3", ls="--", lw=1.5,
96
+ label=f"CV λ = {lam_prop_cv:.4g}")
97
+ ax_prop.set_xlabel("λ (propensity)")
98
+ ax_prop.set_ylabel("Mean CV logistic deviance")
99
+ ax_prop.set_title("Propensity CV (A ~ W, binomial)")
100
+ ax_prop.legend(loc="best", fontsize=8)
101
+ ax_prop.grid(True, alpha=0.3)
102
+
103
+ ax_out.semilogx(lo, so, "o-", color="C2", lw=1.5, ms=5)
104
+ ax_out.axvline(lam_out_cv, color="C3", ls="--", lw=1.5,
105
+ label=f"CV λ = {lam_out_cv:.4g}")
106
+ ax_out.set_xlabel("λ (outcome)")
107
+ ax_out.set_ylabel("Mean CV MSE")
108
+ ax_out.set_title("Outcome CV (Y ~ (A,W), gaussian)")
109
+ ax_out.legend(loc="best", fontsize=8)
110
+ ax_out.grid(True, alpha=0.3)
111
+
112
+ tv = np.asarray(traj_lambdas, dtype=float)
113
+ yv = np.asarray(traj_abs_mean_eif, dtype=float)
114
+ ok = np.isfinite(tv) & np.isfinite(yv) & (tv > 0)
115
+ tv, yv = tv[ok], yv[ok]
116
+ order = np.argsort(tv)
117
+ tv, yv = tv[order], yv[order]
118
+
119
+ if tv.size:
120
+ ax_traj.semilogx(tv, yv, "o-", color="C0", lw=2, ms=6,
121
+ label=r"$|\mathrm{mean}(\mathrm{EIF}_{\mathrm{ATE}})|$")
122
+ ax_traj.fill_between(tv, 0, threshold, alpha=0.12, color="gray")
123
+ else:
124
+ ax_traj.text(
125
+ 0.5, 0.5, "No valid outcome fits on λ grid",
126
+ transform=ax_traj.transAxes, ha="center", va="center",
127
+ )
128
+ ax_traj.axhline(threshold, color="gray", lw=2, alpha=0.85,
129
+ label=r"Threshold $\sigma_{\mathrm{CV}}/(\sqrt{n}\log n)$")
130
+ ax_traj.axvline(lam_out_cv, color="C3", ls="--", lw=1.8,
131
+ label=f"Outcome CV λ = {lam_out_cv:.4g}")
132
+ ax_traj.axvline(lam_undersmooth, color="C4", ls="-", lw=2.0,
133
+ label=f"Undersmoothed λ = {lam_undersmooth:.4g}")
134
+ ax_traj.set_xlabel("Outcome λ (undersmoothing grid)")
135
+ ax_traj.set_ylabel(r"$|\mathrm{mean}(\mathrm{EIF})|$")
136
+ ax_traj.set_title("Undersmoothing trajectory (fixed propensity at its CV-λ)")
137
+ ax_traj.legend(loc="best", fontsize=9, ncol=2)
138
+ ax_traj.grid(True, alpha=0.3)
139
+
140
+ fig.suptitle("ate_hapc diagnostics", fontsize=12, y=0.98)
141
+ fig.subplots_adjust(top=0.92, bottom=0.08, hspace=0.4, wspace=0.3)
142
+ plt.show()
143
+
144
+
145
+ def _coerce_binary(A: np.ndarray) -> np.ndarray:
146
+ """Return ``A`` re-encoded as floats in ``{0,1}``.
147
+
148
+ Accepts ``{0,1}``, ``{-1,+1}`` (or any pair where one value is non-positive
149
+ and one positive — falls back to the sign).
150
+ """
151
+ A = np.asarray(A).ravel()
152
+ u = set(np.unique(A).tolist())
153
+ if u.issubset({0, 1, 0.0, 1.0}):
154
+ return A.astype(np.float64)
155
+ if u.issubset({-1, 1, -1.0, 1.0}):
156
+ return ((A > 0).astype(np.float64))
157
+ raise ValueError(
158
+ f"A must be binary in {{0,1}} or {{-1,+1}}; found {sorted(u)}"
159
+ )
160
+
161
+
162
+ def ate_hapc(X: np.ndarray, Y: np.ndarray, A: np.ndarray,
163
+ alpha: float = 0.05,
164
+ max_degree: int = 1,
165
+ npcs: Optional[int] = None,
166
+ log_lambda_prop_min: float = -5,
167
+ log_lambda_prop_max: float = -3,
168
+ grid_length_prop: int = 10,
169
+ log_lambda_out_min: float = -5,
170
+ log_lambda_out_max: float = -3,
171
+ grid_length_out: int = 10,
172
+ nfolds: int = 5,
173
+ norm: str = "sv",
174
+ predict: Optional[np.ndarray] = None,
175
+ max_iter: int = 5000,
176
+ tol: float = 1e-3,
177
+ step_factor: float = 0.8,
178
+ verbose: bool = False,
179
+ crit: str = "grad",
180
+ center: bool = True,
181
+ approx: bool = False,
182
+ ini: str = "1",
183
+ plot_diagnostics: bool = False) -> ATEResult:
184
+ """ATE estimate with HAPC nuisances and outcome undersmoothing.
185
+
186
+ Parameters
187
+ ----------
188
+ X : np.ndarray, shape (n, p)
189
+ Covariate matrix ``W`` (do NOT include the treatment column).
190
+ Y : np.ndarray, shape (n,)
191
+ Continuous outcome.
192
+ A : np.ndarray, shape (n,)
193
+ Binary treatment in ``{0,1}`` or ``{-1,+1}``.
194
+ alpha : float, default 0.05
195
+ Significance level. The returned interval has confidence
196
+ ``1 - alpha``.
197
+ max_degree, npcs, nfolds, norm, predict, max_iter, tol, step_factor,\
198
+ verbose, crit, center, approx, ini :
199
+ Same meaning and defaults as in :func:`hapc.cv_hapc` (except λ grids,
200
+ see below).
201
+ ``predict`` is accepted for signature parity with :func:`cv_hapc` and
202
+ is currently ignored (``ate_hapc`` always evaluates the EIF on the
203
+ training sample).
204
+ log_lambda_prop_min, log_lambda_prop_max, grid_length_prop :
205
+ Equally spaced log-λ grid for **propensity** cross-validation
206
+ (``A ~ W``, binomial), same rule as :func:`cv_hapc`.
207
+ log_lambda_out_min, log_lambda_out_max, grid_length_out :
208
+ Log-λ grid for **outcome** cross-validation ``Y ~ (A, W)`` (gaussian)
209
+ and for the **undersmoothing** scan (same points, evaluated in
210
+ decreasing λ order until ``|mean(EIF)| ≤ τ``).
211
+ plot_diagnostics : bool, default False
212
+ If True, open a matplotlib figure with (1) propensity CV curve
213
+ (logistic deviance vs λ), (2) outcome CV curve (MSE vs λ), and (3)
214
+ the undersmoothing path: ``|mean(EIF)|`` vs outcome λ with the
215
+ threshold line and vertical markers for the CV and selected
216
+ undersmoothed λ. Requires ``matplotlib`` (``pip install matplotlib``).
217
+
218
+ Returns
219
+ -------
220
+ ATEResult
221
+ Named tuple with three fields ``(estimate, lower, upper)``.
222
+
223
+ Notes
224
+ -----
225
+ The procedure is:
226
+
227
+ 1. Cross-validate the propensity ``A ~ W`` (binomial) on its grid and the
228
+ outcome ``Y ~ (A, W)`` (gaussian) on the outcome grid (independently
229
+ specified).
230
+ 2. Fix the propensity at its CV-best λ; refit on the full sample to
231
+ obtain ``π̂(W_i) = P(A=1 | W_i)``.
232
+ 3. At the CV-best outcome λ, compute a **plugin-centered** influence vector
233
+ (same mean as the DR EIF at :math:`\\psi=\\overline{\\mu}_1-\\overline{\\mu}_0`)
234
+ and let ``σ = std(·)``.
235
+ 4. Threshold ``τ = σ / (√n · log n)``.
236
+ 5. Walk the **outcome** λ grid in **decreasing**
237
+ order; pick the first (largest) λ for which
238
+ ``|mean(EIF_diff)| ≤ τ`` — call it ``λ_u``.
239
+ 6. **Doubly robust** point estimate (same nuisances ``(π̂, μ̂₁, μ̂₀)``):
240
+ ``ψ̂ = mean(A/π̂·(Y-μ̂₁)+μ̂₁ - (1-A)/(1-π̂)·(Y-μ̂₀) - μ̂₀)``.
241
+ One-step influence function (centered at ``ψ̂``):
242
+ ``φ_i = A_i/π̂_i·(Y_i-μ̂_{1i}) + μ̂_{1i} - (1-A_i)/(1-π̂_i)·(Y_i-μ̂_{0i})
243
+ - μ̂_{0i} - ψ̂``.
244
+ CI: ``ψ̂ ± z_{1-α/2} · std(φ) / √n``.
245
+
246
+ This contrasts with **plug-in** G-computation ``mean(μ̂₁(W)-μ̂₀(W))``,
247
+ which can be materially biased when both nuisances are estimated on the
248
+ same sample and the outcome regressions are regularized. The DR
249
+ ``ψ̂`` is consistent if **either** the propensity **or** the pair
250
+ ``(μ̂₁, μ̂₀)`` is correctly specified (standard double robustness).
251
+
252
+ Examples
253
+ --------
254
+ >>> import numpy as np
255
+ >>> from hapc import ate_hapc
256
+ >>> rng = np.random.default_rng(0)
257
+ >>> n = 200
258
+ >>> W = np.column_stack([rng.uniform(-2, 2, n), rng.normal(0, 0.5, n)])
259
+ >>> p = 1.0 / (1.0 + np.exp(-(W[:, 0] + 0.5 * W[:, 1])))
260
+ >>> A = rng.binomial(1, p, n)
261
+ >>> Y = 2 * W[:, 0] + 0.5 + rng.normal(0, 0.5, n) # truth: ATE=0
262
+ >>> res = ate_hapc(W, Y, A, alpha=0.05, max_degree=2, npcs=50,
263
+ ... grid_length_prop=4, grid_length_out=4, nfolds=3,
264
+ ... norm="2")
265
+ >>> bool(res.lower <= res.estimate <= res.upper)
266
+ True
267
+ """
268
+ if not (0.0 < alpha < 1.0):
269
+ raise ValueError(f"alpha must be in (0,1); got {alpha}")
270
+
271
+ # --- Coerce inputs ------------------------------------------------------
272
+ X = np.ascontiguousarray(np.asarray(X, dtype=np.float64))
273
+ if X.ndim != 2:
274
+ raise ValueError(f"X must be 2-D; got shape {X.shape}")
275
+ Y = np.asarray(Y, dtype=np.float64).ravel()
276
+ A01 = _coerce_binary(A)
277
+ n, _p = X.shape
278
+ if Y.size != n or A01.size != n:
279
+ raise ValueError("X, Y, A must all have the same number of rows.")
280
+
281
+ if npcs is None:
282
+ npcs = int(n)
283
+
284
+ lambdas_out = np.exp(
285
+ np.linspace(log_lambda_out_min, log_lambda_out_max, grid_length_out))
286
+
287
+ cv_kwargs_base = dict(
288
+ max_degree=max_degree, npcs=npcs, nfolds=nfolds, norm=norm,
289
+ max_iter=max_iter, tol=tol, step_factor=step_factor,
290
+ verbose=verbose, crit=crit, center=center, approx=approx, ini=ini,
291
+ )
292
+
293
+ # --- 1. CV propensity (binomial) ---------------------------------------
294
+ cv_prop = cv_hapc(
295
+ X, A01, family="binomial",
296
+ log_lambda_min=log_lambda_prop_min,
297
+ log_lambda_max=log_lambda_prop_max,
298
+ grid_length=grid_length_prop,
299
+ **cv_kwargs_base,
300
+ )
301
+ lam_prop_cv = float(cv_prop.best_lambda)
302
+
303
+ # Refit propensity at CV λ on full data, predict in-sample probabilities.
304
+ prop = _hapc(
305
+ X, A01, family="binomial", max_degree=max_degree, npcs=npcs,
306
+ lambda_=lam_prop_cv, norm=norm, predict=X,
307
+ max_iter=max_iter, tol=tol, step_factor=step_factor,
308
+ verbose=verbose, crit=crit, center=center, approx=approx, ini=ini,
309
+ )
310
+ pi1 = np.clip(np.asarray(prop.probabilities).ravel(), 1e-8, 1 - 1e-8)
311
+
312
+ # --- 2. CV outcome (gaussian on [A,W]) ---------------------------------
313
+ Xout = np.column_stack([A01, X])
314
+ cv_out = cv_hapc(
315
+ Xout, Y, family="gaussian",
316
+ log_lambda_min=log_lambda_out_min,
317
+ log_lambda_max=log_lambda_out_max,
318
+ grid_length=grid_length_out,
319
+ **cv_kwargs_base,
320
+ )
321
+ lam_out_cv = float(cv_out.best_lambda)
322
+
323
+ # Stacked design for one-shot prediction at both arms.
324
+ Xmu1 = np.column_stack([np.ones(n), X])
325
+ Xmu0 = np.column_stack([np.zeros(n), X])
326
+ Xeval = np.vstack([Xmu1, Xmu0])
327
+
328
+ def _mu_pair(lam: float):
329
+ """Refit outcome at λ on full data, return (μ̂_1, μ̂_0) on training W."""
330
+ res = _hapc(
331
+ Xout, Y, family="gaussian", max_degree=max_degree, npcs=npcs,
332
+ lambda_=float(lam), norm=norm, predict=Xeval,
333
+ max_iter=max_iter, tol=tol, step_factor=step_factor,
334
+ verbose=verbose, crit=crit, center=center, approx=approx, ini=ini,
335
+ )
336
+ p = np.asarray(res.predictions).ravel()
337
+ if p.size != 2 * n:
338
+ raise RuntimeError(
339
+ f"Outcome predict returned {p.size} values, expected {2 * n}."
340
+ )
341
+ return p[:n], p[n:]
342
+
343
+ def _eif_plugin_centered(mu1: np.ndarray, mu0: np.ndarray) -> np.ndarray:
344
+ """Plugin-centered influence vector (undersmoothing gate only).
345
+
346
+ Its mean matches the DR EIF evaluated at plug-in
347
+ :math:`\\psi=\\overline{\\mu}_1-\\overline{\\mu}_0`. The returned ATE
348
+ uses ``_psi_dr`` / ``_eif_dr`` instead.
349
+ """
350
+ eif1 = (A01 / pi1) * (Y - mu1) - (mu1 - mu1.mean())
351
+ eif0 = ((1.0 - A01) / (1.0 - pi1)) * (Y - mu0) - (mu0 - mu0.mean())
352
+ return eif1 - eif0
353
+
354
+ def _psi_dr(mu1: np.ndarray, mu0: np.ndarray) -> float:
355
+ return float(
356
+ np.mean(
357
+ (A01 / pi1) * (Y - mu1)
358
+ + mu1
359
+ - ((1.0 - A01) / (1.0 - pi1)) * (Y - mu0)
360
+ - mu0
361
+ )
362
+ )
363
+
364
+ def _eif_dr(mu1: np.ndarray, mu0: np.ndarray, psi: float) -> np.ndarray:
365
+ return (
366
+ (A01 / pi1) * (Y - mu1)
367
+ + mu1
368
+ - ((1.0 - A01) / (1.0 - pi1)) * (Y - mu0)
369
+ - mu0
370
+ - psi
371
+ )
372
+
373
+ # --- 3. σ at CV configuration → threshold τ ----------------------------
374
+ mu1_cv, mu0_cv = _mu_pair(lam_out_cv)
375
+ eif_cv = _eif_plugin_centered(mu1_cv, mu0_cv)
376
+ sigma_cv = float(np.std(eif_cv, ddof=0))
377
+ threshold = sigma_cv / (np.sqrt(n) * np.log(n))
378
+
379
+ # --- 4. Undersmoothing sweep: largest λ → smallest --------------------
380
+ lam_und: Optional[float] = None
381
+ mu1_und = mu0_und = None
382
+ for lam in np.sort(lambdas_out)[::-1]:
383
+ try:
384
+ mu1, mu0 = _mu_pair(float(lam))
385
+ except Exception:
386
+ continue
387
+ eif = _eif_plugin_centered(mu1, mu0)
388
+ if abs(eif.mean()) <= threshold:
389
+ lam_und = float(lam)
390
+ mu1_und, mu0_und = mu1, mu0
391
+ break
392
+
393
+ if lam_und is None:
394
+ # Threshold never met → fall back to the smallest λ in the grid.
395
+ lam_und = float(lambdas_out.min())
396
+ mu1_und, mu0_und = _mu_pair(lam_und)
397
+
398
+ if plot_diagnostics:
399
+ t_lams: list[float] = []
400
+ t_abs: list[float] = []
401
+ for lam in np.sort(lambdas_out):
402
+ try:
403
+ mu1, mu0 = _mu_pair(float(lam))
404
+ except Exception:
405
+ continue
406
+ eif = _eif_plugin_centered(mu1, mu0)
407
+ t_lams.append(float(lam))
408
+ t_abs.append(float(np.abs(eif.mean())))
409
+ _plot_ate_diagnostics(
410
+ cv_prop, cv_out,
411
+ np.asarray(t_lams), np.asarray(t_abs),
412
+ lam_prop_cv, lam_out_cv, lam_und, threshold,
413
+ )
414
+
415
+ # --- 5. Doubly robust point estimate + (1 - alpha) Wald CI --------------
416
+ psi = _psi_dr(mu1_und, mu0_und)
417
+ eif_dr = _eif_dr(mu1_und, mu0_und, psi)
418
+ sigma_und = float(np.std(eif_dr, ddof=0))
419
+ z = float(_normal.ppf(1.0 - alpha / 2.0))
420
+ half = z * sigma_und / np.sqrt(n)
421
+
422
+ return ATEResult(estimate=psi, lower=psi - half, upper=psi + half)
423
+
424
+
425
+ __all__ = ["ATEResult", "ate_hapc"]