Repository: OxfordIonTrapGroup/oitg

Issue:
Account for spontaneous shelving and deshelving in threshold.py
As per title.

Currently [threshold.py](https://github.com/OxfordIonTrapGroup/oitg/blob/master/oitg/threshold.py) assumes poissonian statistics. This is innacurate for longer readout durations.

Comments:

Patch:
diff --git a/oitg/threshold.py b/oitg/threshold.py
index b74c080..7d4bf82 100644
--- a/oitg/threshold.py
+++ b/oitg/threshold.py
@@ -1,114 +1,156 @@
 import numpy as np
-from scipy.optimize import least_squares
+from scipy.optimize import minimize
 from scipy.stats import poisson
+from scipy.special import gammainc
 
 
-def calc_target_bin_time(bright_rate, dark_rate, p_error_target, p_bright=0.5):
-    """Calculate optimal threshold bin time for target error chance
+def optimise_readout(bright_rate, dark_rate, dark_to_bright_rate=(1 / 1.168),
+                     p_bright=0.5):
+    """Calculate optimal threshold bin time & threshold count
 
     The calculation assumes both bright and dark counts are
     Poisson distributed and gives a threshold minimising the error probability
 
-    The calculation neglects de-shelving and accidental shelving during the
-    readout bin time. It is therefore not suitable for P(error) < 2e-4.
+    The calculation accounts for de-shelving during the readout bin time.
     See thesis: Alice Burrell, 2010
 
     :param bright_rate: expected bright count rate in $s^-1$
     :param dark_rate: expected dark count ratein $s^-1$
-    :param p_error_target: target error probability
+    :param dark_to_bright_rate: dark state decay to become bright in $s^-1$.
+        As this function seeks the global minimum error, the rate must not be
+        zero. Default of 1/1.168 is the Calcium D5/2 shelf decay rate.
     :param p_bright: probability of encountering a bright state (default=0.5)
 
-    :returns: (target_t_bin [s], threshold_rate [$s^-1$])"""
-    def residuals(t_bin):
-        return (calc_p_error(bright_rate, dark_rate, t_bin[0], p_bright=p_bright) -
-                p_error_target)
+    :returns: (target_t_bin [s], threshold_count, p_error)"""
+    thresh = 1
+    error_last = 1.
+    t_bin, error = optimise_t_bin(
+        bright_rate, dark_rate, thresh, dark_to_bright_rate, p_bright)
+    while error_last > error:
+        error_last = error
+        t_bin_last = t_bin
+        thresh += 1
+        t_bin, error = optimise_t_bin(
+            bright_rate, dark_rate, thresh, dark_to_bright_rate, p_bright)
 
-    t_bin_init = 1e-3
-    result = least_squares(residuals,
-                           t_bin_init,
-                           bounds=(0, np.inf),
-                           x_scale=(1e-3, ),
-                           f_scale=(0.1 * p_error_target)**2)
+    return (t_bin_last, thresh - 1, error_last)
 
-    thresh_rate = calc_thresh_rate(bright_rate,
-                                   dark_rate,
-                                   p_bright=p_bright,
-                                   t_bin=result['x'][0])
 
-    return result['x'][0], thresh_rate
+def optimise_treshold(bright_rate, dark_rate, t_bin, dark_to_bright_rate=0.,
+                      p_bright=0.5):
+    """Calculate optimal threshold threshold count for a given bin time
 
+    The calculation accounts for de-shelving during the readout bin time.
+    See thesis: Alice Burrell, 2010
 
-def calc_p_error(bright_rate, dark_rate, t_bin, p_bright=0.5):
-    """Assumes exact threshold count is evaluated as dark
+    :param bright_rate: expected bright count rate in $s^-1$
+    :param dark_rate: expected dark count ratein $s^-1$
+    :param dark_to_bright_rate: dark state decay to become bright in $s^-1$
+    :param p_bright: probability of encountering a bright state (default=0.5)
 
-    The calculation neglects de-shelving and accidental shelving during the
-    readout bin time. It is therefore not suitable for P(error) < 2e-4.
+    :returns: (target_t_bin [s], threshold_count, p_error)"""
+    # de-shelving increases the count rate of dark states (non-poissonian)
+    # The no-de-shelving threshold is therefore a lower bound
+    thresh_min = poisson_optimal_thresh_count(
+        bright_rate * t_bin, dark_rate * t_bin, p_bright)
+    thresh_max = bright_rate * t_bin
+
+    # thresholds are discrete - could implement discreet optimisation, but the
+    # range of values is small in practice. Therefore we can try all options
+    thresh_vec = np.arange(thresh_min, thresh_max + 1, dtype=np.int_)
+    err_vec = np.array([calc_p_error(bright_rate, dark_rate, t_bin,
+                                     i, dark_to_bright_rate,
+                                     p_bright=p_bright)
+                        for i in thresh_vec])
+
+    min_idx = np.argmin(err_vec)
+    return thresh_vec[min_idx], err_vec[min_idx]
+
+
+def optimise_t_bin(bright_rate, dark_rate, thresh_count, dark_to_bright_rate=0.,
+                   p_bright=0.5):
+    """Calculate optimal threshold bin time for a given threshold count
+
+    The calculation accounts for de-shelving during the readout bin time.
     See thesis: Alice Burrell, 2010
 
     :param bright_rate: expected bright count rate in $s^-1$
     :param dark_rate: expected dark count ratein $s^-1$
+    :param dark_to_bright_rate: dark state decay to become bright in $s^-1$
+    :param p_bright: probability of encountering a bright state (default=0.5)
+
+    :returns: (target_t_bin [s], threshold_count, p_error)"""
+    t_bin_init = thresh_count * (0.5 / dark_rate + 0.5 / bright_rate)
+    t_scale = 1e-4
+    err_scale = 1e-6
+
+    def p_error(x):
+        return calc_p_error(bright_rate, dark_rate, x[0] * t_scale,
+                            thresh_count, dark_to_bright_rate,
+                            p_bright=p_bright) / err_scale
+
+    result = minimize(p_error,
+                      np.array([t_bin_init / t_scale]),
+                      bounds=((1 / bright_rate / t_scale, np.inf),))
+
+    return result.x[0] * t_scale, result.fun * err_scale
+
+
+def calc_p_error(bright_rate, dark_rate, t_bin, thresh_count,
+                 dark_to_bright_rate=0., p_bright=0.5):
+    """Calculate error probability for Poisson statistics with de-shelving
+
+    See thesis: Alice Burrell, 2010
+
+    :param bright_rate: expected bright count rate in $s^-1$
+    :param dark_rate: expected dark count rate in $s^-1$
     :param t_bin: integration time in s.
+    :param thresh_count: threshold count for discriminating bright/dark state
+        (Assumes the exact threshold count is evaluated as dark)
+    :param dark_to_bright_rate: dark state decay to become bright in $s^-1$
     :param p_bright: probability of encountering a bright state (default=0.5)
     """
-    thresh_rate = calc_thresh_rate(bright_rate,
-                                   dark_rate,
-                                   p_bright=p_bright,
-                                   t_bin=t_bin)
-    thresh_count = np.ceil(thresh_rate * t_bin).astype(np.int_)
+    def p_n_given_dark(n):
+        "Burrell Eqns 3.2 & 3.6"
+        if dark_to_bright_rate == 0.:
+            return poisson.pmf(n, mu=dark_rate * t_bin)
+
+        rb_tau = (bright_rate - dark_rate) / dark_to_bright_rate
+        eps = bright_rate * t_bin / rb_tau
+
+        x_n = np.exp(-eps) * np.power(rb_tau / (rb_tau - 1), n) / (rb_tau - 1)
+        x_n *= (gammainc(n + 1, eps * (rb_tau - 1))
+                - gammainc(n + 1, dark_rate * t_bin / rb_tau * (rb_tau - 1)))
+
+        d_n = poisson.pmf(n, mu=dark_rate * t_bin) * np.exp(
+            -t_bin * dark_to_bright_rate) + x_n
+        return d_n
 
-    n_vec = np.arange(thresh_count + 1, dtype=np.int_)
+    # up-to & including thresh_count
+    n_vec = np.arange(np.round(thresh_count) + 1, dtype=np.int_)
 
-    p_error = (1 - p_bright) * (1 -
-                                np.sum(poisson.pmf(n_vec[:-1], mu=dark_rate * t_bin)))
+    p_error = (1 - p_bright) * (1 - np.sum(p_n_given_dark(n_vec[:-1])))
     p_error += p_bright * np.sum(poisson.pmf(n_vec, mu=bright_rate * t_bin))
     return p_error
 
 
-def calc_thresh_rate(bright_rate, dark_rate, t_bin=1e-3, p_bright=0.5):
-    """Optimal threshold rate for distinguishing bright and dark states
+def poisson_optimal_thresh_count(mean_bright, mean_dark, p_bright=0.5):
+    """Optimal threshold rate in the absence of de-shelving
 
-    The calculation assumes both bright and dark counts are
-    Poisson distributed and gives a threshold minimising the error probability.
+    The calculation assumes both bright and dark counts are Poisson distributed
+    and gives a threshold minimising the error probability.
 
     The calculation neglects de-shelving and accidental shelving during the
     readout bin time. It is therefore not suitable for P(error) < 2e-4.
     See thesis: Alice Burrell, 2010
 
-    :param bright_rate: expected bright count rate in $s^-1$
-    :param dark_rate: expected dark count ratein $s^-1$
-    :param t_bin: integration time in s. Only relevant for p_bright!=0.5
-        (default=1e-3)
+    :param mean_bright: expected counts if the ion started in a bright state
+    :param mean_dark: expected counts if the ion started in a dark state
     :param p_bright: probability of encountering a bright state
         (default=0.5)
 
-    :returns: threshold_rate /$s^-1$
+    :returns: threshold_count
     """
-    thresh = np.log(p_bright / (1 - p_bright)) / t_bin + bright_rate - dark_rate
-    thresh /= np.log(bright_rate / dark_rate)
-    return thresh
-
-
-if __name__ == "__main__":
-    test_threshhold = True
-    test_calc_p_error = True
-    test_bin_time = True
-
-    bright = 4e4
-    dark = 2e4
-    p_bright = 0.5
-    error_target = 1e-3
-    t_bin = 2e-3
-
-    if test_threshhold:
-        print(calc_thresh_rate(bright, dark, t_bin=t_bin, p_bright=p_bright))
-
-    if test_calc_p_error:
-        print("p_error_calc", calc_p_error(bright, dark, t_bin, p_bright))
-
-    if test_bin_time:
-        t_bin, thresh_rate = calc_target_bin_time(bright,
-                                                  dark,
-                                                  error_target,
-                                                  p_bright=p_bright)
-        print("t_bin, thresh_rate", t_bin, thresh_rate)
-        print("p_error for this bin", calc_p_error(bright, dark, t_bin, p_bright))
+    thresh = np.log(p_bright / (1 - p_bright)) + mean_bright - mean_dark
+    thresh /= np.log(mean_bright / mean_dark)
+    return np.round(thresh)


Test patch:
diff --git a/test/test_threshold.py b/test/test_threshold.py
index 8d9f1dd..af7eb79 100755
--- a/test/test_threshold.py
+++ b/test/test_threshold.py
@@ -3,36 +3,32 @@ from oitg.threshold import *
 
 
 class ThresholdTest(unittest.TestCase):
-    def test_threshold(self):
-        bright = 4e4
-        dark = 2e4
-        t_bin = 2e-3
+    def test_optimise_readout(self):
+        bright_rate = 9.e4
+        dark_rate = 3e4
         p_bright = 0.5
-        calc_thresh_rate(bright, dark, t_bin=t_bin, p_bright=p_bright)
-
-    def test_calc_p_error(self):
-        bright = 4e4
-        dark = 2e4
-        t_bin = 2e-3
-        p_bright = 0.5
-        calc_p_error(bright, dark, t_bin, p_bright)
-
-    def test_bin_time(self):
-        bright = 4e4
-        dark = 2e4
+        dark_to_bright_rate = 1 / 1.168
+        t_bin, threshold, p_error = optimise_readout(
+            bright_rate, dark_rate, dark_to_bright_rate=dark_to_bright_rate,
+            p_bright=p_bright,)
+        self.assertAlmostEqual(0.0010438517809917116, t_bin,
+                               delta=0.01 * 0.0010438517809917116)
+        self.assertEqual(58, threshold)
+        self.assertAlmostEqual(0.000285363868745659, p_error,
+                               delta=0.01 * 0.000285363868745659)
+
+    def test_optimise_t_bin(self):
+        bright_rate = 9.e4
+        dark_rate = 3e4
         p_bright = 0.5
-        error_target = 1e-3
-
-        t_bin, thresh_rate = calc_target_bin_time(bright,
-                                                  dark,
-                                                  error_target,
-                                                  p_bright=p_bright)
-
-        self.assertAlmostEqual(error_target,
-                               calc_p_error(bright, dark, t_bin, p_bright),
-                               delta=0.01 * error_target)
-        self.assertAlmostEqual(
-            thresh_rate, calc_thresh_rate(bright, dark, t_bin=t_bin, p_bright=p_bright))
+        dark_to_bright_rate = 1 / 1.168
+        t_bin, p_error = optimise_t_bin(
+            bright_rate, dark_rate, 58,
+            dark_to_bright_rate=dark_to_bright_rate, p_bright=p_bright,)
+        self.assertAlmostEqual(0.0010438517809917116, t_bin,
+                               delta=0.01 * 0.0010438517809917116)
+        self.assertAlmostEqual(0.000285363868745659, p_error,
+                               delta=0.01 * 0.000285363868745659)
 
 
 if __name__ == '__main__':