DeepDoubleB · jmduarte · Apr 3, 2019 · May 23, 2019 · Oct 26, 2019 · Feb 4, 2020
diff --git a/NtupleAK8/python/DeepNtuplizerAK8_cfi.py b/NtupleAK8/python/DeepNtuplizerAK8_cfi.py
@@ -14,9 +14,9 @@
                                 genParticles = cms.InputTag("prunedGenParticles"),
                                 muons        = cms.InputTag("slimmedMuons"),
                                 electrons    = cms.InputTag("slimmedElectrons"),
-                                jetPtMin     = cms.untracked.double(200),
+                                jetPtMin     = cms.untracked.double(100),
                                 jetPtMax     = cms.untracked.double(-1),
-                                jetAbsEtaMax = cms.untracked.double(2.4),
+                                jetAbsEtaMax = cms.untracked.double(3),
 #                                 tagInfoName = cms.string('pfDeepCSV'),
                                 fjTagInfoName = cms.string('pfBoostedDoubleSVAK8'),
                                 bDiscriminators = cms.vstring(),

diff --git a/NtupleAK8/run/crab_status.py b/NtupleAK8/run/crab_status.py
@@ -4,15 +4,6 @@
 dirs = os.listdir("crab_projects")
 
 for d in dirs:
-	os.system('crab status crab_projects/'+d)
-	continue
-	o = os.popen('crab status crab_projects/'+d).read().split("\n")
-	for i, l in enumerate(o):
-		if l.startswith("CRAB project directory:"): print l
-		if l.startswith("Jobs status"):
-			for j in range(5):
-				if len(o[i+j]) < 2:continue
-				if  any(s in o[i+j] for s in ['unsubmitted', 'idle', 'finished','running','transferred', 'transferring', 'failed']):
-					print o[i+j]
+    os.system('crab status crab_projects/'+d)
 
 
diff --git a/NtupleAK8/run/samples_94x/Zqq.conf b/NtupleAK8/run/samples_94x/Zqq.conf
@@ -1,6 +1,7 @@
 fjKeepFlavors=0,2,3,4
 
 /ZJetsToQQ_HT-800toInf_TuneCP5_13TeV-madgraphMLM-pythia8/RunIIFall17MiniAODv2-PU2017_12Apr2018_94X_mc2017_realistic_v14-v1/MINIAODSIM
+/ZJetsToQQ_HT600to800_3j_TuneCP5_13TeV-madgraphMLM-pythia8/RunIIFall17MiniAODv2-PU2017_12Apr2018_94X_mc2017_realistic_v14-v1/MINIAODSIM
 /ZJetsToQQ_HT400to600_TuneCP5_13TeV-madgraphMLM-pythia8/RunIIFall17MiniAODv2-PU2017_12Apr2018_94X_mc2017_realistic_v14-v1/MINIAODSIM
 
 
diff --git a/NtupleAK8/run/samples_94x/multimass.conf b/NtupleAK8/run/samples_94x/multimass.conf
@@ -0,0 +1,8 @@
+fjKeepFlavors=0,2,3,4
+
+/BulkGravitonToHHTo4Q_MX-600to6000_MH-15to250_part1_TuneCP5_13TeV-madgraph_pythia8/RunIIFall17MiniAODv2-PU2017_12Apr2018_multigridpack_94X_mc2017_realistic_v14-v1/MINIAODSIM
+/BulkGravitonToHHTo4Q_MX-600to6000_MH-15to250_part2_TuneCP5_13TeV-madgraph_pythia8/RunIIFall17MiniAODv2-PU2017_12Apr2018_multigridpack_94X_mc2017_realistic_v14-v1/MINIAODSIM
+/BulkGravitonToHHTo4Q_MX-600to6000_MH-15to250_part3_TuneCP5_13TeV-madgraph_pythia8/RunIIFall17MiniAODv2-PU2017_12Apr2018_multigridpack_94X_mc2017_realistic_v14-v1/MINIAODSIM
+
+
+
diff --git a/NtupleAK8/run/samples_94x/qcd.conf b/NtupleAK8/run/samples_94x/qcd.conf
@@ -1,16 +1,6 @@
 fjKeepFlavors=0,2,3,4
 
