finish code for normal_pdf below

You should be able to write your own function based on the following formula, and you are **NOT allowed** to use outside resource packages other than those we provided.

**(for undergrads only) normalPDF**

Using the covariance matrix as a diagonal matrix with variances of the individual variables appearing on the main diagonal of the matrix and zeros everywhere else means that we assume the features are independent. In this case, the multivariate normal density function simplifies to the expression below:
$$\mathcal{N}(x: \mu, \Sigma) = \prod_{i=1}^D \frac{1}{\sqrt{2\pi\sigma_i^2}}\exp{\left( -\frac{1}{2\sigma_i^2} (x_i-\mu_i)^2\right)}$$
where $\sigma^2_i$ is the variance for the $i^{th}$ feature, which is the diagonal element of the covariance matrix.


**Hints:**

- If you encounter "LinAlgError", you can mitigate your number/array by summing a small value before taking the operation, e.g. np.linalg.inv($\Sigma_k$ + SIGMA_CONST). You can arrest and handle such error by using [Try and Exception Block](https://realpython.com/python-exceptions/#the-try-and-except-block-handling-exceptions) in Python. Please only add SIGMA_CONST to all elements when sigma_i is not invertible.

- In the above calculation, you must avoid computing a $(N,N)$ matrix. Using the above equation for large N will crash your kernel and/or give you a memory error on Gradescope. Instead, you can do this same operation by calculating $(X-\mu)\Sigma^{-1}$, a $(N,D)$ matrix, transpose it to be a $(D,N)$ matrix and do an element-wise multiplication with $(X-\mu)^T$, which is also a $(D,N)$ matrix. Lastly, you will need to sum over the 0 axis to get a $(1,N)$ matrix before proceeding with the rest of the calculation. This uses the fact that doing an element-wise multiplication and summing over the 0 axis is the same as taking the diagonal of the $(N,N)$ matrix from the matrix multiplication.

- In Numpy implementation for each individual $\mu$, you can either use a 2-D array with dimension $(1,D)$ for each Gaussian Distribution, or a 1-D array with length $D$. Same to other array parameters. Both ways should be acceptable but pay attention to the shape mismatch problem and be **consistent all the time** when you implement such arrays.
import numpy as np
from kmeans import KMeans
from numpy.linalg import LinAlgError
from tqdm import tqdm

SIGMA_CONST = 1e-06
LOG_CONST = 1e-32
FULL_MATRIX = False


class GMM(object):

def __init__(self, X, K, max_iters=100):
"""
Args:
X: the observations/datapoints, N x D numpy array
K: number of clusters/components
max_iters: maximum number of iterations (used in EM implementation)
"""
self.points = X
self.max_iters = max_iters
self.N = self.points.shape[0]
self.D = self.points.shape[1]
self.K = K

def softmax(self, logit):
"""
Args:
logit: N x D numpy array
Return:
prob: N x D numpy array. See the above function.
Hint:
Add keepdims=True in your np.sum() function to avoid broadcast error.
"""
logit_sub = logit - logit.max(axis=1, keepdims=True)
exp_logit = np.exp(logit_sub)
sum_exp_logit = np.sum(exp_logit, axis=1, keepdims=True)
return exp_logit / sum_exp_logit

def logsumexp(self, logit):
"""
Args:
logit: N x D numpy array
Return:
s: N x 1 array where s[i,0] = logsumexp(logit[i,:]). See the above function
Hint:
The keepdims parameter could be handy
"""
logit_max = np.max(logit, axis=1, keepdims=True)
return logit_max + np.log(
np.sum(np.exp(logit - logit_max), axis=1, keepdims=True)
)

def normalPDF(self, points, mu_i, sigma_i):
"""
Args:
points: N x D numpy array
mu_i: (D,) numpy array, the center for the ith gaussian.
sigma_i: DxD numpy array, the covariance matrix of the ith gaussian.
Return:
pdf: (N,) numpy array, the probability density value of N data for the ith gaussian

Hint:
np.diagonal() should be handy.
"""

def multinormalPDF(self, points, mu_i, sigma_i):
"""
Args:
points: N x D numpy array
mu_i: (D,) numpy array, the center for the ith gaussian.
sigma_i: DxD numpy array, the covariance matrix of the ith gaussian.
Return:
normal_pdf: (N,) numpy array, the probability density value of N data for the ith gaussian

Hint:
1. np.linalg.det() and np.linalg.inv() should be handy.
2. Note the value in self.D may be outdated and not correspond to the current dataset.
3. You may wanna check if the matrix is singular before implementing calculation process.
"""
raise NotImplementedError

def create_pi(self):
"""
Initialize the prior probabilities
Args:
Return:
pi: numpy array of length K, prior
"""
return np.ones(self.K) / self.K

def create_mu(self):
"""
Intialize random centers for each gaussian
Args:
Return:
mu: KxD numpy array, the center for each gaussian.
"""
random_indices = np.random.choice(self.N, self.K, replace=False)
mu = self.points[random_indices, :]
return mu

def create_sigma(self):
"""
Initialize the covariance matrix with np.eye() for each k. For grads, you can also initialize the
by K diagonal matrices.
Args:
Return:
sigma: KxDxD numpy array, the diagonal standard deviation of each gaussian.
You will have KxDxD numpy array for full covariance matrix case
"""
return np.repeat(np.eye(self.D)[np.newaxis, :, :], self.K, axis=0)

def _init_components(self, **kwargs):
"""
Args:
kwargs: any other arguments you want
Return:
pi: numpy array of length K, prior
mu: KxD numpy array, the center for each gaussian.
sigma: KxDxD numpy array, the diagonal standard deviation of each gaussian.
You will have KxDxD numpy array for full covariance matrix case

Hint: np.random.seed(5) must be used at the start of this function to ensure consistent outputs.
"""
np.random.seed(5)
pi = self.create_pi()
mu = self.create_mu()
sigma = self.create_sigma()
return pi, mu, sigma

def _ll_joint(self, pi, mu, sigma, full_matrix=FULL_MATRIX, **kwargs):
"""
Args:
pi: np array of length K, the prior of each component
mu: KxD numpy array, the center for each gaussian.
sigma: KxDxD numpy array, the diagonal standard deviation of each gaussian. You will have KxDxD numpy
array for full covariance matrix case
full_matrix: whether we use full covariance matrix in Normal PDF or not. Default is True.

Return:
ll(log-likelihood): NxK array, where ll(i, k) = log pi(k) + log NormalPDF(points_i | mu[k], sigma[k])
"""
N, D = self.points.shape
ll = np.zeros((N, self.K))
for k in range(self.K):
if not full_matrix:
pdf = self.normalPDF(self.points, mu[k], sigma[k])
ll[:, k] = np.log(pi[k] + LOG_CONST) + np.log(pdf + LOG_CONST)
return ll

def _E_step(self, pi, mu, sigma, full_matrix=FULL_MATRIX, **kwargs):
"""
Args:
pi: np array of length K, the prior of each component
mu: KxD numpy array, the center for each gaussian.
sigma: KxDxD numpy array, the diagonal standard deviation of each gaussian.You will have KxDxD numpy
array for full covariance matrix case
full_matrix: whether we use full covariance matrix in Normal PDF or not. Default is True.
Return:
tau: NxK array, the posterior distribution (a.k.a, the soft cluster assignment) for each observation.

Hint:
You should be able to do this with just a few lines of code by using _ll_joint() and softmax() defined above.
"""
ll_joint = self._ll_joint(pi, mu, sigma, **kwargs)
tau = self.softmax(ll_joint)
return tau

def _M_step(self, tau, full_matrix=FULL_MATRIX, **kwargs):
"""
Args:
tau: NxK array, the posterior distribution (a.k.a, the soft cluster assignment) for each observation.
full_matrix: whether we use full covariance matrix in Normal PDF or not. Default is True.
Return:
pi: np array of length K, the prior of each component
mu: KxD numpy array, the center for each gaussian.
sigma: KxDxD numpy array, the diagonal standard deviation of each gaussian. You will have KxDxD numpy
array for full covariance matrix case

Hint:
There are formulas in the slides and in the Jupyter Notebook.
Undergrads: To simplify your calculation in sigma, make sure to only take the diagonal terms in your covariance matrix
"""
N_k = np.sum(tau, axis=0)
pi = N_k / self.N
mu = np.dot(tau.T, self.points) / N_k[:, np.newaxis]
sigma = np.zeros((self.K, self.D, self.D))

for k in range(self.K):
diff = self.points - mu[k]
sigma_k = np.dot(tau[:, k] * diff.T, diff) / N_k[k]
sigma[k] = sigma_k + np.eye(self.D) * SIGMA_CONST

return pi, mu, sigma

def __call__(
self, full_matrix=FULL_MATRIX, abs_tol=1e-16, rel_tol=1e-16, **kwargs
): # No need to change
"""
Args:
abs_tol: convergence criteria w.r.t absolute change of loss
rel_tol: convergence criteria w.r.t relative change of loss
kwargs: any additional arguments you want

Return:
tau: NxK array, the posterior distribution (a.k.a, the soft cluster assignment) for each observation.
(pi, mu, sigma): (1xK np array, KxD numpy array, KxDxD numpy array)

Hint:
You do not need to change it. For each iteration, we process E and M steps, then update the paramters.
"""
pi, mu, sigma = self._init_components(**kwargs)
pbar = tqdm(range(self.max_iters))

prev_loss = None
for it in pbar:
# E-step
tau = self._E_step(pi, mu, sigma, full_matrix)

# M-step
pi, mu, sigma = self._M_step(tau, full_matrix)

# calculate the negative log-likelihood of observation
joint_ll = self._ll_joint(pi, mu, sigma, full_matrix)
loss = -np.sum(self.logsumexp(joint_ll))
if it:
diff = np.abs(prev_loss - loss)
if diff < abs_tol and diff / prev_loss < rel_tol:
break
prev_loss = loss
pbar.set_description("iter %d, loss: %.4f" % (it, loss))
return tau, (pi, mu, sigma)


def cluster_pixels_gmm(image, K, max_iters=10, full_matrix=True):
"""
Clusters pixels in the input image

Each pixel can be considered as a separate data point (of length 3),
which you can then cluster using GMM. Then, process the outputs into
the shape of the original image, where each pixel is its most likely value.

Args:
image: input image of shape(H, W, 3)
K: number of components
max_iters: maximum number of iterations in GMM. Default is 10
full_matrix: whether we use full covariance matrix in Normal PDF or not. Default is True.
Return:
clustered_img: image of shape(H, W, 3) after pixel clustering

Hints:
What do mu and tau represent?
"""
raise NotImplementedError