← Back to branch summary

~christopher-hunt08/maus/maus_integrated_kalman

« back to all changes in this revision

Viewing changes to src/common_py/analysis/inspectors.py

MAUS-v2.3.0

Show diffs side-by-side

added added

removed removed

Lines of Context:
src/common_py/analysis/inspectors.py
 
1
#!/usr/bin/env python
 
2
 
 
3
# This file is part of MAUS: http://micewww.pp.rl.ac.uk/projects/maus
 
4
#
 
5
# MAUS is free software: you can redistribute it and/or modify
 
6
# it under the terms of the GNU General Public License as published by
 
7
# the Free Software Foundation, either version 3 of the License, or
 
8
# (at your option) any later version.
 
9
#
 
10
# MAUS is distributed in the hope that it will be useful,
 
11
# but WITHOUT ANY WARRANTY; without even the implied warranty of
 
12
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 
13
# GNU General Public License for more details.
 
14
#
 
15
# You should have received a copy of the GNU General Public License
 
16
# along with MAUS.  If not, see <http://www.gnu.org/licenses/>.
 
17
#
 
18
 
 
19
# pylint: disable = W0311, E1101, W0613, C0111, R0911, W0621, C0103, R0902
 
20
# pylint: disable = W0102
 
21
 
 
22
 
 
23
import analysis
 
24
import analysis.covariances
 
25
 
 
26
import ROOT
 
27
 
 
28
 
 
29
class PhaseSpace2DInspector() :
 
30
 
 
31
  def __init__(self, plane_id, ensemble_size) :
 
32
    self.plane = plane_id
 
33
    self.position = 0.0
 
34
    self.ensemble_size = ensemble_size
 
35
    self.covariance = analysis.covariances.CovarianceMatrix()
 
36
 
 
37
    self.position_plot = ROOT.TH2F(\
 
38
           'inspected_position_{0}'.format(self.plane), 'Beam Position', \
 
39
                                        400, -200.0, 200.0, 400, -200.0, 200.0)
 
40
    self.momentum_plot = ROOT.TH2F(\
 
41
           'inspected_momentum_{0}'.format(self.plane), 'Beam Momentum', \
 
42
                                        400, -200.0, 200.0, 400, -200.0, 200.0)
 
43
    self.x_phasespace_plot = ROOT.TH2F(\
 
44
        'inspected_x_phasespace_{0}'.format(self.plane), 'X-Phasespace', \
 
45
                                        400, -200.0, 200.0, 400, -200.0, 200.0)
 
46
    self.y_phasespace_plot = ROOT.TH2F(\
 
47
        'inspected_y_phasespace_{0}'.format(self.plane), 'Y-Phasespace', \
 
48
                                        400, -200.0, 200.0, 400, -200.0, 200.0)
 
49
    self.xy_phasespace_plot = ROOT.TH2F(\
 
50
        'inspected_xpy_phasespace_{0}'.format(self.plane), 'X-Py-Phasespace', \
 
51
                                        400, -200.0, 200.0, 400, -200.0, 200.0)
 
52
    self.yx_phasespace_plot = ROOT.TH2F(\
 
53
        'inspected_ypx_phasespace_{0}'.format(self.plane), 'Y-Px-Phasespace', \
 
54
                                        400, -200.0, 200.0, 400, -200.0, 200.0)
 
55
 
 
56
    self.pz_plot = ROOT.TH1F(\
 
57
              'inspected_pz_{0}'.format(self.plane), 'p_{z}', 400, 0.0, 400.0 )
 
58
    self.p_plot = ROOT.TH1F(\
 
59
                   'inspected_p_{0}'.format(self.plane), 'p', 400, 0.0, 400.0 )
 
60
 
 
61
    
 
62
    self.emittance_plot = ROOT.TH1F(\
 
63
              'inspected_emittance_{0}'.format(self.plane), 'Emittance', \
 
64
                                                               1000, 0.0, 20.0)
 
65
    self.emittance_x_plot = ROOT.TH1F(\
 
66
           'inspected_emittance_x_{0}'.format(self.plane), 'X Emittance', \
 
67
                                                            1000, 0.0, 20.0)
 
