Merge pull request #11 from BQSKit/states

States
BQSKit · Aug 14, 2023 · 379dca0 · 379dca0
2 parents 707c12a + b7a63aa
commit 379dca0
Show file tree

Hide file tree

Showing 12 changed files with 427 additions and 52 deletions.
diff --git a/.github/workflows/build-wheels.yml b/.github/workflows/build-wheels.yml
@@ -19,11 +19,20 @@ jobs:
         with:
           submodules: 'recursive'
 
+      - uses: actions/setup-python@v4
+        with:
+          python-version: 3.8
+          architecture: x64
+
       - name: Build Manylinux Wheels
-        run: >
-          docker run -e OPENBLAS_ARGS="DYNAMIC_ARCH=1" --rm -v $(pwd):/io
-          edyounis/bqskitrs-manylinux:1.1 build  --release
-          --features=openblas --compatibility=manylinux2014
+        uses: PyO3/maturin-action@v1
+        with:
+          target: x86_64
+          manylinux: 2014
+          args: --release --features=openblas
+          docker-options: -e DYNAMIC_ARCH=1
+          before-script-linux: |
+            yum install -y cmake3 eigen3-devel llvm-toolset-7 openssl-devel
       
       - name: Upload Wheel Artifacts
         uses: actions/upload-artifact@v3
@@ -44,7 +53,7 @@ jobs:
       - name: Setup Python
         uses: actions/setup-python@v3
         with:
-          python-version: '3.10'
+          python-version: '3.11'
 
       - name: Install Python Dependencies
         run: pip install -U setuptools wheel maturin
@@ -89,7 +98,7 @@ jobs:
       - name: Setup Python
         uses: actions/setup-python@v3
         with:
-          python-version: '3.10'
+          python-version: '3.11'
 
       - name: Install Python Dependencies
         run: pip install -U setuptools wheel maturin
@@ -111,7 +120,7 @@ jobs:
         run: brew install gcc ceres-solver eigen lapack
 
       - name: Build macOS x86_64 Wheels
-        run: MACOSX_DEPLOYMENT_TARGET=11.0 maturin build --features="accelerate" --release --target x86_64-apple-darwin
+        run: MACOSX_DEPLOYMENT_TARGET=10.9 maturin build --features="accelerate" --release --target x86_64-apple-darwin
 
       - name: Build macOS aarch64 Wheels
         run: maturin build --features="accelerate" --release --target aarch64-apple-darwin

