123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165 |
- 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
|