Chebyshev pseudoprime test

CR228847.numpy

#!/usr/bin/python3

import timeit
from numpy import mod
from numpy.polynomial.polynomial import Polynomial


def is_prime_naive(n):
    if n <= 2 or n % 2 == 0:
        return n == 2
    return all(n % i != 0 for i in range(3, int(n ** 0.5 + 1), 2))


def is_pseudoprime_chebyshev(n):
    r = smallest_r(n)
    if r == 0:
        return n == 2

    xp, _ = tpair_mod_n_xpowrm1(n, n, r)
    return xp.trim().has_samecoef(Polynomial([0] * (n % r) + [1]))


def smallest_r(n):
    if n <= 1 or n % 2 == 0:
        return 0

    for r in range(5, 1000000, 2):
        if not is_prime_naive(r):
            continue

        u = n % r
        if u == 0 and r < u:
            return 0
        if u > 1 and u < r - 1:
            return r


def tpair_mod_n_xpowrm1(m, n, r):
    """Returns the tuple (T_m, T_{m+1}) modulo (n, x^r - 1)
       where T_i is the ith Chebyshev polynomial"""
    if m == 0:
        return (Polynomial([1]), Polynomial([0, 1]))

    T_k, T_kinc = tpair_mod_n_xpowrm1(m >> 1, n, r)

    T_2kinc = reduce_poly_mod_n_xpowrm1(2 * T_k * T_kinc + Polynomial([0, n-1]), n, r)

    if m & 1:
        # m = 2k+1, m+1 = 2k+2
        T_m = T_2kinc
        T_minc = reduce_poly_mod_n_xpowrm1(2 * T_kinc * T_kinc + Polynomial([n-1]), n, r)
    else:
        # m = 2k, m+1 = 2k+1
        T_m = reduce_poly_mod_n_xpowrm1(2 * T_k * T_k + Polynomial([n-1]), n, r)
        T_minc = T_2kinc
    return (T_m, T_minc)


def reduce_poly_mod_n_xpowrm1(poly, n, r):
    reduced = poly
    if (len(reduced.coef) > r):
        reduced = Polynomial(reduced.coef[0:r]) + Polynomial(reduced.coef[r:])
    reduced.coef = mod(reduced.coef, n)
    return reduced


if __name__ == "__main__":
    # Sanity checks
    for n in range(1000):
        assert is_pseudoprime_chebyshev(n) == is_prime_naive(n)

    # Timing
    print(timeit.timeit('for n in range(1000): is_prime_naive(n)',
                        setup="from __main__ import is_prime_naive",
                        number=10))
    print(timeit.timeit('for n in range(1000): is_pseudoprime_chebyshev(n)',
                        setup="from __main__ import is_pseudoprime_chebyshev",
                        number=10))

CR228847.pypy

#!/usr/bin/pypy3

import timeit


def is_prime_naive(n):
    if n <= 2 or n % 2 == 0:
        return n == 2
    return all(n % i != 0 for i in range(3, int(n ** 0.5 + 1), 2))


def is_pseudoprime_chebyshev(n):
    r = smallest_r(n)
    if r == 0:
        return n == 2

    # We work with polynomials represented in little-endian format
    # (poly[i] is the coefficient of x^i)
    xp, _ = tpair_mod_n_xpowrm1(n, n, r)
    return normalise_poly(xp) == [0] * (n % r) + [1]


def smallest_r(n):
    if n <= 1 or n % 2 == 0:
        return 0

    for r in range(5, 1000000, 2):
        if not is_prime_naive(r):
            continue

        u = n % r
        if u == 0 and r < u:
            return 0
        if u > 1 and u < r - 1:
            return r


def tpair_mod_n_xpowrm1(m, n, r):
    """Returns the tuple (T_m, T_{m+1}) modulo (n, x^r - 1)
       where T_i is the ith Chebyshev polynomial"""
    if m == 0:
        return ([1], [0, 1])

    T_k, T_kinc = tpair_mod_n_xpowrm1(m >> 1, n, r)

    # T_2kinc = 2 * T_k * T_kinc - x
    T_2kinc = [2 * coeff % n
               for coeff in convolve_mod_n_xpowrm1(T_k, T_kinc, n, r)]
    T_2kinc = poly_add_mod_n(T_2kinc, [0, n - 1], n)

    if m & 1:
        # m = 2k+1, m+1 = 2k+2
        T_m = T_2kinc
        # T_minc = 2 * T_kinc * T_kinc - 1
        # TODO Optimised polynomial squaring
        T_minc = [2 * coeff % n
                  for coeff in convolve_mod_n_xpowrm1(T_kinc, T_kinc, n, r)]
        T_minc = poly_add_mod_n(T_minc, [n - 1], n)
    else:
        # m = 2k, m+1 = 2k+1
        # T_m = 2 * T_k * T_k - 1
        T_m = [2 * coeff % n
               for coeff in convolve_mod_n_xpowrm1(T_k, T_k, n, r)]
        T_m = poly_add_mod_n(T_m, [n - 1], n)
        T_minc = T_2kinc
    return (T_m, T_minc)


def convolve_mod_n_xpowrm1(p, q, n, r):
    """ Calculates the convolution p * q modulo (n, x^r - 1) """
    if p == [] or q == []:
        return []

    conv = [0] * min(len(p) + len(q) - 1, r)
    for i, p_i in enumerate(p):
        for j, q_j in enumerate(q):
            idx = (i + j) % r
            conv[idx] = (conv[idx] + p_i * q_j) % n

    return conv


def poly_add_mod_n(p, q, n):
    r = [0] * max(len(p), len(q))

    for i, p_i in enumerate(p):
        r[i] = p_i

    for j, q_j in enumerate(q):
        r[j] = (r[j] + q_j) % n

    return r


def normalise_poly(p):
    while len(p) > 0 and p[-1] == 0:
        p = p[:-1]

    return p


if __name__ == "__main__":
    # Sanity checks
    for n in range(1000):
        assert is_pseudoprime_chebyshev(n) == is_prime_naive(n)

    # Timing
    print(timeit.timeit('for n in range(1000): is_prime_naive(n)',
                        setup="from __main__ import is_prime_naive",
                        number=10))
    print(timeit.timeit('for n in range(1000): is_pseudoprime_chebyshev(n)',
                        setup="from __main__ import is_pseudoprime_chebyshev",
                        number=10))