<filename>tes/analysis.py
"""
Calibration and modeling
"""
import matplotlib.pyplot as plt
import matplotlib as mpl
import numpy as np
from numpy import (
    bitwise_and as and_b, logical_and as and_l, logical_or as or_l,
    logical_not as not_l, bitwise_or as or_b
)
from scipy import signal
from scipy import stats
from scipy.optimize import brentq
from numba import jit
from collections import namedtuple
import qutip as qt
from lmfit import Model
from pathlib import Path
from scipy.optimize import minimize


@jit
def maxima(f, thresh=10):
    """
    Find local maxima of f using rising zero crossings of the gradient of f.

    :param f: function to find maxima for
    :param thresh: Only return a maxima if f[maxima]>thresh

    :return: array of maxima.
    :Notes: Used to find the maxima of the smoothed measurement histogram.
    """
    g = np.gradient(f)
    pos = g > 0
    xings = (pos[:-1] & ~pos[1:]).nonzero()[0] + 1  # indices of the maxima
    return xings[np.where(f[xings] > thresh)]


Guess = namedtuple(
    'Guess', 'hist f bin_c max_i thresholds'
)


def guess_thresholds(data, bins=16000, win_length=200, maxima_threshold=10):
    """
    Guess the photon number thresholds using a smoothed histogram.

    :param ndarray data: the measurement data.
    :param bins: argument to numpy.histogram.
    :param int win_length: Hanning window length used for smoothing.
    :param int maxima_threshold: passed to maxima function.

    :return: (hist, f, bin_c, max_i, thresholds) as a named
            tuple. Where hist is the histogram
            generated by np.histogram(), f is the smoothed histogram, bin_c is a
            ndarray of bin centers, max_i are the indices of the maxima in f,
            and thresholds is a ndarray of threshold guesses.

    :notes: The smoothing is achieved by convolving the Hanning window with a
            histogram generated using numpy.histogram. The guess still needs
            some human sanity checks, especially multiple maxima per photon
            peak. Use plot_guess() and adjust window and maxima_threshold to
            remove them.
    """

    hist, edges = np.histogram(data, bins=bins)
    bin_w = edges[1] - edges[0]
    bin_c = edges[:-1] + bin_w / 2

    # smooth the histogram
    win = signal.hann(win_length)
    f = signal.convolve(hist, win, mode='same') / sum(win)

    max_i = maxima(f, thresh=maxima_threshold)  # indices of the maxima
    m = [0.0] + list(bin_c[max_i])
    t = [m[i - 1] + (m[i] - m[i - 1]) / 2 for i in range(2, len(m))]
    thresholds = np.array([0.0] + t + [m[-1] + (m[-1] - t[-1])*1.10])

    return Guess(hist, f, bin_c, max_i, thresholds)


def plot_guess(
        hist, f, x, max_i, init_t, figsize=None,
        measurement_name=r'\texttt{area}'
):
    if figsize is not None:
        fig = plt.figure(figsize=figsize)
    else:
        fig = plt.figure()
    ax = fig.add_axes([0.12, 0.12, 0.87, 0.87])
    ax.plot(
        x[hist.nonzero()[0]], hist[hist.nonzero()[0]], 's',
        markersize=3, markerfacecolor='none', mew=0.5,
        label='histogram bin', color='gray'
    )
    ax.plot(x, f, 'k', lw=1, label='smoothed histogram')
    ax.plot(
        x[max_i], f[max_i], 'o', label='maxima', markerfacecolor='none', mew=1,
        markersize=5, color='r'
    )
    for t in init_t:
        plt.plot([t, t], [0, max(hist)], ':k', lw=0.5)

    fig.canvas.draw_idle()
    offset = ax.xaxis.get_major_formatter().get_offset()
    ax.xaxis.get_major_formatter().set_useOffset(False)
    if offset != '':
        ax.set_xlabel('{} measurement ({})'.format(measurement_name, offset))
    else:
        ax.set_xlabel('{} measurement'.format(measurement_name))
    ax.set_ylabel('count')
    # ax.set_title('Histogram and initial threshold guess')
    plt.legend()
    return fig


# @jit
def retime(events, tdat, fidx, tidx):
    # TODO generalise
    # remove incorporate tick relative times
    #     events = edat.view(event_dt)
    event_time = events['time']
    event_i = 0
    event_time[0] = 0xFFFF
    # tidx[0] is the first tick period ending with tick 1
    for i in range(len(tidx) - 1):
        # i+1 is the tick number
        changed = fidx[
                  tidx[i]['start'] + 1:tidx[i]['stop'] + 1]['changed'
        ].nonzero()[0]
        if len(changed):
            raise NotImplementedError(
                'tick:{} is not homogeneous'.format(i)
            )
        event_count = int(
            (
                    fidx[tidx[i]['stop']]['payload'] +
                    fidx[tidx[i]['stop']]['length'] -
                    fidx[tidx[i]['start']]['payload']
            ) / (fidx[tidx[i]['start']]['event_size'] * 8)
        )
        #         print(event_count)
        event_i += event_count
        time = int(tdat[i + 1]['time']) + int(event_time[event_i])
        #     print(event_i)
        if time > 0xFFFF:
            time = 0xFFFF
        event_time[event_i] = time
    return event_time


def normalised_pdf(x, params, dist=stats.gamma):
    """
    Convenience function, calculate normalised pdf at each x for dist.

    :param union[numpy.ndarray, int] x: calculate the pdf a each x
    :param iterable params: fitted distribution parameters
    :param dist: type of distribution (see scipy.stats)
    :return: ndarray containing pdf(x)
    """
    return dist.pdf(x, *params[:-1])


def scaled_pdf(x, params, dist=stats.gamma):
    """
    Convenience function, calculate pdf at each x for dist, normalisation is
    given by params[-1].

    :param union[numpy.ndarray, int] x: calculate the pdf a each x
    :param iterable params: fitted distribution parameters
    :param dist: type of distribution (see scipy.stats)
    :return: ndarray containing pdf(x)
    """
    return params[-1]*dist.pdf(x, *params[:-1])


class MixtureModel(
    namedtuple(
        'MixtureModel', 'param_list thresholds log_likelihood '
        'converged dist has_noise_threshold'
     )
):
    """
    Subclass of namedtuple used to represent a mixture model.

    fields:
        :param_list: list of parameters for each distribution in the mixture.
        :thresholds: the intersection of neighbouring distributions in the
                     mixture.
        :zero_loc: Fixed location parameter used to fit the first distribution,
                   or None if the location parameter was fitted.
        :log_likelihood: the log of the likelihood that the model could produce
                         the data used to construct it.
        :converged: boolean indicating that the expectation maximisation
                    algorithm terminated normally when fitting the model.
        :dist: the type of distribution forming the components of the mixture,
               (see scipy.stats).

    """

    __slots__ = ()

    def save(self, filename):
        """
        save model as .npz file.

        :param filename: filename to save as, excluding extension"
        :return: None
        """
        np.savez(
            filename,
            param_list=self.param_list,
            thresholds=self.thresholds,
            # zero_loc=self.zero_loc,
            log_likelihood=self.log_likelihood,
            converged=self.converged,
            dist=self.dist.name,
            has_noise_threshold=self.has_noise_threshold
        )

    @staticmethod
    def load(filename):
        """
        load model from .npz file.

        :param filename: filename excluding extension"
        :return: instance of MixtureModel.
        """
        d = np.load(Path(filename).with_suffix('.npz'))
        return MixtureModel(
            d['param_list'],
            d['thresholds'],
            # d['zero_loc'][()],
            d['log_likelihood'][()],
            d['converged'][()],
            getattr(stats, str(d['dist'])),
            d['has_noise_threshold'][()]
        )

    def _eval(self, x, d=None, scale=False, func='pdf'):
        """
        evaluate the function of the distribution(s) selected by d at x.

        :param x: value(s) where to the function is evaluated.
        :param d: selects element(s) of param_list that parameterise the
                  distributions. If None all distributions in param_list are
                  used.
        :param bool scale: when true return the scaled pdf, the scale is always
                           param[-1].
        :param str func: the name of the distribution function to evaluate
                             (see scipy.stats.rv_continuous)
        :return: ndarray shape (len(d), len(x)) containing the pdf(s) or a
                 single value if only 1 point in 1 pdf is selected.
        """

        if d is None:
            # if vaccum:
            p_list = self.param_list
            # else:
            # p_list = self.param_list[1:]
        else:
            p_list = self.param_list[d]

        try:
            xl = len(x)
        except TypeError:
            xl = 1

        f = getattr(self.dist, func)

        # print(p_list.shape)
        if p_list.shape[0] == 1 or len(p_list.shape) == 1:
            if scale:
                return f(x, *p_list[:-1])*p_list[-1]
            else:
                return f(x, *p_list[:-1])
        else:
            f_evals = np.zeros((len(p_list), xl))
            i = 0
            for p in p_list:
                if scale:
                    f_evals[i, :] = f(x, *p[:-1])*p[-1]
                else:
                    f_evals[i, :] = f(x, *p[:-1])
                i += 1
            return f_evals

    def pdf(self, x, d=None, scale=False):
        """
        evaluate the pdf(s) of the distribution(s) selected by d at x.

        :param x: value(s) where to the function is evaluated.
        :param d: selects element(s) of param_list that parameterise the
                  distributions. If None all distributions in param_list are
                  used.
        :param bool scale: when true return the scaled pdf, the scale is always
                           param[-1].
        :return: ndarray shape (len(d), len(x)) containing the pdf(s) or a
                 single value if only 1 point in 1 distribution is selected.
        """
        return self._eval(x, d, scale=scale, func='pdf')

    def cdf(self, x, d=None, scale=False):
        """
        evaluate the pdf(s) of the distribution(s) selected by d at x.

        :param x: value(s) where to the function is evaluated.
        :param d: selects element(s) of param_list that parameterise the
                  distributions. If None all distributions in param_list are
                  used.
        :param bool scale: when true return the scaled pdf, the scale is always
                           param[-1].
        :return: ndarray shape (len(d), len(x)) containing the pdf(s) or a
                 single value if only 1 point in 1 distribution is selected.
        """
        return self._eval(x, d, scale=scale, func='cdf')


@jit
def expectation_maximisation(
    data, initial_thresholds, has_noise_threshold=False,  # fix_noise=None,
    dist=stats.gamma, tol=0.01, max_iter=30, verbose=True
):
    """
    Fit a mixture of distributions to data.

    :param ndarray data: the measurement data.
    :param initial_thresholds: initial thresholds that divide data into
           individual distributions in the mixture.
    :param bool fix_noise: When True keep threshold[1] fixed.
    :param dist: the type of distribution to use (see scipy.stats)
    :param float tol: termination tolerance for change in log_likelihood.
    :param int max_iter: max iterations of expectation maximisation to use.
    :param bool verbose: print progress during optimisation.
    :param bool normalise: passed to maximise step, thresholds are calculated
           using the intersection of the normalised distributions.

    :return: (param_list, zero_loc, log_likelihood, converged) as a named tuple.
            Where param_list is a list of parameter values for distribution in
            the mixture, zero_loc is the arg used to fit, log_likelihood is the
            log_likelihood of the fit, converged is a bool indicating normal
            termination.

    :note: Uses The Expectation maximisation algorithm with hard assignment of
           responsibilities.

    """
    # trim data
    # if has_noise_threshold:
    #     fix_noise = initial_thresholds[1]
    # else:
    fix_noise = None

    fit_data = data[
        and_l(data > initial_thresholds[0], data <= initial_thresholds[-1])
    ]
    thresholds = np.copy(initial_thresholds)
    # last_threshold = initial_thresholds[-1]
    # data = data[data <= last_threshold]

    i = 1
    if verbose:
        print(
            'Starting expectation maximisation, {} {} distributions'
            .format(len(initial_thresholds)-1, dist.name)
        )
        print('Iteration:{} maximisation'.format(i))

    param_list = hard_maximisation(
        fit_data, thresholds, dist=dist,
        # has_noise_threshold=has_noise_threshold
    )
    # return param_list
    if verbose:
        print('Iteration:{} expectation '.format(i))
    new_thresh = hard_expectation(
        param_list, dist=dist, normalise=False
    )

    new_thresh.append(initial_thresholds[-1])
    # return param_list, new_thresh

    ll = mixture_model_ll(
        fit_data, param_list, dist=dist
    )
    # return new_thresh, param_list, ll
    if verbose:
        print('Threshold changes:{!r}'.format(initial_thresholds - new_thresh))
        print('log likelihood:{}'.format(ll))

    converged = False
    thresholds = np.array(new_thresh)
    while i < max_iter:
        i += 1
        if verbose:
            print('Iteration:{} maximisation'.format(i))
        new_param_list = hard_maximisation(
            fit_data, thresholds, dist=dist #,
            #has_noise_threshold=has_noise_threshold
        )
        if verbose:
            print('Iteration:{} expectation '.format(i))
        new_thresh = hard_expectation(
            new_param_list, dist=dist, normalise=False
        )
        new_thresh.append(initial_thresholds[-1])
        new_ll = mixture_model_ll(
            fit_data, new_param_list, #has_noise_threshold=has_noise_threshold,
            dist=dist
        )
        if verbose:
            print(
                'Threshold changes:{!r}'.format(thresholds - new_thresh)
            )
            print('log likelihood:{} change:{}'.format(new_ll, new_ll-ll))

        # terminate if tol met or ll goes down. FIXME is this appropriate?
        if abs(new_ll - ll) <= tol or new_ll < ll:
            if new_ll > ll:
                ll = new_ll
                param_list = new_param_list
            converged = True
            break

        ll = new_ll
        param_list = new_param_list
        thresholds = np.array(new_thresh)

    if verbose:
        if converged:
            print('Converged')
        else:
            print('Maximum iterations reached')

    # calculate the normalised thresholds
    thresholds = hard_expectation(
        param_list, normalise=True, dist=dist
    )
    # add the terminal threshold which is never fitted but marks the end of the
    # data that was used in fitting the model
    thresholds.append(initial_thresholds[-1])

    # calculate the threshold masks

    return MixtureModel(
        np.array(param_list), np.array(thresholds), ll, converged,
        dist, has_noise_threshold
    )


def hard_expectation(
        param_list,
        # fix_noise=None,
        normalise=False, dist=stats.gamma
):
    """
    Calculate the thresholds for a mixture model that define the data
    partitions.

    :param list param_list: list of parameters for the distributions in the
           mixture model, returned by expectation_maximisation().
    :param fix_noise: Fix threshold[1] at this value.
    :param bool normalise: calculate thresholds based on a mixture of
           normalised pdf's.
    :param dist: type of distribution
    :return: ndarray of threshold values.
    """

    # if fix_noise is None:
    #     t = [0]
    #     start = 1
    # else:
    #     t = [0, fix_noise]
    #     start = 2
    t = [0]
    # for m in range(start, len(param_list)):
    for m in range(1, len(param_list)):
        left = param_list[m - 1]
        right = param_list[m]

        if normalise:
            pdf = normalised_pdf
        else:
            pdf = scaled_pdf

        def f(x):
            return pdf(x, left, dist) - pdf(x, right, dist)

        t.append(brentq(f, dist.median(*left[:-1]), dist.median(*right[:-1])))

    return t


def hard_maximisation(
        data, thresholds, dist=stats.gamma, verbose=True
        #has_noise_threshold=False
):
    """
    Fit distributions to the data partitioned by thresholds.

    :param data: the data to model.
    :param thresholds: thresholds that divide data into separate distributions.
    :param dist: type of distribution to fit (scipy.stats).
    :param bool verbose: Print fitted distribution parameters.
    :return: list of parameters for each distribution.
    """
    param_list = []
    #     print('len thresholds',len(thresholds))
    for i in range(1, len(thresholds)):
        part = partition(data, thresholds, i)
        print('{}: {}'.format(i, len(part)))
        if len(part) != 0:
            # if has_noise_threshold and i == 1:
            #     fit = dist.fit(part, floc=0)
            # else:
            fit = dist.fit(part)
            if verbose:
                print('{} distribution:{} params:{}'.format(dist.name, i, fit))
        else:
            fit = []
            if verbose:
                print('{} distribution:{} empty'.format(dist.name, i))

        param_list.append(list(fit) + [len(part) / len(data)])
    return param_list


