xboa
_ellipse_closed_orbit_finder.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 """
18 Closed orbit finder
19 """
20 
21 import math
22 import json
23 import copy
24 try:
25  import numpy
26 except ImportError:
27  pass
28 import xboa.common as common
29 
30 class EllipseClosedOrbitFinder(object):
31  """
32  Find the closed orbit given an ellipse
33  """
34  def __init__(self, tracking, seed_hit, eps_max = 1e6, use_seed = False):
35  """
36  Initialise the closed orbit finder
37  - tracking: object of type TrackingBase used for tracking particles.
38  tracking should return a list of hits all of which sit at the cell
39  end
40  - seed_hit: hit object used as a seed for finding the closed orbit. For
41  tracking, takes e.g. pid from this variable. Uses the dynamical
42  variables as seeds for closed orbit finding.
43  - eps_max: ignore particles with amplitude greater than eps_max.
44  - use_seed: set to True to use the seed as the first point in the
45  ellipse when finding ellipses
46 
47  Requires numpy
48  """
49  common.has_numpy()
50  self.tracking = tracking
51  self.seed = seed_hit
52  self.eps_max = eps_max
53  self.use_seed = use_seed
54 
55  def scan_generator(self, step_start_dict, step_end_dict, step_size_dict):
56  """
57  Make a generator to scan over the variables, looking for closed orbits.
58  - step_start_dict dict defining the start position for the scan
59  - step_end_dict dict defining the end position for the variable scan
60  - step_size_dict dict of step sizes for the variable scan
61  Each dictionary should have the same set of keys. The keys should be
62  hit set_variable keys.
63 
64  Returns an iterable of EllipseClosedOrbitFinderIterations, one iteration
65  for each point in the scan; e.g. call next() against the return value to
66  get the next element in the list, or put it in a loop like
67  my_scan = [x for x in yield_scan(start, end, step)]
68  """
69  next_hit = self.seed.deepcopy()
70  step = copy.deepcopy(step_start_dict)
71  keys = sorted(step.keys())
72  while True:
73  for key in keys:
74  if step[key] >= step_end_dict[key]:
75  raise StopIteration("Run out of steps at "+json.dumps(step))
76  next_hit[key] = step[key]
77  hit_list = self.tracking.track_one(next_hit)
78  x_list = [[hit[key] for key in keys] for hit in hit_list]
80  x_list,
81  self.eps_max)
82  i = len(keys)-1
83  key = keys[i]
84  step[key] += step_size_dict[key]
85  while step[key] > step_end_dict[key] and i > 0:
86  step[key] = step_start_dict[key]
87  i -= 1
88  key = keys[i]
89  step[key] += step_size_dict[key]
90 
91  def find_closed_orbit_generator(self, closed_orbit_variable_list,
92  number_of_points, max_iterations=100):
93  """
94  Make a generator that finds the closed orbit
95  - closed orbit variable_list: list of string variable names over which
96  the finder will run to find the closed orbit
97  - number_of_points; require a certain number of points for the ellipse
98  fit. If a call to tracking does not generate the required number
99  of points, take the last output from tracking and throw it back
100  through. If this does not return any new hits (i.e. the particle
101  has fallen out of the accelerator acceptance) then raise a Runtime
102  Error
103  - max_iterations: maximum number of iterations to make. If set to None,
104  keep iterating until we find a closed orbit. This can go forever.
105 
106  Returns a generator object that will continue iterating until the
107  tracking noise is large. If the closed orbit finder fails to fit an
108  ellipse (for example because the test particle is going out of bounds or
109  is singular) then find_closed_orbit raises a RuntimeError.
110 
111  The caller is invited to break out of the loop early if, for example,
112  she does not require convergence to the precision of the tracking.
113 
114  Note that this can make a warning from numpy like
115  "Warning: invalid value encountered in double_scalars"
116  """
117  keys = closed_orbit_variable_list
118  next_hit = self.seed.deepcopy()
119  noise_mean = None
120  iteration = 0
121  while iteration < max_iterations or max_iterations == None:
122  iteration += 1
123  x_list = self._get_points(keys, next_hit, number_of_points)
124  try:
125  my_co = EllipseClosedOrbitFinderIteration(keys, x_list,
126  self.eps_max, True)
127  except numpy.linalg.linalg.LinAlgError:
128  if len(x_list) > number_of_points:
129  raise StopIteration("Closed orbit has reached numerical "+\
130  "precision limit. Nice!")
131  else:
132  raise ValueError("Failed to find ellipse determinant.")
133 
134  yield my_co
135  if my_co.get_mean_noise() < my_co.get_sigma_noise():
136  raise StopIteration("Closed orbit finder has reached "+\
137  "convergence limit. Tracking noise is "+\
138  "greater than ellipse size.")
139  for i, var in enumerate(keys):
140  next_hit[var] = my_co.centre[i]
141  raise StopIteration("Closed orbit finder has finished iterating without converging.")
142 
143 
144  def check_closed_orbit(self, test_variable_dict, number_of_points):
145  """
146  Check whether test_variables fall on the closed orbit
147  - test_variable_dict: dictionary mapping variable names to variable
148  values. Typically variables should be dynamical variables of the
149  hit. Other variables will be filled from seed.
150  - number_of_points; require a certain number of points for the ellipse
151  fit. If a call to tracking does not generate the required number
152  of points, take the last output from tracking and throw it back
153  through. If this does not return any new hits (i.e. the particle
154  has fallen out of the accelerator acceptance) then raise a Runtime
155  Error
156 
157  Return is a EllipseClosedOrbitFinderIteration
158  """
159  next_hit = self.seed.deepcopy()
160  keys = sorted(test_variable_dict.keys())
161  for var in keys:
162  next_hit[var] = test_variable_dict[var]
163  x_list = self._get_points(keys, next_hit, number_of_points)
164  co_check = EllipseClosedOrbitFinderIteration(keys, x_list, self.eps_max)
165  try:
166  co_check.calculate_ellipse()
167  except numpy.linalg.linalg.LinAlgError:
168  pass
169  return co_check
170 
171  def _get_points(self, keys, next_hit, number_of_points):
172  hit_list = []
173  if self.use_seed:
174  hit_list.append(next_hit)
175  while len(hit_list) < number_of_points:
176  hit_list += self.tracking.track_one(next_hit)[1:]
177  if hit_list[-1] == next_hit:
178  raise RuntimeError("Tracking failed to generate an ellipse")
179  next_hit = hit_list[-1]
180  x_list = [[hit[var] for var in keys] for hit in hit_list]
181  return x_list
182 
184  """
185  Data following a single pass of the closed orbit finder
186 
187  Handle a single iteration of the CO finder. Few utility functions to find
188  the ellipse, look for convergence and return outputs to the user.
189  """
190  def __init__(self, keys, points, eps_max, calculate_ellipse = False):
191  """
192  Initialisation
193  - keys keys that tell us what the variables mean
194  - points list of points. Each point is a list that represents a vector
195  in the space defined by keys
196  - eps_max when fitting the ellipse, ignore points with amplitude greater
197  than eps_max
198  - calculate_ellipse set to true to calculate the ellipse. Otherwise, the
199  ellipse can be calculated by a call to calculate_ellipse()
200 
201  Requires numpy
202  """
203  common.has_numpy()
204  self.keys = keys
205  self.points = points
206  self.eps_max = eps_max
207  self.centre = None
208  self.ellipse = None
209  self.noise = None
210  if calculate_ellipse == True:
211  self.calculate_ellipse()
213  def calculate_ellipse(self):
214  """
215  Calculate the beam ellipse and noise list based on points
217  Raises a ValueError if noise cannot be calculated.
218  """
219  self.centre, self.ellipse = common.fit_ellipse(self.points,
220  self.eps_max,
221  verbose = False)
222  noise_list = []
223  ellipse = copy.deepcopy(self.ellipse)
224  ellipse /= numpy.linalg.det(ellipse)**(1./len(self.centre))
225  ellipse_inv = numpy.linalg.inv(ellipse)
226  for x_vec in self.points:
227  x_numpy = numpy.matrix(x_vec)-self.centre
228  noise = x_numpy * ellipse_inv * x_numpy.transpose()
229  noise_list.append(float(noise[0]))
230  for value in noise_list:
231  if math.isnan(value) or math.isinf(value):
232  raise ValueError("Failed to calculate noise; "+\
233  "ellipse may not be well-conditioned")
234  self.noise = noise_list
235 
236  def get_mean_noise(self):
237  """
238  Return the mean of the noise list.
239 
240  If noise is None, will attempt to calculate noise by calling
241  calculate_ellipse. Raises a ValueError if noise cannot be calculated.
242  """
243  if self.noise == None:
244  self.calculate_ellipse()
245  if len(self.noise) == 0:
246  raise ValueError("Failed to calculate noise - ellipse singular?")
247  n_noise = float(len(self.noise))
248  mean_noise = sum(self.noise)/n_noise
249  return mean_noise
250 
251 
252  def get_sigma_noise(self):
253  """
254  Return the standard deviation of the noise list
255 
256  If noise is None, will attempt to calculate noise by calling
257  calculate_ellipse. Raises a ValueError if noise cannot be calculated.
258  """
259  if self.noise == None:
260  self.calculate_ellipse()
261  if len(self.noise) == 0:
262  raise ValueError("Failed to calculate noise - ellipse singular?")
263  mean_noise = self.get_mean_noise()
264  noise_squ = [noise*noise for noise in self.noise]
265  var_noise = sum(noise_squ)/len(self.noise)-mean_noise**2
266  if var_noise <= 0.:
267  return 0.
268  return var_noise**0.5
270  def plot_ellipse(self, x_axis_string, y_axis_string,
271  x_axis_units, y_axis_units,
272  marker_style=4,
273  title_string='fit', canvas=None):
274  """
275  Plot the beam ellipse for a set of points
276  - x_axis_string: string name of the variable to go on the x_axis
277  - y_axis_string: string name of the variable to go on the y_axis
278  - x_axis_units: string units on the x axis
279  - y_axis_units: string units on the y axis
280  - title_string: title of the histogram
281 
282  Return value is a tuple of (canvas, histogram, ellipse, graph) where
283  canvas is an object of type ROOT.TCanvas, histogram is an object of type
284  ROOT.TH2D, ellipse is an object of type ROOT.TF2 and graph is an object
285  of type ROOT.TGraph
286  """
287  common.has_root()
288  x_var = self.keys.index(x_axis_string)
289  x_list = [x[x_var]*common.units[x_axis_units] for x in self.points]
290  y_var = self.keys.index(y_axis_string)
291  y_list = [x[y_var]*common.units[y_axis_units] for x in self.points]
292  if x_axis_units != '':
293  x_axis_string += ' ['+x_axis_units+']'
294  if y_axis_units != '':
295  y_axis_string += ' ['+y_axis_units+']'
296  hist, graph = common.make_root_graph(title_string,
297  x_list, x_axis_string, y_list, y_axis_string)
298  if canvas == None:
299  canvas = common.make_root_canvas(title_string)
300  hist.Draw()
301  canvas.cd()
302  x_min = hist.GetXaxis().GetXmin()
303  x_max = hist.GetXaxis().GetXmax()
304  y_min = hist.GetYaxis().GetXmin()
305  y_max = hist.GetYaxis().GetXmax()
306  graph.SetMarkerStyle(marker_style)
307  graph.Draw('p')
308  ell_centre = [self.centre[x_var], self.centre[y_var]]
309  ell_matrix = [[self.ellipse[x_var, x_var], self.ellipse[x_var, y_var]],
310  [self.ellipse[x_var, y_var], self.ellipse[y_var, y_var]]]
311  ell_matrix = numpy.array(ell_matrix)
312  contours = [self.get_mean_noise()]
313  ellipse = common.make_root_ellipse_function(ell_centre, ell_matrix,
314  [2.], x_min, x_max, y_min, y_max)
315  ellipse.Draw("SAME")
316  canvas.Update()
317  return canvas, hist, ellipse, graph
318 
319  def json_repr(self):
320  """
321  Represent the iteration as a json object
322  - points: list of float points in the iteration
323  - keys: list of keys that defines the points
324  - eps_max: maximum error on the ellipse fit
325 
326  (To get the noise, etc call __init__ against this data)
327  """
328  return {
329  "points":self.points,
330  "keys":self.keys,
331  "eps_max":self.eps_max
332  }
common module defines common utility data and functions that are used elsewhere
Definition: __init__.py:1
def scan_generator
Make a generator to scan over the variables, looking for closed orbits.