source: TI03-DataExtractor/branches/titania_install/bin/splitGrib.py @ 1520

Subversion URL: http://proj.badc.rl.ac.uk/svn/ndg/TI03-DataExtractor/branches/titania_install/bin/splitGrib.py@1709
Revision 1520, 17.2 KB checked in by astephen, 14 years ago (diff)

This is the live version on titania - changes have been made so safest to SVN it.

  • Property svn:executable set to *
Line 
1#!/usr/bin/env python
2
3"""
4splitGrib.py
5============
6
7Script for splitting up Grib files based on parameter, level, date, hour
8and/or time.
9
10Usage:
11======
12
13  splitGrib.py [-p] [-d] [-h] [-f] [-l] [-a] [-r] [-t <filename_template>] <gribfile1> [<gribfile2> etc]
14
15Where:
16======
17
18        -p      split all records by parameter
19        -l      split all records by level
20        -d      split all records by date
21        -h      split all records by hour
22        -f      split all records by forecast step
23        -a      split all records by all of the above
24        -r      read all files in before splitting and treat as one logical file.
25        <filename_template>  is a filename template following the syntax as follows:
26                :PARAM: - parameter code
27                :CHAR_PARAM: - paramater mapped to character code
28                :LEVEL: - level code
29                :DATE: - date (8 digits)
30                :YEAR: - year (4 digits)
31                :MONTH: - month (2 digits)
32                :DAY: - day (2 digits)
33                :HOUR: - hour (2 digits)
34                :FCSTEP: - forecast step (n-digits)
35                :FCSTEP1: - forecast step (1 digit)
36                :FCSTEP2: - forecast step (2 digits)
37                :FCSTEP3: - forecast step (3 digits)
38
39                E.g. You might specify a <filename_template> as:
40                "/badc/ecmwf-e40/data/gg/as/:YEAR:/:MONTH:/:DAY:/ggas:DATE::HOUR::CHAR_PARAM:.grb"
41                This might result in a file called:
42                "/badc/ecmwf-e40/data/gg/as/1999/01/01/ggas1999010112tcc.grb"
43               
44               
45"""
46
47# import standard library modules
48import os, sys, struct, getopt, re
49
50# Set up global variables
51gribCharCodeDict={27:"lvc", 28:"hvc", 29:"lvt", 30:"hvt", 31:"ci", 32:"asn", 33:"rsn", 34:"sstk", 35:"istl1", 36:"istl2", 37:"istl3", 38:"istl4", 39:"swvl1", 40:"swvl2", 41:"swvl3", 42:"swvl4", 43:"slt", 44:"es", 45:"smlt", 47:"dsrp", 49:"10fg", 50:"lspf", 51:"mx2t24", 52:"mn2t24", 53:"mont", 54:"pres", 60:"pv", 127:"at", 128:"bv", 129:"z", 130:"t", 131:"u", 132:"v", 133:"q", 134:"sp", 135:"w", 136:"tcw", 137:"tcwv", 138:"vo", 139:"stl1", 140:"swl1", 141:"sd", 142:"lsp", 143:"cp", 144:"sf", 145:"bld", 146:"sshf", 147:"slhf", 148:"chnk", 149:"snr", 150:"tnr", 151:"msl", 152:"lnsp", 153:"swhr", 154:"lwhr", 155:"d", 156:"gh", 157:"r", 158:"tsp", 159:"blh", 160:"sdor", 161:"isor", 162:"anor", 163:"slor", 164:"tcc", 165:"10u", 166:"10v", 167:"2t", 168:"2d", 169:"ssrd", 170:"stl2", 171:"swl2", 172:"lsm", 173:"sr", 174:"al", 175:"strd", 176:"ssr", 177:"str", 178:"tsr", 179:"ttr", 180:"ewss", 181:"nsss", 182:"e", 183:"stl3", 184:"swl3", 185:"ccc", 186:"lcc", 187:"mcc", 188:"hcc", 189:"sund", 190:"ewov", 191:"nsov", 192:"nwov", 193:"neov", 194:"btmp", 195:"lgws", 196:"mgws", 197:"gwd", 198:"src", 199:"veg", 200:"vso", 201:"mx2t", 202:"mn2t", 203:"o3", 204:"paw", 205:"ro", 206:"tco3", 207:"10si", 208:"tsrc", 209:"ttrc", 210:"ssrc", 211:"strc", 212:"si", 214:"dhr", 215:"dhvd", 216:"dhcc", 217:"dhlc", 218:"vdzw", 219:"vdmw", 220:"ewgd", 221:"nsgd", 222:"ctzw", 223:"ctmw", 224:"vdh", 225:"htcc", 226:"htlc", 227:"crnh", 228:"tp", 229:"iews", 230:"inss", 231:"ishf", 232:"ie", 233:"asq", 234:"lsrh", 235:"skt", 236:"stl4", 237:"swl4", 238:"tsn", 239:"csf", 240:"lsf", 241:"acf", 242:"alw", 243:"fal", 244:"fsr", 245:"flsr", 246:"clwc", 247:"ciwc", 248:"cc", 249:"aiw", 250:"ice", 251:"atte", 252:"athe", 253:"atze"}
52
53
54# Deprecated test stuff...
55import StringIO
56inf=StringIO.StringIO("""14987 Data
5727395 Diti
5893749 Doto
5993745 Dyti
6093742 Doti
6193752 Daty
6210294 Ditq
63""")
64
65def exitNicely(msg):
66    """
67    Exits nicely with a useful error message.
68    """
69    print __doc__
70    print msg
71    sys.exit()
72
73class SplitGrib:
74    """
75    Class container for grib splitting system.
76    """
77   
78    def __init__(self, gribFiles, splitby=["all"], group="no", template=None):
79        """
80        Controls the process of splitting grib files into other files.
81        """
82        if "all" in splitby: splitby="all"       
83        self.fileDict={}
84        nrecs=0
85        outputFiles=[]
86       
87        for f in gribFiles:
88            morerecs=self._readFile(f)
89            nrecs=nrecs+morerecs
90           
91            if group=="no":
92                fileList=self._writeGribFiles(self.fileDict, splitby, template)
93                outputFiles=outputFiles+fileList
94                self.fileDict={}
95               
96        if group=="yes": # if all read in before splitting to output files
97            outputFiles=self._writeGribFiles(self.fileDict, splitby, template)
98                         
99        print "\nOutput files were:"
100        print "\t"+("\n\t".join(outputFiles))
101        print "\nWorked on %s grib records in total." % nrecs
102        print "Created %s output files in total." % len(outputFiles)
103
104       
105    def _readFile(self, gribFile):
106        """
107        Reads in a single GRIB file and adds the contents (metadata and
108        each record) to a dictionary in memory.
109        """
110        infile=open(gribFile, "rb")
111        print "Reading in:", gribFile
112
113        recCount=0
114        pos=0
115        for i in range(1, 1000000):
116            rt=self._getRecordDetails(infile, i, pos, gribFile)
117            if rt=="End of file":
118                break
119
120            (size, param, level, date, hour, fcStep)=rt
121            print (size, param, level, date, hour, fcStep)
122
123            pointer=self.fileDict
124           
125            for x in (param, level, date, hour, fcStep):
126                if not pointer.has_key(x): pointer[x]={}
127                pointer=pointer[x]
128            pointer["data"]=infile.read(size)
129            pos=pos+size
130            print pos
131            recCount=recCount+1           
132             
133        infile.close()
134        return recCount
135
136
137    def _getRecordDetails(self, fileHandle, gcount, pos, filename):
138        """
139        Reads the metadata from the next record in the GRIB file.
140        fileHandle is the open file handle, gcount is the GRIB record
141        number and pos is the current position in the file.
142
143        Returns tuple of: (param, level, date, hour, fcStep)
144       
145        Sections of interest are:
146          Section 0:
147            1-4 'GRIB'
148            5-7 size
149          Section 1:
150            9 Param and units
151            10 lev type
152            11-12 Height, pressure
153            13,14,15,16,17 year,month,day,hour,minute
154            18 forecast time unit
155            19 P1 -period of time 1
156            20 P2 - time between forecasts for averaging or accum 
157        """
158        # Read and check section 0
159        try:
160            GRIB=""
161            while GRIB!="GRIB":
162                fileHandle.seek(pos)
163                GRIB="".join(list(struct.unpack("cccc", fileHandle.read(4))))
164                pos=pos+1         
165           
166        except:
167            return "End of file"
168       
169        if GRIB!="GRIB":
170            exitNicely("Grib Record number %s could not be decoded as valid GRIB from file '%s'." % (gcount, filename))
171       
172        sizes=struct.unpack("BBB", fileHandle.read(3))
173
174        size=(sizes[0]<<16)+(sizes[1]<<8)+(sizes[2])
175        print size
176        # Now start reading section 1
177        fileHandle.read(9)
178       
179        param=struct.unpack("B", fileHandle.read(1))[0]
180        levtype=struct.unpack("B", fileHandle.read(1))[0]
181
182        level1=struct.unpack("B", fileHandle.read(1))[0]
183        level2=struct.unpack("B", fileHandle.read(1))[0]
184        level=level1*256+level2
185
186        year=struct.unpack("B", fileHandle.read(1))[0]
187        month=struct.unpack("B", fileHandle.read(1))[0]
188        day=struct.unpack("B", fileHandle.read(1))[0]
189        hour=struct.unpack("B", fileHandle.read(1))[0]
190        minute=struct.unpack("B", fileHandle.read(1))[0]
191       
192        if year<100:
193            if year<50:
194                year=2000+year
195            else:
196                year=1900+year
197               
198        date="%.4d%.2d%.2d" % (year,month,day)
199
200        fcTimeUnit=struct.unpack("B", fileHandle.read(1))[0]
201        fcStep=struct.unpack("B", fileHandle.read(1))[0]
202        fcStep2=struct.unpack("B", fileHandle.read(1))[0]     
203       
204        fileHandle.seek(pos-1)
205         
206        return (size, param, level, date, hour, fcStep)
207
208
209    def _generateOutputFileName(self, splitDict, template):
210        """
211        Returns an output file generated from a dictionary of values provided,
212        along with a filename template.
213        The rules are:
214        <filename_template>  is a filename template following the syntax as follows:
215                :PARAM: - parameter code
216                :CHAR_PARAM: - paramater mapped to character code
217                :LEVEL: - level code
218                :DATE: - date (8 digits)
219                :YEAR: - year (4 digits)
220                :MONTH: - month (2 digits)
221                :DAY: - day (2 digits)
222                :HOUR: - hour (2 digits)
223                :FCSTEP: - forecast step (n-digits)
224                :FCSTEP1: - forecast step (1 digit)
225                :FCSTEP2: - forecast step (2 digits)
226                :FCSTEP3: - forecast step (3 digits)
227
228                E.g. You might specify a <filename_template> as:
229                "/badc/ecmwf-e40/data/gg/as/:YEAR:/:MONTH:/:DAY:/ggas:DATE::HOUR::CHAR_PARAM:.grb"
230                This might result in a file called:
231                "/badc/ecmwf-e40/data/gg/as/1999/01/01/ggas1999010112tcc.grb"       
232        """
233        #paramMap={2:"blh", 9:"tcc", 1:"edd"}
234        itemList="PARAM CHAR_PARAM LEVEL DATE YEAR MONTH DAY HOUR FCSTEP FCSTEP1 FCSTEP2 FCSTEP3".split()
235        fname=template
236
237        for key in splitDict.keys():
238            if key in itemList:
239                match=re.search("(.*):%s:(.*)" % key, fname)
240                if match:
241                    value=splitDict[key]
242                    groups=match.groups()
243                    if key=="CHAR_PARAM":
244                        replacer=gribCharCodeDict[int(value)]
245                    elif len(key)>6 and key[:6]=="FCSTEP":
246                        intLen=int(key[6:])
247                        replacer=value[0:intLen]
248                    else:
249                        replacer=value
250                    fname=groups[0]+replacer+groups[1]
251
252        # Since you might get the same code more than once we need a recursive call
253        # if there is still a colon in the template
254        if fname.find(":")>-1:
255            fname=self._generateOutputFileName(splitDict, fname)
256
257        return fname
258
259
260    def _writeGribFiles(self, dataDict, splitby, template):
261        """
262        Writes output files from the data dictionary already built.
263        It uses the splitting rules to decide how to output them
264        and the template to generate appropriate output file paths.
265        """
266        d=dataDict
267        openFileDict={}
268        ofd=openFileDict
269
270        # Sort params
271        params=d.keys()
272        params.sort()
273       
274        for param in params:
275             pdict=d[param]
276
277             # Sorting levels
278             print "Check sorting of levels follows perl and existing data!"
279             levels=pdict.keys()
280             levels.sort()
281             
282             for level in levels:
283                 ldict=pdict[level]
284
285                 # Sorting dates
286                 dates=ldict.keys()
287                 dates.sort()
288                 
289                 for date in dates:
290                     ddict=ldict[date]
291
292                     # Sorting hours
293                     hours=ddict.keys()
294                     hours.sort()
295                     
296                     for hour in hours:
297                         hdict=ddict[hour]
298
299                         # Sorting fcSteps
300                         fcSteps=hdict.keys()
301                         fcSteps.sort()
302
303                         for fcStep in fcSteps:
304                             fdict=hdict[fcStep]
305                             record=fdict["data"]
306
307                             if splitby=="all":
308                                 if template:
309                                     splitDict={}
310                                     for arg in ("param", "level", "date", "hour", "fcStep"):
311                                         splitDict[arg.upper()]=eval(str(arg))
312                                     fname=self._generateOutputFileName(splitDict, template)
313                                 else:
314                                     fname="output_grib/%s.%s.%s.%s.%s.txt" % (param, level, date, hour, fcStep)
315                                     
316                             else:  # if only splitting by some arguments
317                                 splitters="param level date hour fcStep".split()
318                                 fname=""
319                                 if template:
320                                     splitDict={}
321                                     for arg in ("param", "level", "date", "hour", "fcStep"):
322                                         try:
323                                             value=eval(arg)
324                                             splitDict[arg.upper()]=str(value)                           
325                                         except:
326                                             pass
327                                         
328                                     fname=self._generateOutputFileName(splitDict, template)
329                                 else:
330                                     for splitter in splitters:
331                                         if splitter in splitby:
332                                             splitterString=eval(splitter)
333                                             fname=fname+"."+splitterString
334                                     fname="output_grib/%s.txt" % fname[1:]
335
336                                     
337
338                             if not ofd.has_key(fname):
339                                 ofd[fname]=open(fname, "wb")
340                             ofd[fname].write(record)
341
342        print "Closing all output files:"
343        outputFiles=ofd.keys()
344        outputFiles.sort()
345        for f in outputFiles:
346            print "\tClosing %s" % f
347            ofd[f].close()
348
349        return outputFiles 
350                                     
351                                 
352                                 
353
354                             
355if __name__=="__main__":
356
357    #sys.argv="splitGrib.py -p -d -t output_grib/:PARAM:.:DATE:.txt infile.grib".split()
358    #print "ARGS:", sys.argv
359    args=sys.argv[1:]
360    argList,gribFiles=getopt.getopt(args, "pldhfart:")
361   
362    splitby=[]
363    group="no"
364    template=None
365    for arg,value in argList:
366        if arg=="-a":
367            splitby.append("all")
368        elif arg=="-p":
369            splitby.append("param")
370        elif arg=="-l":
371            splitby.append("level")
372        elif arg=="-d":
373            splitby.append("date")
374        elif arg=="-h":
375            splitby.append("hour")
376        elif arg=="-f":
377            splitby.append("fcStep")
378        elif arg=="-r":
379            group="yes"
380        elif arg=="-t":
381            template=value
382
383    if splitby==[] or gribFiles==[]:
384        exitNicely("Please provide at least one split setting and at least one input grib file.")
385
386    SplitGrib(gribFiles, splitby, group, template)
387
388
389
390
391"""
392def getLevelAndParam(fileHandle, gcount, pos):
393    try:
394        GRIB=""
395        while GRIB!="GRIB":
396            fileHandle.seek(pos)
397            GRIB="".join(list(struct.unpack("cccc", fileHandle.read(4))))
398            pos=pos+1     
399       
400    except:
401        return "End of file"
402   
403    if GRIB!="GRIB":#"".join(GRIB)!="GRIB": 
404        exitNicely("Grib Record number %s could not be decoded as valid GRIB." % gcount)
405   
406    lengthSoFar=4
407    size=struct.unpack("H", fileHandle.read(2))[0]+ \
408         struct.unpack("B", fileHandle.read(1))[0]
409 
410    fileHandle.read(9)
411    param=struct.unpack("B", fileHandle.read(1))[0]
412    levtype=struct.unpack("B", fileHandle.read(1))[0]
413
414    level1=struct.unpack("B", fileHandle.read(1))[0]
415    level2=struct.unpack("B", fileHandle.read(1))[0]
416
417    level=level1*256+level2
418    fileHandle.seek(pos-1)
419   
420    return (level, param, size) 
421
422
423def filterGrib(inputFile, outputFile, levels, params):
424    print "Filtering '%s'..." % inputFile
425    infile=open(inputFile, "rb")
426    outfile=open(outputFile, "wb")
427    written=0
428   
429    pos=0
430    for i in range(1, 10000000):
431        rt=getLevelAndParam(infile, i, pos)
432        if rt=="End of file":
433            break
434        (level, param, size)=rt
435        #print (level, param, size)
436        keepp="no"
437        keepl="no"
438        if params!="all":
439            if params[0]=="range":
440                if param>=params[1] and param<=params[2]:
441                    keepp="yes"
442            elif param in params:
443                keepp="yes"
444        else:
445            keepp="yes"
446       
447        if levels!="all":
448            if levels[0]=="range":
449                if level>=levels[1] and level<=levels[2]:
450                    keepl="yes"
451            elif level in levels:
452                keepl="yes"
453        else:
454            keepl="yes"
455
456        data=infile.read(size)
457        pos=pos+size
458       
459        if keepl=="yes" and keepp=="yes":
460            outfile.write(data)
461            written=written+1
462                           
463    infile.close()
464    outfile.close()
465
466    print "Completed file: %s records found." % (i-1)
467    print "%s filtered records written to: %s" % (written, outputFile)
468   
469   
470
471"""
Note: See TracBrowser for help on using the repository browser.