ENH: add internals.ComplexBlock

Author: tkf

pandas/core/internals.py

@@ -250,6 +250,12 @@ def should_store(self, value):
         # unnecessarily
         return issubclass(value.dtype.type, np.floating)
 
+class ComplexBlock(Block):
+    _can_hold_na = True
+
+    def should_store(self, value):
+        return issubclass(value.dtype.type, np.complexfloating)
+
 class IntBlock(Block):
     _can_hold_na = False
 
@@ -267,7 +273,8 @@ class ObjectBlock(Block):
 
     def should_store(self, value):
         return not issubclass(value.dtype.type,
-                              (np.integer, np.floating, np.bool_))
+                              (np.integer, np.floating, np.complexfloating,
+                               np.bool_))
 
 class DatetimeBlock(IntBlock):
     _can_hold_na = True
@@ -279,6 +286,8 @@ def make_block(values, items, ref_items, do_integrity_check=False):
 
     if issubclass(vtype, np.floating):
         klass = FloatBlock
+    elif issubclass(vtype, np.complexfloating):
+        klass = ComplexBlock
     elif issubclass(vtype, np.datetime64):
         klass = DatetimeBlock
     elif issubclass(vtype, np.integer):
@@ -423,7 +432,7 @@ def is_consolidated(self):
 
     def get_numeric_data(self, copy=False):
         num_blocks = [b for b in self.blocks
-                      if isinstance(b, (IntBlock, FloatBlock))]
+                      if isinstance(b, (IntBlock, FloatBlock, ComplexBlock))]
 
         indexer = np.sort(np.concatenate([b.ref_locs for b in num_blocks]))
         new_items = self.items.take(indexer)
@@ -1103,19 +1112,22 @@ def _interleaved_dtype(blocks):
     have_bool = counts[BoolBlock] > 0
     have_object = counts[ObjectBlock] > 0
     have_float = counts[FloatBlock] > 0
+    have_complex = counts[ComplexBlock] > 0
     have_dt64 = counts[DatetimeBlock] > 0
-    have_numeric = have_float or have_int
+    have_numeric = have_float or have_complex or have_int
 
     if have_object:
         return np.object_
     elif have_bool and have_numeric:
         return np.object_
     elif have_bool:
         return np.bool_
-    elif have_int and not have_float:
+    elif have_int and not have_float and not have_complex:
         return np.int64
-    elif have_dt64 and not have_float:
+    elif have_dt64 and not have_float and not have_complex:
         return np.datetime64
+    elif have_complex:
+        return np.complex64
     else:
         return np.float64
 

pandas/src/tseries.pyx

@@ -156,7 +156,7 @@ cdef double INF = <double> np.inf
 cdef double NEGINF = -INF
 
 cpdef checknull(object val):
-    if util.is_float_object(val):
+    if util.is_float_object(val) or util.is_complex_object(val):
         return val != val or val == INF or val == NEGINF
     elif util.is_datetime64_object(val):
         return val.view('i8') == NaT

pandas/tests/test_internals.py

@@ -19,13 +19,17 @@ def assert_block_equal(left, right):
 def get_float_mat(n, k):
     return np.repeat(np.atleast_2d(np.arange(k, dtype=float)), n, axis=0)
 
-TEST_COLS = ['a', 'b', 'c', 'd', 'e', 'f', 'g']
+TEST_COLS = ['a', 'b', 'c', 'd', 'e', 'f', 'g', 'h']
 N = 10
 
 def get_float_ex(cols=['a', 'c', 'e']):
     floats = get_float_mat(N, 3).T
     return make_block(floats, cols, TEST_COLS)
 
+def get_complex_ex(cols=['h']):
+    complexes = (get_float_mat(N, 1).T * 1j).astype(np.complex64)
+    return make_block(complexes, cols, TEST_COLS)
+
 def get_obj_ex(cols=['b', 'd']):
     mat = np.empty((N, 2), dtype=object)
     mat[:, 0] = 'foo'
@@ -44,6 +48,7 @@ class TestBlock(unittest.TestCase):
 
     def setUp(self):
         self.fblock = get_float_ex()
+        self.cblock = get_complex_ex()
         self.oblock = get_obj_ex()
         self.bool_block = get_bool_ex()
         self.int_block = get_int_ex()
@@ -60,6 +65,7 @@ def _check(blk):
             assert_block_equal(blk, unpickled)
 
         _check(self.fblock)
+        _check(self.cblock)
         _check(self.oblock)
         _check(self.bool_block)
 
@@ -175,7 +181,8 @@ def setUp(self):
         self.blocks = [get_float_ex(),
                        get_obj_ex(),
                        get_bool_ex(),
-                       get_int_ex()]
+                       get_int_ex(),
+                       get_complex_ex()]
         self.mgr = BlockManager.from_blocks(self.blocks, np.arange(N))
 
     def test_constructor_corner(self):
@@ -198,13 +205,13 @@ def test_is_indexed_like(self):
         self.assert_(not self.mgr._is_indexed_like(mgr2))
 
     def test_block_id_vector_item_dtypes(self):
-        expected = [0, 1, 0, 1, 0, 2, 3]
+        expected = [0, 1, 0, 1, 0, 2, 3, 4]
         result = self.mgr.block_id_vector
         assert_almost_equal(expected, result)
 
         result = self.mgr.item_dtypes
         expected = ['float64', 'object', 'float64', 'object', 'float64',
-                    'bool', 'int64']
+                    'bool', 'int64', 'complex64']
         self.assert_(np.array_equal(result, expected))
 
     def test_union_block_items(self):
@@ -298,6 +305,7 @@ def test_consolidate_ordering_issues(self):
         self.mgr.set('d', randn(N))
         self.mgr.set('b', randn(N))
         self.mgr.set('g', randn(N))
+        self.mgr.set('h', randn(N))
 
         cons = self.mgr.consolidate()
         self.assertEquals(cons.nblocks, 1)