kapture/heb/plot.py

import numpy as np
import matplotlib.pyplot as plt
import scipy.constants


try:
    from numpy.fft import rfftfreq
except ImportError:
    # This is a backport of the rfftfreq function present as of 2012 but not
    # included in Ubuntu 12.04.
    #
    # Licensed under BSD-new.
    def rfftfreq(n, d=1.0):
        """
        Return the Discrete Fourier Transform sample frequencies
        (for usage with rfft, irfft).

        The returned float array `f` contains the frequency bin centers in cycles
        per unit of the sample spacing (with zero at the start).  For instance, if
        the sample spacing is in seconds, then the frequency unit is cycles/second.

        Given a window length `n` and a sample spacing `d`::

          f = [0, 1, ...,     n/2-1,     n/2] / (d*n)   if n is even
          f = [0, 1, ..., (n-1)/2-1, (n-1)/2] / (d*n)   if n is odd

        Unlike `fftfreq` (but like `scipy.fftpack.rfftfreq`)
        the Nyquist frequency component is considered to be positive.

        Parameters
        ----------
        n : int
            Window length.
        d : scalar, optional
            Sample spacing (inverse of the sampling rate). Defaults to 1.

        Returns
        -------
        f : ndarray
            Array of length ``n//2 + 1`` containing the sample frequencies.

        Examples
        --------
        >>> signal = np.array([-2, 8, 6, 4, 1, 0, 3, 5, -3, 4], dtype=float)
        >>> fourier = np.fft.rfft(signal)
        >>> n = signal.size
        >>> sample_rate = 100
        >>> freq = np.fft.fftfreq(n, d=1./sample_rate)
        >>> freq
        array([  0.,  10.,  20.,  30.,  40., -50., -40., -30., -20., -10.])
        >>> freq = np.fft.rfftfreq(n, d=1./sample_rate)
        >>> freq
        array([  0.,  10.,  20.,  30.,  40.,  50.])

        """
        val = 1.0/(n*d)
        N = n//2 + 1
        results = np.arange(0, N, dtype=int)
        return results * val


H = 184
FREV = scipy.constants.c / 110.4    # lightspeed divided by circumference
TREV = 1 / FREV


# This is a backport of the polyval function present as of 2011 but not
# included in Ubuntu 12.04.
#
# Licensed under BSD-new.
def polyval(x, c, tensor=True):
    """
    Evaluate a polynomial at points x.

    If `c` is of length `n + 1`, this function returns the value

    .. math:: p(x) = c_0 + c_1 * x + ... + c_n * x^n

    The parameter `x` is converted to an array only if it is a tuple or a
    list, otherwise it is treated as a scalar. In either case, either `x`
    or its elements must support multiplication and addition both with
    themselves and with the elements of `c`.

    If `c` is a 1-D array, then `p(x)` will have the same shape as `x`.  If
    `c` is multidimensional, then the shape of the result depends on the
    value of `tensor`. If `tensor` is true the shape will be c.shape[1:] +
    x.shape. If `tensor` is false the shape will be c.shape[1:]. Note that
    scalars have shape (,).

    Trailing zeros in the coefficients will be used in the evaluation, so
    they should be avoided if efficiency is a concern.

    Parameters
    ----------
    x : array_like, compatible object
        If `x` is a list or tuple, it is converted to an ndarray, otherwise
        it is left unchanged and treated as a scalar. In either case, `x`
        or its elements must support addition and multiplication with
        with themselves and with the elements of `c`.
    c : array_like
        Array of coefficients ordered so that the coefficients for terms of
        degree n are contained in c[n]. If `c` is multidimensional the
        remaining indices enumerate multiple polynomials. In the two
        dimensional case the coefficients may be thought of as stored in
        the columns of `c`.
    tensor : boolean, optional
        If True, the shape of the coefficient array is extended with ones
        on the right, one for each dimension of `x`. Scalars have dimension 0
        for this action. The result is that every column of coefficients in
        `c` is evaluated for every element of `x`. If False, `x` is broadcast
        over the columns of `c` for the evaluation.  This keyword is useful
        when `c` is multidimensional. The default value is True.

        .. versionadded:: 1.7.0

    Returns
    -------
    values : ndarray, compatible object
        The shape of the returned array is described above.

    See Also
    --------
    polyval2d, polygrid2d, polyval3d, polygrid3d

    Notes
    -----
    The evaluation uses Horner's method.

    Examples
    --------
    >>> from numpy.polynomial.polynomial import polyval
    >>> polyval(1, [1,2,3])
    6.0
    >>> a = np.arange(4).reshape(2,2)
    >>> a
    array([[0, 1],
           [2, 3]])
    >>> polyval(a, [1,2,3])
    array([[  1.,   6.],
           [ 17.,  34.]])
    >>> coef = np.arange(4).reshape(2,2) # multidimensional coefficients
    >>> coef
    array([[0, 1],
           [2, 3]])
    >>> polyval([1,2], coef, tensor=True)
    array([[ 2.,  4.],
           [ 4.,  7.]])
    >>> polyval([1,2], coef, tensor=False)
    array([ 2.,  7.])

    """
    c = np.array(c, ndmin=1, copy=0)
    if c.dtype.char in '?bBhHiIlLqQpP':
        # astype fails with NA
        c = c + 0.0
    if isinstance(x, (tuple, list)):
        x = np.asarray(x)
    if isinstance(x, np.ndarray) and tensor:
        c = c.reshape(c.shape + (1,)*x.ndim)

    c0 = c[-1] + x*0
    for i in range(2, len(c) + 1) :
        c0 = c[-i] + c0*x
    return c0


