Merge pull request #13 from BQSKit/residual-fix

Fixes, new gate, and 0.4.1

Author: edyounis

pyproject.toml

@@ -8,7 +8,7 @@ build-backend = "maturin"
 
 [project]
 name = "bqskitrs"
-version = "0.4.0"
+version = "0.4.1"
 maintainers = [
     {name = "Ethan Smith", email = "[email protected]"},
     {name = "Ed Younis", email = "[email protected]"},

src/ir/gates/mod.rs

@@ -38,6 +38,7 @@ pub enum Gate {
     CRX(CRXGate),
     CRY(CRYGate),
     CRZ(CRZGate),
+    RZSubGate(RZSubGate),
     VariableUnitary(VariableUnitaryGate),
     Dynamic(Arc<dyn DynGate + Send + Sync>),
 }
@@ -59,6 +60,7 @@ impl Unitary for Gate {
             Gate::CRX(_) => 1,
             Gate::CRY(_) => 1,
             Gate::CRZ(_) => 1,
+            Gate::RZSubGate(_) => 1,
             Gate::VariableUnitary(v) => v.num_params(),
             Gate::Dynamic(d) => d.num_params(),
         }
@@ -80,6 +82,7 @@ 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::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),
         }
@@ -103,6 +106,7 @@ 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::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),
         }
@@ -128,6 +132,7 @@ 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::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),
         }
@@ -151,6 +156,7 @@ impl Size for Gate {
             Gate::CRX(_) => 2,
             Gate::CRY(_) => 2,
             Gate::CRZ(_) => 2,
+            Gate::RZSubGate(_) => 1,
             Gate::VariableUnitary(v) => v.num_qudits(),
             Gate::Dynamic(d) => d.num_qudits(),
         }
@@ -174,6 +180,7 @@ 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::RZSubGate(z) => z.optimize(env_matrix),
             Gate::VariableUnitary(v) => v.optimize(env_matrix),
             Gate::Dynamic(d) => d.optimize(env_matrix),
         }

src/ir/gates/parameterized/mod.rs

@@ -12,6 +12,7 @@ mod u2;
 mod u3;
 mod u8;
 mod variable;
+mod rzsub;
 
 pub use self::u8::U8Gate;
 pub use crx::CRXGate;
@@ -26,4 +27,5 @@ pub use rzz::RZZGate;
 pub use u1::U1Gate;
 pub use u2::U2Gate;
 pub use u3::U3Gate;
+pub use rzsub::RZSubGate;
 pub use variable::VariableUnitaryGate;

src/ir/gates/parameterized/rzsub.rs

@@ -0,0 +1,88 @@
+use std::f64::consts::PI;
+
+use crate::i;
+use crate::ir::gates::utils::{rot_z, rot_z_jac};
+use crate::ir::gates::{Gradient, Size};
+use crate::ir::gates::{Optimize, Unitary};
+
+use ndarray::{Array2, Array3, ArrayViewMut2};
+use ndarray_linalg::c64;
+
+/// Arbitrary Y rotation single qubit gate
+#[derive(Copy, Clone, Debug, PartialEq, Default)]
+pub struct RZSubGate{
+    radix: usize,
+    level1: usize,
+    level2: usize,
+}
+
+impl RZSubGate {
+    pub fn new(radix: usize, level1: usize, level2: usize) -> Self {
+        RZSubGate {
+            radix: radix,
+            level1: level1,
+            level2: level2,
+        }
+    }
+}
+
+impl Unitary for RZSubGate {
+    fn num_params(&self) -> usize {
+        1
+    }
+
+    fn get_utry(&self, params: &[f64], _constant_gates: &[Array2<c64>]) -> Array2<c64> {
+        let pexp = i!(0.5 * params[0]).exp();
+        let nexp = i!(-0.5 * params[0]).exp();
+
+        let mut unitary = Array2::eye(self.radix);
+        unitary[[self.level1, self.level1]] = nexp;
+        unitary[[self.level2, self.level2]] = pexp;
+        unitary
+    }
+}
+
+impl Gradient for RZSubGate {
+    fn get_grad(&self, params: &[f64], _const_gates: &[Array2<c64>]) -> Array3<c64> {
+        let dpexp = i!(0.5) * i!(0.5 * params[0]).exp();
+        let dnexp = i!(-0.5) * i!(-0.5 * params[0]).exp();
+
+        let mut grad = Array3::zeros((1, self.radix, self.radix));
+        grad[[0, self.level1, self.level1]] = dnexp;
+        grad[[0, self.level2, self.level2]] = dpexp;
+        grad
+    }
+
+    fn get_utry_and_grad(
+        &self,
+        params: &[f64],
+        _const_gates: &[Array2<c64>],
+    ) -> (Array2<c64>, Array3<c64>) {
+        let pexp = i!(0.5 * params[0]).exp();
+        let nexp = i!(-0.5 * params[0]).exp();
+        let dpexp = i!(0.5) * pexp;
+        let dnexp = i!(-0.5) * nexp;
+
+        let mut unitary = Array2::eye(self.radix);
+        unitary[[self.level1, self.level1]] = nexp;
+        unitary[[self.level2, self.level2]] = pexp;
+
+        let mut grad = Array3::zeros((1, self.radix, self.radix));
+        grad[[0, self.level1, self.level1]] = dnexp;
+        grad[[0, self.level2, self.level2]] = dpexp;
+
+        (unitary, grad)
+    }
+}
+
+impl Size for RZSubGate {
+    fn num_qudits(&self) -> usize {
+        1
+    }
+}
+
+impl Optimize for RZSubGate {
+    fn optimize(&self, env_matrix: ArrayViewMut2<c64>) -> Vec<f64> {
+        unimplemented!()
+    }
+}

