Thanks to visit codestin.com
Credit goes to github.com

Skip to content

Commit 015a1a2

Browse files
ganesh-k13ganesh-k13
committed
ENH: Added countbits (popcount)
ENH, DOC: Added countbits (popcount) ENH: Popcount implementation ENH: Add popcount to umath ENH: Added countbits (popcount) to umath `__all__` ENH: Refined popcount logic DOC: Added `bit_count` Co-authored-by: Eric Wieser <[email protected]> MAINT: Renamed `countbits` to `bit_count` MAINT: Fixed 4 1s magic number DOC: Added `popcount` to docstring ENH: Added bit_count annotations ENH: Added GNU/CLANG popcount DOC: Added `popcount` language example ENH, BUG: Moved `bitcount` to npy_math.h as `popcount` | Fixed final right shift ENH: Enable `popcount` for signed TST: Tests for `bit_count` BUG, DOC: (BUG) Added missing typecast causing an unwanted upcast (DOC) Added more details on `popcount` implementation MAINT, BUG: (MAINT) Refined `popcount` TC to use typecode (BUG) Fixed ufunc.ntypes to include signed ints ENH: Added windows builtin support ENH: Added `popcount` implementation for big python ints natively [1/2] `popcount` object loop changes ENH: Object loop for `bit_count` [2/2] `popcount` object loop changes TST: Refined `bit_count` tests and added object type ENH: Added `bit_count` to `np.int*` DOC: Added `np.bit_count` (#19355) MAINT: Various linting and minor fixes: 1. Fixed passing all args to _internals umath bitcount. Note: We use kwargs here that might hinder performance 2. Fixed linting errors. 3. Improved verbosity of logs 4. Made a generic TO_BITS_LEN macro to accomdate more length based functions in future BENCH: Added bit_count (popcount) MAINT: Style nits | Added signed case DOC, MAINT: Improved example ENH: Added annotations for bit_count TST: Added annotations tests for bit_count MAINT: Fixed linting errors MAINT: Moved Magic constants to npy_math_internal MAINT: Remove python implementation | Added 3.10 check to tests DOC: Added abs value usage to doc MAINT: Resolved merge conflicts
1 parent 4e667cb commit 015a1a2

File tree

17 files changed

+172
-8
lines changed

17 files changed

+172
-8
lines changed

‎benchmarks/benchmarks/bench_ufunc.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -4,7 +4,7 @@
44

55

66
ufuncs = ['abs', 'absolute', 'add', 'arccos', 'arccosh', 'arcsin', 'arcsinh',
7-
'arctan', 'arctan2', 'arctanh', 'bitwise_and', 'bitwise_not',
7+
'arctan', 'arctan2', 'arctanh', 'bit_count', 'bitwise_and', 'bitwise_not',
88
'bitwise_or', 'bitwise_xor', 'cbrt', 'ceil', 'conj', 'conjugate',
99
'copysign', 'cos', 'cosh', 'deg2rad', 'degrees', 'divide', 'divmod',
1010
'equal', 'exp', 'exp2', 'expm1', 'fabs', 'float_power', 'floor',

‎doc/release/upcoming_changes/19355.new_feature.rst

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,12 @@
1+
`np.bit_count` to compute the number of 1-bits in an integer
2+
------------------------------------------------------------
3+
4+
This new function counts the number of 1-bits in a number.
5+
These work on all the numpy integer types, as well as the
6+
builtin arbitrary-precision `Decimal` and `long` types.
7+
8+
.. code-block:: python
9+
10+
>>> a = np.array([2**i - 1 for i in range(16)])
11+
>>> np.bit_count(a)
12+
array([ 0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15])

‎doc/source/reference/routines.math.rst

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -180,3 +180,5 @@ Miscellaneous
180180
real_if_close
181181

182182
interp
183+
184+
bit_count

‎numpy/__init__.pyi

Lines changed: 23 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2471,6 +2471,17 @@ class ndarray(_ArrayOrScalarCommon, Generic[_ShapeType, _DType_co]):
24712471
def __dlpack__(self: NDArray[number[Any]], *, stream: None = ...) -> _PyCapsule: ...
24722472
def __dlpack_device__(self) -> tuple[int, L[0]]: ...
24732473

2474+
def bit_count(
2475+
self,
2476+
out: None | NDArray[Any] = ...,
2477+
*,
2478+
where: _ArrayLikeBool_co = ...,
2479+
casting: _CastingKind = ...,
2480+
order: _OrderKACF = ...,
2481+
dtype: DTypeLike = ...,
2482+
subok: bool = ...,
2483+
) -> NDArray[Any]: ...
2484+
24742485
# Keep `dtype` at the bottom to avoid name conflicts with `np.dtype`
24752486
@property
24762487
def dtype(self) -> _DType_co: ...
@@ -2614,6 +2625,17 @@ class generic(_ArrayOrScalarCommon):
26142625
self: _ScalarType, *shape: SupportsIndex, order: _OrderACF = ...
26152626
) -> ndarray[Any, _dtype[_ScalarType]]: ...
26162627

2628+
def bit_count(
2629+
self,
2630+
out: None | NDArray[Any] = ...,
2631+
*,
2632+
where: _ArrayLikeBool_co = ...,
2633+
casting: _CastingKind = ...,
2634+
order: _OrderKACF = ...,
2635+
dtype: DTypeLike = ...,
2636+
subok: bool = ...,
2637+
) -> Any: ...
2638+
26172639
def squeeze(
26182640
self: _ScalarType, axis: L[0] | tuple[()] = ...
26192641
) -> _ScalarType: ...
@@ -3190,6 +3212,7 @@ arcsinh: _UFunc_Nin1_Nout1[L['arcsinh'], L[8], None]
31903212
arctan2: _UFunc_Nin2_Nout1[L['arctan2'], L[5], None]
31913213
arctan: _UFunc_Nin1_Nout1[L['arctan'], L[8], None]
31923214
arctanh: _UFunc_Nin1_Nout1[L['arctanh'], L[8], None]
3215+
bit_count: _UFunc_Nin1_Nout1[L['bit_count'], L[11], None]
31933216
bitwise_and: _UFunc_Nin2_Nout1[L['bitwise_and'], L[12], L[-1]]
31943217
bitwise_not: _UFunc_Nin1_Nout1[L['invert'], L[12], None]
31953218
bitwise_or: _UFunc_Nin2_Nout1[L['bitwise_or'], L[12], L[0]]

‎numpy/core/_methods.py

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,7 @@
1919
umr_minimum = um.minimum.reduce
2020
umr_sum = um.add.reduce
2121
umr_prod = um.multiply.reduce
22+
umr_bit_count = um.bit_count
2223
umr_any = um.logical_or.reduce
2324
umr_all = um.logical_and.reduce
2425

@@ -290,3 +291,8 @@ def _dump(self, file, protocol=2):
290291

291292
def _dumps(self, protocol=2):
292293
return pickle.dumps(self, protocol=protocol)
294+
295+
def _bit_count(a, out=None, *, where=True, casting='same_kind',
296+
order='K', dtype=None, subok=True):
297+
return umr_bit_count(a, out, where=where, casting=casting,
298+
order=order, dtype=dtype, subok=subok)

‎numpy/core/code_generators/generate_umath.py

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -997,6 +997,13 @@ def english_upper(s):
997997
TD(ints),
998998
TD('O', f='npy_ObjectLCM'),
999999
),
1000+
'bit_count':
1001+
Ufunc(1, 1, None,
1002+
docstrings.get('numpy.core.umath.bit_count'),
1003+
None,
1004+
TD(ints),
1005+
TD('O', f='npy_ObjectPopCount'),
1006+
),
10001007
'matmul' :
10011008
Ufunc(2, 1, None,
10021009
docstrings.get('numpy.core.umath.matmul'),

‎numpy/core/code_generators/ufunc_docstrings.py

Lines changed: 34 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -4200,3 +4200,37 @@ def add_newdoc(place, name, doc):
42004200
array([ 0, 20, 20, 60, 20, 20])
42014201
42024202
""")
4203+
4204+
add_newdoc('numpy.core.umath', 'bit_count',
4205+
"""
4206+
Computes the number of 1-bits in the absolute value of ``x``.
4207+
Analogous to the builtin `int.bit_count` or ``popcount`` in C++.
4208+
4209+
Parameters
4210+
----------
4211+
x : array_like, unsigned int
4212+
Input array.
4213+
$PARAMS
4214+
4215+
Returns
4216+
-------
4217+
y : ndarray
4218+
The corresponding number of 1-bits in the input.
4219+
$OUT_SCALAR_1
4220+
4221+
References
4222+
----------
4223+
.. [1] https://stackoverflow.com/a/109025/5671364
4224+
4225+
.. [2] Wikipedia, "Hamming weight",
4226+
https://en.wikipedia.org/wiki/Hamming_weight
4227+
4228+
Examples
4229+
--------
4230+
>>> np.bit_count(1023)
4231+
10
4232+
>>> a = np.array([2**i - 1 for i in range(16)])
4233+
>>> np.bit_count(a)
4234+
array([ 0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15])
4235+
4236+
""")

‎numpy/core/src/multiarray/methods.c

Lines changed: 9 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -354,6 +354,12 @@ array_ptp(PyArrayObject *self, PyObject *args, PyObject *kwds)
354354
NPY_FORWARD_NDARRAY_METHOD("_ptp");
355355
}
356356

357+
static PyObject *
358+
array_bit_count(PyArrayObject *self, PyObject *args, PyObject *kwds)
359+
{
360+
NPY_FORWARD_NDARRAY_METHOD("_bit_count");
361+
}
362+
357363

358364
static PyObject *
359365
array_swapaxes(PyArrayObject *self, PyObject *args)
@@ -3052,9 +3058,11 @@ NPY_NO_EXPORT PyMethodDef array_methods[] = {
30523058
{"__dlpack__",
30533059
(PyCFunction)array_dlpack,
30543060
METH_FASTCALL | METH_KEYWORDS, NULL},
3055-
30563061
{"__dlpack_device__",
30573062
(PyCFunction)array_dlpack_device,
30583063
METH_NOARGS, NULL},
3064+
{"bit_count",
3065+
(PyCFunction)array_bit_count,
3066+
METH_VARARGS | METH_KEYWORDS, NULL},
30593067
{NULL, NULL, 0, NULL} /* sentinel */
30603068
};

‎numpy/core/src/multiarray/scalartypes.c.src

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1566,7 +1566,7 @@ gentype_byteswap(PyObject *self, PyObject *args, PyObject *kwds)
15661566
* std, var, sum, cumsum, prod, cumprod, compress, sort, argsort,
15671567
* round, argmax, argmin, max, min, ptp, any, all, astype, resize,
15681568
* reshape, choose, tostring, tobytes, copy, searchsorted, view,
1569-
* flatten, ravel, squeeze#
1569+
* flatten, ravel, squeeze, bit_count#
15701570
*/
15711571
static PyObject *
15721572
gentype_@name@(PyObject *self, PyObject *args, PyObject *kwds)
@@ -2187,6 +2187,9 @@ static PyMethodDef gentype_methods[] = {
21872187
{"sum",
21882188
(PyCFunction)gentype_sum,
21892189
METH_VARARGS | METH_KEYWORDS, NULL},
2190+
{"bit_count",
2191+
(PyCFunction)gentype_bit_count,
2192+
METH_VARARGS | METH_KEYWORDS, NULL},
21902193
{"cumsum",
21912194
(PyCFunction)gentype_cumsum,
21922195
METH_VARARGS | METH_KEYWORDS, NULL},

‎numpy/core/src/npymath/npy_math_internal.h.src

Lines changed: 0 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -853,7 +853,6 @@ npy_rshift@u@@c@(npy_@u@@type@ a, npy_@u@@type@ b)
853853
/**end repeat1**/
854854
/**end repeat**/
855855

856-
857856
#define __popcnt32 __popcnt
858857
/**begin repeat
859858
*

‎numpy/core/src/umath/funcs.inc.src

Lines changed: 15 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -267,6 +267,21 @@ npy_ObjectClip(PyObject *arr, PyObject *min, PyObject *max) {
267267
return o;
268268
}
269269

270+
static PyObject *
271+
npy_ObjectPopCount(PyObject *obj) {
272+
PyObject *result = NULL;
273+
274+
/* Try to use inbuilt popcount if available */
275+
static PyObject *builtin_popcount_func = NULL;
276+
builtin_popcount_func = PyObject_GetAttrString(obj, "bit_count");
277+
278+
if (builtin_popcount_func != NULL) {
279+
result = PyObject_CallFunction(builtin_popcount_func, NULL);
280+
}
281+
282+
return result;
283+
}
284+
270285
/*
271286
*****************************************************************************
272287
** COMPLEX FUNCTIONS **

‎numpy/core/src/umath/loops.c.src

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -570,6 +570,15 @@ NPY_NO_EXPORT void
570570
UNARY_LOOP_FAST(@type@, @type@, *out = +in);
571571
}
572572

573+
NPY_NO_EXPORT void
574+
@TYPE@_bit_count(char **args, npy_intp const *dimensions, npy_intp const *steps, void *NPY_UNUSED(func))
575+
{
576+
UNARY_LOOP {
577+
const @type@ in1 = *(@type@ *)ip1;
578+
*((@type@ *)op1) = npy_popcount@c@(in1);
579+
}
580+
}
581+
573582
/**begin repeat1
574583
* #isa = , _avx2#
575584
* #ISA = , AVX2#

‎numpy/core/src/umath/loops.h.src

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -179,8 +179,10 @@ NPY_NO_EXPORT void
179179
@S@@TYPE@_@kind@(char **args, npy_intp const *dimensions, npy_intp const *steps, void *NPY_UNUSED(func));
180180
/**end repeat2**/
181181

182-
/**end repeat1**/
182+
NPY_NO_EXPORT void
183+
@S@@TYPE@_bit_count(char **args, npy_intp const *dimensions, npy_intp const *steps, void *NPY_UNUSED(func));
183184

185+
/**end repeat1**/
184186
/**end repeat**/
185187

186188
/*

‎numpy/core/tests/test_umath.py

Lines changed: 33 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -2034,7 +2034,15 @@ def test_reduce(self):
20342034

20352035
class TestBitwiseUFuncs:
20362036

2037-
bitwise_types = [np.dtype(c) for c in '?' + 'bBhHiIlLqQ' + 'O']
2037+
_all_ints_bits = [
2038+
np.dtype(c).itemsize * 8 for c in np.typecodes["AllInteger"]]
2039+
bitwise_types = [
2040+
np.dtype(c) for c in '?' + np.typecodes["AllInteger"] + 'O']
2041+
bitwise_bits = [
2042+
2, # boolean type
2043+
*_all_ints_bits, # All integers
2044+
max(_all_ints_bits) + 1, # Object_ type
2045+
]
20382046

20392047
def test_values(self):
20402048
for dt in self.bitwise_types:
@@ -2115,6 +2123,30 @@ def test_reduction(self):
21152123
btype = np.array([True], dtype=object)
21162124
assert_(type(f.reduce(btype)) is bool, msg)
21172125

2126+
@pytest.mark.parametrize("input_dtype_obj, bitsize",
2127+
zip(bitwise_types, bitwise_bits))
2128+
def test_popcount(self, input_dtype_obj, bitsize):
2129+
input_dtype = input_dtype_obj.type
2130+
2131+
# bit_count is only in-built in 3.10+
2132+
if sys.version_info < (3, 10) and input_dtype == np.object_:
2133+
pytest.skip()
2134+
2135+
for i in range(1, bitsize):
2136+
num = 2**i - 1
2137+
msg = f"bit_count for {num}"
2138+
assert i == np.bit_count(input_dtype(num)), msg
2139+
if np.issubdtype(
2140+
input_dtype, np.signedinteger) or input_dtype == np.object_:
2141+
assert i == np.bit_count(input_dtype(-num)), msg
2142+
2143+
a = np.array([2**i-1 for i in range(1, bitsize)], dtype=input_dtype)
2144+
bit_count_a = np.bit_count(a)
2145+
expected = np.arange(1, bitsize, dtype=input_dtype)
2146+
2147+
msg = f"array bit_count for {input_dtype}"
2148+
assert all(bit_count_a == expected), msg
2149+
21182150

21192151
class TestInt:
21202152
def test_logical_not(self):

‎numpy/core/umath.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -22,7 +22,7 @@
2222
'UFUNC_PYVALS_NAME', '_add_newdoc_ufunc', 'absolute', 'add',
2323
'arccos', 'arccosh', 'arcsin', 'arcsinh', 'arctan', 'arctan2', 'arctanh',
2424
'bitwise_and', 'bitwise_or', 'bitwise_xor', 'cbrt', 'ceil', 'conj',
25-
'conjugate', 'copysign', 'cos', 'cosh', 'deg2rad', 'degrees', 'divide',
25+
'conjugate', 'copysign', 'cos', 'cosh', 'bit_count', 'deg2rad', 'degrees', 'divide',
2626
'divmod', 'e', 'equal', 'euler_gamma', 'exp', 'exp2', 'expm1', 'fabs',
2727
'floor', 'floor_divide', 'float_power', 'fmax', 'fmin', 'fmod', 'frexp',
2828
'frompyfunc', 'gcd', 'geterrobj', 'greater', 'greater_equal', 'heaviside',

‎numpy/matrixlib/tests/test_defmatrix.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -286,7 +286,7 @@ def test_instance_methods(self):
286286
'partition', 'argpartition',
287287
'take', 'tofile', 'tolist', 'tostring', 'tobytes', 'all', 'any',
288288
'sum', 'argmax', 'argmin', 'min', 'max', 'mean', 'var', 'ptp',
289-
'prod', 'std', 'ctypes', 'itemset',
289+
'prod', 'std', 'ctypes', 'itemset', 'bit_count',
290290
]
291291
for attrib in dir(a):
292292
if attrib.startswith('_') or attrib in excluded_methods:

‎numpy/typing/tests/data/reveal/ufuncs.pyi

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
import numpy as np
22
import numpy.typing as npt
33

4+
i8: np.int64
45
f8: np.float64
56
AR_f8: npt.NDArray[np.float64]
67
AR_i8: npt.NDArray[np.int64]
@@ -66,3 +67,14 @@ reveal_type(np.matmul.signature) # E: Literal['(n?,k),(k,m?)->(n?,m?)']
6667
reveal_type(np.matmul.identity) # E: None
6768
reveal_type(np.matmul(AR_f8, AR_f8)) # E: Any
6869
reveal_type(np.matmul(AR_f8, AR_f8, axes=[(0, 1), (0, 1), (0, 1)])) # E: Any
70+
71+
reveal_type(np.bit_count.__name__) # E: Literal['bit_count']
72+
reveal_type(np.bit_count.ntypes) # E: Literal[11]
73+
reveal_type(np.bit_count.identity) # E: None
74+
reveal_type(np.bit_count.nin) # E: Literal[1]
75+
reveal_type(np.bit_count.nout) # E: Literal[1]
76+
reveal_type(np.bit_count.nargs) # E: Literal[2]
77+
reveal_type(np.bit_count.signature) # E: None
78+
reveal_type(np.bit_count.identity) # E: None
79+
reveal_type(np.bit_count(i8)) # E: Any
80+
reveal_type(np.bit_count(AR_i8)) # E: Any

0 commit comments

Comments
 (0)