ashlr-plugin 1.36.4

package/bench/refs/python-lib/pandas/core/_numba/kernels/mean_.py

@@ -0,0 +1,198 @@
+"""
+Numba 1D mean kernels that can be shared by
+* Dataframe / Series
+* groupby
+* rolling / expanding
+
+Mirrors pandas/_libs/window/aggregation.pyx
+"""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+import numba
+import numpy as np
+
+from pandas.core._numba.kernels.shared import is_monotonic_increasing
+from pandas.core._numba.kernels.sum_ import grouped_kahan_sum
+
+if TYPE_CHECKING:
+    from pandas._typing import npt
+
+
+@numba.jit(nogil=True, parallel=False)
+def add_mean(
+    val: float,
+    nobs: int,
+    sum_x: float,
+    neg_ct: int,
+    compensation: float,
+    num_consecutive_same_value: int,
+    prev_value: float,
+) -> tuple[int, float, int, float, int, float]:
+    if not np.isnan(val):
+        nobs += 1
+        y = val - compensation
+        t = sum_x + y
+        compensation = t - sum_x - y
+        sum_x = t
+        if val < 0:
+            neg_ct += 1
+
+        if val == prev_value:
+            num_consecutive_same_value += 1
+        else:
+            num_consecutive_same_value = 1
+        prev_value = val
+
+    return nobs, sum_x, neg_ct, compensation, num_consecutive_same_value, prev_value
+
+
+@numba.jit(nogil=True, parallel=False)
+def remove_mean(
+    val: float, nobs: int, sum_x: float, neg_ct: int, compensation: float
+) -> tuple[int, float, int, float]:
+    if not np.isnan(val):
+        nobs -= 1
+        y = -val - compensation
+        t = sum_x + y
+        compensation = t - sum_x - y
+        sum_x = t
+        if val < 0:
+            neg_ct -= 1
+    return nobs, sum_x, neg_ct, compensation
+
+
+@numba.jit(nogil=True, parallel=False)
+def sliding_mean(
+    values: np.ndarray,
+    result_dtype: np.dtype,
+    start: np.ndarray,
+    end: np.ndarray,
+    min_periods: int,
+) -> tuple[np.ndarray, list[int]]:
+    N = len(start)
+    nobs = 0
+    sum_x = 0.0
+    neg_ct = 0
+    compensation_add = 0.0
+    compensation_remove = 0.0
+
+    is_monotonic_increasing_bounds = is_monotonic_increasing(
+        start
+    ) and is_monotonic_increasing(end)
+
+    output = np.empty(N, dtype=result_dtype)
+
+    for i in range(N):
+        s = start[i]
+        e = end[i]
+        if i == 0 or not is_monotonic_increasing_bounds:
+            prev_value = values[s]
+            num_consecutive_same_value = 0
+
+            for j in range(s, e):
+                val = values[j]
+                (
+                    nobs,
+                    sum_x,
+                    neg_ct,
+                    compensation_add,
+                    num_consecutive_same_value,
+                    prev_value,
+                ) = add_mean(
+                    val,
+                    nobs,
+                    sum_x,
+                    neg_ct,
+                    compensation_add,
+                    num_consecutive_same_value,
+                    prev_value,  # pyright: ignore[reportArgumentType]
+                )
+        else:
+            for j in range(start[i - 1], s):
+                val = values[j]
+                nobs, sum_x, neg_ct, compensation_remove = remove_mean(
+                    val, nobs, sum_x, neg_ct, compensation_remove
+                )
+
+            for j in range(end[i - 1], e):
+                val = values[j]
+                (
+                    nobs,
+                    sum_x,
+                    neg_ct,
+                    compensation_add,
+                    num_consecutive_same_value,
+                    prev_value,
+                ) = add_mean(
+                    val,
+                    nobs,
+                    sum_x,
+                    neg_ct,
+                    compensation_add,
+                    num_consecutive_same_value,
+                    prev_value,  # pyright: ignore[reportArgumentType]
+                )
+
+        if nobs >= min_periods and nobs > 0:
+            result = sum_x / nobs
+            if num_consecutive_same_value >= nobs:
+                result = prev_value
+            elif neg_ct == 0 and result < 0:
+                result = 0
+            elif neg_ct == nobs and result > 0:
+                result = 0
+        else:
+            result = np.nan
+
+        output[i] = result
+
+        if not is_monotonic_increasing_bounds:
+            nobs = 0
+            sum_x = 0.0
+            neg_ct = 0
+            compensation_remove = 0.0
+
+    # na_position is empty list since float64 can already hold nans
+    # Do list comprehension, since numba cannot figure out that na_pos is
+    # empty list of ints on its own
+    na_pos = [0 for i in range(0)]
+    return output, na_pos
+
+
+@numba.jit(nogil=True, parallel=False)
+def grouped_mean(
+    values: np.ndarray,
+    result_dtype: np.dtype,
+    labels: npt.NDArray[np.intp],
+    ngroups: int,
+    min_periods: int,
+    skipna: bool,
+) -> tuple[np.ndarray, list[int]]:
+    output, nobs_arr, comp_arr, consecutive_counts, prev_vals = grouped_kahan_sum(
+        values, result_dtype, labels, ngroups, skipna
+    )
+
+    # Post-processing, replace sums that don't satisfy min_periods
+    for lab in range(ngroups):
+        nobs = nobs_arr[lab]
+        num_consecutive_same_value = consecutive_counts[lab]
+        prev_value = prev_vals[lab]
+        sum_x = output[lab]
+        if nobs >= min_periods:
+            if num_consecutive_same_value >= nobs:
+                result = prev_value * nobs
+            else:
+                result = sum_x
+        else:
+            result = np.nan
+        result /= nobs
+        output[lab] = result
+
+    # na_position is empty list since float64 can already hold nans
+    # Do list comprehension, since numba cannot figure out that na_pos is
+    # empty list of ints on its own
+    na_pos = [0 for i in range(0)]
+    return output, na_pos

