Adding Rust Support for MPRX and Diagonal gates

src/ir/gates/mod.rs

@@ -38,6 +38,9 @@ pub enum Gate {
     CRX(CRXGate),
     CRY(CRYGate),
     CRZ(CRZGate),
+    MPRZ(MPRZGate),
+    MPRY(MPRYGate),
+    DiagonalGate(DiagonalGate),
     RZSubGate(RZSubGate),
     VariableUnitary(VariableUnitaryGate),
     Dynamic(Arc<dyn DynGate + Send + Sync>),
@@ -60,6 +63,9 @@ impl Unitary for Gate {
             Gate::CRX(_) => 1,
             Gate::CRY(_) => 1,
             Gate::CRZ(_) => 1,
+            Gate::MPRY(v) => v.num_params(),
+            Gate::MPRZ(v) => v.num_params(),
+            Gate::DiagonalGate(v) => v.num_params(),
             Gate::RZSubGate(_) => 1,
             Gate::VariableUnitary(v) => v.num_params(),
             Gate::Dynamic(d) => d.num_params(),
@@ -82,6 +88,9 @@ impl Unitary for Gate {
             Gate::CRX(x) => x.get_utry(params, const_gates),
             Gate::CRY(y) => y.get_utry(params, const_gates),
             Gate::CRZ(z) => z.get_utry(params, const_gates),
+            Gate::MPRY(y) => y.get_utry(params, const_gates),
+            Gate::MPRZ(z) => z.get_utry(params, const_gates),
+            Gate::DiagonalGate(u) => u.get_utry(params, const_gates),
             Gate::RZSubGate(z) => z.get_utry(params, const_gates),
             Gate::VariableUnitary(v) => v.get_utry(params, const_gates),
             Gate::Dynamic(d) => d.get_utry(params, const_gates),
@@ -106,6 +115,9 @@ impl Gradient for Gate {
             Gate::CRX(x) => x.get_grad(params, const_gates),
             Gate::CRY(y) => y.get_grad(params, const_gates),
             Gate::CRZ(z) => z.get_grad(params, const_gates),
+            Gate::MPRY(y) => y.get_grad(params, const_gates),
+            Gate::MPRZ(z) => z.get_grad(params, const_gates),
+            Gate::DiagonalGate(u) => u.get_grad(params, const_gates),
             Gate::RZSubGate(z) => z.get_grad(params, const_gates),
             Gate::VariableUnitary(v) => v.get_grad(params, const_gates),
             Gate::Dynamic(d) => d.get_grad(params, const_gates),
@@ -132,6 +144,9 @@ impl Gradient for Gate {
             Gate::CRX(x) => x.get_utry_and_grad(params, const_gates),
             Gate::CRY(y) => y.get_utry_and_grad(params, const_gates),
             Gate::CRZ(z) => z.get_utry_and_grad(params, const_gates),
+            Gate::MPRY(y) => y.get_utry_and_grad(params, const_gates),
+            Gate::MPRZ(z) => z.get_utry_and_grad(params, const_gates),
+            Gate::DiagonalGate(u) => u.get_utry_and_grad(params, const_gates),
             Gate::RZSubGate(z) => z.get_utry_and_grad(params, const_gates),
             Gate::VariableUnitary(v) => v.get_utry_and_grad(params, const_gates),
             Gate::Dynamic(d) => d.get_utry_and_grad(params, const_gates),
@@ -156,6 +171,9 @@ impl Size for Gate {
             Gate::CRX(_) => 2,
             Gate::CRY(_) => 2,
             Gate::CRZ(_) => 2,
+            Gate::MPRY(v) => v.num_qudits(),
+            Gate::MPRZ(v) => v.num_qudits(),
+            Gate::DiagonalGate(u) => u.num_qudits(),
             Gate::RZSubGate(_) => 1,
             Gate::VariableUnitary(v) => v.num_qudits(),
             Gate::Dynamic(d) => d.num_qudits(),
@@ -180,6 +198,9 @@ impl Optimize for Gate {
             Gate::CRX(x) => x.optimize(env_matrix),
             Gate::CRY(y) => y.optimize(env_matrix),
             Gate::CRZ(z) => z.optimize(env_matrix),
+            Gate::MPRY(y) => y.optimize(env_matrix),
+            Gate::MPRZ(z) => z.optimize(env_matrix),
+            Gate::DiagonalGate(y) => y.optimize(env_matrix),
             Gate::RZSubGate(z) => z.optimize(env_matrix),
             Gate::VariableUnitary(v) => v.optimize(env_matrix),
             Gate::Dynamic(d) => d.optimize(env_matrix),

src/ir/gates/parameterized/diagonal.rs

@@ -0,0 +1,99 @@
+use crate::ir::gates::utils::imag_exp;
+use crate::ir::gates::{Gradient, Size};
+use crate::ir::gates::{Optimize, Unitary};
+
+use ndarray::{Array2, Array3, ArrayViewMut2};
+use ndarray_linalg::c64;
+
+
+
+/// A gate representing a Diagonal Gate on n-qubits
+#[derive(Copy, Clone, Debug, PartialEq, Default)]
+pub struct DiagonalGate {
+    size: usize,
+    dim: usize,
+    shape: (usize, usize),
+    num_parameters: usize,
+}
+
+impl DiagonalGate {
+    pub fn new(size: usize) -> Self {
+        let base: u32 = 2;
+        let dim = base.pow(size as u32) as usize;
+        let num_params = (base.pow((size) as u32) - 1) as usize;
+        DiagonalGate {
+            size,
+            dim,
+            shape: (dim, dim),
+            num_parameters: num_params,
+        }
+    }
+}
+
+impl Unitary for DiagonalGate {
+    fn num_params(&self) -> usize {
+        self.num_parameters
+    }
+
+    fn get_utry(&self, params: &[f64], _constant_gates: &[Array2<c64>]) -> Array2<c64> {
+        let mut arr: Array2<c64> = Array2::zeros((self.dim, self.dim));
+        let mut i: usize = 1;
+        arr[[0, 0]] = c64::new(1 as f64, 0 as f64);
+        for param in params {
+            arr[[i, i]] = imag_exp(*param);
+            i += 1;
+        }
+        arr
+    }
+}
+
+impl Gradient for DiagonalGate {
+    fn get_grad(&self, _params: &[f64], _const_gates: &[Array2<c64>]) -> Array3<c64> {
+        let mut grad: Array3<c64> = Array3::zeros((_params.len(), self.dim, self.dim));
+
+        for (i, param) in _params.iter().enumerate() {
+            let ind = i + 1;
+            grad[[i, ind, ind]] = c64::new(0.0, 1.0) * imag_exp(*param);
+        }
+
+        grad
+    }
+
+    fn get_utry_and_grad(
+        &self,
+        _params: &[f64],
+        _const_gates: &[Array2<c64>],
+    ) -> (Array2<c64>, Array3<c64>) {
+        let utry = self.get_utry(_params, _const_gates);
+        let grad = self.get_grad(_params, _const_gates);
+        (utry, grad)
+    }
+}
+
+impl Size for DiagonalGate {
+    fn num_qudits(&self) -> usize {
+        self.size
+    }
+}
+
+impl Optimize for DiagonalGate {
+    fn optimize(&self, env_matrix: ArrayViewMut2<c64>) -> Vec<f64> {
+        let mut thetas: Vec<f64> = Vec::new();
+        let mut i: usize = 1;
+        let real_0 = env_matrix[[0, 0]].re;
+        let imag_0 = env_matrix[[0, 0]].im;
+        // We want the top left element to be 1, so the current top left
+        // element is the global phase offset
+        let global_phase = (imag_0 / real_0).atan();
+        // Calculate each theta and subtract from the global phase
+        while i < self.dim {
+            let real = env_matrix[[i, i]].re;
+            let imag = env_matrix[[i, i]].im;
+            let phase = (imag / real).atan();
+            let a = global_phase - phase;
+            thetas.push(a);
+            i = i + 1;
+        }
+        thetas
+    }
+}

src/ir/gates/parameterized/mod.rs

@@ -1,6 +1,9 @@
 mod crx;
 mod cry;
 mod crz;
+mod mpry;
+mod mprz;
+mod diagonal;
 mod rx;
 mod rxx;
 mod ry;
@@ -27,5 +30,8 @@ pub use rzz::RZZGate;
 pub use u1::U1Gate;
 pub use u2::U2Gate;
 pub use u3::U3Gate;
+pub use mpry::MPRYGate;
+pub use mprz::MPRZGate;
+pub use diagonal::DiagonalGate;
 pub use rzsub::RZSubGate;
 pub use variable::VariableUnitaryGate;

src/ir/gates/parameterized/mpry.rs

@@ -0,0 +1,115 @@
+use crate::ir::gates::utils::{rot_y, get_indices, svd};
+use crate::ir::gates::{Gradient, Size};
+use crate::ir::gates::{Optimize, Unitary};
+use crate::{i, r};
+
+use ndarray::{Array2, Array3, ArrayViewMut2};
+use ndarray_linalg::c64;
+use crate::squaremat::Matmul;
+
+/// A gate representing a multiplexed Y rotation with a target qudit
+/// and n -1 select qudits
+#[derive(Copy, Clone, Debug, PartialEq, Default)]
+pub struct MPRYGate {
+    size: usize,
+    dim: usize,
+    target_qudit: usize,
+    shape: (usize, usize),
+    num_parameters: usize,
+}
+
+impl MPRYGate {
+    pub fn new(size: usize, target_qudit: usize) -> Self {
+        let base: u32 = 2;
+        let dim = base.pow(size as u32) as usize;
+        let num_params = base.pow((size - 1) as u32) as usize;
+        MPRYGate {
+            size,
+            dim,
+            target_qudit,
+            shape: (dim, dim),
+            num_parameters: num_params,
+        }
+    }
+}
+
+impl Unitary for MPRYGate {
+    fn num_params(&self) -> usize {
+        self.num_parameters
+    }
+
+    fn get_utry(&self, params: &[f64], _constant_gates: &[Array2<c64>]) -> Array2<c64> {
+        // Return the unitary matrix for the MPRY gate
+        // If the target qudit is the bottom most qubit, then the matrix
+        // is block diagonal of RY rotations
+        // Otherwise, use the get_indices to permute the matrix appropriately
+        let mut arr: Array2<c64> = Array2::zeros((self.dim, self.dim));
+        let mut i: usize = 0;
+        for param in params {
+            let block = rot_y(*param);
+            let (x1, x2) = get_indices(i, self.target_qudit, self.size);
+            arr[[x1, x1]] = block[[0, 0]];
+            arr[[x2, x2]] = block[[1, 1]];
+            arr[[x2, x1]] = block[[1, 0]];
+            arr[[x1, x2]] = block[[0, 1]];
+            i += 1;
+        }
+        let (u, vt) = svd(arr.view_mut());
+        u.matmul(vt.view())
+    }
+}
+
+impl Gradient for MPRYGate {
+    fn get_grad(&self, _params: &[f64], _const_gates: &[Array2<c64>]) -> Array3<c64> {
+        // Each theta is independent, so the gradient is simply
+        // the derivative of the rotation matrix with respect to theta
+        let mut grad: Array3<c64> = Array3::zeros((_params.len(), self.dim, self.dim));
+
+        for (i, &param) in _params.iter().enumerate() {
+            let dcos = -((param / 2.0).sin()) / 2.0;
+            let dsin = -c64::new(0.0, 1.0) * ((param / 2.0).cos()) / 2.0;
+
+            let (x1, x2) = get_indices(i, self.target_qudit, self.size);
+
+            grad[[i, x1, x1]] = c64::new(dcos, 0.0);
+            grad[[i, x2, x2]] = c64::new(dcos, 0.0);
+            grad[[i, x2, x1]] = dsin;
+            grad[[i, x1, x2]] = -dsin;
+        }
+        grad
+    }
+
+    fn get_utry_and_grad(
+        &self,
+        _params: &[f64],
+        _const_gates: &[Array2<c64>],
+    ) -> (Array2<c64>, Array3<c64>) {
+        let utry = self.get_utry(_params, _const_gates);
+        let grad = self.get_grad(_params, _const_gates);
+        (utry, grad)
+    }
+}
+
+impl Size for MPRYGate {
+    fn num_qudits(&self) -> usize {
+        self.size
+    }
+}
+
+impl Optimize for MPRYGate {
+    fn optimize(&self, env_matrix: ArrayViewMut2<c64>) -> Vec<f64> {
+        let mut thetas: Vec<f64> = Vec::new();
+        let mut i: usize = 0;
+        // Each variable is independent, so you can optimize each Ry separately
+        while i < self.dim {
+            let (x1, x2) = get_indices(i, self.target_qudit, self.size);
+            let a = (env_matrix[[x1, x1]] + env_matrix[[x2, x2]]).re;
+            let b = (env_matrix[[x2, x1]] - env_matrix[[x1, x2]]).re;
+            let mut theta = 2.0 * (a / (a.powi(2) + b.powi(2)).sqrt()).acos();
+            theta *= if b > 0.0 { -1.0 } else { 1.0 };
+            thetas.push(theta);
+            i = i + 2;
+        }
+        thetas
+    }
+}