\noindent Note: when doing simulations or draws from a random distribution in \texttt{R}, please set your seed using \texttt{set.seed(2018)}.

# EM for clustering

The EM algorithm can be seen an unsupervised clustering method based on mixture models. It follows an iterative approach, sub-optimal, which tries to find the parameters of the probability distribution that has the maximum likelihood of its attributes in the presence of missing/latent data.

The algorithm's input are the data set $X$, the total number of clusters/models $K$, the accepted error to converge $\epsilon$ and the maximum number of iterations.

For each iteration, first it is executed what’s called the Expectation Step (E-step), that estimates the probability of each point belonging to each model, followed by the Maximization step (M-step), that re-estimates the parameter vector of the probability distribution of each model. The algorithm finishes when the distribution parameters converges or reach the maximum number of iterations. Convergence is assured since the algorithm increases the likelihood at each iteration until it reaches the (eventually local) maximum.

We generate a dataset such that observations are drawn from one of two linear processes $C_1$ and $C_2$:

```{r}
set.seed(2018)

slope1 <- -.3; intercept1 <- 1.5    # generating data from C_1
xs1 <- sample(seq(-2,2,len=201), 40)
ys1 <- intercept1 + slope1*xs1 + rnorm(length(xs1),0,.15) # add some noise

slope2 <- 1.2; intercept2 <- -.4    # generating data from C_2
xs2 <- sample(seq(-2,2,len=201), 40)
ys2 <- intercept2 + slope2*xs2 + rnorm(length(xs1),0,.15)

mydata <- rbind( cbind(xs1,ys1), cbind(xs2,ys2) ) 
plot(mydata, pch=19, xlab="X", ylab="Y")
```

You need to do two things,

1) Find the parameters (slope and intercept) of the two models
2) Assignment of each datapoint $x_i$ to the correct DGP

Knowing one makes the other easy to find. If we know the parameters, we can apply MLE to find which model is more probable to have generated $x_i$, ie compute $p(x_i|C_j)$. If we know the classifications we could just perform two linear regressions to find their respective parameters.

Recall the basic structure of the EM algorithm:

1) Init step: Assign random values to the model's parameters.
2) E step: Assign points to the model that fits each one best (these assignments are continuous, not binary)
3) M step: Update the parameters using the points assigned in the previous step
4) Iterate until parameter values converge (e.g., Frobenius norm of old parameter vector - new parameters vector < tolerance) or it reaches the maximum number of iterations.

First, initialize your parameters by running the following codes:

```{r}
i1 <- s1 <- i2 <- s2 <- 0 # model parameters for slope and intersect
init_params <- function() {
  i1 <<- 2*runif(1)
  s1 <<- 2*runif(1)
  i2 <<- 2*runif(1)
  s2 <<- 2*runif(1)
  c(i1,s1,i2,s2)
}

params <- init_params()
```

## Question 2: E-step

In this question, we will write a function for E-step. Read the following carefully and write a function called `e.step`.

Assume the parameters are given for each model at this step. For each point $x_i$, we can compute two weights $w_{i1}$ and $w_{i2}$ (defined below) for the soft assignments of model 1 and model 2. Let $j \in \{1,2\}$ index model, $\alpha_j$ denote the intercept and $\beta_j$ denote the slope. The residuals $r_{ij}$ of the two models for each $x_i$ is:
\[
r_{ij}= y_i - \alpha_j - \beta_j x_i
\]
The weights $w_{ij}$ is defined as:
\[
w_{ij}=\frac{p(r_{ij}|x_i\in C_j)}{p(r_{i1}|x_i\in C_1)+p(r_{i2}|x_i\in C_2)}
\]
Assume that normal distributions are assigned to these probabilities:
\[
r_{ij}|x_i \in C_j \sim \mathcal{N}(0,\sigma^2)
\]
which means
\[
p(r_{ij}|x_i\in C_j)=\frac{1}{\sqrt{2\pi\sigma}}\exp{\{-\frac{r_{ij}^2}{2\sigma^2}\}}
\]
Plugging in distributions into $w_{ij}$ formula yields:
\[
w_{ij}=\frac{\exp\{-\frac{r_{ij}^2}{2\sigma^2}\}}{\exp\{-\frac{r_{i1}^2}{2\sigma^2}\}+\exp\{-\frac{r_{i2}^2}{2\sigma^2}\}}
\]
Note that $\sigma$ is a free parameter and corresponds to the expected noise in the sample. Assume it is equal to 0.5 for our purposes here.

Write a function `e.step` that takes data and parameters as inputs and returns $w_{i1}$ and $w_{i2}$ values.


```{r}
# params is [s1,i1,s2,i2]
e.step <- function(mydata, params, sigma=0.5) {
  w1 <- rep(NA, nrow(mydata))
  w2 <- rep(NA, nrow(mydata))
  
  #TO DO: write a for loop to calculate w1 and w2
  for (i in 1:nrow(mydata)) {
    
  }
  
  cbind(w1,w2)  
}

# TO DO: run a sample E-step
ws <- e.step(mydata, params)
head(ws)
```

## Question 3: M-step

Next, we will write a function `wls` for M-step.

Assume the weights are given for each data point $x_i$ at this step. To estimate the parameters of each process use weighted least squares. Recall that the OLS estimator of the parameters $\theta$ is given by:
\[
\hat{\theta} = (X^\top X)^{-1}X^\top Y
\]
To estimate $\theta_j$ for the $j$-th model, we add weights $w_j$ to the estimator. Specifically, make a diagonal matrix $W_j$ with these values and compute:
\[
\hat{\theta}_j = (X^\top W_jX)^{-1}X^\top W_j Y
\]

Write a function `wls` that takes $X$, $Y$ and $W$ as inputs and solves for the weighted least squares $\hat{\theta}_j$. Write a function `m.step` that returns the $\hat{\theta}_j$ for both models and has two arguments: the data, and the $w_{i1}$ and $w_{i2}$ output from the E-step.


```{r}
# wls - weighted least squares
wls <- function(X,Y,W) {
  #TO DO: produce analytic solutions for wls
}

m.step <- function(mydata, ws) {
  X <- cbind(rep(1, nrow(mydata)), mydata[,1])
  Y <- as.matrix(mydata[,2], ncol=1)
  
  #TO DO: solve updated parameters 
  
}

#TO DO: run a sample M-step
params <- m.step(mydata, ws)
params
```

## Question 4: EM together

Put together the E and M steps into an EM algorithm wrapped in a function called `em.2lines`. The function takes three arguments: your data, the tolerance level (set to `1e-2` as a default), and the maximum number of iterations (set to `1e3` for now as a default). Return the estimated parameters, the final weights for each data point $x_i$, and the predicted class label for each data point (e.g., if $w_{i1}>w_{i2}$ then the class label of $i$-th data point is $1$ and $0$ otherwise).

Plot the results and color the datapoints according to their predicted class; draw appropriately colored linear regressions through the two sets of data using your EM estimated slopes and intercepts.

```{r}
em.2lines <- function(mydata, tol=1e-2, max.step=1e3) {
  step <- 0
  
  s1 <- i1 <- s2 <- i2 <- 0 # model parameters for slope and intersect
  params <- init_params()
  
  #TO DO: write an iterative function to conduct EM. You should save the final estimations for the parameters, the weights for each datapoint x^i, and the classification for each datapoint
  
  
}

```
