import numpy as np
from numpy.testing import assert_allclose

def orthonormalize(X, rows=True):
    if rows: X = X.T
    Q, R = np.linalg.qr(X)
    if rows: Q = Q.T
    return Q


def random_orthonormal(n_samples, n_features, rseed=None):
    assert(n_samples <= n_features)
    rand = np.random.RandomState(rseed)
    X = rand.randn(n_samples, n_features)
    return orthonormalize(X)


def solve_weighted(A, b, w):
    """solve Ax = b with weights w"""
    return np.linalg.solve(np.dot(A.T * w ** 2, A),
                           np.dot(A.T * w ** 2, b))


def _Estep(eigvec, data, weights, coeff):
    for i in range(data.shape[0]):
        coeff[i] = solve_weighted(eigvec.T, data[i], weights[i])
    return coeff


def _Mstep(eigvec, data, weights, coeff):
    data = data.copy()
    for i in range(eigvec.shape[0]):
        c = coeff[:, i]
        for j in range(data.shape[1]):
            w = weights[:, j]
            x = data[:, j]
            eigvec[i, j] = np.dot(x, w ** 2 * c) / np.dot(c, w ** 2 * c)

        # remove this vector from the data
        data -= np.outer(coeff[:, i], eigvec[i])

    return orthonormalize(eigvec)


def empca(data, weights, nvec, niter=25, rseed=None):
    assert data.shape == weights.shape
    assert nvec <= data.shape[1]
    eigvec = random_orthonormal(nvec, data.shape[1], rseed=rseed)
    coeff = np.zeros((data.shape[0], nvec))

    for k in range(niter):
        coeff = _Estep(eigvec, data, weights, coeff)
        eigvec = _Mstep(eigvec, data, weights, coeff)
    coeff = _Estep(eigvec, data, weights, coeff)

    return eigvec, coeff


def pca(data, nvec):
    U, s, VT = np.linalg.svd(data)
    return VT[:, :nvec], U[:, :nvec] * s[:nvec]

@@ -55,16 +55,17 @@
 lclose%0A%0A
+%0A
 def orth
@@ -98,24 +98,45 @@
     
-if rows: X = X.T
+orient = lambda X: X.T if rows else X
 %0A   
@@ -160,31 +160,18 @@
 .qr(
-X)%0A    if rows: Q = Q.T
+orient(X))
 %0A   
@@ -178,17 +178,25 @@
  return 
-Q
+orient(Q)
 %0A%0A%0Adef r
@@ -461,19 +461,214 @@
 eights w
-%22%22%22
+%0A%0A    Parameters%0A    ----------%0A    A : array-like %5BN, M%5D%0A    b : array-like %5BN%5D%0A    w : array-like %5BN%5D%0A%0A    Returns%0A    -------%0A    x : ndarray, %5BM%5D%0A    %22%22%22%0A    A, b, w = map(np.asarray, (A, b, w))
 %0A    ret
@@ -804,24 +804,43 @@
 ts, coeff):%0A
+    # Update coeff%0A
     for i in
@@ -931,25 +931,8 @@
 i%5D)%0A
-    return coeff%0A
 %0A%0Ade
@@ -979,26 +979,45 @@
     
-data = data.copy()
+# Update eigvec%0A    w2 = weights ** 2
 %0A   
@@ -1058,20 +1058,34 @@
         
-c
+d
  = 
+data - np.dot(
 coeff%5B:,
@@ -1089,62 +1089,33 @@
 %5B:, 
+:
 i%5D
-%0A        for j in range(data.shape%5B1%5D):%0A    
+, eigvec%5B:i%5D)%0A
         
 w = 
@@ -1114,57 +1114,30 @@
     
-w = weights%5B:, j%5D%0A            x = data%5B:, j%5D%0A    
+c = coeff%5B:, i:i + 1%5D%0A
     
@@ -1152,11 +1152,8 @@
 ec%5Bi
-, j
 %5D = 
@@ -1163,21 +1163,19 @@
 dot(
-x, w ** 
+c.T, w
 2 * 
-c
+d
 ) / 
@@ -1186,31 +1186,28 @@
 ot(c
+.T
 , w
- ** 
 2 * c)%0A
-%0A
         
 # re
@@ -1206,103 +1206,24 @@
     
-# remove this vector from the data%0A        data -= np.outer(coeff%5B:, i%5D, eigvec%5Bi%5D)%0A%0A    return
+eigvec%5B:i + 1%5D =
  ort
@@ -1240,16 +1240,25 @@
 e(eigvec
+%5B: i + 1%5D
 )%0A%0A%0Adef 
@@ -1524,24 +1524,16 @@
 %0A       
- coeff =
  _Estep(
@@ -1569,25 +1569,16 @@
 %0A       
- eigvec =
  _Mstep(
@@ -1610,24 +1610,16 @@
 eff)%0A   
- coeff =
  _Estep(

make empca more efficient