def mixture_model_ll(
        data, param_list,
        # has_noise_threshold=False,
        dist=stats.gamma
):
    """
    Calculate the log likelihood of a mixture model.

    :param ndarray data: the data used to construct the model.
    :param param_list: list of parameters for each distribution in the model.
    :param dist: the type of distribution to use (see scipy.stats).
    :return: the log likelihood.
    """

    ll = np.empty((len(param_list), len(data)), dtype=np.float64)
    # if has_noise_threshold:
    #     start = 1
    #     params = param_list
    # else:
    #     start = 0

    for i, params in enumerate(param_list):
        p = param_list[i][:-1]
        ll[i, :] = dist.pdf(data, *p) * param_list[i][-1]
    return np.sum(np.log(np.sum(ll, 0)))


@jit
def partition(data, thresholds, i):
    """
    Return the ith partition of the data as defined by thresholds.

    :param ndarray data: the data to partition.
    :param iterable thresholds: the threshold values that partition data.
    :param int i: the partition to return, the ith partition is defined as
                  threshold[i] < data <= threshold[i+1]. The last partition,
                  when i = len(thresholds-1), is defined as
    :return: ndarray of the data in the partition.
    :note: This is used to assign a hard responsibility in the expectation
           maximisation algorithm used to fit data to a mixture model.
    """
    print('part:{} > {} <= {}'.format(i, thresholds[i-1], thresholds[i]))
    return data[
        np.bitwise_and(data > thresholds[i-1], data <= thresholds[i])
    ]


def plot_mixture_model(
    model, data=None, counts=None, xrange=None, normalised=False,
    bins=500, bar=False, measurement_name=r'\texttt{area}', last_threshold=None,
    figsize=None
):
    """
    Plot a mixture model optionally including a histogram of the modeled data.
    if no data is None x must not be None.

    :param MixtureModel model: the mixture model.
    :param ndarray data: the data to histogram.
    :param tuple xrange: the x range to plot.
    :param bool normalised: normalise the model distributions.
    :param bins: passed to numpy.histogram().
    :param bool bar: Plot bars instead of markers.
    :param figsize: passed to matplotlib.pyplot.figure().
    :return: the figure handle.
    """

    points = 10000  # number of x points for plotting pdfs
    if data is not None:
        data_len = len(data)
    if data is None and xrange is None:
        raise AttributeError('Either data or x must be supplied')

    if figsize is None:
        fig = plt.figure()
    else:
        fig = plt.figure(figsize=figsize)

    ax = fig.add_axes([0.12, 0.12, 0.82, 0.85])
    ax.spines['right'].set_visible(False)
    ax.spines['top'].set_visible(False)

    if data is not None:
        if last_threshold is not None:
            """ 
            Introduced because under high zoom matplotlibs pgf back end writes 
            files with dimensions out of range errors when processed by Tex.
            """
            data = data[data < model.thresholds[last_threshold]]

        hist, edges = np.histogram(data, bins=bins)
        w = edges[1]-edges[0]
        bin_centers = edges[:-1]+w/2
        x = np.linspace(edges[0], edges[-1], points)

        # max_p = max(hist)
        # s = sum(hist) * w
        s = data_len * w
        if not normalised:
            cdfs = model.cdf(edges, scale=True)*data_len
            mhist = np.sum(cdfs[:, 1:]-cdfs[:, :-1], 0)

            if bar:
                data_handle = ax.bar(
                    bin_centers, hist, w, align='center',
                    facecolor=(1, 0, 0, 0.2),
                    edgecolor=(1, 1, 1, 1),
                    lw=1, label='data bin'
                )
                model_data_handle = ax.bar(
                    bin_centers, mhist, w, align='center',
                    facecolor=(0, 0, 1, 0.2),
                    edgecolor=(1, 1, 1, 1),
                    lw=1, label='model bin'
                )
            else:
                data_handle, = ax.plot(
                    bin_centers, hist, 's',
                    markeredgecolor='gray',
                    markerfacecolor='none',
                    markersize=4,
                    # markeredgealpha=0.2,
                    # alpha=0.2,
                    lw=1, label='data bin',
                    mew=0.5
                )

        else:
            s = 1
            mhist = None
            hist = None
            bin_centers = None
    else:
        mhist = None
        hist = None
        bin_centers = None
        s = 1
        x = np.linspace(*xrange, 10000)

    ax.set_prop_cycle(None)  # reset colors

    pdfs = np.zeros((len(model.param_list), points), np.float64)
    pdf_handles = []
    if last_threshold is not None:
        num_pdfs = last_threshold
    else:
        num_pdfs = len(model.param_list)
    # if model.has_noise_threshold:
    #     num_pdfs -= 1
    #     start = 1
    # else:
    #     start = 0

    for i in range(0, num_pdfs):

        if normalised:
            pdfs[i, :] = model.pdf(x, d=i)
            lw = 1
        else:
            pdfs[i, :] = model.pdf(x, d=i, scale=True)*s
            lw = 3

        h, = ax.plot(x, pdfs[i, :], lw=lw, label='pdf{}'.format(i))
        pdf_handles.append(h)

    # if data is None or normalised:
    #     ax.set_xlim([-10, max(data)])
    if model.has_noise_threshold:
        max_p = np.max(np.max(pdfs[1:]))
    else:
        max_p = np.max(np.max(pdfs))

    if data is None:
        ax.set_xlim([-10, max(x)])
    else:
        ax.set_xlim([-10, max(data)])
        ax.set_ylim([0, max_p*1.1])


    # ylim = ax.get_ylim()
    # ax.set_ylim([0, ylim[1]])

    if not normalised:
        model_handle, = ax.plot(x, np.sum(pdfs, 0), 'k', lw=1, label='model')

    if not normalised:
        # # FIXME fix_noise??
        # if model.has_noise_threshold:
        #     fix_noise = model.thresholds[1]
        # else:
        #     fix_noise = None
        thresh = np.zeros(len(model.thresholds))
        thresh[:-1] = hard_expectation(
            model.param_list, normalise=False, dist=model.dist
        )
        thresh[-1] = model.thresholds[-1]
    else:
        thresh = model.thresholds

    # if last_threshold is not None:
    #     thresh = thresh[:last_threshold+1]

    # if model.has_noise_threshold:
    #     thresh = thresh[1:]

    # print(thresh)

    for t in thresh[:-1]:
        ax.axvline(
            t, color='k', linestyle='--', lw=0.5, label='thr{}'.format(t)
        )

    if not normalised:
        ax.axvline(
            thresh[-1], color='k', linestyle='-', lw=1, label='thr{}'.format(t)
        )
        # ax.plot([t, t], [0, max_p], ':k', lw=1)

    legend_labels = []
    if data is not None:
        for t in range(len(pdf_handles)):
            if model.has_noise_threshold:
                p = t
            else:
                p = t+1
            if t == 0 and model.has_noise_threshold:
                label = 'noise'
            elif t == len(pdf_handles)-1:
                # print('there')
                # label = '{}+ photons'
                label = '{} photon distribution'
            elif p == 1:
                label = '{} photon distribution'
            else:
                # label = '{} photons'
                label = '{} photon distribution'

            legend_labels.append(label.format(p))

    if not normalised:
        legend_handles = pdf_handles + [model_handle]
    else:
        legend_handles = pdf_handles

    legend_labels.append('model')
    # if counts is not None:
    #     legend_labels.append('model {:,}'.format(sum(counts.count)))
    # else:
    #     legend_labels.append('model')

    if data is not None and not normalised:
        legend_handles += [data_handle]
        legend_labels += ['data bin']
        if bar:
            legend_handles += [model_data_handle]
            legend_labels += ['model bin']

    fig.legend(
        legend_handles, legend_labels, frameon=False,
        bbox_to_anchor=(1, 0.98), bbox_transform=fig.transFigure
    )

    ax.set_xticks(list(thresh))
    exp = int(np.floor(np.log10(model.thresholds[-1])))
    tick_values = model.thresholds/10**exp
    tick_labels = [r'${:.2f}$'.format(tick) for tick in tick_values]
    ax.set_xticklabels(tick_labels)
    ax.set_xlabel(
        r'{} measurement $(\times 10^{{ {} }})$'
        .format(measurement_name, exp)
    )

    ytvalues = ax.get_yticks()
    exp = int(np.floor(np.log10(ytvalues[-1])))
    tick_values = ytvalues/10**exp
    tick_labels = [r'${:1.1f}$'.format(tick) for tick in tick_values]
    ax.set_yticklabels(tick_labels)

    if normalised:
        ax.set_ylabel(r'probability $(\times 10^{{ {} }})$'.format(exp))
    else:
        ax.set_ylabel(r'count $(\times 10^{{ {} }})$'.format(exp))


    return fig, legend_handles, legend_labels, bin_centers, mhist, hist


def threshold_masks(data, model):
    """

    :param data:
    :param model:
    :return:
    """
    masks = np.zeros((len(model.thresholds), len(data)), np.bool)
    for t in range(len(model.thresholds[:-1])):
        if t == 0:
            masks[0, :] = data <= model.thresholds[1]
        else:
            masks[t, :] = and_l(
                data > model.thresholds[t], data <= model.thresholds[t+1]
            )
    masks[-1, :] = data > model.thresholds[-2]
    return masks


"""
timing analysis
"""


@jit
def xcor(s1, s2):
    """
    Count correlations between timestamp sequences.

    :param ndarray s1: Monotonically increasing timestamp sequence.
    :param ndarray s2: Monotonically increasing timestamp sequence.
    :return: correlation count.

    :notes: iterates over s1 and performs a linear search of  s2 for the s1
            timestamp.
    """
    i2 = 0
    c = 0
    for i in range(len(s1)):
        while i2 < len(s2) and s1[i] > s2[i2]:
            i2 += 1
        if i2 < len(s2):
            c += s1[i] == s2[i2]
    return c


@jit
def x_correlation(s1, s2, r):
    """
    Cross correlation of timestamp sequences s1 and s2 over the delay range r.

    :param ndarray s1: Monotonically increasing timestamp sequence.
    :param ndarray s2: Monotonically increasing timestamp sequence.
    :param r: range of delays added to s1
    :return: ndarray representing the cross correlation function.
    """
    x_corr = np.zeros((len(r),), np.uint64)
    for l in r:
        x_corr[l - r[0]] = xcor(s1 + l, s2)
    return x_corr


@jit
def drive_correlation(
    data, model, abs_time, detection_mask, r, last_threshold=-1,
    verbose=False
):
    """
    Calculate the temporal cross-correlation between a channel measuring a
    heralding signal, ie the laser drive pulse, and a channel detecting photons.
    The correlation is calculated for the different photon numbers determined by
    the thresholds parameter.


    :param ndarray abs_time: absolute timestamp sequence.
    :param ndarray detection_mask: boolean mask that identifies the entries in
                                   abs_time belonging to the photon channel,
                                   other entries are assumed to be the heralding
                                   channel.
    :param ndarray data: energy measurement data for the photon channel.
    :param ndarray masks: boolean masks that are appied to
                          abs_time[detection_mask] and indicate which detections
                          are assigned which threshold.
                          Shape(N, len(abs_time[detection_mask]))
    :param r: the range of heralding channel delays to cross correlate.
    :param bool verbose: Print progress message as each threshold is processed.
    :return: list of ndarrays representing the cross correlation.

    :note: abs_time[mask] and data must be the same length.
           TODO speed up the algorithm.
    """

    if last_threshold > len(model.thresholds):
        raise AttributeError('last_threshold must be < len(model.thresholds)')
    if last_threshold < 0:
        last_threshold = len(model.thresholds)+last_threshold

    xc = []
    if last_threshold == 0:
        if verbose:
            print('Cross correlating all events')
        xc = x_correlation(
            abs_time[not_l(detection_mask)], abs_time[detection_mask], r
        )
    else:
        for t in range(last_threshold+1):
            if verbose:
                print(
                    'Cross correlating partition {} of {}'
                    .format(t+1, last_threshold+1))

                if t != last_threshold:
                    xc.append(
                        x_correlation(
                            abs_time[not_l(detection_mask)],
                            abs_time[detection_mask][
                                and_l(
                                    data > model.thresholds[t],
                                    data <= model.thresholds[t+1]

                                )

                            ],
                            r
                        )
                    )
                else:
                    xc.append(
                        x_correlation(
                            abs_time[not_l(detection_mask)],
                            abs_time[detection_mask][data > model.thresholds[t]]
                            , r
                        )
                    )

    return np.array(xc)


@jit
def window(i, abs_time, low, high):
    """
    get offsets from the current index i in abs_time that are in the relative
    time window defined by low and high.

    :param int i: current abs_time index.
    :param ndarray abs_time: absolute times.
    :param low: low end of relative window.
    :param high: high end of relative window.
    :return: (low_o, high_o) coincident times are abs_time[low_o:i+high_o]
    """

    length = len(abs_time)
    low_o = 0  # low index offset
    high_o = 0  # high index offset
    now = abs_time[i]
    low_t = now + low
    high_t = now + high

    if low_t < now:
        while (i + low_o >= 0) and abs_time[i + low_o] >= low_t: low_o -= 1;
    else:
        while (low_o + i < length) and abs_time[i + low_o] < low_t: low_o += 1;

    if high_t < now:
        while (i + high_o >= 0) and abs_time[i + high_o] > high_t:
            high_o -= 1
    else:
        while (high_o + i < length) and abs_time[i + high_o] <= high_t:
            high_o += 1
    return low_o, high_o  # offset indexs marking abs_times in the window


def coincidence(abs_time, mask, low, high):
    """
    Find coincidences between two channels returning indices of partner events.

    :param ndarray abs_time: absolute times.
    :param ndarray mask: ndarray of bool that identifies the channels in
                         abs_time. For abs_times where mask is False, time t is
                         coincident if abs_time+low <= t <= abs_time+high. If
                         the mask is True t is coincident if
                         abs_time-high <= t <= abs_time-low.
    :param float low: the low side of the coincidence window.
    :param float high:  the high side of the coincidence whindw.
    :return: (coinc, coinc_mask) where coinc is a ndarray of indices of the
             coincident event in the other channel. When more than one event is
             found in the window the negated value of the first index is
             entered.
             coinc_mask is a ndarray of bool indicating where exactly one
             event was found in the window.
    """

    coinc = np.zeros(len(abs_time), np.int32)
    coinc_mask = np.zeros(len(abs_time), np.bool)
    for t in range(len(abs_time)):
        if t % 10000000 == 0:
            print(t)
        if not mask[t]:
            low_o, high_o = window(t, abs_time, low, high)
            o = low_o
        else:
            low_o, high_o = window(t, abs_time, -high, -low)
            o = high_o
        coinc_length = high_o - low_o
        if high_o and low_o:
            if coinc_length == 1:
                coinc[t] = t + o
                coinc_mask[t] = True
            elif coinc_length > 1:
                coinc[t] = -(t + o)
                coinc_mask[t] = False
    return coinc, coinc_mask


Counts = namedtuple(
    'Counts', 'count uncorrelated vacuum has_noise_threshold'
)


