Add some counting routines and tests.

yt/utilities/lib/bitarray.pxd

@@ -78,6 +78,13 @@ cdef inline void ba_set_range(np.uint8_t *buf, np.uint64_t start_ind,
     else:
         buf[buf_stop] &= ~bitmask
 
+
+cdef inline np.uint8_t _num_set_bits( np.uint8_t b ):
+    # https://stackoverflow.com/questions/30688465/how-to-check-the-number-of-set-bits-in-an-8-bit-unsigned-char
+    b = b - ((b >> 1) & 0x55)
+    b = (b & 0x33) + ((b >> 2) & 0x33)
+    return (((b + (b >> 4)) & 0x0F) * 0x01)
+
 cdef class bitarray:
     cdef np.uint8_t *buf
     cdef np.uint64_t size
@@ -87,4 +94,4 @@ cdef class bitarray:
     cdef void _set_value(self, np.uint64_t ind, np.uint8_t val)
     cdef np.uint8_t _query_value(self, np.uint64_t ind)
     cdef void _set_range(self, np.uint64_t start, np.uint64_t stop, np.uint8_t val)
-    #cdef int query_range(self, np.uint64_t ind, np.uint64_t count, int *val)
+    cdef np.uint64_t _count(self)

yt/utilities/lib/bitarray.pyx

@@ -196,3 +196,33 @@ cdef class bitarray:
 
         """
         ba_set_range(self.buf, start, stop, val)
+
+    @cython.boundscheck(False)
+    @cython.wraparound(False)
+    @cython.cdivision(True)
+    cdef np.uint64_t _count(self):
+        cdef np.uint64_t count = 0
+        cdef np.uint64_t i
+        for i in range(self.buf_size):
+            count += _num_set_bits(self.buf[i])
+        return count
+
+
+    @cython.boundscheck(False)
+    @cython.wraparound(False)
+    @cython.cdivision(True)
+    def count(self):
+        r"""Count the number of values set in the array.
+
+        Parameters
+        ----------
+
+        Examples
+        --------
+
+        >>> arr_in = np.array([True, True, False, True, True, False])
+        >>> a = ba.bitarray(arr = arr_in)
+        >>> a.count()
+
+        """
+        return self._count()

yt/utilities/lib/tests/test_bitarray.py

@@ -26,6 +26,7 @@ def test_inout_bitarray():
     b = ba.bitarray(arr=arr_in)
     arr_out = b.as_bool_array()
     assert_equal(arr_in, arr_out)
+    assert_equal(b.count(), arr_in.sum())
 
     # Let's check we can do something interesting.
     arr_in1 = np.random.random(32**3) > 0.5
@@ -34,6 +35,8 @@ def test_inout_bitarray():
     b2 = ba.bitarray(arr=arr_in2)
     b3 = ba.bitarray(arr=(arr_in1 & arr_in2))
     assert_equal((b1.ibuf & b2.ibuf), b3.ibuf)
+    assert_equal(b1.count(), arr_in1.sum())
+    assert_equal(b2.count(), arr_in2.sum())
 
     b = ba.bitarray(10)
     for i in range(10):
@@ -61,6 +64,7 @@ def test_set_range():
     comparison_array[4:65] = 1
     arr = b.as_bool_array().astype("uint8")
     assert_array_equal(arr, comparison_array)
+    assert_equal(b.count(), comparison_array.sum())
 
     # Test when we start and stop in the same byte
     b = ba.bitarray(127)
@@ -69,6 +73,7 @@ def test_set_range():
     comparison_array[4:6] = 1
     arr = b.as_bool_array().astype("uint8")
     assert_array_equal(arr, comparison_array)
+    assert_equal(b.count(), comparison_array.sum())
 
     # Test now where we're in the middle of start
     b = ba.bitarray(64)
@@ -77,6 +82,7 @@ def test_set_range():
     comparison_array[33:36] = 1
     arr = b.as_bool_array().astype("uint8")
     assert_array_equal(arr, comparison_array)
+    assert_equal(b.count(), comparison_array.sum())
 
     # Now we test when we end on a byte edge, but we have 65 entries
     b = ba.bitarray(65)
@@ -85,6 +91,7 @@ def test_set_range():
     comparison_array[32:64] = 1
     arr = b.as_bool_array().astype("uint8")
     assert_array_equal(arr, comparison_array)
+    assert_equal(b.count(), comparison_array.sum())
 
     # Let's do the inverse
     b = ba.bitarray(127)
@@ -97,6 +104,7 @@ def test_set_range():
     comparison_array[3:9] = 1
     arr = b.as_bool_array().astype("uint8")
     assert_array_equal(arr, comparison_array)
+    assert_equal(b.count(), comparison_array.sum())
 
     # Now let's overlay some zeros
     b.set_range(7, 10, 0)