kapture/heb/__init__.py

import logging
import threading
import numpy as np
import time


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.info("GPU fft: {} s".format(time.time() - start))
        return out_dev.get()[:, :width / 2 + 1]

except ImportError:
    logging.info("Failed to import reikna package. Falling back to Numpy FFT.")
    def _fft(array):
        start = time.time()
        freqs = np.fft.rfft(_pad_array(array))
        logging.info("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):
        self.filename = filename
        self.array = array
        self._heatmaps = {}
        self._ffts = {}

    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):
        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)
        return _fft(heatmap)