def count(
        data, measurement_model, vacuum=False, has_noise_threshold=False,
        coinc_mask=None, herald_mask=None
):
    """

    :param ndarray data: measurement data.
    :param MixtureModel measurement_model: model created using
                                           expectation_maximisation.
    :param bool vacuum: estimate vacuum terms using coinc_mask and herald_mask.
    :param coinc_mask: ndarray of bool indicating coincidence.
    :param herald_mask: ndarray of bool indicating which events are in the
                        heralding channel.
    :param has_noise_threshold: When True model.thresholds[1] represents the
                                boundary between noise and 1 photon.
    :return: Counts(count, uncorrelated, vacuum). Where count is an ndarray of
             counts for each photon number
             and len(count)=len(measurement_model.thresholds). uncorrelated is a
             ndarray the same shape as count and counts
             events not correlated with the herald. uncorrelated is only
             calculated when coinc_mask and herald_mask are supplied.
             vacuum replicates the parameter value and indicates that count[0]
             contains the vacuum count.
    """

    if vacuum and (coinc_mask is None or herald_mask is None):
        raise AttributeError(
            'Neither coinc_mask or herald_mask cannot be None when vacuum=True'
        )

    if (
        (herald_mask is not None and coinc_mask is None) or
        (herald_mask is None and coinc_mask is not None)
    ):
        raise AttributeError(
            'coinc_mask or herald_mask both be specified or both None'
        )

    if herald_mask is not None:
        data_mask = not_l(herald_mask)
        uncorrelated = []
    else:
        data_mask = None
        uncorrelated = None

    counts = []
# thresholds
    for t in range(1, len(measurement_model.thresholds)):
        counts.append(
            len(
                and_l(
                    (data > measurement_model.thresholds[t - 1]),
                    (data <= measurement_model.thresholds[t])
                ).nonzero()[0]
            )
        )
        if uncorrelated is not None:
            uncorrelated.append(
                len(
                    and_l(
                        data[not_l(coinc_mask)[data_mask]] >
                        measurement_model.thresholds[t-1],
                        data[not_l(coinc_mask)[data_mask]] <=
                        measurement_model.thresholds[t]
                    ).nonzero()[0]
                )
            )

    # last threshold overflow
    counts[-1] += len((data > measurement_model.thresholds[-1]).nonzero()[0])
    if uncorrelated is not None:
        uncorrelated[-1] += (
            len(
                (data[not_l(coinc_mask)[data_mask]] >
                    measurement_model.thresholds[-1]).nonzero()[0]
            )
        )

    # counting vacuum
    # uncorrelated heralds + correlated data <= the first threshold
    if vacuum:
        uncorrelated_heralds = len(
            (not_l(coinc_mask)[herald_mask]).nonzero()[0]
        )
        if has_noise_threshold:
            noise = len(
                (
                    data[coinc_mask[data_mask]] <=
                    measurement_model.thresholds[1]
                ).nonzero()[0]
            )
            counts[0] = noise+uncorrelated_heralds
            # note these are correlated events less that the noise threshold
            if uncorrelated is not None:
                uncorrelated[0] = uncorrelated_heralds
        else:
            counts.insert(0, uncorrelated_heralds)
            if uncorrelated is not None:
                uncorrelated.insert(0, uncorrelated_heralds)

    if uncorrelated is not None:
        uncorrelated = np.array(uncorrelated)

    return Counts(
        np.array(counts), uncorrelated, vacuum, has_noise_threshold
    )


Povm = namedtuple(
    'povm', 'elements vacuum'
)


# FIXME remove counts
# def povm_elements(measurement_model, counts):
#     """
#     estimate the elements of the system povm from the measurement model and the
#     counting data.
#
#     :param array like data: measurement data.
#     :param mixturemodel measurement_model: model created using
#                                            expectation_maximisation.
#     :param Counts counts: the count data returned by count().
#     :return: (elements count vacuum). where counts is a ndarray with
#              shape (len(n)) containing the count for each measurement outcome
#              in n. elements is a ndarray with
#              shape(len(n), len(measurement_model.thresholds)) the first index is
#              the fock state, the second is the outcome. The contents of elements
#              is the probability of the outcome when measuring fock state.
#
#     :notes: POVM elements is indexed by (fock state, measurement outcome).
#             The last threshold in the measurement model marks the boundary of
#             the data that was used to create the model, it is *not* altered by
#             the expectation maximisation algorithm and is essentially a guess.
#             The second last threshold is the boundary of the overflow bin,
#             this last photon number bin is counts detections of >=
#             len(measurement_model.param_list)-2 photons.
#             TODO expand and clarify description.
#     """
#
#     # t thresholds
#     # t-1 distributions
#     # t+1 counts if vacuum else t count[-1] is overflow past last threshold
#
#
#     num_elements = len(counts.count)
#     # cumulative density functions at the thresholds, the last value is not used
#     # normalise to a convex mixture excluding the noise distribution.
#     cdfs = measurement_model.cdf(measurement_model.thresholds)
#
#     elems = cdfs[:, 1:] - cdfs[:, :-1]
#     # Adjust the first elements, which should not have a proceeding threshold.
#     elems[:, 0] = cdfs[:, 1]
#
#     # Adjust the last elements to account for the truncation at the last
#     # threshold (limit of data used to create the model).
#     elems[:, -1] = 1 - cdfs[:, -2]
#
#     # elements should always have a vacuum term.
#     elements = np.zeros((num_elements, num_elements))
#
#     # the majority of the POVM elements are given by the difference in the CDFs
#     # at the threshold values
#
#     if counts.vacuum:
#         if measurement_model.has_noise_threshold:
#             elements = elems
#         else:
#             elements[1:, 1:] = elems
#     else:
#         elements = elems
#
#     # vacuum terms
#     # FIXME check this is correct when has_noise_threshold
#     if counts.vacuum:
#         # probability of counting vacuum when there is a vacuum input
#         # TODO is it appropriate to do same for m > 1 ??
#         for m in range(1, 2):  # range(1, len(counts.uncorrelated)):
#             elements[0, 0] = (
#                     (counts.count[0]-counts.uncorrelated[m])/counts.count[0]
#             )
#             elements[0, m] = counts.uncorrelated[m]/counts.count[0]
#
#          # elements[m, 0] = correlatedsingle/single
#
#         if measurement_model.has_noise_threshold:
#             # look for uncrrelated ???
#             pass
#
#     else:
#         elements[0, 0] = 1
#         elements[0, 1] = 0
#
#     # c = np.array(counts.count[:-1])
#     # c[-1] += counts.count[-1]
#
#     return Povm(elements, counts.vacuum, counts.has_noise_threshold)


def fit_state_least_squares(
        counts, povm, max_outcome=None, vacuum=True, thermal=True, N=150
):

    # FIXME check this model with new POVM remove call to outcomes
    def qutip_d_thermal(n, nbar, alpha, N):
        # n [low, high]
        t = qt.thermal_dm(N, nbar)
        D = qt.displace(N, alpha)
        s = D * t * D.dag()

        #         print('n in model',n)
        state = [
                    (qt.states.fock(N, num) *
                     qt.states.fock(N, num).dag() * s).tr()
                    for num in np.arange(0, N)
                ]
        # print('state sum',np.sum(state))

        # print('state.shape', state.shape)
        num_elements = min(povm.shape[1], n[1]-n[0]+1)
        # print('num elements', num_elements, 'N', N)
        truncated_state = np.real(state[n[0]:num_elements + 1])
        truncated_state[-1] += np.real(np.sum(state[num_elements + 1:]))
        # print('nbar', nbar, 'alpha', alpha)
        # print('sum state', np.sum(state[n[0]:]))
        # print('sum t state', np.sum(truncated_state[n[0]:]))
        # print(
        #     'trunc test',
        #     np.real(np.sum(state[n[0]:]))-np.sum(truncated_state[n[0]:])
        # )
        if num_elements < povm.shape[0]:
            truncated_povm = povm[:, n[0]:num_elements]
            truncated_povm[-1] += np.sum(povm[num_elements:])
        else:
            truncated_povm = povm
        outcomes = np.zeros_like(truncated_povm)
        #         print('outcomes',outcomes.shape)
        for e in range(num_elements):
            outcomes[e, :] += truncated_state[:-1] * truncated_povm[e, :]
        out = np.sum(outcomes, 0)
        out[-1] += truncated_state[-1]
        # if n[0] != 0:
        #     out = out / np.sum(out)
        #       print('out',out.shape)
        return out

    n = np.zeros(2, np.uint32)
    if vacuum:
        n[0] = 0
    else:
        n[0] = 1

    if max_outcome is not None:
        n[1] = max_outcome
        c = counts[n[0]:n[1] + 1]
        c[-1] += np.sum(counts[n[1] + 1:])
    else:
        n[1] = povm.shape[0]+1
        c = counts[n[0]:]

    print('fit c shape', c.shape)

    state_model = Model(qutip_d_thermal)
    # alpha = 0.1
    pars = state_model.make_params(nbar=0.1, alpha=0.1, N=N)
    pars['nbar'].set(min=0, max=0.5)
    pars['alpha'].set(min=0, max=10)
    pars['nbar'].set(vary=thermal)
    if not thermal:
        pars['nbar'].set(value=0.0)
    pars['N'].set(vary=False)
    #     print('n',n)
    #     pars['povm_elements'].set(vary=False)
    fit = state_model.fit(c / np.sum(c), pars, n=n)
    print('fit n', n)
    print('fit c', c)
    print('fit sum(c)', np.sum(c))
    print('fit best_fit', fit.best_fit)
    print('fit sum(best_fit)', np.sum(fit.best_fit))
    print('fit best_fit counts', fit.best_fit * np.sum(c))
    return c, n, fit


def resize_vector(a, max_index, copy=True):

    if max_index+1 > len(a):
        raise AttributeError('max_index must be <= len(a)')
    if max_index+1 == len(a):
        return a

    if copy:
        o = np.copy(a[:max_index+1])
    else:
        o = a[:max_index+1]
    o[-1] += np.sum(a[max_index+1:])
    return o


def displaced_thermal(max_photon_number, nbar, alpha, N=200):
    # n [low, high]
    t = qt.thermal_dm(N, nbar)
    D = qt.displace(N, alpha)
    s = D * t * D.dag()

    state = [
        (qt.states.fock(N, num)*qt.states.fock(N, num).dag()*s).tr()
        for num in np.arange(0, N)
    ]

    truncated_state = resize_vector(np.real(state), max_photon_number)

    return truncated_state #/np.sum(truncated_state)


def displaced_vacuum(max_photon_number, alpha, N=200):
    # n [low, high]
    t = qt.thermal_dm(N, 0)
    D = qt.displace(N, alpha)
    s = D * t * D.dag()

    state = [
        (qt.states.fock(N, num)*qt.states.fock(N, num).dag()*s).tr()
        for num in np.arange(0, N)
    ]

    truncated_state = resize_vector(np.real(state), max_photon_number)

    return truncated_state #/np.sum(truncated_state)


def outcome_probabilities(state, povm, max_photon_number):

    if max_photon_number is None:
        max_photon_number = len(povm.elements) - 2
    elif max_photon_number > len(povm.elements) - 2:
        raise AttributeError(
            'max_photon_number must be at least 2 less than the number of '
            'POVM elements ({})'.format(len(povm.elements) - 2)
        )

    outcomes = np.zeros((max_photon_number+1, max_photon_number+1))

    for i in range(max_photon_number+1):
        outcomes[i, :] = resize_vector(
            state*povm.elements[i, :], max_photon_number
        )
    return outcomes


def thermal_neg_ll(x, counts, povm, max_photon_number):

    c = resize_vector(counts.count, max_photon_number)
    state = displaced_thermal(max_photon_number, *x, N=100)
    if povm is not None:
        outcomes = povm.elements*state
        if len(np.sum(outcomes, 0).nonzero()[0]) != outcomes.shape[1]:
            print('inf')
            ll = -np.inf
        else:
            ll = np.sum(np.log(np.sum(outcomes, 0))*c)
    else:
        if len(state.nonzero()[0]) != len(state):
            print('inf')
            ll = -np.inf
        else:
            ll = np.sum(np.log(state)*c)

    return -ll


def coherent_neg_ll(x, counts, povm, max_photon_number):
    c = resize_vector(counts.count, max_photon_number)
    state = displaced_vacuum(max_photon_number, *x, N=100)
    if povm is not None:
        outcomes = povm.elements*state
        if len(np.sum(outcomes, 0).nonzero()[0]) != outcomes.shape[1]:
            print('inf')
            ll = -np.inf
        else:
            ll = np.sum(np.log(np.sum(outcomes, 0))*c)
    else:
        if len(state.nonzero()[0]) != len(state):
            print('inf')
            ll = -np.inf
        else:
            ll = np.sum(np.log(state)*c)

    return -ll


def thermal_ml(counts, povm, max_photon_number, x0=None):
    if x0 is None:
        x0 = [0.3, 3.0]
    r = minimize(
        thermal_neg_ll, x0, args=(counts, povm, max_photon_number),
        method='L-BFGS-B', tol=1e-15,
        bounds=[(0.0, 2.0), (0.1, 6.0)],
        options={
            'disp': True, 'eps': 1e-6, 'ftol': 1e-15, 'gtol': 1e-15, 'maxls': 50
        }
    )
    return r


def coherent_ml(counts, povm, max_photon_number, x0=None):
    if x0 is None:
        x0 = [3.0]
    r = minimize(
        coherent_neg_ll, x0, args=(counts, povm, max_photon_number),
        method='L-BFGS-B', tol=1e-15,
        bounds=[(0.1, 6.0)],
        options={
            'disp': True, 'eps': 1e-6, 'ftol': 1e-15, 'gtol': 1e-15, 'maxls': 50
        }
    )
    return r


def plot_state_fit(nbar, alpha, counts, povm=None, figsize=None):
    width = 0.46   # the width of the bars
    if figsize is not None:
        fig = plt.figure(figsize=figsize)
    else:
        fig = plt.figure()

    ax = fig.add_axes([0.12, 0.12, 0.84, 0.83])

    max_photon_number = len(counts.count)-1

    state = displaced_thermal(max_photon_number, nbar, alpha, N=200)
    if povm is not None:
        model_counts = np.sum(povm.elements*state*np.sum(counts.count), 1)
    else:
        model_counts = state*np.sum(counts.count)

    fock = np.arange(0, max_photon_number+1)
    # print(fock.shape, fock)

    ax.bar(fock-width/2, counts.count, width=width, color='r', label='data')
    ax.bar(
        fock+width/2, model_counts, width=width, color='b', label='model'
    )
    # chi2 = sum(np.power(c - model_counts, 2) / model_counts)
    ax.set_ylabel('count')
    ax.set_xlabel('Fock state')
    ax.set_xticks(fock)
    xlabels = [r'$\vert {} \rangle$'.format(f) for f in fock]
    xlabels[-1] = r'$\vert {}\rangle{}$'.format(fock[-1], '+')
    ax.set_xticklabels(xlabels)

    ax.legend(frameon=False)
    # plt.plot(n,nc,'.')
    # plt.plot(n,r.best_fit,'o',markerfacecolor='none')
    # ax.set_ylim(0,11000000)
    ty = 0.8
    tx = 0.01
    ls = 0.05
    ax.text(tx, ty, r'Displaced thermal state', transform=ax.transAxes)
    # ax.text(
    #     tx, ty - ls, r'$\alpha={:.4f}\pm{:.4f}$'
    #         .format(
    #         state_fit.params['alpha'].value, state_fit.params['alpha'].stderr
    #     ),
    #     transform=ax.transAxes, horizontalalignment='left'
    # )
    # ax.text(
    #     tx, ty - 2 * ls, r'${}={:.4f}\pm{:.4f}$'
    #         .format(
    #         r'\bar{n}', state_fit.params['nbar'].value,
    #         state_fit.params['nbar'].stderr
    #     ),
    #     transform=ax.transAxes, horizontalalignment='left'
    # )
    # ax.text(
    #     tx, ty - 3 * ls, r'$\chi^2/{}={:.3f}$'
    #         .format(
    #         state_fit.nfree - 1, chi2 / (state_fit.nfree - 1)
    #     ),
    #     transform=ax.transAxes, horizontalalignment='left'
    # )
    ax.legend(frameon=False)
    return fig


def povm_elements(model):
    cdfs = model.cdf(model.thresholds)
    elems = np.zeros((cdfs.shape[0]+1, cdfs.shape[1]))
    elems[1:, 1:] = cdfs[:, 1:] - cdfs[:, :-1]
    elems[1:, 0] = cdfs[:, 0]
    elems[1:, -1] = 1-cdfs[:, -2]
    elems[0, 0] = 1
    return elems


