---
title: "Linear ITR: Simulation"
author: "Andong Wang"
date: "2/22/2023"
output: html_document
---

```{r setup, include=FALSE}
knitr::opts_chunk$set(echo = TRUE)
library(tidyverse)
library(nleqslv)
library(rgenoud)

# Required functions to run beforehand
source("./TLITR Functions.R")
```

# Generate simulation study data for source population and target population.

```{r gendata}
set.seed(2023)

### population pool
N <- 50000

### Covariates (height and age)
height <- runif(N, min = 12, max = 84)
age <- runif(N, min = 0, max = 100)
X <- cbind(height, age)
X <- as.data.frame(X)

### RWD target population 
m <- 10000
rwd_id <- sample(1:N,m,replace = FALSE) # RWD is a random sample of population pool
rwd <- X[rwd_id,]
rwd <- as.data.frame(rwd)
colnames(rwd) <- c("Height", "Age")

### Treatment assignment for RWD
b1 <- .01
b2 <- 1
b3 <- 1
lps_true <- b1 + b2 * scale(rwd$Height) + b3 * scale(rwd$Age) # true linearized propensity score
ps_true <- 1 / (1 + exp(-lps_true)) # true propensity score 
A <- rbinom(m, 1, ps_true) # for observation data
rwd <- rwd %>% mutate(Assignment = A)

### RCT source population
temp <- X[-rwd_id,]
temp <- as.data.frame(temp)
colnames(temp) <- c("Height", "Age")
c1 <- -2.5
c2 <- 1
c3 <- 2
lss_true <- c1 + c2 * scale(temp$Height) + c3 * scale(temp$Age) # true linearized sampling score
ss_true <- 1 / (1 + exp(-lss_true)) # true sampling score 
S <- rbinom(N-m, 1, ss_true)
rct <- temp[S==1,]
n <- as.numeric(dim(rct)[1])

### Treatment assignment for RCT
A <- rbinom(n, 1, 0.5)
rct <- rct %>% mutate(Assignment = A)

### Potential Outcome (Y1 and Y0) -- Patients' health ratings
Y0_rct = 2*rct$Height - rct$Age + 100 + rnorm(n, 0, 1)
tau_rct = 2 * (rct$Height > 55 & rct$Age < 41) - 1
tau_rct = tau_rct* 7 # scale up tau (i.e. Individual Treatment Effect) to emphasize strong effect
Y1_rct = Y0_rct + tau_rct
Y_rct <- Y1_rct*rct$Assignment+Y0_rct*(1-rct$Assignment) # Observed outcome
rct <- rct %>% mutate(Y0 = Y0_rct, Y1= Y1_rct, Y=Y_rct)

Y0_rwd = 2*rwd$Height - rwd$Age + 100 + rnorm(m, 0, 1)
tau = 2 * (rwd$Height > 55 & rwd$Age < 41) - 1
tau = tau* 7 
Y1_rwd = Y0_rwd + tau
Y_rwd <- Y1_rwd*rwd$Assignment+Y0_rwd*(1-rwd$Assignment) 
rwd <- rwd %>% mutate(Y0 = Y0_rwd, Y1= Y1_rwd, Y=Y_rwd)
```

# Compute Calibration Weights

```{r cw}
Xd = cbind(1, rwd[,1:2])
Xd = as.matrix(Xd)
x_rct <- as.matrix(rct[,1:2])
x_rwd <- as.matrix(rwd[,1:2])

M2 <- c(mean(rct$Height), mean(rct$Age))
mu_g0 <- c(mean(rwd$Height), mean(rwd$Age))

lamFun <- function(lam, moments, moments.bar){ 
  qi <- (exp(moments%*%lam)/sum(exp(moments%*%lam)))[,1]
  colSums(qi*moments) - moments.bar
}

lam.hat <- nleqslv::searchZeros(matrix(rnorm(length(mu_g0)*10,0.5),nrow = 10),
                                lamFun, moments = x_rct, moments.bar = mu_g0)$x[1,]

wts <- exp(x_rct%*%lam.hat)/sum(exp(x_rct%*%lam.hat))
```

