Make search for polynomials of fixed-term much more efficient

src/galois/_polys/_irreducible.py

@@ -5,11 +5,11 @@
 
 import functools
 import random
-from typing import TYPE_CHECKING, Iterator, Type
+from typing import TYPE_CHECKING, Callable, Iterable, Iterator, Type
 
 from typing_extensions import Literal
 
-from .._domains import _factory
+from .._domains import Array, _factory
 from .._helper import export, verify_isinstance
 from .._prime import is_prime_power
 from ._poly import Poly
@@ -108,8 +108,9 @@ def irreducible_poly(
         else:
             poly = _random_search(order, degree, terms)
     except StopIteration as e:
+        terms_str = "any" if terms is None else str(terms)
         raise RuntimeError(
-            f"No monic irreducible polynomial of degree {degree} over GF({order}) with {terms} terms exists."
+            f"No monic irreducible polynomial of degree {degree} over GF({order}) with {terms_str} terms exists."
         ) from e
 
     return poly
@@ -194,23 +195,25 @@ def irreducible_polys(
     if terms is not None and not 1 <= terms <= degree + 1:
         raise ValueError(f"Argument 'terms' must be at least 1 and at most {degree + 1}, not {terms}.")
 
-    field = _factory.FIELD_FACTORY(order)
-
-    # Only search monic polynomials of degree m over GF(q)
-    start = order**degree
-    stop = 2 * order**degree
-    step = 1
-
-    if reverse:
-        start, stop, step = stop - 1, start - 1, -1
-
-    while True:
-        poly = _deterministic_search(field, start, stop, step, terms)
-        if poly is not None:
-            start = int(poly) + step
-            yield poly
-        else:
-            break
+    if terms is None:
+        # Iterate over and test all monic polynomials of degree m over GF(q).
+        start = order**degree
+        stop = 2 * order**degree
+        step = 1
+        if reverse:
+            start, stop, step = stop - 1, start - 1, -1
+        field = _factory.FIELD_FACTORY(order)
+
+        while True:
+            poly = _deterministic_search(field, start, stop, step)
+            if poly is not None:
+                start = int(poly) + step
+                yield poly
+            else:
+                break
+    else:
+        # Iterate over and test monic polynomials of degree m over GF(q) with `terms` non-zero terms.
+        yield from _deterministic_search_fixed_terms(order, degree, terms, "is_irreducible", reverse)
 
 
 @functools.lru_cache(maxsize=4096)
@@ -219,21 +222,81 @@ def _deterministic_search(
     start: int,
     stop: int,
     step: int,
-    terms: int | None,
 ) -> Poly | None:
     """
     Searches for an irreducible polynomial in the range using the specified deterministic method.
     """
     for element in range(start, stop, step):
         poly = Poly.Int(element, field=field)
-        if terms is not None and poly.nonzero_coeffs.size != terms:
-            continue
         if poly.is_irreducible():
             return poly
 
     return None
 
 
+def _deterministic_search_fixed_terms(
+    order: int,
+    degree: int,
+    terms: int,
+    test: str,
+    reverse: bool = False,
+) -> Iterator[Poly]:
+    """
+    Iterates over all polynomials of the given degree and number of non-zero terms in lexicographical
+    order, only yielding those that pass the specified test (either 'is_irreducible()' or 'is_primitive()').
+    """
+    assert test in ["is_irreducible", "is_primitive"]
+    field = _factory.FIELD_FACTORY(order)
+
+    # A wrapper function around range to iterate forwards or backwards.
+    def direction(x):
+        if reverse:
+            return reversed(x)
+        return x
+
+    # Initialize the search by setting the first term to x^m with coefficient 1. This function will
+    # recursively add the remaining terms, with the last term being x^0.
+    yield from _deterministic_search_fixed_terms_recursive([degree], [1], terms - 1, test, field, direction)
+
+
+def _deterministic_search_fixed_terms_recursive(
+    degrees: Iterable[int],
+    coeffs: Iterable[int],
+    terms: int,
+    test: str,
+    field: Type[Array],
+    direction: Callable[[Iterable[int]], Iterable[int]],
+) -> Iterator[Poly]:
+    """
+    Recursively finds all polynomials having non-zero coefficients `coeffs` with degree `degrees` with `terms`
+    additional non-zero terms. The polynomials are found in lexicographical order, only yielding those that pass
+    the specified test (either 'is_irreducible()' or 'is_primitive()').
+    """
+    if terms == 0:
+        # There are no more terms, yield the polynomial.
+        poly = Poly.Degrees(degrees, coeffs, field=field)
+        if getattr(poly, test)():
+            yield poly
+    elif terms == 1:
+        # The last term must be the x^0 term, so we don't need to loop over possible degrees.
+        for coeff in direction(range(1, field.order)):
+            next_degrees = (*degrees, 0)
+            next_coeffs = (*coeffs, coeff)
+            yield from _deterministic_search_fixed_terms_recursive(
+                next_degrees, next_coeffs, terms - 1, test, field, direction
+            )
+    else:
+        # Find the next term's degree. It must be at least terms - 1 so that the polynomial can have the specified
+        # number of terms of lesser degree. It must also be less than the degree of the previous term.
+        for degree in direction(range(terms - 1, degrees[-1])):
+            for coeff in direction(range(1, field.order)):
+                next_degrees = (*degrees, degree)
+                next_coeffs = (*coeffs, coeff)
+                yield from _deterministic_search_fixed_terms_recursive(
+                    next_degrees, next_coeffs, terms - 1, test, field, direction
+                )
+
+
 def _random_search(order: int, degree: int, terms: int | None) -> Poly:
     """
     Searches for a random irreducible polynomial.

src/galois/_polys/_primitive.py

@@ -13,6 +13,7 @@
 from .._domains import _factory
 from .._helper import export, verify_isinstance
 from .._prime import is_prime, is_prime_power
+from ._irreducible import _deterministic_search_fixed_terms
 from ._poly import Poly
 
 if TYPE_CHECKING:
@@ -127,8 +128,9 @@ def primitive_poly(
         else:
             poly = _random_search(order, degree, terms)
     except StopIteration as e:
+        terms_str = "any" if terms is None else str(terms)
         raise RuntimeError(
-            f"No monic primitive polynomial of degree {degree} over GF({order}) with {terms} terms exists."
+            f"No monic primitive polynomial of degree {degree} over GF({order}) with {terms_str} terms exists."
         ) from e
 
     return poly
@@ -215,23 +217,25 @@ def primitive_polys(
     if terms is not None and not 1 <= terms <= degree + 1:
         raise ValueError(f"Argument 'terms' must be at least 1 and at most {degree + 1}, not {terms}.")
 
-    field = _factory.FIELD_FACTORY(order)
-
-    # Only search monic polynomials of degree m over GF(q)
-    start = order**degree
-    stop = 2 * order**degree
-    step = 1
-
-    if reverse:
-        start, stop, step = stop - 1, start - 1, -1
-
-    while True:
-        poly = _deterministic_search(field, start, stop, step, terms)
-        if poly is not None:
-            start = int(poly) + step
-            yield poly
-        else:
-            break
+    if terms is None:
+        # Iterate over and test all monic polynomials of degree m over GF(q).
+        start = order**degree
+        stop = 2 * order**degree
+        step = 1
+        if reverse:
+            start, stop, step = stop - 1, start - 1, -1
+        field = _factory.FIELD_FACTORY(order)
+
+        while True:
+            poly = _deterministic_search(field, start, stop, step)
+            if poly is not None:
+                start = int(poly) + step
+                yield poly
+            else:
+                break
+    else:
+        # Iterate over and test monic polynomials of degree m over GF(q) with `terms` non-zero terms.
+        yield from _deterministic_search_fixed_terms(order, degree, terms, "is_primitive", reverse)
 
 
 @functools.lru_cache(maxsize=4096)
@@ -240,15 +244,12 @@ def _deterministic_search(
     start: int,
     stop: int,
     step: int,
-    terms: int | None,
 ) -> Poly | None:
     """
     Searches for a primitive polynomial in the range using the specified deterministic method.
     """
     for element in range(start, stop, step):
         poly = Poly.Int(element, field=field)
-        if terms is not None and poly.nonzero_coeffs.size != terms:
-            continue
         if poly.is_primitive():
             return poly
 

tests/polys/test_irreducible_polys.py

@@ -127,15 +127,15 @@ def test_irreducible_polys(order, degree, polys):
     assert [f.coeffs.tolist() for f in galois.irreducible_polys(order, degree)] == polys
 
 
-def test_specific_terms():
-    degree = 8
+@pytest.mark.parametrize("order,degree,polys", PARAMS)
+def test_specific_terms(order, degree, polys):
     all_polys = []
     for terms in range(1, degree + 2):
-        polys = list(galois.irreducible_polys(2, degree, terms=terms))
-        assert all(p.nonzero_coeffs.size == terms for p in polys)
-        all_polys += polys
+        new_polys = list(galois.irreducible_polys(order, degree, terms=terms))
+        assert all(p.nonzero_coeffs.size == terms for p in new_polys)
+        all_polys += new_polys
     all_polys = [p.coeffs.tolist() for p in sorted(all_polys, key=int)]
-    assert all_polys == IRREDUCIBLE_POLYS_2_8
+    assert all_polys == polys
 
 
 def test_specific_terms_none_found():

tests/polys/test_primitive_polys.py

@@ -131,15 +131,15 @@ def test_primitive_polys(order, degree, polys):
     assert [f.coeffs.tolist() for f in galois.primitive_polys(order, degree)] == polys
 
 
-def test_specific_terms():
-    degree = 8
+@pytest.mark.parametrize("order,degree,polys", PARAMS)
+def test_specific_terms(order, degree, polys):
     all_polys = []
     for terms in range(1, degree + 2):
-        polys = list(galois.primitive_polys(2, degree, terms=terms))
-        assert all(p.nonzero_coeffs.size == terms for p in polys)
-        all_polys += polys
+        new_polys = list(galois.primitive_polys(order, degree, terms=terms))
+        assert all(p.nonzero_coeffs.size == terms for p in new_polys)
+        all_polys += new_polys
     all_polys = [p.coeffs.tolist() for p in sorted(all_polys, key=int)]
-    assert all_polys == PRIMITIVE_POLYS_2_8
+    assert all_polys == polys
 
 
 def test_specific_terms_none_found():