CRAB/python/LumiList.py

#!/usr/bin/env python

"""
Handle lists of lumi sections. Constuct in several different formats and filter
(mask) a secondary list of lumis.
This class can also handle ranges of events as the structure is identical
or could be subclassed renaming a function or two.
"""

__revision__ = "$Id: LumiList.py,v 1.3 2010/05/04 17:24:58 spiga Exp $"
__version__ = "$Revision: 1.3 $"

try: # FUTURE: Python 2.6, prior to 2.6 requires simplejson
    import json
except ImportError:
    import simplejson as json

class LumiList(object):
    """
    Deal with lists of lumis in several different forms:
    Compact list:
        {
        '1': [[1, 33], [35, 35], [37, 47], [49, 75], [77, 130], [133, 136]],
        '2':[[1,45],[50,80]]
        }
        where the first key is the run number, subsequent pairs are
        ranges of lumis within that run that are desired
    Runs and lumis:
        {
        '1': [1,2,3,4,6,7,8,9,10],
        '2': [1,4,5,20]
        }
        where the first key is the run number and the list is a list of
        individual lumi sections
    Run  lumi pairs:
        [[1,1], [1,2],[1,4], [2,1], [2,5], [1,10]]
        where each pair in the list is an individual run&lumi
    CMSSW representation:
        '1:1-1:33,1:35,1:37-1:47,2:1-2:45,2:50-2:80'
        The string used by CMSSW in lumisToProcess or lumisToSkip
        is a subset of the compactList example above
    """


    def __init__(self, filename = None, lumis = None, runsAndLumis = None, runs = None):
        """
        Constructor takes filename (JSON), a list of run/lumi pairs,
        or a dict with run #'s as the keys and a list of lumis as the values, or just a list of runs
        """
        self.compactList = {}
        if filename:
            self.filename = filename
            jsonFile = open(self.filename,'r')
            self.compactList = json.load(jsonFile)
        elif lumis:
            runsAndLumis = {}
            for (run, lumi) in lumis:
                run = str(run)
                if not runsAndLumis.has_key(run):
                    runsAndLumis[run] = []
                runsAndLumis[run].append(lumi)

        if runsAndLumis:
            for run in runsAndLumis.keys():
                runString = str(run)
                lastLumi = -1000
                self.compactList[runString] = []
                lumiList = runsAndLumis[run]
                for lumi in sorted(lumiList):
                    if lumi == lastLumi:
                        pass # Skip duplicates
                    elif lumi != lastLumi + 1: # Break in lumi sequence
                        self.compactList[runString].append([lumi, lumi])
                    else:
                        nRange =  len(self.compactList[runString])
                        self.compactList[runString][nRange-1][1] = lumi
                    lastLumi = lumi
        if runs:
            for run in runs:
                runString = str(run)
                self.compactList[runString] = [[1,0xFFFFFFF]]

    def filterLumis(self, lumiList):
        """
        Return a list of lumis that are in compactList.
        lumilist is of the simple form
        [(run1,lumi1),(run1,lumi2),(run2,lumi1)]
        """
        filteredList = []
        for (run, lumi) in lumiList:
            runsInLumi = self.compactList.get(str(run), [[0,-1]])
            for (first, last) in runsInLumi:
                if lumi >= first and lumi <= last:
                    filteredList.append((run, lumi))
                    break
        return filteredList


    def getCompactList(self):
        """
        Return the compact list representation
        """
        return self.compactList


    def getLumis(self):
        """
        Return the list of pairs representation
        """
        theList = []
        runs = self.compactList.keys()
        runs.sort(key=int)
        for run in runs:
            lumis = self.compactList[run]
            for lumiPair in sorted(lumis):
                for lumi in range(lumiPair[0], lumiPair[1]+1):
                    theList.append((int(run), lumi))

        return theList


    def getCMSSWString(self):
        """
        Turn compactList into a list of the format
        R1:L1,R2:L2-R2:L3 which is acceptable to CMSSW LumiBlockRange variable
        """

        parts = []
        runs = self.compactList.keys()
        runs.sort(key=int)
        for run in runs:
            lumis = self.compactList[run]
            for lumiPair in sorted(lumis):
                if lumiPair[0] == lumiPair[1]:
                    parts.append("%s:%s" % (run, lumiPair[0]))
                else:
                    parts.append("%s:%s-%s:%s" %
                                 (run, lumiPair[0], run, lumiPair[1]))

        output = ','.join(parts)
        return output