# Compute Unweighted ITR, Calibration-Weighted ITR, and True ITR for the Target Population

```{r itr}
## TRUE ITR for the RWD population
w = rep(1, m)
tau = 2 * (rwd$Height > 55 & rwd$Age < 41) - 1
tau = tau* 7
Xd = cbind(1, rwd[,1:2])
Xd = as.matrix(Xd)
eta1 = get_eta(Xd, w, tau)
eta1 = - eta1 / eta1[3]
eta1

## TRUE ITR for the RCT population (unweighted ITR for RWD)
w = rep(1, n)
tau = 2 * (rct$Height > 55 & rct$Age < 41) - 1
tau = tau* 7
Xd = cbind(1, rct[,1:2])
Xd = as.matrix(Xd)
eta0 = get_eta(Xd, w, tau)
eta0 =  - eta0 / eta0[3]
eta0

## Calibration-Weighted ITR for the RWD population
w = wts
Xd = cbind(1, rct[,1:2])
Xd = as.matrix(Xd)
# Estimating Tau
tau_hat <- est.contrast(a.in.rct = rct$Assignment, x.in.rct=Xd, y.in.rct = rct$Y, Qm = 1, meth = 3)
tau = tau_hat
eta_w = get_eta(Xd, w, tau)
eta_w = - eta_w / eta_w[3]
eta_w
```

# Plot to demonstrate effect of calibration-weighting

```{r format}
## Format data
rwd <- rwd%>% mutate(Population = 0)
rct <- rct%>% mutate(Population = 1)
df <- rbind(rwd, rct)
ITE = df$Height > 55 & df$Age < 41
df = cbind(df, ITE)
df$Assignment <- as.factor(df$Assignment)
df$Population <- as.factor(df$Population)
df$Assignment = factor(df$Assignment, levels = c(0, 1), labels = c("controlled", "treated"))
df$Population = factor(df$Population, levels = c(0, 1), labels = c("Target", "Source"))
df$ITE = factor(df$ITE, levels = c(F, T), labels = c("negative treatment effect", "positive treatment effect"))
df_samp <- df[sample(1:dim(df)[1],round(dim(df)[1]/7)),] # Thinning the population
```

```{r plot1}
## Create plot
lineseg = data.frame(x1 = c(55, 55), x2 = c(55, 84), y1 = c(0, 41), y2 = c(41, 41))

p1 <- ggplot(data = df_samp, aes(x = Height, y = Age)) + geom_point(aes(color = Population, pch = ITE), size = 1) +
  geom_abline(slope = eta1[2], intercept = eta1[1], color = "firebrick3", size = 1) +
  geom_abline(slope = eta0[2], intercept = eta0[1], color = "blue", size = 1) +
  geom_abline(slope = eta_w[2], intercept = eta_w[1], color = "purple3", size = 1) +
  geom_segment(aes(x = x1, y = y1, xend = x2, yend = y2), data = lineseg, color = "black", size = 1) +
  scale_shape_manual(values = c(4,15)) +
  scale_x_continuous(breaks = seq(12, 84, by = 12)) +
  scale_y_continuous(breaks = seq(0, 100, by = 20)) +
  ggtitle("Simulation: Unweighted vs. Weighted vs. True-Optimal ITR") +
  theme(
    legend.text = element_text(size = 14),
    legend.title = element_text(size = 14, face="bold"),
    axis.title.x = element_text(size = 14),
    axis.title.y = element_text(size = 14),
    plot.title = element_text(size = 14)) 
```

# Correct Classfication Rate Comparison

```{r CCR}
class <- cbind(rwd, ITE = 2*sign(rwd$Height > 55 & rwd$Age < 41) -1)

class$Score1 <- eta1[1] + eta1[2] * class$Height + eta1[3] * class$Age
class$CorrectClass1 <- (sign(class$ITE*class$Score1) + 1)/2
mean(class$CorrectClass1)

class$Score0 <- eta0[1] + eta0[2] * class$Height + eta0[3] * class$Age
class$CorrectClass0 <- (sign(class$ITE*class$Score0) + 1)/2
mean(class$CorrectClass0)

class$Score_w <- eta_w[1] + eta_w[2] * class$Height + eta_w[3] * class$Age
class$CorrectClass_w <- (sign(class$ITE*class$Score_w) + 1)/2
mean(class$CorrectClass_w)
```

