1/// Some basic MLE utilities
  2use ark_ff::PrimeField;
  3use ark_poly::{DenseMultilinearExtension, SparseMultilinearExtension};
  4use ark_std::log2;
  5
  6use super::vec::SparseMatrix;
  7
  8/// Pad matrix so that its columns and rows are powers of two
  9pub fn pad_matrix<F: PrimeField>(m: &SparseMatrix<F>) -> SparseMatrix<F> {
 10    let mut r = m.clone();
 11    r.n_rows = m.n_rows.next_power_of_two();
 12    r.n_cols = m.n_cols.next_power_of_two();
 13    r
 14}
 15
 16/// Returns the dense multilinear extension from the given matrix, without modifying the original
 17/// matrix.
 18pub fn matrix_to_dense_mle<F: PrimeField>(matrix: SparseMatrix<F>) -> DenseMultilinearExtension<F> {
 19    let n_vars: usize = (log2(matrix.n_rows) + log2(matrix.n_cols)) as usize; // n_vars = s + s'
 20
 21    // Matrices might need to get padded before turned into an MLE
 22    let padded_matrix = pad_matrix(&matrix);
 23
 24    // build dense vector representing the sparse padded matrix
 25    let mut v: Vec<F> = vec![F::zero(); padded_matrix.n_rows * padded_matrix.n_cols];
 26    for (row_i, row) in padded_matrix.coeffs.iter().enumerate() {
 27        for &(value, col_i) in row.iter() {
 28            v[(padded_matrix.n_cols * row_i) + col_i] = value;
 29        }
 30    }
 31
 32    // convert the dense vector into a mle
 33    vec_to_dense_mle(n_vars, &v)
 34}
 35
 36/// Takes the n_vars and a dense vector and returns its dense MLE.
 37pub fn vec_to_dense_mle<F: PrimeField>(n_vars: usize, v: &[F]) -> DenseMultilinearExtension<F> {
 38    let v_padded: Vec<F> = if v.len() != (1 << n_vars) {
 39        // pad to 2^n_vars
 40        [
 41            v.to_owned(),
 42            std::iter::repeat(F::zero())
 43                .take((1 << n_vars) - v.len())
 44                .collect(),
 45        ]
 46        .concat()
 47    } else {
 48        v.to_owned()
 49    };
 50    DenseMultilinearExtension::<F>::from_evaluations_vec(n_vars, v_padded)
 51}
 52
 53/// Returns the sparse multilinear extension from the given matrix, without modifying the original
 54/// matrix.
 55pub fn matrix_to_mle<F: PrimeField>(m: SparseMatrix<F>) -> SparseMultilinearExtension<F> {
 56    let n_rows = m.n_rows.next_power_of_two();
 57    let n_cols = m.n_cols.next_power_of_two();
 58    let n_vars: usize = (log2(n_rows) + log2(n_cols)) as usize; // n_vars = s + s'
 59
 60    // build the sparse vec representing the sparse matrix
 61    let mut v: Vec<(usize, F)> = Vec::new();
 62    for (i, row) in m.coeffs.iter().enumerate() {
 63        for (val, j) in row.iter() {
 64            v.push((i * n_cols + j, *val));
 65        }
 66    }
 67
 68    // convert the dense vector into a mle
 69    vec_to_mle(n_vars, &v)
 70}
 71
 72/// Takes the n_vars and a sparse vector and returns its sparse MLE.
 73pub fn vec_to_mle<F: PrimeField>(n_vars: usize, v: &[(usize, F)]) -> SparseMultilinearExtension<F> {
 74    SparseMultilinearExtension::<F>::from_evaluations(n_vars, v)
 75}
 76
 77/// Takes the n_vars and a dense vector and returns its dense MLE.
 78pub fn dense_vec_to_dense_mle<F: PrimeField>(
 79    n_vars: usize,
 80    v: &[F],
 81) -> DenseMultilinearExtension<F> {
 82    // Pad to 2^n_vars
 83    let v_padded: Vec<F> = [
 84        v.to_owned(),
 85        std::iter::repeat(F::zero())
 86            .take((1 << n_vars) - v.len())
 87            .collect(),
 88    ]
 89    .concat();
 90    DenseMultilinearExtension::<F>::from_evaluations_vec(n_vars, v_padded)
 91}
 92
 93/// Takes the n_vars and a dense vector and returns its sparse MLE.
 94pub fn dense_vec_to_mle<F: PrimeField>(n_vars: usize, v: &[F]) -> SparseMultilinearExtension<F> {
 95    let v_sparse = v
 96        .iter()
 97        .enumerate()
 98        .map(|(i, v_i)| (i, *v_i))
 99        .collect::<Vec<(usize, F)>>();
100    SparseMultilinearExtension::<F>::from_evaluations(n_vars, &v_sparse)
101}
102
103#[cfg(test)]
104mod tests {
105    use super::*;
106    use crate::{
107        arith::ccs::tests::get_test_z,
108        utils::multilinear_polynomial::fix_variables,
109        utils::multilinear_polynomial::tests::fix_last_variables,
110        utils::{hypercube::BooleanHypercube, vec::tests::to_F_matrix},
111    };
112    use ark_poly::MultilinearExtension;
113    use ark_std::Zero;
114
115    use ark_pallas::Fr;
116
117    #[test]
118    fn test_matrix_to_mle() {
119        let A = to_F_matrix::<Fr>(vec![
120            vec![2, 3, 4, 4],
121            vec![4, 11, 14, 14],
122            vec![2, 8, 17, 17],
123            vec![420, 4, 2, 0],
124        ]);
125
126        let A_mle = matrix_to_mle(A);
127        assert_eq!(A_mle.evaluations.len(), 15); // 15 non-zero elements
128        assert_eq!(A_mle.num_vars, 4); // 4x4 matrix, thus 2bit x 2bit, thus 2^4=16 evals
129
130        let A = to_F_matrix::<Fr>(vec![
131            vec![2, 3, 4, 4, 1],
132            vec![4, 11, 14, 14, 2],
133            vec![2, 8, 17, 17, 3],
134            vec![420, 4, 2, 0, 4],
135            vec![420, 4, 2, 0, 5],
136        ]);
137        let A_mle = matrix_to_mle(A.clone());
138        assert_eq!(A_mle.evaluations.len(), 23); // 23 non-zero elements
139        assert_eq!(A_mle.num_vars, 6); // 5x5 matrix, thus 3bit x 3bit, thus 2^6=64 evals
140
141        // check that the A_mle evaluated over the boolean hypercube equals the matrix A_i_j values
142        let bhc = BooleanHypercube::new(A_mle.num_vars);
143        let A_padded = pad_matrix(&A);
144        let A_padded_dense = A_padded.to_dense();
145        for (i, A_row) in A_padded_dense.iter().enumerate() {
146            for (j, _) in A_row.iter().enumerate() {
147                let s_i_j = bhc.at_i(i * A_row.len() + j);
148                assert_eq!(A_mle.evaluate(&s_i_j).unwrap(), A_padded_dense[i][j]);
149            }
150        }
151    }
152
153    #[test]
154    fn test_vec_to_mle() {
155        let z = get_test_z::<Fr>(3);
156        let n_vars = 3;
157        let z_mle = dense_vec_to_mle(n_vars, &z);
158
159        // check that the z_mle evaluated over the boolean hypercube equals the vec z_i values
160        let bhc = BooleanHypercube::new(z_mle.num_vars);
161        for (i, z_i) in z.iter().enumerate() {
162            let s_i = bhc.at_i(i);
163            assert_eq!(z_mle.evaluate(&s_i).unwrap(), z_i.clone());
164        }
165        // for the rest of elements of the boolean hypercube, expect it to evaluate to zero
166        for i in (z.len())..(1 << z_mle.num_vars) {
167            let s_i = bhc.at_i(i);
168            assert_eq!(z_mle.evaluate(&s_i).unwrap(), Fr::zero());
169        }
170    }
171
172    #[test]
173    fn test_fix_variables() {
174        let A = to_F_matrix(vec![
175            vec![2, 3, 4, 4],
176            vec![4, 11, 14, 14],
177            vec![2, 8, 17, 17],
178            vec![420, 4, 2, 0],
179        ]);
180
181        let A_mle = matrix_to_dense_mle(A.clone());
182        let A = A.to_dense();
183        let bhc = BooleanHypercube::new(2);
184        for (i, y) in bhc.enumerate() {
185            // First check that the arkworks and espresso funcs match
186            let expected_fix_left = A_mle.fix_variables(&y); // try arkworks fix_variables
187            let fix_left = fix_variables(&A_mle, &y); // try espresso fix_variables
188            assert_eq!(fix_left, expected_fix_left);
189
190            // Check that fixing first variables pins down a column
191            // i.e. fixing x to 0 will return the first column
192            //      fixing x to 1 will return the second column etc.
193            let column_i: Vec<Fr> = A.clone().iter().map(|x| x[i]).collect();
194            assert_eq!(fix_left.evaluations, column_i);
195
196            // Now check that fixing last variables pins down a row
197            // i.e. fixing y to 0 will return the first row
198            //      fixing y to 1 will return the second row etc.
199            let row_i: Vec<Fr> = A[i].clone();
200            let fix_right = fix_last_variables(&A_mle, &y);
201            assert_eq!(fix_right.evaluations, row_i);
202        }
203    }
204}