p3_interpolation/
lib.rs

Help
1//! Tools for Lagrange interpolation.
2
3#![no_std]
4
5extern crate alloc;
6
7use alloc::vec::Vec;
8
9use p3_field::{
10    batch_multiplicative_inverse, scale_vec, two_adic_coset_zerofier, ExtensionField, TwoAdicField,
11};
12use p3_matrix::Matrix;
13use p3_maybe_rayon::prelude::*;
14use p3_util::log2_strict_usize;
15
16/// Given evaluations of a batch of polynomials over the canonical power-of-two subgroup, evaluate
17/// the polynomials at `point`.
18///
19/// This assumes the point is not in the subgroup, otherwise the behavior is undefined.
20pub fn interpolate_subgroup<F, EF, Mat>(subgroup_evals: &Mat, point: EF) -> Vec<EF>
21where
22    F: TwoAdicField,
23    EF: ExtensionField<F> + TwoAdicField,
24    Mat: Matrix<F>,
25{
26    interpolate_coset(subgroup_evals, F::ONE, point, None)
27}
28
29/// Given evaluations of a batch of polynomials over the given coset of the canonical power-of-two
30/// subgroup, evaluate the polynomials at `point`.
31///
32/// This assumes the point is not in the coset, otherwise the behavior is undefined.
33/// If available, reuse denominator diffs that is `1 / (x_i-z)` to avoid batch inversion.
34pub fn interpolate_coset<F, EF, Mat>(
35    coset_evals: &Mat,
36    shift: F,
37    point: EF,
38    diff_invs: Option<&[EF]>,
39) -> Vec<EF>
40where
41    F: TwoAdicField,
42    EF: ExtensionField<F> + TwoAdicField,
43    Mat: Matrix<F>,
44{
45    // Slight variation of this approach: https://hackmd.io/@vbuterin/barycentric_evaluation
46
47    let height = coset_evals.height();
48    let log_height = log2_strict_usize(height);
49
50    let g = F::two_adic_generator(log_height).powers().take(height);
51    let col_scale: Vec<_> = if let Some(diff_invs) = diff_invs {
52        g.zip(diff_invs)
53            .map(|(sg, &diff_inv)| diff_inv * sg)
54            .collect()
55    } else {
56        let subgroup = g.collect::<Vec<_>>();
57        let diffs: Vec<EF> = subgroup
58            .par_iter()
59            .map(|&subgroup_i| point - subgroup_i * shift)
60            .collect();
61        let diff_invs = batch_multiplicative_inverse(&diffs);
62        subgroup
63            .par_iter()
64            .zip(diff_invs)
65            .map(|(&sg, diff_inv)| diff_inv * sg)
66            .collect()
67    };
68    let sum = coset_evals.columnwise_dot_product(&col_scale);
69
70    let zerofier = two_adic_coset_zerofier::<EF>(log_height, EF::from_base(shift), point);
71    let denominator = F::from_canonical_usize(height) * shift.exp_u64(height as u64 - 1);
72    scale_vec(zerofier * denominator.inverse(), sum)
73}
74
75#[cfg(test)]
76mod tests {
77    use alloc::vec;
78    use alloc::vec::Vec;
79
80    use p3_baby_bear::BabyBear;
81    use p3_field::{batch_multiplicative_inverse, Field, FieldAlgebra};
82    use p3_matrix::dense::RowMajorMatrix;
83    use p3_util::log2_strict_usize;
84
85    use crate::{interpolate_coset, interpolate_subgroup};
86
87    #[test]
88    fn test_interpolate_subgroup() {
89        // x^2 + 2 x + 3
90        type F = BabyBear;
91        let evals = [
92            6, 886605102, 1443543107, 708307799, 2, 556938009, 569722818, 1874680944,
93        ]
94        .map(F::from_canonical_u32);
95        let evals_mat = RowMajorMatrix::new(evals.to_vec(), 1);
96        let point = F::from_canonical_u32(100);
97        let result = interpolate_subgroup(&evals_mat, point);
98        assert_eq!(result, vec![F::from_canonical_u32(10203)]);
99    }
100
101    #[test]
102    fn test_interpolate_coset() {
103        // x^2 + 2 x + 3
104        type F = BabyBear;
105        let shift = F::GENERATOR;
106        let evals = [
107            1026, 129027310, 457985035, 994890337, 902, 1988942953, 1555278970, 913671254,
108        ]
109        .map(F::from_canonical_u32);
110        let evals_mat = RowMajorMatrix::new(evals.to_vec(), 1);
111        let point = F::from_canonical_u32(100);
112        let result = interpolate_coset(&evals_mat, shift, point, None);
113        assert_eq!(result, vec![F::from_canonical_u32(10203)]);
114
115        use p3_field::TwoAdicField;
116        let n = evals.len();
117        let k = log2_strict_usize(n);
118
119        let denom: Vec<_> = F::two_adic_generator(k)
120            .shifted_powers(shift)
121            .take(n)
122            .map(|w| point - w)
123            .collect();
124
125        let denom = batch_multiplicative_inverse(&denom);
126        let result = interpolate_coset(&evals_mat, shift, point, Some(&denom));
127        assert_eq!(result, vec![F::from_canonical_u32(10203)]);
128    }
129}
p3_interpolation/lib.rs

p3_interpolation/
lib.rs
