T1 constant calculation

cirq-core/cirq/experiments/t1_decay_experiment.py

@@ -14,9 +14,13 @@
 
 from typing import Any, Optional, TYPE_CHECKING
 
+import warnings
 import pandas as pd
 import sympy
 from matplotlib import pyplot as plt
+import numpy as np
+from scipy import optimize
+
 
 from cirq import circuits, ops, study, value
 from cirq._compat import proper_repr
@@ -103,12 +107,40 @@ def data(self) -> pd.DataFrame:
         """A data frame with delay_ns, false_count, true_count columns."""
         return self._data
 
-    def plot(self, ax: Optional[plt.Axes] = None, **plot_kwargs: Any) -> plt.Axes:
+    @property
+    def constant(self) -> float:
+        """The t1 decay constant."""
+
+        def exp_decay(x, t1):
+            return np.exp(-x / t1)
+
+        xs = self._data['delay_ns']
+        ts = self._data['true_count']
+        fs = self._data['false_count']
+        probs = ts / (fs + ts)
+
+        # Find the point closest to probability of 1/e
+        guess_index = np.argmin(np.abs(probs - 1.0 / np.e))
+        t1_guess = xs[guess_index]
+
+        # Fit to exponential decay to find the t1 constant
+        try:
+            popt, _ = optimize.curve_fit(exp_decay, xs, probs, p0=[t1_guess])
+            t1 = popt[0]
+            return t1
+        except RuntimeError:
+            warnings.warn("Optimal parameters could not be found for curve fit", RuntimeWarning)
+            return np.nan
+
+    def plot(
+        self, ax: Optional[plt.Axes] = None, include_fit: bool = False, **plot_kwargs: Any
+    ) -> plt.Axes:
         """Plots the excited state probability vs the amount of delay.
 
         Args:
             ax: the plt.Axes to plot on. If not given, a new figure is created,
                 plotted on, and shown.
+            include_fit: boolean to include exponential decay fit on graph
             **plot_kwargs: Arguments to be passed to 'plt.Axes.plot'.
 
         Returns:
@@ -125,6 +157,11 @@ def plot(self, ax: Optional[plt.Axes] = None, **plot_kwargs: Any) -> plt.Axes:
         fs = self._data['false_count']
 
         ax.plot(xs, ts / (fs + ts), 'ro-', **plot_kwargs)
+
+        if include_fit and not np.isnan(self.constant):
+            ax.plot(xs, np.exp(-xs / self.constant), label='curve fit')
+            plt.legend()
+
         ax.set_xlabel(r"Delay between initialization and measurement (nanoseconds)")
         ax.set_ylabel('Excited State Probability')
         ax.set_title('T1 Decay Experiment Data')

cirq-core/cirq/experiments/t1_decay_experiment_test.py

@@ -18,6 +18,8 @@
 
 import cirq
 
+import numpy as np
+
 
 def test_init_result():
     data = pd.DataFrame(
@@ -165,6 +167,67 @@ def test_all_off_results():
     )
 
 
+@pytest.mark.parametrize(
+    't1, data',
+    [
+        (100, [[100.0, 6, 4], [400.0, 10, 0], [700.0, 10, 0], [1000.0, 10, 0]]),
+        (400, [[100.0, 0, 10], [400.0, 6, 4], [700.0, 10, 0], [1000.0, 10, 0]]),
+        (
+            200,
+            [
+                [time, int(np.exp(-time / 200)), 10 - int(np.exp(-time / 200))]
+                for time in np.linspace(0, 1000, 100)
+            ],
+        ),
+    ],
+)
+def test_constant(t1, data):
+    result = cirq.experiments.T1DecayResult(
+        data=pd.DataFrame(
+            columns=['delay_ns', 'false_count', 'true_count'],
+            index=range(len(data)),
+            data=data,
+        )
+    )
+
+    assert np.isclose(result.constant, t1, 5)
+
+
+def test_curve_fit_plot_works():
+    good_fit = cirq.experiments.T1DecayResult(
+        data=pd.DataFrame(
+            columns=['delay_ns', 'false_count', 'true_count'],
+            index=range(4),
+            data=[
+                [100.0, 6, 4],
+                [400.0, 10, 0],
+                [700.0, 10, 0],
+                [1000.0, 10, 0],
+            ],
+        )
+    )
+
+    good_fit.plot(include_fit=True)
+
+
+def test_curve_fit_plot_warning():
+    bad_fit = cirq.experiments.T1DecayResult(
+        data=pd.DataFrame(
+            columns=['delay_ns', 'false_count', 'true_count'],
+            index=range(4),
+            data=[
+                [100.0, 10, 0],
+                [400.0, 10, 0],
+                [700.0, 10, 0],
+                [1000.0, 10, 0],
+            ],
+        )
+    )
+
+    with pytest.warns(RuntimeWarning, match='Optimal parameters could not be found for curve fit'):
+        bad_fit.plot(include_fit=True)
+
+
 def test_bad_args():
     with pytest.raises(ValueError, match='repetitions <= 0'):
         _ = cirq.experiments.t1_decay(
@@ -197,7 +260,6 @@ def test_bad_args():
 
 
 def test_str():
-
     result = cirq.experiments.T1DecayResult(
         data=pd.DataFrame(
             columns=['delay_ns', 'false_count', 'true_count'],