diff --git a/Cargo.toml b/Cargo.toml
@@ -29,7 +29,7 @@ mkl = [
 ndarray = { version = "0.15.6", features = ["blas"] }
 ndarray-linalg = "0.16.0"
 
-openblas-src = { version = "0.10.5", optional = true }
+openblas-src = { version = "0.10", optional = true }
 accelerate-src = { version = "0.3.2", optional = true }
 intel-mkl-src = { version = "0.8.1", optional = true }
 

diff --git a/Dockerfile b/Dockerfile
@@ -2,7 +2,7 @@ FROM quay.io/pypa/manylinux2014_x86_64
 # This docker file is based on the official maturin docker file https://github.com/PyO3/maturin/blob/master/Dockerfile
 ENV PATH /root/.cargo/bin:$PATH
 # Add all supported python versions
-ENV PATH /opt/python/cp38-cp38/bin/:/opt/python/cp39-cp39/bin/:/opt/python/cp310-cp310/bin/:$PATH
+ENV PATH /opt/python/cp38-cp38/bin/:/opt/python/cp39-cp39/bin/:/opt/python/cp310-cp310/bin/:/opt/python/cp311-cp311/bin/:$PATH
 # Otherwise `cargo new` errors
 ENV USER root
 
@@ -23,6 +23,10 @@ RUN yum install -y git cmake3 eigen3-devel llvm-toolset-7 && \
     yum clean all && \
     rm -rf /var/cache/yum
 
+RUN yum install -y openssl openssl-libs openssl-devel && \
+    yum clean all && \
+    rm -rf /var/cache/yum
+
 WORKDIR /io
 
 ENTRYPOINT ["/usr/bin/maturin"]
diff --git a/pyproject.toml b/pyproject.toml
@@ -1,14 +1,14 @@
 [build-system]
 requires = [
-    "maturin>=0.13,<0.14",
+    "maturin>=0.14,<0.15",
     "wheel",
     "oldest-supported-numpy",
 ]
 build-backend = "maturin"
 
 [project]
 name = "bqskitrs"
-version = "0.3.0"
+version = "0.4.0rc5"
 maintainers = [
     {name = "Ethan Smith", email = "[email protected]"},
     {name = "Ed Younis", email = "[email protected]"},

diff --git a/src/ir/circuit.rs b/src/ir/circuit.rs
@@ -4,7 +4,7 @@ use super::Operation;
 use itertools::Itertools;
 
 use itertools::izip;
-use ndarray::{Array2, Array3, ArrayD, ArrayView2, Ix2};
+use ndarray::{Array2, Array3, ArrayD, ArrayView2, Ix2, Array1};
 use ndarray_linalg::c64;
 use crate::squaremat::*;
 use crate::permutation_matrix::calc_permutation_matrix;
@@ -84,6 +84,30 @@ impl Circuit {
             param_idx += op.num_params();
         }
     }
+
+    pub fn get_state(&self, params: &[f64], const_gates: &[Array2<c64>]) -> Array1<c64> {
+        let mut zero = Array1::zeros((self.dim,));
+        zero[0] = c64::new(1.0, 0.0);
+        self.get_utry(params, const_gates).dot(&zero.view())
+        // TODO: Calculate state via Matrix-Vector vs Matrix-Matrix mul
+    }
+
+    pub fn get_state_and_grads(&self, params: &[f64], const_gates: &[Array2<c64>]) -> (Array1<c64>, Array2<c64>) {
+        let mut zero = Array1::zeros((self.dim,));
+        zero[0] = c64::new(1.0, 0.0);
+        let (utry, grads) = self.get_utry_and_grad(params, const_gates);
+        let state = utry.dot(&zero.view());
+
+        let mut out_grad = Array2::zeros((
+            params.len(),
+            self.dim,
+        ));
+
+        for (mut arr, grad) in out_grad.outer_iter_mut().zip(grads.outer_iter()) {
+            arr.assign(&grad.dot(&zero.view()));
+        }
+        (state, out_grad)
+    }
 }
 
 impl Unitary for Circuit {

diff --git a/src/ir/inst/minimizers/cost_fn.rs b/src/ir/inst/minimizers/cost_fn.rs
@@ -1,11 +1,11 @@
 use crate::{
     ir::circuit::Circuit,
     ir::gates::{Gradient, Unitary},
-    utils::{matrix_distance_squared, matrix_distance_squared_jac},
+    utils::{matrix_distance_squared, matrix_distance_squared_jac, state_infidelity, state_infidelity_jac, matrix_distance_system_squared, matrix_distance_system_squared_jac},
 };
 
 use enum_dispatch::enum_dispatch;
-use ndarray::Array2;
+use ndarray::{Array2, Array1};
 use ndarray_linalg::c64;
 
 /// Trait defining the signature of a cost function used by minimizers.
@@ -69,15 +69,98 @@ impl DifferentiableCostFn for HilbertSchmidtCostFn {
     }
 }
 