-#/QCD_Pt_300to470_TuneCP5_13TeV_pythia8/RunIIFall17MiniAOD-94X_mc2017_realistic_v10-v1/MINIAODSIM
-#/QCD_Pt_470to600_TuneCP5_13TeV_pythia8/RunIIFall17MiniAOD-94X_mc2017_realistic_v10-v1/MINIAODSIM
-#/QCD_Pt_600to800_TuneCP5_13TeV_pythia8/RunIIFall17MiniAOD-94X_mc2017_realistic_v10-v1/MINIAODSIM
-#/QCD_Pt_800to1000_TuneCP5_13TeV_pythia8/RunIIFall17MiniAOD-94X_mc2017_realistic_v10-v2/MINIAODSIM
-#/QCD_Pt_1000to1400_TuneCP5_13TeV_pythia8/RunIIFall17MiniAOD-94X_mc2017_realistic_v10-v1/MINIAODSIM
-#/QCD_Pt_1400to1800_TuneCP5_13TeV_pythia8/RunIIFall17MiniAOD-94X_mc2017_realistic_v10-v1/MINIAODSIM
-#/QCD_Pt_1800to2400_TuneCP5_13TeV_pythia8/RunIIFall17MiniAOD-94X_mc2017_realistic_v10-v1/MINIAODSIM
-#/QCD_Pt_2400to3200_TuneCP5_13TeV_pythia8/RunIIFall17MiniAOD-94X_mc2017_realistic_v10-v2/MINIAODSIM
-#/QCD_Pt_3200toInf_TuneCP5_13TeV_pythia8/RunIIFall17MiniAOD-94X_mc2017_realistic_v10-v1/MINIAODSIM
-
-#/QCD_Pt_170to300_TuneCP5_13TeV_pythia8/RunIIFall17MiniAODv2-PU2017_12Apr2018_94X_mc2017_realistic_v14-v1/MINIAODSIM
+/QCD_Pt_170to300_TuneCP5_13TeV_pythia8/RunIIFall17MiniAODv2-PU2017_12Apr2018_94X_mc2017_realistic_v14-v1/MINIAODSIM
 /QCD_Pt_300to470_TuneCP5_13TeV_pythia8/RunIIFall17MiniAODv2-PU2017_12Apr2018_94X_mc2017_realistic_v14-v1/MINIAODSIM
 /QCD_Pt_470to600_TuneCP5_13TeV_pythia8/RunIIFall17MiniAODv2-PU2017_12Apr2018_94X_mc2017_realistic_v14-v1/MINIAODSIM
 /QCD_Pt_600to800_TuneCP5_13TeV_pythia8/RunIIFall17MiniAODv2-PU2017_12Apr2018_94X_mc2017_realistic_v14-v1/MINIAODSIM

diff --git a/NtupleAK8/run/status.py b/NtupleAK8/run/status.py
@@ -4,11 +4,17 @@
 dirs = os.listdir("crab_projects")
 
 for d in dirs:
-	if "4B" not in d: continue
+	if "MX" not in d: continue
+	#if "1000to1400" in d: continue
+	#if "Inf" in d: continue
 	#os.system('crab resubmit crab_projects/'+d+" --sitewhitelist=T2_DE_RWTH")
-	os.system('crab resubmit crab_projects/'+d)
+	#os.system('crab resubmit crab_projects/'+d)
+#	continue
+	print 'crab status crab_projects/'+d
+	os.system('crab status crab_projects/'+d)
+	#o = os.popen('crab status crab_projects/'+d).read().split("\n")
+	print o
 	continue
