arraykit 1.9.0__tar.gz → 1.11.0__tar.gz

{arraykit-1.9.0/arraykit.egg-info → arraykit-1.11.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: arraykit
-Version: 1.9.0
+Version: 1.11.0
 Summary: Array utilities for StaticFrame
 Author: Christopher Ariza, Brandt Bucher, Charles Burkland
 License: MIT
@@ -64,6 +64,18 @@ ArrayKit requires the following:
 What is New in ArrayKit
 -------------------------
 
+1.11.0
+............
+
+Added ``group_reduce()``.
+
+
+1.10.0
+............
+
+Added ``TriMap.register_pairs()``.
+
+
 1.9.0
 ............
 

{arraykit-1.9.0 → arraykit-1.11.0}/README.rst

@@ -35,6 +35,18 @@ ArrayKit requires the following:
 What is New in ArrayKit
 -------------------------
 
+1.11.0
+............
+
+Added ``group_reduce()``.
+
+
+1.10.0
+............
+
+Added ``TriMap.register_pairs()``.
+
+
 1.9.0
 ............
 

arraykit-1.11.0/VERSION

@@ -0,0 +1,2 @@
+1.11.0
+

{arraykit-1.9.0 → arraykit-1.11.0/arraykit.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: arraykit
-Version: 1.9.0
+Version: 1.11.0
 Summary: Array utilities for StaticFrame
 Author: Christopher Ariza, Brandt Bucher, Charles Burkland
 License: MIT
@@ -64,6 +64,18 @@ ArrayKit requires the following:
 What is New in ArrayKit
 -------------------------
 
+1.11.0
+............
+
+Added ``group_reduce()``.
+
+
+1.10.0
+............
+
+Added ``TriMap.register_pairs()``.
+
+
 1.9.0
 ............
 

{arraykit-1.9.0 → arraykit-1.11.0}/arraykit.egg-info/SOURCES.txt

@@ -39,6 +39,7 @@ test/test_delimited_to_arrays_property.py
 test/test_factorize.py
 test/test_fill_directional.py
 test/test_group_ordering.py
+test/test_group_reduce.py
 test/test_nonzero_1d.py
 test/test_nonzero_1d_property.py
 test/test_objectable.py

{arraykit-1.9.0 → arraykit-1.11.0}/src/__init__.py

@@ -27,6 +27,7 @@ from ._arraykit import get_new_indexers_and_screen as get_new_indexers_and_scree
 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

{arraykit-1.9.0 → arraykit-1.11.0}/src/__init__.pyi

@@ -54,6 +54,7 @@ class TriMap:
     def __repr__(self) -> str: ...
     def register_one(self, /, src_from: int, dst_from: int) -> None: ...
     def register_many_from_one(self, __dst_pos: np.ndarray) -> None: ...
+    def register_pairs(self, __src_pos: np.ndarray, __dst_pos: np.ndarray) -> None: ...
     def register_unmatched_dst(self) -> None: ...
     def register_many(self, /, src_from: int, dst_from: np.ndarray) -> None: ...
     def finalize(self) -> None: ...
@@ -232,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,

{arraykit-1.9.0 → arraykit-1.11.0}/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,

{arraykit-1.9.0 → arraykit-1.11.0}/src/methods.c

@@ -8,6 +8,7 @@
 # include "numpy/arrayscalars.h"
 # include "numpy/halffloat.h"
 # include <string.h>
+# include <math.h>
 
 # ifdef _WIN32
 #  include <io.h>
@@ -1128,6 +1129,280 @@ fail:
     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;
+        }
+    }
+}
+
+// Accumulate `n` uint64 into `out[size]` per group.
+static void
+AK_group_reduce_u64(
+        const npy_uint64 *v,
+        const npy_intp *codes,
+        npy_intp n,
+        npy_uint64 *out,
+        npy_intp size,
+        AK_GroupReduceOp op) {
+    npy_uint64 init;
+    switch (op) {
+        case GR_PROD: init = 1; break;
+        case GR_MIN:  init = NPY_MAX_UINT64; break;
+        case GR_MAX:  init = 0; 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_uint64 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). Values are accumulated at a 64-bit width by kind: signed
+// integers -> int64, unsigned integers -> uint64, floats -> float64; the result uses
+// that accumulator dtype (the caller casts to a narrower output as needed -- min/max are
+// selected elements so such a cast is lossless). Integer overflow wraps, matching numpy.
+// float16/float32 sum/prod are rejected: accumulating them at float64 would not match
+// numpy's native-width result, so the caller should fall back for those. 'count' returns
+// int64 group sizes and ignores the values dtype. 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;
+    }
+
+    // Choose a 64-bit accumulator by input kind. Integers accumulate exactly at their
+    // signed/unsigned 64-bit width (wrapping on overflow, like numpy); floats at float64.
+    int acc_type;
+    if (PyArray_ISSIGNED(values)) {
+        acc_type = NPY_INT64;
+    }
+    else if (PyArray_ISUNSIGNED(values)) {
+        acc_type = NPY_UINT64;
+    }
+    else if (vtype == NPY_HALF || vtype == NPY_FLOAT || vtype == NPY_DOUBLE) {
+        acc_type = NPY_DOUBLE;
+    }
+    else {
+        PyErr_SetString(PyExc_ValueError,
+                "values must be an integer or float (float16/32/64) dtype");
+        return NULL;
+    }
+    // float16/float32 sum/prod cannot be accumulated at float64 without diverging from
+    // numpy's native-width result; the caller falls back to a per-group reduction.
+    if ((op == GR_SUM || op == GR_PROD)
+            && acc_type == NPY_DOUBLE && vtype != NPY_DOUBLE) {
+        PyErr_SetString(PyExc_ValueError,
+                "float16/float32 sum/prod is not supported; cast to float64");
+        return NULL;
+    }
+
+    // upcast the input to the accumulator dtype (no copy when already 64-bit wide)
+    PyArrayObject *v64;
+    if (vtype == acc_type) {
+        Py_INCREF(values);
+        v64 = values;
+    }
+    else {
+        v64 = (PyArrayObject*)PyArray_Cast(values, acc_type);
+        if (v64 == NULL) {
+            return NULL;
+        }
+    }
+
+    PyObject *out_arr = PyArray_EMPTY(1, dims, acc_type, 0);
+    if (!out_arr) {
+        Py_DECREF(v64);
+        return NULL;
+    }
+    switch (acc_type) {
+        case NPY_DOUBLE:
+            AK_group_reduce_f64(
+                    (npy_float64*)PyArray_DATA(v64), codes_buffer, n,
+                    (npy_float64*)PyArray_DATA((PyArrayObject*)out_arr), size, op);
+            break;
+        case NPY_UINT64:
+            AK_group_reduce_u64(
+                    (npy_uint64*)PyArray_DATA(v64), codes_buffer, n,
+                    (npy_uint64*)PyArray_DATA((PyArrayObject*)out_arr), size, op);
+            break;
+        default: // NPY_INT64
+            AK_group_reduce_i64(
+                    (npy_int64*)PyArray_DATA(v64), codes_buffer, n,
+                    (npy_int64*)PyArray_DATA((PyArrayObject*)out_arr), size, op);
+            break;
+    }
+    Py_DECREF(v64);
+    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;

