---
title: "Method details and design decisions"
author: "Daijiro Kabata"
output: rmarkdown::html_vignette
vignette: >
  %\VignetteIndexEntry{Method details and design decisions}
  %\VignetteEngine{knitr::rmarkdown}
  %\VignetteEncoding{UTF-8}
---

```{r setup, include = FALSE}
knitr::opts_chunk$set(
  collapse = TRUE,
  comment = "#>"
)
```

This vignette documents the method implemented by `psave()` at the level of formulas, records the small number of places where the package deliberately deviates from the paper's reference code (with justification), and situates `psAve` among related software. Throughout, "the paper" is Kabata, Stuart & Shintani (2024).

## Notation

For subjects $i = 1, \dots, n$: covariates $X_i$ (a $p$-column numeric matrix after full dummy expansion of factors), binary treatment $A_i \in \{0, 1\}$, and outcome $Y_i$. The propensity score is $e(X_i) = \Pr(A_i = 1 \mid X_i)$. The prognostic score is $g(0, X_i)$, the predicted outcome under the untreated condition, estimated from untreated units only (Hansen 2008).

`psave()` fits $M$ candidate propensity score models $\hat e_1, \dots, \hat e_M$ (on all $n$ units) and $K$ candidate prognostic models $\hat g_1, \dots, \hat g_K$ (on the $n_0$ untreated units, predicted for all $n$), then selects convex mixing weights $\gamma$ (prognostic) and $\lambda$ (propensity) on a simplex grid.

## The simplex grid and tie-breaking

Both $\gamma$ and $\lambda$ range over the probability simplex discretized with increment `step` (default 0.05). The exported function `simplex_grid(M, step)` enumerates the grid in **integer arithmetic**: with $S = 1/\text{step}$ steps, it lists every integer composition $(c_1, \dots, c_M)$ with $\sum_m c_m = S$ and $c_m \ge 0$, returned as $c/S$. The number of grid points is exactly $\binom{S + M - 1}{M - 1}$:

```{r}
library(psAve)

nrow(simplex_grid(2, step = 0.05))   # choose(21, 1) = 21
nrow(simplex_grid(3, step = 0.05))   # choose(22, 2) = 231
nrow(simplex_grid(4, step = 0.05))   # choose(23, 3) = 1771
```

The enumeration order is a documented part of the method because it defines the tie-breaking rule. The grid is generated by recursive descent, with $c_1$ running from $S$ down to 0, then $c_2$ descending over the remainder, and so on:

```{r}
head(simplex_grid(3, step = 0.25))
tail(simplex_grid(3, step = 0.25), 3)
```

