Merge pull request #5 from rsholmes/ancestors

ancestors bank
maureeungaro · Jan 30, 2019 · dcea6f3 · dcea6f3
2 parents cb0729d + 7870bb5
commit dcea6f3
Show file tree

Hide file tree

Showing 3 changed files with 140 additions and 1 deletion.
diff --git a/src/banks.cc b/src/banks.cc
@@ -518,6 +518,108 @@ vector<generatedParticle> getGenerated(evioDOMTree& EDT, gBank bank, double verb
 }
 
 
+vector<ancestorInfo> getAncestors(evioDOMTree& EDT, gBank bank, double verbosity)
+{
+  vector<ancestorInfo> trajs;
+  vector<int> pid;
+  vector<int> tid;
+  vector<int> mtid;
+  vector<double> trackE;
+  vector<double> px;
+  vector<double> py;
+  vector<double> pz;
+  vector<double> vx;
+  vector<double> vy;
+  vector<double> vz;
+
+  unsigned long sizeOfBank = bank.name.size();
+
+  evioDOMNodeListP thisBankNode = EDT.getNodeList(tagNumEquals(bank.idtag, 0));
+  for(evioDOMNodeList::const_iterator iter=thisBankNode->begin(); iter!=thisBankNode->end(); iter++)
+    {
+      const evioDOMNodeP node = *iter;
+      if(node->isContainer())
+	{
+	  evioDOMNodeList *variablesNodes = node->getChildList();
+	  for(evioDOMNodeList::const_iterator cIter=variablesNodes->begin(); cIter!=variablesNodes->end(); cIter++)
+	    {
+	      const evioDOMNodeP variable = *cIter;
+	      int vnum = variable->num;
+	      for(unsigned i=0; i<sizeOfBank; i++)
+		{
+		  if(bank.gid[i] == vnum && variable->isLeaf() && bank.name[i] == "pid")
+		    {
+		      const vector<int> *vec = (*cIter)->getVector<int>();
+		      for(unsigned int h=0; h<vec->size(); h++) pid.push_back((*vec)[h]);
+		    }
+		  if(bank.gid[i] == vnum && variable->isLeaf() && bank.name[i] == "tid")
+		    {
+		      const vector<int> *vec = (*cIter)->getVector<int>();
+		      for(unsigned int h=0; h<vec->size(); h++) tid.push_back((*vec)[h]);
+		    }
+		  if(bank.gid[i] == vnum && variable->isLeaf() && bank.name[i] == "mtid")
+		    {
+		      const vector<int> *vec = (*cIter)->getVector<int>();
+		      for(unsigned int h=0; h<vec->size(); h++) mtid.push_back((*vec)[h]);
+		    }
+		  if(bank.gid[i] == vnum && variable->isLeaf() && bank.name[i] == "trackE")
+		    {
+		      const vector<double> *vec = (*cIter)->getVector<double>();
+		      for(unsigned int h=0; h<vec->size(); h++) trackE.push_back((*vec)[h]);
+		    }
+		  if(bank.gid[i] == vnum && variable->isLeaf() && bank.name[i] == "px")
+		    {
+		      const vector<double> *vec = (*cIter)->getVector<double>();
+		      for(unsigned int h=0; h<vec->size(); h++) px.push_back((*vec)[h]);
+		    }
+		  if(bank.gid[i] == vnum && variable->isLeaf() && bank.name[i] == "py")
+		    {
+		      const vector<double> *vec = (*cIter)->getVector<double>();
+		      for(unsigned int h=0; h<vec->size(); h++) py.push_back((*vec)[h]);
+		    }
+		  if(bank.gid[i] == vnum && variable->isLeaf() && bank.name[i] == "pz")
+		    {
+		      const vector<double> *vec = (*cIter)->getVector<double>();
+		      for(unsigned int h=0; h<vec->size(); h++) pz.push_back((*vec)[h]);
+		    }
+		  if(bank.gid[i] == vnum && variable->isLeaf() && bank.name[i] == "vx")
+		    {
+		      const vector<double> *vec = (*cIter)->getVector<double>();
+		      for(unsigned int h=0; h<vec->size(); h++) vx.push_back((*vec)[h]);
+		    }
+		  if(bank.gid[i] == vnum && variable->isLeaf() && bank.name[i] == "vy")
+		    {
+		      const vector<double> *vec = (*cIter)->getVector<double>();
+		      for(unsigned int h=0; h<vec->size(); h++) vy.push_back((*vec)[h]);
+		    }
+		  if(bank.gid[i] == vnum && variable->isLeaf() && bank.name[i] == "vz")
+		    {
+		      const vector<double> *vec = (*cIter)->getVector<double>();
+		      for(unsigned int h=0; h<vec->size(); h++) vz.push_back((*vec)[h]);
+		    }
+		}
+	    }
+	}
+    }
+
+  for(unsigned p=0; p<pid.size(); p++)
+    {
+      ancestorInfo part;
+
+      part.pid = pid[p];
+      part.tid = tid[p];
+      part.mtid = mtid[p];
+      part.trackE = trackE[p];
+      part.p      = G4ThreeVector(px[p], py[p], pz[p]);
+      part.vtx    = G4ThreeVector(vx[p], vy[p], vz[p]);
+
+      trajs.push_back(part);
+    }
+
+  return trajs;
+}
+
+
 
 
 

