|
|
@@ -0,0 +1,301 @@
|
|
|
+#| # Evaluation Representation of Polynomials and FFT optimizations
|
|
|
+#| In addition to the coefficient-based representation of polynomials used
|
|
|
+#| in babysnark.py, for performance we will also use an alternative
|
|
|
+#| representation where the polynomial is evaluated at a fixed set of points.
|
|
|
+#| Some operations, like multiplication and division, are significantly more
|
|
|
+#| efficient in this form.
|
|
|
+#| We can use FFT-based tools for efficiently converting
|
|
|
+#| between coefficient and evaluation representation.
|
|
|
+#|
|
|
|
+#| This library provides:
|
|
|
+#| - Fast fourier transform for finite fields
|
|
|
+#| - Interpolation and evaluation using FFT
|
|
|
+
|
|
|
+from finite_fields.finitefield import FiniteField
|
|
|
+from finite_fields.polynomial import polynomialsOver
|
|
|
+from finite_fields.euclidean import extendedEuclideanAlgorithm
|
|
|
+import random
|
|
|
+from finite_fields.numbertype import typecheck, memoize, DomainElement
|
|
|
+from functools import reduce
|
|
|
+import numpy as np
|
|
|
+
|
|
|
+#| ## Fast Fourier Transform on Finite Fields
|
|
|
+def fft_helper(a, omega, field):
|
|
|
+ """
|
|
|
+ Given coefficients A of polynomial this method does FFT and returns
|
|
|
+ the evaluation of the polynomial at [omega^0, omega^(n-1)]
|
|
|
+
|
|
|
+ If the polynomial is a0*x^0 + a1*x^1 + ... + an*x^n then the coefficients
|
|
|
+ list is of the form [a0, a1, ... , an].
|
|
|
+ """
|
|
|
+ n = len(a)
|
|
|
+ assert not (n & (n - 1)), "n must be a power of 2"
|
|
|
+
|
|
|
+ if n == 1:
|
|
|
+ return a
|
|
|
+
|
|
|
+ b, c = a[0::2], a[1::2]
|
|
|
+ b_bar = fft_helper(b, pow(omega, 2), field)
|
|
|
+ c_bar = fft_helper(c, pow(omega, 2), field)
|
|
|
+ a_bar = [field(1)] * (n)
|
|
|
+ for j in range(n):
|
|
|
+ k = j % (n // 2)
|
|
|
+ a_bar[j] = b_bar[k] + pow(omega, j) * c_bar[k]
|
|
|
+ return a_bar
|
|
|
+
|
|
|
+
|
|
|
+#| ## Representing a polynomial by evaluation at fixed points
|
|
|
+@memoize
|
|
|
+def make_polynomial_evalrep(field, omega, n):
|
|
|
+ assert n & n - 1 == 0, "n must be a power of 2"
|
|
|
+
|
|
|
+ # Check that omega is an n'th primitive root of unity
|
|
|
+ assert type(omega) is field
|
|
|
+ omega = field(omega)
|
|
|
+ assert omega**(n) == 1
|
|
|
+ _powers = [omega**i for i in range(n)]
|
|
|
+ assert len(set(_powers)) == n
|
|
|
+
|
|
|
+ _poly_coeff = polynomialsOver(field)
|
|
|
+
|
|
|
+ class PolynomialEvalRep(object):
|
|
|
+
|
|
|
+ def __init__(self, xs, ys):
|
|
|
+ # Each element of xs must be a power of omega.
|
|
|
+ # There must be a corresponding y for every x.
|
|
|
+ if type(xs) is not tuple:
|
|
|
+ xs = tuple(xs)
|
|
|
+ if type(ys) is not tuple:
|
|
|
+ ys = tuple(ys)
|
|
|
+
|
|
|
+ assert len(xs) <= n+1
|
|
|
+ assert len(xs) == len(ys)
|
|
|
+ for x in xs:
|
|
|
+ assert x in _powers
|
|
|
+ for y in ys:
|
|
|
+ assert type(y) is field
|
|
|
+
|
|
|
+ self.evalmap = dict(zip(xs, ys))
|
|
|
+
|
|
|
+ @classmethod
|
|
|
+ def from_coeffs(cls, poly):
|
|
|
+ assert type(poly) is _poly_coeff
|
|
|
+ assert poly.degree() <= n
|
|
|
+ padded_coeffs = poly.coefficients + [field(0)] * (n - len(poly.coefficients))
|
|
|
+ ys = fft_helper(padded_coeffs, omega, field)
|
|
|
+ xs = [omega**i for i in range(n) if ys[i] != 0]
|
|
|
+ ys = [y for y in ys if y != 0]
|
|
|
+ return cls(xs, ys)
|
|
|
+
|
|
|
+ def to_coeffs(self):
|
|
|
+ # To convert back to the coefficient form, we use polynomial interpolation.
|
|
|
+ # The non-zero elements stored in self.evalmap, so we fill in the zero values
|
|
|
+ # here.
|
|
|
+ ys = [self.evalmap[x] if x in self.evalmap else field(0) for x in _powers]
|
|
|
+ coeffs = [b / field(n) for b in fft_helper(ys, 1 / omega, field)]
|
|
|
+ return _poly_coeff(coeffs)
|
|
|
+
|
|
|
+ _lagrange_cache = {}
|
|
|
+ def __call__(self, x):
|
|
|
+ if type(x) is int:
|
|
|
+ x = field(x)
|
|
|
+ assert type(x) is field
|
|
|
+ xs = _powers
|
|
|
+
|
|
|
+ def lagrange(x, xi):
|
|
|
+ # Let's cache lagrange values
|
|
|
+ if (x,xi) in PolynomialEvalRep._lagrange_cache:
|
|
|
+ return PolynomialEvalRep._lagrange_cache[(x,xi)]
|
|
|
+
|
|
|
+ mul = lambda a,b: a*b
|
|
|
+ num = reduce(mul, [x - xj for xj in xs if xj != xi], field(1))
|
|
|
+ den = reduce(mul, [xi - xj for xj in xs if xj != xi], field(1))
|
|
|
+ PolynomialEvalRep._lagrange_cache[(x,xi)] = num / den
|
|
|
+ return PolynomialEvalRep._lagrange_cache[(x,xi)]
|
|
|
+
|
|
|
+ y = field(0)
|
|
|
+ for xi, yi in self.evalmap.items():
|
|
|
+ y += yi * lagrange(x, xi)
|
|
|
+ return y
|
|
|
+
|
|
|
+ def __mul__(self, other):
|
|
|
+ # Scale by integer
|
|
|
+ if type(other) is int:
|
|
|
+ other = field(other)
|
|
|
+ if type(other) is field:
|
|
|
+ return PolynomialEvalRep(self.evalmap.keys(),
|
|
|
+ [other * y for y in self.evalmap.values()])
|
|
|
+
|
|
|
+ # Multiply another polynomial in the same representation
|
|
|
+ if type(other) is type(self):
|
|
|
+ xs = []
|
|
|
+ ys = []
|
|
|
+ for x, y in self.evalmap.items():
|
|
|
+ if x in other.evalmap:
|
|
|
+ xs.append(x)
|
|
|
+ ys.append(y * other.evalmap[x])
|
|
|
+ return PolynomialEvalRep(xs, ys)
|
|
|
+
|
|
|
+ @typecheck
|
|
|
+ def __iadd__(self, other):
|
|
|
+ # Add another polynomial to this one in place.
|
|
|
+ # This is especially efficient when the other polynomial is sparse,
|
|
|
+ # since we only need to add the non-zero elements.
|
|
|
+ for x, y in other.evalmap.items():
|
|
|
+ if x not in self.evalmap:
|
|
|
+ self.evalmap[x] = y
|
|
|
+ else:
|
|
|
+ self.evalmap[x] += y
|
|
|
+ return self
|
|
|
+
|
|
|
+ @typecheck
|
|
|
+ def __add__(self, other):
|
|
|
+ res = PolynomialEvalRep(self.evalmap.keys(), self.evalmap.values())
|
|
|
+ res += other
|
|
|
+ return res
|
|
|
+
|
|
|
+ def __sub__(self, other): return self + (-other)
|
|
|
+ def __neg__(self): return PolynomialEvalRep(self.evalmap.keys(),
|
|
|
+ [-y for y in self.evalmap.values()])
|
|
|
+
|
|
|
+ def __truediv__(self, divisor):
|
|
|
+ # Scale by integer
|
|
|
+ if type(divisor) is int:
|
|
|
+ other = field(divisor)
|
|
|
+ if type(divisor) is field:
|
|
|
+ return self * (1/divisor)
|
|
|
+ if type(divisor) is type(self):
|
|
|
+ res = PolynomialEvalRep((),())
|
|
|
+ for x, y in self.evalmap.items():
|
|
|
+ assert x in divisor.evalmap
|
|
|
+ res.evalmap[x] = y / divisor.evalmap[x]
|
|
|
+ return res
|
|
|
+ return NotImplemented
|
|
|
+
|
|
|
+ def __copy__(self):
|
|
|
+ return PolynomialEvalRep(self.evalmap.keys(), self.evalmap.values())
|
|
|
+
|
|
|
+ def __repr__(self):
|
|
|
+ return f'PolyEvalRep[{hex(omega.n)[:15]}...,{n}]({len(self.evalmap)} elements)'
|
|
|
+
|
|
|
+ @classmethod
|
|
|
+ def divideWithCoset(cls, p, t, c=field(3)):
|
|
|
+ """
|
|
|
+ This assumes that p and t are polynomials in coefficient representation,
|
|
|
+ and that p is divisible by t.
|
|
|
+ This function is useful when t has roots at some or all of the powers of omega,
|
|
|
+ in which case we cannot just convert to evalrep and use division above
|
|
|
+ (since it would cause a divide by zero.
|
|
|
+ Instead, we evaluate p(X) at powers of (c*omega) for some constant cofactor c.
|
|
|
+ To do this efficiently, we create new polynomials, pc(X) = p(cX), tc(X) = t(cX),
|
|
|
+ and evaluate these at powers of omega. This conversion can be done efficiently
|
|
|
+ on the coefficient representation.
|
|
|
+ See also: cosetFFT in libsnark / libfqfft.
|
|
|
+ https://github.com/scipr-lab/libfqfft/blob/master/libfqfft/evaluation_domain/domains/extended_radix2_domain.tcc
|
|
|
+ """
|
|
|
+ assert type(p) is _poly_coeff
|
|
|
+ assert type(t) is _poly_coeff
|
|
|
+ # Compute p(cX), t(cX) by multiplying coefficients
|
|
|
+ c_acc = field(1)
|
|
|
+ pc = _poly_coeff(list(p.coefficients)) # make a copy
|
|
|
+ for i in range(p.degree() + 1):
|
|
|
+ pc.coefficients[-i-1] *= c_acc
|
|
|
+ c_acc *= c
|
|
|
+ c_acc = field(1)
|
|
|
+ tc = _poly_coeff(list(t.coefficients)) # make a copy
|
|
|
+ for i in range(t.degree() + 1):
|
|
|
+ tc.coefficients[-i-1] *= c_acc
|
|
|
+ c_acc *= c
|
|
|
+
|
|
|
+ # Divide using evalrep
|
|
|
+ pc_rep = cls.from_coeffs(pc)
|
|
|
+ tc_rep = cls.from_coeffs(tc)
|
|
|
+ hc_rep = pc_rep / tc_rep
|
|
|
+ hc = hc_rep.to_coeffs()
|
|
|
+
|
|
|
+ # Compute h(X) from h(cX) by dividing coefficients
|
|
|
+ c_acc = field(1)
|
|
|
+ h = _poly_coeff(list(hc.coefficients)) # make a copy
|
|
|
+ for i in range(hc.degree() + 1):
|
|
|
+ h.coefficients[-i-1] /= c_acc
|
|
|
+ c_acc *= c
|
|
|
+
|
|
|
+ # Correctness checks
|
|
|
+ # assert pc == tc * hc
|
|
|
+ # assert p == t * h
|
|
|
+ return h
|
|
|
+
|
|
|
+
|
|
|
+ return PolynomialEvalRep
|
|
|
+
|
|
|
+#| ## Sparse Matrix
|
|
|
+#| In our setting, we have O(m*m) elements in the matrix, and expect the number of
|
|
|
+#| elements to be O(m).
|
|
|
+#| In this setting, it's appropriate to use a rowdict representation - a dense
|
|
|
+#| array of dictionaries, one for each row, where the keys of each dictionary
|
|
|
+#| are column indices.
|
|
|
+
|
|
|
+class RowDictSparseMatrix():
|
|
|
+ # Only a few necessary methods are included here.
|
|
|
+ # This could be replaced with a generic sparse matrix class, such as scipy.sparse,
|
|
|
+ # but this does not work as well with custom value types like Fp
|
|
|
+
|
|
|
+ def __init__(self, m, n, zero=None):
|
|
|
+ self.m = m
|
|
|
+ self.n = n
|
|
|
+ self.shape = (m,n)
|
|
|
+ self.zero = zero
|
|
|
+ self.rowdicts = [dict() for _ in range(m)]
|
|
|
+
|
|
|
+ def __setitem__(self, key, v):
|
|
|
+ i, j = key
|
|
|
+ self.rowdicts[i][j] = v
|
|
|
+
|
|
|
+ def __getitem__(self, key):
|
|
|
+ i, j = key
|
|
|
+ return self.rowdicts[i][j] if j in self.rowdicts[i] else self.zero
|
|
|
+
|
|
|
+ def items(self):
|
|
|
+ for i in range(self.m):
|
|
|
+ for j, v in self.rowdicts[i].items():
|
|
|
+ yield (i,j), v
|
|
|
+
|
|
|
+ def dot(self, other):
|
|
|
+ if isinstance(other, np.ndarray):
|
|
|
+ assert other.dtype == 'O'
|
|
|
+ assert other.shape in ((self.n,),(self.n,1))
|
|
|
+ ret = np.empty((self.m,), dtype='O')
|
|
|
+ ret.fill(self.zero)
|
|
|
+ for i in range(self.m):
|
|
|
+ for j, v in self.rowdicts[i].items():
|
|
|
+ ret[i] += other[j] * v
|
|
|
+ return ret
|
|
|
+
|
|
|
+ def to_dense(self):
|
|
|
+ mat = np.empty((self.m, self.n), dtype='O')
|
|
|
+ mat.fill(self.zero)
|
|
|
+ for (i,j), val in self.items():
|
|
|
+ mat[i,j] = val
|
|
|
+ return mat
|
|
|
+
|
|
|
+ def __repr__(self): return repr(self.rowdicts)
|
|
|
+
|
|
|
+#-
|
|
|
+# Examples
|
|
|
+if __name__ == '__main__':
|
|
|
+ import misc
|
|
|
+
|
|
|
+ Fp = FiniteField(52435875175126190479447740508185965837690552500527637822603658699938581184513,1) # (# noqa: E501)
|
|
|
+ Poly = polynomialsOver(Fp)
|
|
|
+
|
|
|
+ n = 8
|
|
|
+ omega = misc.get_omega(Fp, n)
|
|
|
+ PolyEvalRep = make_polynomial_evalrep(Fp, omega, n)
|
|
|
+
|
|
|
+ f = Poly([1,2,3,4,5])
|
|
|
+ xs = tuple([omega**i for i in range(n)])
|
|
|
+ ys = tuple(map(f, xs))
|
|
|
+ # print('xs:', xs)
|
|
|
+ # print('ys:', ys)
|
|
|
+
|
|
|
+ assert f == PolyEvalRep(xs, ys).to_coeffs()
|