# Unit test code
'''
import unittest
from WMCore.DataStructs.LumiList import LumiList


class LumiListTest(unittest.TestCase):
    """
    _LumiListTest_

    """

    def testRead(self):
        """
        Test reading from JSON
        """
        exString = "1:1-1:33,1:35,1:37-1:47,2:49-2:75,2:77-2:130,2:133-2:136"
        exDict   = {'1': [[1, 33], [35, 35], [37, 47]],
                    '2': [[49, 75], [77, 130], [133, 136]]}

        jsonList = LumiList(filename = 'lumiTest.json')
        lumiString = jsonList.getCMSSWString()
        lumiList = jsonList.getCompactList()

        self.assertTrue(lumiString == exString)
        self.assertTrue(lumiList   == exDict)

    def testList(self):
        """
        Test constucting from list of pairs
        """

        listLs1 = range(1, 34) + [35] + range(37, 48)
        listLs2 = range(49, 76) + range(77, 131) + range(133, 137)
        lumis = zip([1]*100, listLs1) + zip([2]*100, listLs2)

        jsonLister = LumiList(filename = 'lumiTest.json')
        jsonString = jsonLister.getCMSSWString()
        jsonList = jsonLister.getCompactList()

        pairLister = LumiList(lumis = lumis)
        pairString = pairLister.getCMSSWString()
        pairList = pairLister.getCompactList()

        self.assertTrue(jsonString == pairString)
        self.assertTrue(jsonList   == pairList)


    def testRuns(self):
        """
        Test constucting from run and list of lumis
        """
        runsAndLumis = {
            1: range(1, 34) + [35] + range(37, 48),
            2: range(49, 76) + range(77, 131) + range(133, 137)
        }
        runsAndLumis2 = {
            '1': range(1, 34) + [35] + range(37, 48),
            '2': range(49, 76) + range(77, 131) + range(133, 137)
        }

        jsonLister = LumiList(filename = 'lumiTest.json')
        jsonString = jsonLister.getCMSSWString()
        jsonList   = jsonLister.getCompactList()

        runLister = LumiList(runsAndLumis = runsAndLumis)
        runString = runLister.getCMSSWString()
        runList   = runLister.getCompactList()

        runLister2 = LumiList(runsAndLumis = runsAndLumis2)
        runList2 = runLister2.getCompactList()

        self.assertTrue(jsonString == runString)
        self.assertTrue(jsonList   == runList)
        self.assertTrue(runList2   == runList)


    def testFilter(self):
        """
        Test filtering of a list of lumis
        """
        runsAndLumis = {
            1: range(1, 34) + [35] + range(37, 48),
            2: range(49, 76) + range(77, 131) + range(133, 137)
        }

        completeList = zip([1]*150, range(1, 150)) + \
                       zip([2]*150, range(1, 150)) + \
                       zip([3]*150, range(1, 150))

        smallList    = zip([1]*50,  range(1, 10)) + zip([2]*50, range(50, 70))
        overlapList  = zip([1]*150, range(30, 40)) + \
                       zip([2]*150, range(60, 80))
        overlapRes   = zip([1]*9,   range(30, 34)) + [(1, 35)] + \
                       zip([1]*9,   range(37, 40)) + \
                       zip([2]*30,  range(60, 76)) + \
                       zip([2]*9,   range(77, 80))

        runLister = LumiList(runsAndLumis = runsAndLumis)

        # Test a list to be filtered which is a superset of constructed list
        filterComplete = runLister.filterLumis(completeList)
        # Test a list to be filtered which is a subset of constructed list
        filterSmall    = runLister.filterLumis(smallList)
        # Test a list to be filtered which is neither
        filterOverlap  = runLister.filterLumis(overlapList)

        self.assertTrue(filterComplete == runLister.getLumis())
        self.assertTrue(filterSmall    == smallList)
        self.assertTrue(filterOverlap  == overlapRes)

    def testDuplicates(self):
        """
        Test a list with lots of duplicates
        """
        result = zip([1]*100, range(1, 34) + range(37, 48))
        lumis  = zip([1]*100, range(1, 34) + range(37, 48) + range(5, 25))

        lister = LumiList(lumis = lumis)
        self.assertTrue(lister.getLumis() == result)

    def testNull(self):
        """
        Test a null list
        """

        runLister = LumiList(lumis = None)

        self.assertTrue(runLister.getCMSSWString() == '')
        self.assertTrue(runLister.getLumis() == [])
        self.assertTrue(runLister.getCompactList() == {})


if __name__ == '__main__':
    unittest.main()
'''

