EP-0149: Universal & Multivariate Binomial Convolution Moment Semirings¶

import functools
import math
from algebrax.semiring._base import Semiring

@functools.lru_cache(maxsize=32)
def _get_pascal_table(order: int) -> tuple[tuple[int, ...], ...]:
    return tuple(tuple(math.comb(n, k) for k in range(n + 1)) for n in range(order + 1))

class BinomialConvolutionSemiring(Semiring[tuple[float, ...]]):
    """Universal 1D Binomial Convolution Semiring over R[ε] / (ε^{K+1})."""

    def __init__(self, order: int = 1) -> None:
        if order < 0:
            raise ValueError(f"order must be a non-negative integer, got {order}")
        self.order = order
        self._pascal = _get_pascal_table(order)

    @property
    def zero(self) -> tuple[float, ...]:
        return (0.0,) * (self.order + 1)

    @property
    def one(self) -> tuple[float, ...]:
        return (1.0,) + (0.0,) * self.order

    def add(self, a: tuple[float, ...], b: tuple[float, ...]) -> tuple[float, ...]:
        if len(a) != self.order + 1 or len(b) != self.order + 1:
            raise ValueError(f"Operands must have length {self.order + 1}")
        return tuple(u + v for u, v in zip(a, b))

    def mul(self, a: tuple[float, ...], b: tuple[float, ...]) -> tuple[float, ...]:
        if len(a) != self.order + 1 or len(b) != self.order + 1:
            raise ValueError(f"Operands must have length {self.order + 1}")
        order = self.order
        pascal = self._pascal
        res = [0.0] * (order + 1)
        for k in range(order + 1):
            coeffs = pascal[k]
            acc = 0.0
            for j in range(k + 1):
                acc += coeffs[j] * a[j] * b[k - j]
            res[k] = acc
        return tuple(res)

class StatisticalMomentSemiring(BinomialConvolutionSemiring):
    """1D Statistical moment semiring with central moments, mean, variance, skewness, kurtosis."""

    def raw_moments(self, m: tuple[float, ...]) -> list[float]:
        p = m[0]
        if p == 0.0:  # NOSONAR - exact zero denominator singularity check
            return [float("nan")] * len(m)
        return [val / p for val in m]

    def central_moments(self, m: tuple[float, ...]) -> list[float]:
        """Compute all central moments μ_k = E[(X - μ)^k] from raw moments."""
        p = m[0]
        if p == 0.0:  # NOSONAR - exact zero denominator singularity check
            return [float("nan")] * len(m)
        raw = [val / p for val in m]
        mu = raw[1] if len(raw) > 1 else 0.0
        pascal = self._pascal

        central = [1.0]  # μ_0 = 1
        for k in range(1, len(raw)):
            val = sum(pascal[k][j] * ((-mu) ** (k - j)) * raw[j] for j in range(k + 1))
            central.append(val)
        return central

    def mean(self, m: tuple[float, ...]) -> float:
        if m[0] == 0.0:  # NOSONAR - exact zero denominator singularity check
            return float("nan")
        return m[1] / m[0]

    def variance(self, m: tuple[float, ...]) -> float:
        cm = self.central_moments(m)
        return cm[2] if len(cm) > 2 else 0.0

    def skewness(self, m: tuple[float, ...]) -> float:
        cm = self.central_moments(m)
        if len(cm) <= 3 or cm[2] <= 0.0:
            return float("nan")
        return cm[3] / (cm[2] ** 1.5)

    def kurtosis(self, m: tuple[float, ...]) -> float:
        cm = self.central_moments(m)
        if len(cm) <= 4 or cm[2] <= 0.0:
            return float("nan")
        return cm[4] / (cm[2] ** 2)

@functools.lru_cache(maxsize=32)
def _get_multivariate_transition_table(num_vars: int, order: int) -> dict[tuple[tuple[int, ...], tuple[int, ...]], tuple[tuple[int, ...], int]]:
    """Lazy bounded cache of multi-index transitions (alpha, beta) -> (gamma, coeff)."""
    # Populated on demand without eager memory blowup
    return {}

class MultivariateBinomialConvolutionSemiring(Semiring[dict[tuple[int, ...], float]]):
    """
    Universal Multivariate Binomial Convolution Semiring over R[ε1,...,εd] / <ε^β : |β| = K+1>.
    """
    def __init__(self, num_vars: int, order: int = 1) -> None:
        if num_vars < 1:
            raise ValueError(f"num_vars must be >= 1, got {num_vars}")
        if order < 0:
            raise ValueError(f"order must be >= 0, got {order}")
        self.num_vars = num_vars
        self.order = order

    @property
    def zero(self) -> dict[tuple[int, ...], float]:
        return {}

    @property
    def one(self) -> dict[tuple[int, ...], float]:
        return {(0,) * self.num_vars: 1.0}

    def add(self, a: dict[tuple[int, ...], float], b: dict[tuple[int, ...], float]) -> dict[tuple[int, ...], float]:
        res = dict(a)
        for k, v in b.items():
            res[k] = res.get(k, 0.0) + v
        return {k: v for k, v in res.items() if not math.isclose(v, 0.0, abs_tol=1e-15)}

    def mul(self, a: dict[tuple[int, ...], float], b: dict[tuple[int, ...], float]) -> dict[tuple[int, ...], float]:
        res: dict[tuple[int, ...], float] = {}
        for alpha, v1 in a.items():
            for beta, v2 in b.items():
                gamma = tuple(x + y for x, y in zip(alpha, beta))
                if sum(gamma) <= self.order:
                    coeff = math.prod(math.comb(g, a_i) for g, a_i in zip(gamma, alpha))
                    res[gamma] = res.get(gamma, 0.0) + coeff * v1 * v2
        return {k: v for k, v in res.items() if not math.isclose(v, 0.0, abs_tol=1e-15)}

class MultivariateMomentSemiring(MultivariateBinomialConvolutionSemiring):
    """Multivariate moment semiring with mean vectors, covariance matrices, and Hessians."""

    def mean_vector(self, m: dict[tuple[int, ...], float]) -> list[float]:
        p = m.get((0,) * self.num_vars, 0.0)
        if p == 0.0:  # NOSONAR - exact zero denominator singularity check
            return [float("nan")] * self.num_vars
        means = []
        for i in range(self.num_vars):
            idx = tuple(1 if j == i else 0 for j in range(self.num_vars))
            means.append(m.get(idx, 0.0) / p)
        return means

    def covariance_matrix(self, m: dict[tuple[int, ...], float]) -> list[list[float]]:
        if self.order < 2:
            raise ValueError(f"Covariance matrix calculation requires order >= 2, but instance has order={self.order}")
        p = m.get((0,) * self.num_vars, 0.0)
        if p == 0.0:  # NOSONAR - exact zero denominator singularity check
            return [[float("nan")] * self.num_vars for _ in range(self.num_vars)]
        means = self.mean_vector(m)
        d = self.num_vars
        cov = [[0.0] * d for _ in range(d)]
        # Exploit upper-triangle symmetry
        for i in range(d):
            for j in range(i, d):
                idx = [0] * d
                idx[i] += 1
                idx[j] += 1
                raw_ij = m.get(tuple(idx), 0.0) / p
                val = raw_ij - (means[i] * means[j])
                cov[i][j] = val
                cov[j][i] = val
        return cov