Fixed a bit in is_contigous

npstructures/bitarray.py

@@ -1,3 +1,6 @@
+from functools import lru_cache
+from typing import Union
+
 import numpy as np
 from .util import np
 from numbers import Number
@@ -11,11 +14,11 @@ class BitArray:
     def __init__(self, data: np.ndarray, bit_stride: int, shape: tuple, offset: int = 0):
         self._data = data
         self._bit_stride = self._dtype(bit_stride)
-        self._mask = self._dtype(2**bit_stride-1)
+        self._mask = self._dtype(2 ** bit_stride - 1)
         self._shape = shape
         self._offset = self._dtype(offset)
-        self._n_entries_per_register = self._register_size//self._bit_stride
-        self._shifts = self._dtype(bit_stride)*np.arange(self._n_entries_per_register, dtype=self._dtype)
+        self._n_entries_per_register = self._register_size // self._bit_stride
+        self._shifts = self._dtype(bit_stride) * np.arange(self._n_entries_per_register, dtype=self._dtype)
 
     @classmethod
     def pack(cls, array: np.ndarray, bit_stride: int) -> "BitArray":
@@ -40,12 +43,12 @@ def pack(cls, array: np.ndarray, bit_stride: int) -> "BitArray":
         """
         assert cls._register_size % bit_stride == 0
         n_entries_per_register = cls._register_size // cls._dtype(bit_stride)
-        shifts = cls._dtype(bit_stride)*np.arange(n_entries_per_register, dtype=cls._dtype)
+        shifts = cls._dtype(bit_stride) * np.arange(n_entries_per_register, dtype=cls._dtype)
         bits = array[0::n_entries_per_register].astype(cls._dtype)
         for i, shift in enumerate(shifts[1:], 1):
             size = array[i::n_entries_per_register].size
             bits[:size] |= (array[i::n_entries_per_register].astype(cls._dtype) << shift.astype(cls._dtype))
-        
+
         return cls(bits, bit_stride, array.shape)
 
     def unpack(self) -> np.ndarray:
@@ -61,15 +64,15 @@ def unpack(self) -> np.ndarray:
 
         return values[:self._shape[0]]
 
-    def __getitem__(self, idx: list):
+    def __getitem__(self, idx: Union[int, np.ndarray, list]):
         if isinstance(idx, list):
             idx = np.asanyarray(idx)
-        register_idx = self._dtype(idx+self._offset) // (self._n_entries_per_register)
+        register_idx = self._dtype(idx + self._offset) // (self._n_entries_per_register)
         register_offset = self._dtype(idx + self._offset) % (self._n_entries_per_register)
         if isinstance(idx, Number):
-            return (self._data[register_idx] >> (register_offset*self._bit_stride)) & self._mask
+            return (self._data[register_idx] >> (register_offset * self._bit_stride)) & self._mask
         if isinstance(idx, np.ndarray):
-            array = self._data[register_idx] >> (register_offset*self._bit_stride) & self._mask
+            array = self._data[register_idx] >> (register_offset * self._bit_stride) & self._mask
             return self.pack(array, self._bit_stride)
 
     def sliding_window(self, window_size: int) -> np.ndarray:
@@ -90,9 +93,34 @@ def sliding_window(self, window_size: int) -> np.ndarray:
 
 
         """
-        mask = (~self._dtype(0)) >> self._dtype(self._register_size-window_size * self._bit_stride)
+        mask = (~self._dtype(0)) >> self._dtype(self._register_size - window_size * self._bit_stride)
         rev_shifts = self._shifts[::-1] + self._bit_stride
         res = self._data[:, None] >> self._shifts
         res[:-1] |= self._data[1:, None] << rev_shifts
         res &= mask
-        return res.ravel()[:self._shape[0]-window_size+1]
+        return res.ravel()[:self._shape[0] - window_size + 1]
+
+
+class BitMask(BitArray):
+    _dtype = np.uint8
+    _register_size = _dtype(8)
+
+    @classmethod
+    def zeros(cls, shape):
+        data_size = (shape + cls._register_size - 1) // cls._register_size
+        return cls(np.zeros(data_size, dtype=cls._dtype), 1, shape)
+
+    @property
+    @lru_cache()
+    def _bit_masks(self):
+        return self._dtype(1) << np.arange(self._register_size, dtype=self._dtype)
+
+
+    def __setitem__(self, idx, value):
+        assert value in (True, False)
+        idx = np.asanyarray(idx)
+        register_idx = self._dtype(idx + self._offset) // (self._n_entries_per_register)
+        register_offset = self._dtype(idx + self._offset) % (self._n_entries_per_register)
+        if value is True:
+            self._data[register_idx] |= self._bit_masks[register_offset]
+

npstructures/npdataclasses.py

@@ -4,7 +4,7 @@
 from itertools import accumulate
 import numpy as np
 
-    
+
 def shallow_tuple(obj):
     return tuple(
         getattr(obj, field.name) for field in dataclasses.fields(obj)
@@ -106,8 +106,13 @@ def __array_function__(self, func, types, args, kwargs):
         if func == np.concatenate:
             objects = args[0]
             tuples = [shallow_tuple(o) for o in objects]
-            return self.__class__(*(np.concatenate(list(t))
-                                    for t in zip(*tuples)))
+            columns = []
+            for t in zip(*tuples):
+                assert all(id(type(t[0])) == id(type(i)) for i in t), ([type(v) for v in t], [id(type(v)) for v in t])
+                columns.append(np.concatenate(list(t)))
+            return self.__class__(*columns)
+            # np.concatenate(list(t))
+            #                    for t in zip(*tuples)))
         if func == np.equal:
             one, other = args
             return all(
@@ -170,7 +175,7 @@ def __str__(self):
                 lines.append("".join(f"{str(col)[:col_length-2]:>{col_length}}" for col in cols))
             return "\n".join(lines)
 
-    
+
         __repr__ = __str__
 
 

npstructures/raggedarray/base.py

@@ -5,14 +5,17 @@
 
 class RaggedBase:
     '''
-    Base class for ragged arrays. 
+    Base class for ragged arrays.
     Handles evertything to do with the underlying data buffer.
     '''
 
     def __init__(self, data, shape):
         self.__data = data
         self._shape = shape
-        self.is_contigous = True
+        if isinstance(shape, (RaggedView, RaggedView2)):
+            self.is_contigous = False
+        else:
+            self.is_contigous = True
 
     @property
     def size(self) -> int:
@@ -24,7 +27,7 @@ def dtype(self) -> npt.DTypeLike:
 
     def _change_view(self, new_view):
         ret = self.__class__(self.__data, new_view)
-        ret.is_contigous = False
+        #ret.is_contigous = False
         return ret
 
     def _flatten_myself(self):
@@ -35,11 +38,11 @@ def _flatten_myself(self):
             self.__data = self.__data[idx]
             self._shape = shape
             self.is_contigous = True
-        else:
-            assert False, self._shape
+        # else:
+        #    #assert False, self._shape
 
     def ravel(self) -> npt.ArrayLike:
-        """Return a flattened view of the data. 
+        """Return a flattened view of the data.
 
         For now it makes the data contigous on this call. Changes the
         state of the array
@@ -48,7 +51,7 @@ def ravel(self) -> npt.ArrayLike:
         -------
         npt.ArrayLike
         """
-        
+
         if not self.is_contigous:
             self._flatten_myself()
         return self.__data