def bunch(data, number, adc=None, frm=0, to=-1, **kwargs):
    bdata = data.bunch(number)[frm:to,1:]
    xs = np.arange(frm, max(to, bdata.shape[0]))

    if adc is None:
        plt.plot(xs, np.mean(bdata, axis=1),
                 label='Averaged ADC count, bunch {}'.format(number),
                 **kwargs)
    else:
        plt.plot(xs, bdata[:,adc-1],
                 label='ADC {}, bunch {}'.format(adc, number),
                 **kwargs)


def train(data, axis, adc=None, frm=0, to=-1, **kwargs):
    pdata = data.array[frm:to, :-1]

    if adc is None:
        axis.plot(np.mean(pdata, axis=1), **kwargs)
    else:
        axis.plot(pdata[:,adc-1], **kwargs)

    axis.set_xlabel('Sample point')

    if adc:
        axis.set_title('ADC {}'.format(adc))


def combined(data, axis, adcs, frm=0, to=-1, show_reconstructed=True):
    array = data.array[frm:to, :]

    N_s = 16
    N_p = array.shape[0]

    # Remove bunch number and flatten array
    array = array[:,:4].reshape((-1, 1))

    # plot original data
    orig_xs = np.arange(0, array.shape[0])
    axis.plot(orig_xs, array, '.')

    # 4 rows for each ADC
    ys = array.reshape((4, -1), order='F')

    # multi-fit for each column of 4 ADCs, unfortunately, xs' are always the
    # same, so we have to offset them later when computing the value
    xs = [1, 2, 3, 4]
    c = np.polynomial.polynomial.polyfit(xs, ys, 6)

    smooth_xs = np.linspace(1, 4, N_s)
    fitted_ys = polyval(smooth_xs, c, tensor=True)

    if True:
        # Output maxima
        ys = np.max(fitted_ys, axis=1)
        xs = 4 * ((np.argmax(fitted_ys, axis=1) + 1) / float(N_s)) + orig_xs[::4] - .5

        axis.plot(xs, ys, '.', color="red")

    if True:
        # Debug output
        ys = fitted_ys.reshape((-1, 1))
        xs = np.repeat(np.arange(0, 4 * N_p, 4), N_s) + np.tile(np.linspace(0, 3, N_s), N_p)
        axis.plot(xs, ys, '.', alpha=0.3)


def heatmap(heatmap, axis):
    image = axis.imshow(heatmap, interpolation='nearest', aspect='auto')
    axis.invert_yaxis()
    axis.set_xlabel('Turn')
    axis.set_ylabel('Bunch position')
    return image


def fft(data, axis, adc, frm=0, to=-1, skipped=0):
    transform = data.fft(adc, frm=frm, to=to)
    freqs = rfftfreq(transform.shape[1], TREV * (skipped + 1))

    transform = transform[:, 1:]  # Throw away column 1
    # Due to the nature of the signal being divided into 'buckets', frequency '1' has an excessively large value and
    # skews the proportions for the other signals. It does not hold any useful information either...
    # Therefore, we remove it from the results.

    magnitude = np.log(np.abs(transform))
    image = axis.imshow(magnitude,
                        interpolation='none',
                        aspect='auto',
                        extent=[freqs[1], freqs[-1], 0, H],  # Start from freq[1] since we removed freq[0] from the data
                        origin='lower')

    axis.set_xlabel('Frequency')
    axis.set_ylabel('Bunch position')
    return image