def model_outcomes(x, N, thermal=False, povm=None, qtN=100):
    """
    Model the measurement outcome probabilities assuming the input is a
    displaced thermal state.

    :param x iterable: The state parameters
    :param N int: Dimension of the Fock space, max photon number is N-1
    :param thermal: When True model input as displaced thermal state otherwise
                    model as coherent state (displaced thermal).
    :param povm: The detector povm
    :param qtN: The dimension of the Fock space to use for qutip models.
    """
    #     print(x)
    D = qt.displace(qtN, x[0])
    if thermal:
        #         print('thermal')
        t = qt.thermal_dm(qtN, x[1])
    else:
        t = qt.thermal_dm(qtN, 0)

    s = D * t * D.dag()

    state = np.real(np.array([
        (qt.states.fock(qtN, n) * qt.states.fock(qtN, n).dag() * s).tr()
        for n in np.arange(0, qtN)
    ]))
    #     print(np.sum(state))
    if povm is not None:
        n_ele = povm.shape[1]  # number of povm elements
        if N > n_ele:
            raise AttributeError(
                'N must be less than or equal to the number of POVM elements'
            )
        o = np.sum(state[:n_ele] * povm.transpose(), 1)  #
        o[-1] += np.sum(state[n_ele:])  # truncate to number of elements
        if N < n_ele:  # truncate to N if necessary
            ot = np.array(o[:N], copy=True)
            ot[-1] += np.sum(o[N:])
            o = ot
    else:
        o = np.array(state[:N], copy=True)
        o[-1] += np.sum(state[N:])
    return o


def fit_input_ls(
        counts, N, povm=None, thermal=False, qtN=150
):
    def dt_model(N, alpha, nbar):
        if thermal:
            x = [alpha, nbar]
        else:
            x = [alpha]

        o = model_outcomes(x, N, thermal=thermal, povm=povm, qtN=qtN)

        return o * total

    if N > len(counts):
        raise AttributeError('N cannot be greater than the number of counts')
    if N < len(counts):
        c = np.array(counts[:N], copy=True)
        c[-1] += np.sum(counts[N:])
        counts = c

    model = Model(dt_model, nan_policy='omit')

    pars = model.make_params(nbar=0.1, alpha=1.0)
    pars['nbar'].set(min=0, max=0.5)
    pars['alpha'].set(min=0, max=10)
    pars['nbar'].set(vary=thermal)

    if not thermal:
        pars['nbar'].set(value=0.0)

    total = np.sum(counts)
    return model.fit(
        counts, pars, N=N, weights=np.sqrt(1.0 / counts)
    )


def xcor_plot(xc, r, has_noise_threshold=False, figsize=None):

    if figsize is not None:
        fig = plt.figure(figsize=figsize)
    else:
        fig = plt.figure()

    ax = fig.add_axes([0.12, 0.12, 0.85, 0.85])
    ax.spines['right'].set_visible(False)
    ax.spines['top'].set_visible(False)

    sxc = np.sum(xc, 0)

    ax.bar(
        r[sxc != 0] * 4, sxc[sxc != 0], 4,
        align='center',
        facecolor=(0.01, 0.01, 0.01, 0.1),
        edgecolor=(0, 0, 0, 0.5),
        lw=0.5, label=r'Total'
    )

    if xc.shape[0] > 1:
        for x in range(xc.shape[0] - 1):
            r_nz = r[xc[x] != 0] * 4
            xc_nz = xc[x][xc[x] != 0]
            if has_noise_threshold:
                p = x
            else:
                p = x + 1
            if x == 0 and has_noise_threshold:
                label = 'noise'
            elif x == xc.shape[0] - 2:
                r_nz = r[np.sum(xc[x:]) != 0] * 4
                xc_nz = np.sum(xc[x:])[np.sum(xc[x:]) != 0]
                label = '{}+ photons'.format(p)
            elif p == 1:
                label = '{} photon'.format(p)
            else:
                label = '{} photons'.format(p)
            ax.plot(
                r[xc[x] != 0] * 4, xc[x][xc[x] != 0], '-', markersize=5,
                markerfacecolor='none', mew=0.5, label=label, lw=1
            )

    ax.set_xlabel(r'detection time relative to herald (ns)')
    ax.set_ylabel(r'count')
    #     ax.set_xlim([200,800])
    ax.frameon = False
    fig.legend(
        bbox_to_anchor=(0.33, 0.98), bbox_transform=fig.transFigure,
        frameon=False
    )
    return fig


def plot_ls_state(r, counts, N, thermal=False, figsize=None):
    if figsize is not None:
        fig = plt.figure(figsize=figsize)
    else:
        fig = plt.figure()

    width = 0.48  # the width of the bars
    ax = fig.add_axes([0.12, 0.12, 0.84, 0.83])
    n = np.arange(0, r.ndata)
    #     print(len(counts))
    if N > len(counts):
        raise AttributeError('N cannot be greater than len(counts)')
    if N < len(counts):
        c = np.array(counts[:N], copy=True)
        c[-1] += np.sum(counts[N:])
    else:
        c = counts

    #     print(n[0],len(n),len(c))
    rects1 = ax.bar(n - width / 2, c, width, color='r')
    rects2 = ax.bar(n + width / 2, r.best_fit, width, color='b'
                    )
    #     chi2=sum(np.power(counts-r.best_fit,2)/model_counts)
    ax.set_ylabel('count')
    ax.set_xlabel('Fock state')
    ax.set_xticks(n)
    xlabels = [r'$\vert {} \rangle$'.format(p) for p in n]
    xlabels[-1] = r'$\vert {}\rangle{}$'.format(n[-1], '+')
    ax.set_xticklabels(xlabels)

    ax.legend((rects1[0], rects2[0]), ('data', 'model'), frameon=False)
    # plt.plot(n,nc,'.')
    # plt.plot(n,r.best_fit,'o',markerfacecolor='none')
    # ax.set_ylim(0,11000000)
    ty = 0.8
    tx = 0.7
    ls = 0.05
    if thermal:
        ax.text(tx, ty, r'Displaced thermal state', transform=ax.transAxes)
    else:
        ax.text(tx, ty, r'Coherent state', transform=ax.transAxes)
    ax.text(
        tx, ty - ls, r'$\alpha={:.4f}\pm{:.4f}$'
            .format(
            r.params['alpha'].value, r.params['alpha'].stderr
        ),
        transform=ax.transAxes, horizontalalignment='left'
    )
    if thermal:
        ax.text(
            tx, ty - 2 * ls, r'${}={:.4f}\pm{:.4f}$'
                .format(
                r'\bar{n}', r.params['nbar'].value, r.params['nbar'].stderr
            ),
            transform=ax.transAxes, horizontalalignment='left'
        )
    ax.text(
        tx, ty - 3 * ls, r'$\chi^2/{}={:.3f}$'
            .format(
            r.ndata - r.nvarys - 1, r.chisqr / (r.ndata - r.nvarys - 1)
        ),
        transform=ax.transAxes, horizontalalignment='left'
    )
    ax.legend(
        (rects1[0], rects2[0]), ('model', 'data'),
        bbox_to_anchor=(tx, ty - 0.34), loc='lower left',
        frameon=False
    )
    return fig

<filename>tes/analysis.py
"""
Calibration and modeling
"""
import matplotlib.pyplot as plt
import matplotlib as mpl
import numpy as np
from numpy import (
    bitwise_and as and_b, logical_and as and_l, logical_or as or_l,
    logical_not as not_l, bitwise_or as or_b
)
from scipy import signal
from scipy import stats
from scipy.optimize import brentq
from numba import jit
from collections import namedtuple
import qutip as qt
from lmfit import Model
from pathlib import Path
from scipy.optimize import minimize


@jit
def maxima(f, thresh=10):
    """
    Find local maxima of f using rising zero crossings of the gradient of f.

    :param f: function to find maxima for
    :param thresh: Only return a maxima if f[maxima]>thresh

    :return: array of maxima.
    :Notes: Used to find the maxima of the smoothed measurement histogram.
    """
    g = np.gradient(f)
    pos = g > 0
    xings = (pos[:-1] & ~pos[1:]).nonzero()[0] + 1  # indices of the maxima
    return xings[np.where(f[xings] > thresh)]


Guess = namedtuple(
    'Guess', 'hist f bin_c max_i thresholds'
)


def guess_thresholds(data, bins=16000, win_length=200, maxima_threshold=10):
    """
    Guess the photon number thresholds using a smoothed histogram.

    :param ndarray data: the measurement data.
    :param bins: argument to numpy.histogram.
    :param int win_length: Hanning window length used for smoothing.
    :param int maxima_threshold: passed to maxima function.

    :return: (hist, f, bin_c, max_i, thresholds) as a named
            tuple. Where hist is the histogram
            generated by np.histogram(), f is the smoothed histogram, bin_c is a
            ndarray of bin centers, max_i are the indices of the maxima in f,
            and thresholds is a ndarray of threshold guesses.

    :notes: The smoothing is achieved by convolving the Hanning window with a
            histogram generated using numpy.histogram. The guess still needs
            some human sanity checks, especially multiple maxima per photon
            peak. Use plot_guess() and adjust window and maxima_threshold to
            remove them.
    """

    hist, edges = np.histogram(data, bins=bins)
    bin_w = edges[1] - edges[0]
    bin_c = edges[:-1] + bin_w / 2

    # smooth the histogram
    win = signal.hann(win_length)
    f = signal.convolve(hist, win, mode='same') / sum(win)

    max_i = maxima(f, thresh=maxima_threshold)  # indices of the maxima
    m = [0.0] + list(bin_c[max_i])
    t = [m[i - 1] + (m[i] - m[i - 1]) / 2 for i in range(2, len(m))]
    thresholds = np.array([0.0] + t + [m[-1] + (m[-1] - t[-1])*1.10])

    return Guess(hist, f, bin_c, max_i, thresholds)


def plot_guess(
        hist, f, x, max_i, init_t, figsize=None,
        measurement_name=r'\texttt{area}'
):
    if figsize is not None:
        fig = plt.figure(figsize=figsize)
    else:
        fig = plt.figure()
    ax = fig.add_axes([0.12, 0.12, 0.87, 0.87])
    ax.plot(
        x[hist.nonzero()[0]], hist[hist.nonzero()[0]], 's',
        markersize=3, markerfacecolor='none', mew=0.5,
        label='histogram bin', color='gray'
    )
    ax.plot(x, f, 'k', lw=1, label='smoothed histogram')
    ax.plot(
        x[max_i], f[max_i], 'o', label='maxima', markerfacecolor='none', mew=1,
        markersize=5, color='r'
    )
    for t in init_t:
        plt.plot([t, t], [0, max(hist)], ':k', lw=0.5)

    fig.canvas.draw_idle()
    offset = ax.xaxis.get_major_formatter().get_offset()
    ax.xaxis.get_major_formatter().set_useOffset(False)
    if offset != '':
        ax.set_xlabel('{} measurement ({})'.format(measurement_name, offset))
    else:
        ax.set_xlabel('{} measurement'.format(measurement_name))
    ax.set_ylabel('count')
    # ax.set_title('Histogram and initial threshold guess')
    plt.legend()
    return fig


# @jit
def retime(events, tdat, fidx, tidx):
    # TODO generalise
    # remove incorporate tick relative times
    #     events = edat.view(event_dt)
    event_time = events['time']
    event_i = 0
    event_time[0] = 0xFFFF
    # tidx[0] is the first tick period ending with tick 1
    for i in range(len(tidx) - 1):
        # i+1 is the tick number
        changed = fidx[
                  tidx[i]['start'] + 1:tidx[i]['stop'] + 1]['changed'
        ].nonzero()[0]
        if len(changed):
            raise NotImplementedError(
                'tick:{} is not homogeneous'.format(i)
            )
        event_count = int(
            (
                    fidx[tidx[i]['stop']]['payload'] +
                    fidx[tidx[i]['stop']]['length'] -
                    fidx[tidx[i]['start']]['payload']
            ) / (fidx[tidx[i]['start']]['event_size'] * 8)
        )
        #         print(event_count)
        event_i += event_count
        time = int(tdat[i + 1]['time']) + int(event_time[event_i])
        #     print(event_i)
        if time > 0xFFFF:
            time = 0xFFFF
        event_time[event_i] = time
    return event_time


def normalised_pdf(x, params, dist=stats.gamma):
    """
    Convenience function, calculate normalised pdf at each x for dist.

    :param union[numpy.ndarray, int] x: calculate the pdf a each x
    :param iterable params: fitted distribution parameters
    :param dist: type of distribution (see scipy.stats)
    :return: ndarray containing pdf(x)
    """
    return dist.pdf(x, *params[:-1])


def scaled_pdf(x, params, dist=stats.gamma):
    """
    Convenience function, calculate pdf at each x for dist, normalisation is
    given by params[-1].

    :param union[numpy.ndarray, int] x: calculate the pdf a each x
    :param iterable params: fitted distribution parameters
    :param dist: type of distribution (see scipy.stats)
    :return: ndarray containing pdf(x)
    """
    return params[-1]*dist.pdf(x, *params[:-1])


class MixtureModel(
    namedtuple(
        'MixtureModel', 'param_list thresholds log_likelihood '
        'converged dist has_noise_threshold'
     )
):
    """
    Subclass of namedtuple used to represent a mixture model.

    fields:
        :param_list: list of parameters for each distribution in the mixture.
        :thresholds: the intersection of neighbouring distributions in the
                     mixture.
        :zero_loc: Fixed location parameter used to fit the first distribution,
                   or None if the location parameter was fitted.
        :log_likelihood: the log of the likelihood that the model could produce
                         the data used to construct it.
        :converged: boolean indicating that the expectation maximisation
                    algorithm terminated normally when fitting the model.
        :dist: the type of distribution forming the components of the mixture,
               (see scipy.stats).

    """

    __slots__ = ()

    def save(self, filename):
        """
        save model as .npz file.

        :param filename: filename to save as, excluding extension"
        :return: None
        """
        np.savez(
            filename,
            param_list=self.param_list,
            thresholds=self.thresholds,
            # zero_loc=self.zero_loc,
            log_likelihood=self.log_likelihood,
            converged=self.converged,
            dist=self.dist.name,
            has_noise_threshold=self.has_noise_threshold
        )

    @staticmethod
    def load(filename):
        """
        load model from .npz file.

        :param filename: filename excluding extension"
        :return: instance of MixtureModel.
        """
        d = np.load(Path(filename).with_suffix('.npz'))
        return MixtureModel(
            d['param_list'],
            d['thresholds'],
            # d['zero_loc'][()],
            d['log_likelihood'][()],
            d['converged'][()],
            getattr(stats, str(d['dist'])),
            d['has_noise_threshold'][()]
        )

    def _eval(self, x, d=None, scale=False, func='pdf'):
        """
        evaluate the function of the distribution(s) selected by d at x.

        :param x: value(s) where to the function is evaluated.
        :param d: selects element(s) of param_list that parameterise the
                  distributions. If None all distributions in param_list are
                  used.
        :param bool scale: when true return the scaled pdf, the scale is always
                           param[-1].
        :param str func: the name of the distribution function to evaluate
                             (see scipy.stats.rv_continuous)
        :return: ndarray shape (len(d), len(x)) containing the pdf(s) or a
                 single value if only 1 point in 1 pdf is selected.
        """

        if d is None:
            # if vaccum:
            p_list = self.param_list
            # else:
            # p_list = self.param_list[1:]
        else:
            p_list = self.param_list[d]

        try:
            xl = len(x)
        except TypeError:
            xl = 1

        f = getattr(self.dist, func)

        # print(p_list.shape)
        if p_list.shape[0] == 1 or len(p_list.shape) == 1:
            if scale:
                return f(x, *p_list[:-1])*p_list[-1]
            else:
                return f(x, *p_list[:-1])
        else:
            f_evals = np.zeros((len(p_list), xl))
            i = 0
            for p in p_list:
                if scale:
                    f_evals[i, :] = f(x, *p[:-1])*p[-1]
                else:
                    f_evals[i, :] = f(x, *p[:-1])
                i += 1
            return f_evals

    def pdf(self, x, d=None, scale=False):
        """
        evaluate the pdf(s) of the distribution(s) selected by d at x.

        :param x: value(s) where to the function is evaluated.
        :param d: selects element(s) of param_list that parameterise the
                  distributions. If None all distributions in param_list are
                  used.
        :param bool scale: when true return the scaled pdf, the scale is always
                           param[-1].
        :return: ndarray shape (len(d), len(x)) containing the pdf(s) or a
                 single value if only 1 point in 1 distribution is selected.
        """
        return self._eval(x, d, scale=scale, func='pdf')

    def cdf(self, x, d=None, scale=False):
        """
        evaluate the pdf(s) of the distribution(s) selected by d at x.

        :param x: value(s) where to the function is evaluated.
        :param d: selects element(s) of param_list that parameterise the
                  distributions. If None all distributions in param_list are
                  used.
        :param bool scale: when true return the scaled pdf, the scale is always
                           param[-1].
        :return: ndarray shape (len(d), len(x)) containing the pdf(s) or a
                 single value if only 1 point in 1 distribution is selected.
        """
        return self._eval(x, d, scale=scale, func='cdf')


