Add Pauli twirl + apply channels

forest/benchmarking/superoperator_tools.py

@@ -1,4 +1,5 @@
-"""A module for converting between different representations of superoperators.
+"""A module containing tools for working with superoperators. Eg. converting between different
+representations of superoperators.
 
 We have arbitrarily decided to use a column stacking convention.
 
@@ -27,8 +28,11 @@
 """
 from typing import Sequence, Tuple, List
 import numpy as np
-from forest.benchmarking.utils import n_qubit_pauli_basis
+from forest.benchmarking.utils import n_qubit_pauli_basis, partial_trace
 
+# ==================================================================================================
+# Superoperator conversion tools
+# ==================================================================================================
 
 def vec(matrix: np.ndarray) -> np.ndarray:
     """
@@ -363,3 +367,103 @@ def computational2pauli_basis_matrix(dim) -> np.ndarray:
     :return: A dim**2 by dim**2 basis transform matrix
     """
     return pauli2computational_basis_matrix(dim).conj().T / dim
+
+
+# ==================================================================================================
+# Channel and Superoperator approximation tools
+# ==================================================================================================
+def pauli_twirl_chi_matrix(chi_matrix: np.ndarray) -> np.ndarray:
+    r"""
+    Implements a Pauli twirl of a chi matrix (aka a process matrix).
+
+    See the folloiwng reference for more details
+
+    [SPICC] Scalable protocol for identification of correctable codes
+            Silva et al.,
+            PRA 78, 012347 
2008
+            http://dx.doi.org/10.1103/PhysRevA.78.012347
+            https://arxiv.org/abs/0710.1900
+
+    Note: Pauli twirling a quantum channel can give rise to a channel that is less noisy; use with
+    care.
+
+    :param chi_matrix:  a dim**2 by dim**2 chi or process matrix
+    :return: dim**2 by dim**2 chi or process matrix
+    """
+    return np.diag(chi_matrix.diagonal())
+
+
+#TODO Honest approximations for Channels that act on one or MORE qubits.
+
+# ==================================================================================================
+# Apply channel
+# ==================================================================================================
+def apply_kraus_ops_2_state(kraus_ops: Sequence[np.ndarray], state: np.ndarray) -> np.ndarray:
+    r"""
+    Apply a quantum channel, specified by Kraus operators, to state.
+
+    The Kraus operators need not be square.
+
+    :param kraus_ops: A list or tuple of N Kraus operators, each operator is M by dim ndarray
+    :param state: A dim by dim ndarray which is the density matrix for the state
+    :return: M by M ndarray which is the density matrix for the state after the action of kraus_ops
+    """
+    if isinstance(kraus_ops, np.ndarray):  # handle input of single kraus op
+        if len(kraus_ops[0].shape) < 2:  # first elem is not a matrix
+            kraus_ops = [kraus_ops]
+
+    dim, _ = state.shape
+    rows, cols = kraus_ops[0].shape
+
+    if dim != cols:
+        raise ValueError("Dimensions of state and Kraus operator are incompatible")
+
+    new_state = np.zeros((rows, rows))
+    for M in kraus_ops:
+        new_state += M @ state @ np.transpose(M.conj())
+
+    return new_state
+
+
+def apply_choi_matrix_2_state(choi: np.ndarray, state: np.ndarray) -> np.ndarray:
+    r"""
+    Apply a quantum channel, specified by a Choi matrix (using the column stacking convention),
+    to a state.
+
+    The Choi matrix is a dim**2 by dim**2 matrix and the state rho is a dim by dim matrix. The
+    output state is
+
+    rho_{out} = Tr_{A}[(rho^T \otimes Id) Choi_matrix ],
+
+    where Tr_{A} is the partial trace over hilbert space H_A with respect to the tensor product
+    H_A \otimes H_B, and T denotes transposition.
+
+
+    :param choi: a dim**2 by dim**2 matrix
+    :param state: A dim by dim ndarray which is the density matrix for the state
+    :return: a dim by dim matrix.
+    """
+    dim = int(np.sqrt(np.asarray(choi).shape[0]))
+    dims = [dim, dim]
+    tot_matrix = np.kron(state.transpose(), np.identity(dim)) @ choi
+    return partial_trace(tot_matrix, [1], dims)
+
+
+def apply_lioville_matrix_2_state(superoperator: np.ndarray, state: np.ndarray) -> np.ndarray:
+    r"""
+    Apply a quantum channel, specified by a superoperator aka the Lioville or Pauli-Lioville
+    representation, to a state by matrix multiplication.
+
+    Note: you must ensure the state and superoperator are represented in the same basis.
+
+    :param superoperator: a dim**2 by dim**2 matrix.
+    :param state: A dim**2 ndarray which is the vectorized density matrix for the state
+    :return: A dim**2 ndarray which is the vectorized density matrix of state after action of
+        superoperator
+    """
+    dim = int(np.sqrt(np.asarray(superoperator).shape[0]))
+    rows, cols = state.shape
+
+    if dim**2 != rows:
+        raise ValueError("Dimensions of the vectorized state and superoperator are incompatible")
+    return superoperator @ state

forest/benchmarking/tests/test_superoperator_tools.py

@@ -1,7 +1,7 @@
 import numpy as np
 from forest.benchmarking.utils import *
 from forest.benchmarking.superoperator_tools import *
-
+import pytest
 # Test philosophy:
 # Using the by hand calculations found in the docs we check conversion
 # between one qubit channels with one Kraus operator (Hadamard) and two
@@ -72,6 +72,28 @@ def amplitude_damping_choi(p):
                           [ 1,  1,  1, -1],
                           [-1, -1, -1,  1]])
 
+# Single Qubit Pauli Channel
+def one_q_pauli_channel_chi(px,py,pz):
+    p = (px + py + pz)
+    pp_chi = np.asarray([[  1-p,    0,     0,   0],
+                         [    0,   px,     0,   0],
+                         [    0,    0,    py,   0],
+                         [    0,    0,     0,  pz]])
+    return pp_chi
+
+
+# Pauli twirled Amplitude damping channel
+def analytical_pauli_twirl_of_AD_chi(p):
+    # see equation 7 of  https://arxiv.org/pdf/1701.03708.pdf
+    poly1 = (2 + 2*np.sqrt(1 - p) - p) / 4
+    poly2 = p / 4
+    poly3 = (2 - 2*np.sqrt(1 - p) - p) / 4
+    pp_chi = np.asarray([[poly1,    0,     0,     0],
+                         [    0,poly2,     0,     0],
+                         [    0,    0, poly2,     0],
+                         [    0,    0,     0, poly3]])
+    return pp_chi
+
 
 # I \otimes Z channel or gate (two qubits)
 two_qubit_paulis = n_qubit_pauli_basis(2)
@@ -92,6 +114,10 @@ def amplitude_damping_choi(p):
 rho_out = np.matmul(np.matmul(Ad0, ONE_STATE), Ad0.transpose().conj()) + \
           np.matmul(np.matmul(Ad1, ONE_STATE), Ad1.transpose().conj())
 
+#===================================================================================================
+# Test superoperator conversion tools
+#===================================================================================================
+
 def test_vec():
     A = np.asarray([[1, 2], [3, 4]])
     B = np.asarray([[1, 2, 5], [3, 4, 6]])
@@ -239,3 +265,62 @@ def test_choi_pl_bijectivity():
     h_superop = kraus2superop(HADAMARD)
     assert np.allclose(choi2superop(choi2superop(h_choi)), h_choi)
     assert np.allclose(superop2choi(superop2choi(h_superop)), h_superop)
+
+
+# ===================================================================================================
+# Test channel and superoperator approximation tools
+# ===================================================================================================
+
+
+def test_pauli_twirl_of_pauli_channel():
+    # diagonal channel so should not change anything
+    px = np.random.rand()
+    py = np.random.rand()
+    pz = np.random.rand()
+    pauli_chan_chi_matrix = one_q_pauli_channel_chi(px, py, pz)
+    pauli_twirled_chi_matrix = pauli_twirl_chi_matrix(pauli_chan_chi_matrix)
+    assert np.allclose(pauli_chan_chi_matrix, pauli_twirled_chi_matrix)
+
+
+def test_pauli_twirl_of_amp_damp():
+    p = np.random.rand()
+    ana_chi = analytical_pauli_twirl_of_AD_chi(p)
+    num_chi = pauli_twirl_chi_matrix(amplitude_damping_chi(p))
+    assert np.allclose(ana_chi, num_chi)
+
+# ==================================================================================================
+# Test apply channel
+# ==================================================================================================
+
+
+def test_apply_kraus_ops_2_state():
+    AD_kraus = amplitude_damping_kraus(0.1)
+    assert np.allclose(rho_out, apply_kraus_ops_2_state(AD_kraus, ONE_STATE))
+
+
+def test_apply_non_square_kraus_ops_2_state():
+    Id = np.eye(2)
+    bra_zero = np.asarray([[1], [0]])
+    bra_one = np.asarray([[0], [1]])
+    state_one = np.kron(Id / 2, ONE_STATE)
+    state_zero = np.kron(Id / 2, ZERO_STATE)
+    Kraus1 = np.kron(Id, bra_one.transpose())
+    Kraus0 = np.kron(Id, bra_zero.transpose())
+    assert np.allclose(apply_kraus_ops_2_state(Kraus0, state_zero), Id / 2)
+    assert np.allclose(apply_kraus_ops_2_state(Kraus1, state_one), Id / 2)
+    assert np.allclose(apply_kraus_ops_2_state(Kraus0, state_one), 0)
+
+
+def test_apply_choi_matrix_2_state():
+    choi = amplitude_damping_choi(0.1)
+    assert np.allclose(rho_out, apply_choi_matrix_2_state(choi, ONE_STATE))
+
+
+def test_apply_lioville_matrix_2_state():
+    super = amplitude_damping_super(0.1)
+    try:
+            apply_lioville_matrix_2_state(super, ONE_STATE)
+    except ValueError as e:
+        assert str(e) == "Dimensions of the vectorized state and superoperator are incompatible"
+    rho_out_s = apply_lioville_matrix_2_state(super, vec(ONE_STATE))
+    assert np.allclose(vec(rho_out), rho_out_s)