-	o = os.popen('crab status crab_projects/'+d).read().split("\n")
 	for i, l in enumerate(o):
 		if l.startswith("CRAB project directory:"): print l
 		if l.startswith("Jobs status"):

diff --git a/NtupleAK8/run/template_runCrab.py b/NtupleAK8/run/template_runCrab.py
@@ -23,15 +23,17 @@
 #config.Data.unitsPerJob = 200
 #config.Data.outLFNDirBase = '/store/user/%s/_outLFNDirBase_' % (getUsernameFromSiteDB())
 config.Data.outLFNDirBase = '_outLFNDirBase_'
+config.Data.allowNonValidInputDataset = True
+config.Data.publication = False
 
 # uncomment this part to use CRAB to submit to FNAL_LPC
 # config.Data.ignoreLocality = True
 # config.Site.whitelist = ['T3_US_FNALLPC']
 # config.Site.ignoreGlobalBlacklist = True
 
 # Run in DESY/RWTH
-#config.Site.whitelist = ['T2_CH_CERN', 'T2_DE_DESY', 'T2_DE_RWTH']
-config.Site.whitelist = ['T2_DE_RWTH']
+config.Site.whitelist = ['T2_CH_CERN', 'T2_DE_DESY', 'T2_DE_RWTH']
+#config.Site.whitelist = ['T2_DE_RWTH']
 #config.Site.ignoreGlobalBlacklist = True
 #config.Site.blacklist = ['T3_US_UMiss']
 config.User.voGroup = 'dcms'
@@ -40,9 +42,10 @@
 #config.Site.storageSite = '_storageSite_'
 config.Site.storageSite = "T2_DE_RWTH"
 
+config.JobType.allowUndistributedCMSSW = True
 
 # uncomment this part to use CRAB to submit to FNAL LPC
-#config.Data.ignoreLocality = True
+config.Data.ignoreLocality = True
 #config.Site.whitelist = ['T3_US_FNALLPC']
 #config.Site.ignoreGlobalBlacklist = True
 
@@ -51,4 +54,4 @@
 #config.Site.whitelist = ['T2_CH_CERN', 'T2_DE_DESY', 'T2_DE_RWTH']
 #config.Site.ignoreGlobalBlacklist = True
 
