From 7054295b90e7bda9737916a7ea069928113534e8 Mon Sep 17 00:00:00 2001 From: Amir Ebrahimi Date: Fri, 13 Dec 2024 16:52:31 -0800 Subject: [PATCH] Override np.packbits and np.unpackbits --- src/galois/_domains/_function.py | 9 ++- src/galois/_fields/_gf2.py | 96 ++++++++++++++++++++------------ 2 files changed, 66 insertions(+), 39 deletions(-) diff --git a/src/galois/_domains/_function.py b/src/galois/_domains/_function.py index eaf7d4a9f..c206cf707 100644 --- a/src/galois/_domains/_function.py +++ b/src/galois/_domains/_function.py @@ -304,8 +304,6 @@ class FunctionMixin(np.ndarray, metaclass=ArrayMeta): _UNSUPPORTED_FUNCTIONS = [ # Unary - np.packbits, - np.unpackbits, np.unwrap, np.around, np.round, @@ -334,6 +332,8 @@ class FunctionMixin(np.ndarray, metaclass=ArrayMeta): np.convolve: "_convolve", np.fft.fft: "_fft", np.fft.ifft: "_ifft", + np.packbits: "_packbits", + np.unpackbits: "_unpackbits", } _convolve: Function @@ -353,7 +353,10 @@ def __array_function__(self, func, types, args, kwargs): field = type(self) if func in field._OVERRIDDEN_FUNCTIONS: - output = getattr(field, field._OVERRIDDEN_FUNCTIONS[func])(*args, **kwargs) + try: + output = getattr(field, field._OVERRIDDEN_FUNCTIONS[func])(*args, **kwargs) + except AttributeError: + output = super().__array_function__(func, types, args, kwargs) elif func in field._UNSUPPORTED_FUNCTIONS: raise NotImplementedError( diff --git a/src/galois/_fields/_gf2.py b/src/galois/_fields/_gf2.py index 2ef21c5c3..72e91659e 100644 --- a/src/galois/_fields/_gf2.py +++ b/src/galois/_fields/_gf2.py @@ -5,7 +5,7 @@ from __future__ import annotations import numpy as np -from typing_extensions import Literal, Self +from typing_extensions import Literal, Self, Optional from .._domains._lookup import ( add_ufunc, @@ -84,6 +84,30 @@ class sqrt(sqrt_ufunc): def implementation(self, a: FieldArray) -> FieldArray: return a.copy() +def packbits(a, axis=None, bitorder='big'): + if isinstance(a, GF2BP): + return a + + if not isinstance(a, GF2): + 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]) + return packed + + +def unpackbits(a, axis=None, count=None, bitorder='big'): + if isinstance(a, GF2): + return a + + if not isinstance(a, GF2BP): + raise TypeError("Bit-unpacking is only supported on instances of GF2BP.") + + if axis is None: + axis = -1 + + return GF2(np.unpackbits(a.view(np.ndarray), axis=axis, count=a._axis_count if count is None else count, bitorder=bitorder)) + class UFuncMixin_2_1(UFuncMixin): """ @@ -101,6 +125,8 @@ def __init_subclass__(cls) -> None: cls._power = power(cls) cls._log = log(cls) cls._sqrt = sqrt(cls) + cls._packbits = packbits + cls._unpackbits = unpackbits class add_ufunc_bitpacked(add_ufunc): @@ -110,7 +136,7 @@ class add_ufunc_bitpacked(add_ufunc): def __call__(self, ufunc, method, inputs, kwargs, meta): output = super().__call__(ufunc, method, inputs, kwargs, meta) - output._unpacked_shape = inputs[0]._unpacked_shape + output._axis_count = inputs[0]._axis_count return output @@ -121,7 +147,7 @@ class subtract_ufunc_bitpacked(subtract_ufunc): def __call__(self, ufunc, method, inputs, kwargs, meta): output = super().__call__(ufunc, method, inputs, kwargs, meta) - output._unpacked_shape = inputs[0]._unpacked_shape + output._axis_count = inputs[0]._axis_count return output @@ -132,7 +158,7 @@ class multiply_ufunc_bitpacked(multiply_ufunc): def __call__(self, ufunc, method, inputs, kwargs, meta): output = super().__call__(ufunc, method, inputs, kwargs, meta) - output._unpacked_shape = inputs[0]._unpacked_shape + output._axis_count = inputs[0]._axis_count return output @@ -143,7 +169,7 @@ class divide_ufunc_bitpacked(divide): def __call__(self, ufunc, method, inputs, kwargs, meta): output = super().__call__(ufunc, method, inputs, kwargs, meta) - output._unpacked_shape = inputs[0]._unpacked_shape + output._axis_count = inputs[0]._axis_count return output @@ -157,16 +183,16 @@ def __call__(self, ufunc, method, inputs, kwargs, meta): assert isinstance(a, GF2BP) and isinstance(b, GF2BP) - # bit-packed matrices have rows packed by default, so unpack the second operand and repack to columns + # bit-packed matrices have columns packed by default, so unpack the second operand and repack to rows field = self.field - unpacked_shape = b._unpacked_shape + row_axis_count = b.shape[0] b = field._view( np.packbits( - np.unpackbits(b.view(np.ndarray), axis=-1, count=b._unpacked_shape[-1]), + np.unpackbits(b.view(np.ndarray), axis=-1, count=b._axis_count), axis=0, ) ) - b._unpacked_shape = unpacked_shape + b._axis_count = row_axis_count # Make sure the inner dimensions match (e.g. (M, N) x (N, P) -> (M, P)) assert a.shape[-1] == b.shape[0] @@ -182,9 +208,11 @@ def __call__(self, ufunc, method, inputs, kwargs, meta): # matrix-matrix multiplication output = GF2.Zeros(final_shape) for i in range(b.shape[-1]): + # TODO: Include alternate path for numpy < v2 + # output[:, i] = np.bitwise_xor.reduce(np.unpackbits((a & b[:, i]).view(np.ndarray), axis=-1), axis=-1) output[:, i] = np.bitwise_xor.reduce(np.bitwise_count((a & b[:, i]).view(np.ndarray)), axis=-1) % 2 output = field._view(np.packbits(output.view(np.ndarray), axis=-1)) - output._unpacked_shape = final_shape + output._axis_count = final_shape[-1] return output @@ -192,7 +220,6 @@ def __call__(self, ufunc, method, inputs, kwargs, meta): def not_implemented(*args, **kwargs): return NotImplemented - class UFuncMixin_2_1_BitPacked(UFuncMixin): """ A mixin class that provides explicit calculation arithmetic for GF(2). @@ -209,6 +236,8 @@ def __init_subclass__(cls) -> None: cls._power = power(cls) cls._log = log(cls) cls._sqrt = sqrt(cls) + cls._packbits = packbits + cls._unpackbits = unpackbits @classmethod def _assign_ufuncs(cls): @@ -268,12 +297,6 @@ class factory :func:`~galois.GF`. galois-fields """ - def astype(self, dtype, **kwargs): - if dtype is GF2BP: - return GF2BP(self) # bits are packed in initialization - - return super().astype(dtype, **kwargs) - @export class GF2BP( @@ -322,39 +345,40 @@ class factory :func:`~galois.GF`. def __new__( cls, x: ElementLike | ArrayLike, + axis_element_count: Optional[int] = None, dtype: DTypeLike | None = None, copy: bool = True, order: Literal["K", "A", "C", "F"] = "K", ndmin: int = 0, ) -> Self: - if isinstance(x, np.ndarray): - dtype = cls._get_dtype(dtype) + # axis_element_count is required, but by making it optional it allows us to catch uses of the class that are not + # supported (e.g. Random) + if isinstance(x, np.ndarray) and axis_element_count is not None: + # NOTE: I'm not sure that we want to change the dtype specifically for the bit-packed version or how we verify + # dtype = cls._get_dtype(dtype) + # x = cls._verify_array_like_types_and_values(x) - x = cls._verify_array_like_types_and_values(x) - array = cls._view(np.packbits(np.array(x, dtype=dtype, copy=copy, order=order, ndmin=ndmin).view(np.ndarray), axis=-1)) - array._unpacked_shape = x.shape + array = cls._view(np.array(x, dtype=dtype, copy=copy, order=order, ndmin=ndmin)) + array._axis_count = axis_element_count - # Perform view without verification since the elements were verified in _verify_array_like_types_and_values() return array raise NotImplementedError( "GF2BP is a custom bit-packed GF2 class with limited functionality. " "If you were using an alternate constructor (e.g. Random), then use the GF2 class and convert it to the " - "bit-packed version by using `.astype(GF2BP)`." + "bit-packed version by using `np.packbits`." ) - def astype(self, dtype, **kwargs): - if dtype is GF2: - return GF2( - np.unpackbits( - self.view(np.ndarray), - axis=-1, - count=self._unpacked_shape[-1], - ) - ) - - return super().astype(dtype, **kwargs) - + def __init__( + self, + x: ElementLike | ArrayLike, + axis_element_count: Optional[int] = None, + dtype: DTypeLike | None = None, + copy: bool = True, + order: Literal["K", "A", "C", "F"] = "K", + ndmin: int = 0, + ): + pass GF2._default_ufunc_mode = "jit-calculate" GF2._ufunc_modes = ["jit-calculate", "python-calculate"]