Repository: rigetticomputing/pyquil
Pull request #286 (2018-02-02 17:10:05)

Feature/readout noise compensation
Implement #285 

ntezak: @ncrubin The API is complete, so far only the docstrings explain the different components. I could add an example to the `ReadoutNoise.ipynb` notebook and the docs, what do you think, Nick?

Patch:
diff --git a/pyquil/noise.py b/pyquil/noise.py
index c8e535e..e9475a3 100644
--- a/pyquil/noise.py
+++ b/pyquil/noise.py
@@ -513,3 +513,141 @@ def add_decoherence_noise(prog, T1=30e-6, T2=30e-6, gate_time_1q=50e-9, gate_tim
         ro_fidelity=ro_fidelity
     )
     return _apply_noise_model(prog, noise_model)
+
+
+def _bitstring_probs_by_qubit(p):
+    """
+    Ensure that an array ``p`` with bitstring probabilities has a separate axis for each qubit such
+    that ``p[i,j,...,k]`` gives the estimated probability of bitstring ``ij...k``.
+
+    This should not allocate much memory if ``p`` is already in ``C``-contiguous order (row-major).
+
+    :param np.array p: An array that enumerates bitstring probabilities. When flattened out
+        ``p = [p_00...0, p_00...1, ...,p_11...1]``. The total number of elements must therefore be a
+        power of 2.
+    :return: A reshaped view of ``p`` with a separate length-2 axis for each bit.
+    """
+    p = np.asarray(p, order="C")
+    num_qubits = int(round(np.log2(p.size)))
+    return p.reshape((2,) * num_qubits)
+
+
+def estimate_bitstring_probs(results):
+    """
+    Given an array of single shot results estimate the probability distribution over all bitstrings.
+
+    :param np.array results: A 2d array where the outer axis iterates over shots
+        and the inner axis over bits.
+    :return: An array with as many axes as there are qubit and normalized such that it sums to one.
+        ``p[i,j,...,k]`` gives the estimated probability of bitstring ``ij...k``.
+    :rtype: np.array
+    """
+    nshots, nq = np.shape(results)
+    outcomes = np.array([int("".join(map(str, r)), 2) for r in results])
+    probs = np.histogram(outcomes, bins=np.arange(-.5, 2 ** nq, 1))[0] / float(nshots)
+    return _bitstring_probs_by_qubit(probs)
+
+
+_CHARS = 'klmnopqrstuvwxyzabcdefgh0123456789'
+
+
+def _apply_local_transforms(p, ts):
+    """
+    Given a 2d array of single shot results (outer axis iterates over shots, inner axis over bits)
+    and a list of assignment probability matrices (one for each bit in the readout, ordered like
+    the inner axis of results) apply local 2x2 matrices to each bit index.
+
+    :param np.array p: An array that enumerates a function indexed by bitstrings::
+
+            f(ijk...) = p[i,j,k,...]
+
+    :param Sequence[np.array] ts: A sequence of 2x2 transform-matrices, one for each bit.
+    :return: ``p_transformed`` an array with as many dimensions as there are bits with the result of
+        contracting p along each axis by the corresponding bit transformation.
+
+            p_transformed[ijk...] = f'(ijk...) = sum_lmn... ts[0][il] ts[1][jm] ts[2][kn] f(lmn...)
+
+    :rtype: np.array
+    """
+    p_corrected = _bitstring_probs_by_qubit(p)
+    nq = p_corrected.ndim
+    for idx, trafo_idx in enumerate(ts):
+
+        # this contraction pattern looks like
+        # 'ij,abcd...jklm...->abcd...iklm...' so it properly applies a "local"
+        # transformation to a single tensor-index without changing the order of
+        # indices
+        einsum_pat = ('ij,' + _CHARS[:idx] + 'j' + _CHARS[idx:nq - 1] +
+                      '->' + _CHARS[:idx] + 'i' + _CHARS[idx:nq - 1])
+        p_corrected = np.einsum(einsum_pat, trafo_idx, p_corrected)
+
+    return p_corrected
+
+
+def corrupt_bitstring_probs(p, assignment_probabilities):
+    """
+    Given a 2d array of true bitstring probabilities (outer axis iterates over shots, inner axis
+    over bits) and a list of assignment probability matrices (one for each bit in the readout,
+    ordered like the inner axis of results) compute the corrupted probabilities.
+
+    :param np.array p: An array that enumerates bitstring probabilities. When
+        flattened out ``p = [p_00...0, p_00...1, ...,p_11...1]``. The total number of elements must
+        therefore be a power of 2. The canonical shape has a separate axis for each qubit, such that
+        ``p[i,j,...,k]`` gives the estimated probability of bitstring ``ij...k``.
+    :param List[np.array] assignment_probabilities: A list of assignment probability matrices
+        per qubit. Each assignment probability matrix is expected to be of the form::
+
+            [[p00 p01]
+             [p10 p11]]
+
+    :return: ``p_corrected`` an array with as many dimensions as there are qubits that contains
+        the noisy-readout-corrected estimated probabilities for each measured bitstring, i.e.,
+        ``p[i,j,...,k]`` gives the estimated probability of bitstring ``ij...k``.
+    :rtype: np.array
+    """
+    return _apply_local_transforms(p, assignment_probabilities)
+
+
+def correct_bitstring_probs(p, assignment_probabilities):
+    """
+    Given a 2d array of corrupted bitstring probabilities (outer axis iterates over shots, inner
+    axis over bits) and a list of assignment probability matrices (one for each bit in the readout)
+    compute the corrected probabilities.
+
+    :param np.array p: An array that enumerates bitstring probabilities. When
+        flattened out ``p = [p_00...0, p_00...1, ...,p_11...1]``. The total number of elements must
+        therefore be a power of 2. The canonical shape has a separate axis for each qubit, such that
+        ``p[i,j,...,k]`` gives the estimated probability of bitstring ``ij...k``.
+    :param List[np.array] assignment_probabilities: A list of assignment probability matrices
+        per qubit. Each assignment probability matrix is expected to be of the form::
+
+            [[p00 p01]
+             [p10 p11]]
+
+    :return: ``p_corrected`` an array with as many dimensions as there are qubits that contains
+        the noisy-readout-corrected estimated probabilities for each measured bitstring, i.e.,
+        ``p[i,j,...,k]`` gives the estimated probability of bitstring ``ij...k``.
+    :rtype: np.array
+    """
+    return _apply_local_transforms(p, (np.linalg.inv(ap) for ap in assignment_probabilities))
+
+
+def bitstring_probs_to_z_moments(p):
+    """
+    Convert between bitstring probabilities and joint Z moment expectations.
+
+    :param np.array p: An array that enumerates bitstring probabilities. When
+        flattened out ``p = [p_00...0, p_00...1, ...,p_11...1]``. The total number of elements must
+        therefore be a power of 2. The canonical shape has a separate axis for each qubit, such that
+        ``p[i,j,...,k]`` gives the estimated probability of bitstring ``ij...k``.
+    :return: ``z_moments``, an np.array with one length-2 axis per qubit which contains the
+        expectations of all monomials in ``{I, Z_0, Z_1, ..., Z_{n-1}}``. The expectations of each
+        monomial can be accessed via::
+
+            <Z_0^j_0 Z_1^j_1 ... Z_m^j_m> = z_moments[j_0,j_1,...,j_m]
+
+    :rtype: np.array
+    """
+    zmat = np.array([[1, 1],
+                     [1, -1]])
+    return _apply_local_transforms(p, (zmat for _ in range(p.ndim)))