@jit
def expectation_maximisation(
    data, initial_thresholds, has_noise_threshold=False,  # fix_noise=None,
    dist=stats.gamma, tol=0.01, max_iter=30, verbose=True
):
    """
    Fit a mixture of distributions to data.

    :param ndarray data: the measurement data.
    :param initial_thresholds: initial thresholds that divide data into
           individual distributions in the mixture.
    :param bool fix_noise: When True keep threshold[1] fixed.
    :param dist: the type of distribution to use (see scipy.stats)
    :param float tol: termination tolerance for change in log_likelihood.
    :param int max_iter: max iterations of expectation maximisation to use.
    :param bool verbose: print progress during optimisation.
    :param bool normalise: passed to maximise step, thresholds are calculated
           using the intersection of the normalised distributions.

    :return: (param_list, zero_loc, log_likelihood, converged) as a named tuple.
            Where param_list is a list of parameter values for distribution in
            the mixture, zero_loc is the arg used to fit, log_likelihood is the
            log_likelihood of the fit, converged is a bool indicating normal
            termination.

    :note: Uses The Expectation maximisation algorithm with hard assignment of
           responsibilities.

    """
    # trim data
    # if has_noise_threshold:
    #     fix_noise = initial_thresholds[1]
    # else:
    fix_noise = None

    fit_data = data[
        and_l(data > initial_thresholds[0], data <= initial_thresholds[-1])
    ]
    thresholds = np.copy(initial_thresholds)
    # last_threshold = initial_thresholds[-1]
    # data = data[data <= last_threshold]

    i = 1
    if verbose:
        print(
            'Starting expectation maximisation, {} {} distributions'
            .format(len(initial_thresholds)-1, dist.name)
        )
        print('Iteration:{} maximisation'.format(i))

    param_list = hard_maximisation(
        fit_data, thresholds, dist=dist,
        # has_noise_threshold=has_noise_threshold
    )
    # return param_list
    if verbose:
        print('Iteration:{} expectation '.format(i))
    new_thresh = hard_expectation(
        param_list, dist=dist, normalise=False
    )

    new_thresh.append(initial_thresholds[-1])
    # return param_list, new_thresh

    ll = mixture_model_ll(
        fit_data, param_list, dist=dist
    )
    # return new_thresh, param_list, ll
    if verbose:
        print('Threshold changes:{!r}'.format(initial_thresholds - new_thresh))
        print('log likelihood:{}'.format(ll))

    converged = False
    thresholds = np.array(new_thresh)
    while i < max_iter:
        i += 1
        if verbose:
            print('Iteration:{} maximisation'.format(i))
        new_param_list = hard_maximisation(
            fit_data, thresholds, dist=dist #,
            #has_noise_threshold=has_noise_threshold
        )
        if verbose:
            print('Iteration:{} expectation '.format(i))
        new_thresh = hard_expectation(
            new_param_list, dist=dist, normalise=False
        )
        new_thresh.append(initial_thresholds[-1])
        new_ll = mixture_model_ll(
            fit_data, new_param_list, #has_noise_threshold=has_noise_threshold,
            dist=dist
        )
        if verbose:
            print(
                'Threshold changes:{!r}'.format(thresholds - new_thresh)
            )
            print('log likelihood:{} change:{}'.format(new_ll, new_ll-ll))

        # terminate if tol met or ll goes down. FIXME is this appropriate?
        if abs(new_ll - ll) <= tol or new_ll < ll:
            if new_ll > ll:
                ll = new_ll
                param_list = new_param_list
            converged = True
            break

        ll = new_ll
        param_list = new_param_list
        thresholds = np.array(new_thresh)

    if verbose:
        if converged:
            print('Converged')
        else:
            print('Maximum iterations reached')

    # calculate the normalised thresholds
    thresholds = hard_expectation(
        param_list, normalise=True, dist=dist
    )
    # add the terminal threshold which is never fitted but marks the end of the
    # data that was used in fitting the model
    thresholds.append(initial_thresholds[-1])

    # calculate the threshold masks

    return MixtureModel(
        np.array(param_list), np.array(thresholds), ll, converged,
        dist, has_noise_threshold
    )


def hard_expectation(
        param_list,
        # fix_noise=None,
        normalise=False, dist=stats.gamma
):
    """
    Calculate the thresholds for a mixture model that define the data
    partitions.

    :param list param_list: list of parameters for the distributions in the
           mixture model, returned by expectation_maximisation().
    :param fix_noise: Fix threshold[1] at this value.
    :param bool normalise: calculate thresholds based on a mixture of
           normalised pdf's.
    :param dist: type of distribution
    :return: ndarray of threshold values.
    """

    # if fix_noise is None:
    #     t = [0]
    #     start = 1
    # else:
    #     t = [0, fix_noise]
    #     start = 2
    t = [0]
    # for m in range(start, len(param_list)):
    for m in range(1, len(param_list)):
        left = param_list[m - 1]
        right = param_list[m]

        if normalise:
            pdf = normalised_pdf
        else:
            pdf = scaled_pdf

        def f(x):
            return pdf(x, left, dist) - pdf(x, right, dist)

        t.append(brentq(f, dist.median(*left[:-1]), dist.median(*right[:-1])))

    return t


def hard_maximisation(
        data, thresholds, dist=stats.gamma, verbose=True
        #has_noise_threshold=False
):
    """
    Fit distributions to the data partitioned by thresholds.

    :param data: the data to model.
    :param thresholds: thresholds that divide data into separate distributions.
    :param dist: type of distribution to fit (scipy.stats).
    :param bool verbose: Print fitted distribution parameters.
    :return: list of parameters for each distribution.
    """
    param_list = []
    #     print('len thresholds',len(thresholds))
    for i in range(1, len(thresholds)):
        part = partition(data, thresholds, i)
        print('{}: {}'.format(i, len(part)))
        if len(part) != 0:
            # if has_noise_threshold and i == 1:
            #     fit = dist.fit(part, floc=0)
            # else:
            fit = dist.fit(part)
            if verbose:
                print('{} distribution:{} params:{}'.format(dist.name, i, fit))
        else:
            fit = []
            if verbose:
                print('{} distribution:{} empty'.format(dist.name, i))

        param_list.append(list(fit) + [len(part) / len(data)])
    return param_list


def mixture_model_ll(
        data, param_list,
        # has_noise_threshold=False,
        dist=stats.gamma
):
    """
    Calculate the log likelihood of a mixture model.

    :param ndarray data: the data used to construct the model.
    :param param_list: list of parameters for each distribution in the model.
    :param dist: the type of distribution to use (see scipy.stats).
    :return: the log likelihood.
    """

    ll = np.empty((len(param_list), len(data)), dtype=np.float64)
    # if has_noise_threshold:
    #     start = 1
    #     params = param_list
    # else:
    #     start = 0

    for i, params in enumerate(param_list):
        p = param_list[i][:-1]
        ll[i, :] = dist.pdf(data, *p) * param_list[i][-1]
    return np.sum(np.log(np.sum(ll, 0)))


@jit
def partition(data, thresholds, i):
    """
    Return the ith partition of the data as defined by thresholds.

    :param ndarray data: the data to partition.
    :param iterable thresholds: the threshold values that partition data.
    :param int i: the partition to return, the ith partition is defined as
                  threshold[i] < data <= threshold[i+1]. The last partition,
                  when i = len(thresholds-1), is defined as
    :return: ndarray of the data in the partition.
    :note: This is used to assign a hard responsibility in the expectation
           maximisation algorithm used to fit data to a mixture model.
    """
    print('part:{} > {} <= {}'.format(i, thresholds[i-1], thresholds[i]))
    return data[
        np.bitwise_and(data > thresholds[i-1], data <= thresholds[i])
    ]


def plot_mixture_model(
    model, data=None, counts=None, xrange=None, normalised=False,
    bins=500, bar=False, measurement_name=r'\texttt{area}', last_threshold=None,
    figsize=None
):
    """
    Plot a mixture model optionally including a histogram of the modeled data.
    if no data is None x must not be None.

    :param MixtureModel model: the mixture model.
    :param ndarray data: the data to histogram.
    :param tuple xrange: the x range to plot.
    :param bool normalised: normalise the model distributions.
    :param bins: passed to numpy.histogram().
    :param bool bar: Plot bars instead of markers.
    :param figsize: passed to matplotlib.pyplot.figure().
    :return: the figure handle.
    """

    points = 10000  # number of x points for plotting pdfs
    if data is not None:
        data_len = len(data)
    if data is None and xrange is None:
        raise AttributeError('Either data or x must be supplied')

    if figsize is None:
        fig = plt.figure()
    else:
        fig = plt.figure(figsize=figsize)

    ax = fig.add_axes([0.12, 0.12, 0.82, 0.85])
    ax.spines['right'].set_visible(False)
    ax.spines['top'].set_visible(False)

    if data is not None:
        if last_threshold is not None:
            """ 
            Introduced because under high zoom matplotlibs pgf back end writes 
            files with dimensions out of range errors when processed by Tex.
            """
            data = data[data < model.thresholds[last_threshold]]

        hist, edges = np.histogram(data, bins=bins)
        w = edges[1]-edges[0]
        bin_centers = edges[:-1]+w/2
        x = np.linspace(edges[0], edges[-1], points)

        # max_p = max(hist)
        # s = sum(hist) * w
        s = data_len * w
        if not normalised:
            cdfs = model.cdf(edges, scale=True)*data_len
            mhist = np.sum(cdfs[:, 1:]-cdfs[:, :-1], 0)

            if bar:
                data_handle = ax.bar(
                    bin_centers, hist, w, align='center',
                    facecolor=(1, 0, 0, 0.2),
                    edgecolor=(1, 1, 1, 1),
                    lw=1, label='data bin'
                )
                model_data_handle = ax.bar(
                    bin_centers, mhist, w, align='center',
                    facecolor=(0, 0, 1, 0.2),
                    edgecolor=(1, 1, 1, 1),
                    lw=1, label='model bin'
                )
            else:
                data_handle, = ax.plot(
                    bin_centers, hist, 's',
                    markeredgecolor='gray',
                    markerfacecolor='none',
                    markersize=4,
                    # markeredgealpha=0.2,
                    # alpha=0.2,
                    lw=1, label='data bin',
                    mew=0.5
                )

        else:
            s = 1
            mhist = None
            hist = None
            bin_centers = None
    else:
        mhist = None
        hist = None
        bin_centers = None
        s = 1
        x = np.linspace(*xrange, 10000)

    ax.set_prop_cycle(None)  # reset colors

    pdfs = np.zeros((len(model.param_list), points), np.float64)
    pdf_handles = []
    if last_threshold is not None:
        num_pdfs = last_threshold
    else:
        num_pdfs = len(model.param_list)
    # if model.has_noise_threshold:
    #     num_pdfs -= 1
    #     start = 1
    # else:
    #     start = 0

    for i in range(0, num_pdfs):

        if normalised:
            pdfs[i, :] = model.pdf(x, d=i)
            lw = 1
        else:
            pdfs[i, :] = model.pdf(x, d=i, scale=True)*s
            lw = 3

        h, = ax.plot(x, pdfs[i, :], lw=lw, label='pdf{}'.format(i))
        pdf_handles.append(h)

    # if data is None or normalised:
    #     ax.set_xlim([-10, max(data)])
    if model.has_noise_threshold:
        max_p = np.max(np.max(pdfs[1:]))
    else:
        max_p = np.max(np.max(pdfs))

    if data is None:
        ax.set_xlim([-10, max(x)])
    else:
        ax.set_xlim([-10, max(data)])
        ax.set_ylim([0, max_p*1.1])


    # ylim = ax.get_ylim()
    # ax.set_ylim([0, ylim[1]])

    if not normalised:
        model_handle, = ax.plot(x, np.sum(pdfs, 0), 'k', lw=1, label='model')

    if not normalised:
        # # FIXME fix_noise??
        # if model.has_noise_threshold:
        #     fix_noise = model.thresholds[1]
        # else:
        #     fix_noise = None
        thresh = np.zeros(len(model.thresholds))
        thresh[:-1] = hard_expectation(
            model.param_list, normalise=False, dist=model.dist
        )
        thresh[-1] = model.thresholds[-1]
    else:
        thresh = model.thresholds

    # if last_threshold is not None:
    #     thresh = thresh[:last_threshold+1]

    # if model.has_noise_threshold:
    #     thresh = thresh[1:]

    # print(thresh)

    for t in thresh[:-1]:
        ax.axvline(
            t, color='k', linestyle='--', lw=0.5, label='thr{}'.format(t)
        )

    if not normalised:
        ax.axvline(
            thresh[-1], color='k', linestyle='-', lw=1, label='thr{}'.format(t)
        )
        # ax.plot([t, t], [0, max_p], ':k', lw=1)

    legend_labels = []
    if data is not None:
        for t in range(len(pdf_handles)):
            if model.has_noise_threshold:
                p = t
            else:
                p = t+1
            if t == 0 and model.has_noise_threshold:
                label = 'noise'
            elif t == len(pdf_handles)-1:
                # print('there')
                # label = '{}+ photons'
                label = '{} photon distribution'
            elif p == 1:
                label = '{} photon distribution'
            else:
                # label = '{} photons'
                label = '{} photon distribution'

            legend_labels.append(label.format(p))

    if not normalised:
        legend_handles = pdf_handles + [model_handle]
    else:
        legend_handles = pdf_handles

    legend_labels.append('model')
    # if counts is not None:
    #     legend_labels.append('model {:,}'.format(sum(counts.count)))
    # else:
    #     legend_labels.append('model')

    if data is not None and not normalised:
        legend_handles += [data_handle]
        legend_labels += ['data bin']
        if bar:
            legend_handles += [model_data_handle]
            legend_labels += ['model bin']

    fig.legend(
        legend_handles, legend_labels, frameon=False,
        bbox_to_anchor=(1, 0.98), bbox_transform=fig.transFigure
    )

    ax.set_xticks(list(thresh))
    exp = int(np.floor(np.log10(model.thresholds[-1])))
    tick_values = model.thresholds/10**exp
    tick_labels = [r'${:.2f}$'.format(tick) for tick in tick_values]
    ax.set_xticklabels(tick_labels)
    ax.set_xlabel(
        r'{} measurement $(\times 10^{{ {} }})$'
        .format(measurement_name, exp)
    )

    ytvalues = ax.get_yticks()
    exp = int(np.floor(np.log10(ytvalues[-1])))
    tick_values = ytvalues/10**exp
    tick_labels = [r'${:1.1f}$'.format(tick) for tick in tick_values]
    ax.set_yticklabels(tick_labels)

    if normalised:
        ax.set_ylabel(r'probability $(\times 10^{{ {} }})$'.format(exp))
    else:
        ax.set_ylabel(r'count $(\times 10^{{ {} }})$'.format(exp))


    return fig, legend_handles, legend_labels, bin_centers, mhist, hist


