← Back to branch summary

~jan.greis/maus/straighttemp

« back to all changes in this revision

Viewing changes to src/common_cpp/Recon/Kalman/MAUSTrackWrapper.cc

  • Committer: Jan Greis
  • Date: 2015-06-29 09:07:42 UTC
  • mfrom: (697.2.8 release)
  • Revision ID: j.r.greis@warwick.ac.uk-20150629090742-3mnfledv3r24bka9
Merge from 0.9.6

Show diffs side-by-side

added added

removed removed

Lines of Context:
src/common_cpp/Recon/Kalman/MAUSTrackWrapper.cc
 
1
/* This file is part of MAUS: http://micewww.pp.rl.ac.uk:8080/projects/maus
 
2
 *
 
3
 * MAUS 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
 * MAUS 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 MAUS.  If not, see <http://www.gnu.org/licenses/>.
 
15
 *
 
16
 */
 
17
 
 
18
#include "src/common_cpp/Recon/Kalman/MAUSTrackWrapper.hh"
 
19
 
 
20
#include <vector>
 
21
 
 
22
#include "TMatrixD.h"
 
23
 
 
24
namespace MAUS {
 
25
 
 
26
  template<class ELEMENT>
 
27
  bool SortByZ(const ELEMENT *a, const ELEMENT *b) {
 
28
      return (a->get_position().z() > b->get_position().z());
 
29
  }
 
30
 
 
31
 
 
32
 
 
33
  Kalman::Track BuildTrack(SciFiClusterPArray cluster_array, const SciFiGeometryHelper* geom) {
 
34
    Kalman::Track new_track(1);
 
35
    size_t numbclusters = cluster_array.size();
 
36
 
 
37
    if (numbclusters < 1) return new_track;
 
38
    int tracker = cluster_array[0]->get_tracker();
 
39
 
 
40
    const SciFiPlaneMap& geom_map = geom->GeometryMap().find(tracker)->second.Planes;
 
41
    int tracker_const = (tracker == 0 ? -1 : 1);
 
42
 
 
43
    for (SciFiPlaneMap::const_iterator iter = geom_map.begin(); iter != geom_map.end(); ++iter) {
 
44
      int id = iter->first * tracker_const;
 
45
      Kalman::State new_state = Kalman::State(1, iter->second.Position.z());
 
46
      new_state.SetId(id);
 
47
      new_track.Append(new_state);
 
48
    }
 
49
 
 
50
    for (size_t j = 0; j < numbclusters; ++j) {
 
51
      SciFiCluster* cluster = cluster_array[j];
 
52
 
 
53
      int id = (cluster->get_station() - 1)*3 + cluster->get_plane(); // Actually (id - 1)
 
54
 
 
55
      // TODO :
 
56
      // - APPLY GEOMETRY CORRECTIONS!
 
57
      // - Fill covariance matrix correctly!
 
58
      TMatrixD state_vector(1, 1);
 
59
      TMatrixD covariance(1, 1);
 
60
 
 
61
      state_vector(0, 0) = cluster->get_alpha();
 
62
      covariance(0, 0) = 0.0;
 
63
 
 
64
      new_track[id].SetVector(state_vector);
 
65
      new_track[id].SetCovariance(covariance);
 
66
    }
 
67
    return new_track;
 
68
  }
 
69
 
 
70
 
 
71
 
 
72
 
 
73
  Kalman::State ComputeSeed(SciFiHelicalPRTrack* h_track, const SciFiGeometryHelper* geom,
 
74
                                                       bool correct_energy_loss, double seed_cov) {
 
75
    TMatrixD vector(5, 1);
 
76
    TMatrixD covariance(5, 5);
 
77
 
 
78
    int tracker = h_track->get_tracker();
 
79
    int seed_id = (tracker == 0 ? -15 : 15);
 
80
    double seed_pos = geom->GetPlanePosition(tracker, 5, 2);
 
81
    double length =  seed_pos;
 
82
    double c  = CLHEP::c_light;
 
83
    double particle_charge = h_track->get_charge();
 
84
    double Bz = geom->GetFieldValue(tracker);
 
85
 
 
86
    // Get h_track values.
 
87
    double r  = h_track->get_R();
 
88
    double pt = - c*Bz*r*particle_charge;
 
89
    double dsdz = - h_track->get_dsdz();
 
90
    double x0 = h_track->get_circle_x0(); // Circle Center x
 
91
    double y0 = h_track->get_circle_y0(); // Circle Center y
 
92
    double s = (h_track->get_line_sz_c() - length*dsdz); // Path length at start plane
 
93
    double phi_0 = s / r; // Phi at start plane
 
94
    double phi = phi_0 + TMath::PiOver2(); // Direction of momentum
 
95
 
 
96
    // TODO: Actually propagate the track parrameters and covariance matrix back to start plane.
 
97
    //       This is an approximation.
 
98
    ThreeVector patrec_momentum(-pt*sin(phi_0), pt*cos(phi_0), - pt/dsdz);
 
99
 
 
100
    if (correct_energy_loss) {
 
101
      double P = patrec_momentum.mag();
 
102
      double patrec_bias; // Account for two planes of energy loss
 
103
      if (tracker == 0) {
 
104
        patrec_bias = (P + 1.4) / P;
 
105
      } else {
 
106
        patrec_bias = (P + 1.4) / P;
 
107
      }
 
108
      patrec_momentum = patrec_bias * patrec_momentum;
 
109
    }
 
110
 
 
111
    double x = x0 + r*cos(phi_0);
 
112
    double px = patrec_momentum.x();
 
113
    double y = y0 + r*sin(phi_0);
 
114
    double py = patrec_momentum.y();
 
115
    double kappa = particle_charge / patrec_momentum.z();
 
116
 
 
117
    vector(0, 0) = x;
 
118
    vector(1, 0) = px;
 
119
    vector(2, 0) = y;
 
120
    vector(3, 0) = py;
 
121
    vector(4, 0) = kappa;
 
122
 
 
123
//
 
124
// METHOD = ED SANTOS
 
125
//
 
126
    if (seed_cov > 0.0) {
 
127
      covariance.Zero();
 
128
      for (int i = 0; i < 5; ++i) {
 
129
        covariance(i, i) = seed_cov;
 
130
      }
 
131
    } else {
 
132
//
 
133
// METHOD = CHRISTOPHER HUNT
 
134
//
 
135
      std::vector<double> cov = h_track->get_covariance();
 
136
      TMatrixD patrec_covariance(5, 5);
 
137
      if (cov.size() != 25) {
 
138
        throw MAUS::Exception(MAUS::Exception::recoverable,
 
139
                              "Dimension of covariance matrix does not match the state vector",
 
140
                              "KalmanSeed::ComputeInitalCovariance(Helical)");
 
141
      }
 
142
 
 
143
      double mc = particle_charge*c*Bz; // Magnetic constant
 
144
      double sin = std::sin(phi_0);
 
145
      double cos = std::cos(phi_0);
 
146
      double sin_plus = std::sin(phi);
 
147
      double cos_plus = std::cos(phi);
 
148
 
 
149
      TMatrixD jacobian(5, 5);
 
150
      jacobian(0, 0) = 1.0;
 
151
      jacobian(0, 2) = cos + phi*sin;
 
152
      jacobian(0, 3) = -sin;
 
153
 
 
154
      jacobian(1, 2) = mc*cos_plus + mc*phi*sin_plus;
 
155
      jacobian(1, 3) = -mc*sin_plus;
 
156
 
 
157
      jacobian(2, 1) = 1.0;
 
158
      jacobian(2, 2) = sin - phi*cos;
 
159
      jacobian(2, 3) = cos;
 
160
 
 
161
      jacobian(3, 2) = mc*sin_plus - mc*phi*cos_plus;
 
162
      jacobian(1, 3) = mc*cos_plus;
 
163
 
 
164
      jacobian(4, 3) = -dsdz / (mc*r*r);
 
165
      jacobian(4, 4) = 1.0 / (mc*r);
 
166
 
 
167
      TMatrixD jacobianT(5, 5);
 
168
      jacobianT.Transpose(jacobian);
 
169
 
 
170
      for (int i = 0; i < 5; ++i) {
 
171
        for (int j = 0; j < 5; ++j) {
 
172
          patrec_covariance(i, j) = cov.at(i*5 + j);
 
173
        }
 
174
      }
 
175
      covariance = jacobian*patrec_covariance*jacobianT;
 
176
    }
 
177
 
 
178
    Kalman::State seed_state(vector, covariance, seed_pos);
 
179
    seed_state.SetId(seed_id);
 
180
 
 
181
    return seed_state;
 
182
  }
 
183
 
 
184
 
 
185
  Kalman::State ComputeSeed(SciFiStraightPRTrack* s_track, const SciFiGeometryHelper* geom,
 
186
                                                                                 double seed_cov) {
 
187
    TMatrixD vector(4, 1);
 
188
    TMatrixD covariance(4, 4);
 
189
 
 
190
    int tracker = s_track->get_tracker();
 
191
    double seed_pos = geom->GetPlanePosition(tracker, 5, 2);
 
192
    double length =  seed_pos;
 
193
    int seed_id = (tracker == 0 ? -15 : 15);
 
194
 
 
195
    double x0 = s_track->get_x0();
 
196
    double y0 = s_track->get_y0();
 
197
    double mx = s_track->get_mx();
 
198
    double my = s_track->get_my();
 
199
 
 
200
    // Get position at the starting end of tracker
 
201
    double x = x0 + mx*length;
 
202
    double y = y0 + my*length;
 
203
 
 
204
    vector(0, 0) = x;
 
205
    vector(1, 0) = mx;
 
206
    vector(2, 0) = y;
 
207
    vector(3, 0) = my;
 
208
 
 
209
//
 
210
// METHOD = ED SANTOS
 
211
//
 
212
    if (seed_cov > 0.0) {
 
213
      covariance.Zero();
 
214
      for (int i = 0; i < 4; ++i) {
 
215
        covariance(i, i) = seed_cov;
 
216
      }
 
217
    } else {
 
218
//
 
219
// METHOD = CHRISTOPHER HUNT
 
220
//
 
221
      std::vector<double> cov = s_track->get_covariance();
 
222
      TMatrixD patrec_covariance(4, 4);
 
223
 
 
224
      if (cov.size() != (unsigned int)16) {
 
225
        throw MAUS::Exception(MAUS::Exception::recoverable,
 
226
                              "Dimension of covariance matrix does not match the state vector",
 
227
                              "KalmanSeed::ComputeInitalCovariance(Straight)");
 
228
      }
 
229
 
 
230
      TMatrixD jacobian(4, 4);
 
231
      jacobian(0, 0) = 1.0;
 
232
      jacobian(1, 1) = 1.0;
 
233
      jacobian(2, 2) = 1.0;
 
234
      jacobian(3, 3) = 1.0;
 
235
      jacobian(0, 1) = length;
 
236
      jacobian(2, 3) = length;
 
237
 
 
238
      TMatrixD jacobianT(4, 4);
 
239
      jacobianT.Transpose(jacobian);
 
240
 
 
241
      for (int i = 0; i < 4; ++i) {
 
242
        for (int j = 0; j < 4; ++j) {
 
243
          patrec_covariance(i, j) = cov.at(i*4 + j);
 
244
        }
 
245
      }
 
246
 
 
247
      covariance = jacobian*patrec_covariance*jacobianT;
 
248
    }
 
249
 
 
250
    Kalman::State seed_state(vector, covariance, seed_pos);
 
251
    seed_state.SetId(seed_id);
 
252
    return seed_state;
 
253
  }
 
254
 
 
255
 
 
256
  SciFiTrack* ConvertToSciFiTrack(const Kalman::TrackFit* fitter,
 
257
                                                                 const SciFiGeometryHelper* geom) {
 
258
    SciFiTrack* new_track = new SciFiTrack();
 
259
    const Kalman::Track& smoothed = fitter->Smoothed();
 
260
//    const Kalman::Track& smoothed = fitter->Filtered();
 
261
//    const Kalman::Track& filtered = fitter->Filtered();
 
262
//    const Kalman::Track& predicted = fitter->Predicted();
 
263
    const Kalman::Track& data = fitter->Data();
 
264
    Kalman::State seed = fitter->GetSeed();
 
265
    double default_mom = geom->GetDefaultMomentum();
 
266
 
 
267
    if (smoothed.GetLength() < 1)
 
268
      throw MAUS::Exception(MAUS::Exception::recoverable,
 
269
                            "Not enough points in Kalman Track",
 
270
                            "ConvertToSciFiTrack()");
 
271
 
 
272
    double chi_squared = fitter->CalculateChiSquared(smoothed);
 
273
    int NDF = fitter->GetNDF();
 
274
    double p_value = TMath::Prob(chi_squared, NDF);
 
275
 
 
276
    int tracker;
 
277
    if (smoothed[0].GetId() > 0) {
 
278
      tracker = 1;
 
279
    } else {
 
280
      tracker = 0;
 
281
    }
 
282
    new_track->set_tracker(tracker);
 
283
 
 
284
    int dimension = smoothed.GetDimension();
 
285
    if (dimension == 4) {
 
286
      new_track->SetAlgorithmUsed(SciFiTrack::kalman_straight);
 
287
    } else if (dimension == 5) {
 
288
      new_track->SetAlgorithmUsed(SciFiTrack::kalman_helical);
 
289
    } else {
 
290
      throw MAUS::Exception(MAUS::Exception::recoverable,
 
291
                            "Unexpected dimension of Kalman::Track",
 
292
                            "ConvertToSciFiTrack()");
 
293
    }
 
294
 
 
295
//    const SciFiGeometryMap& geom_map = geom->GeometryMap();
 
296
 
 
297
    ThreeVector reference_pos = geom->GetReferencePosition(tracker);
 
298
    HepRotation reference_rot = geom->GetReferenceRotation(tracker);
 
299
    int charge = 0;
 
300
 
 
301
    for (unsigned int i = 0; i < smoothed.GetLength(); ++i) {
 
302
      const Kalman::State& smoothed_state = smoothed[i];
 
303
//      const Kalman::State& filtered_state = filtered[i];
 
304
//      const Kalman::State& predicted_state = predicted[i];
 
305
      const Kalman::State& data_state = data[i];
 
306
 
 
307
      SciFiTrackPoint* new_point = new SciFiTrackPoint();
 
308
 
 
309
      new_point->set_tracker(tracker);
 
310
 
 
311
      int id = abs(smoothed_state.GetId());
 
312
      new_point->set_station(((id-1)/3)+1);
 
313
      new_point->set_plane((id-1)%3);
 
314
 
 
315
      ThreeVector pos;
 
316
      ThreeVector mom;
 
317
 
 
318
      TMatrixD state_vector = smoothed_state.GetVector();
 
319
 
 
320
      if (dimension == 4) {
 
321
        pos.setZ(smoothed_state.GetPosition());
 
322
        pos.setX(state_vector(0, 0));
 
323
        mom.setX(state_vector(1, 0)*default_mom);
 
324
        pos.setY(state_vector(2, 0));
 
325
        mom.setY(state_vector(3, 0)*default_mom);
 
326
 
 
327
        pos *= reference_rot;
 
328
        pos += reference_pos;
 
329
 
 
330
        mom *= reference_rot;
 
331
        if (tracker == 0) mom *= -1.0;
 
332
        mom.setZ(default_mom); // MeV/c
 
333
      } else if (dimension == 5) {
 
334
        pos.setX(state_vector(0, 0));
 
335
        mom.setX(state_vector(1, 0));
 
336
        pos.setY(state_vector(2, 0));
 
337
        mom.setY(state_vector(3, 0));
 
338
        pos.setZ(smoothed_state.GetPosition());
 
339
 
 
340
        pos *= reference_rot;
 
341
        pos += reference_pos;
 
342
 
 
343
        mom *= reference_rot;
 
344
        if (mom.z() < 0.0) {
 
345
          charge = -1;
 
346
        } else {
 
347
          charge = 1;
 
348
        }
 
349
        mom.setZ(fabs(1.0/state_vector(4, 0)));
 
350
      }
 
351
 
 
352
//      if (mom.z() < 0.0) {
 
353
//        mom.setZ(- mom.z());
 
354
//      }
 
355
//      if (tracker == 0) {
 
356
//        mom.setZ(fabs(mom.z()));
 
357
//        mom.setY(-1.0*mom.y());
 
358
//      }
 
359
 
 
360
      new_point->set_pos(pos);
 
361
      new_point->set_mom(mom);
 
362
 
 
363
      // TODO
 
364
    //  _pull              = kalman_site->residual(KalmanState::Projected)(0, 0);
 
365
    //  _residual          = kalman_site->residual(KalmanState::Filtered)(0, 0);
 
366
    //  _smoothed_residual = kalman_site->residual(KalmanState::Smoothed)(0, 0);
 
367
    //  AND CHARGE!
 
368
      if (data_state) {
 
369
//        new_point->set_pull(sqrt(fitter->CalculatePredictedResidual(i).GetVector().E2Norm()));
 
370
//        new_point->set_residual(sqrt(fitter->CalculateFilteredResidual(i).GetVector().E2Norm()));
 
371
//        new_point->set_smoothed_residual(
 
372
//                                sqrt(fitter->CalculateSmoothedResidual(i).GetVector().E2Norm()));
 
373
        new_point->set_pull(fitter->CalculatePredictedResidual(i).GetVector()(0, 0));
 
374
        new_point->set_residual(fitter->CalculateFilteredResidual(i).GetVector()(0, 0));
 
375
        new_point->set_smoothed_residual(fitter->CalculateSmoothedResidual(i).GetVector()(0, 0));
 
376
      } else {
 
377
        new_point->set_pull(0.0);
 
378
        new_point->set_residual(0.0);
 
379
        new_point->set_smoothed_residual(0.0);
 
380
      }
 
381
 
 
382
 
 
383
      TMatrixD C = smoothed_state.GetCovariance();
 
384
      int size = C.GetNrows();
 
385
      int num_elements = size*size;
 
386
 
 
387
      double* matrix_elements = C.GetMatrixArray();
 
388
      std::vector<double> covariance(num_elements);
 
389
      for (int i = 0; i < num_elements; ++i) {
 
390
        covariance[i] = matrix_elements[i];
 
391
      }
 
392
      new_point->set_covariance(covariance);
 
393
      std::vector<double> errors(size);
 
394
      for (int i = 0; i < size; ++i) {
 
395
        errors[i] = std::sqrt(fabs(C(i, i)));
 
396
      }
 
397
      new_point->set_errors(errors);
 
398
 
 
399
// TODO
 
400
//    _cluster = new TRef(kalman_site->cluster());
 
401
      new_point->set_cluster(new TRef());
 
402
      new_track->add_scifitrackpoint(new_point);
 
403
    }
 
404
 
 
405
    new_track->set_charge(charge);
 
406
    new_track->set_chi2(chi_squared);
 
407
    new_track->set_ndf(NDF);
 
408
    new_track->set_P_value(p_value);
 
409
 
 
410
    TMatrixD seed_vector = seed.GetVector();
 
411
    ThreeVector seed_pos;
 
412
    ThreeVector seed_mom;
 
413
    if (dimension == 4) {
 
414
      seed_pos.setZ(seed.GetPosition());
 
415
      seed_pos.setX(seed_vector(0, 0));
 
416
      seed_mom.setX(seed_vector(1, 0)*default_mom);
 
417
      seed_pos.setY(seed_vector(2, 0));
 
418
      seed_mom.setY(seed_vector(3, 0)*default_mom);
 
419
    } else if (dimension == 5) {
 
420
      seed_pos.setX(seed_vector(0, 0));
 
421
      seed_mom.setX(seed_vector(1, 0));
 
422
      seed_pos.setY(seed_vector(2, 0));
 
423
      seed_mom.setY(seed_vector(3, 0));
 
424
      seed_pos.setZ(seed.GetPosition());
 
425
    }
 
426
    seed_pos *= reference_rot;
 
427
    seed_pos += reference_pos;
 
428
 
 
429
    seed_mom *= reference_rot;
 
430
 
 
431
    if (dimension == 4) {
 
432
      seed_mom.setZ(default_mom);
 
433
    } else if (dimension == 5) {
 
434
      seed_mom.setZ(fabs(1.0/seed_vector(4, 0)));
 
435
    }
 
436
 
 
437
    new_track->SetSeedPosition(seed_pos);
 
438
    new_track->SetSeedMomentum(seed_mom);
 
439
    new_track->SetSeedCovariance(seed.GetCovariance().GetMatrixArray(), dimension*dimension);
 
440
 
 
441
// TODO:
 
442
// - Set Cluster
 
443
// - Calculate p-value
 
444
// - Set Algorithm used
 
445
// - Set Seed Info
 
446
// - Init track before the fit?
 
447
 
 
448
    return new_track;
 
449
  }
 
450
 
 
451
 
 
452
  Kalman::Track BuildSpacepointTrack(SciFiSpacePointPArray spacepoints,
 
453
                                    const SciFiGeometryHelper* geom, int plane_num, double smear) {
 
454
//    TRandom3 rand;
 
455
 
 
456
    Kalman::Track new_track(2);
 
457
    int tracker = (*spacepoints.begin())->get_tracker();
 
458
 
 
459
    for (unsigned int i = 0; i < 5; ++i) {
 
460
      Kalman::State new_state(2, geom->GetPlanePosition(tracker, i+1, plane_num));
 
461
      new_state.SetId((1 + i*3 + plane_num)*(tracker == 0 ? -1 : 1));
 
462
      new_track.Append(new_state);
 
463
    }
 
464
 
 
465
    for (SciFiSpacePointPArray::iterator it = spacepoints.begin(); it != spacepoints.end(); ++it) {
 
466
      int station = (*it)->get_station();
 
467
      TMatrixD vec(2, 1);
 
468
      TMatrixD cov(2, 2);
 
469
      cov.Zero();
 
470
      vec(0, 0) = (*it)->get_position().x();// * (1.0 + rand.Gaus(0.0, smear));
 
471
      vec(1, 0) = (*it)->get_position().y();// * (1.0 + rand.Gaus(0.0, smear));
 
472
      cov(0, 0) = 0.0;
 
473
      cov(1, 1) = 0.0;
 
474
      new_track[station-1].SetVector(vec);
 
475
      new_track[station-1].SetCovariance(cov);
 
476
      new_track[station-1].SetPosition((*it)->get_position().z());
 
477
    }
 
478
 
 
479
    return new_track;
 
480
  }
 
481
}
 
482