src/python/circuit.rs

@@ -36,37 +36,58 @@ fn pygate_to_native(pygate: &PyAny, constant_gates: &mut Vec<Array2<c64>>) -> Py
         "U2Gate" => Ok(U2Gate::new().into()),
         "U3Gate" => Ok(U3Gate::new().into()),
         "U8Gate" => Ok(U8Gate::new().into()),
+        "EmbeddedGate" => {
+            let egate = pygate.getattr("gate")?;
+            let egate_cls = egate.getattr("__class__")?;
+            let egate_dunder_name = egate_cls.getattr("__name__")?;
+            let egate_name = egate_dunder_name.extract::<&str>()?;
+
+            if egate_name == "RZGate" {
+                let level_maps = pygate.getattr("level_maps")?.extract::<Vec<Vec<usize>>>()?;
+                let level1 = level_maps[0][0];
+                let level2 = level_maps[0][1];
+                let radix = pygate.getattr("dim")?.extract::<usize>()?;
+                Ok(RZSubGate::new(radix, level1, level2).into())
+            } else {
+                extract_dynamic_gate(pygate, constant_gates, name)
+                // TODO: Generalize
+            }
+        },
         "VariableUnitaryGate" => {
             let size = pygate.getattr("num_qudits")?.extract::<usize>()?;
             let radixes = pygate.getattr("radixes")?.extract::<Vec<usize>>()?;
             Ok(VariableUnitaryGate::new(size, radixes).into())
-        }
+        },
         _ => {
-            if pygate.getattr("num_params")?.extract::<usize>()? == 0 {
-                let args: Vec<f64> = vec![];
-                let pyobj = pygate.call_method("get_unitary", (args,), None)?;
-                let pymat = pyobj.getattr("numpy")?.extract::<&PyArray2<c64>>()?;
-                let mat = pymat.to_owned_array();
-                let gate_size = pygate.getattr("num_qudits")?.extract::<usize>()?;
-                let index = constant_gates.len();
-                constant_gates.push(mat);
-                Ok(ConstantGate::new(index, gate_size).into())
-            } else if pygate.hasattr("get_unitary")?
-                && ((pygate.hasattr("get_grad")? && pygate.hasattr("get_unitary_and_grad")?)
-                    || pygate.hasattr("optimize")?)
-            {
-                let dynamic: Arc<dyn DynGate + Send + Sync> = Arc::new(PyGate::new(pygate.into()));
-                Ok(Gate::Dynamic(dynamic))
-            } else {
-                Err(exceptions::PyValueError::new_err(format!(
-                    "Gate {} does not implement the necessary methods for optimization.",
-                    name
-                )))
-            }
+            extract_dynamic_gate(pygate, constant_gates, name)
         }
     }
 }
 
+fn extract_dynamic_gate(pygate: &PyAny, constant_gates: &mut Vec<Array2<c64>>, name: &str) -> Result<Gate, PyErr> {
+    if pygate.getattr("num_params")?.extract::<usize>()? == 0 {
+        let args: Vec<f64> = vec![];
+        let pyobj = pygate.call_method("get_unitary", (args,), None)?;
+        let pymat = pyobj.getattr("numpy")?.extract::<&PyArray2<c64>>()?;
+        let mat = pymat.to_owned_array();
+        let gate_size = pygate.getattr("num_qudits")?.extract::<usize>()?;
+        let index = constant_gates.len();
+        constant_gates.push(mat);
+        Ok(ConstantGate::new(index, gate_size).into())
+    } else if pygate.hasattr("get_unitary")?
+        && ((pygate.hasattr("get_grad")? && pygate.hasattr("get_unitary_and_grad")?)
+            || pygate.hasattr("optimize")?)
+    {
+        let dynamic: Arc<dyn DynGate + Send + Sync> = Arc::new(PyGate::new(pygate.into()));
+        Ok(Gate::Dynamic(dynamic))
+    } else {
+        Err(exceptions::PyValueError::new_err(format!(
+            "Gate {} does not implement the necessary methods for optimization.",
+            name
+        )))
+    }
+}
+
 impl<'source> FromPyObject<'source> for Circuit {
     fn extract(ob: &'source PyAny) -> PyResult<Self> {
         let gil = Python::acquire_gil();

src/python/minimizers/residual_fn.rs

@@ -52,11 +52,14 @@ impl ResidualFn for PyResidualFn {
         let py = gil.python();
         let parameters = PyArray1::from_slice(py, params);
         let args = PyTuple::new(py, &[parameters]);
-        match self.cost_fn.call_method1(py, "get_cost", args) {
+        match self.cost_fn.call_method1(py, "get_residuals", args) {
             Ok(val) => val
                 .extract::<Vec<f64>>(py)
-                .expect("Return type of get_cost was not a sequence of floats."),
-            Err(..) => panic!("Failed to call 'get_cost' on passed ResidualFunction."), // TODO: make a Python exception?
+                .expect("Return type of get_residuals was not a sequence of floats."),
+            Err(e) => {
+                println!("{:?}, {:?}, {:?}", e.get_type(py), e.value(py), e.traceback(py));
+                panic!("Failed to call 'get_residuals' on passed ResidualFunction."); // TODO: make a Python exception?
+            },
         }
     }
 }