-config.Site.storageSite = '_storageSite_'
+#config.Site.storageSite = '_storageSite_'
diff --git a/NtupleAK8/src/FatJetInfoFiller.cc b/NtupleAK8/src/FatJetInfoFiller.cc
@@ -113,6 +113,19 @@ void FatJetInfoFiller::book() {
   //double-b
   data.add<float>("fj_doubleb", 0);
 
+  // DeepAK8, DDX 
+  data.add<float>("fj_DDBvL", 0);
+  data.add<float>("fj_DDCvL", 0);
+  data.add<float>("fj_DDCvB", 0);
+  data.add<float>("fj_DDBvL_mi", 0);
+  data.add<float>("fj_DDCvL_mi", 0);
+  data.add<float>("fj_DDCvB_mi", 0);
+
+  data.add<float>("fj_DeepAK8bbvslight_mi", 0);
+  data.add<float>("fj_DeepAK8ccvslight_mi", 0);
+  data.add<float>("fj_DeepAK8HbbvsQCD_mi", 0);
+  data.add<float>("fj_DeepAK8HccvsQCD_mi", 0);
+
   //flavor info
   data.add<int>("fj_isBB", 0);
   data.add<int>("fj_isNonBB", 0);
@@ -273,6 +286,17 @@ bool FatJetInfoFiller::fill(const pat::Jet& jet, size_t jetidx, const JetHelper&
   const auto &vars = bdsvTagInfo->taggingVariables();
 
   data.fill<float>("fj_doubleb", jet.bDiscriminator("pfBoostedDoubleSecondaryVertexAK8BJetTags"));
+  data.fill<float>("fj_DDBvL", jet.bDiscriminator("pfDeepDoubleBvLJetTags:probHbb"));
+  data.fill<float>("fj_DDCvL", jet.bDiscriminator("pfDeepDoubleCvLJetTags:probHcc"));
+  data.fill<float>("fj_DDCvB", jet.bDiscriminator("pfDeepDoubleCvBJetTags:probHcc"));
+  data.fill<float>("fj_DDBvL_mi", jet.bDiscriminator("pfMassIndependentDeepDoubleBvLJetTags:probHbb"));
+  data.fill<float>("fj_DDCvL_mi", jet.bDiscriminator("pfMassIndependentDeepDoubleCvLJetTags:probHcc"));
+  data.fill<float>("fj_DDCvB_mi", jet.bDiscriminator("pfMassIndependentDeepDoubleCvBJetTags:probHcc"));
+
+  data.fill<float>("fj_DeepAK8bbvslight_mi", jet.bDiscriminator("pfMassDecorrelatedDeepBoostedDiscriminatorsJetTags:bbvsLight"));
+  data.fill<float>("fj_DeepAK8ccvslight_mi", jet.bDiscriminator("pfMassDecorrelatedDeepBoostedDiscriminatorsJetTags:ccvsLight"));
+  data.fill<float>("fj_DeepAK8HbbvsQCD_mi", jet.bDiscriminator("pfMassDecorrelatedDeepBoostedDiscriminatorsJetTags:ZHbbvsQCD"));
+  data.fill<float>("fj_DeepAK8HccvsQCD_mi", jet.bDiscriminator("pfMassDecorrelatedDeepBoostedDiscriminatorsJetTags:ZHccvsQCD"));
 
   //flavor info
   data.fill<int>("fj_Hflavour", jet.hadronFlavour());  

diff --git a/NtupleAK8/src/JetInfoFillerAK8.cc b/NtupleAK8/src/JetInfoFillerAK8.cc
@@ -13,9 +13,9 @@
 namespace deepntuples {
 
 void JetInfoFillerAK8::readConfig(const edm::ParameterSet& iConfig, edm::ConsumesCollector && cc) {
-  minPt_ = iConfig.getUntrackedParameter<double>("jetPtMin", 150);
+  minPt_ = iConfig.getUntrackedParameter<double>("jetPtMin", 50);
   maxPt_ = iConfig.getUntrackedParameter<double>("jetPtMax", -1);
-  maxAbsEta_ = iConfig.getUntrackedParameter<double>("jetAbsEtaMax", 2.4);
+  maxAbsEta_ = iConfig.getUntrackedParameter<double>("jetAbsEtaMax", 3);
   btag_discriminators_ = iConfig.getParameter<std::vector<std::string>>("bDiscriminators");
 
   vtxToken_ = cc.consumes<reco::VertexCollection>(iConfig.getParameter<edm::InputTag>("vertices"));

diff --git a/NtupleAK8/src/PFCandidateFiller.cc b/NtupleAK8/src/PFCandidateFiller.cc
@@ -6,6 +6,7 @@
  */
 
 #include "DeepNTuples/NtupleAK8/interface/PFCandidateFiller.h"
+#include "DeepNTuples/NtupleCommons/interface/sorting_modules.h"
 
 namespace deepntuples {
 
@@ -60,16 +61,51 @@ void PFCandidateFiller::book() {
 
 bool PFCandidateFiller::fill(const pat::Jet& jet, size_t jetidx, const JetHelper& jet_helper) {
 
-  const auto& jetConstituents = jet_helper.getJetConstituents();
+  std::vector<const pat::PackedCandidate*> neutralPFCands;
+  std::unordered_map<const pat::PackedCandidate*, double> drMinSvMap;
+  //std::unordered_map<const pat::PackedCandidate*, double> ptRelMap;
 
-  int nConstituents = 0;
+  std::vector<sorting::sortingClass<size_t>> sortedneutral;
+
+  const float jet_uncorr_pt=jet.correctedJet("Uncorrected").pt();
+
+  unsigned int i = 0;
+  for (const auto * pfcand : jet_helper.getJetConstituents()){
+    if (!pfcand) continue;
+    if (pfcand->pt() < minPt_) continue;
+    if (pfcand->charge() == 0) {
+      neutralPFCands.push_back(pfcand);
+      drMinSvMap[pfcand];
+      double minDR = 0.8;
+      for (const auto &sv : *SVs){
+        double dr = reco::deltaR(*pfcand, sv);
+        if (dr < minDR) minDR = dr;
+      }
+      drMinSvMap[pfcand] = minDR;
+      sortedneutral.push_back(sorting::sortingClass<size_t>
+			      (i, -1,
+			       -minDR, pfcand->pt()/jet_uncorr_pt));
+      i++;
+    }
+  }
+
+  // sort by ABCInv
+  std::sort(sortedneutral.begin(),sortedneutral.end(),sorting::sortingClass<size_t>::compareByABCInv);
+  std::vector<size_t> sortedneutralindices;
+  sortedneutralindices=sorting::invertSortingVector(sortedneutral);
+
+  data.fill<int>("n_pfcands", neutralPFCands.size());
+  data.fill<float>("npfcands", neutralPFCands.size());
 
   float etasign = jet.eta()>0 ? 1 : -1;
 
-  for (const auto *pfcand : jetConstituents){
+
+  for (i = 0; i < neutralPFCands.size(); i++){
+
+    const auto *pfcand = neutralPFCands.at(sortedneutral.at(i).get());    
+    //  const auto *pfcand = neutralPFCands.at(i);
 
     if (pfcand->pt() < minPt_) continue;
-    nConstituents += 1;
     // basic kinematics, valid for both charged and neutral
     data.fillMulti<float>("pfcand_ptrel", pfcand->pt()/jet.pt());
     data.fillMulti<float>("pfcand_erel", pfcand->energy()/jet.energy());
@@ -79,12 +115,12 @@ bool PFCandidateFiller::fill(const pat::Jet& jet, size_t jetidx, const JetHelper
     data.fillMulti<float>("pfcand_puppiw", pfcand->puppiWeight());
     data.fillMulti<float>("pfcand_mass", pfcand->mass());
 
-    double minDR = 999;
+    double minDR = 0.8;
     for (const auto &sv : *SVs){
       double dr = reco::deltaR(*pfcand, sv);
       if (dr < minDR) minDR = dr;
     }
-    data.fillMulti<float>("pfcand_drminsv", minDR==999 ? -1 : minDR);
+    data.fillMulti<float>("pfcand_drminsv", catchInfsAndBound(drMinSvMap.at(pfcand),0,-0.8,0,-0.8));
 
     const auto& subjets = jet_helper.getSubJets();
     data.fillMulti<float>("pfcand_drsubjet1", subjets.size()>0 ? reco::deltaR(*pfcand, *subjets.at(0)) : -1);
@@ -118,9 +154,6 @@ bool PFCandidateFiller::fill(const pat::Jet& jet, size_t jetidx, const JetHelper
     }
 
   }
-  data.fill<int>("n_pfcands", nConstituents);
-  data.fill<float>("npfcands", nConstituents);
-
 
   return true;
 }

diff --git a/NtupleAK8/test/DeepNtuplizerAK8.py b/NtupleAK8/test/DeepNtuplizerAK8.py
@@ -8,17 +8,18 @@
 options.outputFile = 'output.root'
 options.inputFiles = [
 #'file:85820ACA-657B-BC44-AC74-AACD6D54B348.root'
-'root://cmsxrootd.fnal.gov//store/relval/CMSSW_10_3_0_pre3/RelValTTbar_13/MINIAODSIM/103X_upgrade2018_realistic_v4-v1/20000/D3C5A8CB-17C6-8548-B133-177A42C8B012.root'
+#'root://cmsxrootd.fnal.gov//store/relval/CMSSW_10_3_0_pre3/RelValTTbar_13/MINIAODSIM/103X_upgrade2018_realistic_v4-v1/20000/D3C5A8CB-17C6-8548-B133-177A42C8B012.root'
 #'root://cmsxrootd.fnal.gov//store/mc/RunIISummer17MiniAOD/BulkGravTohhTohbbhbb_narrow_M-1000_13TeV-madgraph/MINIAODSIM/92X_upgrade2017_realistic_v10-v2/90000/04DF3196-3B99-E711-AE12-008CFAC93CE8.root'
 #'root://cmsxrootd.fnal.gov//store/mc/RunIIFall17MiniAODv2/GluGluToBulkGravitonToHHTo4C_M-1000_narrow_13TeV-madgraph-pythia8/MINIAODSIM/PU2017_12Apr2018_94X_mc2017_realistic_v14-v1/10000/86032B22-7943-E811-AF6F-D0946626135C.root'
 #'root://cmsxrootd.fnal.gov//store/mc/RunIIFall17MiniAOD/GluGluHToBB_M125_13TeV_powheg_pythia8/MINIAODSIM/94X_mc2017_realistic_v10-v1/20000/C8932584-5006-E811-9840-141877410512.root'
 #'root://cmsxrootd.fnal.gov//store/mc/RunIIFall17MiniAOD/GluGluToBulkGravitonToHHTo4C_M-1000_narrow_13TeV-madgraph-pythia8/MINIAODSIM/PU2017_94X_mc2017_realistic_v11-v1/90000/0820742A-B729-E811-AC15-24BE05C6E711.root', 'root://cmsxrootd.fnal.gov//store/mc/RunIIFall17MiniAOD/GluGluToBulkGravitonToHHTo4C_M-1000_narrow_13TeV-madgraph-pythia8/MINIAODSIM/PU2017_94X_mc2017_realistic_v11-v1/90000/3290B112-B129-E811-9DE9-008CFAF55422.root'
+#'/store/mc/RunIIAutumn18MiniAOD/GluGluToBulkGravitonToHHTo4B_M-1000_narrow_TuneCP5_PSweights_13TeV-madgraph_pythia8/MINIAODSIM/102X_upgrade2018_realistic_v15-v1/20000/4FF5E1EB-773A-4B42-9BFE-30C27EFDF3F3.root'
 #'root://cmsxrootd.fnal.gov//store/mc/RunIIFall17MiniAODv2/GluGluToBulkGravitonToHHTo4C_M-1000_narrow_13TeV-madgraph-pythia8/MINIAODSIM/PU2017_12Apr2018_94X_mc2017_realistic_v14-v1/10000/86032B22-7943-E811-AF6F'
 #'root://cmsxrootd.fnal.gov//store/mc/RunIISummer16MiniAODv2/BulkGravTohhTohbbhbb_narrow_M-2500_13TeV-madgraph/MINIAODSIM/PUMoriond17_80X_mcRun2_asymptotic_2016_TrancheIV_v6-v1/80000/0A83E4E2-34B6-E611-89A0-549F35AE4FA2.root',
 #'file:/eos/user/a/anovak/022C3683-D4AB-E611-AC4D-3417EBE70078.root'  #include file: for local files, for catalogues /store..
+'/store/mc/RunIIFall17MiniAODv2/GluGluHToCC_M125_13TeV_powheg_pythia8/MINIAODSIM/PU2017_12Apr2018_94X_mc2017_realistic_v14-v2/30000/72164088-CB67-E811-9D0D-008CFA197AC4.root'
 ]
-options.maxEvents = -1
-#options.maxEvents = 1000
+options.maxEvents = -1#10000
 
 options.register('inputScript', '', VarParsing.multiplicity.singleton, VarParsing.varType.string, "input Script")
 options.register('skipEvents', 0, VarParsing.multiplicity.singleton, VarParsing.varType.int, "skip N events")
@@ -42,7 +43,7 @@
 process = cms.Process("DNNFiller")
 
 process.load('FWCore.MessageService.MessageLogger_cfi')
-process.MessageLogger.cerr.FwkReport.reportEvery = 10       
+process.MessageLogger.cerr.FwkReport.reportEvery = 1000       
 if not options.inputScript:  # this is probably for testing
 	process.MessageLogger.cerr.FwkReport.reportEvery = 10
 
@@ -63,6 +64,9 @@
     skipEvents=cms.untracked.uint32(options.skipEvents)
 )
 
+#process.source.eventsToProcess = cms.untracked.VEventRange("1:3127296-1:3127297")
+
+
 if options.inputScript:
     process.load(options.inputScript)
 
@@ -108,7 +112,20 @@
 	'pfDeepCSVJetTags:probbb',
 	 #'pfDeepCSVJetTags:probcc',
 
-	'pfBoostedDoubleSecondaryVertexAK8BJetTags'
+	'pfBoostedDoubleSecondaryVertexAK8BJetTags',
+
+
+      'pfDeepDoubleBvLJetTags:probHbb',
+      'pfDeepDoubleCvLJetTags:probHcc',
+      'pfDeepDoubleCvBJetTags:probHcc',
+      'pfMassIndependentDeepDoubleBvLJetTags:probHbb',
+      'pfMassIndependentDeepDoubleCvLJetTags:probHcc',
+      'pfMassIndependentDeepDoubleCvBJetTags:probHcc',
+
+    "pfMassDecorrelatedDeepBoostedDiscriminatorsJetTags:bbvsLight",
+    "pfMassDecorrelatedDeepBoostedDiscriminatorsJetTags:ccvsLight",
+    "pfMassDecorrelatedDeepBoostedDiscriminatorsJetTags:ZHccvsQCD",
+    "pfMassDecorrelatedDeepBoostedDiscriminatorsJetTags:ZHbbvsQCD"
 ]
 
 jetCorrectionsAK8 = ('AK8PFchs', ['L1FastJet', 'L2Relative', 'L3Absolute'], 'None')

diff --git a/makeLocalSampleLists.py b/makeLocalSampleLists.py
@@ -1,18 +1,28 @@
 import os
+import argparse
 
-dirs = ['Hcc', 'Hbb', "qcd"]
-dirs = ['Z']
+parser = argparse.ArgumentParser()
+parser.add_argument("-a", '--train', help="Number of files for training dataset", type=int)
+parser.add_argument("-b", '--test', help="Number of files for testing dataset", type=int)
+args = parser.parse_args()
+
+# dirs = ['Hcc', 'Hbb', "qcd"]
+dirs = ['qcd']
+#dirs = ['multi']
+maind = '/net/scratch_cms3a/novak/102REv3/'
 for d in dirs:
-	tests = os.listdir(d+"/test")
-	f = open("list_test_{}.txt".format(d), "w")
-	for line in tests:
-		f.write(d+"/test/"+line+"\n")
-	f.close()
-
-	trains = os.listdir(d)
-	f = open("list_{}.txt".format(d), "w")
-	for line in trains:
-		if not line.endswith(".root"): continue
-		f.write(d+"/"+line+"\n")
-	f.close()
-
+    flist = os.listdir(maind + d)
+    flist = [f for f in flist if f.endswith(".root")]
+    ntot = len(flist)
+    assert ntot > args.train + args.test, "Not enough files for specified split"
+
+    f = open("list_{}.txt".format(d), "w")
+    for line in flist[:args.train]:
+        f.write(maind + d + "/" + line + "\n")
+    f.close()
+
+    f = open("list_test_{}.txt".format(d), "w")
+    for line in flist[args.train : args.train+args.test]:
+        f.write(maind + d + "/" + line + "\n")
+    f.close()
+