def threshold_masks(data, model):
    """

    :param data:
    :param model:
    :return:
    """
    masks = np.zeros((len(model.thresholds), len(data)), np.bool)
    for t in range(len(model.thresholds[:-1])):
        if t == 0:
            masks[0, :] = data <= model.thresholds[1]
        else:
            masks[t, :] = and_l(
                data > model.thresholds[t], data <= model.thresholds[t+1]
            )
    masks[-1, :] = data > model.thresholds[-2]
    return masks


"""
timing analysis
"""


@jit
def xcor(s1, s2):
    """
    Count correlations between timestamp sequences.

    :param ndarray s1: Monotonically increasing timestamp sequence.
    :param ndarray s2: Monotonically increasing timestamp sequence.
    :return: correlation count.

    :notes: iterates over s1 and performs a linear search of  s2 for the s1
            timestamp.
    """
    i2 = 0
    c = 0
    for i in range(len(s1)):
        while i2 < len(s2) and s1[i] > s2[i2]:
            i2 += 1
        if i2 < len(s2):
            c += s1[i] == s2[i2]
    return c


@jit
def x_correlation(s1, s2, r):
    """
    Cross correlation of timestamp sequences s1 and s2 over the delay range r.

    :param ndarray s1: Monotonically increasing timestamp sequence.
    :param ndarray s2: Monotonically increasing timestamp sequence.
    :param r: range of delays added to s1
    :return: ndarray representing the cross correlation function.
    """
    x_corr = np.zeros((len(r),), np.uint64)
    for l in r:
        x_corr[l - r[0]] = xcor(s1 + l, s2)
    return x_corr


@jit
def drive_correlation(
    data, model, abs_time, detection_mask, r, last_threshold=-1,
    verbose=False
):
    """
    Calculate the temporal cross-correlation between a channel measuring a
    heralding signal, ie the laser drive pulse, and a channel detecting photons.
    The correlation is calculated for the different photon numbers determined by
    the thresholds parameter.


    :param ndarray abs_time: absolute timestamp sequence.
    :param ndarray detection_mask: boolean mask that identifies the entries in
                                   abs_time belonging to the photon channel,
                                   other entries are assumed to be the heralding
                                   channel.
    :param ndarray data: energy measurement data for the photon channel.
    :param ndarray masks: boolean masks that are appied to
                          abs_time[detection_mask] and indicate which detections
                          are assigned which threshold.
                          Shape(N, len(abs_time[detection_mask]))
    :param r: the range of heralding channel delays to cross correlate.
    :param bool verbose: Print progress message as each threshold is processed.
    :return: list of ndarrays representing the cross correlation.

    :note: abs_time[mask] and data must be the same length.
           TODO speed up the algorithm.
    """

    if last_threshold > len(model.thresholds):
        raise AttributeError('last_threshold must be < len(model.thresholds)')
    if last_threshold < 0:
        last_threshold = len(model.thresholds)+last_threshold

    xc = []
    if last_threshold == 0:
        if verbose:
            print('Cross correlating all events')
        xc = x_correlation(
            abs_time[not_l(detection_mask)], abs_time[detection_mask], r
        )
    else:
        for t in range(last_threshold+1):
            if verbose:
                print(
                    'Cross correlating partition {} of {}'
                    .format(t+1, last_threshold+1))

                if t != last_threshold:
                    xc.append(
                        x_correlation(
                            abs_time[not_l(detection_mask)],
                            abs_time[detection_mask][
                                and_l(
                                    data > model.thresholds[t],
                                    data <= model.thresholds[t+1]

                                )

                            ],
                            r
                        )
                    )
                else:
                    xc.append(
                        x_correlation(
                            abs_time[not_l(detection_mask)],
                            abs_time[detection_mask][data > model.thresholds[t]]
                            , r
                        )
                    )

    return np.array(xc)


@jit
def window(i, abs_time, low, high):
    """
    get offsets from the current index i in abs_time that are in the relative
    time window defined by low and high.

    :param int i: current abs_time index.
    :param ndarray abs_time: absolute times.
    :param low: low end of relative window.
    :param high: high end of relative window.
    :return: (low_o, high_o) coincident times are abs_time[low_o:i+high_o]
    """

    length = len(abs_time)
    low_o = 0  # low index offset
    high_o = 0  # high index offset
    now = abs_time[i]
    low_t = now + low
    high_t = now + high

    if low_t < now:
        while (i + low_o >= 0) and abs_time[i + low_o] >= low_t: low_o -= 1;
    else:
        while (low_o + i < length) and abs_time[i + low_o] < low_t: low_o += 1;

    if high_t < now:
        while (i + high_o >= 0) and abs_time[i + high_o] > high_t:
            high_o -= 1
    else:
        while (high_o + i < length) and abs_time[i + high_o] <= high_t:
            high_o += 1
    return low_o, high_o  # offset indexs marking abs_times in the window


def coincidence(abs_time, mask, low, high):
    """
    Find coincidences between two channels returning indices of partner events.

    :param ndarray abs_time: absolute times.
    :param ndarray mask: ndarray of bool that identifies the channels in
                         abs_time. For abs_times where mask is False, time t is
                         coincident if abs_time+low <= t <= abs_time+high. If
                         the mask is True t is coincident if
                         abs_time-high <= t <= abs_time-low.
    :param float low: the low side of the coincidence window.
    :param float high:  the high side of the coincidence whindw.
    :return: (coinc, coinc_mask) where coinc is a ndarray of indices of the
             coincident event in the other channel. When more than one event is
             found in the window the negated value of the first index is
             entered.
             coinc_mask is a ndarray of bool indicating where exactly one
             event was found in the window.
    """

    coinc = np.zeros(len(abs_time), np.int32)
    coinc_mask = np.zeros(len(abs_time), np.bool)
    for t in range(len(abs_time)):
        if t % 10000000 == 0:
            print(t)
        if not mask[t]:
            low_o, high_o = window(t, abs_time, low, high)
            o = low_o
        else:
            low_o, high_o = window(t, abs_time, -high, -low)
            o = high_o
        coinc_length = high_o - low_o
        if high_o and low_o:
            if coinc_length == 1:
                coinc[t] = t + o
                coinc_mask[t] = True
            elif coinc_length > 1:
                coinc[t] = -(t + o)
                coinc_mask[t] = False
    return coinc, coinc_mask


Counts = namedtuple(
    'Counts', 'count uncorrelated vacuum has_noise_threshold'
)


def count(
        data, measurement_model, vacuum=False, has_noise_threshold=False,
        coinc_mask=None, herald_mask=None
):
    """

    :param ndarray data: measurement data.
    :param MixtureModel measurement_model: model created using
                                           expectation_maximisation.
    :param bool vacuum: estimate vacuum terms using coinc_mask and herald_mask.
    :param coinc_mask: ndarray of bool indicating coincidence.
    :param herald_mask: ndarray of bool indicating which events are in the
                        heralding channel.
    :param has_noise_threshold: When True model.thresholds[1] represents the
                                boundary between noise and 1 photon.
    :return: Counts(count, uncorrelated, vacuum). Where count is an ndarray of
             counts for each photon number
             and len(count)=len(measurement_model.thresholds). uncorrelated is a
             ndarray the same shape as count and counts
             events not correlated with the herald. uncorrelated is only
             calculated when coinc_mask and herald_mask are supplied.
             vacuum replicates the parameter value and indicates that count[0]
             contains the vacuum count.
    """

    if vacuum and (coinc_mask is None or herald_mask is None):
        raise AttributeError(
            'Neither coinc_mask or herald_mask cannot be None when vacuum=True'
        )

    if (
        (herald_mask is not None and coinc_mask is None) or
        (herald_mask is None and coinc_mask is not None)
    ):
        raise AttributeError(
            'coinc_mask or herald_mask both be specified or both None'
        )

    if herald_mask is not None:
        data_mask = not_l(herald_mask)
        uncorrelated = []
    else:
        data_mask = None
        uncorrelated = None

    counts = []
# thresholds
    for t in range(1, len(measurement_model.thresholds)):
        counts.append(
            len(
                and_l(
                    (data > measurement_model.thresholds[t - 1]),
                    (data <= measurement_model.thresholds[t])
                ).nonzero()[0]
            )
        )
        if uncorrelated is not None:
            uncorrelated.append(
                len(
                    and_l(
                        data[not_l(coinc_mask)[data_mask]] >
                        measurement_model.thresholds[t-1],
                        data[not_l(coinc_mask)[data_mask]] <=
                        measurement_model.thresholds[t]
                    ).nonzero()[0]
                )
            )

    # last threshold overflow
    counts[-1] += len((data > measurement_model.thresholds[-1]).nonzero()[0])
    if uncorrelated is not None:
        uncorrelated[-1] += (
            len(
                (data[not_l(coinc_mask)[data_mask]] >
                    measurement_model.thresholds[-1]).nonzero()[0]
            )
        )

    # counting vacuum
    # uncorrelated heralds + correlated data <= the first threshold
    if vacuum:
        uncorrelated_heralds = len(
            (not_l(coinc_mask)[herald_mask]).nonzero()[0]
        )
        if has_noise_threshold:
            noise = len(
                (
                    data[coinc_mask[data_mask]] <=
                    measurement_model.thresholds[1]
                ).nonzero()[0]
            )
            counts[0] = noise+uncorrelated_heralds
            # note these are correlated events less that the noise threshold
            if uncorrelated is not None:
                uncorrelated[0] = uncorrelated_heralds
        else:
            counts.insert(0, uncorrelated_heralds)
            if uncorrelated is not None:
                uncorrelated.insert(0, uncorrelated_heralds)

    if uncorrelated is not None:
        uncorrelated = np.array(uncorrelated)

    return Counts(
        np.array(counts), uncorrelated, vacuum, has_noise_threshold
    )


Povm = namedtuple(
    'povm', 'elements vacuum'
)


# FIXME remove counts
# def povm_elements(measurement_model, counts):
#     """
#     estimate the elements of the system povm from the measurement model and the
#     counting data.
#
#     :param array like data: measurement data.
#     :param mixturemodel measurement_model: model created using
#                                            expectation_maximisation.
#     :param Counts counts: the count data returned by count().
#     :return: (elements count vacuum). where counts is a ndarray with
#              shape (len(n)) containing the count for each measurement outcome
#              in n. elements is a ndarray with
#              shape(len(n), len(measurement_model.thresholds)) the first index is
#              the fock state, the second is the outcome. The contents of elements
#              is the probability of the outcome when measuring fock state.
#
#     :notes: POVM elements is indexed by (fock state, measurement outcome).
#             The last threshold in the measurement model marks the boundary of
#             the data that was used to create the model, it is *not* altered by
#             the expectation maximisation algorithm and is essentially a guess.
#             The second last threshold is the boundary of the overflow bin,
#             this last photon number bin is counts detections of >=
#             len(measurement_model.param_list)-2 photons.
#             TODO expand and clarify description.
#     """
#
#     # t thresholds
#     # t-1 distributions
#     # t+1 counts if vacuum else t count[-1] is overflow past last threshold
#
#
#     num_elements = len(counts.count)
#     # cumulative density functions at the thresholds, the last value is not used
#     # normalise to a convex mixture excluding the noise distribution.
#     cdfs = measurement_model.cdf(measurement_model.thresholds)
#
#     elems = cdfs[:, 1:] - cdfs[:, :-1]
#     # Adjust the first elements, which should not have a proceeding threshold.
#     elems[:, 0] = cdfs[:, 1]
#
#     # Adjust the last elements to account for the truncation at the last
#     # threshold (limit of data used to create the model).
#     elems[:, -1] = 1 - cdfs[:, -2]
#
#     # elements should always have a vacuum term.
#     elements = np.zeros((num_elements, num_elements))
#
#     # the majority of the POVM elements are given by the difference in the CDFs
#     # at the threshold values
#
#     if counts.vacuum:
#         if measurement_model.has_noise_threshold:
#             elements = elems
#         else:
#             elements[1:, 1:] = elems
#     else:
#         elements = elems
#
#     # vacuum terms
#     # FIXME check this is correct when has_noise_threshold
#     if counts.vacuum:
#         # probability of counting vacuum when there is a vacuum input
#         # TODO is it appropriate to do same for m > 1 ??
#         for m in range(1, 2):  # range(1, len(counts.uncorrelated)):
#             elements[0, 0] = (
#                     (counts.count[0]-counts.uncorrelated[m])/counts.count[0]
#             )
#             elements[0, m] = counts.uncorrelated[m]/counts.count[0]
#
#          # elements[m, 0] = correlatedsingle/single
#
#         if measurement_model.has_noise_threshold:
#             # look for uncrrelated ???
#             pass
#
#     else:
#         elements[0, 0] = 1
#         elements[0, 1] = 0
#
#     # c = np.array(counts.count[:-1])
#     # c[-1] += counts.count[-1]
#
#     return Povm(elements, counts.vacuum, counts.has_noise_threshold)


def fit_state_least_squares(
        counts, povm, max_outcome=None, vacuum=True, thermal=True, N=150
):

    # FIXME check this model with new POVM remove call to outcomes
    def qutip_d_thermal(n, nbar, alpha, N):
        # n [low, high]
        t = qt.thermal_dm(N, nbar)
        D = qt.displace(N, alpha)
        s = D * t * D.dag()

        #         print('n in model',n)
        state = [
                    (qt.states.fock(N, num) *
                     qt.states.fock(N, num).dag() * s).tr()
                    for num in np.arange(0, N)
                ]
        # print('state sum',np.sum(state))

        # print('state.shape', state.shape)
        num_elements = min(povm.shape[1], n[1]-n[0]+1)
        # print('num elements', num_elements, 'N', N)
        truncated_state = np.real(state[n[0]:num_elements + 1])
        truncated_state[-1] += np.real(np.sum(state[num_elements + 1:]))
        # print('nbar', nbar, 'alpha', alpha)
        # print('sum state', np.sum(state[n[0]:]))
        # print('sum t state', np.sum(truncated_state[n[0]:]))
        # print(
        #     'trunc test',
        #     np.real(np.sum(state[n[0]:]))-np.sum(truncated_state[n[0]:])
        # )
        if num_elements < povm.shape[0]:
            truncated_povm = povm[:, n[0]:num_elements]
            truncated_povm[-1] += np.sum(povm[num_elements:])
        else:
            truncated_povm = povm
        outcomes = np.zeros_like(truncated_povm)
        #         print('outcomes',outcomes.shape)
        for e in range(num_elements):
            outcomes[e, :] += truncated_state[:-1] * truncated_povm[e, :]
        out = np.sum(outcomes, 0)
        out[-1] += truncated_state[-1]
        # if n[0] != 0:
        #     out = out / np.sum(out)
        #       print('out',out.shape)
        return out

    n = np.zeros(2, np.uint32)
    if vacuum:
        n[0] = 0
    else:
        n[0] = 1

    if max_outcome is not None:
        n[1] = max_outcome
        c = counts[n[0]:n[1] + 1]
        c[-1] += np.sum(counts[n[1] + 1:])
    else:
        n[1] = povm.shape[0]+1
        c = counts[n[0]:]

    print('fit c shape', c.shape)

    state_model = Model(qutip_d_thermal)
    # alpha = 0.1
    pars = state_model.make_params(nbar=0.1, alpha=0.1, N=N)
    pars['nbar'].set(min=0, max=0.5)
    pars['alpha'].set(min=0, max=10)
    pars['nbar'].set(vary=thermal)
    if not thermal:
        pars['nbar'].set(value=0.0)
    pars['N'].set(vary=False)
    #     print('n',n)
    #     pars['povm_elements'].set(vary=False)
    fit = state_model.fit(c / np.sum(c), pars, n=n)
    print('fit n', n)
    print('fit c', c)
    print('fit sum(c)', np.sum(c))
    print('fit best_fit', fit.best_fit)
    print('fit sum(best_fit)', np.sum(fit.best_fit))
    print('fit best_fit counts', fit.best_fit * np.sum(c))
    return c, n, fit