package/bench/refs/python-lib/pandas/core/_numba/kernels/min_max_.py

@@ -0,0 +1,178 @@
+"""
+Numba 1D min/max kernels that can be shared by
+* Dataframe / Series
+* groupby
+* rolling / expanding
+
+Mirrors pandas/_libs/window/aggregation.pyx
+"""
+
+from __future__ import annotations
+
+from typing import (
+    TYPE_CHECKING,
+    Any,
+)
+
+import numba
+import numpy as np
+
+if TYPE_CHECKING:
+    from pandas._typing import npt
+
+
+@numba.njit(nogil=True, parallel=False)
+def bisect_left(a: list[Any], x: Any, lo: int = 0, hi: int = -1) -> int:
+    """Same as https://docs.python.org/3/library/bisect.html; not in numba yet!"""
+    if hi == -1:
+        hi = len(a)
+    while lo < hi:
+        mid = (lo + hi) // 2
+        if a[mid] < x:
+            lo = mid + 1
+        else:
+            hi = mid
+    return lo
+
+
+@numba.jit(nogil=True, parallel=False)
+def sliding_min_max(
+    values: np.ndarray,
+    result_dtype: np.dtype,
+    start: np.ndarray,
+    end: np.ndarray,
+    min_periods: int,
+    is_max: bool,
+) -> tuple[np.ndarray, list[int]]:
+    # Basic idea of the algorithm: https://stackoverflow.com/a/12239580
+    # It was generalized to work with an arbitrary list of any window size and position
+    # by adding the Dominators stack.
+
+    N = len(start)
+    na_pos = []
+    output = np.empty(N, dtype=result_dtype)
+
+    def cmp(a: Any, b: Any, is_max: bool) -> bool:
+        if is_max:
+            return a >= b
+        else:
+            return a <= b
+
+    # Indices of bounded extrema in `values`. `candidates[i]` is always increasing.
+    # `values[candidates[i]]` is decreasing for max and increasing for min.
+    candidates: list[int] = []  # this is a queue
+    # Indices of largest windows that "cover" preceding windows.
+    dominators: list[int] = []  # this is a stack
+
+    if min_periods < 1:
+        min_periods = 1
+
+    if N > 2:
+        i_next = N - 1  # equivalent to i_next = i+1 inside the loop
+        for i in range(N - 2, -1, -1):
+            next_dominates = start[i_next] < start[i]
+            if next_dominates and (
+                not dominators or start[dominators[-1]] > start[i_next]
+            ):
+                dominators.append(i_next)
+            i_next = i
+
+    # NaN tracking to guarantee min_periods
+    valid_start = -min_periods
+
+    last_end = 0
+    last_start = -1
+
+    for i in range(N):
+        this_start = start[i].item()
+        this_end = end[i].item()
+
+        if dominators and dominators[-1] == i:
+            dominators.pop()
+
+        if not (
+            this_end > last_end or (this_end == last_end and this_start >= last_start)
+        ):
+            raise ValueError(
+                "Start/End ordering requirement is violated at index " + str(i)
+            )
+
+        stash_start = (
+            this_start if not dominators else min(this_start, start[dominators[-1]])
+        )
+        while candidates and candidates[0] < stash_start:
+            candidates.pop(0)
+
+        for k in range(last_end, this_end):
+            if not np.isnan(values[k]):
+                valid_start += 1
+                while valid_start >= 0 and np.isnan(values[valid_start]):
+                    valid_start += 1
+                while candidates and cmp(values[k], values[candidates[-1]], is_max):
+                    candidates.pop()  # Q.pop_back()
+                candidates.append(k)  # Q.push_back(k)
+
+        if not candidates or (this_start > valid_start):
+            if values.dtype.kind != "i":
+                output[i] = np.nan
+            else:
+                na_pos.append(i)
+        elif candidates[0] >= this_start:
+            # ^^ This is here to avoid costly bisection for fixed window sizes.
+            output[i] = values[candidates[0]]
+        else:
+            q_idx = bisect_left(candidates, this_start, lo=1)
+            output[i] = values[candidates[q_idx]]
+        last_end = this_end
+        last_start = this_start
+
+    return output, na_pos
+
+
+@numba.jit(nogil=True, parallel=False)
+def grouped_min_max(
+    values: np.ndarray,
+    result_dtype: np.dtype,
+    labels: npt.NDArray[np.intp],
+    ngroups: int,
+    min_periods: int,
+    is_max: bool,
+    skipna: bool = True,
+) -> tuple[np.ndarray, list[int]]:
+    N = len(labels)
+    nobs = np.zeros(ngroups, dtype=np.int64)
+    na_pos = []
+    output = np.empty(ngroups, dtype=result_dtype)
+
+    for i in range(N):
+        lab = labels[i]
+        val = values[i]
+        if lab < 0 or (not skipna and nobs[lab] >= 1 and np.isnan(output[lab])):
+            continue
+
+        if values.dtype.kind == "i" or not np.isnan(val):
+            nobs[lab] += 1
+        else:
+            if not skipna:
+                # If skipna is False and we encounter a NaN,
+                # both min and max of the group will be NaN
+                output[lab] = np.nan
+            continue
+
+        if nobs[lab] == 1:
+            # First element in group, set output equal to this
+            output[lab] = val
+            continue
+
+        if is_max:
+            if val > output[lab]:
+                output[lab] = val
+        elif val < output[lab]:
+            output[lab] = val
+
+    # Set labels that don't satisfy min_periods as np.nan
+    for lab, count in enumerate(nobs):
+        if count < min_periods:
+            na_pos.append(lab)
+
+    return output, na_pos