{arraykit-1.9.0 → arraykit-1.11.0}/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);
 

{arraykit-1.9.0 → arraykit-1.11.0}/src/tri_map.c

@@ -310,6 +310,60 @@ TriMap_register_many_from_one(TriMapObject *self, PyObject *arg) {
     Py_RETURN_NONE;
 }
 
+// Bulk pair registration: given two equal-length int64 arrays `src_pos` and `dst_pos`,
+// register the pair (src_pos[i], dst_pos[i]) for each i in a single C loop -- equivalent
+// to calling register_one(src_pos[i], dst_pos[i]) for each i, but without per-element
+// Python overhead. Either position may be -1 to indicate an unmatched (fill) side. This
+// generalizes register_many_from_one (where src is the implicit range 0..src_len) to
+// arbitrary src positions, as needed for many-to-many joins.
+PyObject *
+TriMap_register_pairs(TriMapObject *self, PyObject *args) {
+    PyObject* src_arg;
+    PyObject* dst_arg;
+    if (!PyArg_ParseTuple(args,
+            "OO:register_pairs",
+            &src_arg,
+            &dst_arg)) {
+        return NULL;
+    }
+    if (self->finalized) {
+        PyErr_SetString(PyExc_RuntimeError, "Cannot register post finalization");
+        return NULL;
+    }
+    if (!PyArray_Check(src_arg) || !PyArray_Check(dst_arg)) {
+        PyErr_SetString(PyExc_TypeError, "Must provide arrays");
+        return NULL;
+    }
+    PyArrayObject* src_a = (PyArrayObject*)src_arg;
+    PyArrayObject* dst_a = (PyArrayObject*)dst_arg;
+    if (PyArray_TYPE(src_a) != NPY_INT64 || PyArray_TYPE(dst_a) != NPY_INT64) {
+        PyErr_SetString(PyExc_ValueError, "Arrays must be of type int64");
+        return NULL;
+    }
+    if (PyArray_NDIM(src_a) != 1 || PyArray_NDIM(dst_a) != 1) {
+        PyErr_SetString(PyExc_ValueError, "Arrays must be 1-dimensional");
+        return NULL;
+    }
+    if (!PyArray_IS_C_CONTIGUOUS(src_a) || !PyArray_ISALIGNED(src_a)
+            || !PyArray_IS_C_CONTIGUOUS(dst_a) || !PyArray_ISALIGNED(dst_a)) {
+        PyErr_SetString(PyExc_ValueError, "Arrays must be contiguous");
+        return NULL;
+    }
+    npy_intp n = PyArray_SIZE(src_a);
+    if (n != PyArray_SIZE(dst_a)) {
+        PyErr_SetString(PyExc_ValueError, "Arrays must be the same length");
+        return NULL;
+    }
+    const npy_int64* s = (npy_int64*)PyArray_DATA(src_a);
+    const npy_int64* d = (npy_int64*)PyArray_DATA(dst_a);
+    for (npy_intp i = 0; i < n; i++) {
+        if (AK_TM_register_one(self, (Py_ssize_t)s[i], (Py_ssize_t)d[i])) {
+            return NULL;
+        }
+    }
+    Py_RETURN_NONE;
+}
+
 PyObject *
 TriMap_register_unmatched_dst(TriMapObject *self) {
     if (self->finalized) {
@@ -1400,6 +1454,7 @@ TriMap_map_dst_fill(TriMapObject *self, PyObject *args) {
 static PyMethodDef TriMap_methods[] = {
     {"register_one", (PyCFunction)TriMap_register_one, METH_VARARGS, NULL},
     {"register_many_from_one", (PyCFunction)TriMap_register_many_from_one, METH_O, NULL},
+    {"register_pairs", (PyCFunction)TriMap_register_pairs, METH_VARARGS, NULL},
     {"register_unmatched_dst", (PyCFunction)TriMap_register_unmatched_dst, METH_NOARGS, NULL},
     {"register_many", (PyCFunction)TriMap_register_many, METH_VARARGS, NULL},
     {"finalize", (PyCFunction)TriMap_finalize, METH_NOARGS, NULL},

arraykit-1.11.0/test/test_group_reduce.py

@@ -0,0 +1,217 @@
+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_all_int_widths(self) -> None:
+        # every signed width accumulates at int64; every unsigned width at uint64.
+        # values are exact (selected elements / small sums), only the dtype widens.
+        codes = np.array([0, 1, 0, 2, 1, 0], dtype=np.intp)
+        raw = [1, 10, 2, 100, 20, 3]
+        for dt in ('int8', 'int16', 'int32', 'int64'):
+            values = np.array(raw, dtype=dt)
+            for op, exp in (
+                ('sum', [6, 30, 100]),
+                ('min', [1, 10, 100]),
+                ('max', [3, 20, 100]),
+            ):
+                post = group_reduce(codes, 3, values, op)
+                self.assertEqual(post.dtype, np.dtype(np.int64), (dt, op))
+                self.assertEqual(post.tolist(), exp, (dt, op))
+        for dt in ('uint8', 'uint16', 'uint32', 'uint64'):
+            values = np.array(raw, dtype=dt)
+            for op, exp in (
+                ('sum', [6, 30, 100]),
+                ('min', [1, 10, 100]),
+                ('max', [3, 20, 100]),
+            ):
+                post = group_reduce(codes, 3, values, op)
+                self.assertEqual(post.dtype, np.dtype(np.uint64), (dt, op))
+                self.assertEqual(post.tolist(), exp, (dt, op))
+
+    def test_group_reduce_float_widths(self) -> None:
+        # min/max/count work for every float width (result is float64); the
+        # selected element is exact
+        codes = np.array([0, 1, 0, 2, 1, 0], dtype=np.intp)
+        raw = [1.0, 10.0, 2.0, 100.0, 20.0, 3.0]
+        for dt in ('float16', 'float32', 'float64'):
+            values = np.array(raw, dtype=dt)
+            for op, exp in (('min', [1.0, 10.0, 100.0]), ('max', [3.0, 20.0, 100.0])):
+                post = group_reduce(codes, 3, values, op)
+                self.assertEqual(post.dtype, np.dtype(np.float64), (dt, op))
+                self.assertEqual(post.tolist(), exp, (dt, op))
+
+    def test_group_reduce_narrow_float_sum_rejected(self) -> None:
+        # float16/float32 sum/prod cannot match numpy's native-width result at
+        # float64, so they are rejected (caller falls back); float64 is fine
+        codes = np.array([0, 1], dtype=np.intp)
+        for dt in ('float16', 'float32'):
+            values = np.array([1.0, 2.0], dtype=dt)
+            for op in ('sum', 'prod'):
+                with self.assertRaises(ValueError):
+                    group_reduce(codes, 2, values, op)
+        # float64 sum/prod is supported
+        self.assertEqual(
+            group_reduce(codes, 2, np.array([1.0, 2.0]), 'sum').tolist(), [1.0, 2.0]
+        )
+
+    def test_group_reduce_integer_overflow_wraps(self) -> None:
+        # integer sum overflow wraps silently, matching numpy (not raising)
+        codes = np.array([0, 0], dtype=np.intp)
+        big = np.full(2, np.iinfo(np.int64).max, dtype=np.int64)
+        self.assertEqual(
+            group_reduce(codes, 1, big, 'sum').tolist(), [int(np.sum(big))]
+        )
+        ubig = np.full(2, np.iinfo(np.uint64).max, dtype=np.uint64)
+        self.assertEqual(
+            group_reduce(codes, 1, ubig, 'sum').tolist(), [int(np.sum(ubig))]
+        )
+
+    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 (complex)
+            group_reduce(codes, 2, np.array([1, 2, 3], dtype=np.complex128), 'sum')
+        with self.assertRaises(ValueError):  # values unsupported dtype (datetime)
+            group_reduce(
+                codes, 2, np.array([1, 2, 3], dtype='datetime64[s]'), 'max'
+            )
+        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()

{arraykit-1.9.0 → arraykit-1.11.0}/test/test_tri_map.py

@@ -1508,3 +1508,157 @@ class TestUnit(unittest.TestCase):
         tm.finalize()
         with self.assertRaises(RuntimeError):  # post-finalize
             tm.register_many_from_one(np.array([0, 1], dtype=np.int64))
+
+    # ---------------------------------------------------------------------------
+    # register_pairs (bulk arbitrary (src, dst) pairs)
+
+    def test_tri_map_register_pairs_a(self) -> None:
+        # a many-to-many mapping: src 0 -> dst 0 and 1; src 1 -> dst 1
+        tm = TriMap(2, 2)
+        tm.register_pairs(
+            np.array([0, 0, 1], dtype=np.int64),
+            np.array([0, 1, 1], dtype=np.int64),
+        )
+        tm.finalize()
+        self.assertTrue(tm.is_many())
+        self.assertTrue(tm.src_no_fill())
+        self.assertTrue(tm.dst_no_fill())
+        src = np.array([10, 20])
+        dst = np.array([100, 200])
+        self.assertEqual(tm.map_src_no_fill(src).tolist(), [10, 10, 20])
+        self.assertEqual(tm.map_dst_no_fill(dst).tolist(), [100, 200, 200])
+
+    def test_tri_map_register_pairs_unmatched(self) -> None:
+        # -1 on either side marks that side as a fill row
+        tm = TriMap(3, 2)
+        tm.register_pairs(
+            np.array([0, 1, 2], dtype=np.int64),
+            np.array([0, 1, -1], dtype=np.int64),
+        )
+        tm.finalize()
+        self.assertTrue(tm.src_no_fill())  # every row has a src
+        self.assertFalse(tm.dst_no_fill())  # row 2 has no dst
+        dst = np.array([100, 200])
+        self.assertEqual(
+            tm.map_dst_fill(dst, -9, np.dtype(np.int64)).tolist(), [100, 200, -9]
+        )
+
+    def test_tri_map_register_pairs_unmatched_src(self) -> None:
+        # -1 on the src side (as register_unmatched_dst would produce for OUTER)
+        tm = TriMap(2, 3)
+        tm.register_pairs(
+            np.array([0, 1, -1], dtype=np.int64),
+            np.array([0, 1, 2], dtype=np.int64),
+        )
+        tm.finalize()
+        self.assertFalse(tm.src_no_fill())  # row 2 has no src
+        self.assertTrue(tm.dst_no_fill())  # every row has a dst
+        src = np.array([10, 20])
+        self.assertEqual(
+            tm.map_src_fill(src, -9, np.dtype(np.int64)).tolist(), [10, 20, -9]
+        )
+
+    def test_tri_map_register_pairs_is_many(self) -> None:
+        # a single output row per src/dst -> not many
+        tm = TriMap(2, 2)
+        tm.register_pairs(
+            np.array([0, 1], dtype=np.int64),
+            np.array([1, 0], dtype=np.int64),
+        )
+        tm.finalize()
+        self.assertFalse(tm.is_many())
+
+    def test_tri_map_register_pairs_equivalence_loop(self) -> None:
+        # register_pairs == a register_one loop over the same pairs
+        rng = np.random.RandomState(0)
+        for _ in range(50):
+            src_len = int(rng.randint(1, 12))
+            dst_len = int(rng.randint(1, 10))
+            n = int(rng.randint(1, 20))
+            src_pos = rng.randint(-1, src_len, size=n).astype(np.int64)
+            dst_pos = rng.randint(-1, dst_len, size=n).astype(np.int64)
+            tb = TriMap(src_len, dst_len)
+            tb.register_pairs(src_pos, dst_pos)
+            tb.finalize()
+            tl = TriMap(src_len, dst_len)
+            for i in range(n):
+                tl.register_one(int(src_pos[i]), int(dst_pos[i]))
+            tl.finalize()
+            src = np.arange(100, 100 + src_len)
+            dst = np.arange(200, 200 + dst_len)
+            self.assertEqual(tb.is_many(), tl.is_many())
+            self.assertEqual(
+                tb.map_src_fill(src, -1, np.dtype(np.int64)).tolist(),
+                tl.map_src_fill(src, -1, np.dtype(np.int64)).tolist(),
+            )
+            self.assertEqual(
+                tb.map_dst_fill(dst, -1, np.dtype(np.int64)).tolist(),
+                tl.map_dst_fill(dst, -1, np.dtype(np.int64)).tolist(),
+            )
+
+    def test_tri_map_register_pairs_equivalence_many_from_one(self) -> None:
+        # with src = arange(src_len), register_pairs matches register_many_from_one
+        rng = np.random.RandomState(1)
+        for _ in range(20):
+            src_len = int(rng.randint(1, 12))
+            dst_len = int(rng.randint(1, 10))
+            dst_pos = rng.randint(-1, dst_len, size=src_len).astype(np.int64)
+            src_pos = np.arange(src_len, dtype=np.int64)
+            tp = TriMap(src_len, dst_len)
+            tp.register_pairs(src_pos, dst_pos)
+            tp.finalize()
+            tm = TriMap(src_len, dst_len)
+            tm.register_many_from_one(dst_pos)
+            tm.finalize()
+            src = np.arange(100, 100 + src_len)
+            dst = np.arange(200, 200 + dst_len)
+            self.assertEqual(tp.is_many(), tm.is_many())
+            self.assertEqual(
+                tp.map_src_fill(src, -1, np.dtype(np.int64)).tolist(),
+                tm.map_src_fill(src, -1, np.dtype(np.int64)).tolist(),
+            )
+            self.assertEqual(
+                tp.map_dst_fill(dst, -1, np.dtype(np.int64)).tolist(),
+                tm.map_dst_fill(dst, -1, np.dtype(np.int64)).tolist(),
+            )
+
+    def test_tri_map_register_pairs_empty(self) -> None:
+        tm = TriMap(3, 3)
+        tm.register_pairs(
+            np.array([], dtype=np.int64), np.array([], dtype=np.int64)
+        )
+        tm.finalize()
+        self.assertEqual(tm.map_src_no_fill(np.arange(3)).tolist(), [])
+
+    def test_tri_map_register_pairs_errors(self) -> None:
+        with self.assertRaises(ValueError):  # length mismatch
+            TriMap(3, 3).register_pairs(
+                np.array([0, 1], dtype=np.int64), np.array([0], dtype=np.int64)
+            )
+        with self.assertRaises(ValueError):  # wrong dtype
+            TriMap(3, 3).register_pairs(
+                np.array([0, 1], dtype=np.int32), np.array([0, 1], dtype=np.int64)
+            )
+        with self.assertRaises(ValueError):  # 2d
+            TriMap(3, 3).register_pairs(
+                np.array([[0, 1]], dtype=np.int64), np.array([[0, 1]], dtype=np.int64)
+            )
+        with self.assertRaises(ValueError):  # out of bounds src
+            TriMap(3, 3).register_pairs(
+                np.array([9], dtype=np.int64), np.array([0], dtype=np.int64)
+            )
+        with self.assertRaises(ValueError):  # out of bounds dst
+            TriMap(3, 3).register_pairs(
+                np.array([0], dtype=np.int64), np.array([9], dtype=np.int64)
+            )
+        with self.assertRaises(TypeError):  # not an array
+            TriMap(3, 3).register_pairs([0, 1], np.array([0, 1], dtype=np.int64))
+        tm = TriMap(2, 2)
+        tm.register_pairs(
+            np.array([0, 1], dtype=np.int64), np.array([0, 1], dtype=np.int64)
+        )
+        tm.finalize()
+        with self.assertRaises(RuntimeError):  # post-finalize
+            tm.register_pairs(
+                np.array([0], dtype=np.int64), np.array([0], dtype=np.int64)
+            )

arraykit-1.9.0/VERSION

@@ -1,2 +0,0 @@
-1.9.0
-