def resize_vector(a, max_index, copy=True):

    if max_index+1 > len(a):
        raise AttributeError('max_index must be <= len(a)')
    if max_index+1 == len(a):
        return a

    if copy:
        o = np.copy(a[:max_index+1])
    else:
        o = a[:max_index+1]
    o[-1] += np.sum(a[max_index+1:])
    return o


def displaced_thermal(max_photon_number, nbar, alpha, N=200):
    # n [low, high]
    t = qt.thermal_dm(N, nbar)
    D = qt.displace(N, alpha)
    s = D * t * D.dag()

    state = [
        (qt.states.fock(N, num)*qt.states.fock(N, num).dag()*s).tr()
        for num in np.arange(0, N)
    ]

    truncated_state = resize_vector(np.real(state), max_photon_number)

    return truncated_state #/np.sum(truncated_state)


def displaced_vacuum(max_photon_number, alpha, N=200):
    # n [low, high]
    t = qt.thermal_dm(N, 0)
    D = qt.displace(N, alpha)
    s = D * t * D.dag()

    state = [
        (qt.states.fock(N, num)*qt.states.fock(N, num).dag()*s).tr()
        for num in np.arange(0, N)
    ]

    truncated_state = resize_vector(np.real(state), max_photon_number)

    return truncated_state #/np.sum(truncated_state)


def outcome_probabilities(state, povm, max_photon_number):

    if max_photon_number is None:
        max_photon_number = len(povm.elements) - 2
    elif max_photon_number > len(povm.elements) - 2:
        raise AttributeError(
            'max_photon_number must be at least 2 less than the number of '
            'POVM elements ({})'.format(len(povm.elements) - 2)
        )

    outcomes = np.zeros((max_photon_number+1, max_photon_number+1))

    for i in range(max_photon_number+1):
        outcomes[i, :] = resize_vector(
            state*povm.elements[i, :], max_photon_number
        )
    return outcomes


def thermal_neg_ll(x, counts, povm, max_photon_number):

    c = resize_vector(counts.count, max_photon_number)
    state = displaced_thermal(max_photon_number, *x, N=100)
    if povm is not None:
        outcomes = povm.elements*state
        if len(np.sum(outcomes, 0).nonzero()[0]) != outcomes.shape[1]:
            print('inf')
            ll = -np.inf
        else:
            ll = np.sum(np.log(np.sum(outcomes, 0))*c)
    else:
        if len(state.nonzero()[0]) != len(state):
            print('inf')
            ll = -np.inf
        else:
            ll = np.sum(np.log(state)*c)

    return -ll


def coherent_neg_ll(x, counts, povm, max_photon_number):
    c = resize_vector(counts.count, max_photon_number)
    state = displaced_vacuum(max_photon_number, *x, N=100)
    if povm is not None:
        outcomes = povm.elements*state
        if len(np.sum(outcomes, 0).nonzero()[0]) != outcomes.shape[1]:
            print('inf')
            ll = -np.inf
        else:
            ll = np.sum(np.log(np.sum(outcomes, 0))*c)
    else:
        if len(state.nonzero()[0]) != len(state):
            print('inf')
            ll = -np.inf
        else:
            ll = np.sum(np.log(state)*c)

    return -ll


def thermal_ml(counts, povm, max_photon_number, x0=None):
    if x0 is None:
        x0 = [0.3, 3.0]
    r = minimize(
        thermal_neg_ll, x0, args=(counts, povm, max_photon_number),
        method='L-BFGS-B', tol=1e-15,
        bounds=[(0.0, 2.0), (0.1, 6.0)],
        options={
            'disp': True, 'eps': 1e-6, 'ftol': 1e-15, 'gtol': 1e-15, 'maxls': 50
        }
    )
    return r


def coherent_ml(counts, povm, max_photon_number, x0=None):
    if x0 is None:
        x0 = [3.0]
    r = minimize(
        coherent_neg_ll, x0, args=(counts, povm, max_photon_number),
        method='L-BFGS-B', tol=1e-15,
        bounds=[(0.1, 6.0)],
        options={
            'disp': True, 'eps': 1e-6, 'ftol': 1e-15, 'gtol': 1e-15, 'maxls': 50
        }
    )
    return r


def plot_state_fit(nbar, alpha, counts, povm=None, figsize=None):
    width = 0.46   # the width of the bars
    if figsize is not None:
        fig = plt.figure(figsize=figsize)
    else:
        fig = plt.figure()

    ax = fig.add_axes([0.12, 0.12, 0.84, 0.83])

    max_photon_number = len(counts.count)-1

    state = displaced_thermal(max_photon_number, nbar, alpha, N=200)
    if povm is not None:
        model_counts = np.sum(povm.elements*state*np.sum(counts.count), 1)
    else:
        model_counts = state*np.sum(counts.count)

    fock = np.arange(0, max_photon_number+1)
    # print(fock.shape, fock)

    ax.bar(fock-width/2, counts.count, width=width, color='r', label='data')
    ax.bar(
        fock+width/2, model_counts, width=width, color='b', label='model'
    )
    # chi2 = sum(np.power(c - model_counts, 2) / model_counts)
    ax.set_ylabel('count')
    ax.set_xlabel('Fock state')
    ax.set_xticks(fock)
    xlabels = [r'$\vert {} \rangle$'.format(f) for f in fock]
    xlabels[-1] = r'$\vert {}\rangle{}$'.format(fock[-1], '+')
    ax.set_xticklabels(xlabels)

    ax.legend(frameon=False)
    # plt.plot(n,nc,'.')
    # plt.plot(n,r.best_fit,'o',markerfacecolor='none')
    # ax.set_ylim(0,11000000)
    ty = 0.8
    tx = 0.01
    ls = 0.05
    ax.text(tx, ty, r'Displaced thermal state', transform=ax.transAxes)
    # ax.text(
    #     tx, ty - ls, r'$\alpha={:.4f}\pm{:.4f}$'
    #         .format(
    #         state_fit.params['alpha'].value, state_fit.params['alpha'].stderr
    #     ),
    #     transform=ax.transAxes, horizontalalignment='left'
    # )
    # ax.text(
    #     tx, ty - 2 * ls, r'${}={:.4f}\pm{:.4f}$'
    #         .format(
    #         r'\bar{n}', state_fit.params['nbar'].value,
    #         state_fit.params['nbar'].stderr
    #     ),
    #     transform=ax.transAxes, horizontalalignment='left'
    # )
    # ax.text(
    #     tx, ty - 3 * ls, r'$\chi^2/{}={:.3f}$'
    #         .format(
    #         state_fit.nfree - 1, chi2 / (state_fit.nfree - 1)
    #     ),
    #     transform=ax.transAxes, horizontalalignment='left'
    # )
    ax.legend(frameon=False)
    return fig


def povm_elements(model):
    cdfs = model.cdf(model.thresholds)
    elems = np.zeros((cdfs.shape[0]+1, cdfs.shape[1]))
    elems[1:, 1:] = cdfs[:, 1:] - cdfs[:, :-1]
    elems[1:, 0] = cdfs[:, 0]
    elems[1:, -1] = 1-cdfs[:, -2]
    elems[0, 0] = 1
    return elems


def model_outcomes(x, N, thermal=False, povm=None, qtN=100):
    """
    Model the measurement outcome probabilities assuming the input is a
    displaced thermal state.

    :param x iterable: The state parameters
    :param N int: Dimension of the Fock space, max photon number is N-1
    :param thermal: When True model input as displaced thermal state otherwise
                    model as coherent state (displaced thermal).
    :param povm: The detector povm
    :param qtN: The dimension of the Fock space to use for qutip models.
    """
    #     print(x)
    D = qt.displace(qtN, x[0])
    if thermal:
        #         print('thermal')
        t = qt.thermal_dm(qtN, x[1])
    else:
        t = qt.thermal_dm(qtN, 0)

    s = D * t * D.dag()

    state = np.real(np.array([
        (qt.states.fock(qtN, n) * qt.states.fock(qtN, n).dag() * s).tr()
        for n in np.arange(0, qtN)
    ]))
    #     print(np.sum(state))
    if povm is not None:
        n_ele = povm.shape[1]  # number of povm elements
        if N > n_ele:
            raise AttributeError(
                'N must be less than or equal to the number of POVM elements'
            )
        o = np.sum(state[:n_ele] * povm.transpose(), 1)  #
        o[-1] += np.sum(state[n_ele:])  # truncate to number of elements
        if N < n_ele:  # truncate to N if necessary
            ot = np.array(o[:N], copy=True)
            ot[-1] += np.sum(o[N:])
            o = ot
    else:
        o = np.array(state[:N], copy=True)
        o[-1] += np.sum(state[N:])
    return o


def fit_input_ls(
        counts, N, povm=None, thermal=False, qtN=150
):
    def dt_model(N, alpha, nbar):
        if thermal:
            x = [alpha, nbar]
        else:
            x = [alpha]

        o = model_outcomes(x, N, thermal=thermal, povm=povm, qtN=qtN)

        return o * total

    if N > len(counts):
        raise AttributeError('N cannot be greater than the number of counts')
    if N < len(counts):
        c = np.array(counts[:N], copy=True)
        c[-1] += np.sum(counts[N:])
        counts = c

    model = Model(dt_model, nan_policy='omit')

    pars = model.make_params(nbar=0.1, alpha=1.0)
    pars['nbar'].set(min=0, max=0.5)
    pars['alpha'].set(min=0, max=10)
    pars['nbar'].set(vary=thermal)

    if not thermal:
        pars['nbar'].set(value=0.0)

    total = np.sum(counts)
    return model.fit(
        counts, pars, N=N, weights=np.sqrt(1.0 / counts)
    )


def xcor_plot(xc, r, has_noise_threshold=False, figsize=None):

    if figsize is not None:
        fig = plt.figure(figsize=figsize)
    else:
        fig = plt.figure()

    ax = fig.add_axes([0.12, 0.12, 0.85, 0.85])
    ax.spines['right'].set_visible(False)
    ax.spines['top'].set_visible(False)

    sxc = np.sum(xc, 0)

    ax.bar(
        r[sxc != 0] * 4, sxc[sxc != 0], 4,
        align='center',
        facecolor=(0.01, 0.01, 0.01, 0.1),
        edgecolor=(0, 0, 0, 0.5),
        lw=0.5, label=r'Total'
    )

    if xc.shape[0] > 1:
        for x in range(xc.shape[0] - 1):
            r_nz = r[xc[x] != 0] * 4
            xc_nz = xc[x][xc[x] != 0]
            if has_noise_threshold:
                p = x
            else:
                p = x + 1
            if x == 0 and has_noise_threshold:
                label = 'noise'
            elif x == xc.shape[0] - 2:
                r_nz = r[np.sum(xc[x:]) != 0] * 4
                xc_nz = np.sum(xc[x:])[np.sum(xc[x:]) != 0]
                label = '{}+ photons'.format(p)
            elif p == 1:
                label = '{} photon'.format(p)
            else:
                label = '{} photons'.format(p)
            ax.plot(
                r[xc[x] != 0] * 4, xc[x][xc[x] != 0], '-', markersize=5,
                markerfacecolor='none', mew=0.5, label=label, lw=1
            )

    ax.set_xlabel(r'detection time relative to herald (ns)')
    ax.set_ylabel(r'count')
    #     ax.set_xlim([200,800])
    ax.frameon = False
    fig.legend(
        bbox_to_anchor=(0.33, 0.98), bbox_transform=fig.transFigure,
        frameon=False
    )
    return fig


def plot_ls_state(r, counts, N, thermal=False, figsize=None):
    if figsize is not None:
        fig = plt.figure(figsize=figsize)
    else:
        fig = plt.figure()

    width = 0.48  # the width of the bars
    ax = fig.add_axes([0.12, 0.12, 0.84, 0.83])
    n = np.arange(0, r.ndata)
    #     print(len(counts))
    if N > len(counts):
        raise AttributeError('N cannot be greater than len(counts)')
    if N < len(counts):
        c = np.array(counts[:N], copy=True)
        c[-1] += np.sum(counts[N:])
    else:
        c = counts

    #     print(n[0],len(n),len(c))
    rects1 = ax.bar(n - width / 2, c, width, color='r')
    rects2 = ax.bar(n + width / 2, r.best_fit, width, color='b'
                    )
    #     chi2=sum(np.power(counts-r.best_fit,2)/model_counts)
    ax.set_ylabel('count')
    ax.set_xlabel('Fock state')
    ax.set_xticks(n)
    xlabels = [r'$\vert {} \rangle$'.format(p) for p in n]
    xlabels[-1] = r'$\vert {}\rangle{}$'.format(n[-1], '+')
    ax.set_xticklabels(xlabels)

    ax.legend((rects1[0], rects2[0]), ('data', 'model'), frameon=False)
    # plt.plot(n,nc,'.')
    # plt.plot(n,r.best_fit,'o',markerfacecolor='none')
    # ax.set_ylim(0,11000000)
    ty = 0.8
    tx = 0.7
    ls = 0.05
    if thermal:
        ax.text(tx, ty, r'Displaced thermal state', transform=ax.transAxes)
    else:
        ax.text(tx, ty, r'Coherent state', transform=ax.transAxes)
    ax.text(
        tx, ty - ls, r'$\alpha={:.4f}\pm{:.4f}$'
            .format(
            r.params['alpha'].value, r.params['alpha'].stderr
        ),
        transform=ax.transAxes, horizontalalignment='left'
    )
    if thermal:
        ax.text(
            tx, ty - 2 * ls, r'${}={:.4f}\pm{:.4f}$'
                .format(
                r'\bar{n}', r.params['nbar'].value, r.params['nbar'].stderr
            ),
            transform=ax.transAxes, horizontalalignment='left'
        )
    ax.text(
        tx, ty - 3 * ls, r'$\chi^2/{}={:.3f}$'
            .format(
            r.ndata - r.nvarys - 1, r.chisqr / (r.ndata - r.nvarys - 1)
        ),
        transform=ax.transAxes, horizontalalignment='left'
    )
    ax.legend(
        (rects1[0], rects2[0]), ('model', 'data'),
        bbox_to_anchor=(tx, ty - 0.34), loc='lower left',
        frameon=False
    )
    return fig

