Simplify 'sum_by_group'

[dad616610]

This PR replaces the _sum_by_group function. The main goals are to improve readability and maintainability; performance was improved as well

Readability and maintenance

IMO, it became more clear what the code does (maybe I became more familiar with the code and my judgement is skewed). Also all the work is now being done by one function, not scattered around: this should ease maintenance, since there's less code now

I added tests as well. This better documents the function and makes future changes more confident.

Performance

How I measured

I collected all the indices and values passed to _sum_by_group_np and made this my testing data.

Testing data collection

Assuming we're on a commit prior current one (a83733b), so we still have the old _sum_by_group_np.
The code below would save indices and values under "data" directory. I'm getting ~28k small files: 14k for indices and 14k for values.

from pathlib import Path

p = Path(__file__)
root = Path(*p.parts[: p.parts.index("pandapipes") + 1])
data_dir = root / "data"
data_dir.mkdir(exist_ok=True)
n_inputs = 0

def _sum_by_group_np(indices, *values):
    global n_inputs

    np.save(data_dir / f"{n_inputs:05}_inds.npy", indices)
    np.save(data_dir / f"{n_inputs:05}_vals.npy", np.vstack(values))
    n_inputs += 1

Then I created a file under the project root to test different function implementations, let's call it "test_sum_by_group.py". It's content is under the spoiler below. I test the function on all inputs 10 times (this is one round), do 3 such rounds and take a minimum.

Performance test file

The function implementations are shown in the next collapsible sections.

import numpy as np
import numba as nb

def ind_val():
    """Collect our data from previous step"""
    from collections import defaultdict

    root = Path("data")
    data = defaultdict(lambda: [None, None])

    for f in root.iterdir():
        idx, is_val = f.stem.split("_")
        data_idx = is_val == "vals"
        data[int(idx)][data_idx] = np.load(f)

    inds = []
    vals = []
    for ind, val in data.values():
        inds.append(ind)
        vals.append(val)

    return inds, vals

def main():
    """Main function where the testing takes place"""
    from timeit import repeat
    inds, vals = ind_val()

    def tf(f):
        """A helper function to time all inputs"""
        for _ind, _val in zip(inds, vals):
            f(_ind, _val)

    fs = (
        old,
        new_bincount,
        new_cumsum,
        nmb,
    )
    for f in fs:
        res = repeat(
            f"tf({f.__name__})",
            repeat=3,
            number=10,
            globals={"tf": tf} | {f.__name__: f for f in fs},
        )
        print(f"{f.__name__}: {'. '.join(map(lambda v: f'{v:.5f}', res))}")


main()

Functions:

old

This is the previous _sum_by_group_np

def _sum_by_group_sorted(indices, *values):
    """Auxiliary function to sum up values by some given indices (both as numpy arrays). Expects the
    indices and values to already be sorted.

    :param indices:
    :type indices:
    :param values:
    :type values:
    :return:
    :rtype:
    """
    # Index defines whether a specific index has already appeared in the index array before.
    index = np.ones(len(indices), "bool")
    index[:-1] = indices[1:] != indices[:-1]

    # make indices unique for output
    indices = indices[index]

    val = list(values)
    for i, _ in enumerate(val):
        # sum up values, chose only those with unique indices and then subtract the previous sums
        # --> this way for each index the sum of all values belonging to this index is returned
        nans = np.isnan(val[i])
        if np.any(nans):
            np.nan_to_num(val[i], copy=False)
            np.cumsum(val[i], out=val[i])
            val[i] = val[i][index]
            still_na = nans[index]
            val[i][1:] = val[i][1:] - val[i][:-1]
            val[i][still_na] = np.nan
        else:
            np.cumsum(val[i], out=val[i])
            val[i] = val[i][index]
            val[i][1:] = val[i][1:] - val[i][:-1]
    return [indices] + val


def old(indices, values):
    # sort indices and values by indices
    order = np.argsort(indices)
    inds = indices[order]
    val = list(values)
    for i, _ in enumerate(val):
        val[i] = val[i][order]

    res = _sum_by_group_sorted(inds, *val)
    return res

new_bincount

def new_bincount(indices, values):
    if indices.size == 0:
        return np.empty(0, dtype=int), [np.empty(0, dtype=val.dtype) for val in values]

    groups = indices.astype(int, copy=False)
    unique_indices = np.bincount(groups).nonzero()[0]
    vals = [
        np.bincount(groups, weights=val)[unique_indices].astype(val.dtype, copy=False)
        for val in values
    ]
    return unique_indices, vals