Test patch:
diff --git a/pyquil/tests/test_noise.py b/pyquil/tests/test_noise.py
index ec9f25a..ff94926 100644
--- a/pyquil/tests/test_noise.py
+++ b/pyquil/tests/test_noise.py
@@ -6,8 +6,10 @@ import pytest
 from pyquil.gates import CZ, RZ, RX, I, H
 from pyquil.noise import (damping_kraus_map, dephasing_kraus_map, tensor_kraus_maps,
                           _get_noisy_names, _get_program_gates, _decoherence_noise_model,
-                          add_decoherence_noise, combine_kraus_maps, damping_after_dephasing, INFINITY,
-                          _apply_noise_model, _noise_model_program_header, KrausModel, NoiseModel)
+                          add_decoherence_noise, combine_kraus_maps, damping_after_dephasing,
+                          INFINITY, _apply_noise_model, _noise_model_program_header, KrausModel,
+                          NoiseModel, corrupt_bitstring_probs, correct_bitstring_probs,
+                          estimate_bitstring_probs, bitstring_probs_to_z_moments)
 from pyquil.quil import Pragma, Program
 
 
@@ -141,3 +143,37 @@ def test_noise_model():
     assert nm == NoiseModel.from_dict(nm.to_dict())
     assert nm.gates_by_name("I") == [km1]
     assert nm.gates_by_name("RX") == [km2]
+
+
+def test_readout_compensation():
+    np.random.seed(1234124)
+    p = np.random.rand(2, 2, 2, 2, 2, 2)
+    p /= p.sum()
+
+    aps = [np.eye(2) + .2 * (np.random.rand(2, 2) - 1) for _ in range(p.ndim)]
+    for ap in aps:
+        ap.flat[ap.flat < 0] = 0.
+        ap /= ap.sum()
+        assert np.alltrue(ap >= 0)
+
+    assert np.alltrue(p >= 0)
+
+    p_corrupted = corrupt_bitstring_probs(p, aps)
+    p_restored = correct_bitstring_probs(p_corrupted, aps)
+    assert np.allclose(p, p_restored)
+
+    results = [[0, 0, 0]] * 100 + [[0, 1, 1]] * 200
+    p1 = estimate_bitstring_probs(results)
+    assert np.isclose(p1[0, 0, 0], 1. / 3.)
+    assert np.isclose(p1[0, 1, 1], 2. / 3.)
+    assert np.isclose(p1.sum(), 1.)
+
+    zm = bitstring_probs_to_z_moments(p1)
+    assert np.isclose(zm[0, 0, 0], 1)
+    assert np.isclose(zm[1, 0, 0], 1)
+    assert np.isclose(zm[0, 1, 0], -1. / 3)
+    assert np.isclose(zm[0, 0, 1], -1. / 3)
+    assert np.isclose(zm[0, 1, 1], 1.)
+    assert np.isclose(zm[1, 1, 0], -1. / 3)
+    assert np.isclose(zm[1, 0, 1], -1. / 3)
+    assert np.isclose(zm[1, 1, 1], 1.)