Calibration and modeling Find local maxima of f using rising zero crossings of the gradient of f.      :param f: function to find maxima for     :param thresh: Only return a maxima if f[maxima]>thresh      :return: array of maxima.     :Notes: Used to find the maxima of the smoothed measurement histogram. # indices of the maxima Guess the photon number thresholds using a smoothed histogram.      :param ndarray data: the measurement data.     :param bins: argument to numpy.histogram.     :param int win_length: Hanning window length used for smoothing.     :param int maxima_threshold: passed to maxima function.      :return: (hist, f, bin_c, max_i, thresholds) as a named             tuple. Where hist is the histogram             generated by np.histogram(), f is the smoothed histogram, bin_c is a             ndarray of bin centers, max_i are the indices of the maxima in f,             and thresholds is a ndarray of threshold guesses.      :notes: The smoothing is achieved by convolving the Hanning window with a             histogram generated using numpy.histogram. The guess still needs             some human sanity checks, especially multiple maxima per photon             peak. Use plot_guess() and adjust window and maxima_threshold to             remove them. # smooth the histogram # indices of the maxima # ax.set_title('Histogram and initial threshold guess') # @jit # TODO generalise # remove incorporate tick relative times #     events = edat.view(event_dt) # tidx[0] is the first tick period ending with tick 1 # i+1 is the tick number #         print(event_count) #     print(event_i) Convenience function, calculate normalised pdf at each x for dist.      :param union[numpy.ndarray, int] x: calculate the pdf a each x     :param iterable params: fitted distribution parameters     :param dist: type of distribution (see scipy.stats)     :return: ndarray containing pdf(x) Convenience function, calculate pdf at each x for dist, normalisation is     given by params[-1].      :param union[numpy.ndarray, int] x: calculate the pdf a each x     :param iterable params: fitted distribution parameters     :param dist: type of distribution (see scipy.stats)     :return: ndarray containing pdf(x) Subclass of namedtuple used to represent a mixture model.      fields:         :param_list: list of parameters for each distribution in the mixture.         :thresholds: the intersection of neighbouring distributions in the                      mixture.         :zero_loc: Fixed location parameter used to fit the first distribution,                    or None if the location parameter was fitted.         :log_likelihood: the log of the likelihood that the model could produce                          the data used to construct it.         :converged: boolean indicating that the expectation maximisation                     algorithm terminated normally when fitting the model.         :dist: the type of distribution forming the components of the mixture,                (see scipy.stats). save model as .npz file.          :param filename: filename to save as, excluding extension"         :return: None # zero_loc=self.zero_loc, load model from .npz file.          :param filename: filename excluding extension"         :return: instance of MixtureModel. # d['zero_loc'][()], evaluate the function of the distribution(s) selected by d at x.          :param x: value(s) where to the function is evaluated.         :param d: selects element(s) of param_list that parameterise the                   distributions. If None all distributions in param_list are                   used.         :param bool scale: when true return the scaled pdf, the scale is always                            param[-1].         :param str func: the name of the distribution function to evaluate                              (see scipy.stats.rv_continuous)         :return: ndarray shape (len(d), len(x)) containing the pdf(s) or a                  single value if only 1 point in 1 pdf is selected. # if vaccum: # else: # p_list = self.param_list[1:] # print(p_list.shape) evaluate the pdf(s) of the distribution(s) selected by d at x.          :param x: value(s) where to the function is evaluated.         :param d: selects element(s) of param_list that parameterise the                   distributions. If None all distributions in param_list are                   used.         :param bool scale: when true return the scaled pdf, the scale is always                            param[-1].         :return: ndarray shape (len(d), len(x)) containing the pdf(s) or a                  single value if only 1 point in 1 distribution is selected. evaluate the pdf(s) of the distribution(s) selected by d at x.          :param x: value(s) where to the function is evaluated.         :param d: selects element(s) of param_list that parameterise the                   distributions. If None all distributions in param_list are                   used.         :param bool scale: when true return the scaled pdf, the scale is always                            param[-1].         :return: ndarray shape (len(d), len(x)) containing the pdf(s) or a                  single value if only 1 point in 1 distribution is selected. # fix_noise=None, Fit a mixture of distributions to data.      :param ndarray data: the measurement data.     :param initial_thresholds: initial thresholds that divide data into            individual distributions in the mixture.     :param bool fix_noise: When True keep threshold[1] fixed.     :param dist: the type of distribution to use (see scipy.stats)     :param float tol: termination tolerance for change in log_likelihood.     :param int max_iter: max iterations of expectation maximisation to use.     :param bool verbose: print progress during optimisation.     :param bool normalise: passed to maximise step, thresholds are calculated            using the intersection of the normalised distributions.      :return: (param_list, zero_loc, log_likelihood, converged) as a named tuple.             Where param_list is a list of parameter values for distribution in             the mixture, zero_loc is the arg used to fit, log_likelihood is the             log_likelihood of the fit, converged is a bool indicating normal             termination.      :note: Uses The Expectation maximisation algorithm with hard assignment of            responsibilities. # trim data # if has_noise_threshold: #     fix_noise = initial_thresholds[1] # else: # last_threshold = initial_thresholds[-1] # data = data[data <= last_threshold] # has_noise_threshold=has_noise_threshold # return param_list # return param_list, new_thresh # return new_thresh, param_list, ll #, #has_noise_threshold=has_noise_threshold #has_noise_threshold=has_noise_threshold, # terminate if tol met or ll goes down. FIXME is this appropriate? # calculate the normalised thresholds # add the terminal threshold which is never fitted but marks the end of the # data that was used in fitting the model # calculate the threshold masks # fix_noise=None, Calculate the thresholds for a mixture model that define the data     partitions.      :param list param_list: list of parameters for the distributions in the            mixture model, returned by expectation_maximisation().     :param fix_noise: Fix threshold[1] at this value.     :param bool normalise: calculate thresholds based on a mixture of            normalised pdf's.     :param dist: type of distribution     :return: ndarray of threshold values. # if fix_noise is None: #     t = [0] #     start = 1 # else: #     t = [0, fix_noise] #     start = 2 # for m in range(start, len(param_list)): #has_noise_threshold=False Fit distributions to the data partitioned by thresholds.      :param data: the data to model.     :param thresholds: thresholds that divide data into separate distributions.     :param dist: type of distribution to fit (scipy.stats).     :param bool verbose: Print fitted distribution parameters.     :return: list of parameters for each distribution. #     print('len thresholds',len(thresholds)) # if has_noise_threshold and i == 1: #     fit = dist.fit(part, floc=0) # else: # has_noise_threshold=False, Calculate the log likelihood of a mixture model.      :param ndarray data: the data used to construct the model.     :param param_list: list of parameters for each distribution in the model.     :param dist: the type of distribution to use (see scipy.stats).     :return: the log likelihood. # if has_noise_threshold: #     start = 1 #     params = param_list # else: #     start = 0 Return the ith partition of the data as defined by thresholds.      :param ndarray data: the data to partition.     :param iterable thresholds: the threshold values that partition data.     :param int i: the partition to return, the ith partition is defined as                   threshold[i] < data <= threshold[i+1]. The last partition,                   when i = len(thresholds-1), is defined as     :return: ndarray of the data in the partition.     :note: This is used to assign a hard responsibility in the expectation            maximisation algorithm used to fit data to a mixture model. Plot a mixture model optionally including a histogram of the modeled data.     if no data is None x must not be None.      :param MixtureModel model: the mixture model.     :param ndarray data: the data to histogram.     :param tuple xrange: the x range to plot.     :param bool normalised: normalise the model distributions.     :param bins: passed to numpy.histogram().     :param bool bar: Plot bars instead of markers.     :param figsize: passed to matplotlib.pyplot.figure().     :return: the figure handle. # number of x points for plotting pdfs Introduced because under high zoom matplotlibs pgf back end writes              files with dimensions out of range errors when processed by Tex. # max_p = max(hist) # s = sum(hist) * w # markeredgealpha=0.2, # alpha=0.2, # reset colors # if model.has_noise_threshold: #     num_pdfs -= 1 #     start = 1 # else: #     start = 0 # if data is None or normalised: #     ax.set_xlim([-10, max(data)]) # ylim = ax.get_ylim() # ax.set_ylim([0, ylim[1]]) # # FIXME fix_noise?? # if model.has_noise_threshold: #     fix_noise = model.thresholds[1] # else: #     fix_noise = None # if last_threshold is not None: #     thresh = thresh[:last_threshold+1] # if model.has_noise_threshold: #     thresh = thresh[1:] # print(thresh) # ax.plot([t, t], [0, max_p], ':k', lw=1) # print('there') # label = '{}+ photons' # label = '{} photons' # if counts is not None: #     legend_labels.append('model {:,}'.format(sum(counts.count))) # else: #     legend_labels.append('model') :param data:     :param model:     :return: timing analysis Count correlations between timestamp sequences.      :param ndarray s1: Monotonically increasing timestamp sequence.     :param ndarray s2: Monotonically increasing timestamp sequence.     :return: correlation count.      :notes: iterates over s1 and performs a linear search of  s2 for the s1             timestamp. Cross correlation of timestamp sequences s1 and s2 over the delay range r.      :param ndarray s1: Monotonically increasing timestamp sequence.     :param ndarray s2: Monotonically increasing timestamp sequence.     :param r: range of delays added to s1     :return: ndarray representing the cross correlation function. Calculate the temporal cross-correlation between a channel measuring a     heralding signal, ie the laser drive pulse, and a channel detecting photons.     The correlation is calculated for the different photon numbers determined by     the thresholds parameter.       :param ndarray abs_time: absolute timestamp sequence.     :param ndarray detection_mask: boolean mask that identifies the entries in                                    abs_time belonging to the photon channel,                                    other entries are assumed to be the heralding                                    channel.     :param ndarray data: energy measurement data for the photon channel.     :param ndarray masks: boolean masks that are appied to                           abs_time[detection_mask] and indicate which detections                           are assigned which threshold.                           Shape(N, len(abs_time[detection_mask]))     :param r: the range of heralding channel delays to cross correlate.     :param bool verbose: Print progress message as each threshold is processed.     :return: list of ndarrays representing the cross correlation.      :note: abs_time[mask] and data must be the same length.            TODO speed up the algorithm. get offsets from the current index i in abs_time that are in the relative     time window defined by low and high.      :param int i: current abs_time index.     :param ndarray abs_time: absolute times.     :param low: low end of relative window.     :param high: high end of relative window.     :return: (low_o, high_o) coincident times are abs_time[low_o:i+high_o] # low index offset # high index offset # offset indexs marking abs_times in the window Find coincidences between two channels returning indices of partner events.      :param ndarray abs_time: absolute times.     :param ndarray mask: ndarray of bool that identifies the channels in                          abs_time. For abs_times where mask is False, time t is                          coincident if abs_time+low <= t <= abs_time+high. If                          the mask is True t is coincident if                          abs_time-high <= t <= abs_time-low.     :param float low: the low side of the coincidence window.     :param float high:  the high side of the coincidence whindw.     :return: (coinc, coinc_mask) where coinc is a ndarray of indices of the              coincident event in the other channel. When more than one event is              found in the window the negated value of the first index is              entered.              coinc_mask is a ndarray of bool indicating where exactly one              event was found in the window. :param ndarray data: measurement data.     :param MixtureModel measurement_model: model created using                                            expectation_maximisation.     :param bool vacuum: estimate vacuum terms using coinc_mask and herald_mask.     :param coinc_mask: ndarray of bool indicating coincidence.     :param herald_mask: ndarray of bool indicating which events are in the                         heralding channel.     :param has_noise_threshold: When True model.thresholds[1] represents the                                 boundary between noise and 1 photon.     :return: Counts(count, uncorrelated, vacuum). Where count is an ndarray of              counts for each photon number              and len(count)=len(measurement_model.thresholds). uncorrelated is a              ndarray the same shape as count and counts              events not correlated with the herald. uncorrelated is only              calculated when coinc_mask and herald_mask are supplied.              vacuum replicates the parameter value and indicates that count[0]              contains the vacuum count. # thresholds # last threshold overflow # counting vacuum # uncorrelated heralds + correlated data <= the first threshold # note these are correlated events less that the noise threshold # FIXME remove counts # def povm_elements(measurement_model, counts): #     """ #     estimate the elements of the system povm from the measurement model and the #     counting data. # #     :param array like data: measurement data. #     :param mixturemodel measurement_model: model created using #                                            expectation_maximisation. #     :param Counts counts: the count data returned by count(). #     :return: (elements count vacuum). where counts is a ndarray with #              shape (len(n)) containing the count for each measurement outcome #              in n. elements is a ndarray with #              shape(len(n), len(measurement_model.thresholds)) the first index is #              the fock state, the second is the outcome. The contents of elements #              is the probability of the outcome when measuring fock state. # #     :notes: POVM elements is indexed by (fock state, measurement outcome). #             The last threshold in the measurement model marks the boundary of #             the data that was used to create the model, it is *not* altered by #             the expectation maximisation algorithm and is essentially a guess. #             The second last threshold is the boundary of the overflow bin, #             this last photon number bin is counts detections of >= #             len(measurement_model.param_list)-2 photons. #             TODO expand and clarify description. #     """ # #     # t thresholds #     # t-1 distributions #     # t+1 counts if vacuum else t count[-1] is overflow past last threshold # # #     num_elements = len(counts.count) #     # cumulative density functions at the thresholds, the last value is not used #     # normalise to a convex mixture excluding the noise distribution. #     cdfs = measurement_model.cdf(measurement_model.thresholds) # #     elems = cdfs[:, 1:] - cdfs[:, :-1] #     # Adjust the first elements, which should not have a proceeding threshold. #     elems[:, 0] = cdfs[:, 1] # #     # Adjust the last elements to account for the truncation at the last #     # threshold (limit of data used to create the model). #     elems[:, -1] = 1 - cdfs[:, -2] # #     # elements should always have a vacuum term. #     elements = np.zeros((num_elements, num_elements)) # #     # the majority of the POVM elements are given by the difference in the CDFs #     # at the threshold values # #     if counts.vacuum: #         if measurement_model.has_noise_threshold: #             elements = elems #         else: #             elements[1:, 1:] = elems #     else: #         elements = elems # #     # vacuum terms #     # FIXME check this is correct when has_noise_threshold #     if counts.vacuum: #         # probability of counting vacuum when there is a vacuum input #         # TODO is it appropriate to do same for m > 1 ?? #         for m in range(1, 2):  # range(1, len(counts.uncorrelated)): #             elements[0, 0] = ( #                     (counts.count[0]-counts.uncorrelated[m])/counts.count[0] #             ) #             elements[0, m] = counts.uncorrelated[m]/counts.count[0] # #          # elements[m, 0] = correlatedsingle/single # #         if measurement_model.has_noise_threshold: #             # look for uncrrelated ??? #             pass # #     else: #         elements[0, 0] = 1 #         elements[0, 1] = 0 # #     # c = np.array(counts.count[:-1]) #     # c[-1] += counts.count[-1] # #     return Povm(elements, counts.vacuum, counts.has_noise_threshold) # FIXME check this model with new POVM remove call to outcomes # n [low, high] #         print('n in model',n) # print('state sum',np.sum(state)) # print('state.shape', state.shape) # print('num elements', num_elements, 'N', N) # print('nbar', nbar, 'alpha', alpha) # print('sum state', np.sum(state[n[0]:])) # print('sum t state', np.sum(truncated_state[n[0]:])) # print( #     'trunc test', #     np.real(np.sum(state[n[0]:]))-np.sum(truncated_state[n[0]:]) # ) #         print('outcomes',outcomes.shape) # if n[0] != 0: #     out = out / np.sum(out) #       print('out',out.shape) # alpha = 0.1 #     print('n',n) #     pars['povm_elements'].set(vary=False) # n [low, high] #/np.sum(truncated_state) # n [low, high] #/np.sum(truncated_state) # the width of the bars # print(fock.shape, fock) # chi2 = sum(np.power(c - model_counts, 2) / model_counts) # plt.plot(n,nc,'.') # plt.plot(n,r.best_fit,'o',markerfacecolor='none') # ax.set_ylim(0,11000000) # ax.text( #     tx, ty - ls, r'$\alpha={:.4f}\pm{:.4f}$' #         .format( #         state_fit.params['alpha'].value, state_fit.params['alpha'].stderr #     ), #     transform=ax.transAxes, horizontalalignment='left' # ) # ax.text( #     tx, ty - 2 * ls, r'${}={:.4f}\pm{:.4f}$' #         .format( #         r'\bar{n}', state_fit.params['nbar'].value, #         state_fit.params['nbar'].stderr #     ), #     transform=ax.transAxes, horizontalalignment='left' # ) # ax.text( #     tx, ty - 3 * ls, r'$\chi^2/{}={:.3f}$' #         .format( #         state_fit.nfree - 1, chi2 / (state_fit.nfree - 1) #     ), #     transform=ax.transAxes, horizontalalignment='left' # ) Model the measurement outcome probabilities assuming the input is a     displaced thermal state.      :param x iterable: The state parameters     :param N int: Dimension of the Fock space, max photon number is N-1     :param thermal: When True model input as displaced thermal state otherwise                     model as coherent state (displaced thermal).     :param povm: The detector povm     :param qtN: The dimension of the Fock space to use for qutip models. #     print(x) #         print('thermal') #     print(np.sum(state)) # number of povm elements # # truncate to number of elements # truncate to N if necessary #     ax.set_xlim([200,800]) # the width of the bars #     print(len(counts)) #     print(n[0],len(n),len(c)) #     chi2=sum(np.power(counts-r.best_fit,2)/model_counts) # plt.plot(n,nc,'.') # plt.plot(n,r.best_fit,'o',markerfacecolor='none') # ax.set_ylim(0,11000000)