lcsim-contrib/src/main/java/org/lcsim/contrib/SteveMagill
diff -N MyCheatReconDriver.java
--- /dev/null 1 Jan 1970 00:00:00 -0000
+++ MyCheatReconDriver.java 14 Feb 2012 19:19:58 -0000 1.1
@@ -0,0 +1,538 @@
+/*
+ * MyCheatReconDriver.java
+ *
+ * Created on April 27, 2007, 6:12 AM
+ * modified digisim names for DR cal
+ *
+ * Driver to do a cheat reconstruction. Output is lists of FSparticles, clusters,
+ * tracks, and reconstructed particles
+ */
+
+package org.lcsim.contrib.SteveMagill;
+
+import org.lcsim.event.EventHeader;
+import org.lcsim.util.Driver;
+import org.lcsim.event.MCParticle;
+import org.lcsim.event.ReconstructedParticle;
+import org.lcsim.event.Track;
+import org.lcsim.event.Cluster;
+import org.lcsim.event.base.BaseTrackMC;
+import org.lcsim.event.base.MCReconstructedParticle;
+import org.lcsim.recon.cluster.cheat.CheatClusterDriver;
+import org.lcsim.recon.cluster.cheat.CheatCluster;
+import org.lcsim.recon.cluster.util.TraceOrigin;
+import org.lcsim.recon.tracking.cheat.CheatTrackDriver;
+import org.lcsim.event.util.CreateFinalStateMCParticleList;
+import org.lcsim.geometry.Detector;
+import org.lcsim.geometry.compact.Subdetector;
+import org.lcsim.detector.solids.Tube;
+import org.lcsim.mc.fast.tracking.MCFastTracking;
+import java.util.HashMap;
+import java.util.List;
+import java.util.ArrayList;
+import java.util.Map;
+import org.lcsim.util.lcio.LCIOConstants;
+import hep.physics.vec.Hep3Vector;
+import org.lcsim.util.aida.AIDA;
+import org.lcsim.event.LCRelation;
+import org.lcsim.digisim.MyLCRelation;
+import org.lcsim.recon.cheater.AddSmearing;
+/**
+ *
+ * @author cassell
+ */
+public class MyCheatReconDriver extends Driver
+{
+ private AIDA aida = AIDA.defaultInstance();
+ String inFSname = "GenFinalStateParticles";
+ String inTname = "RefinedCheatTracks";
+ String tempTname = "TempCheatTracks";
+ String tempPname = "TempCheatParticles";
+ String inCname = "RefinedCheatClusters";
+ String outRname;
+ String outVRname;
+ String outGenRname;
+ String outGenVRname;
+ String outFSname;
+ String outTrFSname;
+ String outTname;
+ String outFSTname;
+ String outCname;
+ String outFastMC = "FastMCTracks";
+// String[] diginames;
+//
+// Default parameters for defining "trackable" track
+//
+ double mintfromin = 4;
+ double mintfromout = 3;
+ double mintlength = 100.;
+ double ptcut = .1;
+// Ask Jeremy how to do this
+ double emR = 1270.;
+ double emZ = 1680.;
+// Too close cut
+ double tccut = 200.;
+ Map<MCParticle,Track> tmap;
+ TraceOrigin to;
+ boolean init;
+ /** Creates a new instance of CheatSimReconParticleDriver */
+ public MyCheatReconDriver()
+ {
+ outRname = "ReconPerfectReconParticles";
+// outVRname = "MCFSRecoParticles";
+ outVRname = "ReconPerfectVisReconParticles";
+ outGenRname = "GenPerfectReconParticles";
+ outGenVRname = "GenPerfectVisReconParticles";
+ outFSname = "ReconFSParticles";
+ outTrFSname = "TrackedReconFSParticles";
+ outTname = "ReconTracks";
+// outTname = "PerfectTracks";
+ outFSTname = "FSReconTracks";
+ outCname = "ReconClusters";
+//
+// Make an initial list of MCParticles to be considered
+//
+ CreateFinalStateMCParticleList cfs = new CreateFinalStateMCParticleList("Gen");
+ cfs.setCollectionName(inFSname);
+ add(cfs);
+//
+// Make cheating TrackerHit Clusters
+//
+ add(new CheatTrackDriver(tempTname, tempPname));
+//
+// Make FastMC smeared tracks
+//
+ MCFastTracking mcft = new MCFastTracking();
+ mcft.setFSList(tempPname);
+ mcft.setOutputList(outFastMC);
+ add(mcft);
+//
+// Add smearing and error matrix to tracks
+//
+ add(new AddSmearing(tempTname,outFastMC,inTname));
+//
+// Make CalorimeterHit Cheat Clusters
+//
+// add(new DigiPackageDriver());
+// init = false;
+// diginames = new String[4];
+ String[] diginames = {"Edep_EcalBarrDigiHits","Edep_EcalEndcapDigiHits","Edep_HcalBarrDigiHits","Edep_HcalEndcapDigiHits"};
+ // output of DigiSim is hit collections - EcalBarrDigiHits, EcalEndcapDigiHits, HcalBarrDigiHits, HcalEndcapDigiHits
+ add(new CheatClusterDriver(diginames, inCname));
+ }
+ public void setCheatTrackOutputName(String s)
+ {outTname = s;}
+ public void setCheatClusterOutputName(String s)
+ {outCname = s;}
+ public void setCheatFSParticleOutputName(String s)
+ {outFSname = s;}
+ public void setCheatTrackedFSParticleOutputName(String s)
+ {outTrFSname = s;}
+ public void setCheatFSTrackOutputName(String s)
+ {outFSTname = s;}
+ public void setCheatReconstructedParticleOutputName(String s)
+ {outRname = s;}
+ public void setCheatVisReconstructedParticleOutputName(String s)
+ {outVRname = s;}
+ public void setGenReconstructedParticleOutputName(String s)
+ {outGenRname = s;}
+ public void setGenVisReconstructedParticleOutputName(String s)
+ {outGenVRname = s;}
+ public void setTooCloseCut(double x)
+ {tccut = x;}
+ public void process(EventHeader event)
+ {
+// if(!init)
+// {
+// init = true;
+// CalorimeterInformation ci = CalorimeterInformation.instance();
+// diginames[0] = "EM_BARREL";
+// diginames[1] = "EM_ENDCAP";
+// diginames[2] = "HAD_BARREL";
+// diginames[3] = "HAD_BARREL";
+// }
+ super.process(event);
+ List<MCParticle> initfs = event.get(MCParticle.class,inFSname);
+ List<Track> inittracks = event.get(Track.class,inTname);
+ List<Cluster> initclusters = event.get(Cluster.class,inCname);
+ List<MCParticle> reconfs = new ArrayList<MCParticle>();
+ List<MCParticle> trreconfs = new ArrayList<MCParticle>();
+ List<Track> fstracks = new ArrayList<Track>();
+ List<Track> tracks = new ArrayList<Track>();
+ List<Track> forvtxingt = new ArrayList<Track>();
+ List<Cluster> fsclusters = new ArrayList<Cluster>();
+ List<Cluster> nonfsclusters = new ArrayList<Cluster>();
+ List<ReconstructedParticle> perfect = new ArrayList<ReconstructedParticle>();
+ List<ReconstructedParticle> visperfect = new ArrayList<ReconstructedParticle>();
+ List<ReconstructedParticle> genperfect = new ArrayList<ReconstructedParticle>();
+ List<ReconstructedParticle> genvisperfect = new ArrayList<ReconstructedParticle>();
+ List<ReconstructedParticle> forvtxingperfect = new ArrayList<ReconstructedParticle>();
+ tmap = new HashMap<MCParticle,Track>();
+//
+// Make a map of trackable tracks, and add LCRelations to event:
+// "ReconTracksToMCP" -> from=Track, to = MCParticle
+// "TracksToParent" -> from=Track, to=Track, linking tracks to parent tracks. Most to Tracks are null
+//
+ List<LCRelation> lcrelationList = new ArrayList<LCRelation>();
+ List<LCRelation> parentage = new ArrayList<LCRelation>();
+ for(Track t:inittracks)
+ {
+ MCParticle p = ((BaseTrackMC) t).getMCParticle();
+ if(isTrackable(p,t))
+ {
+ tmap.put(p,t);
+ tracks.add(t);
+ lcrelationList.add(new MyLCRelation(t,p));
+ }
+ }
+ for(MCParticle p:tmap.keySet())
+ {
+ Track t = tmap.get(p);
+ Track pt = getParentTrack(p);
+ parentage.add(new MyLCRelation(t,pt));
+ if(pt == null)forvtxingt.add(t);
+ }
+ event.put("ReconTracksToMCP", lcrelationList, LCRelation.class, 0);
+ event.put("TracksToParent", parentage, LCRelation.class, 0);
+//
+// Start with the generator final state particles
+//
+ double Einit = 0.;
+ double Eadd = 0.;
+ double Esub = 0.;
+ for(MCParticle p:initfs)
+ {
+ Einit += p.getEnergy();
+//
+// For the majority of generator fs particles, they originate near the ip
+// and end near or beyond the Ecal. Just keep them
+//
+ Hep3Vector vtx = p.getOrigin();
+ Hep3Vector ep = p.getEndPoint();
+ if(vtx.magnitude() < tccut)
+ {
+ if( (Math.sqrt(ep.x()*ep.x()+ep.y()*ep.y())>emR-tccut)||(Math.abs(ep.z())>emZ-tccut) )
+ {
+ reconfs.add(p);
+ continue;
+ }
+ }
+//
+// Check if we really want the parent: Only if the vtx is close to Ecal
+//
+ if( (Math.sqrt(vtx.x()*vtx.x()+vtx.y()*vtx.y())>emR-tccut)||(Math.abs(vtx.z())>emZ-tccut) )
+ {
+ MCParticle par = p.getParents().get(0);
+ Hep3Vector pvtx = par.getOrigin();
+ while( (Math.sqrt(pvtx.x()*pvtx.x()+pvtx.y()*pvtx.y())>emR-tccut)||(Math.abs(pvtx.z())>emZ-tccut) )
+ {
+ par = par.getParents().get(0);
+ pvtx = par.getOrigin();
+ }
+// System.out.println("Parent replacement: vtx close to Ecal");
+ if(!reconfs.contains(par))
+ {
+ reconfs.add(par);
+ Eadd += par.getEnergy();
+// System.out.println("Adding to FS list "+par.getType().getName()+": E = "+par.getEnergy());
+ }
+ Esub += p.getEnergy();
+// System.out.println("Subtracting from FS list "+p.getType().getName()+": E = "+p.getEnergy());
+ continue;
+ }
+//
+// Done checking for parentage replacement. Now check for offspring replacement
+//
+//
+// If the endpoint is close to or beyond the Ecal, keep it
+//
+ if( (Math.sqrt(ep.x()*ep.x()+ep.y()*ep.y())>emR-tccut)||(Math.abs(ep.z())>emZ-tccut) )
+ {
+ reconfs.add(p);
+ continue;
+ }
+//
+// If the particle is trackable and there are no trackable offspring, keep it.
+// Otherwise, replace it.
+//
+
+ if(tmap.containsKey(p)&&!hasTrackableDaughters(p))
+ {
+ reconfs.add(p);
+ continue;
+ }
+//
+// Now do the offspring replacement. Only keep trackable particles and
+// particles which reach the Ecal.
+//
+ Esub += p.getEnergy();
+// System.out.println("Offspring replacement");
+// System.out.println("Removing "+p.getType().getName()+": E = "+p.getEnergy());
+ List<MCParticle> repl = new ArrayList<MCParticle>(p.getDaughters());
+ while(repl.size() > 0)
+ {
+ MCParticle d = repl.get(0);
+ repl.remove(d);
+ if(d.getSimulatorStatus().isBackscatter())continue;
+ if(d.getSimulatorStatus().vertexIsNotEndpointOfParent())continue;
+ if(tmap.containsKey(d))
+ {
+ reconfs.add(d);
+ Eadd += d.getEnergy();
+// System.out.println("Adding "+d.getType().getName()+": E = "+d.getEnergy());
+ }
+ else
+ {
+ Hep3Vector dep = d.getEndPoint();
+ if( (Math.sqrt(dep.x()*dep.x()+dep.y()*dep.y())>emR)||(Math.abs(dep.z())>emZ) )
+ {
+ reconfs.add(d);
+ Eadd += d.getEnergy();
+// System.out.println("Adding "+d.getType().getName()+": E = "+d.getEnergy());
+ }
+ else
+ {
+ repl.addAll(d.getDaughters());
+ }
+ }
+ }
+ }
+// System.out.println("Done with replacement lists:");
+// System.out.println("E GenFS = "+Einit+": Added E = "+Eadd+": Subtracted E = "+Esub);
+//
+// Now make the recon lists.
+//
+//
+// First the tracks. All trackable and final state trackable
+//
+// System.out.println("Adding this many tracks to MCRPs" + tracks.size());
+ for(Track t:tracks)
+ {
+ MCParticle p = ((BaseTrackMC) t).getMCParticle();
+ if(reconfs.contains(p))
+ {
+ fstracks.add(t);
+ trreconfs.add(p);
+ }
+ }
+//
+// Now the clusters. Trace the cheated clusters back to the reconfs
+// list. Also trace to genfs list.
+//
+ to = new TraceOrigin(reconfs);
+ Map<MCParticle,Cluster> fsclmap = new HashMap<MCParticle,Cluster>();
+ TraceOrigin gento = new TraceOrigin(initfs);
+ Map<MCParticle,Cluster> genfsclmap = new HashMap<MCParticle,Cluster>();
+ for(Cluster c:initclusters)
+ {
+ CheatCluster cc = (CheatCluster) c;
+ MCParticle trmc = to.traceit(cc.getMCParticle());
+ if(trmc == null)
+ {
+ nonfsclusters.add(c);
+ }
+ else
+ {
+ if(fsclmap.containsKey(trmc))
+ {
+ CheatCluster bc = (CheatCluster) fsclmap.get(trmc);
+ bc.addCluster(c);
+ }
+ else
+ {
+ CheatCluster bc = new CheatCluster(trmc);
+ bc.addCluster(c);
+ fsclmap.put(trmc, bc);
+ }
+ }
+ trmc = gento.traceit(cc.getMCParticle());
+ if(trmc != null)
+ {
+ if(genfsclmap.containsKey(trmc))
+ {
+ CheatCluster bc = (CheatCluster) genfsclmap.get(trmc);
+ bc.addCluster(c);
+ }
+ else
+ {
+ CheatCluster bc = new CheatCluster(trmc);
+ bc.addCluster(c);
+ genfsclmap.put(trmc, bc);
+ }
+ }
+ }
+//
+// Have the map, make the list of clusters
+//
+ fsclusters.addAll(fsclmap.values());
+//
+// Output the lists to event
+//
+ event.put(outFSname,reconfs,MCParticle.class,0);
+ event.getMetaData(reconfs).setSubset(true);
+ event.put(outTrFSname,trreconfs,MCParticle.class,0);
+ event.getMetaData(trreconfs).setSubset(true);
+ int flag = 1<<LCIOConstants.TRBIT_HITS;
+ event.put(outTname,tracks,BaseTrackMC.class,flag);
+ event.put(outFSTname,fstracks,BaseTrackMC.class,flag);
+ event.put("ForVtxingTracks",forvtxingt,BaseTrackMC.class,flag);
+ event.getMetaData(fstracks).setSubset(true);
+ flag = 1<<LCIOConstants.CLBIT_HITS;
+ event.put(outCname,fsclusters,CheatCluster.class,flag);
+ event.put("NonFSReconClusters",nonfsclusters,CheatCluster.class,flag);
+ for(MCParticle p:reconfs)
+ {
+//
+// Eliminate the neutrinos
+//
+ int pid = Math.abs(p.getPDGID());
+ if( (pid == 12)||(pid == 14)||(pid == 16) )continue;
+ MCReconstructedParticle rp = new MCReconstructedParticle(p);
+ if(tmap.containsKey(p))rp.addTrack(tmap.get(p));
+ if(fsclmap.containsKey(p))rp.addCluster(fsclmap.get(p));
+ perfect.add(rp);
+ if(tmap.containsKey(p)||fsclmap.containsKey(p))visperfect.add(rp);
+ }
+ for(Track t:forvtxingt)
+ {
+ MCParticle p = ((BaseTrackMC) t).getMCParticle();
+ MCReconstructedParticle rp = new MCReconstructedParticle(p);
+ rp.addTrack(t);
+ if(fsclmap.containsKey(p))rp.addCluster(fsclmap.get(p));
+ forvtxingperfect.add(rp);
+ }
+ event.put(outRname,perfect);
+ event.put(outVRname,visperfect);
+ event.put("ForVtxingPerfect",forvtxingperfect);
+ for(MCParticle p:initfs)
+ {
+//
+// Eliminate the neutrinos
+//
+ int pid = Math.abs(p.getPDGID());
+ if( (pid == 12)||(pid == 14)||(pid == 16) )continue;
+ MCReconstructedParticle rp = new MCReconstructedParticle(p);
+ if(tmap.containsKey(p))rp.addTrack(tmap.get(p));
+ List<MCParticle> dlist = new ArrayList<MCParticle>(p.getDaughters());
+ while(dlist.size() > 0)
+ {
+ MCParticle d = dlist.get(0);
+ dlist.remove(d);
+ if(tmap.containsKey(d))rp.addTrack(tmap.get(d));
+ dlist.addAll(d.getDaughters());
+ }
+ if(genfsclmap.containsKey(p))rp.addCluster(genfsclmap.get(p));
+ genperfect.add(rp);
+ if(tmap.containsKey(p)||genfsclmap.containsKey(p))genvisperfect.add(rp);
+ }
+ event.put(outGenRname,genperfect);
+ event.put(outGenVRname,genvisperfect);
+ }
+ private boolean isTrackable(MCParticle p,Track t)
+ {
+ boolean tr = false;
+//
+// Always return false for neutrals
+//
+ if(p.getCharge() == 0)return tr;
+//
+// Apply a pt cut
+//
+ Hep3Vector m = p.getMomentum();
+ if(Math.sqrt(m.x()*m.x()+m.y()*m.y()) < ptcut)return tr;
+//
+// Apply a minimum length cut
+//
+ Hep3Vector vtx = p.getOrigin();
+ Hep3Vector ep = p.getEndPoint();
+ double d = Math.sqrt( (ep.x()-vtx.x())*(ep.x()-vtx.x()) +
+ (ep.y()-vtx.y())*(ep.y()-vtx.y()) +
+ (ep.z()-vtx.z())*(ep.z()-vtx.z()) );
+ if(d > mintlength)
+ {
+//
+// Apply a minimum #hits cut, which can depend on whether or not
+// the particle makes it to the Ecal.
+//
+ if( (Math.sqrt(ep.x()*ep.x()+ep.y()*ep.y())>emR)||(Math.abs(ep.z())>emZ) )
+ {
+ if(t.getTrackerHits().size() >= mintfromout)tr=true;
+ }
+ else
+ {
+ if(t.getTrackerHits().size() >= mintfromin)tr=true;
+ }
+ }
+ return tr;
+ }
+ private Track getParentTrack(MCParticle p)
+ {
+ Track rtr = null;
+ MCParticle parent = p;
+ while((parent.getParents().size() == 1)&&(rtr == null))
+ {
+ parent = parent.getParents().get(0);
+ if(tmap.containsKey(parent))
+ {
+ rtr = tmap.get(parent);
+ }
+ }
+ return rtr;
+ }
+ private boolean hasTrackableDaughters(MCParticle p)
+ {
+ boolean htd = false;
+ List<MCParticle> check = new ArrayList<MCParticle>();
+ for(MCParticle d:p.getDaughters())
+ {
+ if(d.getSimulatorStatus().isBackscatter())continue;
+ if(d.getSimulatorStatus().vertexIsNotEndpointOfParent())continue;
+ check.add(d);
+ }
+ while(!htd && (check.size() > 0))
+ {
+ if(tmap.containsKey(check.get(0)))htd = true;
+ else
+ {
+ List<MCParticle> dl = check.get(0).getDaughters();
+ check.remove(0);
+ for(MCParticle d:dl)
+ {
+ if(d.getSimulatorStatus().isBackscatter())continue;
+ if(d.getSimulatorStatus().vertexIsNotEndpointOfParent())continue;
+ check.add(d);
+ }
+ }
+ }
+ return htd;
+ }
+ protected void detectorChanged(Detector detector)
+ {
+ super.detectorChanged(detector);
+
+ Map<String,Subdetector> subdets = detector.getSubdetectors();
+ for (Subdetector subdet : subdets.values())
+ {
+ if (subdet.isCalorimeter())
+ {
+ if (subdet.getName().equals("EMBarrel"))
+ {
+ Tube tube = (Tube)subdet.getDetectorElement().getGeometry().getLogicalVolume().getSolid();
+ emR = tube.getInnerRadius();
+ double outerRadius = tube.getOuterRadius();
+ double zHalfLength = tube.getZHalfLength();
+ }
+ if (subdet.getName().equals("EMEndcap"))
+ {
+ Tube tube = (Tube)subdet.getDetectorElement().getChildren().get(0).getGeometry().getLogicalVolume().getSolid();
+ double zHalfLength = tube.getZHalfLength();
+ Hep3Vector v = subdet.getDetectorElement().getChildren().get(0).getGeometry().getPosition();
+ emZ = Math.abs(v.z()) - zHalfLength;
+ }
+ }
+ }
+ System.out.println("Detector change: innerZ = "+emZ+" innerR = "+emR);
+
+ }
+
+}