qip_iterators/iterators/
qubit_iterators.rs

1use num_traits::{One, Zero};
2use std::marker::PhantomData;
3
4use crate::utils::*;
5
6/// Iterator which provides the indices of nonzero columns for a given row of a matrix
7#[derive(Debug)]
8pub struct MatrixOpIterator<'a, P>
9where
10    P: Clone + Zero,
11{
12    n: usize,
13    data: &'a [P],
14    last_col: Option<usize>,
15}
16
17impl<'a, P> MatrixOpIterator<'a, P>
18where
19    P: Clone + Zero,
20{
21    /// Build a new iterator using the row index, the number of qubits in the matrix, and the
22    /// complex values of the matrix.
23    pub fn new(row: usize, n: usize, data: &'a [P]) -> MatrixOpIterator<P> {
24        let lower = get_flat_index(n, row, 0);
25        let upper = get_flat_index(n, row, 1 << n);
26        MatrixOpIterator {
27            n,
28            data: &data[lower..upper],
29            last_col: None,
30        }
31    }
32}
33
34impl<'a, P> Iterator for MatrixOpIterator<'a, P>
35where
36    P: Clone + Zero,
37{
38    type Item = (usize, P);
39
40    fn next(&mut self) -> Option<Self::Item> {
41        let pos = if let Some(last_col) = self.last_col {
42            last_col + 1
43        } else {
44            0
45        };
46        self.last_col = None;
47        for col in pos..(1 << self.n) {
48            let val = &self.data[col];
49            if !val.is_zero() {
50                self.last_col = Some(col);
51                return Some((col, val.clone()));
52            }
53        }
54        None
55    }
56}
57
58/// Iterator which provides the indices of nonzero columns for a given row of a sparse matrix
59#[derive(Debug)]
60pub struct SparseMatrixOpIterator<'a, P>
61where
62    P: Clone,
63{
64    data: &'a [(usize, P)],
65    last_index: Option<usize>,
66}
67
68impl<'a, P> SparseMatrixOpIterator<'a, P>
69where
70    P: Clone,
71{
72    /// Build a new iterator using the row index, and the sparse matrix data.
73    pub fn new(row: usize, data: &'a [Vec<(usize, P)>]) -> SparseMatrixOpIterator<P> {
74        SparseMatrixOpIterator {
75            data: &data[row],
76            last_index: None,
77        }
78    }
79}
80
81impl<'a, P> Iterator for SparseMatrixOpIterator<'a, P>
82where
83    P: Clone,
84{
85    type Item = (usize, P);
86
87    fn next(&mut self) -> Option<Self::Item> {
88        let pos = if let Some(last_index) = self.last_index {
89            last_index + 1
90        } else {
91            0
92        };
93
94        if pos >= self.data.len() {
95            self.last_index = None;
96            None
97        } else {
98            self.last_index = Some(pos);
99            Some(self.data[pos].clone())
100        }
101    }
102}
103
104/// Iterator which provides the indices of nonzero columns for a given row of a COp
105#[derive(Debug)]
106pub struct ControlledOpIterator<P, It>
107where
108    P: Clone + Zero,
109    It: Iterator<Item = (usize, P)>,
110{
111    row: usize,
112    index_threshold: usize,
113    op_iter: Option<It>,
114    last_col: Option<usize>,
115}
116
117impl<P, It> ControlledOpIterator<P, It>
118where
119    P: Clone + Zero,
120    It: Iterator<Item = (usize, P)>,
121{
122    /// Build a new iterator using the row index, the number of controlling indices, the number of
123    /// op indices, and the iterator for the op.
124    pub fn new<F: FnOnce(usize) -> It>(
125        row: usize,
126        n_control_indices: usize,
127        n_op_indices: usize,
128        iter_builder: F,
129    ) -> ControlledOpIterator<P, It> {
130        let n_indices = n_control_indices + n_op_indices;
131        let index_threshold = (1 << n_indices) - (1 << n_op_indices);
132        let op_iter = if row >= index_threshold {
133            Some(iter_builder(row - index_threshold))
134        } else {
135            None
136        };
137        ControlledOpIterator {
138            row,
139            index_threshold,
140            op_iter,
141            last_col: None,
142        }
143    }
144}
145
146impl<P, It> Iterator for ControlledOpIterator<P, It>
147where
148    P: Clone + Zero + One,
149    It: Iterator<Item = (usize, P)>,
150{
151    type Item = (usize, P);
152
153    fn next(&mut self) -> Option<Self::Item> {
154        if let Some(it) = &mut self.op_iter {
155            let cval = it.next();
156            let ret_val = cval.map(|(i, val)| (i + self.index_threshold, val));
157            self.last_col = ret_val.as_ref().map(|(i, _)| *i);
158            ret_val
159        } else {
160            if let Some(last_col) = self.last_col {
161                if last_col < self.row {
162                    self.last_col = Some(self.row);
163                } else {
164                    self.last_col = None;
165                }
166            } else {
167                self.last_col = Some(self.row);
168            }
169            self.last_col.map(|c| (c, P::one()))
170        }
171    }
172}
173
174/// Iterator which provides the indices of nonzero columns for a given row of a SwapOp
175#[derive(Debug)]
176pub struct SwapOpIterator<P>
177where
178    P: One,
179{
180    row: usize,
181    half_n: usize,
182    last_col: Option<usize>,
183    phantom: PhantomData<P>,
184}
185
186impl<P> SwapOpIterator<P>
187where
188    P: One,
189{
190    /// Build a new iterator using the row index, and the total number of qubits being swapped
191    /// (should be a multiple of 2).
192    pub fn new(row: usize, n_qubits: usize) -> SwapOpIterator<P> {
193        SwapOpIterator {
194            row,
195            half_n: n_qubits >> 1,
196            last_col: None,
197            phantom: PhantomData,
198        }
199    }
200}
201
202impl<P> Iterator for SwapOpIterator<P>
203where
204    P: One,
205{
206    type Item = (usize, P);
207
208    fn next(&mut self) -> Option<Self::Item> {
209        if self.last_col.is_none() {
210            let lower_mask: usize = !(std::usize::MAX << self.half_n);
211            let lower = self.row & lower_mask;
212            let upper = self.row >> self.half_n;
213            self.last_col = Some((lower << self.half_n) + upper);
214        } else {
215            self.last_col = None;
216        }
217        self.last_col.map(|col| (col, P::one()))
218    }
219}
220
221/// Iterator which provides the indices of nonzero columns for a given function.
222#[derive(Debug)]
223pub struct FunctionOpIterator<P> {
224    output_n: usize,
225    x: usize,
226    fx_xor_y: usize,
227    theta: P,
228    last_col: Option<usize>,
229    phantom: PhantomData<P>,
230}
231
232impl<P> FunctionOpIterator<P> {
233    /// Build a new iterator using the row index, the number of qubits in the input register, the
234    /// number in the output register, and a function which maps rows to a column and a phase.
235    pub fn new<F: Fn(usize) -> (usize, P)>(
236        row: usize,
237        input_n: usize,
238        output_n: usize,
239        f: F,
240    ) -> FunctionOpIterator<P> {
241        let x = row >> output_n;
242        let (fx, theta) = f(flip_bits(input_n, x));
243        let y = row & ((1 << output_n) - 1);
244        let fx_xor_y = y ^ flip_bits(output_n, fx);
245        FunctionOpIterator {
246            output_n,
247            x,
248            fx_xor_y,
249            theta,
250            last_col: None,
251            phantom: PhantomData,
252        }
253    }
254}
255
256impl<P> Iterator for FunctionOpIterator<P>
257where
258    P: Clone,
259{
260    type Item = (usize, P);
261
262    fn next(&mut self) -> Option<Self::Item> {
263        if self.last_col.is_some() {
264            self.last_col = None;
265        } else {
266            let colbits = (self.x << self.output_n) | self.fx_xor_y;
267            self.last_col = Some(colbits);
268        };
269        self.last_col.map(|col| (col, self.theta.clone()))
270    }
271}
272
273#[cfg(test)]
274mod iterator_tests {
275    use super::*;
276    use num_complex::Complex;
277    use num_traits::One;
278
279    /// Make a vector of complex numbers whose reals are given by `data`
280    fn from_reals<P: Zero + Clone>(data: &[P]) -> Vec<Complex<P>> {
281        data.iter()
282            .map(|x| Complex::<P> {
283                re: x.clone(),
284                im: P::zero(),
285            })
286            .collect()
287    }
288
289    #[test]
290    fn test_mat_iterator() {
291        let n = 1usize;
292        let mat: Vec<Vec<f64>> = (0..1 << n)
293            .map(|i| -> Vec<f64> {
294                let d = from_reals(&[0.0, 1.0, 1.0, 0.0]);
295                let it = MatrixOpIterator::new(i, n, &d);
296                let v: Vec<f64> = (0..1 << n).map(|_| 0.0).collect();
297                it.fold(v, |mut v, (indx, _)| {
298                    v[indx] = 1.0;
299                    v
300                })
301            })
302            .collect();
303
304        let expected = vec![vec![0.0, 1.0], vec![1.0, 0.0]];
305
306        assert_eq!(mat, expected);
307    }
308
309    #[test]
310    fn test_sparse_mat_iterator() {
311        let n = 1usize;
312        let one = Complex::<f64>::one();
313        let mat: Vec<Vec<f64>> = (0..1 << n)
314            .map(|i| -> Vec<f64> {
315                let d = vec![vec![(1, one)], vec![(0, one)]];
316                let it = SparseMatrixOpIterator::new(i, &d);
317                let v: Vec<f64> = (0..1 << n).map(|_| 0.0).collect();
318                it.fold(v, |mut v, (indx, _)| {
319                    v[indx] = 1.0;
320                    v
321                })
322            })
323            .collect();
324
325        let expected = vec![vec![0.0, 1.0], vec![1.0, 0.0]];
326
327        assert_eq!(mat, expected);
328    }
329
330    #[test]
331    fn test_swap_iterator() {
332        let n = 2usize;
333        let mat: Vec<Vec<f64>> = (0..1 << n)
334            .map(|i| -> Vec<f64> {
335                let it = SwapOpIterator::<f64>::new(i, n);
336                let v: Vec<f64> = (0..1 << n).map(|_| 0.0).collect();
337                it.fold(v, |mut v, (indx, _)| {
338                    v[indx] = 1.0;
339                    v
340                })
341            })
342            .collect();
343
344        let expected = vec![
345            vec![1.0, 0.0, 0.0, 0.0],
346            vec![0.0, 0.0, 1.0, 0.0],
347            vec![0.0, 1.0, 0.0, 0.0],
348            vec![0.0, 0.0, 0.0, 1.0],
349        ];
350
351        assert_eq!(mat, expected);
352    }
353
354    #[test]
355    fn test_c_iterator() {
356        let n = 2usize;
357
358        let d = from_reals(&[0.0, 1.0, 1.0, 0.0]);
359        let builder = |r: usize| MatrixOpIterator::new(r, n >> 1, &d);
360
361        let mat: Vec<Vec<f64>> = (0..1 << n)
362            .map(|i| -> Vec<f64> {
363                let it = ControlledOpIterator::new(i, n >> 1, n >> 1, builder);
364                let v: Vec<f64> = (0..1 << n).map(|_| 0.0).collect();
365                it.fold(v, |mut v, (indx, _)| {
366                    v[indx] = 1.0;
367                    v
368                })
369            })
370            .collect();
371
372        let expected = vec![
373            vec![1.0, 0.0, 0.0, 0.0],
374            vec![0.0, 1.0, 0.0, 0.0],
375            vec![0.0, 0.0, 0.0, 1.0],
376            vec![0.0, 0.0, 1.0, 0.0],
377        ];
378        assert_eq!(mat, expected);
379    }
380}