+#[derive(Clone)]
+pub struct HilbertSchmidtStateCostFn {
+    circ: Circuit,
+    target: Array1<c64>,
+}
+
+impl HilbertSchmidtStateCostFn {
+    pub fn new(circ: Circuit, target: Array1<c64>) -> Self {
+        HilbertSchmidtStateCostFn { circ, target }
+    }
+
+    pub fn is_sendable(&self) -> bool {
+        self.circ.is_sendable()
+    }
+}
+
+impl CostFn for HilbertSchmidtStateCostFn {
+    fn get_cost(&self, params: &[f64]) -> f64 {
+        let calculated = self.circ.get_state(params, &self.circ.constant_gates);
+        state_infidelity(self.target.view(), calculated.view())
+    }
+}
+
+impl DifferentiableCostFn for HilbertSchmidtStateCostFn {
+    fn get_grad(&self, params: &[f64]) -> Vec<f64> {
+        let (m, j) = self
+            .circ
+            .get_state_and_grads(params, &self.circ.constant_gates);
+        state_infidelity_jac(self.target.view(), m.view(), j.view()).1
+    }
+
+    fn get_cost_and_grad(&self, params: &[f64]) -> (f64, Vec<f64>) {
+        let (m, j) = self
+            .circ
+            .get_state_and_grads(params, &self.circ.constant_gates);
+        state_infidelity_jac(self.target.view(), m.view(), j.view())
+    }
+}
+
+#[derive(Clone)]
+pub struct HilbertSchmidtSystemCostFn {
+    circ: Circuit,
+    target: Array2<c64>,
+    vec_count: u32,
+}
+
+impl HilbertSchmidtSystemCostFn {
+    pub fn new(circ: Circuit, target: Array2<c64>, vec_count: u32) -> Self {
+        HilbertSchmidtSystemCostFn { circ, target, vec_count }
+    }
+
+    pub fn is_sendable(&self) -> bool {
+        self.circ.is_sendable()
+    }
+}
+
+impl CostFn for HilbertSchmidtSystemCostFn {
+    fn get_cost(&self, params: &[f64]) -> f64 {
+        let calculated = self.circ.get_utry(params, &self.circ.constant_gates);
+        matrix_distance_system_squared(self.target.view(), calculated.view(), self.vec_count)
+    }
+}
+
+impl DifferentiableCostFn for HilbertSchmidtSystemCostFn {
+    fn get_grad(&self, params: &[f64]) -> Vec<f64> {
+        let (m, j) = self
+            .circ
+            .get_utry_and_grad(params, &self.circ.constant_gates);
+        matrix_distance_system_squared_jac(self.target.view(), m.view(), j.view(), self.vec_count).1
+    }
+
+    fn get_cost_and_grad(&self, params: &[f64]) -> (f64, Vec<f64>) {
+        let (m, j) = self
+            .circ
+            .get_utry_and_grad(params, &self.circ.constant_gates);
+        matrix_distance_system_squared_jac(self.target.view(), m.view(), j.view(), self.vec_count)
+    }
+}
+
 pub enum CostFunction {
     HilbertSchmidt(HilbertSchmidtCostFn),
+    HilbertSchmidtState(HilbertSchmidtStateCostFn),
+    HilbertSchmidtSystem(HilbertSchmidtSystemCostFn),
     Dynamic(Box<dyn DifferentiableCostFn>),
 }
 
 impl CostFunction {
     pub fn is_sendable(&self) -> bool {
         match self {
             CostFunction::HilbertSchmidt(hs) => hs.is_sendable(),
+            CostFunction::HilbertSchmidtState(hs) => hs.is_sendable(),
+            CostFunction::HilbertSchmidtSystem(hs) => hs.is_sendable(),
             CostFunction::Dynamic(_) => false,
         }
     }
@@ -87,6 +170,8 @@ impl CostFn for CostFunction {
     fn get_cost(&self, params: &[f64]) -> f64 {
         match self {
             Self::HilbertSchmidt(hs) => hs.get_cost(params),
+            Self::HilbertSchmidtState(hs) => hs.get_cost(params),
+            Self::HilbertSchmidtSystem(hs) => hs.get_cost(params),
             Self::Dynamic(d) => d.get_cost(params),
         }
     }
@@ -96,13 +181,17 @@ impl DifferentiableCostFn for CostFunction {
     fn get_grad(&self, params: &[f64]) -> Vec<f64> {
         match self {
             Self::HilbertSchmidt(hs) => hs.get_grad(params),
+            Self::HilbertSchmidtState(hs) => hs.get_grad(params),
+            Self::HilbertSchmidtSystem(hs) => hs.get_grad(params),
             Self::Dynamic(d) => d.get_grad(params),
         }
     }
 
     fn get_cost_and_grad(&self, params: &[f64]) -> (f64, Vec<f64>) {
         match self {
             Self::HilbertSchmidt(hs) => hs.get_cost_and_grad(params),
+            Self::HilbertSchmidtState(hs) => hs.get_cost_and_grad(params),
+            Self::HilbertSchmidtSystem(hs) => hs.get_cost_and_grad(params),
             Self::Dynamic(d) => d.get_cost_and_grad(params),
         }
     }