package/bench/refs/python-lib/pandas/core/_numba/kernels/sum_.py

@@ -0,0 +1,255 @@
+"""
+Numba 1D sum kernels that can be shared by
+* Dataframe / Series
+* groupby
+* rolling / expanding
+
+Mirrors pandas/_libs/window/aggregation.pyx
+"""
+
+from __future__ import annotations
+
+from typing import (
+    TYPE_CHECKING,
+    Any,
+)
+
+import numba
+from numba.extending import register_jitable
+import numpy as np
+
+if TYPE_CHECKING:
+    from pandas._typing import npt
+
+from pandas.core._numba.kernels.shared import is_monotonic_increasing
+
+
+@numba.jit(nogil=True, parallel=False)
+def add_sum(
+    val: Any,
+    nobs: int,
+    sum_x: Any,
+    compensation: Any,
+    num_consecutive_same_value: int,
+    prev_value: Any,
+) -> tuple[int, Any, Any, int, Any]:
+    if not np.isnan(val):
+        nobs += 1
+        y = val - compensation
+        t = sum_x + y
+        compensation = t - sum_x - y
+        sum_x = t
+
+        if val == prev_value:
+            num_consecutive_same_value += 1
+        else:
+            num_consecutive_same_value = 1
+        prev_value = val
+
+    return nobs, sum_x, compensation, num_consecutive_same_value, prev_value
+
+
+@numba.jit(nogil=True, parallel=False)
+def remove_sum(
+    val: Any, nobs: int, sum_x: Any, compensation: Any
+) -> tuple[int, Any, Any]:
+    if not np.isnan(val):
+        nobs -= 1
+        y = -val - compensation
+        t = sum_x + y
+        compensation = t - sum_x - y
+        sum_x = t
+    return nobs, sum_x, compensation
+
+
+@numba.jit(nogil=True, parallel=False)
+def sliding_sum(
+    values: np.ndarray,
+    result_dtype: np.dtype,
+    start: np.ndarray,
+    end: np.ndarray,
+    min_periods: int,
+) -> tuple[np.ndarray, list[int]]:
+    dtype = values.dtype
+
+    na_val: object = np.nan
+    if dtype.kind == "i":
+        na_val = 0
+
+    N = len(start)
+    nobs = 0
+    sum_x = 0
+    compensation_add = 0
+    compensation_remove = 0
+    na_pos = []
+
+    is_monotonic_increasing_bounds = is_monotonic_increasing(
+        start
+    ) and is_monotonic_increasing(end)
+
+    output = np.empty(N, dtype=result_dtype)
+
+    for i in range(N):
+        s = start[i]
+        e = end[i]
+        if i == 0 or not is_monotonic_increasing_bounds:
+            prev_value = values[s]
+            num_consecutive_same_value = 0
+
+            for j in range(s, e):
+                val = values[j]
+                (
+                    nobs,
+                    sum_x,
+                    compensation_add,
+                    num_consecutive_same_value,
+                    prev_value,
+                ) = add_sum(
+                    val,
+                    nobs,
+                    sum_x,
+                    compensation_add,
+                    num_consecutive_same_value,
+                    prev_value,
+                )
+        else:
+            for j in range(start[i - 1], s):
+                val = values[j]
+                nobs, sum_x, compensation_remove = remove_sum(
+                    val, nobs, sum_x, compensation_remove
+                )
+
+            for j in range(end[i - 1], e):
+                val = values[j]
+                (
+                    nobs,
+                    sum_x,
+                    compensation_add,
+                    num_consecutive_same_value,
+                    prev_value,
+                ) = add_sum(
+                    val,
+                    nobs,
+                    sum_x,
+                    compensation_add,
+                    num_consecutive_same_value,
+                    prev_value,
+                )
+
+        if nobs == 0 == min_periods:
+            result: object = 0
+        elif nobs >= min_periods:
+            if num_consecutive_same_value >= nobs:
+                result = prev_value * nobs
+            else:
+                result = sum_x
+        else:
+            result = na_val
+            if dtype.kind == "i":
+                na_pos.append(i)
+
+        output[i] = result
+
+        if not is_monotonic_increasing_bounds:
+            nobs = 0
+            sum_x = 0
+            compensation_remove = 0
+
+    return output, na_pos
+
+
+# Mypy/pyright don't like the fact that the decorator is untyped
+@register_jitable  # type: ignore[misc]
+def grouped_kahan_sum(
+    values: np.ndarray,
+    result_dtype: np.dtype,
+    labels: npt.NDArray[np.intp],
+    ngroups: int,
+    skipna: bool,
+) -> tuple[
+    np.ndarray, npt.NDArray[np.int64], np.ndarray, npt.NDArray[np.int64], np.ndarray
+]:
+    N = len(labels)
+
+    nobs_arr = np.zeros(ngroups, dtype=np.int64)
+    comp_arr = np.zeros(ngroups, dtype=values.dtype)
+    consecutive_counts = np.zeros(ngroups, dtype=np.int64)
+    prev_vals = np.zeros(ngroups, dtype=values.dtype)
+    output = np.zeros(ngroups, dtype=result_dtype)
+
+    for i in range(N):
+        lab = labels[i]
+        val = values[i]
+
+        if lab < 0 or np.isnan(output[lab]):
+            continue
+
+        if not skipna and np.isnan(val):
+            output[lab] = np.nan
+            nobs_arr[lab] += 1
+            comp_arr[lab] = np.nan
+            consecutive_counts[lab] = 1
+            prev_vals[lab] = np.nan
+            continue
+
+        sum_x = output[lab]
+        nobs = nobs_arr[lab]
+        compensation_add = comp_arr[lab]
+        num_consecutive_same_value = consecutive_counts[lab]
+        prev_value = prev_vals[lab]
+
+        (
+            nobs,
+            sum_x,
+            compensation_add,
+            num_consecutive_same_value,
+            prev_value,
+        ) = add_sum(
+            val,
+            nobs,
+            sum_x,
+            compensation_add,
+            num_consecutive_same_value,
+            prev_value,
+        )
+
+        output[lab] = sum_x
+        consecutive_counts[lab] = num_consecutive_same_value
+        prev_vals[lab] = prev_value
+        comp_arr[lab] = compensation_add
+        nobs_arr[lab] = nobs
+    return output, nobs_arr, comp_arr, consecutive_counts, prev_vals
+
+
+@numba.jit(nogil=True, parallel=False)
+def grouped_sum(
+    values: np.ndarray,
+    result_dtype: np.dtype,
+    labels: npt.NDArray[np.intp],
+    ngroups: int,
+    min_periods: int,
+    skipna: bool,
+) -> tuple[np.ndarray, list[int]]:
+    na_pos = []
+
+    output, nobs_arr, comp_arr, consecutive_counts, prev_vals = grouped_kahan_sum(
+        values, result_dtype, labels, ngroups, skipna
+    )
+
+    # Post-processing, replace sums that don't satisfy min_periods
+    for lab in range(ngroups):
+        nobs = nobs_arr[lab]
+        num_consecutive_same_value = consecutive_counts[lab]
+        prev_value = prev_vals[lab]
+        sum_x = output[lab]
+        if nobs >= min_periods:
+            if num_consecutive_same_value >= nobs:
+                result = prev_value * nobs
+            else:
+                result = sum_x
+        else:
+            result = sum_x  # Don't change val, will be replaced by nan later
+            na_pos.append(lab)
+        output[lab] = result
+
+    return output, na_pos