---
title: "Multi-category Treatments with `personalized`"
author: "Jared Huling"
date: "`r Sys.Date()`"
output: 
    rmarkdown::html_vignette:
        fig_width: 7
        fig_height: 5
        toc: true
        toc_depth: 3
        number_sections: true
        self_contained: true
preamble: >
    \usepackage{amsmath,amssymb,amsfonts,bm,bbm,amsthm}
    \usepackage{longtable}
    \usepackage{booktabs}
    \usepackage{bbm, bm}
    \def\bsx {\boldsymbol x}
    \def\bsX {\boldsymbol X}
    \def\bsT {\boldsymbol T}
    \def\bsbeta {\boldsymbol \beta}
    \def\bsgamma {\boldsymbol \gamma}
    \def\bsW {\boldsymbol W}
    \def\bsy {\boldsymbol y}
    \def\bsY {\boldsymbol Y}
    \def\bsM {\boldsymbol M}
    \def\bfx {\mathbf x}
    \def\bfX {\mathbf X}
    \def\bfT {\mathbf T}
    \def\bfW {\mathbf W}
    \def\bfy {\mathbf y}
    \def\bfY {\mathbf Y}
    \def\bfM {\mathbf M}
    \def\bfU {\mathbf U}
vignette: >
  %\VignetteIndexEntry{Multi-category Treatments with `personalized`}
  %\VignetteEngine{knitr::rmarkdown}
---

# Example with multi-category treatments

To demonstrate how to use the `personalized` package with multi-category treatments, we simulate data with three treatments. The treatment assignments will be based on covariates and hence mimic an observational setting with no unmeasured confounders.

```{r loadpkg, message=FALSE, warning=FALSE}
library(personalized)
```

```{r sim_three_trt_data}
set.seed(123)
n.obs  <- 250
n.vars <- 10
x <- matrix(rnorm(n.obs * n.vars, sd = 3), n.obs, n.vars)

# simulated non-randomized treatment with multiple levels
# based off of a multinomial logistic model
xbetat_1 <- 0.1 + 0.5 * x[,1]  - 0.25 * x[,5]
xbetat_2 <- 0.1 - 0.5 * x[,9] + 0.25 * x[,5]
trt.1.prob <- exp(xbetat_1) / (1 + exp(xbetat_1) + exp(xbetat_2))
trt.2.prob <- exp(xbetat_2) / (1 + exp(xbetat_1) + exp(xbetat_2))
trt.3.prob <- 1 - (trt.1.prob + trt.2.prob)

prob.mat <- cbind(trt.1.prob, trt.2.prob, trt.3.prob)
trt.mat <- apply(prob.mat, 1, function(rr) rmultinom(1, 1, prob = rr))
trt.num <- apply(trt.mat, 2, function(rr) which(rr == 1))
trt <- as.factor(paste0("Trt_", trt.num))

# simulate response

# effect of treatment 1 relative to treatment 3
delta1 <- 2 * (0.5 + x[,2] - 2 * x[,3]  )
# effect of treatment 2 relative to treatment 3
delta2 <- (0.5 + x[,6] - 2 * x[,5] )

# main covariate effects with nonlinearities
xbeta <- x[,1] + x[,9] - 2 * x[,4]^2 + x[,4] + 
    0.5 * x[,5] ^ 2 + 2 * x[,2] - 3 * x[,5]

# create entire functional form of E(Y|T,X)
xbeta <- xbeta + 
    delta1 * ((trt.num == 1) - (trt.num == 3) ) + 
    delta2 * ((trt.num == 2) - (trt.num == 3) )


# simulate continuous outcomes E(Y|T,X)
y <- xbeta + rnorm(n.obs, sd = 2)

```

We will use the \code{factor} version of the treatment status vector in our analysis, however, the integer values vector, i.e. \code{trt.num}, could be used as well.

```{r display_mult_trt_vector}
trt[1:5]
table(trt)
```

Then we construct a propensity score function that takes covariate information and the treatment statuses as input and generate a matrix of probabilities as output. Each row $i$ of the output matrix represents an observation and each column $j$ is the probability that the $i$th patient received the $j$th treatment. The treatment levels are ordered alphabetically (or numerically if the treatment assignment vector is a vector of integers). Our propensity score model in this example will be a multinomial logistic regression model with a lasso penalty for the probability of treatment assignments conditional on covariate information:

```{r define_multi_propens}
propensity.multinom.lasso <- function(x, trt)
{
    if (!is.factor(trt)) trt <- as.factor(trt)
    gfit <- cv.glmnet(y = trt, x = x, family = "multinomial",
                      nfolds = 3)

    # predict returns a matrix of probabilities:
    # one column for each treatment level
    propens <- drop(predict(gfit, newx = x, 
        type = "response", s = "lambda.min"))

    # return the matrix probability of treatment assignments
    probs <- propens[,match(levels(trt), colnames(propens))]

    probs
}
```

