Update T1 experiment (quantumlib#6487)

* T1 experiment: add A and B constants to fit, make wait times log-spaced

* lowercase variable names

* update tests

* update tests

* update tests

* update tests

* update tests

* update tests

* update tests

* informative assert statements in tests

* use float wait time

* typecheck

* nits

cirq/experiments/t1_decay_experiment.py

@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import Any, Optional, TYPE_CHECKING
+from typing import Any, Optional, Sequence, TYPE_CHECKING, cast
 
 import warnings
 import pandas as pd
@@ -77,7 +77,12 @@ def t1_decay(
 
     var = sympy.Symbol('delay_ns')
 
-    sweep = study.Linspace(var, start=min_delay_nanos, stop=max_delay_nanos, length=num_points)
+    if min_delay_nanos == 0:
+        min_delay_nanos = 0.4
+    sweep_vals_ns = np.unique(
+        np.round(np.logspace(np.log10(min_delay_nanos), np.log10(max_delay_nanos), num_points))
+    )
+    sweep = study.Points(var, cast(Sequence[float], sweep_vals_ns))
 
     circuit = circuits.Circuit(
         ops.X(qubit), ops.wait(qubit, nanos=var), ops.measure(qubit, key='output')
@@ -118,8 +123,8 @@ def data(self) -> pd.DataFrame:
     def constant(self) -> float:
         """The t1 decay constant."""
 
-        def exp_decay(x, t1):
-            return np.exp(-x / t1)
+        def exp_decay(x, t1, a, b):
+            return a * np.exp(-x / t1) + b
 
         xs = self._data['delay_ns']
         ts = self._data['true_count']
@@ -132,8 +137,8 @@ def exp_decay(x, t1):
 
         # Fit to exponential decay to find the t1 constant
         try:
-            popt, _ = optimize.curve_fit(exp_decay, xs, probs, p0=[t1_guess])
-            t1 = popt[0]
+            self.popt, _ = optimize.curve_fit(exp_decay, xs, probs, p0=[t1_guess, 1.0, 0.0])
+            t1 = self.popt[0]
             return t1
         except RuntimeError:
             warnings.warn("Optimal parameters could not be found for curve fit", RuntimeWarning)
@@ -166,7 +171,9 @@ def plot(
         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')
+            t1 = self.constant
+            t1, a, b = self.popt
+            ax.plot(xs, a * np.exp(-xs / t1) + b, label='curve fit')
             plt.legend()
 
         ax.set_xlabel(r"Delay between initialization and measurement (nanoseconds)")

cirq/experiments/t1_decay_experiment_test.py

@@ -53,15 +53,15 @@ def noisy_moment(self, moment, system_qubits):
         repetitions=10,
         max_delay=cirq.Duration(nanos=500),
     )
-    results.plot()
+    results.plot(include_fit=True)
 
 
 def test_result_eq():
     eq = cirq.testing.EqualsTester()
     eq.make_equality_group(
         lambda: cirq.experiments.T1DecayResult(
             data=pd.DataFrame(
-                columns=['delay_ns', 'false_count', 'true_count'], index=[0], data=[[100.0, 2, 8]]
+                columns=['delay_ns', 'false_count', 'true_count'], index=[0], data=[[100, 2, 8]]
             )
         )
     )
@@ -103,7 +103,7 @@ def noisy_moment(self, moment, system_qubits):
         data=pd.DataFrame(
             columns=['delay_ns', 'false_count', 'true_count'],
             index=range(4),
-            data=[[100.0, 0, 10], [400.0, 0, 10], [700.0, 10, 0], [1000.0, 10, 0]],
+            data=[[100.0, 0, 10], [215.0, 0, 10], [464.0, 0, 10], [1000.0, 10, 0]],
         )
     )
 
@@ -117,13 +117,14 @@ def test_all_on_results():
         min_delay=cirq.Duration(nanos=100),
         max_delay=cirq.Duration(micros=1),
     )
-    assert results == cirq.experiments.T1DecayResult(
+    desired = cirq.experiments.T1DecayResult(
         data=pd.DataFrame(
             columns=['delay_ns', 'false_count', 'true_count'],
             index=range(4),
-            data=[[100.0, 0, 10], [400.0, 0, 10], [700.0, 0, 10], [1000.0, 0, 10]],
+            data=[[100.0, 0, 10], [215.0, 0, 10], [464.0, 0, 10], [1000.0, 0, 10]],
         )
     )
+    assert results == desired, f'{results.data=} {desired.data=}'
 
 
 def test_all_off_results():
@@ -135,13 +136,14 @@ def test_all_off_results():
         min_delay=cirq.Duration(nanos=100),
         max_delay=cirq.Duration(micros=1),
     )
-    assert results == cirq.experiments.T1DecayResult(
+    desired = 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]],
+            data=[[100.0, 10, 0], [215.0, 10, 0], [464.0, 10, 0], [1000.0, 10, 0]],
         )
     )
+    assert results == desired, f'{results.data=} {desired.data=}'
 
 
 @pytest.mark.usefixtures('closefigures')
@@ -150,28 +152,14 @@ def test_curve_fit_plot_works():
         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]],
+            data=[[100.0, 6, 4], [215.0, 10, 0], [464.0, 10, 0], [1000.0, 10, 0]],
         )
     )
 
     good_fit.plot(include_fit=True)
 
 
-@pytest.mark.usefixtures('closefigures')
-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)
-
-
-@pytest.mark.parametrize('t1', [200, 500, 700])
+@pytest.mark.parametrize('t1', [200.0, 500.0, 700.0])
 def test_noise_model_continous(t1):
     class GradualDecay(cirq.NoiseModel):
         def __init__(self, t1: float):
@@ -196,10 +184,10 @@ def noisy_moment(self, moment, system_qubits):
     results = cirq.experiments.t1_decay(
         sampler=cirq.DensityMatrixSimulator(noise=GradualDecay(t1)),
         qubit=cirq.GridQubit(0, 0),
-        num_points=4,
+        num_points=10,
         repetitions=10,
-        min_delay=cirq.Duration(nanos=100),
-        max_delay=cirq.Duration(micros=1),
+        min_delay=cirq.Duration(nanos=1),
+        max_delay=cirq.Duration(micros=10),
     )
 
     assert np.isclose(results.constant, t1, 50)