@@ -38,14 +38,17 @@ def predict(otrain):
     return norm.inverse_transform(filled.T).T
 
 def main(transpose_inputs=False):
-    train, test = get_train_test()
+    train, test = get_train_test(random_state=12)
     if transpose_inputs:
         train = train.T
         test  = test.T
 
     predicted = predict(train)
     r2 = metrics.r2_score(test[test > 0], predicted[test > 0])
-    print('R2 score (binary neighbours): {:.1%}'.format(r2))
+    print('R2 score (binary {} neighbours): {:.1%}'.format(
+        ('movie' if transpose_inputs else 'user'),
+        r2))
 
 if __name__ == '__main__':
     main()
+    main(transpose_inputs=True)
@@ -22,12 +22,29 @@ def load():
     reviews = sparse.csc_matrix((values, ij.T)).astype(float)
     return reviews.toarray()
 
-def get_train_test():
+def get_train_test(reviews=None, return_valid=False, random_state=None):
+    '''Split data into training & testing
+
+    Parameters
+    ----------
+    reviews : ndarray, optional
+        Input data
+
+    Returns
+    -------
+    train : ndarray
+        training data
+    test : ndarray
+        testing data
+    '''
     import numpy as np
     import random
-    reviews = load()
+    r = random.Random(random_state)
+
+    if reviews is None:
+        reviews = load()
     U,M = np.where(reviews)
-    test_idxs = np.array(random.sample(range(len(U)), len(U)//10))
+    test_idxs = np.array(r.sample(range(len(U)), len(U)//10))
     train = reviews.copy()
     train[U[test_idxs], M[test_idxs]] = 0
 
@@ -38,5 +55,7 @@ def get_train_test():
     valid_M = train.any(1)
     train = train[valid_M].T[valid_U].T
     test  = test [valid_M].T[valid_U].T
+    if return_valid:
+        return train, test, valid_M, valid_U
     return train, test
 
@@ -28,3 +28,25 @@ def inverse_transform(self, features, copy=True):
         features *= self.std
         features += self.mean
         return features
+
+def predict(train):
+    norm = NormalizePositive()
+    train = norm.fit_transform(train)
+    return norm.inverse_transform(train * 0.)
+
+
+def main(transpose_inputs=False):
+    from load_ml100k import get_train_test
+    from sklearn import metrics
+    train,test = get_train_test(random_state=12)
+    if transpose_inputs:
+        train = train.T
+        test = test.T
+    predicted = predict(train)
+    r2 = metrics.r2_score(test[test > 0], predicted[test > 0])
+    print('R2 score ({} normalization): {:.1%}'.format(
+        ('movie' if transpose_inputs else 'user'),
+        r2))
+if __name__ == '__main__':
+    main()
+    main(transpose_inputs=True)
@@ -7,7 +7,6 @@
 
 import numpy as np
 from sklearn.linear_model import ElasticNetCV
-from load_ml100k import get_train_test
 from norm import NormalizePositive
 from sklearn import metrics
 
@@ -25,21 +24,27 @@ def predict(train):
         # remove the current user for training
         curtrain = np.delete(train, u, axis=0)
         bu = binary[u]
-        reg.fit(curtrain[:,bu].T, train[u, bu])
+        if np.sum(bu) > 5:
+            reg.fit(curtrain[:,bu].T, train[u, bu])
 
-	# Fill the values that were not there already
-        filled[u, ~bu] = reg.predict(curtrain[:,~bu].T)
+            # Fill the values that were not there already
+            filled[u, ~bu] = reg.predict(curtrain[:,~bu].T)
     return norm.inverse_transform(filled)
 
 
 def main(transpose_inputs=False):
-    train,test = get_train_test()
+    from load_ml100k import get_train_test
+    train,test = get_train_test(random_state=12)
     if transpose_inputs:
         train = train.T
         test = test.T
     filled = predict(train)
     r2 = metrics.r2_score(test[test > 0], filled[test > 0])
-    print('R2 score (user regression): {:.1%}'.format(r2))
+
+    print('R2 score ({} regression): {:.1%}'.format(
+        ('movie' if transpose_inputs else 'user'),
+        r2))
 
 if __name__ == '__main__':
     main()
+    main(transpose_inputs=True)