Coursera Modular forms 2017 exercise I-2-2

exercise-I-2-2.py

#!/usr/bin/python3

from math import sqrt, factorial
from operator import mul

def binomial(n, k):
    return factorial(n)/(factorial(n-k) * factorial(k))

def dot_prod(a, b):
    return sum(map(mul, a, b))

def square(x):
    return x*x
            
def norm2(v):
    return sum(map(square, v))

def possible_vs(m, n):
    vv = 2 * n
    max_x = int(sqrt(vv))
    d = 2 * max_x + 1
    for i in range(d**m):
        v = decode_v(m, d, i)
        v2 = norm2(v)
        if vv == norm2(v):
            yield v
        
def decode_v(m, d, i):
    """
    Returns the decoded vector v corresponding to the integer i.
    i: an integer encoding the vector v in D_m
    d: odd number which satisfies d >= 3. given by 2*int(sqrt(2*n)) + 1.
    i is converted into d-nary representation, and j-th digit represents
    the j-th component of the vector v, with the mapping
    0, 1, ... d-1 <-> -x, -x+1, ..., 0, 1, ... x where x is the maximum
    value that one component can take, given by int(sqrt(2*n)) (v.v = 2*n).
    m: dimension of the lattice D_m, i.e. the length of the vector v.
    Caution: the algorithm is pretty ineffective, it searches d**m vectors
    exhaustively.
    """
    if d < 3: # exceptional case. happens only when n = 0.
        return [0]*m
    digits = []
    # maximum value that a component can take.
    # following lines just converts i into d-ary number representation.
    x = (d-1)//2
    while i >= d:
        i0 = i
        d0 = d
        i, r = divmod(i, d)
        digits.append(r)
    digits.append(i)
    for j in range(len(digits)):
        digits[j] -= x
    # fills empty digits with minimum values.
    digits.extend([-x] * (m - len(digits)))
    return list(digits)

def jts_coeffs(m, n_max):
    """calculates Jacobi theta series of D_m, with u fixed to
    u = (1, -1, 0, ..., 0) up to the given order n."""
    u = [0]*m
    u[0] = 1
    u[1] = -1
    coeffs = {}
    for n in range(n_max + 1):
        for v in possible_vs(m, n):
            l = dot_prod(u, v)
            if (n, l) in coeffs:
                coeffs[(n, l)] += 1
            else:
                coeffs[(n, l)] = 1
    return coeffs

def convert_to_string(coeffs):
    output = ""
    for exponent in sorted(coeffs):
        n, l = exponent
        a = coeffs[exponent]
        term = '%d*q^%d*r^%d' % (a, n, l)
        if output:
            output = output + ' + ' + term
        else:
            output = term
    return output

#my pre-calculated solutions
a = {(0, 0):(lambda m:1),
     (1, 0):(lambda m:4*binomial(m-2, 2) + 2),
     (1, 1):(lambda m:4*(m-2)),
     (2, 0):(lambda m:2**4 * binomial(m-2, 4) + 2**3 * binomial(m-2, 2)
             + 2*(m-2)),
     (2, 1):(lambda m:2**4 * binomial(m-2, 3)),
     (2, 2):(lambda m:2*2 * binomial(m-2, 2) + 2),
     (3, 0):(lambda m:2**6 * binomial(m-2, 6) + 2**5 * binomial(m-2, 4)
             + 2**2 * factorial(3) * binomial(m-2, 3) + 2**2 * (m-2))
     }

if __name__ == '__main__':
    for n in range(4):
        m_min = 2*n + 2
        for m in range(m_min, m_min + 3):
            coeffs = jts_coeffs(m, n)
            # the follwing lines are to compare the output with
            # precaluculated values.
            # for key in sorted(coeffs):
            #     j, l = key
            #     a0 = coeffs[key]
            #     print('a0(%d, %d) = %d' % (j, l, a0))
            #     if key in a:
            #         a1 = a[key](m)
            #         print('a1(%d, %d) = %d' % (j, l, a1))
            #         if a0 == a1:
            #             print('match')
            #         else:
            #             print('not match')
            print('jacobi_theta(%d, %d) = %s'
                  % (m, n, convert_to_string(coeffs)))