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