Skip to content
Browse files

ENH: add internals.ComplexBlock

  • Loading branch information...
1 parent a351e94 commit 824cb1045a21527b7515a7d0338fd4af36c7d656 @tkf committed Apr 21, 2012
Showing with 30 additions and 10 deletions.
  1. +17 −5 pandas/core/internals.py
  2. +1 −1 pandas/src/tseries.pyx
  3. +12 −4 pandas/tests/test_internals.py
View
22 pandas/core/internals.py
@@ -243,6 +243,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
@@ -260,7 +266,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
@@ -272,6 +279,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):
@@ -416,7 +425,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)
@@ -1094,19 +1103,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
View
2 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
View
16 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)

0 comments on commit 824cb10

Please sign in to comment.
Something went wrong with that request. Please try again.