# Estimate values for propensity score model component of the value function

```{r propensityscore}
prop_model <- glm(formula = Assignment ~ Age + Height, data = rct, family = binomial)
summary(prop_model)
prop_scores <- prop_model$fitted.values # checking propensity score
rct <- rct %>% mutate(PS = prop_scores)
pred_prop <- predict(prop_model, newdata = rct, type="response")
```

# Estimate values for outcome regression model component of the value function

```{r outcomeregression}
OR_model1 <-
  lm(formula = Y ~ Age + Height,
     data = rct %>% filter(Assignment == 1))

OR_model0 <-
  lm(formula = Y ~ Age + Height,
     data = rct %>% filter(Assignment == 0))

pred_OR1 <- predict(OR_model1, newdata = rct)
pred_OR0 <- predict(OR_model0, newdata = rct)
rct <- cbind(rct, pred_OR1, pred_OR0)
pred_OR <- mean(rct$pred_OR1) - mean(rct$pred_OR0)
```

# Define the value function 

```{r valuefunction}
value_weighted_AIPW <- function(eta) {
  Xd <- cbind(intercept = 1, x_rct)
  reg <- (Xd%*%eta>0)*1
  Qd <- pred_OR1 * reg + pred_OR0 * (1-reg)
  Cd <- reg == Assignment
  pid <- pred_prop * reg + p0 * (1-reg)
  value <- Cd * Y / pid - (Cd - pid) / pid * Qd
  return(sum(weights*value))
}
```

# Define global variables for value function 

```{r def-variables}
X <- x_rct
weights <- wts
Assignment <- rct$Assignment
Y <- rct$Y
p0 <- 1-pred_prop
```

# Optimization with GA

```{r GA}
n_iter <- 100
n_var <- 3 # Intercept, Height, Age
itr <- matrix(0, nrow = n_iter, ncol = n_var)
value <- rep(0, n_iter)

start_time <- Sys.time()
for (i in 1:n_iter){
  set.seed(i)
  print(paste0("Iteration: ", i, "; Seed: ", i))
  dom <- matrix(c(rep(-1,n_var),rep(1,n_var)),n_var,2)
  result_opt <- genoud(fn=value_weighted_AIPW,nvars=n_var,max=TRUE,pop.size=3000,starting.values=rep(0,n_var),Domains=dom,optim.method="Nelder-Mead")
  itr[i,] <- result_opt$par 
  value[i] <- result_opt$value
}
end_time <- Sys.time()
end_time - start_time
# seed 2, 27 yield high value 146.9721
```

# Inspect values and select the GA-optimized ITR with the highest value

```{r inspect}
iter_opt <- which(value == max(value)); iter_opt
value_opt <- value[iter_opt[1]]; value_opt
```

# Correct Classfication Rate Comparison

```{r CCR}
eta_opt <- itr[iter_opt[1],]
eta_opt <- - eta_opt / eta_opt[3]
class <- cbind(rwd, ITE = 2*sign(rwd$Height > 55 & rwd$Age < 41) -1)

class$Score0 <- eta0[1] + eta0[2] * class$Height + eta0[3] * class$Age
class$CorrectClass0 <- (sign(class$ITE*class$Score0) + 1)/2
mean(class$CorrectClass0)

class$Score_w <- eta_w[1] + eta_w[2] * class$Height + eta_w[3] * class$Age
class$CorrectClass_w <- (sign(class$ITE*class$Score_w) + 1)/2
mean(class$CorrectClass_w)

class$Score_opt <- eta_opt[1] + eta_opt[2] * class$Height + eta_opt[3] * class$Age
class$CorrectClass_opt <- (sign(class$ITE*class$Score_opt) + 1)/2
mean(class$CorrectClass_opt)

class$Score1 <- eta1[1] + eta1[2] * class$Height + eta1[3] * class$Age
class$CorrectClass1 <- (sign(class$ITE*class$Score1) + 1)/2
mean(class$CorrectClass1)
```

