Instantly share code, notes, and snippets.

# rmcgibbo/wls-em.py Created Oct 30, 2012

What would you like to do?
linear regression mixture model (weighted linear regression) with E-M
 import numpy as np import matplotlib.pyplot as pp from matplotlib import cm def maximization(X, y, weights): """Calculate the beta that maximize the likelihood, given the weights (fuzzy assignment) Parameters ---------- X : np.array, shape=[n_cases, n_features] Inputs y : np.array, shape=[n_cases] Response variable weights : np.array, shape=[n_cases, n_classes] For each case, gives the probability is was generated by each of the classes Returns ------- beta : np.array, shape=[n_classes, n_features] beta[i] gives the regression parameters betahat for class i """ n_cases, n_features = X.shape n_classes = weights.shape[1] assert weights.shape[0] == n_cases beta = np.zeros((n_classes, n_features)) for i in xrange(n_classes): W = np.diag(weights[:,i]) # intermediate computation X_T_W = np.dot(X.T, W) # weighted least squares beta[i, :] = np.dot(np.linalg.pinv(np.dot(X_T_W, X)), np.dot(X_T_W, y)) return beta def expectation(X, y, beta): """Calculate the expected weights (fuzzy assignment) given the values of beta Parameters ---------- X : np.array, shape=[n_cases, n_features] Inputs y : np.array, shape=[n_cases] Response variable beta : np.array, shape=[n_classes, n_features] beta[i] gives the regression parameters betahat for class i Returns ------- weights : np.array, shape=[n_cases, n_classes] For each case, gives the probability is was generated by each of the classes """ prob = np.exp(-np.square(y - np.dot(X, beta.T).T)) return (prob / np.sum(prob, axis=0)).T def plot_2d_example(): a = 1; b=3 X = 10*np.random.randn(100,1) y = np.concatenate((X[:50,0]*a, X[50:,0]*b)) + np.random.randn(100) weights = np.random.randn(100,2)**2 weights = (weights.T / np.sum(weights, axis=1)).T pp.scatter(X[:,0], y, edgecolor='k', facecolors='none') max_iterations = 2 for i in range(max_iterations): beta = maximization(X, y, weights) weights = expectation(X, y, beta) print 'beta %s' % i print beta x = np.linspace(np.min(X), np.max(X)) for j in range(beta.shape[0]): pp.plot(x, beta[j]*x, c=cm.jet(float(i)/max_iterations), label='it %s' % i, lw=2) pp.legend() pp.savefig('wls-em-2d-example.png') if __name__ == '__main__': plot_2d_example()