Kwker

Statistics and data helpers

These calls cover common steps of data and machine-learning pipelines that sit on top of sorting. Each one gives the same result as the pandas, NumPy, SciPy or PyTorch code named in its section, without sorting more than it needs.

Running totals, lags and ranks within groups

group_cumsum(values, groups) gives each row the running sum of its group, in row order. It matches pandas groupby(groups)[values].cumsum(). group_shift(values, groups, periods) gives each row the value periods rows earlier in its group, and group_diff gives the difference to it. rank_by_group(values, groups) ranks each value within its group, as pandas groupby(...).rank().

PythonRuns on your machine.
import numpy as np
import kwker

store = np.array([1, 2, 1, 1, 2])
sales = np.array([10.0, 3.0, 5.0, 2.0, 4.0])
print(kwker.group_cumsum(sales, store))
print(kwker.group_shift(sales, store, 1))
print(kwker.group_diff(sales, store, 1))
print(kwker.rank_by_group(sales, store))
Output
[10.  3. 15. 17.  7.]
[nan nan 10.  5.  3.]
[nan nan -5. -3.  1.]
[3. 1. 2. 1. 2.]

rank_by_group takes the pandas methods "average", "min", "max", "first" and "dense", and pct=True for ranks between 0 and 1.

Expanding medians and quantiles

expanding_median(a) gives, for each position, the median of all values up to it. expanding_quantile(a, q) does the same for any quantile. They match pandas Series.expanding().median() and .quantile(q).

PythonRuns on your machine.
import numpy as np
import kwker

prices = np.array([4.0, 1.0, 3.0, 5.0, 2.0])
print(kwker.expanding_median(prices))
print(kwker.expanding_quantile(prices, 0.25))
Output
[4.  2.5 3.  3.5 3. ]
[4.   1.75 2.   2.5  2.  ]

For a window of fixed length, use rolling_median and rolling_quantile (see Top-k and selection).

Trimmed means and weighted quantiles

trim_mean(a, proportion, axis=0) averages each lane without its most extreme values: it drops the proportion share of smallest values and the same share of largest ones first. It is scipy.stats.trim_mean. One faulty value cannot pull a trimmed mean far, so federated learning uses it to combine client updates.

weighted_quantile(a, q, weights) finds the quantile when each value counts as much as its weight. It gives the same result as NumPy's np.quantile(a, q, weights=weights, method="inverted_cdf"). With counts as weights, it is the quantile of the data with each value repeated that many times.

PythonRuns on your machine.
import numpy as np
import kwker

# 5 clients x 2 parameters: client 4 sends garbage
updates = np.array([[0.10, 0.20], [0.12, 0.18], [0.11, 0.22], [5.00, -4.00], [0.09, 0.21]])
print(kwker.trim_mean(updates, 0.2, axis=0))

values = np.array([10.0, 20.0, 30.0, 40.0])
counts = np.array([1, 5, 1, 1])
print(kwker.weighted_quantile(values, [0.5, 0.9], counts))
Output
[0.11       0.19666667]
[20. 40.]

Ranking metrics

average_precision_score(y_true, y_score) and roc_auc_score(y_true, y_score) score a binary classifier from its labels and scores. They match scikit-learn's functions of the same names. Average precision weights each threshold's precision by the recall it adds. ROC AUC is the chance that a random positive scores above a random negative. Labels can be 0 / 1, booleans, or any values with pos_label naming the positive one.

PythonRuns on your machine.
import numpy as np
import kwker

y = np.array([0, 1, 1, 0, 1])
score = np.array([0.1, 0.8, 0.35, 0.4, 0.9])
print(round(kwker.average_precision_score(y, score), 4))
print(round(kwker.roc_auc_score(y, score), 4))
Output
0.9167
0.8333

Compare two samples

cdf_distance(a, b, kind) measures how far apart two samples are, from their sorted values alone. Use it for drift monitoring and distribution-shift tests.

