routines.sage

# Various routines to be loaded first
from sage.all import *
from itertools import product, combinations
var('t x u n N rho')
R.<x> = PowerSeriesRing(ZZ)

def extract_instanton(order, a_values, max_n):
    """
    Extract c(n) from a(n) = sum_{d|n} c(n/d)/d^order
    
    a_values: list or dict of a(n) values (1-indexed)
    max_n: maximum n to compute
    """
    # Initialize c array
    c = [0] * (max_n + 1)
    
    for n in range(1, max_n + 1):
        # Start with a(n)
        c[n] = a_values[n]
        
        # Subtract contributions from all proper divisors
        # a(n) = c(n)/1 + sum_{d|n, d>1} c(n/d)/d^order
        # So: c(n) = a(n) - sum_{d|n, d>1} c(n/d)/d^order
        
        divs = divisors(n)
        for d in divs:
            if d > 1:  # Skip d=1
                c[n] -= c[n//d] / (d^order)
    
    return c

def count_solutions(L, n, return_points=False):
    """
    Counts integer solutions to: r_0 + r_1 + ... + r_L <= n
    where r_i >= 0.
    """
    # Number of variables is L + 1
    num_vars = L + 1
    
    # Define inequalities in Sage format: [constant, coeff_r0, coeff_r1, ...]
    # 1. The sum constraint: n - sum(r_i) >= 0
    sum_constraint = [n] + [-1] * num_vars
    
    # 2. Non-negativity: r_i >= 0
    non_negative_constraints = []
    for i in range(num_vars):
        row = [0] * (num_vars + 1)
        row[i + 1] = 1
        non_negative_constraints.append(row)
    
    # Combine constraints
    ieqs = [sum_constraint] + non_negative_constraints
    
    # Create the Polyhedron object
    P = Polyhedron(ieqs=ieqs, base_ring=ZZ)
    
    if return_points:
        return P.integral_points()
    else:
        return P.integral_points_count()

# Example usage:
# result = count_solution(L=2, n=5)
# print(f"number of terms to compute: {result}")


def ell_i(i, r, m_sq_i, p_sq):
    """
    Compute ell_i(r) = log(-m_i^2/p^2) + 2*(gamma_E - sum_{k=1}^{r-1} 1/k)
    
    For r = 0: the sum is empty (0), so ell_i(0) = log(m_i^2/p^2) + 2*gamma_E
    For r >= 1: ell_i(r) = log(-m_i^2/p^2) - 2* H_{r-1}
    """
    r = Integer(r)
    m_sq_i = SR(m_sq_i)
    p_sq = SR(p_sq)
    
    log_term = log(-m_sq_i / p_sq)
    if r == 0:
        harmonic_term = 0
    else:
        harmonic_term = harmonic_number(r)
    
    return log_term - 2 * harmonic_term


def harmonic_number(n):
    """Compute H_n = sum_{k=1}^n 1/k for n >= 1, return 0 for n=0."""
    n = Integer(n)
    if n <= 0:
        return SR(0)
    return sum(QQ(1)/k for k in range(1, n+1))

def elementary_symmetric_polynomial(values, k):
    """
    Compute the k-th elementary symmetric polynomial using Sage's SymmetricFunctions.
    """
    n = len(values)
    k = Integer(k)
    
    if k < 0:
        return SR(0)
    if k == 0:
        return SR(1)
    if k > n:
        return SR(0)
    
    # Create symmetric function ring over QQ
    e = SymmetricFunctions(QQ).e()
    e_k = e([k])
    
    # Expand in n variables - this creates a polynomial in x_0, ..., x_{n-1}
    f = e_k.expand(n)
    
    # Get the polynomial ring variables
    # f is in QQ[x_0, x_1, ..., x_{n-1}]
    # We need to map these to our values
    
    # Create a list of substitutions: the i-th variable in the ring -> values[i]
    # The variables are x_0, x_1, ... in order
    R = f.parent()
    gens = R.gens()  # This gives the generators (x_0, x_1, ...)
    
    subs_dict = {}
    for i in range(n):
        if i < len(gens):
            subs_dict[gens[i]] = SR(values[i])
    
    # Substitute and return
    result = f.subs(subs_dict)
    return SR(result)