68
    self.emittance_y_plot = ROOT.TH1F(\
 
69
         'inspected_emittance_y_{0}'.format(self.plane), 'Y Emittance', \
 
70
                                                            1000, 0.0, 20.0)
 
71
    self.alpha_plot = ROOT.TH1F(\
 
72
               'inspected_alpha_{0}'.format(self.plane), 'Alpha Function', \
 
73
                                                             2000, -5.0, 5.0)
 
74
    self.alpha_x_plot = ROOT.TH1F(\
 
75
               'inspected_alpha_x_{0}'.format(self.plane), 'X Alpha Function', \
 
76
                                                             2000, -5.0, 5.0)
 
77
    self.alpha_y_plot = ROOT.TH1F(\
 
78
               'inspected_alpha_y_{0}'.format(self.plane), 'Y Alpha Function', \
 
79
                                                             2000, -5.0, 5.0)
 
80
    self.beta_plot = ROOT.TH1F(\
 
81
               'inspected_beta_{0}'.format(self.plane), 'Beta Function', \
 
82
                                                             3000, 0.0, 3000.0)
 
83
    self.beta_x_plot = ROOT.TH1F(\
 
84
               'inspected_beta_x_{0}'.format(self.plane), 'X Beta Function', \
 
85
                                                             3000, 0.0, 3000.0)
 
86
    self.beta_y_plot = ROOT.TH1F(\
 
87
               'inspected_beta_y_{0}'.format(self.plane), 'Y Beta Function', \
 
88
                                                             3000, 0.0, 3000.0)
 
89
    self.total_momentum_plot = ROOT.TH1F(\
 
90
                'inspected_momentum_{0}'.format(self.plane), 'Momentum', \
 
91
                                                              1000, 0.0, 500.0)
 
92
 
 
93
 
 
94
  def set_covariance_correction(self, correction, axes=['x', 'px', 'y', 'py']) :
 
95
    self.covariance.set_covariance_correction( correction, axes )
 
96
 
 
97
 
 
98
  def add_hit(self, hit) :
 
99
    self.covariance.add_hit(hit)
 
100
 
 
101
    self.position_plot.Fill(hit.get_x(), hit.get_y())
 
102
    self.momentum_plot.Fill(hit.get_px(), hit.get_py())
 
103
    self.x_phasespace_plot.Fill(hit.get_x(), hit.get_px())
 
104
    self.y_phasespace_plot.Fill(hit.get_y(), hit.get_py())
 
105
    self.xy_phasespace_plot.Fill(hit.get_x(), hit.get_py())
 
106
    self.yx_phasespace_plot.Fill(hit.get_y(), hit.get_px())
 
107
    self.pz_plot.Fill(hit.get_pz())
 
108
    self.p_plot.Fill(hit.get_p())
 
109
    self.position = hit.get_z()
 
110
 
 
111
    if self.ensemble_size > 1 and \
 
112
                               self.covariance.length() > self.ensemble_size :
 
113
      self.fill_plots()
 
114
 
 
115
 
 
116
  def fill_plots(self) :
 
117
    if self.covariance.length() < (self.ensemble_size / 2) :
 
118
      return
 
119
 
 
120
    self.emittance_plot.Fill(self.covariance.get_emittance(\
 
121
                                                          ['x','px','y','py']))
 
122
    self.emittance_x_plot.Fill(self.covariance.get_emittance(['x','px']))
 
123
    self.emittance_y_plot.Fill(self.covariance.get_emittance(['y','py']))
 
124
    self.alpha_plot.Fill(self.covariance.get_alpha(['x','y']))
 
125
    self.alpha_x_plot.Fill(self.covariance.get_alpha(['x']))
 
126
    self.alpha_y_plot.Fill(self.covariance.get_alpha(['y']))
 
127
    self.beta_plot.Fill(self.covariance.get_beta(['x','y']))
 
128
    self.beta_x_plot.Fill(self.covariance.get_beta(['x']))
 
129
    self.beta_y_plot.Fill(self.covariance.get_beta(['y']))
 
130
    self.total_momentum_plot.Fill(self.covariance.get_momentum())
 
131
  
 
132
    self.covariance.clear()
 
133
 
 
134
 
 
135
  def get_current_ensemble_size(self) :
 
136
    return self.covariance.length()
 
137
 
 
138
 
 
139
  def get_number_ensembles(self) :
 
