[QGFPoly] Add QGFPoly, a new data type to represent polynomials over galois fields (#1620)

* [QGFPoly] Add a new data type to represent polynomials over galois fields

* Address nits and create function to convert to/from Poly class to array of coefficients

* Fix formatting

Author: tanujkhattar

qualtran/__init__.py

@@ -61,6 +61,7 @@
     BQUInt,
     QMontgomeryUInt,
     QGF,
+    QGFPoly,
 )
 
 # Internal imports: none

qualtran/_infra/data_types.py

@@ -1025,6 +1025,108 @@ def __str__(self):
         return f'QGF({self.characteristic}**{self.degree})'
 
 
+@attrs.frozen
+class QGFPoly(QDType):
+    r"""Univariate Polynomials with coefficients in a Galois Field GF($p^m$).
+
+    This data type represents a degree-$n$ univariate polynomials
+    $f(x)=\sum_{i=0}^{n} a_i x^{i}$ where the coefficients $a_{i}$ of the polynomial
+    belong to a Galois Field $GF(p^{m})$.
+
+    The data type uses the [Galois library](https://mhostetter.github.io/galois/latest/) to
+    perform arithmetic over polynomials defined over Galois Fields using the
+    [galois.Poly](https://mhostetter.github.io/galois/latest/api/galois.Poly/).
+
+    Attributes:
+        degree: The degree $n$ of the univariate polynomial $f(x)$ represented by this type.
+        qgf: An instance of `QGF` that represents the galois field $GF(p^m)$ over which the
+            univariate polynomial $f(x)$ is defined.
+
+    References
+        [Polynomials over finite fields](https://mhostetter.github.io/galois/latest/api/galois.Poly/)
+
+        [Polynomial Arithmetic](https://mhostetter.github.io/galois/latest/basic-usage/poly-arithmetic/)
+    """
+
+    degree: SymbolicInt
+    qgf: QGF
+
+    @cached_property
+    def bitsize(self) -> SymbolicInt:
+        """Bitsize of qubit register required to represent a single instance of this data type."""
+        return self.qgf.bitsize * (self.degree + 1)
+
+    @cached_property
+    def num_qubits(self) -> SymbolicInt:
+        """Number of qubits required to represent a single instance of this data type."""
+        return self.bitsize
+
+    def get_classical_domain(self) -> Iterable[Any]:
+        """Yields all possible classical (computational basis state) values representable
+        by this type."""
+        import itertools
+
+        from galois import Poly
+
+        for it in itertools.product(self.qgf.gf_type.elements, repeat=(self.degree + 1)):
+            yield Poly(self.qgf.gf_type(it), field=self.qgf.gf_type)
+
+    @cached_property
+    def _quint_equivalent(self) -> QUInt:
+        return QUInt(self.num_qubits)
+
+    def to_gf_coefficients(self, f_x: galois.Poly) -> galois.Array:
+        """Returns a big-endian array of coefficients of the polynomial f(x)."""
+        f_x_coeffs = self.qgf.gf_type.Zeros(self.degree + 1)
+        f_x_coeffs[self.degree - f_x.degree :] = f_x.coeffs
+        return f_x_coeffs
+
+    def from_gf_coefficients(self, f_x: galois.Array) -> galois.Poly:
+        """Expects a big-endian array of coefficients that represent a polynomial f(x)."""
+        return galois.Poly(f_x, field=self.qgf.gf_type)
+
+    def to_bits(self, x) -> List[int]:
+        """Returns individual bits corresponding to binary representation of x"""
+        self.assert_valid_classical_val(x)
+        assert isinstance(x, galois.Poly)
+        return self.qgf.to_bits_array(self.to_gf_coefficients(x)).reshape(-1).tolist()
+
+    def from_bits(self, bits: Sequence[int]):
+        """Combine individual bits to form x"""
+        reshaped_bits = np.array(bits).reshape((int(self.degree) + 1, int(self.qgf.bitsize)))
+        return self.from_gf_coefficients(self.qgf.from_bits_array(reshaped_bits))
+
+    def assert_valid_classical_val(self, val: Any, debug_str: str = 'val'):
+        """Raises an exception if `val` is not a valid classical value for this type.
+
+        Args:
+            val: A classical value that should be in the domain of this QDType.
+            debug_str: Optional debugging information to use in exception messages.
+        """
+        if not isinstance(val, galois.Poly):
+            raise ValueError(f"{debug_str} should be a {galois.Poly}, not {val!r}")
+        if val.field is not self.qgf.gf_type:
+            raise ValueError(
+                f"{debug_str} should be defined over {self.qgf.gf_type}, not {val.field}"
+            )
+        if val.degree > self.degree:
+            raise ValueError(f"{debug_str} should have a degree <= {self.degree}, not {val.degree}")
+
+    def is_symbolic(self) -> bool:
+        """Returns True if this qdtype is parameterized with symbolic objects."""
+        return is_symbolic(self.degree, self.qgf)
+
+    def iteration_length_or_zero(self) -> SymbolicInt:
+        """Safe version of iteration length.
+
+        Returns the iteration_length if the type has it or else zero.
+        """
+        return self.qgf.order
+
+    def __str__(self):
+        return f'QGFPoly({self.degree}, {self.qgf!s})'
+
+
 QAnyInt = (QInt, QUInt, BQUInt, QMontgomeryUInt)
 QAnyUInt = (QUInt, BQUInt, QMontgomeryUInt, QGF)
 

