Index: modules/Calorimeter.cc
===================================================================
--- modules/Calorimeter.cc	(revision e2dd4c504d097cdae44190e4a00f9d7a1b57b9c6)
+++ modules/Calorimeter.cc	(revision 00e8dcac59e89e1843c78ef81f47775e780edf2c)
@@ -56,12 +56,19 @@
 Calorimeter::Calorimeter() :
   fECalResolutionFormula(0), fHCalResolutionFormula(0),
-  fItParticleInputArray(0), fItTrackInputArray(0),
-  fTowerTrackArray(0), fItTowerTrackArray(0)
+  fItParticleInputArray(0), fItTrackInputArray(0)
 {
+  Int_t i;
+
   fECalResolutionFormula = new DelphesFormula;
   fHCalResolutionFormula = new DelphesFormula;
 
-  fTowerTrackArray = new TObjArray;
-  fItTowerTrackArray = fTowerTrackArray->MakeIterator();
+  for(i = 0; i < 2; ++i)
+  {
+    fECalTowerTrackArray[i] = new TObjArray;
+    fItECalTowerTrackArray[i] = fECalTowerTrackArray[i]->MakeIterator();
+
+    fHCalTowerTrackArray[i] = new TObjArray;
+    fItHCalTowerTrackArray[i] = fHCalTowerTrackArray[i]->MakeIterator();
+  }
 }
 
@@ -70,9 +77,17 @@
 Calorimeter::~Calorimeter()
 {
+  Int_t i;
+
   if(fECalResolutionFormula) delete fECalResolutionFormula;
   if(fHCalResolutionFormula) delete fHCalResolutionFormula;
 
-  if(fTowerTrackArray) delete fTowerTrackArray;
-  if(fItTowerTrackArray) delete fItTowerTrackArray;
+  for(i = 0; i < 2; ++i)
+  {
+    if(fECalTowerTrackArray[i]) delete fECalTowerTrackArray[i];
+    if(fItECalTowerTrackArray[i]) delete fItECalTowerTrackArray[i];
+
+    if(fHCalTowerTrackArray[i]) delete fHCalTowerTrackArray[i];
+    if(fItHCalTowerTrackArray[i]) delete fItHCalTowerTrackArray[i];
+  }
 }
 
@@ -203,4 +218,5 @@
   Double_t ecalFraction, hcalFraction;
   Double_t ecalEnergy, hcalEnergy;
+  Double_t ecalSigma, hcalSigma;
   Int_t pdgCode;
 
@@ -215,8 +231,8 @@
   DelphesFactory *factory = GetFactory();
   fTowerHits.clear();
-  fTowerECalFractions.clear();
-  fTowerHCalFractions.clear();
-  fTrackECalFractions.clear();
-  fTrackHCalFractions.clear();
+  fECalTowerFractions.clear();
+  fHCalTowerFractions.clear();
+  fECalTrackFractions.clear();
+  fHCalTrackFractions.clear();
 
   // loop over all particles
@@ -239,6 +255,6 @@
     hcalFraction = itFractionMap->second.second;
 
-    fTowerECalFractions.push_back(ecalFraction);
-    fTowerHCalFractions.push_back(hcalFraction);
+    fECalTowerFractions.push_back(ecalFraction);
+    fHCalTowerFractions.push_back(hcalFraction);
 
     if(ecalFraction < 1.0E-9 && hcalFraction < 1.0E-9) continue;
@@ -285,6 +301,6 @@
     hcalFraction = itFractionMap->second.second;
 
-    fTrackECalFractions.push_back(ecalFraction);
-    fTrackHCalFractions.push_back(hcalFraction);
+    fECalTrackFractions.push_back(ecalFraction);
+    fHCalTrackFractions.push_back(hcalFraction);
 
     // find eta bin [1, fEtaBins.size - 1]
@@ -349,14 +365,21 @@
       fTowerEdges[3] = (*phiBins)[phiBin];
 
-      fTowerECalEnergy = 0.0;
-      fTowerHCalEnergy = 0.0;
-
-      fTrackECalEnergy = 0.0;
-      fTrackHCalEnergy = 0.0;
+      fECalTowerEnergy = 0.0;
+      fHCalTowerEnergy = 0.0;
+
+      fECalTrackEnergy[0] = 0.0;
+      fECalTrackEnergy[1] = 0.0;
+
+      fHCalTrackEnergy[0] = 0.0;
+      fHCalTrackEnergy[1] = 0.0;
 
       fTowerTrackHits = 0;
       fTowerPhotonHits = 0;
 
-      fTowerTrackArray->Clear();
+      fECalTowerTrackArray[0]->Clear();
+      fECalTowerTrackArray[1]->Clear();
+
+      fHCalTowerTrackArray[0]->Clear();
+      fHCalTowerTrackArray[1]->Clear();
     }
 