140
    return self.emittance_plot.GetEntries()
 
141
 
 
142
 
 
143
  def get_plot_dictionary(self) :
 
144
    ins_plots = {}
 
145
    ins_plots['x_y'] = self.position_plot
 
146
    ins_plots['px_py'] = self.momentum_plot
 
147
    ins_plots['x_px'] = self.x_phasespace_plot
 
148
    ins_plots['y_py'] = self.y_phasespace_plot
 
149
    ins_plots['x_py'] = self.xy_phasespace_plot
 
150
    ins_plots['y_px'] = self.yx_phasespace_plot
 
151
    ins_plots['pz'] = self.pz_plot
 
152
    ins_plots['p'] = self.p_plot
 
153
    ins_plots['emittance'] = self.emittance_plot
 
154
    ins_plots['emittance_x'] = self.emittance_x_plot
 
155
    ins_plots['emittance_y'] = self.emittance_y_plot
 
156
    ins_plots['alpha'] = self.alpha_plot
 
157
    ins_plots['alpha_x'] = self.alpha_x_plot
 
158
    ins_plots['alpha_y'] = self.alpha_y_plot
 
159
    ins_plots['beta'] = self.beta_plot
 
160
    ins_plots['beta_x'] = self.beta_x_plot
 
161
    ins_plots['beta_y'] = self.beta_y_plot
 
162
    ins_plots['momentum'] = self.total_momentum_plot
 
163
 
 
164
    return ins_plots
 
165
 
 
166
 
 
167
  def get_data_dictionary(self) :
 
168
    ins_data = {}
 
169
    ins_data['position'] = self.position
 
170
    ins_data['x_mean'] = self.position_plot.GetMean(1)
 
171
    ins_data['y_mean'] = self.position_plot.GetMean(2)
 
172
    ins_data['px_mean'] = self.momentum_plot.GetMean(1)
 
173
    ins_data['py_mean'] = self.momentum_plot.GetMean(2)
 
174
    ins_data['emittance_mean'] = self.emittance_plot.GetMean()
 
175
    ins_data['emittance_x_mean'] = self.emittance_x_plot.GetMean()
 
176
    ins_data['emittance_y_mean'] = self.emittance_y_plot.GetMean()
 
177
    ins_data['beta_mean'] = self.beta_plot.GetMean()
 
178
    ins_data['beta_x_mean'] = self.beta_x_plot.GetMean()
 
179
    ins_data['beta_y_mean'] = self.beta_y_plot.GetMean()
 
180
    ins_data['alpha_mean'] = self.alpha_plot.GetMean()
 
181
    ins_data['alpha_x_mean'] = self.alpha_x_plot.GetMean() 
 
182
    ins_data['alpha_y_mean'] = self.alpha_y_plot.GetMean()
 
183
    ins_data['momentum_mean'] = self.momentum_plot.GetMean()
 
184
 
 
185
    ins_data['x_rms'] = self.position_plot.GetRMS(1)
 
186
    ins_data['y_rms'] = self.position_plot.GetRMS(2)
 
187
    ins_data['px_rms'] = self.momentum_plot.GetRMS(1)
 
188
    ins_data['py_rms'] = self.momentum_plot.GetRMS(2)
 
189
    ins_data['emittance_rms'] = self.emittance_plot.GetRMS()
 
190
    ins_data['emittance_x_rms'] = self.emittance_x_plot.GetRMS()
 
191
    ins_data['emittance_y_rms'] = self.emittance_y_plot.GetRMS()
 
192
    ins_data['beta_rms'] = self.beta_plot.GetRMS()
 
193
    ins_data['beta_x_rms'] = self.beta_x_plot.GetRMS()
 
194
    ins_data['beta_y_rms'] = self.beta_y_plot.GetRMS()
 
195
    ins_data['alpha_rms'] = self.alpha_plot.GetRMS()
 
196
    ins_data['alpha_x_rms'] = self.alpha_x_plot.GetRMS() 
 
197
    ins_data['alpha_y_rms'] = self.alpha_y_plot.GetRMS()
 
198
    ins_data['momentum_rms'] = self.momentum_plot.GetRMS()
 
199
 
 
200
    return ins_data
 
201
 
 
202