diff --git a/README.rst b/README.rst index 71b94254..e4a357fb 100644 --- a/README.rst +++ b/README.rst @@ -35,6 +35,12 @@ ArrayKit requires the following: What is New in ArrayKit ------------------------- +1.11.0 +............ + +Added ``group_reduce()``. + + 1.10.0 ............ diff --git a/src/__init__.py b/src/__init__.py index bca91283..b33c2a26 100644 --- a/src/__init__.py +++ b/src/__init__.py @@ -27,6 +27,7 @@ from ._arraykit import write_array_to_file as write_array_to_file from ._arraykit import factorize as factorize from ._arraykit import group_ordering as group_ordering +from ._arraykit import group_reduce as group_reduce from ._arraykit import fill_directional as fill_directional from ._arraykit import count_iteration as count_iteration from ._arraykit import first_true_1d as first_true_1d diff --git a/src/__init__.pyi b/src/__init__.pyi index 04559c28..7fb69338 100644 --- a/src/__init__.pyi +++ b/src/__init__.pyi @@ -233,6 +233,9 @@ def factorize( def group_ordering( codes: np.ndarray, *, size: tp.Optional[int] = ... ) -> tp.Tuple[np.ndarray, np.ndarray]: ... +def group_reduce( + codes: np.ndarray, size: int, values: np.ndarray, op: str +) -> np.ndarray: ... def fill_directional( array: np.ndarray, target: np.ndarray, diff --git a/src/_arraykit.c b/src/_arraykit.c index a3cba481..9d81d172 100644 --- a/src/_arraykit.c +++ b/src/_arraykit.c @@ -78,6 +78,10 @@ static PyMethodDef arraykit_methods[] = { (PyCFunction)group_ordering, METH_VARARGS | METH_KEYWORDS, NULL}, + {"group_reduce", + (PyCFunction)group_reduce, + METH_VARARGS | METH_KEYWORDS, + NULL}, {"fill_directional", (PyCFunction)fill_directional, METH_VARARGS | METH_KEYWORDS, diff --git a/src/methods.c b/src/methods.c index 186984f8..e17625cc 100644 --- a/src/methods.c +++ b/src/methods.c @@ -8,6 +8,7 @@ # include "numpy/arrayscalars.h" # include "numpy/halffloat.h" # include +# include # ifdef _WIN32 # include @@ -1128,6 +1129,202 @@ group_ordering(PyObject *Py_UNUSED(m), PyObject *args, PyObject *kwargs) return NULL; } +typedef enum { + GR_SUM, + GR_PROD, + GR_MIN, + GR_MAX, + GR_COUNT, +} AK_GroupReduceOp; + +static int +AK_group_reduce_op_from_str(const char *op, AK_GroupReduceOp *out) { + if (strcmp(op, "sum") == 0) { *out = GR_SUM; return 0; } + if (strcmp(op, "prod") == 0) { *out = GR_PROD; return 0; } + if (strcmp(op, "min") == 0) { *out = GR_MIN; return 0; } + if (strcmp(op, "max") == 0) { *out = GR_MAX; return 0; } + if (strcmp(op, "count") == 0) { *out = GR_COUNT; return 0; } + PyErr_Format(PyExc_ValueError, + "unknown op '%s'; expected one of sum, prod, min, max, count", op); + return -1; +} + +// Accumulate `n` float64 into `out[size]` per group. NaN propagates for min/max +// (matching np.min/np.max, not the nan-skipping variants). +static void +AK_group_reduce_f64( + const npy_float64 *v, + const npy_intp *codes, + npy_intp n, + npy_float64 *out, + npy_intp size, + AK_GroupReduceOp op) { + npy_float64 init; + switch (op) { + case GR_PROD: init = 1.0; break; + case GR_MIN: init = NPY_INFINITY; break; + case GR_MAX: init = -NPY_INFINITY; break; + default: init = 0.0; break; // GR_SUM + } + for (npy_intp g = 0; g < size; g++) { + out[g] = init; + } + for (npy_intp i = 0; i < n; i++) { + npy_intp g = codes[i]; + npy_float64 x = v[i]; + switch (op) { + case GR_SUM: out[g] += x; break; + case GR_PROD: out[g] *= x; break; + case GR_MIN: if (isnan(x) || x < out[g]) out[g] = x; break; + case GR_MAX: if (isnan(x) || x > out[g]) out[g] = x; break; + default: break; + } + } +} + +// Accumulate `n` int64 into `out[size]` per group. +static void +AK_group_reduce_i64( + const npy_int64 *v, + const npy_intp *codes, + npy_intp n, + npy_int64 *out, + npy_intp size, + AK_GroupReduceOp op) { + npy_int64 init; + switch (op) { + case GR_PROD: init = 1; break; + case GR_MIN: init = NPY_MAX_INT64; break; + case GR_MAX: init = NPY_MIN_INT64; break; + default: init = 0; break; // GR_SUM + } + for (npy_intp g = 0; g < size; g++) { + out[g] = init; + } + for (npy_intp i = 0; i < n; i++) { + npy_intp g = codes[i]; + npy_int64 x = v[i]; + switch (op) { + case GR_SUM: out[g] += x; break; + case GR_PROD: out[g] *= x; break; + case GR_MIN: if (x < out[g]) out[g] = x; break; + case GR_MAX: if (x > out[g]) out[g] = x; break; + default: break; + } + } +} + +static char *group_reduce_kwarg_names[] = { + "codes", + "size", + "values", + "op", + NULL +}; + +// Grouped reduction. Given dense group `codes` in [0, size), a 1D +// `values` array, and an `op` ('sum'/'prod'/'min'/'max'/'count'), return a length- +// `size` array of per-group results in code order. Accumulates directly by code in +// an O(n) pass after validating codes (no sort, no reorder). 'count' returns int64 group sizes and ignores +// the values dtype; other ops return the values dtype (float64 or int64). This is +// the vectorized replacement for a per-group Python reduction loop. +PyObject * +group_reduce(PyObject *Py_UNUSED(m), PyObject *args, PyObject *kwargs) +{ + PyArrayObject *codes = NULL; + Py_ssize_t size = 0; + PyArrayObject *values = NULL; + const char *op_name = NULL; + + if (!PyArg_ParseTupleAndKeywords(args, kwargs, + "O!nO!s:group_reduce", + group_reduce_kwarg_names, + &PyArray_Type, &codes, + &size, + &PyArray_Type, &values, + &op_name + )) { + return NULL; + } + AK_GroupReduceOp op; + if (AK_group_reduce_op_from_str(op_name, &op)) { + return NULL; + } + if (size < 0) { + PyErr_SetString(PyExc_ValueError, "size must be non-negative"); + return NULL; + } + if (PyArray_NDIM(codes) != 1 || PyArray_NDIM(values) != 1) { + PyErr_SetString(PyExc_ValueError, "Arrays must be 1-dimensional"); + return NULL; + } + if (PyArray_TYPE(codes) != NPY_INTP) { + PyErr_SetString(PyExc_ValueError, "codes must be of type intp"); + return NULL; + } + if (!PyArray_IS_C_CONTIGUOUS(codes) || !PyArray_IS_C_CONTIGUOUS(values)) { + PyErr_SetString(PyExc_ValueError, "Arrays must be contiguous"); + return NULL; + } + npy_intp n = PyArray_SIZE(codes); + if (PyArray_SIZE(values) != n) { + PyErr_SetString(PyExc_ValueError, + "codes and values must be the same length"); + return NULL; + } + const npy_intp *codes_buffer = (npy_intp*)PyArray_DATA(codes); + // validate codes are in range before any indexed writes into the output + for (npy_intp i = 0; i < n; i++) { + npy_intp c = codes_buffer[i]; + if (c < 0 || c >= size) { + PyErr_Format(PyExc_ValueError, + "code %zd out of range [0, %zd)", + (Py_ssize_t)c, (Py_ssize_t)size); + return NULL; + } + } + + npy_intp dims[1] = {size}; + int vtype = PyArray_TYPE(values); + + if (op == GR_COUNT) { + PyObject *out_arr = PyArray_ZEROS(1, dims, NPY_INT64, 0); + if (!out_arr) { + return NULL; + } + npy_int64 *out = (npy_int64*)PyArray_DATA((PyArrayObject*)out_arr); + for (npy_intp i = 0; i < n; i++) { + out[codes_buffer[i]]++; + } + PyArray_CLEARFLAGS((PyArrayObject*)out_arr, NPY_ARRAY_WRITEABLE); + return out_arr; + } + + if (vtype != NPY_DOUBLE && vtype != NPY_INT64) { + PyErr_SetString(PyExc_ValueError, + "values must be of type float64 or int64"); + return NULL; + } + PyObject *out_arr = PyArray_EMPTY(1, dims, vtype, 0); + if (!out_arr) { + return NULL; + } + if (vtype == NPY_DOUBLE) { + AK_group_reduce_f64( + (npy_float64*)PyArray_DATA(values), + codes_buffer, n, + (npy_float64*)PyArray_DATA((PyArrayObject*)out_arr), size, op); + } + else { + AK_group_reduce_i64( + (npy_int64*)PyArray_DATA(values), + codes_buffer, n, + (npy_int64*)PyArray_DATA((PyArrayObject*)out_arr), size, op); + } + PyArray_CLEARFLAGS((PyArrayObject*)out_arr, NPY_ARRAY_WRITEABLE); + return out_arr; +} + // Fill one strided lane in place: walk positions in the fill direction, carrying // the most recent non-target value into each target position (subject to `limit` // consecutive fills per run). `elem_base`/`elem_stride` address elements in bytes; diff --git a/src/methods.h b/src/methods.h index 7b425300..e941b178 100644 --- a/src/methods.h +++ b/src/methods.h @@ -72,6 +72,9 @@ first_true_2d(PyObject *Py_UNUSED(m), PyObject *args, PyObject *kwargs); PyObject * group_ordering(PyObject *Py_UNUSED(m), PyObject *args, PyObject *kwargs); +PyObject * +group_reduce(PyObject *Py_UNUSED(m), PyObject *args, PyObject *kwargs); + PyObject * fill_directional(PyObject *Py_UNUSED(m), PyObject *args, PyObject *kwargs); diff --git a/test/test_group_reduce.py b/test/test_group_reduce.py new file mode 100644 index 00000000..2b09984f --- /dev/null +++ b/test/test_group_reduce.py @@ -0,0 +1,149 @@ +import unittest + +import numpy as np +from arraykit import factorize, group_reduce + + +class TestUnit(unittest.TestCase): + # ------------------------------------------------------------------ + # basic behavior + + def test_group_reduce_sum_f64(self) -> None: + codes = np.array([0, 1, 0, 2, 1, 0], dtype=np.intp) + values = np.array([1.0, 10.0, 2.0, 100.0, 20.0, 3.0]) + post = group_reduce(codes, 3, values, 'sum') + self.assertEqual(post.tolist(), [6.0, 30.0, 100.0]) + self.assertEqual(post.dtype, np.dtype(np.float64)) + + def test_group_reduce_all_ops_f64(self) -> None: + codes = np.array([0, 1, 0, 2, 1, 0], dtype=np.intp) + values = np.array([1.0, 10.0, 2.0, 100.0, 20.0, 3.0]) + self.assertEqual(group_reduce(codes, 3, values, 'sum').tolist(), [6.0, 30.0, 100.0]) + self.assertEqual(group_reduce(codes, 3, values, 'prod').tolist(), [6.0, 200.0, 100.0]) + self.assertEqual(group_reduce(codes, 3, values, 'min').tolist(), [1.0, 10.0, 100.0]) + self.assertEqual(group_reduce(codes, 3, values, 'max').tolist(), [3.0, 20.0, 100.0]) + + def test_group_reduce_all_ops_i64(self) -> None: + codes = np.array([0, 1, 0, 2, 1, 0], dtype=np.intp) + values = np.array([1, 10, 2, 100, 20, 3], dtype=np.int64) + for op in ('sum', 'prod', 'min', 'max'): + post = group_reduce(codes, 3, values, op) + self.assertEqual(post.dtype, np.dtype(np.int64)) + self.assertEqual(group_reduce(codes, 3, values, 'sum').tolist(), [6, 30, 100]) + self.assertEqual(group_reduce(codes, 3, values, 'prod').tolist(), [6, 200, 100]) + self.assertEqual(group_reduce(codes, 3, values, 'min').tolist(), [1, 10, 100]) + self.assertEqual(group_reduce(codes, 3, values, 'max').tolist(), [3, 20, 100]) + + def test_group_reduce_count(self) -> None: + # count returns int64 group sizes regardless of values dtype + codes = np.array([0, 1, 0, 2, 1, 0], dtype=np.intp) + post = group_reduce(codes, 3, np.array([1.0, 2, 3, 4, 5, 6]), 'count') + self.assertEqual(post.tolist(), [3, 2, 1]) + self.assertEqual(post.dtype, np.dtype(np.int64)) + # count ignores the values dtype entirely + post = group_reduce(codes, 3, np.array([1, 2, 3, 4, 5, 6], dtype=np.int64), 'count') + self.assertEqual(post.tolist(), [3, 2, 1]) + + def test_group_reduce_nan_propagates(self) -> None: + # min/max propagate NaN, matching np.min/np.max (not the nan-skipping variants) + codes = np.array([0, 1, 0, 2, 1, 0], dtype=np.intp) + values = np.array([1.0, np.nan, 2.0, 5.0, np.nan, 3.0]) + mx = group_reduce(codes, 3, values, 'max') + mn = group_reduce(codes, 3, values, 'min') + self.assertEqual(mx[0], 3.0) + self.assertTrue(np.isnan(mx[1])) + self.assertEqual(mx[2], 5.0) + self.assertEqual(mn[0], 1.0) + self.assertTrue(np.isnan(mn[1])) + # sum also propagates NaN + s = group_reduce(codes, 3, values, 'sum') + self.assertTrue(np.isnan(s[1])) + + def test_group_reduce_single_group(self) -> None: + codes = np.array([0, 0, 0], dtype=np.intp) + values = np.array([1.0, 2.0, 3.0]) + self.assertEqual(group_reduce(codes, 1, values, 'sum').tolist(), [6.0]) + + def test_group_reduce_empty(self) -> None: + codes = np.array([], dtype=np.intp) + values = np.array([], dtype=np.float64) + self.assertEqual(group_reduce(codes, 0, values, 'sum').tolist(), []) + + def test_group_reduce_outputs_immutable(self) -> None: + codes = np.array([0, 1, 0], dtype=np.intp) + values = np.array([1.0, 2.0, 3.0]) + for op in ('sum', 'prod', 'min', 'max', 'count'): + post = group_reduce(codes, 2, values, op) + self.assertFalse(post.flags.writeable) + + # ------------------------------------------------------------------ + # equivalence to a per-group numpy reduction + + def test_group_reduce_equivalence_f64(self) -> None: + rng = np.random.RandomState(0) + for _ in range(20): + size = int(rng.randint(1, 40)) + n = int(rng.randint(size, size + 500)) + codes = rng.randint(0, size, n).astype(np.intp) + values = rng.rand(n) * 100 + for op, npf in (('sum', np.sum), ('min', np.min), ('max', np.max)): + got = group_reduce(codes, size, values, op) + for g in range(size): + mask = codes == g + if np.any(mask): # real usage (factorize) has no empty groups + self.assertTrue(np.isclose(got[g], npf(values[mask])), op) + + def test_group_reduce_equivalence_i64(self) -> None: + rng = np.random.RandomState(1) + for _ in range(20): + size = int(rng.randint(1, 40)) + n = int(rng.randint(size, size + 500)) + codes = rng.randint(0, size, n).astype(np.intp) + values = rng.randint(-1000, 1000, n).astype(np.int64) + for op, npf in (('sum', np.sum), ('min', np.min), ('max', np.max)): + got = group_reduce(codes, size, values, op) + for g in range(size): + mask = codes == g + if np.any(mask): + self.assertEqual(got[g], npf(values[mask]), op) + + def test_group_reduce_with_factorize(self) -> None: + # the intended pipeline: factorize(sort=True) -> group_reduce + key = np.array([30, 10, 20, 10, 30, 20, 10]) + values = np.array([1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0]) + uniques, codes = factorize(key, sort=True) + self.assertEqual(uniques.tolist(), [10, 20, 30]) + post = group_reduce(codes, len(uniques), values, 'sum') + # group 10 -> 2+4+7=13; group 20 -> 3+6=9; group 30 -> 1+5=6 + self.assertEqual(post.tolist(), [13.0, 9.0, 6.0]) + + # ------------------------------------------------------------------ + # errors + + def test_group_reduce_errors(self) -> None: + codes = np.array([0, 1, 0], dtype=np.intp) + values = np.array([1.0, 2.0, 3.0]) + with self.assertRaises(ValueError): # unknown op + group_reduce(codes, 2, values, 'median') + with self.assertRaises(ValueError): # length mismatch + group_reduce(codes, 2, np.array([1.0, 2.0]), 'sum') + with self.assertRaises(ValueError): # code out of range + group_reduce(np.array([0, 5], dtype=np.intp), 2, np.array([1.0, 2.0]), 'sum') + with self.assertRaises(ValueError): # negative code + group_reduce(np.array([0, -1], dtype=np.intp), 2, np.array([1.0, 2.0]), 'sum') + with self.assertRaises(ValueError): # codes wrong dtype (int8 is never intp) + group_reduce(np.array([0, 1], dtype=np.int8), 2, np.array([1.0, 2.0]), 'sum') + with self.assertRaises(ValueError): # values unsupported dtype + group_reduce(codes, 2, np.array([1, 2, 3], dtype=np.int32), 'sum') + with self.assertRaises(ValueError): # negative size + group_reduce(codes, -1, values, 'sum') + with self.assertRaises(ValueError): # 2d codes + group_reduce( + np.array([[0, 1]], dtype=np.intp), 2, np.array([[1.0, 2.0]]), 'sum' + ) + with self.assertRaises(TypeError): # not an array + group_reduce([0, 1, 0], 2, values, 'sum') + + +if __name__ == '__main__': + unittest.main()