xboa
_fft_tune.py
Go to the documentation of this file.
1 # This file is a part of xboa
2 #
3 # xboa is free software: you can redistribute it and/or modify
4 # it under the terms of the GNU General Public License as published by
5 # the Free Software Foundation, either version 3 of the License, or
6 # (at your option) any later version.
7 #
8 # xboa is distributed in the hope that it will be useful,
9 # but WITHOUT ANY WARRANTY; without even the implied warranty of
10 # MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
11 # GNU General Public License for more details.
12 #
13 # You should have received a copy of the GNU General Public License
14 # along with xboa in the doc folder. If not, see
15 # <http://www.gnu.org/licenses/>.
16 
17 import math
18 import sys
19 try:
20  import numpy
21 except ImportError:
22  pass
23 try:
24  import ROOT
25 except ImportError:
26  pass
27 
28 import xboa.Common as common
29 from xboa.algorithms.peak_finder import WindowPeakFinder
30 
31 class FFTTuneFinder(object):
32  """
33  Find the tune using the fast fourier transform technique
34 
35  Apply a small displacement to the reference trajectory and track through a
36  lattice; use the displaced trajectory
37  """
38 
39  def __init__(self, u_data = None, peak_finder = None,
40  use_hanning_filter = False):
41  """
42  Initialise the tune finder
43 
44  - u_data: position data, for use in Fourier Transform calculation
45  - peak_finder: peak finder object, object used to find peaks in the tune
46  diagram. If peak_finder is None, uses a WindowPeakFinder with a
47  window size 10
48  - use_hanning_filer: experimental hanning filter (probably doesnt work)
49  """
50  self._peak_finder = peak_finder
51  if self._peak_finder == None:
52  self._peak_finder = WindowPeakFinder(5, 0., 1)
53  self.u = u_data
54  self.use_hanning_filter = use_hanning_filter
55  self.fractional_tune = None
56  self.peak_x_list = []
57  self.peak_y_list = []
58  self.k_mag_x = None
59  self.k_mag_y = None
60 
61  def get_tune(self, tune_tolerance = None):
62  """
63  Find the fractional tune for a given tracking object
64 
65  - tune_tolerance: tolerance with which the tune finder will attempt to
66  calculate the tune. If set to None, FTTuneFinder will use
67  1/len(u)/100
68 
69  Displaces the reference hit by an amount delta, tracks this hit using
70  the tracking object and runs a one dimensional FFT against the axis
71  variable. Number of samples used in the FFT is determined by the
72  tracking object (and hence accuracy of the tune calculation).
73 
74  Note that u_data must be of odd length for this algorithm; FFTTuneFinder
75  will discard the last element if this is not the case.
76 
77  Returns the principle Fourier frequency that is not 0.
78  """
79  if self.u == None:
80  raise ValueError("No data set for Fourier Transform")
81  if len(self.u) % 2 == 0:
82  del self.u[-1]
83  if tune_tolerance == None:
84  tune_tolerance = 1./len(self.u)/100.
85  self._fft_find_peaks()
86  self._get_max_peak()
87  if len(self.peak_x_list) == 0:
88  self._sft_find_peaks(tune_tolerance)
89  self._get_max_peak()
90  if tune_tolerance < 1./len(self.k_mag_x):
91  for i, k_x in enumerate(self.peak_x_list):
92  try:
93  new_peak_x = self._recursive_refine_peak(k_x, tune_tolerance)
94  self.peak_x_list[i] = new_peak_x
95  index = self.k_mag_x.index(new_peak_x)
96  self.peak_y_list[i] = self.k_mag_y[index]
97  except ValueError:
98  sys.excepthook(*sys.exc_info())
99  self._get_max_peak()
100  return self.fractional_tune
101 
102  def run_tracking(self, axis, delta, reference_hit, tracking, use_hits=None):
103  """
104  Set position data from tracking output
105 
106  - reference_hit: Hit on the closed orbit (for a ring).
107  - axis: string, variable from Hit.get_variables() over which the tune
108  will be found.
109  - delta: float, displacement from the reference_hit.
110  - tracking: xboa.tracking.TrackingBase, tracking object to propagate
111  particles through the lattice.
112  - use_hits: list of integers, only consider hits returned by tracking
113  with an index in use_hits. If set to None, consider all hits.
114  """
115  hit = reference_hit.deepcopy()
116  hit[axis] += delta
117  hits_out = tracking.track_one(hit)
118  if use_hits != None:
119  hits_out = [hit for i, hit in enumerate(hits_out) if i in use_hits]
120  self.u = [hit[axis] for hit in hits_out]
121 
122 
123  def plot_signal(self, title):
124  """
125  Plot the FFT frequency spectrum
126  - title, string used as a title in FFT plots
127  Returns tuple of TCanvas, TH2, TGraph
128  """
129  canvas = common.make_root_canvas(title)
130  x_list = range(len(self.u))
131  hist, graph = common.make_root_graph(title,
132  x_list, 'count', self.u, 'signal')
133  hist.SetTitle(title)
134  hist.Draw()
135  graph.Draw('l')
136  canvas.Update()
137  return canvas, hist, graph
138 
139  def plot_fft(self, title, ymin=None, ymax=None):
140  """
141  Plot the FFT frequency spectrum
142  - title, string used as a title in FFT plots
143  Returns tuple of TCanvas, TH2, TGraph
144  """
145  if self.k_mag_x == None or self.k_mag_y == None:
146  self.get_tune()
147  canvas = common.make_root_canvas(title)
148  if ymin == None:
149  ymin = max([1e-5, min(self.k_mag_y)])
150  hist, graph = common.make_root_graph(title,
151  self.k_mag_x, 'frequency',
152  self.k_mag_y, 'magnitude',
153  ymin=ymin, ymax=ymax)
154  hist.SetTitle(title)
155  hist.Draw()
156  graph.Draw('l')
157  canvas.SetLogy()
158  canvas.Update()
159  return canvas, hist, graph
160 
161  def plot_fft_fit(self, title, peak_k_index,
162  fit_lower_index, fit_upper_index, fit = None):
163  """
164  Draw the FFT spectrum and fit within a window
165  - title, string title for the plot
166  - peak_k_index, integer corresponding to element from self.k_mag which
167  makes the seed for the position of the hit
168  - fit_lower_bound, integer corresponding to the element from self.k_mag
169  which makes the lower edge of the fit window
170  - fit_lower_bound, integer corresponding to the element from self.k_mag
171  which makes the upper edge of the fit window
172  - fit, ROOT.TF1 used for fitting. If set to None, a Gaussian will be
173  used with reasonable parameters. Note that user has responsibility to
174  set the fit window in fit
175 
176  Uses the ROOT library to do the fit; this means drawing the data onto a
177  ROOT Canvas using self.plot_fft. Fits with a Gaussian; in the presence
178  of noisy/low statistics data, this can improve the estimate of tune.
179 
180  Returns a tuple of TCanvas, TH2, TGraph, TF1. Hint:- to get the
181  fractional tune, use TF1::GetParameter(0); to get the estimated error,
182  use TF1::GetParError(0).
183  """
184  canvas, hist, graph = self.plot_fft(title)
185  fit_height = self.k_mag_y[peak_k_index]
186  fit_centre = self.k_mag_x[peak_k_index]
187  fit_low = self.k_mag_x[fit_lower_index]
188  fit_hi = self.k_mag_x[fit_upper_index]
189  if fit == None:
190  fit_function = "[1]*exp(-((x-[0])*[2])**2)"
191  fit = ROOT.TF1(title+" fit", fit_function, fit_low, fit_hi)
192  fit.SetParameter(0, fit_centre)
193  fit.SetParameter(1, fit_height)
194  fit.SetParameter(2, fit_height)
195  graph.Fit(fit)
196  canvas.cd()
197  fit.Draw("SAME")
198  print "Got peak at", fit.GetParameter(0), "with error", fit.GetParError(0)
199  canvas.Update()
200  return canvas, hist, graph, fit
201 
202  def _get_max_peak(self):
203  """
204  Find the maximum peak in self.peak_index_list
205  """
206  peak_index_list = [self.k_mag_x.index(k_x) for k_x in self.peak_x_list]
207  peak_values = [self.k_mag_y[i] for i in peak_index_list]
208  if len(peak_values) == 0:
209  raise ValueError("Can't get max peak - found no peaks at all")
210  peak_max = max(peak_values)
211  peak_index = self.k_mag_y.index(peak_max)
212  self.fractional_tune = self.k_mag_x[peak_index]
213  return peak_index
214 
215  def _sft_find_peaks(self, interval):
216  """
217  Perform a slow Fourier Transform and find any peaks
218  """
219  n_items = int(1./interval/2.)
220  self.k_mag_x = [i*interval for i in range(n_items)]
221  self.k_mag_y = [self._sft(k_x) for k_x in self.k_mag_x]
222  self._find_peaks()
223 
224  def _fft_find_peaks(self):
225  """
226  Perform the Fast Fourier Transform and find any peaks
227  """
228  fft = numpy.fft.rfft(numpy.array(self.u))
229  k_list = [[float(numpy.real(z)), float(numpy.imag(z))] for z in fft]
230  self.k_mag_y = [(z[0]**2+z[1]**2)**0.5 for z in k_list]
231  self.k_mag_x = [i/2./float(len(k_list)) for i, z in enumerate(k_list)]
232  self._find_peaks()
233 
234  def _find_peaks(self):
235  """
236  Run the peak finder
237  """
238  peak_index_list = sorted(self._peak_finder.find_peaks(self.k_mag_y))
239  if len(peak_index_list) == 0:
240  return
241  if peak_index_list[0] == 0:
242  peak_index_list.pop(0)
243  self.peak_x_list = [self.k_mag_x[i] for i in peak_index_list]
244  self.peak_y_list = [self.k_mag_y[i] for i in peak_index_list]
245 
246  def _recursive_refine_peak(self, k_x, x_tolerance):
247  """
248  Use slow fourier transforms in the region of a peak to recursively
249  improve the peak estimate to some tolerance.
250 
251  On each iteration, the new k_mag values are appended to k_mag_x/y.
252 
253  Stop recursing if the new value is < the neighbouring value; we assume
254  this is numerical precision noise and the recursion has converged
255  (implicit that the seed peak was closer than any troughs).
256 
257  Returns the new peak_index
258  """
259  peak_index = self.k_mag_x.index(k_x)
260  if peak_index+1 == len(self.k_mag_y):
261  peak_index -= 1
262  y_values = [self.k_mag_y[peak_index-1],
263  self.k_mag_y[peak_index],
264  self.k_mag_y[peak_index+1]]
265  if y_values[2] > y_values[0]:
266  del y_values[0]
267  x_values = [self.k_mag_x[peak_index],
268  self.k_mag_x[peak_index+1]]
269  index_right = peak_index+1
270  else:
271  del y_values[2]
272  x_values = [self.k_mag_x[peak_index-1],
273  self.k_mag_x[peak_index]]
274  index_right = peak_index
275  new_x = (x_values[0] + x_values[1])/2.
276  new_y = self._sft(new_x)
277  if abs(x_values[1] - x_values[0]) < x_tolerance:
278  return self.k_mag_x[peak_index]
279  self.k_mag_x.insert(index_right, new_x)
280  self.k_mag_y.insert(index_right, new_y)
281  # new index for the peak value;
282  # if new_y is highest, peak index is the new_y index
283  if new_y > y_values[0] and new_y > y_values[1]:
284  new_peak_index = index_right
285  # if new_y is lowest, the peak has split and we split the recursion
286  # (this is a bit dodgy)
287  elif y_values[0] > new_y and y_values[1] > new_y:
288  new_peak_index = self.k_mag_y.index(y_values[0])
289  self._recursive_refine_peak(self.k_mag_x[new_peak_index], x_tolerance)
290  new_peak_index = self.k_mag_y.index(y_values[1])
291  # if y_values[0] is highest, peak index is unchanged
292  elif y_values[0] > y_values[1]:
293  new_peak_index = peak_index
294  # else index of y_values[1], which is now 1 higher thanks to insert
295  else:
296  new_peak_index = peak_index + 1
297  return self._recursive_refine_peak(self.k_mag_x[new_peak_index], x_tolerance)
298 
299  def _sft(self, k_x):
300  """
301  Calculate "Slow" Fourier Transform at k_x
302 
303  - k_x, float, position at which the sft is found
304  - hanning filter, bool, set to True to apply a hanning filter
305 
306  "Slow" Fourier transform means using Sum(A_i cos(...) + A_i sin(...)) to
307  get the FT at a given k value
308  """
309  y_point, z_point = 0., 0.
310  n = float(len(self.u))
311  hanning = 1.
312  for m, a_m in enumerate(self.u):
313  if self.use_hanning_filter:
314  hanning = 2.*math.sin(math.pi*m/n)**2.
315  f = 2.*math.pi*m*k_x*(n+1)/n
316  dy = hanning*a_m*math.cos(f)
317  dz = hanning*a_m*math.sin(f)
318  y_point += dy
319  z_point += dz
320  k_y = (y_point**2+z_point**2)**0.5
321  return k_y
322 
323 
def _fft_find_peaks
Perform the Fast Fourier Transform and find any peaks.
Definition: _fft_tune.py:237
def plot_fft
Plot the FFT frequency spectrum.
Definition: _fft_tune.py:150
def _sft_find_peaks
Perform a slow Fourier Transform and find any peaks.
Definition: _fft_tune.py:227
def _get_max_peak
Find the maximum peak in self.peak_index_list.
Definition: _fft_tune.py:213
def _recursive_refine_peak
Use slow fourier transforms in the region of a peak to recursively improve the peak estimate to some ...
Definition: _fft_tune.py:270
def run_tracking
Set position data from tracking output.
Definition: _fft_tune.py:118
Find the tune using the fast fourier transform technique.
Definition: _fft_tune.py:38
def plot_fft_fit
Draw the FFT spectrum and fit within a window.
Definition: _fft_tune.py:190
def _find_peaks
Run the peak finder.
Definition: _fft_tune.py:248
def __init__
Initialise the tune finder.
Definition: _fft_tune.py:51
def _sft
Calculate "Slow" Fourier Transform at k_x.
Definition: _fft_tune.py:321
Deprecated accessor to xboa.common module.
Definition: Common.py:1
def get_tune
Find the fractional tune for a given tracking object.
Definition: _fft_tune.py:81
Algorithms to find peaks in potentially noisy data.
Definition: __init__.py:1
def plot_signal
Plot the FFT frequency spectrum.
Definition: _fft_tune.py:133