diff --git a/src/banks.h b/src/banks.h
@@ -29,6 +29,9 @@ vector<hitOutput> getDgtIntDataBank(evioDOMTree& EDT, string hitType, map<string
 // generated particle infos
 vector<generatedParticle> getGenerated(evioDOMTree& EDT, gBank bank, double verbosity);
 
+// ancestor info
+vector<ancestorInfo> getAncestors (evioDOMTree& EDT, gBank bank, double verbosity);
+
 
 map<string, vector<hitOutput> > getRawIntDataBanks(evioDOMTree& EDT, vector<string> hitTypes, map<string, gBank> *banksMap, double verbosity);
 map<string, vector<hitOutput> > getDgtIntDataBanks(evioDOMTree& EDT, vector<string> hitTypes, map<string, gBank> *banksMap, double verbosity);

diff --git a/src/evio2root.cc b/src/evio2root.cc
@@ -103,9 +103,20 @@ int main(int argc, char **argv)
 		rTrees["generated"].addVariable(banksMap["generated"].name[i], banksMap["generated"].type[i]);
 
 
+		// ancestors bank definitions
+		rTrees["ancestors"] = rTree(banksMap["ancestors"].bankName, banksMap["ancestors"].bdescription, verbosity);
+		for (unsigned i=0; i<banksMap["ancestors"].name.size(); i++)
+		  rTrees["ancestors"].addVariable(banksMap["ancestors"].name[i], banksMap["ancestors"].type[i]);
+
+
 		// hit banks definitions
 		for(map<string, gBank>::iterator it = banksMap.begin(); it != banksMap.end(); it++) {
-			if(it->first != "header" && it->first != "userHeader" && it->first != "generated" && it->first != "raws" && it->first != "psummary") {
+			if (it->first != "header" 
+			    && it->first != "userHeader" 
+			    && it->first != "generated" 
+			    && it->first != "ancestors" 
+			    && it->first != "raws" 
+			    && it->first != "psummary") {
 				rTrees[it->first] = rTree(banksMap[it->first].bankName, banksMap[it->first].bdescription, verbosity);
 				for(unsigned i=0; i<banksMap[it->first].name.size(); i++)
 				{
@@ -215,6 +226,29 @@ int main(int argc, char **argv)
 				}
 
 
+				// ancestors
+				else if (it->first == "ancestors") 
+				  {
+				    vector<ancestorInfo> trajs = getAncestors (EDT, getBankFromMap(it->first, &banksMap), verbosity);
+
+				    rTrees["ancestors"].init();
+
+				    for (unsigned i=0; i<banksMap["ancestors"].name.size(); i++) 
+				      {
+					string varname = banksMap["ancestors"].name[i];
+					string vartype = banksMap["ancestors"].type[i];
+
+
+					for(unsigned p=0; p<trajs.size(); p++) 
+					  {
+					    rTrees["ancestors"].insertVariable(varname, vartype, trajs[p].getVariableFromStringI(varname));
+					    rTrees["ancestors"].insertVariable(varname, vartype, trajs[p].getVariableFromStringD(varname));
+					  }
+				      }
+				    rTrees["ancestors"].fill();
+				  }
+
+
 				//  hit banks
 				else if(it->first != "psummary" && it->first != "raws") {