new_cumsum

def new_cumsum(indices, values):
    if indices.size == 0:
        return (
            np.empty(0, dtype=indices.dtype),
            [np.empty(0, dtype=val.dtype) for val in values],
        )
    order = np.argsort(indices)
    groups = indices[order]

    unique_idx = np.empty(groups.size, dtype=bool)
    unique_idx[0] = True
    np.not_equal(groups[1:], groups[:-1], out=unique_idx[1:])
    group_idx = unique_idx.nonzero()[0]

    return (
        groups[unique_idx],
        [np.add.reduceat(val[order], group_idx, dtype=val.dtype) for val in values],
    )

Numba-powered version of new_bincount (basically just adding a decorator)

nmb

@nb.njit
def nmb(indices, values):
    if indices.size == 0:
        return np.empty(0, dtype=np.int64), [
            np.empty(0, dtype=val.dtype) for val in values
        ]

    iinds = indices.astype(np.int64)
    idx = np.bincount(iinds)
    unique_indices = np.flatnonzero(idx)
    vals = [
        np.bincount(iinds, weights=val)[unique_indices].astype(val.dtype)
        for val in values
    ]
    return unique_indices, vals

There is no support for np.add.reduceat in numba and the regular np.cumsum with subtracting works even worse, so there's no numba-powered version of new_cumsum.

To run the testing suite I used uv run test_sum_by_group.py

Results

Here're the results I obtained on my machine (measurement unit is seconds):

old: 2.15977, 2.15104, 2.17834
new_bincount: 0.40900, 0.40652, 0.41300
new_cumsum: 0.96587, 0.94974, 0.94659
nmb: 1.58427, 0.30429, 0.30589

• new_bincount gives ~5x speedup; in terms of time-per-call (time / 10 (in one round I run a function 10 times) / 10_000 (roughly a number of inputs) we get 21 us vs. 4 us, which is pretty fast, if you ask me.
• nmb does run faster than the new_bincount by 30%, but the absolute speedup is marginal: 4 us time-per-call vs. 3 us. We should also take the compilation time into account here. It takes ~1.2 sec to compile, meaning numba version would
  outperform non-numba version only after 1.2 / (0.4 - 0.3) = 12 runs of performance test suite!
  This numba over non-numba gain deemed miniscule to me, so I decided not to include the numba version at all.
• While new_cumsum is 2x slower than new_bincount it has a serious advantage: new_bincount needs to allocate an array of indices.max() + 1 size (128 MiB for size 2**24 of dtype np.int64). With arrays this big we're spending more time allocating, than running an algorithm. new_cumsum doesn't suffer from this problem, but even with array of 2**24 size new_bincount was still faster. Arrays of this size (more correctly, indices with such numbers) are unlikely to be used and we would lose lots of performance, so I chose new_bincount over new_cumsum.

Performance of uv run pytest

The testing suite (uv run pytest) became a whopping 1s faster!

I measured 10 runs with hyperfine. This result is without this PR:

Benchmark 1: uv run pytest --quiet --disable-warnings src/pandapipes/test
  Time (mean ± σ):     46.235 s ±  0.232 s    [User: 42.459 s, System: 3.130 s]
  Range (min … max):   45.927 s … 46.635 s    10 runs

And this one is with this commit:

Benchmark 1: uv run pytest --quiet --disable-warnings src/pandapipes/test
  Time (mean ± σ):     44.770 s ±  0.203 s    [User: 41.297 s, System: 2.867 s]
  Range (min … max):   44.445 s … 45.073 s    10 runs

API change

• Made sum_by_group a public function (no particular reason)
• Changed the results signature

The previous version returned both unique indices and grouped sums as one list: [indices] + grouped_sums. I made the new function return a tuple of unique indices and grouped sums as separate objects: (indices, grouped_sums). All internal usages were updated.
See the change from the caller side in these diffs:

-juncts, loads_sum = _sum_by_group(loads.junction.values, mass_flow_loads)
+juncts, (loads_sum,) = sum_by_group(loads.junction.values, mass_flow_loads)

-tn_nodes, tn_eq_sum, tn_deriv_sum = _sum_by_group(
+tn_nodes, (tn_eq_sum, tn_deriv_sum) = sum_by_group(

This was done to ease numba-powering the functions, but IMO it also more clearly separates the output: we get unique_indices as separate entity and grouped_sums as another entity. Though, either returning method is equivalent performance-wise

If any of the API changes are unnecessary, I'm happy to revert them