Repository: qutip/qutip-qip (language: python)

Issue:
Add Variational Quantum Algorithms Module
QuTiP currently lacks any functionality directly related to VQAs. It would be useful to have a class-based implementation that allows for problem specification down to the details of parametrized/fixed Hamiltonians and the optimisation method. 

To outline a possible structure for this module, I have implemented something along these lines in [my fork of this repository](https://github.com/EnBr55/qutip-qip/blob/master/src/qutip_qip/vqa.py). This includes:

1. `Optimization_Result` class
  - A convenient interface for results of the VQA optimisation method.
2. `VQA_Block` class
  - A kind of pseudo-circuit-element. Holds one of either a Hamiltonian, a Unitary, or a pre-defined QuTiP gate (e.g. `qutip_qip.operations.x_gate`).
  - Takes a list of qubit targets
  - Can be either an "initial" element or not - this will determine if the block is repeated in the circuit layers.
  - In the case that a Hamiltonian or Unitary is specified, the action of the block is implemented by adding a new entry to the `QubitCircuit`'s `user_gates` dictionary.
3. `VQA` class
  - Initialised with `n_qubits`, `n_layers`
  - Can specify a cost function in terms of:
    - An observable
    - A state
    - A bitstring (single measurement outcome sampled from the circuit) 
  - `add_block(block: VQA_Block)` appends `VQA_Block` to the circuit and, if necessary, updates the circuit's `user_gates` dictionary to include the new gate.
  - `optimise_parameters` method performs optimisation on the free parameters of the circuit and returns an `Optimization_Result` instance.
  - Methods that serve as a wrapper to `QubitCircuit` to simulate the circuit

A usage example can be found at https://github.com/EnBr55/qutip-vqa-examples/blob/main/notebooks/QAOA.ipynb.

I'm currently working on a `Parameterized_Hamiltonian` model as a `VQA_Block` input that will allow for defining more general quantum control problems. 

Any discussion/suggestions as to these proposed models would be great!


Discussion / hints:

Fix patch:
diff --git a/doc/circ.png b/doc/circ.png
new file mode 100644
index 00000000..7ccb7c06
Binary files /dev/null and b/doc/circ.png differ
diff --git a/doc/source/apidoc.rst b/doc/source/apidoc.rst
index 0f06e8c0..99c0caab 100644
--- a/doc/source/apidoc.rst
+++ b/doc/source/apidoc.rst
@@ -19,6 +19,7 @@ Simulation based on operator-state multiplication.
    qutip_qip.qubits
    qutip_qip.decompose
    qutip_qip.qasm
+   qutip_qip.vqa
 
 Pulse-level simulation
 ----------------------
diff --git a/doc/source/apidoc/qutip_qip.vqa.rst b/doc/source/apidoc/qutip_qip.vqa.rst
new file mode 100644
index 00000000..857c6812
--- /dev/null
+++ b/doc/source/apidoc/qutip_qip.vqa.rst
@@ -0,0 +1,26 @@
+﻿qutip\_qip.vqa
+==============
+
+.. automodule:: qutip_qip.vqa
+   :members:
+   :show-inheritance:
+
+   
+   
+   
+
+   
+   
+   .. rubric:: Classes
+
+   .. autosummary::
+   
+      OptimizationResult
+      ParameterizedHamiltonian
+      VQA
+      VQABlock
+   
+   
+
+   
+   
diff --git a/doc/source/figures/vqa_circuit_with_x.png b/doc/source/figures/vqa_circuit_with_x.png
new file mode 100644
index 00000000..aabb0cfc
Binary files /dev/null and b/doc/source/figures/vqa_circuit_with_x.png differ
diff --git a/doc/source/index.rst b/doc/source/index.rst
index f8176fdb..cce967fe 100644
--- a/doc/source/index.rst
+++ b/doc/source/index.rst
@@ -22,6 +22,7 @@ qutip-qip |version|: QuTiP quantum information processing
     qip-basics.rst
     qip-simulator.rst
     qip-processor.rst
+    qip-vqa.rst
 
 .. toctree::
     :maxdepth: 2
diff --git a/doc/source/qip-vqa.rst b/doc/source/qip-vqa.rst
new file mode 100644
index 00000000..660a2a2b
--- /dev/null
+++ b/doc/source/qip-vqa.rst
@@ -0,0 +1,124 @@
+.. _qip_vqa:
+
+******************************
+Variational Quantum Algorithms
+******************************
+
+Implemented by `Ben Braham <https://benbraham.com>`_ as part of a `Unitary Fund microgrant <https://unitary.fund/grants.html>`_.
+
+Overview
+========
+
+Variational Quantum Algorithms (VQAs) are a hybrid quantum-classical optimization algorithm in which an objective function (usually encoded by a parameterized quantum circuit) is evaluated by quantum computation, and the parameters of this function are updated using classical optimization methods. Such algorithms have been proposed for use in NISQ-era quantum computers as they typically scale well with the number of available qubits, and can function without high fault-tolerance.
+
+In QuTiP, VQAs are represented by a parameterized quantum circuit, and include methods for defining a cost function for the circuit, and finding parameters that minimize this cost.
+
+
+Constructing a VQA circuit
+==========================
+
+The :class:`.VQA` class allows for the construction of a parameterized circuit from :class:`.VQABlock` instances, which act as the gates of the circuit. In the most basic instance, a :class:`.VQA` should have:
+
+====================  ==================================================
+Property                           Description
+====================  ==================================================
+``num_qubits``        Positive integer number of qubits for the circuit.
+``num_layers``        Positive integer number of repetitions of the 
+                      layered elements of the circuit.
+``cost_method``       String referring to the method used to
+                      evaluate the circuit's cost.
+
+                      Either "OBSERVABLE", "BITSTRING", or "STATE".
+====================  ==================================================
+
+For example:
+
+.. code-block::
+
+    from qutip_qip.vqa import VQA
+
+    VQA_circuit = VQA(
+                num_qubits=1,
+                num_layers=1,
+                cost_method="OBSERVABLE",
+            )
+
+After constructing this instance, we are ready to begin adding elements to our parameterized circuit. Circuit elements in this module are represented by :class:`.VQABlock` instances. Fundamentally, the role of this class is to generate an operator for the circuit. To do this, it keeps track of its free parameters, and gives information to the :class:`.VQA` instance on how to update them. The operator itself can be generated by a user-defined function call, a parameterized Hamiltonian to exponentiate, a pre-computed unitary operator, or a string referring to a gate that has already been defined (either by the user already, or a gate native to QuTiP).
+
+In the absence of specification, a VQA block takes a :class:`~.Qobj` as the Hamiltonian :math:`H`, and will generate a unitary with free parameter, :math:`\theta`, as :math:`U(\theta) = e^{-i \theta H}`. For example, 
+
+.. testcode::
+
+    from qutip_qip.vqa import VQA, VQABlock
+    from qutip import tensor, sigmax
+
+    VQA_circuit = VQA(num_qubits=1, num_layers=1)
+
+    R_x_block = VQABlock(
+      sigmax() / 2, name="R_x(\\theta)"
+    )
+
+    VQA_circuit.add_block(R_x_block)
+
+
+We added our block to the ``VQA_circuit`` with the :meth:`.VQA.add_block` method. Calling the :meth:`.VQA.export_image` method renders an image of our circuit in its current form:
+
+.. image:: /figures/vqa_circuit_with_x.png
+   
+
+--------------
+
+
+Optimisation Loop
+=================
+
+After specifying a cost method and function to the :class:`.VQA` instance, there are various options for optimization of the free circuit parameters. Calling :meth:`.VQA.optimize_parameters` will begin the optimization process and return an :class:`.OptimizationResult` instance. By default, the method will randomize initial parameters, and use the non-gradient-based ``COBYLA`` method for parameter optimization. Users can specify:
+
+
+  * **Initial parameters**. Given as a list, with length corresponding to the number of free parameters in the circuit. The number of free parameters can be computed automatically with the :meth:`.VQA.get_free_parameters_num` method. Alternatively, the string 'zeros' will initialize all parameters as 0; and 'random' will initialize parameters randomly between 0 and 1. Defaults to 'random'.
+
+  * **Optimization method**. This can be a string referring to a pre-defined ``SciPy`` method `listed here <https://docs.scipy.org/doc/scipy/reference/generated/scipy.optimize.minimize.html>`_, or a callable function.
+
+  * **Jacobian computation**. A flag will tell the optimization method to compute the Jacobian at each step, which is passed to the optimizer so that it can use gradient information.
+
+  * **Layer-by-layer training**. Optimize parameters for the circuit with only a single layer, and hold these fixed while adding additional layers, up to ``VQA.num_layers``.
+
+  * **Bounds and constraints**. To be passed to the optimizer.
+
+The :class:`.OptimizationResult` class provides information about the completed optimization process. For example, the probability amplitudes of different measurement outcomes of the circuit post-optimization can be plotted with :meth:`.OptimizationResult.plot`.
+
+Below, we run an optimization on a toy circuit, tuning a parameterized :math:`x`-rotation gate to try to maximise the probability amplitude of the :math:`|1\rangle` state.
+
+.. plot::
+  :context:
+
+  >>> from qutip_qip.vqa import VQA, VQABlock
+  >>> from qutip import sigmax, sigmaz
+  >>> circ = VQA(num_qubits=1, cost_method="OBSERVABLE")
+
+  Picking the Pauli Z operator as our cost observable, our circuit's cost function will be: :math:`\langle\psi(t)| \sigma_z | \psi(t)\rangle`
+
+  >>> circ.cost_observable = sigmaz()
+
+
+  Adding a Pauli X operator as a block to the circuit, the operation of the entire circuit becomes: :math:`e^{-i t X /2}`.
+
+
+  >>> circ.add_block(VQABlock(sigmax() / 2))
+
+  We can now try to find a minimum in our cost function using the SciPy in-built L-BFGS-B (L-BFGS with box constraints) method. We specify the bounds so that our parameter is :math:`0 \leq t \leq 4`.
+
+  >>> result = circ.optimize_parameters(method="L-BFGS-B", use_jac=True, bounds=[[0, 4]])
+
+  Accessing ``result.res.x``, we have the array of parameters found during optimization. In our case, we only had one free parameter, so we examine the first element of this array.
+
+  >>> angle = round(result.res.x[0], 2)
+  >>> print(f"Angle found: {angle}")
+  Angle found: 3.14
+
+  Finally, we can plot our the measurement outcome probabilities of our circuit after optimization.
+
+  >>> result.plot()
+
+
+In this simple example, our optimization found that (neglecting phase) :math:`R_x(\pi) |0\rangle = |1\rangle`. Of course, this very basic usage generalizes to circuits on multiple qubits, with more complicated cost functions and optimization procedures.
diff --git a/src/qutip_qip/vqa.py b/src/qutip_qip/vqa.py
new file mode 100644
index 00000000..761b4943
--- /dev/null
+++ b/src/qutip_qip/vqa.py
@@ -0,0 +1,822 @@
+"""Variational Quantum Algorithms generation and optimization"""
+import types
+import random
+import numpy as np
+from qutip import basis, tensor, Qobj, qeye, expect
+from qutip.qip.circuit import QubitCircuit
+from scipy.optimize import minimize
+from scipy.linalg import expm_frechet
+from qutip.qip.operations.gates import gate_sequence_product
+
+
+class VQA:
+    """
+    Optimizes free parameters to generate :obj:`.QubitCircuit` instances
+    based on Variational Quantum Algorithms.
+    Accepts :obj:`.VQABlock` elements instead of :obj:`.Gate` elements,
+    which allows for easy parameterization of user-defined circuit elements.
+    Includes methods for parameter optimization and generators of
+    :obj:`.QubitCircuit` instances.
+
+    Parameters
+    ----------
+    num_qubits: int
+        number of qubits used by the algorithm
+    num_layers: int, optional
+        number of layers used by the algorihtm
+    cost_method: str
+        method used to compute the cost of an instance of the circuit
+        constructed by fixing its free parameters. Can be one of `OBSERVABLE`,
+        `BITSTRING` or `STATE`.
+
+        #.  If `OBSERVABLE` is set, then the attribute
+            ``VQA.cost_observable`` needs to be specified as a ``Qobj``.
+            The cost of the circuit is the expectation value of this observable
+            in the final state.
+        #.  If `STATE` is set, then ``VQA.cost_func`` needs to be specified
+            as a callable that takes in a quantum state, as a ``Qobj``, and
+            returns a float.
+        #.  If `BITSTRING` is set, then ``VQA.cost_func`` needs to be
+            specified as a callable that takes in a bitstring and returns
+            a float.
+    """
+
+    def __init__(self, num_qubits, num_layers=1, cost_method="OBSERVABLE"):
+        self.num_qubits = num_qubits
+        self.num_layers = num_layers
+        self.blocks = []
+        self.user_gates = {}
+        self._cost_methods = ["OBSERVABLE", "STATE", "BITSTRING"]
+        self.cost_method = cost_method
+        self.cost_func = None
+        self.cost_observable = None
+
+        if self.num_qubits < 1:
+            raise ValueError("Expected 1 or more qubits")
+        if not isinstance(self.num_qubits, int):
+            raise TypeError("Expected an integer number of qubits")
+        if self.num_layers < 1:
+            raise ValueError("Expected 1 or more layer")
+        if not isinstance(self.num_layers, int):
+            raise TypeError("Expected an integer number of layers")
+        if self.cost_method not in self._cost_methods:
+            raise ValueError(
+                f"Cost method {self.cost_method} not one of "
+                f"{self._cost_methods}"
+            )
+
+    def get_block_series(self):
+        """
+        Ordered list of circuit blocks, including layer repetitions,
+        from first applied to last applied.
+        """
+        blocks = [*self.blocks]
+        for _ in range(1, self.num_layers):
+            for block in list(filter(lambda b: not b.initial, self.blocks)):
+                blocks.append(block)
+        return blocks
+
+    def add_block(self, block):
+        """
+        Append a :obj:`.VQABlock` instance to the circuit, and update the
+        user_gates dictionary if necessary.
+
+        Parameters
+        ----------
+        block: VQABlock
+        """
+        if not block.name:
+            block.name = "U" + str(len(self.blocks))
+        if block.name in list(map(lambda b: b.name, self.blocks)):
+            raise ValueError("Duplicate Block name in blocks dict")
+        self.blocks.append(block)
+        self.user_gates[block.name] = lambda angles=None: block.get_unitary(
+            angles
+        )
+
+    def get_free_parameters_num(self):
+        """
+        Compute the number of free parameters required
+        to evaluate the circuit.
+
+        Returns
+        -------
+        num_params: int
+            Number of free circuit parameters
+        """
+        initial_blocks = list(filter(lambda b: b.initial, self.blocks))
+        layer_blocks = list(filter(lambda b: not b.initial, self.blocks))
+
+        n_initial_params = sum(
+            list(map(lambda b: b.get_free_parameters_num(), initial_blocks))
+        )
+
+        n_layer_params = (
+            sum(list(map(lambda b: b.get_free_parameters_num(), layer_blocks)))
+            * self.num_layers
+        )
+
+        return n_initial_params + n_layer_params
+
+    def construct_circuit(self, angles):
+        """
+        Construct a circuit by specifying values for each
+        free parameter.
+
+        Parameters
+        ----------
+        angles: list of float
+            A list of dimension (n,) for n free parameters in the circuit
+
+        Returns
+        -------
+        circ: :obj:`.QubitCircuit`
+        """
+        circ = QubitCircuit(self.num_qubits)
+        circ.user_gates = self.user_gates
+        i = 0
+        for layer_num in range(self.num_layers):
+            for block in self.blocks:
+                if block.initial and layer_num > 0:
+                    continue
+                if block.is_native_gate:
+                    circ.add_gate(block.operator, targets=block.targets)
+                else:
+                    n = block.get_free_parameters_num()
+                    circ.add_gate(
+                        block.name,
+                        targets=list(range(self.num_qubits)),
+                        arg_value=angles[i : i + n] if n > 0 else None,
+                    )
+                    i += n
+        return circ
+
+    def get_initial_state(self):
+        """
+        Returns
+        -------
+        initial_state: :obj:`qutip.Qobj`
+            Initial circuit state
+        """
+        initial_state = basis(2, 0)
+        for _ in range(self.num_qubits - 1):
+            initial_state = tensor(initial_state, basis(2, 0))
+        return initial_state
+
+    def get_final_state(self, angles):
+        """
+        Evaluate the circuit by specifying each circuit parameter
+
+        Parameters
+        ----------
+        angles: list of float
+            A list of dimension (n,) for n free parameters in the circuit
+
+        Returns
+        -------
+        final_state: :obj:`qutip.Qobj`.
+            Final state of the circuit after evaluation
+        """
+        circ = self.construct_circuit(angles)
+        initial_state = self.get_initial_state()
+        final_state = circ.run(initial_state)
+        return final_state
+
+    def _sample_bitstring_from_state(self, state):
+        """
+        Use probability amplitudes from state after measurement
+        in computational basis to sample a bitstring.
+
+        Parameters
+        ----------
+        state: :obj:`qutip.Qobj`
+
+        Returns
+        -------
+        bitstring: str
+            Formatted binary string with length corresponding to
+            dimension of input state
+
+        E.g. the state 1/sqrt(2) * (|0> + |1>)
+        would return 0 and 1 with equal probability.
+        """
+        num_qubits = int(np.log2(state.shape[0]))
+        outcome_indices = list(range(2**num_qubits))
+        probs = [abs(i.item()) ** 2 for i in state]
+        outcome_index = np.random.choice(outcome_indices, p=probs)
+        return format(outcome_index, f"0{num_qubits}b")
+
+    def evaluate_parameters(self, angles):
+        """
+        Evaluate a cost for the circuit, based on the VQA cost
+        method defined.
+
+        Parameters
+        ----------
+        angles: list of float
+            A list of dimension (n,) for n free parameters in the circuit
+
+        Returns
+        -------
+        cost: float
+        """
+        final_state = self.get_final_state(angles)
+        if self.cost_method == "BITSTRING":
+            if self.cost_func is None:
+                raise ValueError(
+                    "To use BITSTRING as the cost method, please"
+                    " specify the attribute cost_func"
+                )
+            else:
+                return self.cost_func(
+                    self._sample_bitstring_from_state(final_state)
+                )
+        elif self.cost_method == "STATE":
+            if self.cost_func is None:
+                raise ValueError(
+                    "To use STATE as the cost method, please"
+                    " specify the attribute cost_func"
+                )
+            else:
+                return self.cost_func(final_state)
+        elif self.cost_method == "OBSERVABLE":
+            if self.cost_observable is None:
+                raise ValueError(
+                    "To use OBSERVABLE as the cost method, please"
+                    " specify the attribute cost_observable"
+                )
+            else:
+                return expect(self.cost_observable, final_state)
+        else:
+            raise ValueError(f"Unrecognised cost method: {self.cost_method}")
+
+    def optimize_parameters(
+        self,
+        initial="random",
+        method="COBYLA",
+        use_jac=False,
+        layer_by_layer=False,
+        bounds=None,
+        constraints=(),
+    ):
+        """
+        Run VQA optimization loop
+
+        Parameters
+        ----------
+        initial: str or list of floats, optional
+            Initialization method for the free parameters.
+            If a list of floats of dimensions (n,) for n free parameters
+            in the circuit is given, then these are taken to be the initial
+            conditions for the optimizer. Otherwise if a string is given:
+
+            * (Default) "random" will randomize initial free parameters
+              between 0 and 1.
+
+            * "ones" will set each initial free parameter to a value of 1.
+
+        method: str or callable, optional
+            Method to give to ``scipy.optimize.minimize``
+        use_jac: bool, optional
+            Whether to compute the jacobian or not. If computed, it will be
+            passed to the optimizer chosen by ``method``, regardless of if
+            the method is gradient-based or not.
+            Note that derivatives of unitaries generated by
+            ``ParameterizedHamiltonian`` are calculated with the
+            Frechet derivative of the exponential function,
+            using ``scipy.linalg.expm_frechet``.
+        layer_by_layer: bool, optional
+            Grow the circuit from a single layer to ``VQA.num_layers``. At each
+            step, hold the parameters found for previous layers fixed.
+        bounds: sequence or `scipy.optimize.Bounds`, optional
+            Bounds to be passed to the optimizer. Either
+
+            #. Instance of `scipy.optimize.Bounds`
+            #. Sequence of ``(min, max)`` tuples corresponding to each
+               free parameter.
+        constraints: list of `Constraint`
+            See `scipy.optimize.minimize` documentation.
+
+        Returns
+        -------
+        result: :obj:`.OptimizationResult`
+            The optimized angles and final state.
+        """
+
+        n_free_params = self.get_free_parameters_num()
+        # Set initial circuit parameters
+        if isinstance(initial, str):
+            if initial == "random":
+                angles = [random.random() for i in range(n_free_params)]
+            elif initial == "ones":
+                angles = [1 for i in range(n_free_params)]
+            else:
+                raise ValueError("Invalid initial condition string")
+        elif isinstance(initial, list) or isinstance(initial, np.ndarray):
+            if len(initial) != n_free_params:
+                raise ValueError(
+                    f"Expected {n_free_params} initial parameters"
+                    f"but got {len(initial)}."
+                )
+            angles = initial
+        else:
+            raise ValueError(
+                "Initial conditions were neither a list of values"
+                " nor a string specifying initialization."
+            )
+
+        # Run the scipy minimization function
+        if layer_by_layer:
+            max_layers = self.num_layers
+            n_params = 0
+            params = []
+            for layer_num in range(1, max_layers + 1):
+                print(f"Optimizing layer {layer_num}/{max_layers}")
+                self.num_layers = layer_num
+                n_tot = self.get_free_parameters_num()
+                # subset initialization parameters
+                init = angles[n_params:n_tot]
+
+                def layer(a, p):
+                    return self.evaluate_parameters(np.append(p, a))
+
+                if use_jac:
+
+                    def layer_jac(a, p):
+                        return self.compute_jac(
+                            np.append(p, a), list(range(n_params, n_tot))
+                        )
+
+                else:
+                    layer_jac = None
+                res = minimize(
+                    layer,
+                    init,
+                    args=(params),
+                    method=method,
+                    jac=layer_jac,
+                    bounds=bounds,
+                    constraints=constraints,
+                )
+                params = np.append(params, res.x)
+                n_params += n_tot - n_params
+            angles = params
+        else:
+            res = minimize(
+                self.evaluate_parameters,
+                angles,
+                method=method,
+                jac=self.compute_jac if use_jac else None,
+                bounds=bounds,
+                constraints=constraints,
+                options={"disp": False},
+            )
+            angles = res.x
+        final_state = self.get_final_state(angles)
+        result = OptimizationResult(res, final_state)
+        return result
+
+    def get_unitary_products(self, propagators):
+        """
+        Return two ordered lists of propagators in the circuit.
+        Useful for modifying individual propagators and computing
+        the product with these modifications. For example, to modify
+        U_k in a product of N unitaries, one could take
+        U_prods_back[N - 1 - k] * modified_U_k * U_prods[k]
+
+        Returns
+        -------
+        U_prods: list of :obj:`qutip.Qobj`
+            Ordered list of [identity, U_0, U_1, ... U_N]
+        U_prods_back: list of :obj:`qutip.Qobj`
+            Ordered list of [identity, U_N, U_{N-1}, ... U_0]
+        """
+        U_prods = [qeye([2 for _ in range(self.num_qubits)])]
+        U_prods_back = [qeye([2 for _ in range(self.num_qubits)])]
+        for i, _ in enumerate(propagators):
+            U_prods.append(propagators[i] * U_prods[-1])
+            U_prods_back.append(U_prods_back[-1] * propagators[-i - 1])
+        return U_prods, U_prods_back
+
+    def cost_derivative(self, U, dU):
+        """
+        Returns partial derivative of cost function (in observable
+        mode) with respect to the parameter in the block's unitary.
+        Assuming a block unitary of the form e^{-iH * theta}, this
+        will return d/(d theta) of the cost function in terms of
+        an observable.
+
+        Parameters
+        ----------
+        U: :obj:`qutip.Qobj`
+            Block unitary
+        dU: :obj:`qutip.Qobj`
+            Partial derivative of U with respect to its parameter
+
+        Returns
+        -------
+        dCost: float
+            Partial derivative of cost with respect to block's parameter
+        """
+        if self.cost_observable is None:
+            raise NotImplementedError(
+                "cost_derivative function only "
+                "implemented for observable cost "
+                "functions"
+            )
+        init = self.get_initial_state()
+        obs = self.cost_observable
+        dCost = (init.dag() * dU.dag()) * obs * (U * init) + (
+            init.dag() * U.dag()
+        ) * obs * (dU * init)
+        return dCost[0].item().real
+
+    def compute_jac(self, angles, indices_to_compute=None):
+        """
+        Compute the jacobian for the circuit's cost function,
+        assuming the cost function is in observable mode.
+
+        Parameters
+        ----------
+        angles: list of float
+            Circuit free parameters
+        indicies_to_compute: list of int, optional
+            Block indices for which to use in computing the jacobian.
+            By default, this is every index (every block).
+
+        Returns
+        -------
+        jac: (n,) numpy array of floats
+        """
+        if indices_to_compute is None:
+            indices_to_compute = list(range(len(angles)))
+
+        circ = self.construct_circuit(angles)
+        propagators = circ.propagators()
+        U = gate_sequence_product(propagators)
+        U_prods, U_prods_back = self.get_unitary_products(propagators)
+        # subtract one for the identity matrix
+        n = len(U_prods) - 1
+
+        def modify_unitary(k, U):
+            return U_prods_back[n - 1 - k] * U * U_prods[k]
+
+        jacobian = []
+        i = 0
+        for k, block in enumerate(self.get_block_series()):
+            n_params = block.get_free_parameters_num()
+            if n_params > 0:
+                if i in indices_to_compute:
+                    dBlock = block.get_unitary_derivative(
+                        angles[i : i + n_params]
+                    )
+                    dU = modify_unitary(k, dBlock)
+                    jacobian.append(self.cost_derivative(U, dU))
+                i += n_params
+        return np.array(jacobian)
+
+    def export_image(self, filename="circuit.png"):
+        """
+        Export an image of the circuit.
+
+        Parameters
+        ----------
+        filename: str, optional
+            The name of the exported file
+        """
+        circ = self.construct_circuit(
+            [1 for _ in range(self.get_free_parameters_num())]
+        )
+        f = open(filename, "wb+")
+        f.write(circ.png.data)
+        f.close()
+        print(f"Image saved to ./{filename}")
+
+
+class ParameterizedHamiltonian:
+    """
+    Hamiltonian with 0 or more parameterized terms.
+    In general, computes a unitary as
+    :math:`U = e^{H_0 + p_1 H_1 + P_2 H_2 + \\dots}`
+
+    Parameters
+    ----------
+    parameterized_terms: list of :obj:`qutip.Qobj`
+        Hamiltonian terms which each require a unique parameter
+    constant_term: :obj:`qutip.Qobj`
+        Hamiltonian term which does not require parameters.
+    """
+
+    def __init__(self, parameterized_terms=[], constant_term=None):
+        self.p_terms = parameterized_terms
+        self.c_term = constant_term
+        self.num_parameters = len(parameterized_terms)
+        if len(self.p_terms) == 0 and self.c_term is None:
+            raise ValueError(
+                "Parameterized Hamiltonian " "initialised with no terms given"
+            )
+
+    def get_hamiltonian(self, params):
+        if not len(params) == self.num_parameters:
+            raise ValueError(
+                f"params should be of length {self.num_parameters}"
+                f"but was {len(params)}"
+            )
+
+        # Match each p_term with a parameter
+        H_tot = sum(param * H for param, H in zip(self.p_terms, params)) + (
+            self.c_term if self.c_term else 0
+        )
+        return H_tot
+
+
+class VQABlock:
+    """
+    Component of a :obj:`.VQA`. Can return a unitary, and take
+    derivatives of its own unitary. Forms a :obj:`.Gate` in the
+    :obj:`.QubitCircuit` generated by the :obj:`.VQA`.
+
+    Parameters
+    ----------
+    operator: :obj:`qutip.Qobj` or Callable or str
+        If given as a :obj:`qutip.Qobj`, assumed to be a Hamiltonian with
+        a single global parameter.
+        If given as a Callable, assumed to take in a parameter, and return
+        a unitary operator.
+        If given as a str, assumed to reference a native QuTiP gate from
+        ``qutip_qip.operations``
+    is_unitary: bool, optional
+        Specifies that the operator  was already in Unitary form,
+        and does not need to be exponentiated, or take a parameter.
+    name: str, optional
+        Name of the block. This will be used in the custom
+        ``user_gates`` dict of the circuit. If not provided,
+        a name will be generated as "U"+str(len(VQA.blocks)).
+    targets: list of int, optional
+        The qubits targetted by the gate. By default, applied
+        to all qubits.
+    initial: bool, optional
+        Whether or not to repeat this block in layers. For example,
+        this should be True if this block is only used for
+        circuit initialization.
+    """
+
+    def __init__(
+        self,
+        operator,
+        is_unitary=False,
+        name=None,
+        targets=None,
+        initial=False,
+    ):
+        self.operator = operator
+        self.is_unitary = is_unitary
+        self.name = name
+        self.targets = targets
+        self.initial = initial
+        self.is_native_gate = False
+        self.num_parameters = 0
+
+        if isinstance(operator, Qobj):
+            if not self.is_unitary:
+                self.num_parameters = 1
+        elif isinstance(operator, str):
+            self.is_native_gate = True
+            if targets is None:
+                raise ValueError("Targets must be specified for native gates")
+        elif isinstance(operator, ParameterizedHamiltonian):
+            self.num_parameters = operator.num_parameters
+        elif isinstance(operator, types.FunctionType):
+            self.num_parameters = 1
+        else:
+            raise ValueError(
+                "operator should be either: Qobj | function which"
+                " returns Qobj | ParameterizedHamiltonian"
+                " instance | string referring to gate."
+            )
+
+    def get_free_parameters_num(self):
+        return self.num_parameters
+
+    def get_unitary(self, angles=None):
+        """
+        Return the block unitary.
+
+        Parameters
+        ----------
+        angles: list of float, optional
+            Block free parameters. Required if the block has free parameters.
+        """
+        # Qobj unitary or zero-parmeters function returning Qobj unitary
+        if angles is None:
+            if self.is_unitary:
+                return self.operator
+            raise ValueError("No angles were given and block was not unitary")
+
+        if len(angles) != self.get_free_parameters_num():
+            raise ValueError(
+                f"Expected {self.get_free_parameters_num()} angles"
+                f" but got {len(angles)}."
+            )
+
+        # Case where the operator is a string referring to an existing gate.
+        if self.is_native_gate:
+            raise TypeError("Can't compute unitary of native gate")
+        # Function returning Qobj unitary
+        if isinstance(self.operator, types.FunctionType):
+            # In the future, this could be generalized to multiple angles
+            unitary = self.operator(angles[0])
+            if not isinstance(unitary, Qobj):
+                raise TypeError("Provided function does not return Qobj")
+            return unitary
+        # ParameterizedHamiltonian instance
+        if isinstance(self.operator, ParameterizedHamiltonian):
+            return (-1j * self.operator.get_hamiltonian(angles)).expm()
+
+        # If there's no other specification, treat operator as Hamiltonian
+        if len(angles) != 1:
+            raise ValueError(
+                "Expected one angle for singly-parameterized Hamiltonian."
+            )
+
+        return (-1j * angles[0] * self.operator).expm()
+
+    def get_unitary_derivative(self, angles, term_index=0):
+        """
+        Compute the derivative of the block's unitary with respect to its
+        free parameter, assuming it is of the form :math:`e^{-i \\theta H}`
+        for a free parameter theta. If the block's operator is a
+        :obj:`ParameterizedHamiltonian`, use the Frechet derivative of the
+        exponential function.
+
+        Parameters
+        ----------
+        angle: list of float
+            free parameters to take derivatives with respect to
+        term_index: int, optional
+            Index of Parameterized Hamiltonian term that specifies the matrix
+            direction in which to take the derivative.
+
+        Returns
+        -------
+        derivative: float
+        """
+        if self.is_unitary or self.is_native_gate:
+            raise ValueError(
+                "Can only take derivative of block specified "
+                "by Hamiltonians or ParameterizedHamiltonian instances."
+            )
+        if isinstance(self.operator, ParameterizedHamiltonian):
+            arg = -1j * self.operator.get_hamiltonian(angles)
+            direction = -1j * self.operator.p_terms[term_index]
+            return Qobj(
+                expm_frechet(arg, direction, compute_expm=False),
+                dims=direction.dims,
+            )
+        if len(angles) != 1:
+            raise ValueError(
+                "Expected a single angle for non-"
+                "ParameterizedHamiltonian instance."
+            )
+        return self.get_unitary(angles) * -1j * self.operator
+
+
+class OptimizationResult:
+    """
+    Class for results of :obj:`.VQA` optimization loop.
+
+    Parameters
+    ----------
+    res: scipy results instance
+    final_state: :obj:`qutip.Qobj`
+        Final state of the circuit after optimization.
+    """
+
+    def __init__(self, res, final_state):
+        self.res = res
+        self.angles = res.x
+        self.min_cost = res.fun
+        self.nfev = res.nfev
+        self.final_state = final_state
+
+    def _highest_prob_bitstring(self, state):
+        """
+        Return the bitstring associated with the
+        highest probability amplitude measurement state.
+        """
+        num_qubits = int(np.log2(state.shape[0]))
+        index = np.argmax(abs(state))
+        return format(index, f"0{num_qubits}b")
+
+    def get_top_bitstring(self):
+        """
+        Return the bitstring associated with the highest probability
+        measurement outcome
+
+        Returns
+        -------
+        bitstring: str
+            bitstring in the form :math:`|x_0x_1...x_n>` where each
+            :math:`x_i` is 0 or 1 and n is the number of qubits of the system.
+        """
+        return "|" + self._highest_prob_bitstring(self.final_state) + ">"
+
+    def __str__(self):
+        return (
+            "Optimization Result:\n"
+            + f"\tMinimum cost: {self.min_cost}\n"
+            + f"\tNumber of function evaluations: {self.nfev}\n"
+            + f"\tParameters found: {self.angles}"
+        )
+
+    def _label_to_sets(self, S, bitstring):
+        """
+        Convert bitstring to string representation of
+        two sets containing elements of the problem instance.
+
+        Parameters
+        ----------
+        S: list of float
+            Problem instance
+        bitstring: str
+
+        Returns
+        -------
+        sets: str
+            String representation the two sets.
+        """
+        s1 = []
+        s2 = []
+        for i, c in enumerate(bitstring.strip("|").strip(">")):
+            if c == "0":
+                s1.append(S[i])
+            else:
+                s2.append(S[i])
+        return (str(s1) + " " + str(s2)).replace("[", "{").replace("]", "}")
+
+    def plot(self, S=None, label_sets=False, top_ten=False, display=True):
+        """
+        Plot probability amplitudes of each measurement
+        outcome of a state.
+
+        Parameters
+        ----------
+        S: list of float, optional
+            Problem instance
+        min_cost: str, optional
+            The minimum cost found by optimization
+        label_sets: bool, optional
+            Replace bitstring labels with sets referring to the inferred
+            output of the combinatorial optimization problem. For example
+            a bitstring :math:`|010>` would produce a set with the first and
+            last elements of S, and one with the second element of S.
+        top_ten: bool, optional
+            Only plot the ten highest-probability states.
+        display: bool, optional
+            Display the plot with the pyplot plot.show() method
+        """
+        try:
+            import matplotlib.pyplot as plt
+        except Exception:
+            print("could not import matplotlib.pyplot")
+            quit()
+        state = self.final_state
+        min_cost = self.min_cost
+
+        num_qubits = int(np.log2(state.shape[0]))
+        probs = [abs(i.item()) ** 2 for i in state]
+        bitstrings = [
+            "|" + format(i, f"0{num_qubits}b") + ">"
+            for i in range(2**num_qubits)
+        ]
+        if top_ten and len(probs) > 10:
+            threshold = sorted(probs)[-11]
+            top_probs = []
+            top_bitstrings = []
+            for i, prob in enumerate(probs):
+                if prob > threshold:
+                    top_probs.append(prob)
+                    top_bitstrings.append(bitstrings[i])
+            bitstrings = top_bitstrings
+            probs = top_probs
+        if label_sets:
+            labels = [
+                self._label_to_sets(S, bitstring) for bitstring in bitstrings
+            ]
+        plt.bar(
+            list(range(len(bitstrings))),
+            probs,
+            tick_label=labels if label_sets else bitstrings,
+            width=0.8,
+        )
+        plt.xticks(rotation=30)
+        plt.tight_layout()
+        plt.xlabel("Measurement outcome")
+        plt.ylabel("Probability")
+        plt.title(
+            "Measurement Outcomes after Optimisation. "
+            f"Cost: {round(min_cost, 2)}"
+        )
+        plt.tight_layout()
+        if display:
+            plt.show()


Test patch:
diff --git a/tests/test_vqa.py b/tests/test_vqa.py
new file mode 100644
index 00000000..c9fc359b
--- /dev/null
+++ b/tests/test_vqa.py
@@ -0,0 +1,250 @@
+import pytest
+import functools
+import itertools
+import numpy as np
+import qutip
+from qutip_qip.vqa import (
+    VQA,
+    VQABlock,
+    ParameterizedHamiltonian,
+    OptimizationResult,
+)
+
+
+class TestVQABlock:
+    """
+    Test class for the VQABlock class
+    """
+
+    def test_initial(self):
+        block = VQABlock(qutip.sigmax())
+        assert block.get_unitary([1]) == (-1j * qutip.sigmax()).expm()
+
+    @pytest.mark.parametrize("angle", [1, 2, 3])
+    def test_parameterization(self, angle):
+        block = VQABlock(qutip.sigmax())
+        assert (
+            block.get_unitary([angle]) == (-1j * angle * qutip.sigmax()).expm()
+        )
+        assert block.get_free_parameters_num() == 1
+        block = VQABlock(qutip.sigmaz())
+        assert (
+            block.get_unitary([angle]) == (-1j * angle * qutip.sigmaz()).expm()
+        )
+        block = VQABlock(qutip.sigmay())
+        assert (
+            block.get_unitary([angle]) == (-1j * angle * qutip.sigmay()).expm()
+        )
+
+    @pytest.mark.parametrize("angle", [1, 2, 3])
+    def test_unitary_function(self, angle):
+        block = VQABlock(lambda t: (t * -1j * qutip.sigmax()).expm())
+        assert (
+            block.get_unitary([angle]) == (-1j * angle * qutip.sigmax()).expm()
+        )
+
+
+class TestVQA:
+    """
+    Test class for the VQA class
+    """
+
+    @pytest.mark.parametrize("n", [[[0, 1], [1.5, 1]], [[1, 0], [1, 1.5]]])
+    def test_initialization_bad_input(self, n):
+        """
+        Test numerical inputs for initializing VQA
+        """
+        vqa = VQA(num_qubits=1, num_layers=1)
+        with pytest.raises(ValueError):
+            vqa = VQA(num_qubits=n[0][0], num_layers=n[0][1])
+        with pytest.raises(TypeError):
+            vqa = VQA(num_qubits=n[1][0], num_layers=n[1][1])
+
+    @pytest.mark.parametrize("method", ["OBSERVABLE", "STATE", "BITSTRING"])
+    def test_cost_methods(self, method):
+        vqa = VQA(num_qubits=1, num_layers=1, cost_method=method)
+
+    @pytest.mark.parametrize("method", ["Something else", 1, [0]])
+    def test_invalid_cost_methods(self, method):
+        with pytest.raises(ValueError):
+            vqa = VQA(num_qubits=1, num_layers=1, cost_method=method)
+
+
+class TestVQACircuit:
+    """
+    Test class for circuits made from the VQA and VQABlock combination
+    """
+
+    def test_initialize(self):
+        block = VQABlock(qutip.sigmax(), is_unitary=True)
+        vqa = VQA(num_qubits=1)
+        final_state = vqa.get_final_state([])
+        assert final_state == qutip.basis(2, 0)
+        vqa.add_block(block)
+        final_state = vqa.get_final_state([])
+        assert final_state == qutip.basis(2, 1)
+
+    def test_layer_repetitions(self):
+        """
+        tests layers are ordered and repeated correctly
+        """
+        vqa = VQA(num_qubits=1, num_layers=3)
+        initial_block = VQABlock(qutip.sigmax(), initial=True)
+        h_block = VQABlock("SNOT", targets=[0])
+        vqa.add_block(initial_block)
+        vqa.add_block(h_block)
+        # Expect [X, H, H, H] for the three layers
+        blocks = vqa.get_block_series()
+        assert blocks[0] == initial_block
+        for i in range(1, 4):
+            assert blocks[i].name == h_block.name
+
+    def test_parameterized_circuit(self):
+        block = VQABlock(qutip.sigmax())
+        vqa = VQA(num_qubits=1)
+        final_state = vqa.get_final_state([0])
+        assert final_state == qutip.basis(2, 0)
+        vqa.add_block(block)
+        final_state = vqa.get_final_state([np.pi / 2])
+        assert final_state == qutip.Qobj([[0], [-1j]])
+
+    def test_trivial_optimization(self):
+        """
+        tests trivial optimization case where the initial conditinos
+        are already optimal for the problem
+        """
+        block = VQABlock(qutip.sigmax())
+        # test the STATE type of cost function
+        vqa = VQA(num_qubits=1, cost_method="STATE")
+        vqa.add_block(block)
+        # try to reach the |1> state from the |0> state
+        vqa.cost_func = lambda s: 1 - s.overlap(qutip.basis(2, 1)).real
+        res = vqa.optimize_parameters(initial=[np.pi / 2])
+        assert res.get_top_bitstring() == "|1>"
+
+    def test_layer_by_layer(self):
+        """
+        tests trivial optimization going layer-by-layer
+        """
+        vqa = VQA(num_qubits=1, cost_method="STATE", num_layers=4)
+        block = VQABlock(qutip.sigmax())
+        vqa.add_block(block)
+        vqa.cost_func = lambda s: 1 - s.overlap(qutip.basis(2, 1)).real
+        res = vqa.optimize_parameters(
+            initial=[np.pi / 2, 0, 0, 0], layer_by_layer=True
+        )
+        assert res.get_top_bitstring() == "|1>"
+
+    @pytest.mark.parametrize("use_jac", [True, False])
+    def test_bfgs_optimization(self, use_jac):
+        """
+        test bfgs optimizer, starting from a close initial guess
+        """
+        block = VQABlock(qutip.sigmax())
+        vqa = VQA(num_qubits=1, cost_method="OBSERVABLE")
+        vqa.add_block(block)
+        # try to reach the |1> state from the |0> state
+        vqa.cost_observable = qutip.sigmaz()
+        res = vqa.optimize_parameters(
+            initial=[np.pi / 2 + 0.2], method="BFGS", use_jac=use_jac
+        )
+        assert res.get_top_bitstring() == "|1>"
+        # check we actually found the function minimum
+        assert round(res.res.x[0], 2) == round(np.pi / 2, 2)
+
+    def test_parameterized_hamiltonian_blocks(self):
+        """
+        Test that parameterized hamiltonian blocks correctly
+        apply parameters
+        """
+        # Hamiltonian that looks like (t_1*X  +  t_2*Z)
+        block = VQABlock(
+            ParameterizedHamiltonian([qutip.sigmax(), qutip.sigmaz()])
+        )
+        vqa = VQA(num_qubits=1, num_layers=2)
+        vqa.add_block(block)
+        # Do (pi/2*X + 0*Z) and then (0*X + pi/2*Z)
+        final_state = vqa.get_final_state([np.pi / 2, 0, 0, np.pi / 2])
+        # expect |1>
+        assert final_state == qutip.basis(2, 1)
+
+    def test_parameterized_hamiltonian_frechet_derivative(self):
+        """
+        Test gradient-based optimization on parameterized Hamiltonian blocks
+        """
+        vqa = VQA(num_qubits=1)
+        vqa.cost_observable = qutip.sigmaz()
+        block = VQABlock(ParameterizedHamiltonian([qutip.sigmax()]))
+        vqa.add_block(block)
+        res = vqa.optimize_parameters(
+            initial=[np.pi / 2 + 0.2], method="BFGS", use_jac=True
+        )
+        assert res.get_top_bitstring() == "|1>"
+
+    @pytest.mark.parametrize("todo", [False, True])
+    def test_plot(self, todo):
+        """
+        Check plotting function returns without error
+        """
+        # Only test on environments that have the matplotlib dependency
+        try:
+            import matplotlib.pyplot as plt
+        except Exception:
+            return True
+        vqa = VQA(num_qubits=4, num_layers=1, cost_method="STATE")
+        for i in range(4):
+            vqa.add_block(VQABlock("X", targets=[i]))
+        vqa.cost_func = lambda s: 0
+        res = vqa.optimize_parameters()
+        res.plot(top_ten=todo, display=False)
+
+    def test_bitstring_cost(self):
+        "Check the bitstring sampling function"
+        vqa = VQA(num_qubits=1, cost_method="BITSTRING")
+        vqa.add_block(VQABlock(qutip.sigmax()))
+        # target the |1> state by giving the "1" string a cost of 0
+        vqa.cost_func = lambda s: 1 - int(s)
+        res = vqa.optimize_parameters(initial=[np.pi / 2 + 1e-3])
+        assert res.get_top_bitstring() == "|1>"
+
+    def test_optimization_errors(self):
+        """
+        Tests for value errors relating to optimization procedure
+        """
+        vqa = VQA(num_qubits=1)
+        vqa.add_block(VQABlock(qutip.sigmax()))
+        vqa.cost_observable = qutip.sigmax()
+        with pytest.raises(ValueError):
+            # Invalid initialization string
+            res = vqa.optimize_parameters(initial="something else")
+        with pytest.raises(ValueError):
+            # Incorrect number of parameters
+            res = vqa.optimize_parameters(initial=[1, 1])
+        # Valid initialization string
+        res = vqa.optimize_parameters(initial="ones")
+        assert res is not None
+
+
+class TestOptimizationResult:
+    """
+    Test class for the OptimizationResult class
+    """
+
+    def test_label_to_sets(self):
+        """
+        Test partitioning a set of elements based on a bitstring.
+        e.g. [1, 2, 3] with the bitstring |010> should become
+        '{1, 3} {2}'.
+        """
+        # Fake scipy res object
+        class fakeRes:
+            pass
+
+        fakeScipyRes = fakeRes()
+        fakeRes.x = [1]
+        fakeRes.fun = None
+        fakeRes.nfev = None
+        # Problem instance
+        S = [1, 2, 3]
+        result = OptimizationResult(fakeScipyRes, qutip.basis(8, 1))
+        assert result._label_to_sets(S, "|010>") == "{1, 3} {2}"