From e5355ee5e91d50e524808d2ad1c40ff7d4726251 Mon Sep 17 00:00:00 2001
From: Allen Tran <realallentran@gmail.com>
Date: Sun, 24 Dec 2017 10:07:04 -0800
Subject: [PATCH] add tests, fix typo in README.md

---
 README.md          |  2 +-
 ppca/_ppca.py      |  5 +----
 tests/unittests.py | 28 ++++++++++++++++++++++++++++
 3 files changed, 30 insertions(+), 5 deletions(-)
 create mode 100644 tests/unittests.py

diff --git a/README.md b/README.md
index d80dba9..db0275a 100644
--- a/README.md
+++ b/README.md
@@ -27,7 +27,7 @@ ppca.fit(d=100, verbose=True)
 The model parameters and components will be attached to the ppca object.
 ```
 variance_explained = ppca.var_exp
-components = ppca.X
+components = ppca.data
 model_params = ppca.C
 ```
 If you want the principal components, call `transform`.
diff --git a/ppca/_ppca.py b/ppca/_ppca.py
index 40f1949..c7f4c85 100644
--- a/ppca/_ppca.py
+++ b/ppca/_ppca.py
@@ -13,6 +13,7 @@ def __init__(self):
         self.C = None
         self.means = None
         self.stds = None
+        self.eig_vals = None
 
     def _standardize(self, X):
 
@@ -107,10 +108,6 @@ def fit(self, data, d=None, tol=1e-4, min_obs=10, verbose=False):
         self.eig_vals = vals
         self._calc_var()
 
-        import IPython
-        IPython.embed()
-        assert False
-
     def transform(self, data=None):
 
         if self.C is None:
diff --git a/tests/unittests.py b/tests/unittests.py
new file mode 100644
index 0000000..ac63f44
--- /dev/null
+++ b/tests/unittests.py
@@ -0,0 +1,28 @@
+import unittest
+import random
+
+import numpy as np
+from ppca import PPCA
+
+
+class NaNTest(unittest.TestCase):
+
+    def remove_nan_test(self):
+
+        N = 101
+        k = 23
+        p_nan = 0.02
+        n_components = 3
+
+        data = np.random.random((N, k))
+        for n in xrange(N):
+            for _k in xrange(k):
+                if random.random() < p_nan:
+                    data[n, _k] = np.nan
+
+        pca = PPCA()
+        pca.fit(data, n_components)
+
+        self.assertEqual(pca.data[np.isnan(pca.data)].shape, (0, ))
+        self.assertEqual(pca.C.shape, (k, n_components))
+        self.assertEqual(pca.transform().shape, (N, n_components))