for pandas DataFrame

hoxo-m · Dec 11, 2016 · 9e7abea · 9e7abea
1 parent cd1f987
commit 9e7abea
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 0 deletions.
diff --git a/densratio/helpers.py b/densratio/helpers.py
@@ -13,6 +13,8 @@ def to_numpy_matrix(x):
             return matrix(x).T
         else:
             return matrix(x)
+    elif str(type(x)) == "<class 'pandas.core.frame.DataFrame'>":
+        return x.as_matrix()
     elif not x:
         raise ValueError("Cannot transform to numpy.matrix.")
     else:

diff --git a/tests/test_helpers.py b/tests/test_helpers.py
@@ -4,6 +4,7 @@
 
 from numpy import array, matrix
 from numpy.testing import assert_array_equal
+from pandas import DataFrame
 from .context import helpers
 
 class BasicTestSuite(unittest.TestCase):
@@ -42,6 +43,10 @@ def test_to_numpy_matrix_matrix(self):
         x = helpers.to_numpy_matrix(x)
         assert_array_equal(x, matrix([[1,2],[3,4]]))
 
+    def test_to_numpy_matrix_pandas_DataFrame(self):
+        x = DataFrame([[1,2],[3,4]])
+        x = helpers.to_numpy_matrix(x)
+        assert_array_equal(x, matrix([[1,2],[3,4]]))
 
 if __name__ == '__main__':
     unittest.main()