[feat] GeneralBraOpKet and mutable tensors in GeneralState (#165)

# Description

- Refactor of `GeneralState` so that it uses CuTensorNet's
`NetworkState`
- Support for mutable tensors on `GeneralState`, so that parameterised
circuits can be simulated without optimising the contraction path again
for different parameter values.
- Includes a new `GeneralBraOpKet` operator that can be used to
calculate expectation values, inner products and arbitrary matrix
elements.
- Some misc improvements to the `general_state` module here and there.

# Related issues

#150 

# Checklist

- [x] Write a notebook tutorial.
- [x] I have performed a self-review of my code.
- [x] I have commented hard-to-understand parts of my code.
- [x] I have made corresponding changes to the public API documentation.
- [x] I have added tests that prove my fix is effective or that my
feature works.
- [x] I have updated the changelog with any user-facing changes.

---------

Co-authored-by: Oliver Backhouse <[email protected]>
Co-authored-by: Jake Arkinstall <[email protected]>

.github/workflows/build-with-nix.yml

@@ -39,3 +39,6 @@ jobs:
     - name: Test pytket-cutensornet
       # impure is necessary due to nixgl usage (system-dependent cuda)
       run: nix run .#tests --impure --accept-flake-config
+    - name: Test example notebooks
+      # impure is necessary due to nixgl usage (system-dependent cuda)
+      run: nix run .#example-tests --impure --accept-flake-config

.github/workflows/lint.yml

@@ -30,4 +30,4 @@ jobs:
         black --check .
     - name: Run pylint
       run: |
-        pylint --recursive=y --ignore=ttn_tutorial.py,mps_tutorial.py */
+        pylint --recursive=y pytket/extensions/cutensornet/

docs/api.rst

@@ -1,11 +1,6 @@
 API documentation
 -----------------
 
-.. autoclass:: pytket.extensions.cutensornet.CuTensorNetHandle
-
-    .. automethod:: destroy
-
-
 .. toctree::
     modules/general_state.rst
     modules/structured_state.rst

docs/changelog.rst

@@ -3,6 +3,15 @@
 Changelog
 ~~~~~~~~~
 
+Unreleased
+----------
+
+* New API: ``GeneralBraOpKet`` for exact calculation of arbitrary ``<bra|op|ket>`` values. Can be used to calculate inner products, expectation values and arbitrary matrix elements.
+* New feature: both ``GeneralState`` and ``GeneralBraOpKet`` admit circuits with parameterised gates.
+* New feature: ``GeneralState`` has a new method ``get_amplitude`` to obtain the amplitude of computational basis states.
+* New feature: ``GeneralState`` and ``CuTensorNetShotsBackend`` now support RNG seeds for sampling.
+* Deprecated ``TensorNetwork`` object. It is still available for the sake of backwards compatibility, but it has been removed from doc pages.
+
 0.9.0 (October 2024)
 ---------------------
 

docs/modules/general_state.rst

@@ -5,29 +5,16 @@ General state (exact) simulation
 
 .. autoclass:: pytket.extensions.cutensornet.general_state.GeneralState()
 
-    .. automethod:: __init__
+    .. automethod:: sample
+    .. automethod:: get_amplitude
     .. automethod:: get_statevector
     .. automethod:: expectation_value
-    .. automethod:: sample
     .. automethod:: destroy
 
-cuQuantum `contract` API interface
-~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
-
-.. autoclass:: pytket.extensions.cutensornet.general_state.TensorNetwork
-
-.. autoclass:: pytket.extensions.cutensornet.general_state.PauliOperatorTensorNetwork
-
-.. autoclass:: pytket.extensions.cutensornet.general_state.ExpectationValueTensorNetwork
-
-.. autofunction:: pytket.extensions.cutensornet.general_state.tk_to_tensor_network
-
-.. autofunction:: pytket.extensions.cutensornet.general_state.measure_qubits_state
-
-.. autofunction:: pytket.extensions.cutensornet.general_state.get_operator_expectation_value
-
-.. autofunction:: pytket.extensions.cutensornet.general_state.get_circuit_overlap
+.. autoclass:: pytket.extensions.cutensornet.general_state.GeneralBraOpKet()
 
+    .. automethod:: contract
+    .. automethod:: destroy
 
 Pytket backend
 ~~~~~~~~~~~~~~

docs/modules/structured_state.rst

@@ -3,6 +3,13 @@ Structured state evolution
 
 .. automodule:: pytket.extensions.cutensornet.structured_state
 
+Library handle
+~~~~~~~~~~~~~~
+
+.. autoclass:: pytket.extensions.cutensornet.CuTensorNetHandle
+
+    .. automethod:: destroy
+
 
 Simulation
 ~~~~~~~~~~

examples/check-examples

@@ -9,7 +9,12 @@ do
     p2j -o -t ${name}-gen.ipynb python/${name}.py
     cmp ${name}.ipynb ${name}-gen.ipynb
     rm ${name}-gen.ipynb
-    # TODO, add this when GPU is added to CI
-    # Run script:    
-    # python python/${name}.py
+
+    # run tests are performed in nix, allowing
+    # us to manage the testing environment in a
+    # reproducible way.
+    #
+    # See /nix-support/pytket-cutensornet.nix,
+    # in the derivation called
+    # run-pytket-cutensornet-examples.
 done

examples/ci-tested-notebooks.txt

@@ -1,2 +1,3 @@
 mps_tutorial
-ttn_tutorial
+ttn_tutorial
+general_state_tutorial

examples/general_state_tutorial.ipynb

@@ -0,0 +1 @@
+{"cells": [{"cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": ["import numpy as np\n", "from sympy import Symbol\n", "from scipy.stats import unitary_group  # type: ignore\n", "from pytket.circuit import Circuit, OpType, Unitary2qBox, Qubit, Bit\n", "from pytket.passes import DecomposeBoxes\n", "from pytket.utils import QubitPauliOperator\n", "from pytket._tket.pauli import Pauli, QubitPauliString\n", "from pytket.circuit.display import render_circuit_jupyter"]}, {"cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": ["from pytket.extensions.cutensornet.general_state import (\n", "    GeneralState,\n", "    GeneralBraOpKet,\n", ")\n", "from pytket.extensions.cutensornet.backends import CuTensorNetShotsBackend"]}, {"cell_type": "markdown", "metadata": {}, "source": ["# Introduction<br>\n", "This notebook is a guide on how to use the features provided in the `general_state` submodule of pytket-cutensornet. This submodule is a thin wrapper of CuTensorNet's `NetworkState`, allowing users to convert pytket circuits into tensor networks and use CuTensorNet's contraction path optimisation algorithm.<br>\n", "All simulations realised with this submodule are *exact*. Once the pytket circuit has been converted to a tensor network, the computation has two steps:<br>\n", "  1. *Contraction path optimisation*. Attempts to find an order of contracting pairs of tensors in which the the total number of FLOPs is minimised. No operation on the tensor network occurs at this point. Runs on CPU.<br>\n", "  2. *Tensor network contraction*. Uses the ordering of contractions found in the previous step evaluate the tensor network. Runs on GPU.<br>\n", "<br>\n", "**Reference**: The original contraction path optimisation algorithm that NVIDIA implemented on CuTensorNet: https://arxiv.org/abs/2002.01935"]}, {"cell_type": "markdown", "metadata": {}, "source": ["# `GeneralState`<br>\n", "The class `GeneralState` is used to convert a circuit into a tensor network and query information from the final state. Let's walk through a simple example."]}, {"cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": ["my_circ = Circuit(5)\n", "my_circ.CX(3, 4)\n", "my_circ.H(2)\n", "my_circ.CZ(0, 1)\n", "my_circ.ZZPhase(0.1, 4, 3)\n", "my_circ.TK2(0.3, 0.5, 0.7, 2, 1)\n", "my_circ.Ry(0.2, 0)\n", "my_circ.measure_all()"]}, {"cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": ["render_circuit_jupyter(my_circ)"]}, {"cell_type": "markdown", "metadata": {}, "source": ["The first step is to convert our pytket circuit into a tensor network. This is straightforward:"]}, {"cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": ["tn_state = GeneralState(my_circ)"]}, {"cell_type": "markdown", "metadata": {}, "source": ["The variable `tn_state` now holds a tensor network representation of `my_circ`.<br>\n", "**Note**: Circuits must not have mid-circuit measurements or classical logic. The measurements at the end of the circuit are stripped and only considered when calling `tn_state.sample(n_shots)`.<br>\n", "We can now query information from the state. For instance, let's calculate the probability of in the qubits 0 and 3 agreeing in their outcome."]}, {"cell_type": "markdown", "metadata": {}, "source": ["First, let's generate `|x>` computational basis states where `q[0]` and `q[3]` agree on their values. We can do this with some bitwise operators and list comprehension.<br>\n", "**Note**: Remember that pytket uses \"increasing lexicographic order\" (ILO) for qubits, so `q[0]` is the most significant bit."]}, {"cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": ["selected_states = [\n", "    x\n", "    for x in range(2**my_circ.n_qubits)\n", "    if (  # Iterate over all possible states\n", "        x & int(\"10000\", 2) == 0\n", "        and x & int(\"00010\", 2) == 0  # both qubits are 0 or...\n", "        or x & int(\"10000\", 2) != 0\n", "        and x & int(\"00010\", 2) != 0  # both qubits are 1\n", "    )\n", "]"]}, {"cell_type": "markdown", "metadata": {}, "source": ["We can now query the amplitude of all of these states and calculate the probability by summing their squared absolute values."]}, {"cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": ["amplitudes = []\n", "for x in selected_states:\n", "    amplitudes.append(tn_state.get_amplitude(x))\n", "probability = sum(abs(a) ** 2 for a in amplitudes)\n", "print(f\"Probability: {probability}\")"]}, {"cell_type": "markdown", "metadata": {}, "source": ["Of course, calculating probabilities by considering the amplitudes of all relevant states is not efficient in general, since we may need to calculate a number of amplitudes that scales exponentially with the number of qubits. An alternative is to use expectation values. In particular, all of the states in `selected_states` are +1 eigenvectors of the `ZIIZI` observable and, hence, we can calculate the probability `p` by solving the equation `<ZIIZI> = (+1)p + (-1)(1-p)` using the fact that `ZIIZI` only has +1 and -1 eigenvalues."]}, {"cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": ["string_ZIIZI = QubitPauliString(\n", "    my_circ.qubits, [Pauli.Z, Pauli.I, Pauli.I, Pauli.Z, Pauli.I]\n", ")\n", "observable = QubitPauliOperator({string_ZIIZI: 1.0})\n", "expectation_val = tn_state.expectation_value(observable).real\n", "exp_probability = (expectation_val + 1) / 2\n", "assert np.isclose(probability, exp_probability, atol=0.0001)\n", "print(f\"Probability: {exp_probability}\")"]}, {"cell_type": "markdown", "metadata": {}, "source": ["Alternatively, we can estimate the probability by sampling."]}, {"cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": ["n_shots = 100000\n", "outcomes = tn_state.sample(n_shots)\n", "hit_count = 0\n", "for bit_tuple, count in outcomes.get_counts().items():\n", "    if bit_tuple[0] == bit_tuple[3]:\n", "        hit_count += count\n", "samp_probability = hit_count / n_shots\n", "assert np.isclose(probability, samp_probability, atol=0.01)\n", "print(f\"Probability: {samp_probability}\")"]}, {"cell_type": "markdown", "metadata": {}, "source": ["When we finish doing computations with the `tn_state` we must destroy it to free GPU memory."]}, {"cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": ["tn_state.destroy()"]}, {"cell_type": "markdown", "metadata": {}, "source": ["To avoid forgetting this final step, we recommend users call `GeneralState` (and `GeneralBraOpKet`) as context managers:"]}, {"cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": ["with GeneralState(my_circ) as my_state:\n", "    expectation_val = my_state.expectation_value(observable)\n", "print(expectation_val)"]}, {"cell_type": "markdown", "metadata": {}, "source": ["Using this syntax, `my_state` is automatically destroyed when the code exists the `with ...` block."]}, {"cell_type": "markdown", "metadata": {}, "source": ["# Parameterised circuits<br>\n", "Circuits that only differ on the parameters of their gates have the same tensor network topology and, hence, we may use the same contraction path for all of them."]}, {"cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": ["a, b, c = Symbol(\"a\"), Symbol(\"b\"), Symbol(\"c\")\n", "param_circ1 = Circuit(5)\n", "param_circ1.Ry(a, 3).Ry(0.27, 4).CX(4, 3).Ry(b, 2).Ry(0.21, 3)\n", "param_circ1.Ry(0.12, 0).Ry(a, 1)\n", "param_circ1.add_gate(OpType.CnX, [0, 1, 4]).add_gate(OpType.CnX, [4, 1, 3])\n", "param_circ1.X(0).X(1).add_gate(OpType.CnY, [0, 1, 2]).add_gate(OpType.CnY, [0, 4, 3]).X(\n", "    0\n", ").X(1)\n", "param_circ1.Ry(-b, 0).Ry(-c, 1)\n", "render_circuit_jupyter(param_circ1)"]}, {"cell_type": "markdown", "metadata": {}, "source": ["We can pass a parameterised circuit to `GeneralState`. The value of the parameters is provided when calling methods of `GeneralState`. The contraction path is automatically reused on different calls to the same method."]}, {"cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": ["n_circs = 5\n", "with GeneralState(param_circ1) as param_state:\n", "    for i in range(n_circs):\n", "        symbol_map = {s: np.random.random() for s in [a, b, c]}\n", "        exp_val = param_state.expectation_value(observable, symbol_map=symbol_map)\n", "        print(f\"Expectation value for circuit {i}: {exp_val.real}\")"]}, {"cell_type": "markdown", "metadata": {}, "source": ["# `GeneralBraOpKet`<br>\n", "The `GeneralBraOpKet` can be used to calculate any number that can be represented as the result of some `<bra|op|ket>` where `|bra>` and `|ket>` are the final states of pytket circuits, and `op` is a `QubitPauliOperator`. The circuits for `|bra>` and `|ket>` need not be the same."]}, {"cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": ["x, y, z = Symbol(\"x\"), Symbol(\"y\"), Symbol(\"z\")\n", "param_circ2 = Circuit(5)\n", "param_circ2.H(0)\n", "param_circ2.S(1)\n", "param_circ2.Rz(x * z, 2)\n", "param_circ2.Ry(y + x, 3)\n", "param_circ2.TK1(x, y, z, 4)\n", "param_circ2.TK2(z - y, z - x, (x + y) * z, 1, 3)\n", "symbol_map = {a: 2.1, b: 1.3, c: 0.7, x: 3.0, y: 1.6, z: -8.3}"]}, {"cell_type": "markdown", "metadata": {}, "source": ["We can calculate inner products by providing no `op`:"]}, {"cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": ["with GeneralBraOpKet(bra=param_circ2, ket=param_circ1) as braket:\n", "    inner_prod = braket.contract(symbol_map=symbol_map)\n", "with GeneralBraOpKet(bra=param_circ1, ket=param_circ2) as braket:\n", "    inner_prod_conj = braket.contract(symbol_map=symbol_map)\n", "assert np.isclose(np.conj(inner_prod), inner_prod_conj)\n", "print(f\"<circ_b|circ_a> = {inner_prod}\")\n", "print(f\"<circ_a|circ_b> = {inner_prod_conj}\")"]}, {"cell_type": "markdown", "metadata": {}, "source": ["And we are not constrained to Hermitian operators:"]}, {"cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": ["string_XZIXX = QubitPauliString(\n", "    param_circ2.qubits, [Pauli.X, Pauli.Z, Pauli.I, Pauli.X, Pauli.X]\n", ")\n", "string_IZZYX = QubitPauliString(\n", "    param_circ2.qubits, [Pauli.I, Pauli.Z, Pauli.Z, Pauli.Y, Pauli.X]\n", ")\n", "string_ZIZXY = QubitPauliString(\n", "    param_circ2.qubits, [Pauli.Z, Pauli.I, Pauli.Z, Pauli.X, Pauli.Y]\n", ")\n", "operator = QubitPauliOperator(\n", "    {string_XZIXX: -1.38j, string_IZZYX: 2.36, string_ZIZXY: 0.42j + 0.3}\n", ")\n", "with GeneralBraOpKet(bra=param_circ2, ket=param_circ1) as braket:\n", "    value = braket.contract(operator, symbol_map=symbol_map)\n", "print(value)"]}, {"cell_type": "markdown", "metadata": {}, "source": ["# Backends<br>\n", "We provide a pytket `Backend` to obtain shots using `GeneralState`."]}, {"cell_type": "markdown", "metadata": {}, "source": ["Let's consider a more challenging circuit"]}, {"cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": ["def random_circuit(n_qubits: int, n_layers: int) -> Circuit:\n", "    \"\"\"Random quantum volume circuit.\"\"\"\n", "    c = Circuit(n_qubits, n_qubits)\n", "    for _ in range(n_layers):\n", "        qubits = np.random.permutation([i for i in range(n_qubits)])\n", "        qubit_pairs = [[qubits[i], qubits[i + 1]] for i in range(0, n_qubits - 1, 2)]\n", "        for pair in qubit_pairs:\n", "            # Generate random 4x4 unitary matrix.\n", "            SU4 = unitary_group.rvs(4)  # random unitary in SU4\n", "            SU4 = SU4 / (np.linalg.det(SU4) ** 0.25)\n", "            SU4 = np.matrix(SU4)\n", "            c.add_unitary2qbox(Unitary2qBox(SU4), *pair)\n", "    DecomposeBoxes().apply(c)\n", "    return c"]}, {"cell_type": "markdown", "metadata": {}, "source": ["Let's measure only three of the qubits.<br>\n", "**Note**: The complexity of this simulation increases exponentially with the number of qubits measured. Other factors leading to intractability are circuit depth and qubit connectivity."]}, {"cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": ["n_shots = 1000\n", "quantum_vol_circ = random_circuit(n_qubits=40, n_layers=5)\n", "quantum_vol_circ.Measure(Qubit(0), Bit(0))\n", "quantum_vol_circ.Measure(Qubit(1), Bit(1))\n", "quantum_vol_circ.Measure(Qubit(2), Bit(2))"]}, {"cell_type": "markdown", "metadata": {}, "source": ["The `CuTensorNetShotsBackend` is used in the same way as any other pytket `Backend`."]}, {"cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": ["backend = CuTensorNetShotsBackend()\n", "compiled_circ = backend.get_compiled_circuit(quantum_vol_circ)\n", "results = backend.run_circuit(compiled_circ, n_shots=n_shots)\n", "print(results.get_counts())"]}], "metadata": {"kernelspec": {"display_name": "Python 3", "language": "python", "name": "python3"}, "language_info": {"codemirror_mode": {"name": "ipython", "version": 3}, "file_extension": ".py", "mimetype": "text/x-python", "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3", "version": "3.6.4"}}, "nbformat": 4, "nbformat_minor": 2}