The first row puts all weight on the first candidate; the last row puts all weight on the last. **Ties in any grid search are resolved by taking the first row attaining the minimum, within a 1e-9 relative tolerance of the minimum**, so ties favor learners listed earlier in `ps.methods` / `prog.methods`. This makes "first minimum" a reproducible formula rather than an accident of grid construction. The tolerance is deliberate: the criterion values are computed with floating-point matrix algebra whose lowest-order bits can differ across BLAS implementations, so an exact bitwise `which.min()` would not be reproducible across machines, whereas the tolerant first-minimum rule is. (The reference code used `expand.grid()` order with `which.min()`; the package's order differs but is equally deterministic and is fixed by `simplex_grid()` by construction.)

The grid size grows quickly with $M$: with `step = 0.05`, $M = 6$ already gives `r format(choose(25, 5), big.mark = ",")` points. `psave()` warns when the grid exceeds $10^5$ rows and suggests a coarser `step`; in that case the full criterion path is also not stored.

## Step 1: Candidate models

Candidate propensity score models are fit on all $n$ units and predict $\Pr(A = 1 \mid X)$ in-sample. The built-in engines are pinned to fixed hyperparameters (chosen to mirror the `SuperLearner` wrapper defaults used in the paper, so the direct and `SL.*` routes agree); they can be overridden per learner via `control =`, and the resolved values are stored in `fit$info$learners`. Any `"SL.*"` label is passed through to `SuperLearner` verbatim, which is the exact-replication path for the paper. Alternatively, a user-supplied matrix of candidate scores can be given via `ps.matrix`.

Each candidate column is clipped to `clip = c(0.01, 0.99)` **before** averaging (as in the paper). No re-clipping is applied after averaging: a convex combination of values inside the clipping interval cannot leave it.

Candidate prognostic models use the same learner menu in regression mode (or probability mode for `family = binomial()`), fit **only on untreated units** and predicted for all $n$ units.

## Step 2: $\gamma$ — the model-averaged prognostic score

$\gamma$ minimizes the **unweighted** mean squared prediction error among untreated units:

$$
\hat\gamma = \arg\min_{\gamma \in \mathcal{S}_K} \;
\frac{1}{n_0} \sum_{i : A_i = 0} \Bigl( Y_i - \sum_{k=1}^{K} \gamma_k \, \hat g_k(X_i) \Bigr)^2,
$$

where $\mathcal{S}_K$ is the simplex grid. The model-averaged prognostic score is $\bar g(X) = \sum_k \hat\gamma_k \hat g_k(X)$. For `family = binomial()` the same formula is the Brier score; note the paper's simulations validated continuous outcomes (see Limitations). With a single prognostic candidate, $\gamma = 1$ and the grid search is skipped.

## Step 3: $\lambda$ — the model-averaged propensity score

For each grid point $\lambda$, the averaged score is $\bar e_\lambda(X_i) = \sum_m \lambda_m \hat e_m(X_i)$, and the implied inverse-probability weights are

$$
\text{ATT:}\quad W_i = \begin{cases} 1 & A_i = 1 \\ \dfrac{\bar e_\lambda(X_i)}{1 - \bar e_\lambda(X_i)} & A_i = 0 \end{cases}
\qquad\qquad
\text{ATE:}\quad W_i = \begin{cases} \dfrac{1}{\bar e_\lambda(X_i)} & A_i = 1 \\ \dfrac{1}{1 - \bar e_\lambda(X_i)} & A_i = 0. \end{cases}
$$

$\lambda$ is selected to minimize one of four criteria, evaluated at every grid point.

### `criterion = "logloss"`

The negative Bernoulli log-likelihood of treatment (prediction accuracy, the Xie et al. 2019 lineage):

$$
\mathrm{LL}(\lambda) = -\frac{1}{n} \sum_{i=1}^{n} \Bigl[ A_i \log \bar e_\lambda(X_i) + (1 - A_i) \log\bigl(1 - \bar e_\lambda(X_i)\bigr) \Bigr],
$$

finite by clipping. This is the only criterion that does not require an `outcome`.

### `criterion = "smd"`

The mean over covariate columns $j$ of the weighted absolute standardized mean difference:

$$
\mathrm{wASMD}_j(\lambda) = \frac{\Bigl| \dfrac{\sum_i A_i W_i X_{ij}}{\sum_i A_i W_i} - \dfrac{\sum_i (1 - A_i) W_i X_{ij}}{\sum_i (1 - A_i) W_i} \Bigr|}{s_j},
\qquad
s_j = \mathrm{sd}(X_j \mid A = 1),
$$

where $s_j$ is the plain (unweighted, $n - 1$ denominator) sample standard deviation **in the treated group** — for **both** the ATT and the ATE, per the paper's supplement (not the pooled SD). Under the ATT, $W_i = 1$ for treated units, so the first term is the raw treated mean.

### `criterion = "ks"`

The mean over covariates of the weighted Kolmogorov–Smirnov statistic, using the proper weighted empirical CDF in each arm:

$$
F^w_a(x) = \frac{\sum_{i : A_i = a} W_i \, \mathbf{1}(X_{ij} \le x)}{\sum_{i : A_i = a} W_i},
\qquad
\mathrm{wKS}_j(\lambda) = \max_{x} \bigl| F^w_1(x) - F^w_0(x) \bigr|
$$

over the observed values of $X_j$. For a binary covariate this reduces to the absolute difference in weighted proportions.

### `criterion = "prog"` (the default; the paper's "Prog (Ave)")

The wASMD formula above applied to a **single column**: the model-averaged prognostic score $\bar g$ (when `prog.target = "average"`, the headline recommendation) or a single candidate $\hat g_k$ (when `prog.target` names a learner — the paper's "Prog ($g_k$)" variants). The denominator is $\mathrm{sd}(\bar g \mid A = 1)$, again unweighted and for both estimands.

### Vertex mode

`average = FALSE` runs the identical machinery on the $M$ vertices of the simplex only, i.e., it selects the single best candidate propensity score by the chosen criterion (the "best single learner" variants computed alongside Table 1 of the paper's supplement).

## Criterion vs. display conventions for the SMD denominator

One subtlety deserves explicit documentation. The paper (and its reference implementation) standardizes *every* covariate — including binary ones — by the plain sample SD in the treated group. `cobalt`, by contrast, standardizes binary covariates by $\sqrt{p_1 (1 - p_1)}$ (the population formula) when it standardizes them at all, and its `bal.tab()` display leaves binary covariates as raw differences in proportions by default.

`psAve` resolves this as follows:

* **The selection criteria are paper-faithful.** The `smd` and `prog` criteria use uniform sample-SD standardization for all columns (implemented by passing `bin.vars = FALSE` for every column to `cobalt::col_w_smd()`), so the selected $\lambda$ and the reported `criterion.value` match the published method exactly. For the `ks` criterion the two conventions coincide for binary variables, so auto-detection is harmless there.
* **The display follows cobalt.** The `balance` field of the fitted object, `summary()`, `plot(type = "balance")`, and `bal.tab()` use cobalt's native conventions, so the numbers you see agree with what `cobalt::bal.tab()` reports for the same matched or weighted analysis elsewhere in your workflow.

For `criterion = "prog"` the denominator is a single positive constant across the whole $\lambda$ grid, so the *argmin* — and therefore the selected score — is invariant to this choice; only the reported criterion value depends on it.

## Five documented fixes relative to the paper's reference code

The published algorithm is implemented exactly; the published *reference code* contained implementation artifacts that this package fixes. These are documented so that anyone comparing `psAve` output against the original scripts understands where and why small discrepancies arise.

1. **Integer simplex grid.** The reference code built the grid with `expand.grid()` over $\{0, 0.05, \dots, 1\}^M$ and kept rows with `rowSums(gr) == 1` — an exact floating-point equality test. With $M = 4$ and step 0.05, this silently dropped 187 of the 1,771 valid simplex points (about 10.6%), including potentially optimal mixtures. `simplex_grid()` enumerates integer compositions, so every valid grid point is present by construction.
2. **Proper weighted-eCDF KS.** The reference `Fks` computed `ks.test()` on the covariate values *multiplied by* the weights within each arm, which is not the paper's weighted-eCDF definition (multiplying values by weights changes the distribution's support, not its weighting). The package uses the paper's definition via `cobalt::col_w_ks()`.
3. **`binomial()` family for binary responses.** The reference code fit binary models inside `SuperLearner` with `gaussian(link = "logit")`. The package uses `binomial()` throughout for treatment models (and for binary-outcome prognostic models).
4. **No train/test-inconsistent `scale()`.** The reference code applied `scale()` separately to fitting and prediction sets, standardizing them by different means and SDs. The package passes raw covariates to all engines (the pinned engines are either scale-invariant or fit and predict on identical rows).
5. **Strict complete-case handling.** The reference `Fasmd` applied `na.omit()` to a covariate vector while indexing full-length treatment and weight vectors, silently misaligning rows in the presence of missing data. `psave()` refuses incomplete data outright: any `NA` in the treatment, the covariates, or the outcome (when used) is an error, never a silent drop.

A sixth, minor note: candidate scores are clipped before averaging and never re-clipped afterward (see Step 1); the reference code's clipping constants 0.01/0.99 are exposed as the `clip` argument.

## Relation to other software

No other R package implements propensity score model averaging over a mixing-weight simplex (Xie et al. 2019 published no software). The closest relatives, and how `psAve` differs:

* **`WeightIt::method_super` with `SL.method = "method.balance"`** (the Balance Super Learner; Pirracchio & Carone 2018). Also builds a convex combination of candidate propensity models with weights chosen for balance rather than prediction. Differences: its criterion is *covariate* balance (a user-chosen `cobalt`-computed statistic), whereas `psAve`'s default criterion is *prognostic-score* balance, which targets the covariate directions that matter for outcome bias; the Balance Super Learner optimizes over the continuous simplex via `SuperLearner`'s machinery with cross-validated predictions, whereas `psAve` uses the paper's fixed grid with in-sample candidate predictions and a documented tie-break; and `method_super` lives inside a weighting workflow, whereas `psAve` returns a bare score vector equally usable for matching (`MatchIt::matchit(distance = )`) and weighting.
* **`twang`** (McCaffrey, Ridgeway & Morral 2004) tunes a *single* model class — the number of boosting iterations of a GBM — against covariate balance. It selects within one learner; `psAve` averages across heterogeneous learners.
* **`PSweight`** provides a broad estimand/weighting framework (including overlap weights) with a single user-specified propensity model per analysis; it does model *use*, not model *selection or averaging*.
* **`cobalt`** is not a competitor but the substrate: the balance statistics that define `psAve`'s criteria are computed by `cobalt::col_w_smd()` and `cobalt::col_w_ks()`, and `bal.tab()` works on `psave` objects directly.

## Limitations

Honesty about scope, matching the paper's evidence base:

* **Binary treatment only.** The method is defined, and was evaluated, for a two-arm comparison. Multi-category and continuous treatments are out of scope.
* **ATT and ATE only.** The supplement's criterion formulas are validated for these two estimands; other estimands (ATO, ATM, ...) are refused with an error rather than extrapolated.
* **Continuous outcomes in the validating simulations.** The paper's simulations used continuous $Y$. `family = binomial()` is permitted — the $\gamma$ criterion is then the Brier score, an unchanged formula — but its finite-sample performance was not studied in the paper.
* **In-sample candidate predictions.** As in the paper's worked example, candidate models predict on their own training data (no cross-fitting in v1). Flexible learners can therefore produce optimistically extreme candidate scores; clipping bounds the damage, the `diagnostics` table shows every candidate under every criterion, and `plot(type = "distribution")` makes overfit candidates visible. Balance-targeted selection (unlike log-loss selection) is not *rewarded* for overfit treatment prediction, which is part of the method's rationale.
* **Design-based standard errors downstream do not reflect score estimation.** See the weighting vignette; bootstrap the full pipeline if that matters for your application.

## References

Hansen, B. B. (2008). The prognostic analogue of the propensity score. *Biometrika*, 95(2), 481–488. doi:10.1093/biomet/asn004

Kabata, D., Stuart, E. A., & Shintani, A. (2024). Prognostic score-based model averaging approach for propensity score estimation. *BMC Medical Research Methodology*, 24, 228. doi:10.1186/s12874-024-02350-y

McCaffrey, D. F., Ridgeway, G., & Morral, A. R. (2004). Propensity score estimation with boosted regression for evaluating causal effects in observational studies. *Psychological Methods*, 9(4), 403–425. doi:10.1037/1082-989X.9.4.403

Pirracchio, R., & Carone, M. (2018). The Balance Super Learner: A robust adaptation of the Super Learner to improve estimation of the average treatment effect in the treated based on propensity score matching. *Statistical Methods in Medical Research*, 27(8), 2504–2518. doi:10.1177/0962280216682055

Stuart, E. A., Lee, B. K., & Leacy, F. P. (2013). Prognostic score-based balance measures can be a useful diagnostic for propensity score methods in comparative effectiveness research. *Journal of Clinical Epidemiology*, 66(8 Suppl), S84–S90. doi:10.1016/j.jclinepi.2013.01.013

Xie, Y., Zhu, Y., Cotton, C. A., & Wu, P. (2019). A model averaging approach for estimating propensity scores by optimizing balance. *Statistical Methods in Medical Research*, 28(1), 84–101. doi:10.1177/0962280217715487