# Lecture 5 

# Let's generate some date 
# where the variance of error terms is not constant 
n <- 100
x <- seq(1, 10, length.out = n)
beta0 <- 2
beta1 <- 3
sigma <- 1

eps <- rnorm(n, mean = 0, sd = sigma * x)
y <- beta0 + beta1 * x + eps
par(mfrow = c(1,3))
plot(x, y)
fit <- lm(y ~ x)
abline(fit)
plot(fit, 1)
plot(fit, 2)

# Consider the square root transformation for y
par(mfrow = c(1,3))
plot(x, sqrt(y))
fit2 <- lm(I(sqrt(y)) ~ x)
abline(fit2)
plot(fit2, 1)
plot(fit2, 2)

# GLS ---- estimator 
# Generate some data 
set.seed(3)
n <- 100
x <- seq(1, 10, length.out = n)
beta_1 <- 2; beta_2 <- 3; sigma2 <- 1
omega <- x^2                  # known diagonal elements of Omega

y <- beta_1 + beta_2 * x + rnorm(n, 0, sqrt(sigma2 * omega))

plot(x, y)
abline(c(beta_1, beta_2), col = "red")

# Negative log-likelihood function 
nllik_gls <- function(par, y, x, omega) {
  beta_1 <- par[1]
  beta_2 <- par[2]
  sigma2 <- par[3]
  mu <- beta_1 + beta_2 * x
  n <- length(y)
  n/2 * log(sigma2) + 1/(2*sigma2) * sum((y - mu)^2 / omega)
}

#fit the model using nlminb
fit_gls <- nlminb(start = c(beta_1 = 0, beta_2 = 0, sigma2 = 1),
                  objective = nllik_gls,
                  y = y, x = x, omega = omega,
                  lower = c(-Inf, -Inf, 1e-9))
fit_gls$par

# Compare 
# 1) with the analytical solution

X <- model.matrix(y ~ x) 
Omega <- diag(omega)
Omega_inv <- solve(Omega) # 1/diag(Omega)

beta_gls <- solve(t(X)%*%Omega_inv%*%X) %*%t(X) %*% Omega_inv %*% y


# 2) the solution obtained using lm by specifying the weights
help(lm)


fit_wls <- lm(y ~ x, weights = 1/omega) 
summary(lm_gls)

# 3) the function gls in the nlme package 
library(nlme)
help(gls)

fit_nlme <- gls(y ~ x, weights = varFunc( ~ omega))


# Comparison 
cbind(nlminb      = fit_gls$par[1:2],
      closed_form = drop(beta_gls),
      lm_weights  = coef(fit_wls),
      gls_nlme   = coef(fit_nlme),
      OLS         = coef(lm(y~x)))


cbind(fit_gls$par[3], summary(fit_nlme)$sigma^2)


B <- 1000
est <- matrix(NA, B, 2)
for(i in 1:B){
  y <- beta_1 + beta_2 * x + rnorm(n, 0, sqrt(sigma2 * omega))
  est[i, ] <-   c(OLS = coef(lm(y ~ x))[[2]],
                  GLS = coef(lm(y ~ x, weights = 1/omega))[[2]])
} 

par(mfrow=c(1,1))
boxplot(est)
abline(h = beta_2)