KaptureControlGui/KCG/base/backend/dataset.py

import logging
import time

import numpy as np
from numpy.polynomial.polynomial import polyval

from ... import config


def _pad_array(array):
    height, width = array.shape

    # Miriam uses floor hence the padding is actually a cutting. Will wait for
    # response if this is desired ...
    pwidth = 2**np.floor(np.log2(width))
    padded = np.zeros((height, pwidth))
    padded[:, :width] = array[:, :pwidth]
    return padded

try:
    import reikna.cluda
    import reikna.fft

    _plans = {}
    _in_buffers = {}
    _out_buffers = {}

    _api = reikna.cluda.ocl_api()
    _thr = _api.Thread.create()

    def _fft(array):
        start = time.time()
        padded = _pad_array(array).astype(np.complex64)
        height, width = padded.shape

        if width in _plans:
            fft = _plans[width]
            in_dev = _in_buffers[width]
            out_dev = _out_buffers[width]
        else:
            fft = reikna.fft.FFT(padded, axes=(1,)).compile(_thr)
            in_dev = _thr.to_device(padded)
            out_dev = _thr.empty_like(in_dev)
            _plans[width] = fft
            _in_buffers[width] = in_dev
            _out_buffers[width] = out_dev

        fft(out_dev, in_dev)
        logging.debug("GPU fft: {} s".format(time.time() - start))
        return out_dev.get()[:, :width / 2 + 1]

except ImportError:
    logging.debug("Failed to import reikna package. Falling back to Numpy FFT.")
    def _fft(array):
        start = time.time()
        freqs = np.fft.rfft(_pad_array(array))
        logging.debug("np fft: {} s".format(time.time() - start))
        return freqs


BUNCHES_PER_TURN = 184
HEADER_SIZE_BYTES = 32


class DataSet(object):
    def __init__(self, array, filename, header=None):
        self.filename = filename
        self.array = array
        self._heatmaps = {}
        self._ffts = {}
        self.header = header

    @property
    def skipped_turns(self):
        if self.header:
            return self.header['skipped_turns']
        else:
            return 1

    def bunch(self, number):
        return self.array[self.array[:, 4] == number]

    def num_bunches(self):
        return self.array.shape[0]

    def num_turns(self):
        return self.num_bunches() / BUNCHES_PER_TURN

    def heatmap(self, adc=1, frm=0, to=-1, bunch_frm=0, bunch_to=-1):
        if not 1 <= adc <= 4:
            raise ValueError('adc must be in [1,4]')

        if not adc in self._heatmaps:
            heatmap = self.array[:,adc-1].reshape(-1, BUNCHES_PER_TURN).transpose()
            self._heatmaps[adc] = heatmap

        return self._heatmaps[adc][bunch_frm:bunch_to, frm:to]

    def fft(self, adc=1, frm=0, to=-1, drop_first_bin=False):
        if not 1 <= adc <= 4:
            raise ValueError('adc must be in [1,4]')

        # if not adc in self._ffts:
        #     heatmap = self.heatmap(adc, frm, to)
        #     self._ffts[adc] = np.fft.fft2(heatmap, axes=[1])

        # return self._ffts[adc]
        heatmap = self.heatmap(adc, frm, to)
        if drop_first_bin:
            return _fft(heatmap)[:, 1:]
        else:
            return _fft(heatmap)

    def fft_max_freq(self):
        return (self.num_turns() // 2 + 1) * self.fft_freq_dist()
    def fft_freq_dist(self):
        return 1.0/(self.num_turns() * (self.skipped_turns + 1) * config.tRev)

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

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

        N_s = 16
        N_p = array.shape[0]

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

        # plot original data
        orig_xs = np.arange(0, array.shape[0])
        # axis.plot(orig_xs, array, '.')
        ret = [np.array([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")
            ret.append(np.array([xs, ys]))

        if False:
            # Debug output
            ys = fitted_ys.reshape((-1, 1))
            ys = ys.flatten()
            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)
            ret.append(np.array([xs, ys]))
        return ret