Skip to content

fluxopt.constraints.sparse

Sparse summation helpers for linopy expressions.

Avoids dense broadcast when coefficient arrays are highly sparse (e.g., conversion constraints where each converter only touches 2-3 flows out of hundreds).

Functions:

Name Description
sparse_weighted_sum

Compute (var * coeffs).sum(sum_dim) using only non-zero pairs.

sparse_weighted_sum 

sparse_weighted_sum(
    var: Variable, coeffs: DataArray, sum_dim: str, group_dim: str
) -> LinearExpression

Compute (var * coeffs).sum(sum_dim) using only non-zero pairs.

When coeffs is sparse along (group_dim, sum_dim) — e.g. a (converter, flow) matrix where each converter references only a few flows — the naive dense broadcast creates a huge intermediate linopy expression. This function finds the non-zero (group, sum_dim) pairs and uses groupby().sum() to aggregate, avoiding the dense broadcast.

Parameters:

Name Type Description Default

var

Variable

linopy Variable or LinearExpression with sum_dim.

 required

coeffs

DataArray

DataArray with at least (group_dim, sum_dim). Extra dims (e.g. eq_idx, time) are preserved.

 required

sum_dim

str

Dimension to sum over (e.g. 'flow').

 required

group_dim

str

Dimension to group by (e.g. 'converter').

 required

Returns:

Type Description
LinearExpression

linopy expression with sum_dim removed, group_dim present.
Source code in src/fluxopt/constraints/sparse.py 
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
def sparse_weighted_sum(
    var: linopy.Variable,
    coeffs: xr.DataArray,
    sum_dim: str,
    group_dim: str,
) -> linopy.LinearExpression:
    """Compute ``(var * coeffs).sum(sum_dim)`` using only non-zero pairs.

    When *coeffs* is sparse along ``(group_dim, sum_dim)`` — e.g. a
    ``(converter, flow)`` matrix where each converter references only a few
    flows — the naive dense broadcast creates a huge intermediate linopy
    expression.  This function finds the non-zero ``(group, sum_dim)`` pairs
    and uses ``groupby().sum()`` to aggregate, avoiding the dense broadcast.

    Args:
        var: linopy Variable or LinearExpression with *sum_dim*.
        coeffs: DataArray with at least ``(group_dim, sum_dim)``.
            Extra dims (e.g. ``eq_idx``, ``time``) are preserved.
        sum_dim: Dimension to sum over (e.g. ``'flow'``).
        group_dim: Dimension to group by (e.g. ``'converter'``).

    Returns:
        linopy expression with *sum_dim* removed, *group_dim* present.
    """
    coeffs_values = coeffs.values
    group_ids = list(coeffs.coords[group_dim].values)
    sum_ids = list(coeffs.coords[sum_dim].values)

    group_axis = coeffs.dims.index(group_dim)
    sum_axis = coeffs.dims.index(sum_dim)

    # Collapse extra axes to find any non-zero (group, sum) pair
    reduce_axes = tuple(i for i in range(coeffs_values.ndim) if i not in (group_axis, sum_axis))
    nonzero_2d = np.any(coeffs_values != 0, axis=reduce_axes) if reduce_axes else coeffs_values != 0

    # Ensure (group, sum) axis order
    if group_axis > sum_axis:
        nonzero_2d = nonzero_2d.T
    group_idx, sum_idx = np.nonzero(nonzero_2d)

    if len(group_idx) == 0:
        # All zeros — fall back to dense (cheap when empty)
        return (var * coeffs).sum(sum_dim)

    # Ensure all groups appear in the result (groups with all-zero coefficients
    # still need a zero entry so groupby produces an output for them)
    represented = set(group_idx.tolist())
    for g in range(len(group_ids)):
        if g not in represented:
            group_idx = np.append(group_idx, g)
            sum_idx = np.append(sum_idx, 0)  # coeff is 0 for absent groups

    pair_sum_ids = [sum_ids[s] for s in sum_idx]
    pair_group_ids = [group_ids[g] for g in group_idx]

    # Transpose coeffs so (group, sum, ...) are the leading axes,
    # ensuring NumPy advanced indexing places the pair axis first.
    remaining_dims = [d for d in coeffs.dims if d not in (group_dim, sum_dim)]
    ordered = coeffs.transpose(group_dim, sum_dim, *remaining_dims)
    ordered_values = ordered.values

    # Extract per-pair coefficients — advanced indexing on axes 0, 1
    pair_coeffs_data = ordered_values[group_idx, sum_idx]  # (pair, *remaining)

    remaining_coords = {d: coeffs.coords[d] for d in remaining_dims if d in coeffs.coords}
    pair_coeffs = xr.DataArray(
        pair_coeffs_data,
        dims=['pair', *remaining_dims],
        coords=remaining_coords,
    )

    # Select var for active pairs, multiply, group-sum
    selected = var.sel({sum_dim: xr.DataArray(pair_sum_ids, dims=['pair'])})
    weighted = selected * pair_coeffs

    mapping = xr.DataArray(pair_group_ids, dims=['pair'], name=group_dim)
    result = weighted.groupby(mapping).sum()

    # Restore original group order (groupby sorts alphabetically)
    result = result.sel({group_dim: group_ids})

    # Vectorized sel() leaves sum_dim as a non-dim coord — drop it
    return cast('linopy.LinearExpression', result.drop_vars(sum_dim, errors='ignore'))