qip_iterators/iterators/
ops.rs

1use crate::iterators::{
2    ControlledOpIterator, MatrixOpIterator, SparseMatrixOpIterator, SwapOpIterator,
3};
4use num_traits::{One, Zero};
5use std::fmt;
6use std::iter::Sum;
7use std::ops::Mul;
8
9/// Ops which can be applied to quantum states.
10#[derive(Clone)]
11pub enum MatrixOp<P> {
12    /// Indices, Matrix data
13    Matrix(Vec<usize>, Vec<P>),
14    /// Indices, per row [(col, value)]
15    SparseMatrix(Vec<usize>, Vec<Vec<(usize, P)>>),
16    /// A indices, B indices
17    Swap(Vec<usize>, Vec<usize>),
18    /// Control indices, Op indices, Op
19    Control(Vec<usize>, Vec<usize>, Box<MatrixOp<P>>),
20}
21
22impl<P> MatrixOp<P> {
23    /// Get the number of indices represented by `op`
24    pub fn num_indices(&self) -> usize {
25        match self {
26            MatrixOp::Matrix(indices, _) => indices.len(),
27            MatrixOp::SparseMatrix(indices, _) => indices.len(),
28            MatrixOp::Swap(a, b) => a.len() + b.len(),
29            MatrixOp::Control(cs, os, _) => cs.len() + os.len(),
30        }
31    }
32
33    /// Make a new dense matrix op
34    pub fn new_matrix<Indx, Dat>(indices: Indx, data: Dat) -> Self
35    where
36        Indx: Into<Vec<usize>>,
37        Dat: Into<Vec<P>>,
38    {
39        Self::Matrix(indices.into(), data.into())
40    }
41
42    /// Make a new sparse matrix op
43    pub fn new_sparse<Indx, Dat>(indices: Indx, data: Dat) -> Self
44    where
45        Indx: Into<Vec<usize>>,
46        Dat: Into<Vec<Vec<(usize, P)>>>,
47    {
48        Self::SparseMatrix(indices.into(), data.into())
49    }
50
51    /// Make a new swap matrix op
52    pub fn new_swap<IndxA, IndxB>(a: IndxA, b: IndxB) -> Self
53    where
54        IndxA: Into<Vec<usize>>,
55        IndxB: Into<Vec<usize>>,
56    {
57        Self::Swap(a.into(), b.into())
58    }
59
60    /// Make a new control matrix op
61    pub fn new_control<IndxA, IndxB>(c: IndxA, r: IndxB, op: MatrixOp<P>) -> Self
62    where
63        IndxA: Into<Vec<usize>>,
64        IndxB: Into<Vec<usize>>,
65    {
66        Self::Control(c.into(), r.into(), Box::new(op))
67    }
68}
69
70impl<P> MatrixOp<P>
71where
72    P: Clone + Zero + One + Mul,
73{
74    /// The function `f` maps a column to a complex value (given the `row`) for the op matrix.
75    /// Sums for all nonzero entries for a given `op`
76    pub fn sum_for_op_cols<T, F>(&self, nindices: usize, row: usize, f: F) -> T
77    where
78        T: Sum,
79        F: Fn((usize, P)) -> T,
80    {
81        match &self {
82            MatrixOp::Matrix(_, data) => MatrixOpIterator::new(row, nindices, data).map(f).sum(),
83            MatrixOp::SparseMatrix(_, data) => SparseMatrixOpIterator::new(row, data.as_slice())
84                .map(f)
85                .sum(),
86            MatrixOp::Swap(_, _) => SwapOpIterator::new(row, nindices).map(f).sum(),
87            MatrixOp::Control(c_indices, o_indices, op) => {
88                let n_control_indices = c_indices.len();
89                let n_op_indices = o_indices.len();
90                op.sum_for_control_iterator(row, n_control_indices, n_op_indices, f)
91            }
92        }
93    }
94
95    fn sum_for_control_iterator<T, F>(
96        &self,
97        row: usize,
98        n_control_indices: usize,
99        n_op_indices: usize,
100        f: F,
101    ) -> T
102    where
103        T: Sum,
104        F: Fn((usize, P)) -> T,
105    {
106        match &self {
107            MatrixOp::Matrix(_, data) => {
108                let iter_builder = |row: usize| MatrixOpIterator::new(row, n_op_indices, data);
109                ControlledOpIterator::new(row, n_control_indices, n_op_indices, iter_builder)
110                    .map(f)
111                    .sum()
112            }
113            MatrixOp::SparseMatrix(_, data) => {
114                let iter_builder = |row: usize| SparseMatrixOpIterator::new(row, data);
115                ControlledOpIterator::new(row, n_control_indices, n_op_indices, iter_builder)
116                    .map(f)
117                    .sum()
118            }
119            MatrixOp::Swap(_, _) => {
120                let iter_builder = |row: usize| SwapOpIterator::new(row, n_op_indices);
121                ControlledOpIterator::new(row, n_control_indices, n_op_indices, iter_builder)
122                    .map(f)
123                    .sum()
124            }
125            // Control ops are automatically collapsed if made with helper, but implement this anyway
126            // just to account for the possibility.
127            MatrixOp::Control(c_indices, o_indices, op) => {
128                let n_control_indices = n_control_indices + c_indices.len();
129                let n_op_indices = o_indices.len();
130                op.sum_for_control_iterator(row, n_control_indices, n_op_indices, f)
131            }
132        }
133    }
134}
135
136impl<P> fmt::Debug for MatrixOp<P> {
137    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
138        let (name, indices) = match self {
139            MatrixOp::Matrix(indices, _) => ("Matrix".to_string(), indices.clone()),
140            MatrixOp::SparseMatrix(indices, _) => ("SparseMatrix".to_string(), indices.clone()),
141            MatrixOp::Swap(a_indices, b_indices) => {
142                let indices: Vec<_> = a_indices
143                    .iter()
144                    .cloned()
145                    .chain(b_indices.iter().cloned())
146                    .collect();
147                ("Swap".to_string(), indices)
148            }
149            MatrixOp::Control(indices, _, op) => {
150                let name = format!("C({:?})", *op);
151                (name, indices.clone())
152            }
153        };
154        let int_strings = indices
155            .iter()
156            .map(|x| x.clone().to_string())
157            .collect::<Vec<String>>();
158
159        write!(f, "{}[{}]", name, int_strings.join(", "))
160    }
161}