Oxidize the numeric code in the Isometry gate class (#12197)

* Oxidize the numeric code in the Isometry gate class

This commit ports the numeric portion of the Isometry gate class to
rust. While this will likely improve the performance slightly this move
is more to make isolate this code from blas/lapack in numpy. We're
hitting some stability issues on arm64 mac in CI and moving this code to
rust should hopefully fix this issue. As this is more for functional
reasons no real performance tuning was done on this port, there are
likely several opportunities to improve the runtime performance of the
code.

* Remove unused import

* Use rust impl for small utility functions too

* Oxidize the linalg in UCGate too

The UCGate class is used almost exclusively by the Isometry class to
build up the definition of the isometry circuit. There were also some
linear algebra inside the function which could also be the source of the
stability issues we were seeing on arm64. This commit ports this
function as part of the larger isometry migration.

* Migrate another numeric helper method of UCGate

* Remove now unused code paths

* Remove bitstring usage with bitwise ops

This commit removes the use of bit string manipulations that were
faithfully ported from the original python logic (but left a bad taste
in my mouth) into more efficient bitwise operations (which were
possible in the original python too).

* Mostly replace Vec<u8> usage with bitwise operations

The use of intermediate Vec<u8> as proxy bitstrings was originally
ported nearly exactly from the python implementation. But since
everything is working now this commit switches to use bitwise operations
where it makes sense as this will be more efficient.

* Apply suggestions from code review

Co-authored-by: Jake Lishman <jake@binhbar.com>

* Remove python side call sites

* Fix integer typing in uc_gate.rs

* Simplify basis state bitshift loop logic

* Build set of control labels outside construct_basis_states

* Use 2 element array for reverse_qubit_state

* Micro optimize reverse_qubit_state

* Use 1d numpy arrays for diagonal inputs

* Fix lint

* Update crates/accelerate/src/isometry.rs

Co-authored-by: John Lapeyre <jlapeyre@users.noreply.github.com>

* Add back sqrt() accidentally removed by inline suggestion

* Use a constant for rz pi/2 elements

---------

Co-authored-by: Jake Lishman <jake@binhbar.com>
Co-authored-by: John Lapeyre <jlapeyre@users.noreply.github.com>

crates/accelerate/Cargo.toml

@@ -21,6 +21,7 @@ num-complex = "0.4"
 num-bigint = "0.4"
 rustworkx-core = "0.14"
 faer = "0.18.2"
+itertools = "0.12.1"
 qiskit-circuit.workspace = true
 
 [dependencies.smallvec]

crates/accelerate/src/isometry.rs

@@ -0,0 +1,367 @@
+// This code is part of Qiskit.
+//
+// (C) Copyright IBM 2024
+//
+// 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.
+
+use std::ops::BitAnd;
+
+use approx::abs_diff_eq;
+use num_complex::{Complex64, ComplexFloat};
+use pyo3::prelude::*;
+use pyo3::wrap_pyfunction;
+use pyo3::Python;
+
+use hashbrown::HashSet;
+use itertools::Itertools;
+use ndarray::prelude::*;
+use numpy::{IntoPyArray, PyReadonlyArray1, PyReadonlyArray2};
+
+use crate::two_qubit_decompose::ONE_QUBIT_IDENTITY;
+
+/// Find special unitary matrix that maps [c0,c1] to [r,0] or [0,r] if basis_state=0 or
+/// basis_state=1 respectively
+#[pyfunction]
+pub fn reverse_qubit_state(
+    py: Python,
+    state: [Complex64; 2],
+    basis_state: usize,
+    epsilon: f64,
+) -> PyObject {
+    reverse_qubit_state_inner(&state, basis_state, epsilon)
+        .into_pyarray_bound(py)
+        .into()
+}
+
+#[inline(always)]
+fn l2_norm(vec: &[Complex64]) -> f64 {
+    vec.iter()
+        .fold(0., |acc, elem| acc + elem.norm_sqr())
+        .sqrt()
+}
+
+fn reverse_qubit_state_inner(
+    state: &[Complex64; 2],
+    basis_state: usize,
+    epsilon: f64,
+) -> Array2<Complex64> {
+    let r = l2_norm(state);
+    let r_inv = 1. / r;
+    if r < epsilon {
+        Array2::eye(2)
+    } else if basis_state == 0 {
+        array![
+            [state[0].conj() * r_inv, state[1].conj() * r_inv],
+            [-state[1] * r_inv, state[0] * r_inv],
+        ]
+    } else {
+        array![
+            [-state[1] * r_inv, state[0] * r_inv],
+            [state[0].conj() * r_inv, state[1].conj() * r_inv],
+        ]
+    }
+}
+
+/// This method finds the single-qubit gates for a UCGate to disentangle a qubit:
+/// we consider the n-qubit state v[:,0] starting with k zeros (in the computational basis).
+/// The qubit with label n-s-1 is disentangled into the basis state k_s(k,s).
+
+#[pyfunction]
+pub fn find_squs_for_disentangling(
+    py: Python,
+    v: PyReadonlyArray2<Complex64>,
+    k: usize,
+    s: usize,
+    epsilon: f64,
+    n: usize,
+) -> Vec<PyObject> {
+    let v = v.as_array();
+    let k_prime = 0;
+    let i_start = if b(k, s + 1) == 0 {
+        a(k, s + 1)
+    } else {
+        a(k, s + 1) + 1
+    };
+    let mut output: Vec<Array2<Complex64>> = (0..i_start).map(|_| Array2::eye(2)).collect();
+    let mut squs: Vec<Array2<Complex64>> = (i_start..2_usize.pow((n - s - 1) as u32))
+        .map(|i| {
+            reverse_qubit_state_inner(
+                &[
+                    v[[2 * i * 2_usize.pow(s as u32) + b(k, s), k_prime]],
+                    v[[(2 * i + 1) * 2_usize.pow(s as u32) + b(k, s), k_prime]],
+                ],
+                k_s(k, s),
+                epsilon,
+            )
+        })
+        .collect();
+    output.append(&mut squs);
+    output
+        .into_iter()
+        .map(|x| x.into_pyarray_bound(py).into())
+        .collect()
+}
+
+#[pyfunction]
+pub fn apply_ucg(
+    py: Python,
+    m: PyReadonlyArray2<Complex64>,
+    k: usize,
+    single_qubit_gates: Vec<PyReadonlyArray2<Complex64>>,
+) -> PyObject {
+    let mut m = m.as_array().to_owned();
+    let shape = m.shape();
+    let num_qubits = shape[0].ilog2();
+    let num_col = shape[1];
+    let spacing: usize = 2_usize.pow(num_qubits - k as u32 - 1);
+    for j in 0..2_usize.pow(num_qubits - 1) {
+        let i = (j / spacing) * spacing + j;
+        let gate_index = i / (2_usize.pow(num_qubits - k as u32));
+        for col in 0..num_col {
+            let gate = single_qubit_gates[gate_index].as_array();
+            let a = m[[i, col]];
+            let b = m[[i + spacing, col]];
+            m[[i, col]] = gate[[0, 0]] * a + gate[[0, 1]] * b;
+            m[[i + spacing, col]] = gate[[1, 0]] * a + gate[[1, 1]] * b;
+        }
+    }
+    m.into_pyarray_bound(py).into()
+}
+
+#[inline(always)]
+fn bin_to_int(bin: &[u8]) -> usize {
+    bin.iter()
+        .fold(0_usize, |acc, digit| (acc << 1) + *digit as usize)
+}
+
+#[pyfunction]
+pub fn apply_diagonal_gate(
+    py: Python,
+    m: PyReadonlyArray2<Complex64>,
+    action_qubit_labels: Vec<usize>,
+    diag: PyReadonlyArray1<Complex64>,
+) -> PyResult<PyObject> {
+    let diag = diag.as_slice()?;
+    let mut m = m.as_array().to_owned();
+    let shape = m.shape();
+    let num_qubits = shape[0].ilog2();
+    let num_col = shape[1];
+    for state in std::iter::repeat([0_u8, 1_u8])
+        .take(num_qubits as usize)
+        .multi_cartesian_product()
+    {
+        let diag_index = action_qubit_labels
+            .iter()
+            .fold(0_usize, |acc, i| (acc << 1) + state[*i] as usize);
+        let i = bin_to_int(&state);
+        for j in 0..num_col {
+            m[[i, j]] = diag[diag_index] * m[[i, j]]
+        }
+    }
+    Ok(m.into_pyarray_bound(py).into())
+}
+
+#[pyfunction]
+pub fn apply_diagonal_gate_to_diag(
+    mut m_diagonal: Vec<Complex64>,
+    action_qubit_labels: Vec<usize>,
+    diag: PyReadonlyArray1<Complex64>,
+    num_qubits: usize,
+) -> PyResult<Vec<Complex64>> {
+    let diag = diag.as_slice()?;
+    if m_diagonal.is_empty() {
+        return Ok(m_diagonal);
+    }
+    for state in std::iter::repeat([0_u8, 1_u8])
+        .take(num_qubits)
+        .multi_cartesian_product()
+        .take(m_diagonal.len())
+    {
+        let diag_index = action_qubit_labels
+            .iter()
+            .fold(0_usize, |acc, i| (acc << 1) + state[*i] as usize);
+        let i = bin_to_int(&state);
+        m_diagonal[i] *= diag[diag_index]
+    }
+    Ok(m_diagonal)
+}
+
+/// Helper method for _apply_multi_controlled_gate. This constructs the basis states the MG gate
+/// is acting on for a specific state state_free of the qubits we neither control nor act on
+fn construct_basis_states(
+    state_free: &[u8],
+    control_set: &HashSet<usize>,
+    target_label: usize,
+) -> [usize; 2] {
+    let size = state_free.len() + control_set.len() + 1;
+    let mut e1: usize = 0;
+    let mut e2: usize = 0;
+    let mut j = 0;
+    for i in 0..size {
+        e1 <<= 1;
+        e2 <<= 1;
+        if control_set.contains(&i) {
+            e1 += 1;
+            e2 += 1;
+        } else if i == target_label {
+            e2 += 1;
+        } else {
+            assert!(j <= 1);
+            e1 += state_free[j] as usize;
+            e2 += state_free[j] as usize;
+            j += 1
+        }
+    }
+    [e1, e2]
+}
+
+#[pyfunction]
+pub fn apply_multi_controlled_gate(
+    py: Python,
+    m: PyReadonlyArray2<Complex64>,
+    control_labels: Vec<usize>,
+    target_label: usize,
+    gate: PyReadonlyArray2<Complex64>,
+) -> PyObject {
+    let mut m = m.as_array().to_owned();
+    let gate = gate.as_array();
+    let shape = m.shape();
+    let num_qubits = shape[0].ilog2();
+    let num_col = shape[1];
+    let free_qubits = num_qubits as usize - control_labels.len() - 1;
+    let control_set: HashSet<usize> = control_labels.into_iter().collect();
+    if free_qubits == 0 {
+        let [e1, e2] = construct_basis_states(&[], &control_set, target_label);
+        for i in 0..num_col {
+            let temp: Vec<_> = gate
+                .dot(&aview2(&[[m[[e1, i]]], [m[[e2, i]]]]))
+                .into_iter()
+                .take(2)
+                .collect();
+            m[[e1, i]] = temp[0];
+            m[[e2, i]] = temp[1];
+        }
+        return m.into_pyarray_bound(py).into();
+    }
+    for state_free in std::iter::repeat([0_u8, 1_u8])
+        .take(free_qubits)
+        .multi_cartesian_product()
+    {
+        let [e1, e2] = construct_basis_states(&state_free, &control_set, target_label);
+        for i in 0..num_col {
+            let temp: Vec<_> = gate
+                .dot(&aview2(&[[m[[e1, i]]], [m[[e2, i]]]]))
+                .into_iter()
+                .take(2)
+                .collect();
+            m[[e1, i]] = temp[0];
+            m[[e2, i]] = temp[1];
+        }
+    }
+    m.into_pyarray_bound(py).into()
+}
+
+#[pyfunction]
+pub fn ucg_is_identity_up_to_global_phase(
+    single_qubit_gates: Vec<PyReadonlyArray2<Complex64>>,
+    epsilon: f64,
+) -> bool {
+    let global_phase: Complex64 = if single_qubit_gates[0].as_array()[[0, 0]].abs() >= epsilon {
+        single_qubit_gates[0].as_array()[[0, 0]].finv()
+    } else {
+        return false;
+    };
+    for raw_gate in single_qubit_gates {
+        let gate = raw_gate.as_array();
+        if !abs_diff_eq!(
+            gate.mapv(|x| x * global_phase),
+            aview2(&ONE_QUBIT_IDENTITY),
+            epsilon = 1e-8 // Default tolerance from numpy for allclose()
+        ) {
+            return false;
+        }
+    }
+    true
+}
+
+#[pyfunction]
+fn diag_is_identity_up_to_global_phase(diag: Vec<Complex64>, epsilon: f64) -> bool {
+    let global_phase: Complex64 = if diag[0].abs() >= epsilon {
+        diag[0].finv()
+    } else {
+        return false;
+    };
+    for d in diag {
+        if (global_phase * d - 1.0).abs() >= epsilon {
+            return false;
+        }
+    }
+    true
+}
+
+#[pyfunction]
+pub fn merge_ucgate_and_diag(
+    py: Python,
+    single_qubit_gates: Vec<PyReadonlyArray2<Complex64>>,
+    diag: Vec<Complex64>,
+) -> Vec<PyObject> {
+    single_qubit_gates
+        .iter()
+        .enumerate()
+        .map(|(i, raw_gate)| {
+            let gate = raw_gate.as_array();
+            let res = aview2(&[
+                [diag[2 * i], Complex64::new(0., 0.)],
+                [Complex64::new(0., 0.), diag[2 * i + 1]],
+            ])
+            .dot(&gate);
+            res.into_pyarray_bound(py).into()
+        })
+        .collect()
+}
+
+#[inline(always)]
+#[pyfunction]
+fn k_s(k: usize, s: usize) -> usize {
+    if k == 0 {
+        0
+    } else {
+        let filter = 1 << s;
+        k.bitand(filter) >> s
+    }
+}
+
+#[inline(always)]
+#[pyfunction]
+fn a(k: usize, s: usize) -> usize {
+    k / 2_usize.pow(s as u32)
+}
+
+#[inline(always)]
+#[pyfunction]
+fn b(k: usize, s: usize) -> usize {
+    k - (a(k, s) * 2_usize.pow(s as u32))
+}
+
+#[pymodule]
+pub fn isometry(m: &Bound<PyModule>) -> PyResult<()> {
+    m.add_wrapped(wrap_pyfunction!(diag_is_identity_up_to_global_phase))?;
+    m.add_wrapped(wrap_pyfunction!(find_squs_for_disentangling))?;
+    m.add_wrapped(wrap_pyfunction!(reverse_qubit_state))?;
+    m.add_wrapped(wrap_pyfunction!(apply_ucg))?;
+    m.add_wrapped(wrap_pyfunction!(apply_diagonal_gate))?;
+    m.add_wrapped(wrap_pyfunction!(apply_diagonal_gate_to_diag))?;
+    m.add_wrapped(wrap_pyfunction!(apply_multi_controlled_gate))?;
+    m.add_wrapped(wrap_pyfunction!(ucg_is_identity_up_to_global_phase))?;
+    m.add_wrapped(wrap_pyfunction!(merge_ucgate_and_diag))?;
+    m.add_wrapped(wrap_pyfunction!(a))?;
+    m.add_wrapped(wrap_pyfunction!(b))?;
+    m.add_wrapped(wrap_pyfunction!(k_s))?;
+    Ok(())
+}