# Test JSON file

#{"1": [[1, 33], [35, 35], [37, 47]], "2": [[49, 75], [77, 130], [133, 136]]}
Revision:	1.4
Committed:	Wed May 26 19:46:12 2010 UTC (14 years, 11 months ago) by ewv
Content type:	text/x-python
Branch:	MAIN
CVS Tags:	CRAB_2_7_3, CRAB_2_7_3_pre3, CRAB_2_7_3_pre3_beta, CRAB_2_7_3_pre2, CRAB_2_7_3_pre2_beta, CRAB_2_7_3_pre1, CRAB_2_7_3_beta3, CRAB_2_7_3_beta2, CRAB_2_7_3_beta1, CRAB_2_7_3_beta
Changes since 1.3:	+9 -6 lines
Log Message:	Add runselection along with lumi_mask. Needs Splitter, LumiList, and DataDiscovery
#	User	Rev	Content
1	ewv	1.1	#!/usr/bin/env python
2
3			"""
4			Handle lists of lumi sections. Constuct in several different formats and filter
5			(mask) a secondary list of lumis.
6			This class can also handle ranges of events as the structure is identical
7			or could be subclassed renaming a function or two.
8			"""
9
10	ewv	1.4	__revision__ = "$Id: LumiList.py,v 1.3 2010/05/04 17:24:58 spiga Exp $"
11			__version__ = "$Revision: 1.3 $"
12	ewv	1.1
13			try: # FUTURE: Python 2.6, prior to 2.6 requires simplejson
14			import json
15			except ImportError:
16			import simplejson as json
17
18			class LumiList(object):
19			"""
20			Deal with lists of lumis in several different forms:
21			Compact list:
22			{
23			'1': [[1, 33], [35, 35], [37, 47], [49, 75], [77, 130], [133, 136]],
24			'2':[[1,45],[50,80]]
25			}
26			where the first key is the run number, subsequent pairs are
27			ranges of lumis within that run that are desired
28			Runs and lumis:
29			{
30			'1': [1,2,3,4,6,7,8,9,10],
31			'2': [1,4,5,20]
32			}
33			where the first key is the run number and the list is a list of
34			individual lumi sections
35			Run lumi pairs:
36			[[1,1], [1,2],[1,4], [2,1], [2,5], [1,10]]
37			where each pair in the list is an individual run&lumi
38			CMSSW representation:
39			'1:1-1:33,1:35,1:37-1:47,2:1-2:45,2:50-2:80'
40			The string used by CMSSW in lumisToProcess or lumisToSkip
41			is a subset of the compactList example above
42			"""
43
44
45	ewv	1.4	def __init__(self, filename = None, lumis = None, runsAndLumis = None, runs = None):
46	ewv	1.1	"""
47			Constructor takes filename (JSON), a list of run/lumi pairs,
48	ewv	1.4	or a dict with run #'s as the keys and a list of lumis as the values, or just a list of runs
49	ewv	1.1	"""
50			self.compactList = {}
51			if filename:
52			self.filename = filename
53			jsonFile = open(self.filename,'r')
54			self.compactList = json.load(jsonFile)
55			elif lumis:
56			runsAndLumis = {}
57			for (run, lumi) in lumis:
58			run = str(run)
59			if not runsAndLumis.has_key(run):
60			runsAndLumis[run] = []
61			runsAndLumis[run].append(lumi)
62
63			if runsAndLumis:
64			for run in runsAndLumis.keys():
65			runString = str(run)
66			lastLumi = -1000
67			self.compactList[runString] = []
68			lumiList = runsAndLumis[run]
69			for lumi in sorted(lumiList):
70			if lumi == lastLumi:
71			pass # Skip duplicates
72			elif lumi != lastLumi + 1: # Break in lumi sequence
73			self.compactList[runString].append([lumi, lumi])
74			else:
75			nRange = len(self.compactList[runString])
76			self.compactList[runString][nRange-1][1] = lumi
77			lastLumi = lumi
78	ewv	1.4	if runs:
79			for run in runs:
80			runString = str(run)
81			self.compactList[runString] = [[1,0xFFFFFFF]]
82	ewv	1.1
83			def filterLumis(self, lumiList):
84			"""
85			Return a list of lumis that are in compactList.
86			lumilist is of the simple form
87			[(run1,lumi1),(run1,lumi2),(run2,lumi1)]
88			"""
89			filteredList = []
90			for (run, lumi) in lumiList:
91	ewv	1.2	runsInLumi = self.compactList.get(str(run), [[0,-1]])
92	ewv	1.1	for (first, last) in runsInLumi:
93			if lumi >= first and lumi <= last:
94			filteredList.append((run, lumi))
95			break
96			return filteredList
97
98
99			def getCompactList(self):
100			"""
101			Return the compact list representation
102			"""
103			return self.compactList
104
105
106			def getLumis(self):
107			"""
108			Return the list of pairs representation
109			"""
110			theList = []
111			runs = self.compactList.keys()
112			runs.sort(key=int)
113			for run in runs:
114			lumis = self.compactList[run]
115			for lumiPair in sorted(lumis):
116			for lumi in range(lumiPair[0], lumiPair[1]+1):
117			theList.append((int(run), lumi))
118
119			return theList
120
121
122			def getCMSSWString(self):
123			"""
124			Turn compactList into a list of the format
125			R1:L1,R2:L2-R2:L3 which is acceptable to CMSSW LumiBlockRange variable
126			"""
127
128			parts = []
129			runs = self.compactList.keys()
130			runs.sort(key=int)
131			for run in runs:
132			lumis = self.compactList[run]
133			for lumiPair in sorted(lumis):
134			if lumiPair[0] == lumiPair[1]:
135			parts.append("%s:%s" % (run, lumiPair[0]))
136			else:
137			parts.append("%s:%s-%s:%s" %
138			(run, lumiPair[0], run, lumiPair[1]))
139
140			output = ','.join(parts)
141			return output
142
143	ewv	1.2
144
145	ewv	1.1	# Unit test code
146			'''
147			import unittest
148			from WMCore.DataStructs.LumiList import LumiList
149
150
151			class LumiListTest(unittest.TestCase):
152			"""
153			_LumiListTest_
154
155			"""
156
157			def testRead(self):
158			"""
159			Test reading from JSON
160			"""
161			exString = "1:1-1:33,1:35,1:37-1:47,2:49-2:75,2:77-2:130,2:133-2:136"
162			exDict = {'1': [[1, 33], [35, 35], [37, 47]],
163			'2': [[49, 75], [77, 130], [133, 136]]}
164
165			jsonList = LumiList(filename = 'lumiTest.json')
166			lumiString = jsonList.getCMSSWString()
167			lumiList = jsonList.getCompactList()
168
169			self.assertTrue(lumiString == exString)
170			self.assertTrue(lumiList == exDict)
171
172			def testList(self):
173			"""
174			Test constucting from list of pairs
175			"""
176
177			listLs1 = range(1, 34) + [35] + range(37, 48)
178			listLs2 = range(49, 76) + range(77, 131) + range(133, 137)
179			lumis = zip([1]100, listLs1) + zip([2]100, listLs2)
180
181			jsonLister = LumiList(filename = 'lumiTest.json')
182			jsonString = jsonLister.getCMSSWString()
183			jsonList = jsonLister.getCompactList()
184
185			pairLister = LumiList(lumis = lumis)
186			pairString = pairLister.getCMSSWString()
187			pairList = pairLister.getCompactList()
188
189			self.assertTrue(jsonString == pairString)
190			self.assertTrue(jsonList == pairList)
191
192
193			def testRuns(self):
194			"""
195			Test constucting from run and list of lumis
196			"""
197			runsAndLumis = {
198			1: range(1, 34) + [35] + range(37, 48),
199			2: range(49, 76) + range(77, 131) + range(133, 137)
200			}
201			runsAndLumis2 = {
202			'1': range(1, 34) + [35] + range(37, 48),
203			'2': range(49, 76) + range(77, 131) + range(133, 137)
204			}
205
206			jsonLister = LumiList(filename = 'lumiTest.json')
207			jsonString = jsonLister.getCMSSWString()
208			jsonList = jsonLister.getCompactList()
209
210			runLister = LumiList(runsAndLumis = runsAndLumis)
211			runString = runLister.getCMSSWString()
212			runList = runLister.getCompactList()
213
214			runLister2 = LumiList(runsAndLumis = runsAndLumis2)
215			runList2 = runLister2.getCompactList()
216
217			self.assertTrue(jsonString == runString)
218			self.assertTrue(jsonList == runList)
219			self.assertTrue(runList2 == runList)
220
221
222			def testFilter(self):
223			"""
224			Test filtering of a list of lumis
225			"""
226			runsAndLumis = {
227			1: range(1, 34) + [35] + range(37, 48),
228			2: range(49, 76) + range(77, 131) + range(133, 137)
229			}
230
231			completeList = zip([1]*150, range(1, 150)) + \
232	ewv	1.2	zip([2]*150, range(1, 150)) + \
233	ewv	1.4	zip([3]*150, range(1, 150))
234	ewv	1.2
235	ewv	1.1	smallList = zip([1]50, range(1, 10)) + zip([2]50, range(50, 70))
236			overlapList = zip([1]*150, range(30, 40)) + \
237			zip([2]*150, range(60, 80))
238			overlapRes = zip([1]*9, range(30, 34)) + [(1, 35)] + \
239			zip([1]*9, range(37, 40)) + \
240			zip([2]*30, range(60, 76)) + \
241			zip([2]*9, range(77, 80))
242
243			runLister = LumiList(runsAndLumis = runsAndLumis)
244
245			# Test a list to be filtered which is a superset of constructed list
246			filterComplete = runLister.filterLumis(completeList)
247			# Test a list to be filtered which is a subset of constructed list
248			filterSmall = runLister.filterLumis(smallList)
249			# Test a list to be filtered which is neither
250			filterOverlap = runLister.filterLumis(overlapList)
251
252			self.assertTrue(filterComplete == runLister.getLumis())
253			self.assertTrue(filterSmall == smallList)
254			self.assertTrue(filterOverlap == overlapRes)
255
256			def testDuplicates(self):
257			"""
258			Test a list with lots of duplicates
259			"""
260			result = zip([1]*100, range(1, 34) + range(37, 48))
261			lumis = zip([1]*100, range(1, 34) + range(37, 48) + range(5, 25))
262
263			lister = LumiList(lumis = lumis)
264			self.assertTrue(lister.getLumis() == result)
265
266			def testNull(self):
267			"""
268			Test a null list
269			"""
270
271			runLister = LumiList(lumis = None)
272
273			self.assertTrue(runLister.getCMSSWString() == '')
274			self.assertTrue(runLister.getLumis() == [])
275			self.assertTrue(runLister.getCompactList() == {})
276
277
278
279			if __name__ == '__main__':
280			unittest.main()
281			'''
282
283			# Test JSON file
284
285	ewv	1.2	#{"1": [[1, 33], [35, 35], [37, 47]], "2": [[49, 75], [77, 130], [133, 136]]}