mirror of
https://github.com/darkrenaissance/darkfi.git
synced 2026-01-09 14:48:08 -05:00
302 lines
11 KiB
Python
302 lines
11 KiB
Python
#| # 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()
|