lcsim/src/org/lcsim/recon/util

diff -N RemoveUnusedMcParticles.java
--- /dev/null	1 Jan 1970 00:00:00 -0000
+++ RemoveUnusedMcParticles.java	4 Jul 2012 13:54:36 -0000	1.1
@@ -0,0 +1,137 @@

+package org.lcsim.recon.util;
+
+import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.List;
+
+import hep.physics.event.generator.MCEvent;
+
+import org.lcsim.event.EventHeader;
+import org.lcsim.event.MCParticle;
+import org.lcsim.event.SimCalorimeterHit;
+import org.lcsim.event.SimTrackerHit;
+import org.lcsim.util.Driver;
+
+/**
+ * Driver to filter and remove all mc particles from a given collection that did
+ * not contribute to any hit. All generator level mc particles will be kept by
+ * default. In addition, all ancestors of the mc particles that are kept, are
+ * kept as well.
+ * 
+ * @author <a href="mailto:[log in to unmask]">Christian Grefe</a>
+ * 
+ */
+public class RemoveUnusedMcParticles extends Driver {
+
+	protected String mcParticleCollection;
+	protected List<String> ignoredCollections;
+	protected boolean keepAncestors;
+	protected boolean keepGeneratorParticles;
+
+	public RemoveUnusedMcParticles() {
+		mcParticleCollection = MCEvent.MC_PARTICLES;
+		ignoredCollections = new ArrayList<String>();
+		keepAncestors = true;
+		keepGeneratorParticles = true;
+	}
+
+	public void setMcParticleCollection(String mcParticleCollection) {
+		this.mcParticleCollection = mcParticleCollection;
+	}
+
+	public void setKeepAncestors(boolean keepAncestors) {
+		this.keepAncestors = keepAncestors;
+	}
+
+	public void setKeepGeneratorParticles(boolean keepGeneratorParticles) {
+		this.keepGeneratorParticles = keepGeneratorParticles;
+	}
+
+	/**
+	 * Mc particles that contributed to the given list of collections will not
+	 * be kept.
+	 * 
+	 * @param ignoredCollections
+	 */
+	public void setIgnoredCollections(String[] ignoredCollections) {
+		this.ignoredCollections = Arrays.asList(ignoredCollections);
+	}
+
+	@Override
+	protected void process(EventHeader event) {
+
+		// get the mc particles from the event
+		List<MCParticle> mcParticles = event.get(MCParticle.class,
+				mcParticleCollection);
+
+		// create a new list for the mc particles to keep and always keep the
+		// generator level particles
+		List<MCParticle> mcParticlesToKeep = new ArrayList<MCParticle>();
+
+		if (keepGeneratorParticles) {
+			for (MCParticle mcp : mcParticles) {
+				int genStatus = mcp.getGeneratorStatus();
+				if (genStatus == MCParticle.FINAL_STATE
+						|| genStatus == MCParticle.INTERMEDIATE) {
+					addMcParticleWithParents(mcp, mcParticlesToKeep, keepAncestors);
+				}
+			}
+		}
+
+		// keep all mc particles that contributed to tracker hit
+		List<List<SimTrackerHit>> simTrackerHitCollections = event
+				.get(SimTrackerHit.class);
+		for (List<SimTrackerHit> simTrackerHits : simTrackerHitCollections) {
+			String collectionName = event.getMetaData(simTrackerHits).getName();
+			if (!ignoredCollections.contains(collectionName)) {
+				for (SimTrackerHit hit : simTrackerHits) {
+					addMcParticleWithParents(hit.getMCParticle(),
+							mcParticlesToKeep, keepAncestors);
+				}
+			}
+		}
+
+		// keep all mc particles that contributed to any calorimeter hit
+		List<List<SimCalorimeterHit>> simCalorimeterHitCollections = event
+				.get(SimCalorimeterHit.class);
+		for (List<SimCalorimeterHit> simCalorimeterHits : simCalorimeterHitCollections) {
+			String collectionName = event.getMetaData(simCalorimeterHits)
+					.getName();
+			if (!ignoredCollections.contains(collectionName)) {
+				for (SimCalorimeterHit hit : simCalorimeterHits) {
+					for (int index = 0; index < hit.getMCParticleCount(); index++) {
+						addMcParticleWithParents(hit.getMCParticle(index),
+								mcParticlesToKeep, keepAncestors);
+					}
+				}
+			}
+		}
+		
+		// replace the list of mc particles in the event
+		int flags = event.getMetaData(mcParticles).getFlags();
+		event.remove(mcParticleCollection);
+		event.put(mcParticleCollection, mcParticlesToKeep, MCParticle.class, flags);
+	}
+
+	/**
+	 * Helper method to recursively add an mc particle and all of its ancestors
+	 * 
+	 * @param mcp
+	 *            the mc particle to add to the list
+	 * @param mcpList
+	 *            the list of mc particles to which the mc particle is added
+	 */
+	public static void addMcParticleWithParents(MCParticle mcp,
+			List<MCParticle> mcpList, boolean keepAncestors) {
+		if (!mcpList.contains(mcp)) {
+			mcpList.add(mcp);
+			List<MCParticle> parents = mcp.getParents();
+			if (keepAncestors) {
+				for (MCParticle parent : parents) {
+					addMcParticleWithParents(parent, mcpList, keepAncestors);
+				}
+			}
+		}
+	}
+
+}