# Compare value function values with true value of the target population

```{r truevalue}
d0 <- class$Score0 > 0
YD0 <- d0*class$Y1 + (1-d0)* class$Y0
mean(YD0)

d_w <- class$Score_w > 0
YD_w <- d_w*class$Y1 + (1-d_w)* class$Y0
mean(YD_w)

d_opt <- class$Score_opt >0
YD_opt <- d_opt*class$Y1 + (1-d_opt)* class$Y0
mean(YD_opt)

# d1 <- class$Score1 >0
# YD1 <- d1*class$Y1 + (1-d1)* class$Y0
# mean(YD1)

d_true <- class$ITE ==1
YD_true <- d_true*class$Y1 + (1-d_true)* class$Y0
mean(YD_true)
```

# Estimate variance of value function

```{r varval}
var_value0 <- sum((YD0-mean(YD0))^2)/m^2
mean(YD0)-1.96*sqrt(var_value0)
mean(YD0)+1.96*sqrt(var_value0)

var_value_w <- sum((YD_w-mean(YD_w))^2)/m^2
mean(YD_w)-1.96*sqrt(var_value_w)
mean(YD_w)+1.96*sqrt(var_value_w)

var_value_opt <- sum((YD_opt-mean(YD_opt))^2)/m^2
mean(YD_opt)-1.96*sqrt(var_value_opt)
mean(YD_opt)+1.96*sqrt(var_value_opt)

var_value_true <- sum((YD_true-mean(YD_true))^2)/m^2
mean(YD_true)-1.96*sqrt(var_value_true)
mean(YD_true)+1.96*sqrt(var_value_true)
```

# Other performance metrics

```{r}
agree_mat <- function(d_hat, d_true) {
  TP <- sum(d_hat == 1 & d_true == 1)  # True Positives
  TN <- sum(d_hat == 0 & d_true == 0)  # True Negatives
  FP <- sum(d_hat == 1 & d_true == 0)  # False Positives
  FN <- sum(d_hat == 0 & d_true == 1)  # False Negatives
  FPR <- FP / (TN + FP)  # False positive rate
  FNR <- FN / (TP + FN)  # False negative rate

  return(list(
    FPR = FPR,
    FNR = FNR))
}
agree_mat(d0, d_true)
agree_mat(d_w, d_true)
agree_mat(d_opt, d_true)

```



# Plot to demonstrate effect of GA optimization

```{r plot2}
lineseg = data.frame(x1 = c(55, 55), x2 = c(55, 84), y1 = c(0, 41), y2 = c(41, 41))

p2 <- ggplot(data = df_samp, aes(x = Height, y = Age)) + geom_point(aes(color = Population, pch = ITE), size = 1) +
  geom_abline(slope = eta1[2], intercept = eta1[1], color = "firebrick2", size = 1) +
  geom_abline(slope = eta0[2], intercept = eta0[1], color = "blue", size = 1) +
  geom_abline(slope = eta_w[2], intercept = eta_w[1], color = "purple", size = 1) +
  geom_abline(slope = eta_opt[2], intercept = eta_opt[1], color = "green3", size = 1) +
  geom_segment(aes(x = x1, y = y1, xend = x2, yend = y2), data = lineseg, color = "black", size = 1) +
  scale_shape_manual(values = c(4,15)) +
  scale_x_continuous(breaks = seq(12, 84, by = 12)) +
  scale_y_continuous(breaks = seq(0, 100, by = 20)) +
  # ggtitle("Simulation: Unweighted vs Weighted GA-Optimal vs True-Optimal ITR") +
  theme(
    legend.text = element_text(size = 14),
    legend.title = element_text(size = 14, face="bold"),
    axis.title.x = element_text(size = 14),
    axis.title.y = element_text(size = 14),
    plot.title = element_text(size = 14)) 
```

