Fermat’s two-squares theorem

Every prime \(p \equiv 1 \pmod 4\) is a sum of two squares \(a^2 + b^2\), and no prime \(p \equiv 3 \pmod 4\) is. This script checks the theorem for every prime below 10,000 and plots each representation as a lattice point \((a, b)\) on the circle of radius \(\sqrt p\).

import matplotlib.pyplot as plt
import numpy as np

from mathematicskit.number_theory import sieve_of_eratosthenes, sum_of_two_squares

Check the theorem prime by prime

primes = [int(p) for p in sieve_of_eratosthenes(10000)[1:]]
reps = {p: sum_of_two_squares(p) for p in primes}
ones = [p for p in primes if p % 4 == 1]
threes = [p for p in primes if p % 4 == 3]
print(f"primes = 1 (mod 4): {len(ones)}, all sums of two squares: {all(reps[p] is not None for p in ones)}")
print(f"primes = 3 (mod 4): {len(threes)}, none a sum of two squares: {all(reps[p] is None for p in threes)}")
for p in ones[:8]:
    a, b = reps[p]
    print(f"  {p} = {a}^2 + {b}^2")

primes = 1 (mod 4): 609, all sums of two squares: True
primes = 3 (mod 4): 619, none a sum of two squares: True
  5 = 1^2 + 2^2
  13 = 2^2 + 3^2
  17 = 1^2 + 4^2
  29 = 2^2 + 5^2
  37 = 1^2 + 6^2
  41 = 4^2 + 5^2
  53 = 2^2 + 7^2
  61 = 5^2 + 6^2

The representations as lattice points

Each point sits on the circle \(a^2 + b^2 = p\); reflecting in the diagonal shows both orderings of the pair.

ab = np.array([reps[p] for p in ones])
fig, ax = plt.subplots(figsize=(6, 6))
ax.plot(ab[:, 0], ab[:, 1], ".", ms=3, color="tab:blue", label="(a, b), a <= b")
ax.plot(ab[:, 1], ab[:, 0], ".", ms=3, color="tab:orange", label="(b, a)")
ax.set_xlabel("a")
ax.set_ylabel("b")
ax.set_aspect("equal")
ax.set_title("Primes p = a^2 + b^2 below 10,000")
ax.legend()
plt.show()