diff --git a/benchmarks/benchmarks/bench_ufunc.py b/benchmarks/benchmarks/bench_ufunc.py index f3a600a3279a..42f51ed85eaf 100644 --- a/benchmarks/benchmarks/bench_ufunc.py +++ b/benchmarks/benchmarks/bench_ufunc.py @@ -7,7 +7,7 @@ ufuncs = ['abs', 'absolute', 'add', 'arccos', 'arccosh', 'arcsin', 'arcsinh', - 'arctan', 'arctan2', 'arctanh', 'bitwise_and', 'bitwise_not', + 'arctan', 'arctan2', 'arctanh', 'bitwise_and', 'bitwise_count', 'bitwise_not', 'bitwise_or', 'bitwise_xor', 'cbrt', 'ceil', 'conj', 'conjugate', 'copysign', 'cos', 'cosh', 'deg2rad', 'degrees', 'divide', 'divmod', 'equal', 'exp', 'exp2', 'expm1', 'fabs', 'float_power', 'floor', @@ -310,7 +310,7 @@ def time_astype(self, typeconv): class UFuncSmall(Benchmark): """ Benchmark for a selection of ufuncs on a small arrays and scalars - Since the arrays and scalars are small, we are benchmarking the overhead + Since the arrays and scalars are small, we are benchmarking the overhead of the numpy ufunc functionality """ params = ['abs', 'sqrt', 'cos'] @@ -327,7 +327,7 @@ def setup(self, ufuncname): self.array_int_3 = np.array([1, 2, 3]) self.float64 = np.float64(1.1) self.python_float = 1.1 - + def time_ufunc_small_array(self, ufuncname): self.f(self.array_5) @@ -342,7 +342,7 @@ def time_ufunc_numpy_scalar(self, ufuncname): def time_ufunc_python_float(self, ufuncname): self.f(self.python_float) - + class Custom(Benchmark): def setup(self): @@ -565,7 +565,7 @@ def setup(self): self.b32 = np.random.rand(N).astype(np.float32) self.a64 = np.random.rand(N).astype(np.float64) self.b64 = np.random.rand(N).astype(np.float64) - + def time_pow_32(self): np.power(self.a32, self.b32) diff --git a/benchmarks/benchmarks/bench_ufunc_strides.py b/benchmarks/benchmarks/bench_ufunc_strides.py index 70c076dd7982..929615f0fb83 100644 --- a/benchmarks/benchmarks/bench_ufunc_strides.py +++ b/benchmarks/benchmarks/bench_ufunc_strides.py @@ -164,7 +164,7 @@ class UnaryIntContig(_AbstractUnary): [getattr(np, uf) for uf in ( 'positive', 'square', 'reciprocal', 'conjugate', 'logical_not', 'invert', 'isnan', 'isinf', 'isfinite', - 'absolute', 'sign' + 'absolute', 'sign', 'bitwise_count' )], [1], [1], ['b', 'B', 'h', 'H', 'i', 'I', 'l', 'L', 'q', 'Q'] diff --git a/doc/release/upcoming_changes/19355.new_feature.rst b/doc/release/upcoming_changes/19355.new_feature.rst new file mode 100644 index 000000000000..70acc1d23012 --- /dev/null +++ b/doc/release/upcoming_changes/19355.new_feature.rst @@ -0,0 +1,13 @@ +`np.bitwise_count` to compute the number of 1-bits in an integer array +---------------------------------------------------------------------- + +This new function counts the number of 1-bits in a number. +`np.bitwise_count` works on all the numpy integer types and +integer-like objects. + +.. code-block:: python + + >>> a = np.array([2**i - 1 for i in range(16)]) + >>> np.bitwise_count(a) + array([ 0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15], + dtype=uint8) diff --git a/doc/source/reference/routines.math.rst b/doc/source/reference/routines.math.rst index 6d80d5ad73fe..fb7fb5410273 100644 --- a/doc/source/reference/routines.math.rst +++ b/doc/source/reference/routines.math.rst @@ -181,3 +181,5 @@ Miscellaneous real_if_close interp + + bitwise_count diff --git a/numpy/__init__.py b/numpy/__init__.py index 1128e5ee525b..fce0a00b9a55 100644 --- a/numpy/__init__.py +++ b/numpy/__init__.py @@ -127,8 +127,8 @@ around, array, array2string, array_equal, array_equiv, array_repr, array_str, asanyarray, asarray, ascontiguousarray, asfortranarray, atleast_1d, atleast_2d, atleast_3d, base_repr, binary_repr, - bitwise_and, bitwise_not, bitwise_or, bitwise_xor, block, bool_, - broadcast, busday_count, busday_offset, busdaycalendar, byte, bytes_, + bitwise_and, bitwise_count, bitwise_not, bitwise_or, bitwise_xor, block, + bool_, broadcast, busday_count, busday_offset, busdaycalendar, byte, bytes_, can_cast, cbrt, cdouble, ceil, character, choose, clip, clongdouble, complexfloating, compress, concatenate, conj, conjugate, convolve, copysign, copyto, correlate, cos, cosh, count_nonzero, cross, csingle, diff --git a/numpy/__init__.pyi b/numpy/__init__.pyi index 0591b859ac8f..418bdf6143d4 100644 --- a/numpy/__init__.pyi +++ b/numpy/__init__.pyi @@ -2467,6 +2467,17 @@ class ndarray(_ArrayOrScalarCommon, Generic[_ShapeType, _DType_co]): def __dlpack__(self: NDArray[number[Any]], *, stream: None = ...) -> _PyCapsule: ... def __dlpack_device__(self) -> tuple[int, L[0]]: ... + def bitwise_count( + self, + out: None | NDArray[Any] = ..., + *, + where: _ArrayLikeBool_co = ..., + casting: _CastingKind = ..., + order: _OrderKACF = ..., + dtype: DTypeLike = ..., + subok: bool = ..., + ) -> NDArray[Any]: ... + # Keep `dtype` at the bottom to avoid name conflicts with `np.dtype` @property def dtype(self) -> _DType_co: ... @@ -2614,6 +2625,17 @@ class generic(_ArrayOrScalarCommon): self: _ScalarType, *shape: SupportsIndex, order: _OrderACF = ... ) -> NDArray[_ScalarType]: ... + def bitwise_count( + self, + out: None | NDArray[Any] = ..., + *, + where: _ArrayLikeBool_co = ..., + casting: _CastingKind = ..., + order: _OrderKACF = ..., + dtype: DTypeLike = ..., + subok: bool = ..., + ) -> Any: ... + def squeeze( self: _ScalarType, axis: None | L[0] | tuple[()] = ... ) -> _ScalarType: ... @@ -3138,6 +3160,7 @@ arctan2: _UFunc_Nin2_Nout1[L['arctan2'], L[5], None] arctan: _UFunc_Nin1_Nout1[L['arctan'], L[8], None] arctanh: _UFunc_Nin1_Nout1[L['arctanh'], L[8], None] bitwise_and: _UFunc_Nin2_Nout1[L['bitwise_and'], L[12], L[-1]] +bitwise_count: _UFunc_Nin1_Nout1[L['bitwise_count'], L[11], None] bitwise_not: _UFunc_Nin1_Nout1[L['invert'], L[12], None] bitwise_or: _UFunc_Nin2_Nout1[L['bitwise_or'], L[12], L[0]] bitwise_xor: _UFunc_Nin2_Nout1[L['bitwise_xor'], L[12], L[0]] diff --git a/numpy/core/_methods.py b/numpy/core/_methods.py index 9675f9822aaa..296dfec1b02f 100644 --- a/numpy/core/_methods.py +++ b/numpy/core/_methods.py @@ -21,6 +21,7 @@ umr_minimum = um.minimum.reduce umr_sum = um.add.reduce umr_prod = um.multiply.reduce +umr_bitwise_count = um.bitwise_count umr_any = um.logical_or.reduce umr_all = um.logical_and.reduce @@ -236,3 +237,8 @@ def _dump(self, file, protocol=2): def _dumps(self, protocol=2): return pickle.dumps(self, protocol=protocol) + +def _bitwise_count(a, out=None, *, where=True, casting='same_kind', + order='K', dtype=None, subok=True): + return umr_bitwise_count(a, out, where=where, casting=casting, + order=order, dtype=dtype, subok=subok) diff --git a/numpy/core/code_generators/generate_umath.py b/numpy/core/code_generators/generate_umath.py index 9cb943ac19fe..7b30986739ab 100644 --- a/numpy/core/code_generators/generate_umath.py +++ b/numpy/core/code_generators/generate_umath.py @@ -1122,6 +1122,13 @@ def english_upper(s): TD(ints), TD('O', f='npy_ObjectLCM'), ), +'bitwise_count': + Ufunc(1, 1, None, + docstrings.get('numpy.core.umath.bitwise_count'), + None, + TD(ints, dispatch=[('loops_autovec', ints)], out='B'), + TD(P, f='bit_count'), + ), 'matmul' : Ufunc(2, 1, None, docstrings.get('numpy.core.umath.matmul'), diff --git a/numpy/core/code_generators/ufunc_docstrings.py b/numpy/core/code_generators/ufunc_docstrings.py index 437901c19470..f0a4a80d7036 100644 --- a/numpy/core/code_generators/ufunc_docstrings.py +++ b/numpy/core/code_generators/ufunc_docstrings.py @@ -4214,3 +4214,41 @@ def add_newdoc(place, name, doc): array([ 0, 20, 20, 60, 20, 20]) """) + +add_newdoc('numpy.core.umath', 'bitwise_count', + """ + Computes the number of 1-bits in the absolute value of ``x``. + Analogous to the builtin `int.bit_count` or ``popcount`` in C++. + + Parameters + ---------- + x : array_like, unsigned int + Input array. + $PARAMS + + Returns + ------- + y : ndarray + The corresponding number of 1-bits in the input. + Returns uint8 for all integer types + $OUT_SCALAR_1 + + References + ---------- + .. [1] https://graphics.stanford.edu/~seander/bithacks.html#CountBitsSetParallel + + .. [2] Wikipedia, "Hamming weight", + https://en.wikipedia.org/wiki/Hamming_weight + + .. [3] http://aggregate.ee.engr.uky.edu/MAGIC/#Population%20Count%20(Ones%20Count) + + Examples + -------- + >>> np.bitwise_count(1023) + 10 + >>> a = np.array([2**i - 1 for i in range(16)]) + >>> np.bitwise_count(a) + array([ 0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15], + dtype=uint8) + + """) diff --git a/numpy/core/src/npymath/npy_math_internal.h.src b/numpy/core/src/npymath/npy_math_internal.h.src index c2ef06aafe17..382a144537de 100644 --- a/numpy/core/src/npymath/npy_math_internal.h.src +++ b/numpy/core/src/npymath/npy_math_internal.h.src @@ -678,7 +678,6 @@ npy_rshift@u@@c@(npy_@u@@type@ a, npy_@u@@type@ b) /**end repeat1**/ /**end repeat**/ - #define __popcnt32 __popcnt /**begin repeat * diff --git a/numpy/core/src/umath/loops.c.src b/numpy/core/src/umath/loops.c.src index 139b8c2a48d2..c58b3c17c2ba 100644 --- a/numpy/core/src/umath/loops.c.src +++ b/numpy/core/src/umath/loops.c.src @@ -438,6 +438,7 @@ NPY_NO_EXPORT void *((@type@ *)op1) = 1; } } + /**begin repeat1 * Arithmetic * #kind = add, subtract, multiply, bitwise_and, bitwise_or, bitwise_xor# diff --git a/numpy/core/src/umath/loops.h.src b/numpy/core/src/umath/loops.h.src index cce73aff8504..17c16c227bb0 100644 --- a/numpy/core/src/umath/loops.h.src +++ b/numpy/core/src/umath/loops.h.src @@ -147,7 +147,7 @@ NPY_CPU_DISPATCH_DECLARE(NPY_NO_EXPORT void @TYPE@_@kind@, * subtract, multiply, bitwise_and, bitwise_or, bitwise_xor, * left_shift, right_shift, logical_and, logical_or, * logical_xor, isnan, isinf, isfinite, - * absolute, sign# + * absolute, sign, bitwise_count# */ NPY_CPU_DISPATCH_DECLARE(NPY_NO_EXPORT void @TYPE@_@kind@, (char **args, npy_intp const *dimensions, npy_intp const *steps, void *NPY_UNUSED(func))) @@ -208,6 +208,7 @@ NPY_NO_EXPORT void @S@@TYPE@_lcm(char **args, npy_intp const *dimensions, npy_intp const *steps, void *NPY_UNUSED(func)); /**end repeat2**/ + /**end repeat1**/ /**end repeat**/ diff --git a/numpy/core/src/umath/loops_autovec.dispatch.c.src b/numpy/core/src/umath/loops_autovec.dispatch.c.src index 1656f8e04104..6ccafe577c72 100644 --- a/numpy/core/src/umath/loops_autovec.dispatch.c.src +++ b/numpy/core/src/umath/loops_autovec.dispatch.c.src @@ -103,6 +103,13 @@ NPY_NO_EXPORT void NPY_CPU_DISPATCH_CURFX(@TYPE@_right_shift) } #endif } + +NPY_NO_EXPORT void NPY_CPU_DISPATCH_CURFX(@TYPE@_bitwise_count) +(char **args, npy_intp const *dimensions, npy_intp const *steps, void *NPY_UNUSED(func)) +{ + UNARY_LOOP_FAST(@type@, npy_ubyte, *out = npy_popcount@c@(in)); +} + /**end repeat**/ /* @@ -285,3 +292,6 @@ NPY_NO_EXPORT void NPY_CPU_DISPATCH_CURFX(@TYPE@_isinf) NPY_CPU_DISPATCH_CURFX(ULONGLONG_isinf)(args, dimensions, steps, func); } /**end repeat**/ + + + diff --git a/numpy/core/tests/test_ufunc.py b/numpy/core/tests/test_ufunc.py index 5049b0b69ed2..a200b254409e 100644 --- a/numpy/core/tests/test_ufunc.py +++ b/numpy/core/tests/test_ufunc.py @@ -33,6 +33,9 @@ if isinstance(obj, np.ufunc)] UNARY_OBJECT_UFUNCS = [uf for uf in UNARY_UFUNCS if "O->O" in uf.types] +# Remove functions that do not support `floats` +UNARY_OBJECT_UFUNCS.remove(getattr(np, 'bitwise_count')) + class TestUfuncKwargs: def test_kwarg_exact(self): diff --git a/numpy/core/tests/test_umath.py b/numpy/core/tests/test_umath.py index 79869c7c11a1..55950b5831dd 100644 --- a/numpy/core/tests/test_umath.py +++ b/numpy/core/tests/test_umath.py @@ -2610,7 +2610,15 @@ def test_reduce(self): class TestBitwiseUFuncs: - bitwise_types = [np.dtype(c) for c in '?' + 'bBhHiIlLqQ' + 'O'] + _all_ints_bits = [ + np.dtype(c).itemsize * 8 for c in np.typecodes["AllInteger"]] + bitwise_types = [ + np.dtype(c) for c in '?' + np.typecodes["AllInteger"] + 'O'] + bitwise_bits = [ + 2, # boolean type + *_all_ints_bits, # All integers + max(_all_ints_bits) + 1, # Object_ type + ] def test_values(self): for dt in self.bitwise_types: @@ -2691,6 +2699,30 @@ def test_reduction(self): btype = np.array([True], dtype=object) assert_(type(f.reduce(btype)) is bool, msg) + @pytest.mark.parametrize("input_dtype_obj, bitsize", + zip(bitwise_types, bitwise_bits)) + def test_bitwise_count(self, input_dtype_obj, bitsize): + input_dtype = input_dtype_obj.type + + # bitwise_count is only in-built in 3.10+ + if sys.version_info < (3, 10) and input_dtype == np.object_: + pytest.skip("Required Python >=3.10") + + for i in range(1, bitsize): + num = 2**i - 1 + msg = f"bitwise_count for {num}" + assert i == np.bitwise_count(input_dtype(num)), msg + if np.issubdtype( + input_dtype, np.signedinteger) or input_dtype == np.object_: + assert i == np.bitwise_count(input_dtype(-num)), msg + + a = np.array([2**i-1 for i in range(1, bitsize)], dtype=input_dtype) + bitwise_count_a = np.bitwise_count(a) + expected = np.arange(1, bitsize, dtype=input_dtype) + + msg = f"array bitwise_count for {input_dtype}" + assert all(bitwise_count_a == expected), msg + class TestInt: def test_logical_not(self): diff --git a/numpy/core/tests/test_umath_accuracy.py b/numpy/core/tests/test_umath_accuracy.py index 6ee4d2fee080..bb76d9c20bbd 100644 --- a/numpy/core/tests/test_umath_accuracy.py +++ b/numpy/core/tests/test_umath_accuracy.py @@ -11,7 +11,10 @@ UNARY_UFUNCS = [obj for obj in np.core.umath.__dict__.values() if isinstance(obj, np.ufunc)] UNARY_OBJECT_UFUNCS = [uf for uf in UNARY_UFUNCS if "O->O" in uf.types] + +# Remove functions that do not support `floats` UNARY_OBJECT_UFUNCS.remove(getattr(np, 'invert')) +UNARY_OBJECT_UFUNCS.remove(getattr(np, 'bitwise_count')) IS_AVX = __cpu_features__.get('AVX512F', False) or \ (__cpu_features__.get('FMA3', False) and __cpu_features__.get('AVX2', False)) diff --git a/numpy/core/umath.py b/numpy/core/umath.py index 757bf1e59c9d..bf5ee90dbe40 100644 --- a/numpy/core/umath.py +++ b/numpy/core/umath.py @@ -19,7 +19,7 @@ 'absolute', 'add', 'arccos', 'arccosh', 'arcsin', 'arcsinh', 'arctan', 'arctan2', 'arctanh', 'bitwise_and', 'bitwise_or', 'bitwise_xor', 'cbrt', 'ceil', 'conj', - 'conjugate', 'copysign', 'cos', 'cosh', 'deg2rad', 'degrees', 'divide', + 'conjugate', 'copysign', 'cos', 'cosh', 'bitwise_count', 'deg2rad', 'degrees', 'divide', 'divmod', 'e', 'equal', 'euler_gamma', 'exp', 'exp2', 'expm1', 'fabs', 'floor', 'floor_divide', 'float_power', 'fmax', 'fmin', 'fmod', 'frexp', 'frompyfunc', 'gcd', 'greater', 'greater_equal', 'heaviside', diff --git a/numpy/matrixlib/tests/test_defmatrix.py b/numpy/matrixlib/tests/test_defmatrix.py index 8fcb74759069..1aa88fccb5f6 100644 --- a/numpy/matrixlib/tests/test_defmatrix.py +++ b/numpy/matrixlib/tests/test_defmatrix.py @@ -286,7 +286,7 @@ def test_instance_methods(self): 'partition', 'argpartition', 'newbyteorder', 'take', 'tofile', 'tolist', 'tostring', 'tobytes', 'all', 'any', 'sum', 'argmax', 'argmin', 'min', 'max', 'mean', 'var', 'ptp', - 'prod', 'std', 'ctypes', 'itemset', + 'prod', 'std', 'ctypes', 'itemset', 'bitwise_count', ] for attrib in dir(a): if attrib.startswith('_') or attrib in excluded_methods: diff --git a/numpy/typing/tests/data/reveal/ufuncs.pyi b/numpy/typing/tests/data/reveal/ufuncs.pyi index 5f7d99efd12d..28e189411802 100644 --- a/numpy/typing/tests/data/reveal/ufuncs.pyi +++ b/numpy/typing/tests/data/reveal/ufuncs.pyi @@ -9,6 +9,7 @@ if sys.version_info >= (3, 11): else: from typing_extensions import assert_type +i8: np.int64 f8: np.float64 AR_f8: npt.NDArray[np.float64] AR_i8: npt.NDArray[np.int64] @@ -74,3 +75,14 @@ assert_type(np.matmul.signature, Literal["(n?,k),(k,m?)->(n?,m?)"]) assert_type(np.matmul.identity, None) assert_type(np.matmul(AR_f8, AR_f8), Any) assert_type(np.matmul(AR_f8, AR_f8, axes=[(0, 1), (0, 1), (0, 1)]), Any) + +assert_type(np.bitwise_count.__name__, Literal['bitwise_count']) +assert_type(np.bitwise_count.ntypes, Literal[11]) +assert_type(np.bitwise_count.identity, None) +assert_type(np.bitwise_count.nin, Literal[1]) +assert_type(np.bitwise_count.nout, Literal[1]) +assert_type(np.bitwise_count.nargs, Literal[2]) +assert_type(np.bitwise_count.signature, None) +assert_type(np.bitwise_count.identity, None) +assert_type(np.bitwise_count(i8), Any) +assert_type(np.bitwise_count(AR_i8), npt.NDArray[Any])