1 package org.lcsim.recon.cheater;
2
3 import hep.physics.vec.Hep3Vector;
4 import hep.physics.vec.VecOp;
5
6 import java.util.ArrayList;
7 import java.util.Collections;
8 import java.util.Comparator;
9 import java.util.HashMap;
10 import java.util.List;
11 import java.util.Map;
12
13 import org.lcsim.event.EventHeader;
14 import org.lcsim.event.EventHeader.LCMetaData;
15 import org.lcsim.event.LCRelation;
16 import org.lcsim.event.MCParticle;
17 import org.lcsim.event.RelationalTable;
18 import org.lcsim.event.RelationalTable.Mode;
19 import org.lcsim.event.RelationalTable.Weighting;
20 import org.lcsim.event.base.BaseRelationalTable;
21 import org.lcsim.fit.helicaltrack.HelicalTrackHit;
22 import org.lcsim.util.Driver;
23
24
25
26
27
28
29
30
31
32
33
34
35 public class RemoveMultipleHelicalTrackHitsCheater extends Driver {
36
37 protected String trackHitCollection;
38 protected String trackHitRelations;
39 protected String trackMCRelations;
40 protected double coneAngle;
41
42 public RemoveMultipleHelicalTrackHitsCheater() {
43 trackHitCollection = "HelicalTrackHits";
44 trackHitRelations = "HelicalTrackHitRelations";
45 trackMCRelations = "HelicalTrackMCRelations";
46 coneAngle = Math.PI/18.;
47 }
48
49 public void setTrackHitCollection(String trackHitCollection) {
50 this.trackHitCollection = trackHitCollection;
51 }
52
53 public void setTrackHitRelations(String trackHitRelations) {
54 this.trackHitRelations = trackHitRelations;
55 }
56
57 public void setTrackMCRelations(String trackMCRelations) {
58 this.trackMCRelations = trackMCRelations;
59 }
60
61
62
63
64
65 public void setConeAngle(double coneAngle) {
66 this.coneAngle = coneAngle;
67 }
68
69 @Override
70 protected void process(EventHeader event) {
71
72 List<HelicalTrackHit> trackHits = event.get(HelicalTrackHit.class, trackHitCollection);
73 List<LCRelation> hitRelations = event.get(LCRelation.class, trackHitRelations);
74 List<LCRelation> mcRelations = event.get(LCRelation.class, trackMCRelations);
75
76
77 RelationalTable<HelicalTrackHit, LCRelation> hitToHitRelationMap = new BaseRelationalTable<HelicalTrackHit, LCRelation>(Mode.MANY_TO_MANY, Weighting.UNWEIGHTED);
78 for (LCRelation relation : hitRelations) {
79 hitToHitRelationMap.add((HelicalTrackHit) relation.getFrom(), relation);
80 }
81
82
83 RelationalTable<HelicalTrackHit, LCRelation> hitToMCRelationMap = new BaseRelationalTable<HelicalTrackHit, LCRelation>(Mode.MANY_TO_MANY, Weighting.UNWEIGHTED);
84 for (LCRelation relation : mcRelations) {
85 hitToMCRelationMap.add((HelicalTrackHit) relation.getFrom(), relation);
86 }
87
88
89 Map<String, RelationalTable<HelicalTrackHit, MCParticle>> layerToHitsToParticleMap = new HashMap<String, RelationalTable<HelicalTrackHit,MCParticle>>();
90
91
92 Map<String, List<HelicalTrackHit>> layerToHitsMap = new HashMap<String, List<HelicalTrackHit>>();
93
94
95 Map<String, List<MCParticle>> layerToParticleMap = new HashMap<String, List<MCParticle>>();
96
97
98 for (HelicalTrackHit hit : trackHits) {
99 List<MCParticle> particles = hit.getMCParticles();
100 String identifier = hit.getLayerIdentifier();
101 if (!layerToHitsMap.containsKey(identifier)) {
102 layerToHitsToParticleMap.put(identifier, new BaseRelationalTable<HelicalTrackHit, MCParticle>(Mode.MANY_TO_MANY, Weighting.WEIGHTED));
103 layerToHitsMap.put(identifier, new ArrayList<HelicalTrackHit>());
104 layerToParticleMap.put(identifier, new ArrayList<MCParticle>());
105 }
106 layerToHitsMap.get(identifier).add(hit);
107 for (MCParticle particle : particles) {
108 layerToHitsToParticleMap.get(identifier).add(hit, particle);
109 if (!layerToParticleMap.get(identifier).contains(particle)) {
110 layerToParticleMap.get(identifier).add(particle);
111 }
112 }
113
114 }
115
116
117 for (String identifier : layerToHitsMap.keySet()) {
118 RelationalTable<HelicalTrackHit, MCParticle> layerTable = layerToHitsToParticleMap.get(identifier);
119
120
121 List<HelicalTrackHit> processedHits = new ArrayList<HelicalTrackHit>();
122 for (MCParticle particle : layerToParticleMap.get(identifier)) {
123
124
125 if (layerTable.allTo(particle).size() < 2) continue;
126 List<HelicalTrackHit> particleHits = new ArrayList<HelicalTrackHit>(layerTable.allTo(particle));
127
128
129 Collections.sort(particleHits, new CompareHelicalTrackHitsByDistanceFromIP());
130 for (HelicalTrackHit hit : particleHits) {
131 Hep3Vector v1 = hit.getCorrectedPosition();
132 processedHits.add(hit);
133 for (HelicalTrackHit otherHit : particleHits) {
134 if (!trackHits.contains(otherHit) || processedHits.contains(otherHit)) continue;
135 Hep3Vector v12 = VecOp.sub(otherHit.getCorrectedPosition(), v1);
136
137
138 double deltaPhi = Math.acos(VecOp.dot(v1, v12)/(v1.magnitude()*v12.magnitude()));
139 if (Math.abs(deltaPhi) < coneAngle) {
140
141 trackHits.remove(otherHit);
142 hitRelations.removeAll(hitToHitRelationMap.allFrom(otherHit));
143 mcRelations.removeAll(hitToMCRelationMap.allFrom(otherHit));
144 }
145 }
146 }
147 }
148 }
149 }
150
151
152
153
154 protected class CompareHelicalTrackHitsByDistanceFromIP implements Comparator<HelicalTrackHit> {
155 @Override
156 public int compare(HelicalTrackHit o1, HelicalTrackHit o2) {
157 double r1 = o1.getCorrectedPosition().magnitude();
158 double r2 = o2.getCorrectedPosition().magnitude();
159 if (r1 < r2) {
160 return -1;
161 } else if (r2 < r1) {
162 return 1;
163 } else {
164 return 0;
165 }
166 }
167 }
168 }