$"""$Mixture model for matrix completion"""

from typing import Tuple

import numpy as np

from scipy.special import logsumexp

from common import GaussianMixture

def estep$($X: np$.$ndarray, mixture: GaussianMixture$)->$ Tuple$[$np$.$ndarray, float$]$:

$"""$E$-$step: Softly assigns each datapoint to a gaussian component

Args:

X: $($n$,$ d$)$ array holding the data, with incomplete entries $($set to $0)$

mixture: the current gaussian mixture

Returns:

np$.$ndarray: $($n$,$ K$)$ array holding the soft counts

for all components for all examples

float: log$-$likelihood of the assignment

$"""$

raise NotImplementedError

def mstep$($X: np$.$ndarray, post: np$.$ndarray, mixture: GaussianMixture,

min$\_$variance: float $=\mathrm{.}25)->$ GaussianMixture:

$"""$M$-$step: Updates the gaussian mixture by maximizing the log$-$likelihood

of the weighted dataset

Args:

X: $($n$,$ d$)$ array holding the data, with incomplete entries $($set to $0)$

post: $($n$,$ K$)$ array holding the soft counts

for all components for all examples

mixture: the current gaussian mixture

min$\_$variance: the minimum variance for each gaussian

Returns:

GaussianMixture: the new gaussian mixture

$"""$

raise NotImplementedError

def run$($X: np$.$ndarray, mixture: GaussianMixture,

post: np$.$ndarray$)->$ Tuple$[$GaussianMixture$,$ np$.$ndarray, float$]$:

$"""$Runs the mixture model

Args:

X: $($n$,$ d$)$ array holding the data

post: $($n$,$ K$)$ array holding the soft counts

for all components for all examples

Returns:

GaussianMixture: the new gaussian mixture

np$.$ndarray: $($n$,$ K$)$ array holding the soft counts

for all components for all examples

float: log$-$likelihood of the current assignment

$"""$

raise NotImplementedError

def fill$\_$matrix$($X: np$.$ndarray, mixture: GaussianMixture$)->$ np$.$ndarray:

$"""$Fills an incomplete matrix according to a mixture model

Args:

X: $($n$,$ d$)$ array of incomplete data $($incomplete entries $=0)$

mixture: a mixture of gaussians

Returns

np$.$ndarray: a $($n$,$ d$)$ array with completed data

$"""$

raise NotImplementedError