kind Result Same as
"ks" The largest gap between the two empirical CDFs scipy.stats.ks_2samp(a, b).statistic
"wasserstein" The area between the two CDFs (earth mover's distance) scipy.stats.wasserstein_distance(a, b)
"energy" The energy distance scipy.stats.energy_distance(a, b)
PythonRuns on your machine.
import numpy as np
import kwker

train = np.array([1.0, 2.0, 2.0, 3.0, 4.0])
live = np.array([2.0, 3.0, 3.0, 5.0])
for kind in ("ks", "wasserstein", "energy"):
    print(kind, round(kwker.cdf_distance(train, live, kind), 6))
Output
ks 0.35
wasserstein 0.85
energy 0.674537

Count discordant pairs

count_inversions(a) counts the pairs i < j whose values stand in the wrong order: a[i] > a[j]. Equal values are not counted. For two rankings, sort by the first and count the inversions of the second. That count is Kendall tau's number of discordant pairs.

PythonRuns on your machine.
import numpy as np
import kwker

ranking = np.array([1, 3, 2, 5, 4])
print(kwker.count_inversions(ranking))
print(kwker.count_inversions(ranking, descending=True))
Output
2
8

Histograms

histogram(x, bins, range) counts values in uniform bins and returns the counts and the bin edges. The counts match NumPy np.histogram(x, bins, range=range) bin for bin.

PythonRuns on your machine.
import numpy as np
import kwker

x = np.array([0.1, 0.4, 0.5, 0.9, 1.0, 2.0])
counts, edges = kwker.histogram(x, bins=4, range=(0.0, 1.0))
print(counts)
print(edges)
Output
[1 1 1 2]
[0.   0.25 0.5  0.75 1.  ]

Point-in-time lookups (as-of joins)

asof_indices(left_on, right_on, left_by, right_by, direction) finds, for each left row, the right row with the latest time at or before it. It matches pandas merge_asof and Polars join_asof. A feature store uses it to attach the feature values that were known at each event's time. With left_by and right_by, only rows with equal keys match. The result is the index of the matching right row, or -1 when there is none.

PythonRuns on your machine.
import numpy as np
import kwker

event_time = np.array([5, 12, 30, 7])
event_user = np.array([1, 1, 2, 2])
snapshot_time = np.array([0, 10, 0, 20, 8])
snapshot_user = np.array([1, 1, 2, 2, 2])
print(kwker.asof_indices(event_time, snapshot_time, event_user, snapshot_user))
print(kwker.asof_indices(event_time, snapshot_time, event_user, snapshot_user, direction="forward"))
Output
[0 1 3 2]
[ 1 -1 -1  4]

direction is "backward" (the default), "forward" (the first time at or after) or "nearest". The rows do not need to be sorted. Times can be integers or NumPy datetime64 values.

Sizes of row-wise sets

row_intersect_count(a, b) counts, for each row, the ids that appear in both a and b. Recall@k is this count divided by k. row_unique_count(a) counts the distinct values in each row, for example after merging candidate lists from several retrievers.

PythonRuns on your machine.
import numpy as np
import kwker

retrieved = np.array([[4, 7, 1], [2, 9, 5]])
relevant = np.array([[1, 4, 8], [3, 6, 0]])
print(kwker.row_intersect_count(retrieved, relevant))
candidates = np.array([[3, 1, 3, 2], [5, 5, 5, 5]])
print(kwker.row_unique_count(candidates))
Output
[2 0]
[3 1]

Sampling filters for language models

These filters cut a row of logits down to the tokens worth sampling from, without sorting the vocabulary.

The last three take NumPy arrays or PyTorch tensors, and filter the last axis of each row.

PythonRuns on your machine.
import numpy as np
import kwker

logits = np.array([2.0, 1.0, 0.5, -1.0, 3.0], dtype=np.float32)
print(kwker.top_p(logits, 0.8))
print(kwker.top_n_sigma(logits, 1.0))
print(kwker.min_p_filter(logits, 0.2))

batch = np.array([[2.0, 1.0, 0.5, -1.0, 3.0],
                  [0.0, 4.0, 1.0, 2.0, 3.0]], dtype=np.float32)
print(kwker.top_k_filter(batch, [2, 3]))
Output
[4 0]
[  2. -inf -inf -inf   3.]
[  2. -inf -inf -inf   3.]
[[  2. -inf -inf -inf   3.]
 [-inf   4. -inf   2.   3.]]

Distinct strings

unique_strings(strings) returns the distinct strings of a list, sorted, like np.unique on strings. With return_index=True it also returns where each one first appears.

PythonRuns on your machine.
import kwker

urls = ["b.com/x", "a.com", "b.com/x", "a.com/y", "a.com"]
print(kwker.unique_strings(urls))
print(kwker.unique_strings(urls, return_index=True))
Output
['a.com', 'a.com/y', 'b.com/x']
(['a.com', 'a.com/y', 'b.com/x'], array([1, 3, 0], dtype=uint64))

Notes