lcsim-contrib/src/main/java/org/lcsim/contrib/CosminDeaconu/TauPolarization
diff -N TauTauStdhepRepairer.java
--- /dev/null 1 Jan 1970 00:00:00 -0000
+++ TauTauStdhepRepairer.java 16 Jan 2009 19:27:47 -0000 1.1
@@ -0,0 +1,341 @@
+package org.lcsim.contrib.CosminDeaconu.TauPolarization;
+
+import hep.io.stdhep.StdhepEvent;
+import hep.io.stdhep.StdhepReader;
+import hep.io.stdhep.StdhepRecord;
+import hep.io.stdhep.StdhepRunRecord;
+import hep.io.stdhep.StdhepWriter;
+import java.io.File;
+import java.io.IOException;
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.Map;
+import java.util.Random;
+import java.util.Set;
+
+/**
+ * This program takes a stdhep with no tau movement as input,
+ * calculates a decay for the taus and outputs a new stdhep file.
+ *
+ * The file to work on is taken as the first argument.
+ *
+ * If there is a second argument, it is used as the output file name.
+ * Otherwise, the output file name is constructed from the input and the string
+ * "_repaired";
+ *
+ * A last argument named --clobber may also exist, which will automatically
+ * overwrite existing output files instead of prompting.
+ *
+ * @author Cosmin Deaconu <[log in to unmask]>
+ */
+
+public class TauTauStdhepRepairer {
+
+
+ public static void main(String[] args) {
+
+ //Deal with bad number of arguments
+ if (args.length <1 || args.length > 3) {
+ printUsage();
+ System.exit(1);
+ }
+
+
+ //Figure out input name
+ String inputFileName = args[0];
+ if (!inputFileName.endsWith(".stdhep")) {
+ System.out.println("Please give me a file that ends w/ .stdhep");
+ System.exit(2);
+ }
+
+
+ boolean clobber = false;
+ //Figure out output name / clobbering
+ String outputFileName = inputFileName.replace(".stdhep", "_repaired.stdhep");
+ if (args.length == 2) {
+ if (args[1].toLowerCase().equals("--clobber")) clobber = true;
+ else outputFileName = args[1];
+
+ }
+
+ if (args.length==3) {
+
+ if (args[2].toLowerCase().equals("--clobber")) clobber = true;
+ else {
+ printUsage();
+ System.exit(1);
+ }
+ }
+
+ //Handle clobbering
+ File o = new File(outputFileName);
+ if (o.exists()) {
+ if (clobber) o.delete();
+ else {
+ System.out.println(o.getAbsolutePath()+" already exists. Clobber? (y/n)");
+ String input = System.console().readLine();
+ if (input.toLowerCase().equals("y")) {
+ o.delete();
+ } else {
+ System.out.println("Exiting without doing anything.");
+ System.exit(0);
+ }
+ }
+ }
+
+ //Init reader
+ StdhepReader reader = null;
+ try {
+ reader = new StdhepReader(inputFileName);
+ } catch(IOException ioe) {
+ System.out.println("IO Exception: "+ioe.getMessage());
+ System.exit(3);
+ }
+
+ int numEvents = reader.getNumberOfEvents();
+
+ //Init writer
+ StdhepWriter writer = null;
+ try {
+ writer = new StdhepWriter(outputFileName, reader.getTitle()+"_repaired", reader.getComment() + "with fixed tau endpoints.", numEvents);
+ } catch(IOException ioe) {
+ System.out.println("IO Exception: "+ioe.getMessage());
+ System.exit(3);
+ }
+
+ //copy events
+ while(true) {
+ try {
+ StdhepRecord record = reader.nextRecord();
+ //If it's a begin or end record (whatever those mean), just copy cover
+ if (record instanceof StdhepRunRecord) {
+ writer.writeRecord(record);
+
+ //Otherwise, we write a modified version of the record
+ } else {
+ if (!(record instanceof StdhepEvent)) { //Don't know what to do with these records (if they exist)
+ throw new RuntimeException("Unrecognized record type");
+ }
+ writer.writeRecord(fixEvent((StdhepEvent) record));
+ }
+
+ } catch (IOException ioe) {
+ break;
+ }
+ }
+
+ // Clean up
+ try {
+ writer.close();
+ System.out.println(o.getAbsolutePath() + " successfully written out. Done.");
+ } catch(IOException ioe) {
+ System.out.println("IO Exception: "+ioe.getMessage());
+ System.exit(3);
+ }
+ }
+
+ //Random seed for determining lifetime fraction
+ private static final long RANDOM_SEED = 123;
+ private static Random random = new Random(RANDOM_SEED);
+
+ //Retrieves a random lifetime fraction.
+ private static double getRandomLifeTimeFraction() {
+ //use this loop because random.nextDouble() may return 0.
+ double r = 0.0;
+ while (r==0.0) {
+ r = random.nextDouble();
+ }
+ return -Math.log(r);
+ }
+
+ //Recursively translates all daughters
+ private static void deltaDaughters(int particleIndex, double dx, double dy, double dz, double dt, double[]vhep, Map<Integer,Set<Integer>> parentMap, int recurLength) {
+
+ //System.out.println(particleIndex);
+
+ if (!parentMap.containsKey(particleIndex)) return;
+ for (Integer i : parentMap.get(particleIndex)) {
+
+ writeVertexInfo(i, 0, dx, vhep);
+ writeVertexInfo(i, 1, dy, vhep);
+ writeVertexInfo(i, 2, dz, vhep);
+ writeVertexInfo(i, 3, dt, vhep);
+ deltaDaughters(i,dx,dy,dz,dt,vhep, parentMap,recurLength+1);
+ }
+
+ }
+
+ //Adds delta to the given coordinate of the given particle
+ private static void writeVertexInfo(int particleIndex, int coordIndex,
+ double delta, double[] vhep) {
+ vhep[4*particleIndex+coordIndex]+= delta;
+ }
+
+ private static final double TAU_CT = 87.11; //microns
+ private static final int ABS_TAU_PDGID = 15;
+
+ //For every tau, we calculate a random tau lifetime, figure out the distance
+ //corresponding to it and shift all daughters by that distance.
+ private static StdhepRecord fixEvent(StdhepEvent evt) {
+ int evnNum = evt.getNEVHEP();
+ int numParticles = evt.getNHEP();
+
+
+ //Get usable daughter information
+ Map<Integer,Set<Integer>> parentMap = getParentageInfo(evt);
+
+ //Grab status information
+ int[] isthep = new int[numParticles];
+ for (int i = 0; i < numParticles; i++) isthep[i] = evt.getISTHEP(i);
+
+ //Grab id information
+ int[] idthep = new int[numParticles];
+ for (int i = 0; i < numParticles; i++) idthep[i] = evt.getIDHEP(i);
+
+ //Grab mother information
+ int[] jmohep = new int[2*numParticles];
+ for (int i = 0; i < 2*numParticles; i++) {
+ jmohep[i] = evt.getJMOHEP(i/2, i%2);
+ }
+
+ //Grab daughter information
+ int[] jdahep = new int[2*numParticles];
+ for (int i = 0; i < 2*numParticles; i++) {
+ jdahep[i] = evt.getJDAHEP(i/2, i%2);
+ }
+
+ //Grab vertex info
+ double[] vhep = new double[4*numParticles];
+ for (int i = 0; i < 4*numParticles; i++) {
+ vhep[i] = evt.getVHEP(i/4, i%4);
+ }
+
+ //Grab momentum info
+ double[] phep = new double[5*numParticles];
+ for (int i = 0; i < 5*numParticles; i++) {
+ phep[i] = evt.getPHEP(i/5, i%5);
+ }
+
+ int numTaus = 0;
+
+ for (int i = 0; i < numParticles; i++) {
+ //check for tau
+ int this_pdgid = Math.abs(evt.getIDHEP(i));
+ if (this_pdgid!=ABS_TAU_PDGID) continue;
+ numTaus++;
+
+ //Get random timelife fraction
+ double t_fraction = getRandomLifeTimeFraction();
+
+ //Gather known data from taus
+ double px = evt.getPHEP(i, 0);
+ double py = evt.getPHEP(i, 1);
+ double pz = evt.getPHEP(i, 2);
+ double p = Math.sqrt(px*px+py*py+pz*pz);
+ double e = evt.getPHEP(i,3);
+ double m = evt.getPHEP(i, 4);
+
+ //Calculate deltas for daughters
+ double dx = px/m*t_fraction*TAU_CT/1000.; //Microns to mm (stdhep uses mm, I think).
+ double dy = py/m*t_fraction*TAU_CT/1000.;
+ double dz = pz/m*t_fraction*TAU_CT/1000.;
+ double dt = p/e *t_fraction*TAU_CT/1000.; //Divide by 1000 to get value in mm/c
+
+ //System.out.println();
+ deltaDaughters(i, dx, dy, dz, dt, vhep, parentMap, 1);
+ }
+ if (numTaus!=2) {
+ System.out.println("We expect 2 taus. We found "+numTaus+".");
+ System.out.println("Exiting.");
+ System.exit(9);
+
+ }
+ return new StdhepEvent(evnNum, numParticles, isthep, idthep, jmohep, jdahep, phep, vhep);
+ }
+
+
+ private static void printUsage() {
+ System.out.println("TauTauStdhepRepairer INPUTFILE.stdhep [OUTPUTFILE.stdhep] [--clobber]");
+ System.out.println("If no OUTPUTFILE specified, will append _repaired to original filename");
+ System.out.println("Use clobber as the last argument to automatically clobber existing files with the same name as the output file.");
+ }
+
+
+ //taken directly from StdhepConverter
+ private static int fillIndexVec(int[] vec, int idx1, int idx2)
+ {
+ int l = 0;
+ try
+ {
+ if ( idx1 >= 0 && idx2 >= 0 )
+ {
+ if ( idx1 < idx2 )
+ {
+ for ( int i = idx1; i < (idx2 + 1); i++ )
+ {
+ vec[l++] = i;
+ }
+ }
+ else if ( idx1 > idx2 )
+ {
+ vec[l++] = idx1;
+ vec[l++] = idx2;
+ }
+ // indices are equal
+ else
+ {
+ vec[l++] = idx1;
+ }
+ }
+ else if ( idx1 >= 0 )
+ {
+ vec[l++] = idx1;
+ }
+ }
+ catch (ArrayIndexOutOfBoundsException e)
+ {
+ System.out.println("Out of Bounds Exception: "+e.getMessage());
+ System.exit(7);
+ }
+ return l;
+ }
+
+ //derived from StdhepConverter
+ //Helper function for getParentageInfo
+ private static void checkAndAddDaughter(Map<Integer,Set<Integer>> daughters, int parentID, int childID)
+ {
+ if (parentID == childID) return; // Can't be parent of self
+
+
+ if (daughters.containsKey(parentID)) {
+ daughters.get(parentID).add(childID);
+ } else {
+ HashSet<Integer> newSet = new HashSet<Integer>();
+ newSet.add(childID);
+ daughters.put(parentID,newSet);
+ }
+ }
+
+ //this is derived from StdhepConverter
+ //Creates a map from each particle's index to the set of its daughter's indices.
+ private static Map<Integer,Set<Integer>> getParentageInfo(StdhepEvent evt) {
+
+ int n = evt.getNHEP();
+ int[] vec = new int[n];
+ Map<Integer,Set<Integer>> daughters = new HashMap<Integer,Set<Integer>>(n);
+
+ for (int i=0; i<n; i++)
+ {
+ int idx1 = evt.getJDAHEP(i,0) % 10000 - 1;
+ int idx2 = evt.getJDAHEP(i,1) % 10000 - 1;
+ int l = fillIndexVec(vec,idx1,idx2);
+
+ for (int j=0; j<l; j++)
+ {
+ checkAndAddDaughter(daughters,i,vec[j]);
+ }
+ }
+ return daughters;
+ }
+
+}