Autograd with param shift demo

.github/workflows/unittest.yml

@@ -0,0 +1,28 @@
+name: Unit test
+on:
+  push:
+  pull_request:
+    branches: ['main']
+jobs:
+  unnittest:
+    name: Unit test
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v3
+      - uses: actions/setup-python@v4
+        name: Install Python
+        with:
+          python-version: '3.10'
+
+      - name: Install dependency
+        run: python -m pip install -r requirements.txt && python -m pip install pytest
+
+      # TODO(zhaoyilun): Build seperate package for pyquafu-torch
+      - name: Install torch
+        run: python -m pip install torch torchvision torchaudio
+
+      - name: Install pyquafu
+        run: python -m pip install .
+
+      - name: Run unit tests
+        run: pytest tests/

.github/workflows/wheels.yml

@@ -48,9 +48,6 @@ jobs:
       - name: Build wheels
         run: python -m cibuildwheel --output-dir dist
 
-      - name: Run unit tests
-        run: pip install . && pytest tests/
-
       - name: Publish package
         run: python -m twine upload dist/*.whl
         if: ${{ contains(github.ref, '/tags/') }}

quafu/algorithms/__init__.py

@@ -1,5 +1,5 @@
 """Algorithm module"""
 
-from .hamiltonian import *
+from .hamiltonian import Hamiltonian
 from .ansatz import *
 from .estimator import *

quafu/algorithms/estimator.py

@@ -21,6 +21,17 @@
 from quafu.algorithms.hamiltonian import Hamiltonian
 
 
+def execute_circuit(circ: QuantumCircuit, observables: Hamiltonian):
+    """Execute circuit on quafu simulator"""
+    sim_state = simulate(circ, output="state_vector").get_statevector()
+
+    expectation = np.matmul(
+        np.matmul(sim_state.conj().T, observables.get_matrix()), sim_state
+    ).real
+
+    return expectation
+
+
 class Estimator:
     """Estimate expectation for quantum circuits and observables"""
 
@@ -62,11 +73,12 @@ def _run_real_machine(self, observables: Hamiltonian):
 
     def _run_simulation(self, observables: Hamiltonian):
         """Run using quafu simulator"""
-        sim_state = simulate(self._circ, output="state_vector").get_statevector()
-        expectation = np.matmul(
-            np.matmul(sim_state.conj().T, observables.get_matrix()), sim_state
-        ).real
-        return expectation
+        # sim_state = simulate(self._circ, output="state_vector").get_statevector()
+        # expectation = np.matmul(
+        #     np.matmul(sim_state.conj().T, observables.get_matrix()), sim_state
+        # ).real
+        # return expectation
+        return execute_circuit(self._circ, observables)
 
     def run(self, observables: Hamiltonian, params: List[float]):
         """Calculate estimation for given observables

quafu/algorithms/gradients/__init__.py

@@ -0,0 +1,15 @@
+# (C) Copyright 2023 Beijing Academy of Quantum Information Sciences
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from .param_shift import ParamShift

quafu/algorithms/optimizer.py → quafu/algorithms/gradients/param_shift.py

@@ -11,9 +11,11 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
+"""Quafu parameter shift"""
 
-import numpy as np
 from typing import List
+import numpy as np
+
 from quafu.algorithms import Estimator
 from quafu.algorithms.hamiltonian import Hamiltonian
 
@@ -31,7 +33,7 @@ def __call__(self, obs: Hamiltonian, params: List[float]):
             estimator (Estimator): estimator to calculate expectation values
             params (List[float]): params to optimize
         """
-        return self._grad(obs, params)
+        return self.grad(obs, params)
 
     def _gen_param_shift_vals(self, params):
         """Given a param list with n values, replicate to 2*n param list"""
@@ -42,13 +44,19 @@ def _gen_param_shift_vals(self, params):
         minus_params = params - offsets * np.pi / 2
         return plus_params.tolist() + minus_params.tolist()
 
-    def _grad(self, obs: Hamiltonian, params: List[float]):
+    def grad(self, obs: Hamiltonian, params: List[float]):
+        """grad.
+
+        Args:
+            obs (Hamiltonian): obs
+            params (List[float]): params
+        """
         shifted_params_lists = self._gen_param_shift_vals(params)
 
         res = np.zeros(len(shifted_params_lists))
         for i, shifted_params in enumerate(shifted_params_lists):
             res[i] = self._est.run(obs, shifted_params)
 
-        n = len(res)
-        grads = (res[: n // 2] - res[n // 2 :]) / 2
+        num_shift_params = len(res)
+        grads = (res[: num_shift_params // 2] - res[num_shift_params // 2 :]) / 2
         return grads

quafu/algorithms/hamiltonian.py

@@ -66,7 +66,6 @@ def from_pauli_list(pauli_list: Iterable[tuple[str, complex]]) -> Hamiltonian:
         if size == 0:
             raise QuafuError("Pauli list cannot be empty.")
 
-        num_qubits = len(pauli_list[0][0])
         coeffs = np.zeros(size, dtype=complex)
 
         pauli_str_list = []
@@ -81,31 +80,29 @@ def to_legacy_quafu_pauli_list(self):
         """Transform to legacy quafu pauli list format,
         this is a temperal function and should be deleted later"""
         res = []
-        for pauli in self._pauli_str_list:
-            for i, p in enumerate(pauli[::-1]):
-                if p in ["X", "Y", "Z"]:
-                    res.append([p, [i]])
+        for pauli_str in self._pauli_str_list:
+            for i, pauli in enumerate(pauli_str[::-1]):
+                if pauli in ["X", "Y", "Z"]:
+                    res.append([pauli, [i]])
         return res
 
     def _get_pauli_mat(self, pauli_str: str):
         """Calculate the matrix of a pauli string"""
         mat = None
-        for p in pauli_str[::-1]:
-            mat = PAULI_MAT[p] if mat is None else np.kron(PAULI_MAT[p], mat)
+        for pauli in pauli_str[::-1]:
+            mat = PAULI_MAT[pauli] if mat is None else np.kron(PAULI_MAT[pauli], mat)
         return mat
 
     def matrix_generator(self):
         """Generating matrix for each Pauli str"""
-        for i, p in enumerate(self._pauli_str_list):
-            yield self._coeffs[i] * self._get_pauli_mat(p)
+        for i, pauli_str in enumerate(self._pauli_str_list):
+            yield self._coeffs[i] * self._get_pauli_mat(pauli_str)
 
     def get_matrix(self):
         """Generate matrix of Hamiltonian"""
 
-        mat = None
-        for m in self.matrix_generator():
-            if mat is None:
-                mat = m
-                continue
-            mat += m
-        return mat
+        dim = 2**self.num_qubits
+        matrix = np.zeros((dim, dim), dtype=complex)
+        for mat in self.matrix_generator():
+            matrix += mat
+        return matrix

quafu/algorithms/layers/__init__.py

@@ -0,0 +1 @@
+from .qnode import compute_vjp, jacobian

quafu/algorithms/layers/library.py

@@ -0,0 +1,15 @@
+# (C) Copyright 2023 Beijing Academy of Quantum Information Sciences
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Pre-built quafu circuit blocks"""

quafu/algorithms/layers/qnode.py

@@ -0,0 +1,86 @@
+# (C) Copyright 2023 Beijing Academy of Quantum Information Sciences
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+from typing import List, Optional
+import numpy as np
+from quafu import QuantumCircuit
+from quafu.algorithms import Hamiltonian
+from quafu.algorithms.estimator import Estimator
+from quafu.algorithms.gradients import ParamShift
+
+
+def _generate_expval_z(num_qubits: int):
+    obs_list = []
+    base_pauli = "I" * num_qubits
+    for i in range(num_qubits):
+        pauli = base_pauli[:i] + "Z" + base_pauli[i + 1 :]
+        obs_list.append(Hamiltonian.from_pauli_list([(pauli, 1)]))
+    return obs_list
+
+
+# TODO(zhaoyilun): support more measurement types
+def run_circ(circ: QuantumCircuit, params: Optional[List[float]] = None):
+    """Execute a circuit
+
+    Args:
+        circ (QuantumCircuit): circ
+        params (Optional[List[float]]): params
+    """
+    obs_list = _generate_expval_z(circ.num)
+    estimator = Estimator(circ)
+    if params is None:
+        params = [g.paras for g in circ.parameterized_gates]
+    output = [estimator.run(obs, params) for obs in obs_list]
+    return np.array(output)
+
+
+# TODO(zhaoyilun): support more gradient methods
+def jacobian(circ: QuantumCircuit, params_input: np.ndarray):
+    """Calculate Jacobian matrix
+
+    Args:
+        circ (QuantumCircuit): circ
+        params_input (np.ndarray): params_input, with shape [batch_size, num_params]
+    """
+    batch_size, num_params = params_input.shape
+    obs_list = _generate_expval_z(circ.num)
+    num_outputs = len(obs_list)
+    estimator = Estimator(circ)
+    calc_grad = ParamShift(estimator)
+    output = np.zeros((batch_size, num_outputs, num_params))
+    for i in range(batch_size):
+        grad_list = [
+            np.array(calc_grad(obs, params_input[i, :].tolist())) for obs in obs_list
+        ]
+        output[i, :, :] = np.stack(grad_list)
+    return output
+
+
+def compute_vjp(jac: np.ndarray, dy: np.ndarray):
+    """compute vector-jacobian product
+
+    Args:
+        jac (np.ndarray): jac with shape (batch_size, num_outputs, num_params)
+        dy (np.ndarray): dy with shape (batch_size, num_outputs)
+    """
+    batch_size, num_outputs, num_params = jac.shape
+    assert dy.shape[0] == batch_size and dy.shape[1] == num_outputs
+
+    vjp = np.zeros((batch_size, num_params))
+
+    for i in range(batch_size):
+        vjp[i] = dy[i, :].T @ jac[i, :, :]
+
+    return vjp

quafu/algorithms/layers/torch.py

@@ -0,0 +1,61 @@
+# (C) Copyright 2023 Beijing Academy of Quantum Information Sciences
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""quafu PyTorch quantum layer"""
+
+import torch
+import numpy as np
+from quafu import QuantumCircuit
+from quafu.algorithms.layers.qnode import compute_vjp, jacobian
+
+
+class ExecuteCircuits(torch.autograd.Function):
+    """TODO(zhaoyilun): document"""
+
+    @staticmethod
+    def forward(ctx, parameters, kwargs):
+        ctx.run_fn = kwargs["run_fn"]
+        ctx.circ = kwargs["circ"]
+        ctx.save_for_backward(parameters)
+        parameters = parameters.numpy().tolist()
+        outputs = []
+        for para in parameters:
+            out = ctx.run_fn(ctx.circ, para)
+            outputs.append(out)
+        outputs = np.stack(outputs)
+        outputs = torch.from_numpy(outputs)
+        return outputs
+
+    @staticmethod
+    def backward(ctx, grad_out):
+        (parameters,) = ctx.saved_tensors
+        jac = jacobian(ctx.circ, parameters.numpy())
+        vjp = compute_vjp(jac, grad_out.numpy())
+        vjp = torch.from_numpy(vjp)
+        return vjp, None
+
+
+# TODO(zhaoyilun): doc
+def execute(circ: QuantumCircuit, run_fn, grad_fn, parameters: torch.Tensor):
+    """execute.
+
+    Args:
+        circ:
+        run_fn:
+        grad_fn:
+    """
+
+    kwargs = {"circ": circ, "run_fn": run_fn, "grad_fn": grad_fn}
+
+    return ExecuteCircuits.apply(parameters, kwargs)

quafu/simulators/torch.py

@@ -0,0 +1,15 @@
+# (C) Copyright 2023 Beijing Academy of Quantum Information Sciences
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Simulate the execution of a quantum circuit using pytorch"""

tests/quafu/algorithms/optimizer_test.py → tests/quafu/algorithms/gradient_test.py

@@ -16,7 +16,7 @@
 import sys
 from quafu.algorithms.estimator import Estimator
 from quafu.algorithms.hamiltonian import Hamiltonian
-from quafu.algorithms.optimizer import ParamShift
+from quafu.algorithms.gradients import ParamShift
 from quafu.circuits.quantum_circuit import QuantumCircuit