determine sum of squares for 2467-digit prime, simplified for pycg callgraph analysis (manually transformed to Graphviz)

sq2_cg.py

# pylint: disable=C0103
#                 invalid-name
"""
determine sum of squares for 2467-digit prime
- minimized for pycg analysis
- original: https://github.com/Hermann-SW/RSA_numbers_factored/blob/main/python/sq2.py
- manually converted to Graphviz with splines=ortho
- shortened GraphvizFiddle Share link: https://stamm-wilbrandt.de/sq2_cg.py.gvf.html
"""
from time import time
from sys import stderr, argv
from typing import Tuple

def SECTION1():
    """
    Robert Chapman 2010 code from https://math.stackexchange.com/a/5883/1084297
    with small changes:
    - asserts instead bad case returns
    - renamed root4() to root4m1() indicating which 4th root gets determined
    - made sq2() return tuple with positive numbers; before sq2(13) returned (-3,-2)
    - sq2(p) result can be obtained from sympy.solvers.diophantine.diophantine by diop_DN(-1, p)[0]
    """
    return


def mods(a: int, n: int) -> int:
    """returns "signed" a (mod n), in range -n//2..n//2"""
    assert n > 0
    a = a % n
    if 2 * a > n:
        a -= n
    return a


def powmods(a: int, r: int, n: int) -> int:
    """return "signed" a**r (mod n), in range -n//2..n//2"""
    out = 1
    while r > 0:
        if (r % 2) == 1:
            r -= 1
            out = mods(out * a, n)
        r //= 2
        a = mods(a * a, n)
    return out


def quos(a: int, n: int) -> int:
    """returns equivalent of "a//n" for signed mod"""
    assert n > 0
    return (a - mods(a, n)) // n


def grem(w: Tuple[int, int], z: Tuple[int, int]) -> Tuple[int, int]:
    """returns remainder in Gaussian integers when dividing w by z"""
    (w0, w1) = w
    (z0, z1) = z
    n = z0 * z0 + z1 * z1
    assert n > 0  # and "division by zero"
    u0 = quos(w0 * z0 + w1 * z1, n)
    u1 = quos(w1 * z0 - w0 * z1, n)
    return (w0 - z0 * u0 + z1 * u1, w1 - z0 * u1 - z1 * u0)


def ggcd(w: Tuple[int, int], z: Tuple[int, int]) -> Tuple[int, int]:
    """
    returns greatest common divisor for gaussian integers

    Example:
        Demonstrates how ggcd() can be used to determine sum of squares.
    ```
        >>> powmods(13,2,17)
        -1
        >>> ggcd((17,0),(13,1))
        (4, -1)
        >>> 4**2+(-1)**2
        17
        >>>
    ```
    """
    while z != (0, 0):
        w, z = z, grem(w, z)
    return w


def root4m1(p: int) -> int:
    """returns sqrt(-1) (mod p)"""
    assert p > 1 and (p % 4) == 1
    k = p // 4
    j = 2
    while True:
        a = powmods(j, k, p)
        b = mods(a * a, p)
        if b == -1:
            return a
        assert b == 1  # and "p not prime"
        j += 1


def sq2(p: int) -> Tuple[int, int]:
    """
    Args:
        p: asserts if p is no prime =1 (mod 4).
    Returns:
        _: pair of numbers, their squares summing up to p.
    Example:
        Determine unique sum of two squares for prime 233.
    ```
        >>> sq2(233)
        (13, 8)
        >>>
    ```
    """
    assert p > 1 and p % 4 == 1

    a = root4m1(p)
    x, y = ggcd((p, 0), (a, 1))
    return abs(x), abs(y)


if argv[-1] == "gmpy2":
    from gmpy2 import mpz                  # pylint: disable=no-name-in-module
else:                 mpz = lambda x: x

start = time()
t = sq2(mpz(2 ** 2 ** 13 + 897))
print(str(time() - start) + "s", file=stderr)
print(t[0], t[1])

pi@pi400-64:~ $ python sq2_cg.py gmpy2 > /dev/null
2.4962875843048096s
pi@pi400-64:~ $ pylint sq2_cg.py 

--------------------------------------------------------------------
Your code has been rated at 10.00/10 (previous run: 10.00/10, +0.00)

pi@pi400-64:~ $ 

GraphvizFiddle Share link