An important assumption for the propensity score is that $0 < Pr(T_i = t | \bfX) < 1$ for all $\bfX$ and $t$. This assumption, often called the positivity assumption, is impossible to verify. However, in practice validity of the assumption can be assessed via a visualization of the empirical overlap of our estimated propensity scores to determine if there is any evidence of positivity violations. The \code{check.overlap()} function also allows us to visualize the overlap of our propensity scores for multi-category treatment applications. The following code results in the plot shown Figure \ref{fig:check_overlap_multitreat}.

```{r check_overlap_multitreat, fig.cap = "Propensity score overlap plot for multi-category treatment data."}
check.overlap(x = x, trt = trt, propensity.multinom.lasso)
```

Each plot above is for a different treatment group, e.g. the plot in the first row of plots is the subset of patients who received treatment 1. There seems to be no obvious evidence against the positivity assumption. 

As the outcome is continuous and there is a large number of covariates available for our construction of a benefit score, we will use the squared error loss and a lasso penalty. The model can be fit in the same manner as for the binary treatment setting, however only linear models and the weighting method are available. Here we can also specify the reference treatment (the treatment that the non-reference treatments are compared with by each benefit score).

```{r fit_multi_trt_model}
set.seed(123)
subgrp.multi <- fit.subgroup(x = x, y = y,
    trt = trt, propensity.func = propensity.multinom.lasso,
    reference.trt = "Trt_3",
    loss   = "sq_loss_lasso",
    nfolds = 3)

summary(subgrp.multi)
```

The \code{summary()} function now displays selected variables for each of the two benefit scores and shows the quantiles of each benefit score. We can also plot the empirical observations within the different subgroups using the \code{plot()} function, however now it is slightly more complicated. It appears that the average outcome is higher for those who received the level of the treatment they were recommended than those who received a different treatment than they were recommended. Also note that \code{plot.subgroup_fitted()} returns a \code{ggplot} object \citep{ggplot2} and can thus be modified by the user. The below example yields Figure \ref{fig:plot_multi_trt_model}.

```{r plot_multi_trt_model, warning=FALSE, message=FALSE, fig.width=5, fig.cap = "Individual outcome observations by treatment group and subgroup."}
pl <- plot(subgrp.multi)
pl + theme(axis.text.x = element_text(angle = 90, hjust = 1))
```

To obtain valid estimates of the subgroup-specific treatment effects, we perform the repeated training and testing resample procedure using the \code{validate.subgroup()} function (here we set the number of replications `B=4`, but this should be a much larger number in practice) :

```{r validate_multi_trt_model, eval = TRUE}
set.seed(123)
validation.multi <- validate.subgroup(subgrp.multi, 
    B = 4,  # specify the number of replications
    method = "training_test_replication",
    train.fraction = 0.5)

print(validation.multi, digits = 2, sample.pct = TRUE)
```

Setting the \code{sample.pct} argument above to \code{TRUE} prints out the average percent of all patients which are in each subgroup (as opposed to the average sample sizes). We can see that about 58\% of patients were recommended treatment 1 and among those recommended treatment 1, we expect them to have larger outcomes if they actually receive treatment 1 as opposed to the other treatments. The estimated effects are positive within all three subgroups (meaning those recommended each of the different treatments have a positive benefit from receiving the treatment they are recommended as opposed to receiving any another treatments).

We can visualize the subgroup-specific treatment effects using \code{plot()} as usual with results shown in Figure \ref{fig:plotcomparemultivalidated}:
```{r plotcomparemultivalidated, warning=FALSE, message=FALSE, eval = TRUE, fig.width=5, fig.cap = "Validation results for multi-category treatment data."}
plv <- plot(validation.multi)
plv + theme(axis.text.x = element_text(angle = 90, hjust = 1))
```

# More details on propensity scores for multi-category treatments

For cases with multi-category treatments, the user must specify a propensity function that returns $Pr(T = T_i | \bfX = \bfx)$ for patient $i$. In other words, it should return the probability of receiving the treatment that was actually received for each patient. For example:

```{r multinom_propens}
propensity.func.multinom <- function(x, trt)
{
    df <- data.frame(trt = trt, x)
    mfit <- nnet::multinom(trt ~ . -trt, data = df)
    # predict returns a matrix of probabilities:
    # one column for each treatment level
    propens <- nnet::predict.nnet(mfit, type = "probs")

    if (is.factor(trt))
    {
        values <- levels(trt)[trt]
    } else 
    {
        values <- trt
    }
    # return the probability corresponding to the
    # treatment that was observed
    probs <- propens[cbind(1:nrow(propens), 
        match(values, colnames(propens)))]
    probs
}
```

Optionally the user can specify the function to return a matrix of treatment probabilities, however, the columns *must* be ordered by the levels of \code{trt}. An example of this is the following:
```{r multinom_propens2}
propensity.func.multinom <- function(x, trt)
{
    require(nnet)
    df <- data.frame(trt = trt, x)
    mfit <- multinom(trt ~ . -trt, data = df)
    # predict returns a matrix of probabilities:
    # one column for each treatment level
    propens <- predict(mfit, type = "probs")

    if (is.factor(trt))
    {
        levels <- levels(trt)
    } else 
    {
        levels <- sort(unique(trt))
    }
    # return the probability corresponding to the
    # treatment that was observed
    probs <- propens[,match(levels, colnames(propens))]
    probs
}
```