gf.py

import sympy as sp
import itertools as it
from sympy.physics.quantum import TensorProduct


def companion_matrix(p):
    """Gives companion matrix for polynomial p."""
    p = sp.Poly(p)
    assert p.is_monic

    n = p.degree()
    M = [[0] * (n-1) + [-p.nth(0)]]
    for k in range(1, n):
        M.append([0] * (k-1) + [1] + [0] * (n-1-k) + [-p.nth(k)])

    return sp.Matrix(M)

def normalize_rational_gf(f):
    """Returns two polynomials P, Q, such that f = P/Q and Q has constant term 1."""
    f = sp.cancel(f)
    p, q = f.as_numer_denom()

    # Constant term 1.
    qp = sp.Poly(q)
    mult = qp.EM().as_expr() * qp.EC()
    p *= mult
    q /= mult

    return sp.expand(p), sp.expand(q)

def extract_coeff(p, q, x):
    c = p.coeff(x, 0) / q.coeff(x, 0)
    p -= c*q
    p = sp.expand((p - p.coeff(x, 0)) / x)
    return p, q, c

def rational_gf_to_coeff(f):
    """Given rational generating function f, generates the coefficients for it."""
    p, q = sp.cancel(f).as_numer_denom()
    x, *_ = f.free_symbols
    while True:
        p, q, c = extract_coeff(p, q, x)
        yield c

def rational_gf_to_coeff_n(f, n):
    """Given rational generating function f, generates the first n coefficients for it."""
    return list(it.islice(rational_gf_to_coeff(f), n))

# Algorithm from https://math.stackexchange.com/a/2520666/5558.
def hadamard_product(a, b):
    """Given two rational generating functions a and b, calculates their Hadamard product."""
    assert a.free_symbols == b.free_symbols and len(a.free_symbols) == 1

    x, *_ = a.free_symbols
    p, q = normalize_rational_gf(a)
    r, s = normalize_rational_gf(b)

    finite_part = 0
    n = 0
    while sp.degree(p, x) >= sp.degree(q, x) or sp.degree(r, x) >= sp.degree(s, x):
        p, q, ac = extract_coeff(p, q, x)
        r, s, bc = extract_coeff(r, s, x)
        finite_part += ac*bc*x**n
        n += 1

    cq = sp.expand(x**sp.degree(q, x) * q.subs({x: 1/x}))
    cs = sp.expand(x**sp.degree(s, x) * s.subs({x: 1/x}))
    M = companion_matrix(cq)
    N = companion_matrix(cs)

    va = sp.Matrix(rational_gf_to_coeff_n(p/q, M.shape[0]))
    vb = sp.Matrix(rational_gf_to_coeff_n(r/s, N.shape[0]))
    ua = sp.Matrix([[1] + [0] * (M.shape[0] - 1)])
    ub = sp.Matrix([[1] + [0] * (N.shape[0] - 1)])

    MN = TensorProduct(M, N)
    v = TensorProduct(va, vb)
    u = TensorProduct(ua, ub)

    result = (u * (sp.eye(MN.shape[0]) - MN.T*x).inv() * v)[0,0]

    return sp.cancel(finite_part + result*x**n)


if __name__ == "__main__":
    x = sp.var("x")

    a = x / (1 - x - x**2)
    b = x / (1 - x)
    f = hadamard_product(a, b)

    print(f)
    print(sp.factor(f))
    print(sp.simplify(f))
    print("---")
    print(rational_gf_to_coeff_n(a, 20))
    print(rational_gf_to_coeff_n(b, 20))
    print(rational_gf_to_coeff_n(f, 20))