Source code for boundlab.ops.sparse_poly.sum

"""Coalescing addition of sparse polynomial generators."""

from __future__ import annotations

from typing import TYPE_CHECKING

import torch

from boundlab.ops.sparse_poly.structure import coalesce

if TYPE_CHECKING:
    from boundlab.polysp.generator import Generator




def sum(*generators: "Generator") -> "Generator":
    """Add generators and coalesce identical monomial rows."""
    from boundlab.polysp.generator import Generator

    if not generators:
        raise ValueError("sum requires at least one generator")
    
    gens = list(generators)
    max_order = max(gen.order for gen in gens)
    for i in range(len(gens)):
        gens[i] = gens[i].pad_order(max_order)

    first = gens[0]
    for gen in gens[1:]:
        if gen.order != first.order or gen.error_len != first.error_len:
            raise ValueError("generators must have the same order and error table")
        if gen.data.shape[1:] != first.data.shape[1:]:
            raise ValueError("generators must have the same coefficient shape")

    indices = torch.cat([gen.indices for gen in gens], dim=0)
    data = torch.cat([gen.data for gen in gens], dim=0)
    unique, inverse = torch.unique(indices, dim=0, sorted=True, return_inverse=True)
    output = coalesce(
        data,
        inverse,
        unique.shape[0],
    )
    return Generator(output, unique, first.error_len)



__all__ = ["sum"]