Add CPMG to T2 experiment (#3255)

* Add CPMG to T2 experiment

- This adds Carr-Purcell-Meiboom-Gill pulse sequences to the T2
experiment.
- This sequence has the following sequence:
π/2, t, π, 2t, π, ... 2t, π, t

Since there are two variables (N pulses and t delay), this allows
the user to modify either variable or both to create a 2d scan.

In order to accomplish this efficiently, the pi-pulses are encoded as
parameterized X gates, e.g. X ** sympy.symbol('pulse_3') which is
only turned on for more than 3 pulses.

- Fixes #3001

Author: dstrain115

cirq/experiments/t2_decay_experiment.py

@@ -13,7 +13,7 @@
 # limitations under the License.
 import enum
 
-from typing import Any, Optional, TYPE_CHECKING
+from typing import Any, List, Optional, TYPE_CHECKING, Union
 
 import pandas as pd
 import sympy
@@ -35,17 +35,18 @@ class ExperimentType(enum.Enum):
 _T2_COLUMNS = ['delay_ns', 0, 1]
 
 
-def t2_decay(
-        sampler: work.Sampler,
-        *,
-        qubit: devices.GridQubit,
-        experiment_type: 'ExperimentType' = ExperimentType.RAMSEY,
-        num_points: int,
-        max_delay: 'cirq.DURATION_LIKE',
-        min_delay: 'cirq.DURATION_LIKE' = None,
-        repetitions: int = 1000,
-        delay_sweep: Optional[study.Sweep] = None,
-) -> 'cirq.experiments.T2DecayResult':
+def t2_decay(sampler: work.Sampler,
+             *,
+             qubit: devices.GridQubit,
+             experiment_type: 'ExperimentType' = ExperimentType.RAMSEY,
+             num_points: int,
+             max_delay: 'cirq.DURATION_LIKE',
+             min_delay: 'cirq.DURATION_LIKE' = None,
+             repetitions: int = 1000,
+             delay_sweep: Optional[study.Sweep] = None,
+             num_pulses: List[int] = None
+            ) -> Union['cirq.experiments.T2DecayResult',
+                       List['cirq.experiments.T2DecayResult']]:
     """Runs a t2 transverse relaxation experiment.
 
     Initializes a qubit into a superposition state, evolves the system using
@@ -58,7 +59,7 @@ def t2_decay(
     prepared with a square root Y gate (`cirq.Y ** 0.5`) and then waits for
     a variable amount of time.  After this time, it will do basic state
     tomography to measure the expectation of the Pauli-X and Pauli-Y operators
-    by performing either a `cirq.Y ** -0.5` or `cirq.X ** -0.5`.  The square of
+    by performing either a `cirq.Y ** -0.5` or `cirq.X ** 0.5`.  The square of
     these two measurements is summed to determine the length of the Bloch
     vector. This experiment measures the phase decoherence of the system under
     free evolution.
@@ -68,10 +69,41 @@ def t2_decay(
     However, during the mid-point of the delay time being measured, a pi-pulse
     (`cirq.X`) gate will be applied to cancel out inhomogeneous dephasing.
     The same method of measuring the final state as Ramsey experiment is applied
-    after the second half of the delay period.
-
-    CPMG, or the Carr-Purcell-Meiboom-Gill sequence, is currently not
-    implemented.
+    after the second half of the delay period.  See the animation on the wiki
+    page https://en.wikipedia.org/wiki/Spin_echo for a visual illustration
+    of this experiment.
+
+    CPMG, or the Carr-Purcell-Meiboom-Gill sequence, involves using a sqrt(Y)
+    followed by a sequence of pi pulses (X gates) in a specific timing pattern:
+
+        π/2, t, π, 2t, π, ... 2t, π, t
+
+    The first pulse, a sqrt(Y) gate, will put the qubit's state on the Bloch
+    equator.  After a delay, successive X gates will refocus dehomogenous
+    phase effects by causing them to precess in opposite directions and
+    averaging their effects across the entire pulse train.
+
+    This pulse pattern has two variables that can be adjusted.  The first,
+    denoted as 't' in the above sequence, is delay, which can be specified
+    with `delay_min` and `delay_max` or by using a `delay_sweep`, similar to
+    the other experiments.  The second variable is the number of pi pulses
+    (X gates).  This can be specified as a list of integers using the
+    `num_pulses` parameter.  If multiple different pulses are specified,
+    the data will be presented in a data frame with two
+    indices (delay_ns and num_pulses).
+
+    See the following reference for more information about CPMG pulse trains:
+    Meiboom, S., and D. Gill, “Modified spin-echo method for measuring nuclear
+    relaxation times”, Rev. Sci. Inst., 29, 688–691 (1958).
+    https://doi.org/10.1063/1.1716296
+
+    Note that interpreting T2 data is fairly tricky and subtle, as it can
+    include other effects that need to be accounted for.  For instance,
+    amplitude damping (T1) will present as T2 noise and needs to be
+    appropriately compensated for to find a true measure of T2.  Due to this
+    subtlety and lack of standard way to interpret the data, the fitting
+    of the data to an exponential curve and the extrapolation of an actual
+    T2 time value is left as an exercise to the reader.
 
     Args:
         sampler: The quantum engine or simulator to run the circuits.
@@ -87,6 +119,8 @@ def t2_decay(
              of nanoseconds.  If specified, this will override the max_delay and
              min_delay parameters.  If not specified, the experiment will sweep
              from min_delay to max_delay with linear steps.
+        num_pulses: For CPMG, a list of the number of pulses to use.
+             If multiple pulses are specified, each will be swept on.
     Returns:
         A T2DecayResult object that stores and can plot the data.
     """
@@ -100,11 +134,14 @@ def t2_decay(
         raise ValueError('max_delay < min_delay')
     if min_delay_dur < 0:
         raise ValueError('min_delay < 0')
+    if num_pulses and experiment_type != ExperimentType.CPMG:
+        raise ValueError('num_pulses is only valid for CPMG experiments.')
 
     # Initialize values used in sweeps
     delay_var = sympy.Symbol('delay_ns')
     inv_x_var = sympy.Symbol('inv_x')
     inv_y_var = sympy.Symbol('inv_y')
+    max_pulses = max(num_pulses) if num_pulses else 0
 
     if not delay_sweep:
         delay_sweep = study.Linspace(delay_var,
@@ -124,60 +161,122 @@ def t2_decay(
         circuit = circuits.Circuit(
             ops.Y(qubit)**0.5,
             ops.WaitGate(value.Duration(nanos=delay_var))(qubit),
-            ops.X(qubit)**inv_x_var,
-            ops.Y(qubit)**inv_y_var,
-            ops.measure(qubit, key='output'),
-        )
-        tomography_sweep = study.Zip(
-            study.Points('inv_x', [0.0, -0.5]),
-            study.Points('inv_y', [-0.5, 0.0]),
-        )
-        sweep = study.Product(delay_sweep, tomography_sweep)
-    elif experiment_type == ExperimentType.HAHN_ECHO:
-        # Hahn / Spin Echo T2 experiment
-        # Use sqrt(Y) to flip to the equator.
-        # Evolve the state for half the given amount of delay time
-        # Flip the state using an X gate
-        # Evolve the state for half the given amount of delay time
-        # Then measure the state in both X and Y bases.
-
-        circuit = circuits.Circuit(
-            ops.Y(qubit)**0.5,
-            ops.WaitGate(value.Duration(nanos=0.5 * delay_var))(qubit),
-            ops.X(qubit),
-            ops.WaitGate(value.Duration(nanos=0.5 * delay_var))(qubit),
-            ops.X(qubit)**inv_x_var,
-            ops.Y(qubit)**inv_y_var,
-            ops.measure(qubit, key='output'),
         )
-        tomography_sweep = study.Zip(
-            study.Points('inv_x', [0.0, 0.5]),
-            study.Points('inv_y', [-0.5, 0.0]),
-        )
-        sweep = study.Product(delay_sweep, tomography_sweep)
     else:
-        raise ValueError(f'Experiment type {experiment_type} not supported')
+        if experiment_type == ExperimentType.HAHN_ECHO:
+            # Hahn / Spin Echo T2 experiment
+            # Use sqrt(Y) to flip to the equator.
+            # Evolve the state for the given amount of delay time
+            # Flip the state using an X gate
+            # Evolve the state for the given amount of delay time
+            # Then measure the state in both X and Y bases.
+            num_pulses = [0]
+            # This is equivalent to a CPMG experiment with zero pulses
+            # and will follow the same code path.
+
+        # Carr-Purcell-Meiboom-Gill sequence.
+        # Performs the following sequence
+        # π/2 - wait(t) - π - wait(2t) - ... - π - wait(t)
+        # There will be N π pulses (X gates)
+        # where N sweeps over the values of num_pulses
+        #
+        if not num_pulses:
+            raise ValueError('At least one value must be given '
+                             'for num_pulses in a CPMG experiment')
+        circuit = _cpmg_circuit(qubit, delay_var, max_pulses)
+
+    # Add simple state tomography
+    circuit.append(ops.X(qubit)**inv_x_var)
+    circuit.append(ops.Y(qubit)**inv_y_var)
+    circuit.append(ops.measure(qubit, key='output'))
+    tomography_sweep = study.Zip(
+        study.Points('inv_x', [0.0, 0.5]),
+        study.Points('inv_y', [-0.5, 0.0]),
+    )
+
+    if num_pulses and max_pulses > 0:
+        pulse_sweep = _cpmg_sweep(num_pulses)
+        sweep = study.Product(delay_sweep, pulse_sweep, tomography_sweep)
+    else:
+        sweep = study.Product(delay_sweep, tomography_sweep)
 
     # Tabulate measurements into a histogram
     results = sampler.sample(circuit, params=sweep, repetitions=repetitions)
 
-    y_basis_measurements = results[abs(results.inv_y) > 0]
-    x_basis_measurements = results[abs(results.inv_x) > 0]
-    x_basis_tabulation = pd.crosstab(x_basis_measurements.delay_ns,
-                                     x_basis_measurements.output).reset_index()
-    y_basis_tabulation = pd.crosstab(y_basis_measurements.delay_ns,
-                                     y_basis_measurements.output).reset_index()
+    y_basis_measurements = results[abs(results.inv_y) > 0].copy()
+    x_basis_measurements = results[abs(results.inv_x) > 0].copy()
 
-    # If all measurements are 1 or 0, fill in the missing column with all zeros.
-    for tab in [x_basis_tabulation, y_basis_tabulation]:
-        for col_index, name in [(1, 0), (2, 1)]:
-            if name not in tab:
-                tab.insert(col_index, name, [0] * tab.shape[0])
+    if num_pulses and len(num_pulses) > 1:
+        cols = tuple(f'pulse_{t}' for t in range(max_pulses))
+        x_basis_measurements[
+            'num_pulses'] = x_basis_measurements.loc[:, cols].sum(axis=1)
+        y_basis_measurements[
+            'num_pulses'] = y_basis_measurements.loc[:, cols].sum(axis=1)
+
+    x_basis_tabulation = _create_tabulation(x_basis_measurements)
+    y_basis_tabulation = _create_tabulation(y_basis_measurements)
 
     # Return the results in a container object
     return T2DecayResult(x_basis_tabulation, y_basis_tabulation)
 
 
+def _create_tabulation(measurements: pd.DataFrame) -> pd.DataFrame:
+    """Returns a sum of 0 and 1 results per index from a list of measurements.
+    """
+    if 'num_pulses' in measurements.columns:
+        cols = [measurements.delay_ns, measurements.num_pulses]
+    else:
+        cols = [measurements.delay_ns]
+    tabulation = pd.crosstab(cols, measurements.output).reset_index()
+    # If all measurements are 1 or 0, fill in the missing column with all zeros.
+    for col_index, name in [(1, 0), (2, 1)]:
+        if name not in tabulation:
+            tabulation.insert(col_index, name, [0] * tabulation.shape[0])
+    return tabulation
+
+
+def _cpmg_circuit(qubit: devices.GridQubit, delay_var: sympy.Symbol,
+                  max_pulses: int) -> 'cirq.Circuit':
+    """Creates a CPMG circuit for a given qubit.
+
+    The circuit will look like:
+
+      sqrt(Y) - wait(delay_var) - X - wait(2*delay_var) - ... - wait(delay_var)
+
+    with max_pulses number of X gates.
+
+    The X gates are paramterizd by 'pulse_N' symbols so that pulses can be
+    turned on and off.  This is done to combine circuits with different pulses
+    into the same paramterized circuit.
+    """
+    circuit = circuits.Circuit(
+        ops.Y(qubit)**0.5,
+        ops.WaitGate(value.Duration(nanos=delay_var))(qubit), ops.X(qubit))
+    for n in range(max_pulses):
+        pulse_n_on = sympy.Symbol(f'pulse_{n}')
+        circuit.append(
+            ops.WaitGate(value.Duration(nanos=2 * delay_var *
+                                        pulse_n_on))(qubit))
+        circuit.append(ops.X(qubit)**pulse_n_on)
+    circuit.append(ops.WaitGate(value.Duration(nanos=delay_var))(qubit))
+    return circuit
+
+
+def _cpmg_sweep(num_pulses: List[int]):
+    """Returns a sweep for a circuit created by _cpmg_circuit.
+
+    The circuit in _cpmg_circuit parameterizes the pulses, so this function
+    fills in the parameters for each pulse.  For instance, if we want 3 pulses,
+    pulse_0, pulse_1, and pulse_2 should be 1 and the rest of the pulses should
+    be 0.
+    """
+    pulse_points = []
+    for n in range(max(num_pulses)):
+        pulse_points.append(
+            study.Points(f'pulse_{n}', [1 if p > n else 0 for p in num_pulses]))
+    return study.Zip(*pulse_points)
+
+
 class T2DecayResult:
     """Results from a T2 decay experiment.
 
@@ -206,21 +305,29 @@ def __init__(self, x_basis_data: pd.DataFrame, y_basis_data: pd.DataFrame):
         self._expectation_pauli_x = self._expectation(x_basis_data)
         self._expectation_pauli_y = self._expectation(y_basis_data)
 
-    def _expectation(self, data) -> pd.DataFrame:
+    def _expectation(self, data: pd.DataFrame) -> pd.DataFrame:
         """Calculates the expected value of the Pauli operator.
 
         Assuming that the data is measured in the Pauli basis of the operator,
         then the expectation of the Pauli operator would be +1 if the
         measurement is all ones and -1 if the measurement is all zeros.
 
         Returns:
-            Data frame with two columns 'delay_ns' and 'value'
+            Data frame with columns 'delay_ns', 'num_pulses' and 'value'
+            The num_pulses column will only exist if multiple pulses
+            were requestd in the T2 experiment.
         """
-        xs = data['delay_ns']
+        delay = data['delay_ns']
         ones = data[1]
         zeros = data[0]
         pauli_expectation = (2 * (ones / (ones + zeros))) - 1.0
-        return pd.DataFrame({'delay_ns': xs, 'value': pauli_expectation})
+        if 'num_pulses' in data.columns:
+            return pd.DataFrame({
+                'delay_ns': delay,
+                'num_pulses': data['num_pulses'],
+                'value': pauli_expectation
+            })
+        return pd.DataFrame({'delay_ns': delay, 'value': pauli_expectation})
 
     @property
     def expectation_pauli_x(self) -> pd.DataFrame: