sonobe/src/utils/mle.rs
src/utils/mle.rsRaw
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}