Oxidize two qubit basis decomposer

crates/accelerate/src/convert_2q_block_matrix.rs

@@ -16,33 +16,40 @@ use pyo3::Python;
 
 use num_complex::Complex64;
 use numpy::ndarray::linalg::kron;
-use numpy::ndarray::{Array, Array2, ArrayView2};
+use numpy::ndarray::{aview2, Array2, ArrayView2};
 use numpy::{IntoPyArray, PyArray2, PyReadonlyArray2};
+use smallvec::SmallVec;
+
+static ONE_QUBIT_IDENTITY: [[Complex64; 2]; 2] = [
+    [Complex64::new(1., 0.), Complex64::new(0., 0.)],
+    [Complex64::new(0., 0.), Complex64::new(1., 0.)],
+];
 
 /// Return the matrix Operator resulting from a block of Instructions.
 #[pyfunction]
 #[pyo3(text_signature = "(op_list, /")]
 pub fn blocks_to_matrix(
     py: Python,
-    op_list: Vec<(PyReadonlyArray2<Complex64>, Vec<usize>)>,
+    op_list: Vec<(PyReadonlyArray2<Complex64>, SmallVec<[u8; 2]>)>,
 ) -> PyResult<Py<PyArray2<Complex64>>> {
-    let identity: Array2<Complex64> = Array::eye(2);
+    let identity = aview2(&ONE_QUBIT_IDENTITY);
     let input_matrix = op_list[0].0.as_array();
     let mut matrix: Array2<Complex64> = match op_list[0].1.as_slice() {
         [0] => kron(&identity, &input_matrix),
         [1] => kron(&input_matrix, &identity),
         [0, 1] => input_matrix.to_owned(),
         [1, 0] => change_basis(input_matrix),
-        [] => Array::eye(4),
+        [] => Array2::eye(4),
         _ => unreachable!(),
     };
     for (op_matrix, q_list) in op_list.into_iter().skip(1) {
         let op_matrix = op_matrix.as_array();
+
         let result = match q_list.as_slice() {
             [0] => Some(kron(&identity, &op_matrix)),
             [1] => Some(kron(&op_matrix, &identity)),
             [1, 0] => Some(change_basis(op_matrix)),
-            [] => Some(Array::eye(4)),
+            [] => Some(Array2::eye(4)),
             _ => None,
         };
         matrix = match result {
@@ -54,7 +61,8 @@ pub fn blocks_to_matrix(
 }
 
 /// Switches the order of qubits in a two qubit operation.
-fn change_basis(matrix: ArrayView2<Complex64>) -> Array2<Complex64> {
+#[inline]
+pub fn change_basis(matrix: ArrayView2<Complex64>) -> Array2<Complex64> {
     let mut trans_matrix: Array2<Complex64> = matrix.reversed_axes().to_owned();
     for index in 0..trans_matrix.ncols() {
         trans_matrix.swap([1, index], [2, index]);