1/// Some basic MLE utilities
  2use ark_ff::PrimeField;
  3use ark_poly::DenseMultilinearExtension;
  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_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_mle(n_vars, &v)
 34}
 35
 36/// Takes the n_vars and a dense vector and returns its dense MLE.
 37pub fn vec_to_mle<F: PrimeField>(n_vars: usize, v: &Vec<F>) -> DenseMultilinearExtension<F> {
 38    let v_padded: Vec<F> = if v.len() != (1 << n_vars) {
 39        // pad to 2^n_vars
 40        [
 41            v.clone(),
 42            std::iter::repeat(F::zero())
 43                .take((1 << n_vars) - v.len())
 44                .collect(),
 45        ]
 46        .concat()
 47    } else {
 48        v.clone()
 49    };
 50    DenseMultilinearExtension::<F>::from_evaluations_vec(n_vars, v_padded)
 51}
 52
 53pub fn dense_vec_to_mle<F: PrimeField>(n_vars: usize, v: &Vec<F>) -> DenseMultilinearExtension<F> {
 54    // Pad to 2^n_vars
 55    let v_padded: Vec<F> = [
 56        v.clone(),
 57        std::iter::repeat(F::zero())
 58            .take((1 << n_vars) - v.len())
 59            .collect(),
 60    ]
 61    .concat();
 62    DenseMultilinearExtension::<F>::from_evaluations_vec(n_vars, v_padded)
 63}
 64
 65#[cfg(test)]
 66mod tests {
 67    use super::*;
 68    use crate::{
 69        ccs::tests::get_test_z,
 70        utils::multilinear_polynomial::fix_variables,
 71        utils::multilinear_polynomial::tests::fix_last_variables,
 72        utils::{hypercube::BooleanHypercube, vec::tests::to_F_matrix},
 73    };
 74    use ark_poly::MultilinearExtension;
 75    use ark_std::Zero;
 76
 77    use ark_pallas::Fr;
 78
 79    #[test]
 80    fn test_matrix_to_mle() {
 81        let A = to_F_matrix::<Fr>(vec![
 82            vec![2, 3, 4, 4],
 83            vec![4, 11, 14, 14],
 84            vec![2, 8, 17, 17],
 85            vec![420, 4, 2, 0],
 86        ]);
 87
 88        let A_mle = matrix_to_mle(A);
 89        assert_eq!(A_mle.evaluations.len(), 16); // 4x4 matrix, thus 2bit x 2bit, thus 2^4=16 evals
 90
 91        let A = to_F_matrix::<Fr>(vec![
 92            vec![2, 3, 4, 4, 1],
 93            vec![4, 11, 14, 14, 2],
 94            vec![2, 8, 17, 17, 3],
 95            vec![420, 4, 2, 0, 4],
 96            vec![420, 4, 2, 0, 5],
 97        ]);
 98        let A_mle = matrix_to_mle(A.clone());
 99        assert_eq!(A_mle.evaluations.len(), 64); // 5x5 matrix, thus 3bit x 3bit, thus 2^6=64 evals
100
101        // check that the A_mle evaluated over the boolean hypercube equals the matrix A_i_j values
102        let bhc = BooleanHypercube::new(A_mle.num_vars);
103        let A_padded = pad_matrix(&A);
104        let A_padded_dense = A_padded.to_dense();
105        for (i, A_row) in A_padded_dense.iter().enumerate() {
106            for (j, _) in A_row.iter().enumerate() {
107                let s_i_j = bhc.at_i(i * A_row.len() + j);
108                assert_eq!(A_mle.evaluate(&s_i_j).unwrap(), A_padded_dense[i][j]);
109            }
110        }
111    }
112
113    #[test]
114    fn test_vec_to_mle() {
115        let z = get_test_z::<Fr>(3);
116        let n_vars = 3;
117        let z_mle = dense_vec_to_mle(n_vars, &z);
118
119        // check that the z_mle evaluated over the boolean hypercube equals the vec z_i values
120        let bhc = BooleanHypercube::new(z_mle.num_vars);
121        for (i, z_i) in z.iter().enumerate() {
122            let s_i = bhc.at_i(i);
123            assert_eq!(z_mle.evaluate(&s_i).unwrap(), z_i.clone());
124        }
125        // for the rest of elements of the boolean hypercube, expect it to evaluate to zero
126        for i in (z.len())..(1 << z_mle.num_vars) {
127            let s_i = bhc.at_i(i);
128            assert_eq!(z_mle.evaluate(&s_i).unwrap(), Fr::zero());
129        }
130    }
131
132    #[test]
133    fn test_fix_variables() {
134        let A = to_F_matrix(vec![
135            vec![2, 3, 4, 4],
136            vec![4, 11, 14, 14],
137            vec![2, 8, 17, 17],
138            vec![420, 4, 2, 0],
139        ]);
140
141        let A_mle = matrix_to_mle(A.clone());
142        let A = A.to_dense();
143        let bhc = BooleanHypercube::new(2);
144        for (i, y) in bhc.enumerate() {
145            // First check that the arkworks and espresso funcs match
146            let expected_fix_left = A_mle.fix_variables(&y); // try arkworks fix_variables
147            let fix_left = fix_variables(&A_mle, &y); // try espresso fix_variables
148            assert_eq!(fix_left, expected_fix_left);
149
150            // Check that fixing first variables pins down a column
151            // i.e. fixing x to 0 will return the first column
152            //      fixing x to 1 will return the second column etc.
153            let column_i: Vec<Fr> = A.clone().iter().map(|x| x[i]).collect();
154            assert_eq!(fix_left.evaluations, column_i);
155
156            // Now check that fixing last variables pins down a row
157            // i.e. fixing y to 0 will return the first row
158            //      fixing y to 1 will return the second row etc.
159            let row_i: Vec<Fr> = A[i].clone();
160            let fix_right = fix_last_variables(&A_mle, &y);
161            assert_eq!(fix_right.evaluations, row_i);
162        }
163    }
164}