lcsim-contrib/src/main/java/org/lcsim/contrib/Grefe/overlayEvents
diff -N OverlayEvents.java
--- /dev/null 1 Jan 1970 00:00:00 -0000
+++ OverlayEvents.java 11 Jan 2010 11:14:57 -0000 1.1
@@ -0,0 +1,406 @@
+package org.lcsim.contrib.Grefe.overlayEvents;
+
+import java.io.File;
+import java.io.IOException;
+import java.util.ArrayList;
+import java.util.Collection;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+import java.util.Random;
+
+import org.freehep.record.source.EndOfSourceException;
+import org.lcsim.event.EventHeader;
+import org.lcsim.event.GenericObject;
+import org.lcsim.event.MCParticle;
+import org.lcsim.event.SimCalorimeterHit;
+import org.lcsim.event.SimTrackerHit;
+import org.lcsim.event.EventHeader.LCMetaData;
+import org.lcsim.event.base.BaseSimCalorimeterHit;
+import org.lcsim.util.Driver;
+import org.lcsim.util.lcio.LCIOWriter;
+import org.lcsim.util.loop.LCIOEventSource;
+
+/*TODO At the moment timing issues are investigated using a certain time window, i.e. 10 BX.
+ * In that case the time structure is irrelevant because the time window is defined as the readout time window.
+ * A more realistic way would be to merge signal and background events for the whole bunch train and let the
+ * digitization and reconstruction handle the time stamping. This will require to change the time information
+ * for each hit accordingly.
+ */
+
+/**
+ * Driver to overlay events from different LCIO files over the signal LCIO file.
+ * The output event is an overlay of several bunch crossings with a certain bunch spacing in between.
+ * Every overlay LCIO file can contribute a different number of events per bunch crossing.
+ * If there are insufficient events in an overlay LCIO file it will be rewound and the events will be used multiple times.
+ * Collections can be omitted or mapped to different output collections if compatible.
+ * Currently supported types of collections are {@see org.lcsim.event.MCParticle}, {@see org.lcsim.event.SimTrackerHit} and {@see org.lcsim.event.SimCalorimeterHit}. Other collections are omitted.
+ * @author Christian Grefe <[log in to unmask]>
+ *
+ */
+public class OverlayEvents extends Driver {
+
+ protected int nBunchCrossings;
+ protected double bunchSpacing;
+ protected LCIOWriter writer;
+ protected EventHeader curEvent;
+ protected int curBX;
+ protected Map<File, LCIOEventSource> inputLcioFiles;
+ protected Map<File, Integer> inputFilesNEvents ;
+ protected List<String> ignoredCollections;
+ protected List<String> mergeOnlyCollections;
+ protected Map<String, String> convertCollections;
+ protected Random random;
+
+ /**
+ * Default constructor.
+ */
+ public OverlayEvents() {
+ random = new Random();
+ this.setBunchCrossings(0, 0.5);
+ inputFilesNEvents = new HashMap<File, Integer>();
+ inputLcioFiles = new HashMap<File, LCIOEventSource>();
+ convertCollections = new HashMap<String, String>();
+ ignoredCollections = new ArrayList<String>();
+ mergeOnlyCollections = new ArrayList<String>();
+ }
+
+ protected void startOfData() {
+ // Set up overlay LCIO event sources
+ for (File overlayFile : inputFilesNEvents.keySet() ) {
+ try {
+ LCIOEventSource lcio = new LCIOEventSource(overlayFile);
+ inputLcioFiles.put(overlayFile, lcio);
+ System.out.println("Setting up: "+overlayFile.getPath());
+ } catch (IOException e) {
+ System.err.println(e.getMessage());
+ }
+ }
+ // Set up ignored collections for output file
+ if (writer != null) {
+ for (String collection : ignoredCollections) {
+ writer.addIgnore(collection);
+ }
+ for (String collection : mergeOnlyCollections) {
+ writer.addWriteOnly(collection);
+ }
+ }
+ }
+
+ protected void process(EventHeader event) {
+ curEvent = event;
+ // TODO displace signal hits in time according to bunch crossing.
+ /*
+ if (nBunchCrossings > 1) {
+ int sigBX = random.nextInt(nBunchCrossings);
+ }
+ */
+ for (File overlayFile : inputFilesNEvents.keySet() ) {
+ LCIOEventSource lcio = inputLcioFiles.get(overlayFile);
+ for (int bX = 0; bX != nBunchCrossings; bX++) {
+ curBX = bX;
+ for (int bkgEvt = 0; bkgEvt != inputFilesNEvents.get(overlayFile); bkgEvt++) {
+ EventHeader overlayEvent = this.getNextEvent(lcio);
+ if (overlayEvent != null) {
+ this.mergeEvent(overlayEvent);
+ } else {
+ // remove flawed LCIO file from the list of overlay files
+ inputLcioFiles.remove(overlayFile);
+ inputFilesNEvents.remove(overlayFile);
+ System.err.println("Warning: "+overlayFile.getPath()+" is removed from the list of overlay events.");
+ }
+ }
+ }
+ }
+ // write to output file
+ if (writer != null) {
+ try {
+ writer.write(event);
+ System.out.println("Writing event "+event.getEventNumber());
+ }
+ catch (IOException e) {
+ System.err.println(e.getMessage());
+ }
+ }
+ }
+
+ protected void endOfData() {
+
+ }
+
+ /**
+ * Adds a background LCIO file, which will overlay one of its events on each bunch crossing.
+ * @param pathName The file name of the background LCIO.
+ * @throws IOException
+ */
+ public void addLcioFile(String pathName) throws IOException {
+ addLcioFile(pathName, 1);
+ }
+
+ /**
+ * Adds a background LCIO file, which will overlay one of its events on each bunch crossing.
+ * @param lcioFile The background LCIO file.
+ * @throws IOException
+ */
+ public void addLcioFile(File lcioFile) throws IOException {
+ addLcioFile(lcioFile, 1);
+ }
+
+ /**
+ * Adds a background LCIO file, which will overlay a number of events per bunch crossing.
+ * @param pathName The file name of the background LCIO file.
+ * @param nEventsToMerge The number of background events per bunch crossing.
+ * @throws IOException
+ */
+ public void addLcioFile(String pathName, int nEventsToMerge) throws IOException{
+ File lcioFile = new File(pathName);
+ addLcioFile(lcioFile, nEventsToMerge);
+ }
+
+ /**
+ * Adds a background LCIO file, which will overlay a number of events per bunch crossing.
+ * @param lcioFile The background LCIO file.
+ * @param nEventsToMerge The number of background events per bunch crossing.
+ * @throws IOException
+ */
+ public void addLcioFile(File lcioFile, int nEventsToMerge) throws IOException {
+ if (!lcioFile.exists()) {
+ throw new IOException(lcioFile.getPath()+" does not exist");
+ }
+ if (nEventsToMerge < 0) {
+ nEventsToMerge = 0;
+ }
+ inputFilesNEvents.put(lcioFile, nEventsToMerge);
+ }
+
+ /**
+ * Removes the collection from the output file.
+ * @param collectionName The name of the collection to be ignored.
+ */
+ public void ignoreCollection(String collectionName) {
+ if (!ignoredCollections.contains(collectionName)) {
+ ignoredCollections.add(collectionName);
+ }
+ }
+
+ /**
+ * Removes a list of collections from the output file.
+ * @param collections The list of collection names to be ignored.
+ */
+ public void ignoreCollections(List<String> collections) {
+ for (String collection : collections) {
+ ignoreCollection(collection);
+ }
+ }
+
+ /**
+ * Sets the list of collections which should be merged. Any other collection will be omitted.
+ * Collections in this list will be merged even if they are in the list of ignored collections.
+ * @param collections The list of collection names which will be merged only.
+ */
+ public void mergeOnly(List<String> collections) {
+ mergeOnlyCollections = collections;
+ }
+
+ /**
+ * Hits from the input collection will instead be written to the targetCollection.
+ * If the hits can not be converted to the type of the targetCollection they will be ignored.
+ * @param collection collection of hits which will be converted.
+ * @param targetCollection collection to which the hits will be added if possible.
+ */
+ public void convertCollection(String collection, String targetCollection) {
+ if (convertCollections.containsKey(collection)) {
+ System.err.println("Warning: overriding conversion of "+collection+" into "+targetCollection+".");
+ }
+ convertCollections.put(collection, targetCollection);
+ }
+
+ /**
+ * Sets the output file. If no output file is set the original LCIO will be overridden.
+ * @param pathName The LCIO file were the output is written to.
+ * @throws IOException
+ */
+ public void setOutputFile(String pathName) throws IOException {
+ File outputFile = new File(pathName);
+ this.setOutputFile(outputFile);
+ }
+
+ /**
+ * Sets the output file. If no output file is set the original LCIO will be overridden.
+ * @param outputFile The LCIO file were the output is written to.
+ * @throws IOException
+ */
+ public void setOutputFile(File outputFile) throws IOException {
+ writer = new LCIOWriter(outputFile);
+ System.out.println("Setting output file to "+outputFile.getPath());
+ }
+
+ /**
+ * Sets the number of bunch crossings to overlay over one signal event and the time between two bunch crossings.
+ * By default the signal event is placed randomly at one of the bunch crossings.
+ * @param nBuchCrossings The number of bunch crossings.
+ * @param bunchSpacing The time between two bunch crossings in ns.
+ */
+ public void setBunchCrossings(int nBuchCrossings, double bunchSpacing) {
+ if (nBuchCrossings > 1) {
+ this.nBunchCrossings = nBuchCrossings;
+ } else {
+ this.nBunchCrossings = 1;
+ }
+ if (bunchSpacing > 0.) {
+ this.bunchSpacing = bunchSpacing;
+ } else {
+ this.bunchSpacing = 0.;
+ }
+ }
+
+ /**
+ * Goes to the next event in the LCIOEventSource and returns it. If the end of the source is reached the source is rewound and the first event will be returned.
+ * If any other error occurs, i.e. the source does not exist, null is returned instead.
+ * @param lcio The LCIO source
+ * @return The next event in the LCIO file
+ */
+ protected EventHeader getNextEvent(LCIOEventSource lcio) {
+ EventHeader event = null;
+ boolean rewind = false;
+ try {
+ lcio.next();
+ event = (EventHeader) lcio.getCurrentRecord();
+ } catch (EndOfSourceException e) {
+ rewind = true;
+ } catch (Exception e) {
+ System.err.println(e.getMessage());
+ }
+ if (rewind) {
+ try {
+ lcio.rewind();
+ lcio.next();
+ event = (EventHeader) lcio.getCurrentRecord();
+ } catch (Exception e) {
+ System.err.println(e.getMessage());
+ }
+ }
+ return event;
+ }
+
+ /**
+ * Overlays an event over the current signal event. Merges all collections of both events which are not in the list of ignored collections or only those which are in the list of mergeOnly collections.
+ * @param overlayEvent The background event which will be overlayed over the current signal event.
+ */
+ protected void mergeEvent(EventHeader overlayEvent) {
+ if (!curEvent.getDetectorName().equals(overlayEvent.getDetectorName())) {
+ System.err.println("Warning: detector names do not match. Skipping overlay of signal event: "+curEvent.getEventNumber()+" and background event: "+overlayEvent.getEventNumber()+".");
+ return;
+ }
+ // get all collections from the overlay event
+ Collection<LCMetaData> collections = overlayEvent.getMetaData();
+ for (LCMetaData collection : collections) {
+ String collName = collection.getName();
+ if (!ignoredCollections.contains(collName) && (mergeOnlyCollections.size() == 0 || mergeOnlyCollections.contains(collName))) {
+ this.mergeCollection(collection);
+ }
+ }
+ }
+
+ /**
+ * Merges a collection from the background event with the corresponding collection from the signal event.
+ * Note that each type of collection needs its own implementation of how to merge the collections.
+ * Currently supported collections are {@see org.lcsim.event.MCParticle}, {@see org.lcsim.event.SimTrackerHit} and {@see org.lcsim.event.SimCalorimeterHit}.
+ * @param collection The background collection to be merged with the corresponding signal collection
+ */
+ protected void mergeCollection(LCMetaData collection) {
+ String collName = collection.getName();
+ String targetCollName = null;
+ if (convertCollections.containsKey(collName)) {
+ targetCollName = convertCollections.get(collName);
+ } else {
+ targetCollName = collName;
+ }
+ System.out.println("Merging "+collName);
+ // check if collection with this name exists in signal event, else just copy the collection
+ if (!curEvent.hasItem(targetCollName)) {
+ curEvent.put(targetCollName, (List)collection.getEvent().get(collName), collection.getType(), collection.getFlags(), collection.getStringParameters().get("READOUT_NAME")[0]);
+ }
+
+ // MCParticles
+ if (collection.getType().isAssignableFrom(MCParticle.class)) {
+ List<MCParticle> sigList = curEvent.get(MCParticle.class, targetCollName);
+ List<MCParticle> bkgList = collection.getEvent().get(MCParticle.class, collName);
+ sigList.addAll(bkgList);
+ // TODO include time offset by BX
+ return;
+ }
+
+ // SimTrackerHits
+ if (collection.getType().isAssignableFrom(SimTrackerHit.class)) {
+ List<SimTrackerHit> sigList = curEvent.get(SimTrackerHit.class, targetCollName);
+ List<SimTrackerHit> bkgList = collection.getEvent().get(SimTrackerHit.class, collName);
+ sigList.addAll(bkgList);
+ // TODO include time offset by BX
+ return;
+ }
+
+ // SimCaloHits
+ if (collection.getType().isAssignableFrom(SimCalorimeterHit.class)) {
+ List<SimCalorimeterHit> sigList = curEvent.get(SimCalorimeterHit.class, targetCollName);
+ List<SimCalorimeterHit> bkgList = collection.getEvent().get(SimCalorimeterHit.class, collName);
+ // build map of cells with entries
+ Map<Long,SimCalorimeterHit> hitMap = new HashMap<Long, SimCalorimeterHit>();
+ for (SimCalorimeterHit hit : sigList) {
+ hitMap.put(hit.getCellID(), hit);
+ }
+ for (SimCalorimeterHit hit : bkgList) {
+ if (hitMap.containsKey(hit.getCellID())) {
+ SimCalorimeterHit sigHit = hitMap.get(hit.getCellID());
+ int sigMcP = sigHit.getMCParticleCount();
+ int bkgMcP = hit.getMCParticleCount();
+ int nMcParticles = sigMcP + bkgMcP;
+ // arrays of mc particle contributions to the hit
+ Object[] mcList = new Object[nMcParticles];
+ float[] eneList = new float[nMcParticles];
+ float[] timList = new float[nMcParticles];
+ int[] pdgList = new int[nMcParticles];
+ // fill arrays with values from signal hit
+ for (int i =0; i != sigMcP; i++) {
+ mcList[i] = sigHit.getMCParticle(i);
+ eneList[i] = (float)sigHit.getContributedEnergy(i);
+ timList[i] = (float)sigHit.getContributedTime(i);
+ pdgList[i] = sigHit.getPDG(i);
+ }
+ // add values of background hit
+ for (int i =0; i != bkgMcP; i++) {
+ int j = sigMcP + i;
+ mcList[j] = hit.getMCParticle(i);
+ eneList[j] = (float)hit.getContributedEnergy(i);
+ timList[j] = (float)hit.getContributedTime(i);
+ pdgList[j] = hit.getPDG(i);
+ }
+ // create new combined hit, note "time" is inherited from BaseCalorimeterHit and not used by BaseSimCalorimeterHit, thus it is set to 0.
+ SimCalorimeterHit combinedHit = new BaseSimCalorimeterHit(hit.getCellID(), hit.getRawEnergy()+sigHit.getRawEnergy(), 0., mcList, eneList, timList, pdgList);
+ combinedHit.setMetaData(sigHit.getMetaData());
+ combinedHit.setDetectorElement(sigHit.getDetectorElement());
+ // replace old hit with combined hit
+ sigList.remove(sigHit);
+ sigList.add(combinedHit);
+ } else {
+ sigList.add(hit);
+ }
+ }
+ // TODO include time offset by BX
+ return;
+ }
+
+ // LCGenericObject
+ if (collection.getType().isAssignableFrom(GenericObject.class)) {
+ // need to treat all kinds of generic objects differently
+ if (collName.equals("MCParticleEndPointEnergy")) {
+ List<GenericObject> sigList = curEvent.get(GenericObject.class, targetCollName);
+ List<GenericObject> bkgList = collection.getEvent().get(GenericObject.class, collName);
+ sigList.addAll(bkgList);
+ return;
+ }
+ }
+ System.err.println("Collection "+collName+" can not be merged. Merging collections of type "+collection.getType()+" has not been implemented, yet.");
+ this.ignoreCollection(collName);
+ }
+}