@@ -370,9 +393,6 @@
       position = track->Position;
 
-      ecalEnergy = momentum.E() * fTrackECalFractions[number];
-      hcalEnergy = momentum.E() * fTrackHCalFractions[number];
-
-      fTrackECalEnergy += ecalEnergy;
-      fTrackHCalEnergy += hcalEnergy;
+      ecalEnergy = momentum.E() * fECalTrackFractions[number];
+      hcalEnergy = momentum.E() * fHCalTrackFractions[number];
 
       if(ecalEnergy > fTimingEnergyMin && fTower)
@@ -384,5 +404,32 @@
       }
 
-      fTowerTrackArray->Add(track);
+      if(fECalTrackFractions[number] > 1.0E-9 && fHCalTrackFractions[number] < 1.0E-9 )
+      {
+        ecalSigma = fECalResolutionFormula->Eval(0.0, fTowerEta, 0.0, momentum.E());
+        if(ecalSigma/momentum.E() < track->TrackResolution)
+        {
+          fECalTrackEnergy[0] += ecalEnergy;
+          fECalTowerTrackArray[0]->Add(track);
+        }
+        else
+        {
+          fECalTrackEnergy[1] += ecalEnergy;
+          fECalTowerTrackArray[1]->Add(track);
+        }
+      }
+      else if(fECalTrackFractions[number] < 1.0E-9 && fHCalTrackFractions[number] > 1.0E-9 )
+      {
+        hcalSigma = fHCalResolutionFormula->Eval(0.0, fTowerEta, 0.0, momentum.E());
+        if(hcalSigma/momentum.E() < track->TrackResolution)
+        {
+          fHCalTrackEnergy[0] += hcalEnergy;
+          fHCalTowerTrackArray[0]->Add(track);
+        }
+        else
+        {
+          fHCalTrackEnergy[1] += hcalEnergy;
+          fHCalTowerTrackArray[1]->Add(track);
+        }
+      }
 
       continue;
@@ -397,9 +444,9 @@
 
     // fill current tower
-    ecalEnergy = momentum.E() * fTowerECalFractions[number];
-    hcalEnergy = momentum.E() * fTowerHCalFractions[number];
-
-    fTowerECalEnergy += ecalEnergy;
-    fTowerHCalEnergy += hcalEnergy;
+    ecalEnergy = momentum.E() * fECalTowerFractions[number];
+    hcalEnergy = momentum.E() * fHCalTowerFractions[number];
+
+    fECalTowerEnergy += ecalEnergy;
+    fHCalTowerEnergy += hcalEnergy;
 
     if(ecalEnergy > fTimingEnergyMin && fTower)
@@ -426,8 +473,5 @@
   Double_t ecalEnergy, hcalEnergy;
   Double_t ecalSigma, hcalSigma;
-  Double_t ecalTrkSigma, hcalTrkSigma;
-  Double_t ecalFraction, hcalFraction;
-
-  Int_t pdgCode;
+
   TLorentzVector momentum;
   TFractionMap::iterator itFractionMap;
@@ -438,9 +482,9 @@
 
 
-  ecalSigma = fECalResolutionFormula->Eval(0.0, fTowerEta, 0.0, fTowerECalEnergy);
-  hcalSigma = fHCalResolutionFormula->Eval(0.0, fTowerEta, 0.0, fTowerHCalEnergy);
-
-  ecalEnergy = LogNormal(fTowerECalEnergy, ecalSigma);
-  hcalEnergy = LogNormal(fTowerHCalEnergy, hcalSigma);
+  ecalSigma = fECalResolutionFormula->Eval(0.0, fTowerEta, 0.0, fECalTowerEnergy);
+  hcalSigma = fHCalResolutionFormula->Eval(0.0, fTowerEta, 0.0, fHCalTowerEnergy);
+
+  ecalEnergy = LogNormal(fECalTowerEnergy, ecalSigma);
+  hcalEnergy = LogNormal(fHCalTowerEnergy, hcalSigma);
 
   ecalSigma = fECalResolutionFormula->Eval(0.0, fTowerEta, 0.0, ecalEnergy);
@@ -509,47 +553,50 @@
   // fill energy flow candidates
 
