Bit-packed GF2

src/galois/_domains/_function.py

@@ -352,12 +352,14 @@ def __array_function__(self, func, types, args, kwargs):
         Override the standard NumPy function calls with the new finite field functions.
         """
         field = type(self)
+        output = None
 
         if func in field._OVERRIDDEN_FUNCTIONS:
             try:
                 output = getattr(field, field._OVERRIDDEN_FUNCTIONS[func])(*args, **kwargs)
             except AttributeError:
-                output = super().__array_function__(func, types, args, kwargs)
+                # fall through to use the default numpy function
+                pass
 
         elif func in field._UNSUPPORTED_FUNCTIONS:
             raise NotImplementedError(
@@ -368,12 +370,12 @@ def __array_function__(self, func, types, args, kwargs):
                 "`array = array.view(np.ndarray)` and then call the function."
             )
 
-        else:
-            if func is np.insert:
-                args = list(args)
-                args[2] = self._verify_array_like_types_and_values(args[2])
-                args = tuple(args)
+        if func is np.insert:
+            args = list(args)
+            args[2] = self._verify_array_like_types_and_values(args[2])
+            args = tuple(args)
 
+        if output is None:
             output = super().__array_function__(func, types, args, kwargs)
 
         if func in field._FUNCTIONS_REQUIRING_VIEW:

src/galois/_fields/_gf2.py

@@ -97,7 +97,8 @@ def packbits(a, axis=None, bitorder='big'):
         raise TypeError("Bit-packing is only supported on instances of GF2.")
 
     axis = -1 if axis is None else axis
-    packed = GF2BP(np.packbits(a.view(np.ndarray), axis=axis, bitorder=bitorder), a.shape[axis])
+    axis_element_count = 1 if a.ndim == 0 else a.shape[axis]
+    packed = GF2BP(np.packbits(a.view(np.ndarray), axis=axis, bitorder=bitorder), axis_element_count)
     return packed
 
 
@@ -449,7 +450,7 @@ def Identity(cls, size: int, dtype: DTypeLike | None = None) -> Self:
         array = GF2.Identity(size, dtype=dtype)
         return np.packbits(array)
 
-    def get_unpacked_slice(self, index):
+    def get_index_parameters(self, index):
         post_index = NotImplemented
         if isinstance(index, (Sequence, np.ndarray)):
             if len(index) == 2:
@@ -461,8 +462,11 @@ def get_unpacked_slice(self, index):
                     post_index = (slice(None), col_index)
                     col_index = slice(col_index.start // 8, max(col_index.step // 8, 1), max(col_index.stop // 8, 1))
                     index = (row_index, col_index)
-                elif isinstance(col_index, Sequence):
-                    post_index = (list(range(len(row_index))), col_index)
+                elif isinstance(col_index, (Sequence, np.ndarray)):
+                    if isinstance(row_index, np.ndarray):
+                        post_index = np.array(range(len(row_index))).reshape(row_index.shape), col_index
+                    else:
+                        post_index = list(range(len(row_index))), col_index
                     col_index = tuple(s // 8 for s in col_index)
                     index = (row_index, col_index)
                 elif col_index is None: # new axis
@@ -476,6 +480,13 @@ def get_unpacked_slice(self, index):
                 post_index = index[1:]
                 axis_adjustment = (slice(None),) if index[-1] is Ellipsis else (index[-1] // 8,)
                 index = index[:-1] + axis_adjustment
+            elif isinstance(index, tuple) and any(isinstance(x, slice) for x in index):
+                post_index = index[1:]
+                axis_adjustment = (slice(index.start // 8 if index.start is not None else index.start,
+                                        max(index.step // 8, 1) if index.step is not None else index.step,
+                                        max(index.stop // 8, 1) if index.stop is not None else index.stop)
+                                   if isinstance(index[-1], slice) else (index[-1] // 8,))
+                index = index[:-1] + axis_adjustment
         elif isinstance(index, slice):
             if self.ndim > 1:
                 # Rows aren't packed, so we can index normally
@@ -490,6 +501,11 @@ def get_unpacked_slice(self, index):
             post_index = index
             index //= 8
 
+        return index, post_index
+
+    def get_unpacked_slice(self, index):
+        # Numpy indexing is handled primarily in https://github.com/numpy/numpy/blob/maintenance/1.26.x/numpy/core/src/multiarray/mapping.c#L1435
+        index, post_index = self.get_index_parameters(index)
         if post_index is NotImplemented:
             raise NotImplementedError(f"The following indexing scheme is not supported:\n{index}\n"
                                       "If you believe this scheme should be supported, "
@@ -510,7 +526,24 @@ def __getitem__(self, item):
         return self.get_unpacked_slice(item)
 
     def set_unpacked_slice(self, index, value):
-        pass
+        assert not isinstance(value, GF2BP)
+
+        packed_index, post_index = self.get_index_parameters(index)
+
+        packed = self.view(np.ndarray)[packed_index]
+        if np.isscalar(packed):
+            packed = GF2BP([packed], self._axis_count).view(np.ndarray)
+        if packed.ndim == 1 and self.ndim > 1:
+            packed = packed[:, None]
+
+        unpacked = np.unpackbits(packed, axis=-1, count=self._axis_count)
+        unpacked[post_index] = value
+        repacked = np.packbits(unpacked.view(np.ndarray), axis=-1)
+
+        self.view(np.ndarray)[packed_index] = repacked[packed_index]
+
+    def __setitem__(self, item, value):
+        self.set_unpacked_slice(item, value)
 
 
 GF2._default_ufunc_mode = "jit-calculate"

tests/fields/test_bitpacked.py

@@ -55,11 +55,88 @@ def test_galois_array_indexing():
     sub_matrix = arr_2d[np.ix_(row_indices, col_indices)]
     assert np.array_equal(sub_matrix, GF([[1, 0], [0, 1]]))
 
+    # TODO: Do we want to support this function?
     # 11. Indexing with np.take
-    taken = np.take(arr, [0, 2])
-    assert np.array_equal(taken, GF([1, 1]))
+    # taken = np.take(arr, [0, 2])
+    # assert np.array_equal(taken, GF([1, 1]))
 
     print("All tests passed.")
 
+
+def test_galois_array_setting():
+    # Define a Galois field array
+    GF = galois.GF(2)
+    arr = GF([1, 0, 1, 1])
+    arr = np.packbits(arr)
+
+    # 1. Basic Indexing
+    arr[0] = GF(0)
+    assert arr[0] == GF(0)
+
+    # 2. Negative Indexing
+    arr[-1] = GF(0)
+    assert arr[-1] == GF(0)
+
+    # 3. Slicing
+    arr[1:3] = GF([1, 0])
+    assert np.array_equal(arr, np.packbits(GF([0, 1, 0, 0])))
+
+    # 4. Multidimensional Indexing
+    arr_2d = GF([[1, 0], [0, 1]])
+    arr_2d = np.packbits(arr_2d)
+    arr_2d[0, 1] = GF(1)
+    assert arr_2d[0, 1] == GF(1)
+
+    arr_2d[:, 1] = GF([0, 0])
+    assert np.array_equal(arr_2d[:, 1], GF([0, 0]))
+
+    # 5. Boolean Indexing
+    arr = GF([1, 0, 1, 1])
+    arr = np.packbits(arr)
+    mask = np.array([True, False, True, False])
+    arr[mask] = GF(0)
+    assert np.array_equal(arr, np.packbits(GF([0, 0, 0, 1])))
+
+    # 6. Fancy Indexing
+    arr = GF([1, 0, 1, 1])
+    arr = np.packbits(arr)
+    indices = [0, 2, 3]
+    arr[indices] = GF([0, 0, 0])
+    assert np.array_equal(arr, np.packbits(GF([0, 0, 0, 0])))
+
+    # 7. Ellipsis
+    arr_3d = GF(np.random.randint(0, 2, (2, 3, 4)))
+    arr_3d = np.packbits(arr_3d)
+    arr_3d[0, ..., 1] = GF(1)
+    assert np.array_equal(arr_3d[0, :, 1], GF([1, 1, 1]))
+
+    # 8. Indexing with slice objects
+    arr = GF([1, 0, 1, 1])
+    arr = np.packbits(arr)
+    s = slice(1, 3)
+    arr[s] = GF([0, 0])
+    assert np.array_equal(arr, np.packbits(GF([1, 0, 0, 1])))
+
+    # 9. Using np.newaxis (reshaped array assignment)
+    arr = GF([1, 0, 1, 1])
+    arr = np.packbits(arr)
+    reshaped = arr[:, np.newaxis]  # should this be using arr's data (as would be the case without packbits) or a new array?
+    reshaped = np.packbits(reshaped)
+    reshaped[:, 0] = GF([0, 0, 0, 0])
+    # assert np.array_equal(arr, np.packbits(GF([0, 0, 0, 0])))
+    assert np.array_equal(reshaped, np.packbits(GF([[0], [0], [0], [0]])))
+
+    # 10. Indexing with np.ix_
+    arr_2d = GF([[1, 0], [0, 1]])
+    arr_2d = np.packbits(arr_2d)
+    row_indices = np.array([0, 1])
+    col_indices = np.array([0, 1])
+    arr_2d[np.ix_(row_indices, col_indices)] = GF([[0, 0], [0, 0]])
+    assert np.array_equal(arr_2d, np.packbits(GF([[0, 0], [0, 0]])))
+
+    print("All set-indexing tests passed.")
+
+
 if __name__ == "__main__":
     test_galois_array_indexing()
+    test_galois_array_setting()