---
title: "Getting Started with psAve"
author: "Daijiro Kabata"
output: rmarkdown::html_vignette
vignette: >
  %\VignetteIndexEntry{Getting Started with psAve}
  %\VignetteEngine{knitr::rmarkdown}
  %\VignetteEncoding{UTF-8}
---

```{r setup, include = FALSE}
knitr::opts_chunk$set(
  collapse = TRUE,
  comment = "#>",
  eval = requireNamespace("MatchIt", quietly = TRUE) &&
    requireNamespace("ranger", quietly = TRUE) &&
    requireNamespace("rpart", quietly = TRUE) &&
    requireNamespace("xgboost", quietly = TRUE)
)
```

## Introduction

Propensity score (PS) analyses are only as good as the propensity score model, and in practice the analyst rarely knows which model is right. Logistic regression may miss nonlinearities and interactions; flexible machine learning methods (classification trees, random forests, gradient boosting) capture them but can produce extreme scores and unstable weights. Different candidate models often yield propensity scores that all "look reasonable" yet lead to materially different effect estimates — a form of model dependence that undermines the credibility of the analysis. Rather than committing to a single model, `psAve` constructs a *model-averaged* propensity score: a convex combination $\bar e(X) = \sum_m \lambda_m \hat e_m(X)$ of candidate propensity scores, with the mixing weights $\lambda$ chosen on a simplex grid to optimize a balance criterion, implementing the method of Kabata, Stuart & Shintani (2024).

The distinguishing feature of the method is *what* the mixing weights are asked to balance. Covariate balance criteria treat all covariates as equally important, but for bias in the treatment effect what matters is balance on covariates *as they relate to the outcome*. The prognostic score — the predicted outcome under the untreated condition, $g(0, X)$, estimated from untreated units only (Hansen 2008) — summarizes exactly that relationship, and prognostic-score balance has been shown to be a useful diagnostic for propensity score methods (Stuart, Lee & Leacy 2013). `psAve` takes this one step further and uses the weighted standardized mean difference of a (model-averaged) prognostic score as the *selection criterion* for the propensity score mixing weights. In the simulations of Kabata, Stuart & Shintani (2024), this "Prog (Ave)" strategy gave the lowest and most robust bias and RMSE across 16 scenarios, compared with single-model propensity scores and with model averaging targeted at prediction accuracy or covariate balance. The result of `psave()` is deliberately modest: a numeric vector of propensity scores, designed to be handed to `MatchIt::matchit()` as a distance measure or to `WeightIt::weightit()` as a propensity score, with balance assessment via `cobalt`.

## Installation

`psAve` can be installed from GitHub:

```{r, eval = FALSE}
# install.packages("remotes")
remotes::install_github("kabajiro/psAve")
```

The core of the package requires only `cobalt` (an Import). The candidate learners beyond logistic regression (`rpart`, `ranger`, `xgboost`), as well as `MatchIt` and `WeightIt` for the downstream analysis, are Suggests and are only needed when you actually request them.

## A matching workflow with the `lalonde` data

We illustrate the paper's headline estimator — the model-averaged propensity score selected by balance on the model-averaged prognostic score, "Prog (Ave)" — on the `lalonde` dataset that ships with `MatchIt`. The estimand is the average treatment effect in the treated (ATT), the package default.

```{r, message = FALSE}
library(psAve)
data("lalonde", package = "MatchIt")
```

### Step 1: Estimate the model-averaged propensity score

`psave()` looks deliberately like `matchit()` and `weightit()`: a treatment formula and a data frame. The one addition is the `outcome` argument, which names the outcome variable used to build the prognostic score. A one-sided formula (`~ re78`) reuses the covariates on the right-hand side of `formula` as the prognostic-model predictors; a two-sided formula (`re78 ~ age + educ + ...`) lets you specify a distinct prognostic model.

Two of the default candidate learners (`ranger` and `xgboost`) are stochastic, so set a seed first for reproducibility:

```{r}
set.seed(1234)
fit <- psave(treat ~ age + educ + race + married + nodegree + re74 + re75,
             data = lalonde, outcome = ~ re78)
```

By default this fits four candidate propensity score models (`"glm"`, `"rpart"`, `"ranger"`, `"xgboost"`) and four candidate prognostic models with the same learners, then searches all 1,771 points of the mixing-weight simplex (step 0.05, four candidates). The whole call takes on the order of seconds to a minute on the `lalonde` data (n = 614), with `xgboost` dominating the runtime; for a quick first pass you can restrict `ps.methods` (e.g., `ps.methods = c("glm", "rpart")`).

```{r}
fit
```

The printout shows the estimand and criterion, the selected mixing weights $\lambda$ (for the propensity score candidates) and $\gamma$ (for the prognostic candidates), the achieved criterion value, a short balance preview, and — importantly — the literal next call you would issue to carry the score into `MatchIt`.

### Step 2: Match on the averaged propensity score

`fit$ps` is a plain numeric vector, named by the rownames of `data`, and can be passed directly to `MatchIt::matchit()` as the `distance` argument:

```{r, eval = FALSE}
m <- MatchIt::matchit(treat ~ age + educ + race + married + nodegree + re74 + re75,
                      data = lalonde, distance = fit$ps,
                      method = "nearest", caliper = .2)
```

