Reimplement two_qubit_decompose.num_basis_gates in rust (#11019)

* Use newer rust

* Adding weyl chamber code in rust

* Finish num_basis_gates in rust, but calculated values wrong

* Fix bugs in weyl_coordinates

* Upgrade faer from 0.12 to 0.13

* This will allow copyless convert of numpy complex matrix to faer
* abs2 disappeared probably an error, so we work around here.

* Use copyless conversion of complex matrix from numpy to faer

* Fix thinko bugs so that existing tests pass

* Add comment

* Remove allocation

* Replace loop with iterator

* Silence clippy complaints

* refactor

* Run rustfmt

* Slight refactor and format

* Reorganize new code in crate accelerate

* Moved `eigenvalues` into module `utils`.
* Access num_basis_gates via module two_qubit_decompose

* Remove some allocations

* Run rustfmt

* Raise faer version to 0.13.5, lower rustc to 1.67.0

faer 0.13.5 introduces a specific MSRV

* Fix some conflicts with main branch

* Fix more conflicts

* More conflicts

* More conflicts

* Update crates/accelerate/src/two_qubit_decompose.rs

Co-authored-by: Matthew Treinish <mtreinish@kortar.org>

* Apply suggestions from code review

Apply several suggested edits made in review

Co-authored-by: Matthew Treinish <mtreinish@kortar.org>

* Increase faer version and remove some cruft

* We introduced to_num_complex, but this required upgrading faer. We chose 0.15
* Upgrade our code, in turn, for the upgrade to faer 0.15
* Remove a function exported via pyo3 that was only used for testing
* Remove `def old_num_basis_gate`

* Make __num_basis_gates more idiomatic

* Explain use of `min_by` when `max_by` is expected

* Replace myabs2 with faer_abs2

We recently increased the version of the faer dependency.
The newer version has an abs2 in the api named `faer_abs2`.

So we can delete `myabs2`.

* Use alternative mod implementation following review requests

* Make pyo3 function eigenvalues return ndarray rather than list

Following review suggestion.

* Bump faer version to latest 0.16

This was previously at 0.15 for accidental reasons.

* Use rem_euclid for remainder

* Fix black complaint

* Fix pylint complaints

* Apply suggestions from code review

Co-authored-by: Matthew Treinish <mtreinish@kortar.org>

* Remove nalgebra feature from faer

* Use constants for fractions of pi

---------

Co-authored-by: Matthew Treinish <mtreinish@kortar.org>

crates/accelerate/Cargo.toml

@@ -45,3 +45,10 @@ features = ["rayon"]
 [dependencies.indexmap]
 version = "2.2.1"
 features = ["rayon"]
+
+[dependencies.faer]
+version = "0.16.0"
+features = ["ndarray"]
+
+[dependencies.faer-core]
+version = "0.16.0"

crates/accelerate/src/lib.rs

@@ -31,6 +31,8 @@ mod sabre_swap;
 mod sampled_exp_val;
 mod sparse_pauli_op;
 mod stochastic_swap;
+mod two_qubit_decompose;
+mod utils;
 mod vf2_layout;
 
 #[inline]
@@ -61,6 +63,8 @@ fn _accelerate(_py: Python<'_>, m: &PyModule) -> PyResult<()> {
     m.add_wrapped(wrap_pymodule!(sampled_exp_val::sampled_exp_val))?;
     m.add_wrapped(wrap_pymodule!(sabre_layout::sabre_layout))?;
     m.add_wrapped(wrap_pymodule!(vf2_layout::vf2_layout))?;
+    m.add_wrapped(wrap_pymodule!(two_qubit_decompose::two_qubit_decompose))?;
+    m.add_wrapped(wrap_pymodule!(utils::utils))?;
     m.add_wrapped(wrap_pymodule!(
         euler_one_qubit_decomposer::euler_one_qubit_decomposer
     ))?;

crates/accelerate/src/two_qubit_decompose.rs

@@ -0,0 +1,177 @@
+// This code is part of Qiskit.
+//
+// (C) Copyright IBM 2023
+//
+// This code is licensed under the Apache License, Version 2.0. You may
+// obtain a copy of this license in the LICENSE.txt file in the root directory
+// of this source tree or at http://www.apache.org/licenses/LICENSE-2.0.
+//
+// Any modifications or derivative works of this code must retain this
+// copyright notice, and modified files need to carry a notice indicating
+// that they have been altered from the originals.
+
+// In numpy matrices real and imaginary components are adjacent:
+//   np.array([1,2,3], dtype='complex').view('float64')
+//   array([1., 0., 2., 0., 3., 0.])
+// The matrix faer::Mat<c64> has this layout.
+// faer::Mat<num_complex::Complex<f64>> instead stores a matrix
+// of real components and one of imaginary components.
+// In order to avoid copying we want to use `MatRef<c64>` or `MatMut<c64>`.
+
+use num_complex::Complex;
+use pyo3::prelude::*;
+use pyo3::wrap_pyfunction;
+use pyo3::Python;
+use std::f64::consts::PI;
+
+use faer::IntoFaerComplex;
+use faer::{mat, prelude::*, scale, Mat, MatRef};
+use faer_core::{c64, ComplexField};
+use numpy::PyReadonlyArray2;
+
+use crate::utils;
+
+const PI2: f64 = PI / 2.0;
+const PI4: f64 = PI / 4.0;
+const PI32: f64 = 3.0 * PI2;
+// FIXME: zero and one exist but I cant find the right incantation
+const C0: c64 = c64 { re: 0.0, im: 0.0 };
+const C1: c64 = c64 { re: 1.0, im: 0.0 };
+const C1IM: c64 = c64 { re: 0.0, im: 1.0 };
+
+// FIXME: Find a way to compute these matrices at compile time.
+fn transform_from_magic_basis(unitary: Mat<c64>) -> Mat<c64> {
+    let _b_nonnormalized: Mat<c64> = mat![
+        [C1, C1IM, C0, C0],
+        [C0, C0, C1IM, C1],
+        [C0, C0, C1IM, -C1],
+        [C1, -C1IM, C0, C0]
+    ];
+    let _b_nonnormalized_dagger = scale(c64 { re: 0.5, im: 0.0 }) * _b_nonnormalized.adjoint();
+    _b_nonnormalized_dagger * unitary * _b_nonnormalized
+}
+
+// faer::c64 and num_complex::Complex<f64> are both structs
+// holding two f64's. But several functions are not defined for
+// c64. So we implement them here. These things should be contribute
+// upstream.
+
+pub trait PowF {
+    fn powf(self, pow: f64) -> c64;
+}
+
+impl PowF for c64 {
+    fn powf(self, pow: f64) -> c64 {
+        c64::from(self.to_num_complex().powf(pow))
+    }
+}
+
+pub trait Arg {
+    fn arg(self) -> f64;
+}
+
+impl Arg for c64 {
+    fn arg(self) -> f64 {
+        self.to_num_complex().arg()
+    }
+}
+
+fn __weyl_coordinates(unitary: MatRef<c64>) -> [f64; 3] {
+    let uscaled = scale(C1 / unitary.determinant().powf(0.25)) * unitary;
+    let uup = transform_from_magic_basis(uscaled);
+    let mut darg: Vec<_> = (uup.transpose() * &uup)
+        .complex_eigenvalues()
+        .into_iter()
+        .map(|x: c64| -x.arg() / 2.0)
+        .collect();
+    darg[3] = -darg[0] - darg[1] - darg[2];
+    let mut cs: Vec<_> = (0..3)
+        .map(|i| ((darg[i] + darg[3]) / 2.0).rem_euclid(2.0 * PI))
+        .collect();
+    let cstemp: Vec<f64> = cs
+        .iter()
+        .map(|x| x.rem_euclid(PI2))
+        .map(|x| x.min(PI2 - x))
+        .collect();
+    let mut order = utils::arg_sort(&cstemp);
+    (order[0], order[1], order[2]) = (order[1], order[2], order[0]);
+    (cs[0], cs[1], cs[2]) = (cs[order[0]], cs[order[1]], cs[order[2]]);
+
+    // Flip into Weyl chamber
+    if cs[0] > PI2 {
+        cs[0] -= PI32;
+    }
+    if cs[1] > PI2 {
+        cs[1] -= PI32;
+    }
+    let mut conjs = 0;
+    if cs[0] > PI4 {
+        cs[0] = PI2 - cs[0];
+        conjs += 1;
+    }
+    if cs[1] > PI4 {
+        cs[1] = PI2 - cs[1];
+        conjs += 1;
+    }
+    if cs[2] > PI2 {
+        cs[2] -= PI32;
+    }
+    if conjs == 1 {
+        cs[2] = PI2 - cs[2];
+    }
+    if cs[2] > PI4 {
+        cs[2] -= PI2;
+    }
+    [cs[1], cs[0], cs[2]]
+}
+
+#[pyfunction]
+#[pyo3(text_signature = "(basis_b, basis_fidelity, unitary, /")]
+pub fn _num_basis_gates(
+    basis_b: f64,
+    basis_fidelity: f64,
+    unitary: PyReadonlyArray2<Complex<f64>>,
+) -> usize {
+    let u = unitary.as_array().into_faer_complex();
+    __num_basis_gates(basis_b, basis_fidelity, u)
+}
+
+fn __num_basis_gates(basis_b: f64, basis_fidelity: f64, unitary: MatRef<c64>) -> usize {
+    let [a, b, c] = __weyl_coordinates(unitary);
+    let traces = [
+        c64::new(
+            4.0 * (a.cos() * b.cos() * c.cos()),
+            4.0 * (a.sin() * b.sin() * c.sin()),
+        ),
+        c64::new(
+            4.0 * (PI4 - a).cos() * (basis_b - b).cos() * c.cos(),
+            4.0 * (PI4 - a).sin() * (basis_b - b).sin() * c.sin(),
+        ),
+        c64::new(4.0 * c.cos(), 0.0),
+        c64::new(4.0, 0.0),
+    ];
+    // The originial Python had `np.argmax`, which returns the lowest index in case two or more
+    // values have a common maximum value.
+    // `max_by` and `min_by` return the highest and lowest indices respectively, in case of ties.
+    // So to reproduce `np.argmax`, we use `min_by` and switch the order of the
+    // arguments in the comparison.
+    traces
+        .into_iter()
+        .enumerate()
+        .map(|(idx, trace)| (idx, trace_to_fid(&trace) * basis_fidelity.powi(idx as i32)))
+        .min_by(|(_idx1, fid1), (_idx2, fid2)| fid2.partial_cmp(fid1).unwrap())
+        .unwrap()
+        .0
+}
+
+/// Average gate fidelity is :math:`Fbar = (d + |Tr (Utarget \\cdot U^dag)|^2) / d(d+1)`
+/// M. Horodecki, P. Horodecki and R. Horodecki, PRA 60, 1888 (1999)
+fn trace_to_fid(trace: &c64) -> f64 {
+    (4.0 + trace.faer_abs2()) / 20.0
+}
+
+#[pymodule]
+pub fn two_qubit_decompose(_py: Python, m: &PyModule) -> PyResult<()> {
+    m.add_wrapped(wrap_pyfunction!(_num_basis_gates))?;
+    Ok(())
+}

crates/accelerate/src/utils.rs

@@ -0,0 +1,37 @@
+use pyo3::prelude::*;
+
+use faer::prelude::*;
+use faer::IntoFaerComplex;
+use num_complex::Complex;
+use numpy::{IntoPyArray, PyReadonlyArray2};
+
+/// Return indices that sort partially ordered data.
+/// If `data` contains two elements that are incomparable,
+/// an error will be thrown.
+pub fn arg_sort<T: PartialOrd>(data: &[T]) -> Vec<usize> {
+    let mut indices = (0..data.len()).collect::<Vec<_>>();
+    indices.sort_by(|&a, &b| data[a].partial_cmp(&data[b]).unwrap());
+    indices
+}
+
+/// Return the eigenvalues of `unitary` as a one-dimensional `numpy.ndarray`
+/// with `dtype(complex128)`.
+#[pyfunction]
+#[pyo3(text_signature = "(unitary, /")]
+pub fn eigenvalues(py: Python, unitary: PyReadonlyArray2<Complex<f64>>) -> PyObject {
+    unitary
+        .as_array()
+        .into_faer_complex()
+        .complex_eigenvalues()
+        .into_iter()
+        .map(|x| Complex::<f64>::new(x.re, x.im))
+        .collect::<Vec<_>>()
+        .into_pyarray(py)
+        .into()
+}
+
+#[pymodule]
+pub fn utils(_py: Python, m: &PyModule) -> PyResult<()> {
+    m.add_wrapped(wrap_pyfunction!(eigenvalues))?;
+    Ok(())
+}

qiskit/synthesis/two_qubit/two_qubit_decompose.py

@@ -39,13 +39,13 @@ from qiskit.circuit import QuantumRegister, QuantumCircuit, Gate
 from qiskit.circuit.library.standard_gates import CXGate, RXGate, RYGate, RZGate
 from qiskit.exceptions import QiskitError
 from qiskit.quantum_info.operators import Operator
-from qiskit.synthesis.two_qubit.weyl import weyl_coordinates, transform_to_magic_basis
+from qiskit.synthesis.two_qubit.weyl import transform_to_magic_basis
 from qiskit.synthesis.one_qubit.one_qubit_decompose import (
     OneQubitEulerDecomposer,
     DEFAULT_ATOL,
 )
 from qiskit.utils.deprecation import deprecate_arg
-
+from qiskit._accelerate import two_qubit_decompose
 
 logger = logging.getLogger(__name__)
 
@@ -1413,23 +1413,9 @@ class TwoQubitBasisDecomposer:
         """Computes the number of basis gates needed in
         a decomposition of input unitary
         """
-        unitary = np.asarray(unitary, dtype=complex)
-        a, b, c = weyl_coordinates(unitary)[:]
-        traces = [
-            4
-            * (
-                math.cos(a) * math.cos(b) * math.cos(c)
-                + 1j * math.sin(a) * math.sin(b) * math.sin(c)
-            ),
-            4
-            * (
-                math.cos(np.pi / 4 - a) * math.cos(self.basis.b - b) * math.cos(c)
-                + 1j * math.sin(np.pi / 4 - a) * math.sin(self.basis.b - b) * math.sin(c)
-            ),
-            4 * math.cos(c),
-            4,
-        ]
-        return np.argmax([trace_to_fid(traces[i]) * self.basis_fidelity**i for i in range(4)])
+        return two_qubit_decompose._num_basis_gates(
+            self.basis.b, self.basis_fidelity, np.asarray(unitary, dtype=complex)
+        )
 
 
 class TwoQubitDecomposeUpToDiagonal:

qiskit/synthesis/two_qubit/weyl.py

@@ -64,7 +64,6 @@ def weyl_coordinates(U: np.ndarray) -> np.ndarray:
     Up = transform_to_magic_basis(U, reverse=True)
     # We only need the eigenvalues of `M2 = Up.T @ Up` here, not the full diagonalization.
     D = la.eigvals(Up.T @ Up)
-
     d = -np.angle(D) / 2
     d[3] = -d[0] - d[1] - d[2]
     cs = np.mod((d[:3] + d[3]) / 2, 2 * np.pi)