-  // save as eflowtracks only tracks that have better resolution than calo
-
-  fItTowerTrackArray->Reset();
-  while((track = static_cast<Candidate*>(fItTowerTrackArray->Next())))
-  {
-    pdgCode = TMath::Abs(track->PID);
-
-    itFractionMap = fFractionMap.find(pdgCode);
-    if(itFractionMap == fFractionMap.end())
-    {
-      itFractionMap = fFractionMap.find(0);
-    }
-
-    ecalFraction = itFractionMap->second.first;
-    hcalFraction = itFractionMap->second.second;
-
+  ecalEnergy -= fECalTrackEnergy[1];
+  hcalEnergy -= fHCalTrackEnergy[1];
+
+  fItECalTowerTrackArray[0]->Reset();
+  while((track = static_cast<Candidate*>(fItECalTowerTrackArray[0]->Next())))
+  {
     mother = track;
     track = static_cast<Candidate*>(track->Clone());
     track->AddCandidate(mother);
-    momentum = track->Momentum;
-
-    if(ecalFraction > 1.0E-9 && hcalFraction < 1.0E-9 )
-    {
-      ecalTrkSigma = fECalResolutionFormula->Eval(0.0, fTowerEta, 0.0, momentum.E());
-      if(ecalTrkSigma/momentum.E() < track->TrackResolution)
-      {
-        momentum *= ecalEnergy/fTrackECalEnergy;
-      }
-    }
-    else if(ecalFraction < 1.0E-9 && hcalFraction > 1.0E-9 )
-    {
-      hcalTrkSigma = fHCalResolutionFormula->Eval(0.0, fTowerEta, 0.0, momentum.E());
-      if(hcalTrkSigma/momentum.E() < track->TrackResolution)
-      {
-        momentum *= hcalEnergy/fTrackHCalEnergy;
-      }
-    }
+
+    track->Momentum *= ecalEnergy/fECalTrackEnergy[0];
 
     fEFlowTrackOutputArray->Add(track);
   }
 
-  ecalEnergy -= fTrackECalEnergy;
-  hcalEnergy -= fTrackHCalEnergy;
+  fItECalTowerTrackArray[1]->Reset();
+  while((track = static_cast<Candidate*>(fItECalTowerTrackArray[1]->Next())))
+  {
+    mother = track;
+    track = static_cast<Candidate*>(track->Clone());
+    track->AddCandidate(mother);
+
+    fEFlowTrackOutputArray->Add(track);
+  }
+
+  fItHCalTowerTrackArray[0]->Reset();
+  while((track = static_cast<Candidate*>(fItHCalTowerTrackArray[0]->Next())))
+  {
+    mother = track;
+    track = static_cast<Candidate*>(track->Clone());
+    track->AddCandidate(mother);
+
+    track->Momentum *= ecalEnergy/fECalTrackEnergy[0];
+
+    fEFlowTrackOutputArray->Add(track);
+  }
+
+  fItHCalTowerTrackArray[1]->Reset();
+  while((track = static_cast<Candidate*>(fItHCalTowerTrackArray[1]->Next())))
+  {
+    mother = track;
+    track = static_cast<Candidate*>(track->Clone());
+    track->AddCandidate(mother);
+
+    fEFlowTrackOutputArray->Add(track);
+  }
 
   ecalSigma = fECalResolutionFormula->Eval(0.0, fTowerEta, 0.0, ecalEnergy);
Index: modules/Calorimeter.h
===================================================================
--- modules/Calorimeter.h	(revision e2dd4c504d097cdae44190e4a00f9d7a1b57b9c6)
+++ modules/Calorimeter.h	(revision 00e8dcac59e89e1843c78ef81f47775e780edf2c)
@@ -57,6 +57,6 @@
   Candidate *fTower;
   Double_t fTowerEta, fTowerPhi, fTowerEdges[4];
-  Double_t fTowerECalEnergy, fTowerHCalEnergy;
-  Double_t fTrackECalEnergy, fTrackHCalEnergy;
+  Double_t fECalTowerEnergy, fHCalTowerEnergy;
+  Double_t fECalTrackEnergy[2], fHCalTrackEnergy[2];
 
   Double_t fTimingEnergyMin;
@@ -81,9 +81,9 @@
   std::vector < Long64_t > fTowerHits;
 
-  std::vector < Double_t > fTowerECalFractions;
-  std::vector < Double_t > fTowerHCalFractions;
+  std::vector < Double_t > fECalTowerFractions;
+  std::vector < Double_t > fHCalTowerFractions;
 
-  std::vector < Double_t > fTrackECalFractions;
-  std::vector < Double_t > fTrackHCalFractions;
+  std::vector < Double_t > fECalTrackFractions;
+  std::vector < Double_t > fHCalTrackFractions;
 
   DelphesFormula *fECalResolutionFormula; //!
@@ -103,6 +103,9 @@
   TObjArray *fEFlowNeutralHadronOutputArray; //!
 
-  TObjArray *fTowerTrackArray; //!
-  TIterator *fItTowerTrackArray; //!
+  TObjArray *fECalTowerTrackArray[2]; //!
+  TIterator *fItECalTowerTrackArray[2]; //!
+
+  TObjArray *fHCalTowerTrackArray[2]; //!
+  TIterator *fItHCalTowerTrackArray[2]; //!
 
   void FinalizeTower();