qualtran/_infra/data_types_test.py

@@ -15,6 +15,7 @@
 import random
 from typing import Any, Sequence, Union
 
+import galois
 import numpy as np
 import pytest
 import sympy
@@ -31,6 +32,7 @@
     QDType,
     QFxp,
     QGF,
+    QGFPoly,
     QInt,
     QIntOnesComp,
     QMontgomeryUInt,
@@ -167,13 +169,33 @@ def test_qgf():
     assert is_symbolic(qgf_pm)
 
 
+def test_qgf_poly():
+    qgf_poly_4_8 = QGFPoly(4, QGF(characteristic=2, degree=3))
+    assert str(qgf_poly_4_8) == 'QGFPoly(4, QGF(2**3))'
+    assert qgf_poly_4_8.num_qubits == 5 * 3
+    n, p, m = sympy.symbols('n, p, m', integer=True, positive=True)
+    qgf_poly_n_pm = QGFPoly(n, QGF(characteristic=p, degree=m))
+    assert qgf_poly_n_pm.num_qubits == (n + 1) * ceil(log2(p**m))
+    assert is_symbolic(qgf_poly_n_pm)
+
+
 @pytest.mark.parametrize('qdtype', [QBit(), QInt(4), QUInt(4), BQUInt(3, 5)])
 def test_domain_and_validation(qdtype: QDType):
     for v in qdtype.get_classical_domain():
         qdtype.assert_valid_classical_val(v)
 
 
-@pytest.mark.parametrize('qdtype', [QBit(), QInt(4), QUInt(4), BQUInt(3, 5), QGF(2, 8)])
+@pytest.mark.parametrize(
+    'qdtype',
+    [
+        QBit(),
+        QInt(4),
+        QUInt(4),
+        BQUInt(3, 5),
+        QGF(2, 8),
+        QGFPoly(4, QGF(characteristic=2, degree=2)),
+    ],
+)
 def test_domain_and_validation_arr(qdtype: QDType):
     arr = np.array(list(qdtype.get_classical_domain()))
     qdtype.assert_valid_classical_val_array(arr)
@@ -207,6 +229,11 @@ def test_validation_errs():
     with pytest.raises(ValueError):
         QGF(2, 8).assert_valid_classical_val(2**8)
 
+    with pytest.raises(ValueError):
+        qgf = QGF(2, 3)
+        poly = galois.Poly(qgf.gf_type([1, 2, 3, 4, 5, 6, 7]), field=qgf.gf_type)
+        QGFPoly(4, qgf).assert_valid_classical_val(poly)
+
 
 def test_validate_arrays():
     rs = np.random.RandomState(52)
@@ -319,6 +346,12 @@ def test_qgf_to_and_from_bits():
     assert_to_and_from_bits_array_consistent(qgf_256, gf256([*range(256)]))
 
 
+def test_qgf_poly_to_and_from_bits():
+    qgf_4 = QGF(2, 2)
+    qgf_poly_2_4 = QGFPoly(2, qgf_4)
+    assert_to_and_from_bits_array_consistent(qgf_poly_2_4, [*qgf_poly_2_4.get_classical_domain()])
+
+
 def test_quint_to_and_from_bits():
     quint4 = QUInt(4)
     assert [*quint4.get_classical_domain()] == [*range(0, 16)]