Primes of the form x^2 + ny^2, David Cox, Exercise 2.9. Determine quadratic forms for a given discriminant D < 0

forms.py

from math import sqrt, gcd
import unittest
from typing import Tuple, Set


def discriminant(a, b, c):
    return b*b - 4 * a * c


def get_reduced_forms(D) -> Set[Tuple[int, int, int]]:
    """
    Return the primitive reduced forms for the given discriminant.

    Ranges over a/b using the following bounds:

    (1) a <= sqrt(-D/3).  This is because -D = 4ac - b^2 >= 4a^2 - a^2 = 3a^2
    (2) |b| < a.  From the definition of reduced form.

    For each a/b pair, c must be: -(D - b^2)/4 * a

    We then filter out the following situations:

    * If |b| = a or a = c, b must be positive
    * Forms where gcd(a, b, c) != 1
    """

    assert D < 0

    a_bound = int(sqrt(-D / 3))
    forms = set()

    # Since D < 0, can check a > 0
    for a in range(1, a_bound + 1):
        for b in range(-a_bound, a_bound + 1):
            c = int(-(D - (b * b)) / (4 * a))

            # |b| < a
            if abs(b) > a:
                continue

            # To be reduced, |b| = a or a = c forces b >= 0
            if b < 0 and (abs(b) == a or a == c):
                continue

            # Only want primitive forms
            if gcd(gcd(a, b), c) != 1:
                continue

            if discriminant(a, b, c) == D:
                forms.add((a, b, c))

    return forms


class QuadraticForms(unittest.TestCase):
    def test_verify_table(self):
        """
        Verify entries from table 2.14
        """

        self.assertEqual(get_reduced_forms(D=-4), {(1, 0, 1)})
        self.assertEqual(get_reduced_forms(D=-8), {(1, 0, 2)})
        self.assertEqual(get_reduced_forms(D=-12), {(1, 0, 3)})
        self.assertEqual(get_reduced_forms(D=-20), {
            (1, 0, 5),
            (2, 2, 3),
        })
        self.assertEqual(get_reduced_forms(-28), {(1, 0, 7)})
        self.assertEqual(get_reduced_forms(-56), {
            (1, 0, 14),
            (2, 0, 7),
            (3, 2, 5),
            (3, -2, 5),
        })
        self.assertEqual(get_reduced_forms(-108), {
            (1, 0, 27),
            (4, 2, 7,),
            (4, -2, 7)}
        )

    def test_compute_reduced_forms(self):
        self.assertEqual(get_reduced_forms(-3), {(1, 1, 1)})
        self.assertEqual(get_reduced_forms(-15), {(1, 1, 4), (2, 1, 2)})
        self.assertEqual(get_reduced_forms(-24), {(1, 0, 6), (2, 0, 3)})
        self.assertEqual(get_reduced_forms(-31), {
            (1, 1, 8),
            (2, 1, 4),
            (2, -1, 4)
        })
        self.assertEqual(get_reduced_forms(-52), {(1, 0, 13), (2, 2, 7)})


if __name__ == "__main__":
    unittest.main()