Retyping the formula and the data name creates an opportunity for row misalignment if the two calls do not use literally the same data. The convenience wrapper `psave_match()` removes that hazard by reusing the formula and data stored in the `psave` object; all other arguments are forwarded verbatim to `matchit()`, and the result is an ordinary `matchit` object:

```{r}
m <- psave_match(fit, method = "nearest", caliper = .2)
m
```

### Step 3: Assess balance, including prognostic-score balance

Because the object returned by `psave_match()` is a genuine `matchit` object, the full `cobalt` toolkit applies. Supplying the model-averaged prognostic score through `distance` adds a prognostic-balance row to the balance table — the diagnostic recommended by Stuart, Lee & Leacy (2013):

```{r}
cobalt::bal.tab(m, distance = data.frame(prog = fit$prog))
```

You can also call `cobalt::bal.tab()` directly on the `psave` object itself, which assesses balance for the implied inverse-probability weights at the fitted estimand and automatically includes both the averaged propensity score and the prognostic score as distance measures:

```{r}
cobalt::bal.tab(fit)
```

### Step 4: Estimate the treatment effect

Effect estimation after matching is deliberately *not* part of `psAve` — the matched object is a standard `matchit` object, so all established guidance applies unchanged. We recommend following `MatchIt`'s vignette on estimating effects after matching (`vignette("estimating-effects", package = "MatchIt")`), which uses the `marginaleffects` package to compute the ATT with cluster-robust standard errors on the matched sample. For a weighting analysis instead of matching, including the exact IPW estimator used in the paper, see `vignette("weighting", package = "psAve")`.

## Interpreting the output

### `print()`

`print(fit)` is a one-screen orientation: the estimand and selection criterion, the mixing weights rendered as labeled text bars (so you can see at a glance which candidate models contribute to the average), the criterion value at the selected $\lambda$, a preview of the worst-balanced covariates plus the prognostic score, and the literal next call. A candidate receiving weight 0 was judged not to improve prognostic-score balance; that is informative, not a failure.

### `summary()`

```{r}
summary(fit)
```

`summary(fit)` adds three pieces:

1. **Mixing-weight tables** for $\lambda$ and $\gamma$. The $\gamma$ weights tell you which learners the untreated-outcome model relied on (selected by unweighted mean squared error among untreated units).
2. **The diagnostics table**: for every single candidate propensity score *and* for the selected average, all four criteria are reported (log loss, mean weighted ASMD of covariates, mean weighted KS of covariates, and weighted ASMD of the prognostic score). This is the "was averaging worth it?" table — you can verify directly that the averaged score achieves better prognostic-score balance than any single candidate, and see what it trades away (typically a little log loss: the averaged score is usually a *worse* predictor of treatment than the most flexible single learner, by design).
3. **The full balance table**: unweighted versus weighted SMD and KS statistics for every covariate and for the prognostic score, with the conventional 0.1 threshold marked.

### `plot()`

```{r, fig.width = 7, fig.height = 5}
plot(fit, type = "balance")
```

`type = "balance"` draws a Love plot (via `cobalt::love.plot()`) of covariate and prognostic-score balance before and after weighting by the implied weights.

```{r, fig.width = 7, fig.height = 5}
plot(fit, type = "distribution")
```

`type = "distribution"` shows the distribution of the propensity scores by treatment group — each grey curve is one candidate model, and the colored curve is the selected average. This is the plot to inspect for extreme candidate scores: flexible learners fit in-sample can push scores toward 0 or 1 (they are clipped to `[0.01, 0.99]` by default), and you can see how the averaging tempers them.

```{r, fig.width = 7, fig.height = 5}
plot(fit, type = "criterion")
```

`type = "criterion"` displays the selection criterion over the $\lambda$ grid (exactly for up to three candidates; as one profile per candidate otherwise), so you can judge how sharply the criterion identifies the selected mixture.

## FAQ

### Doesn't using the outcome bias my analysis?

This is the natural first objection, and the design of the method answers it. Two facts matter:

1. **The prognostic model never sees treated outcomes.** Following Hansen (2008), the prognostic score $g(0, X)$ is estimated *only on untreated units* and then predicted for everyone. The $\gamma$ mixing weights are likewise selected using prediction error among untreated units only. No model in the pipeline uses the treated-arm outcomes.
2. **The selection criterion never contrasts outcomes between arms.** The $\lambda$ criterion is a *balance* statistic — the weighted standardized mean difference of the prognostic score between treatment groups. It measures whether the two groups look comparable on an outcome-relevant summary of the covariates; it does not measure, and cannot be driven toward, any particular treatment effect estimate. At no point is a treated-versus-untreated outcome comparison computed during design.

This preserves the separation between the *design* stage and the *analysis* stage (Rubin 2001; Hansen 2008): everything `psave()` does is a function of the covariates, the treatment indicator, and the untreated units' outcomes, exactly like the prognostic-score balance diagnostics advocated by Stuart, Lee & Leacy (2013). It is the same reasoning under which `MatchIt`'s documentation welcomes externally estimated distance measures so long as no post-treatment information about the treated-untreated outcome contrast enters the design. A fuller discussion, with the exact formulas, is in `vignette("method-details", package = "psAve")`.

## 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

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