forked from vonSchlotzkow/CFSR2power
-
Notifications
You must be signed in to change notification settings - Fork 0
/
CFSRwrapper.py
168 lines (154 loc) · 5.68 KB
/
CFSRwrapper.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
import pygrib
import re
from pylab import vector_lengths
import numpy
from optparse import OptionParser
def filenamefromfield(field,year,month,lowres=False,basedir=""):
if lowres:
lowres=".l."
else:
lowres="."
return "%s%i%02i/%s%sgdas.%i%02i.grb2" % (basedir,year,month,field,lowres,year,month)
def deepcopydatatolist(c):
return map(lambda x: x.getdata() , c)
def ReduceToVecLengths(mcl):
"""
Calculate absoulute value of vectors stored in a list of wrapped
grib messages
Danger Will Robinson: This function has side effects on the first
element of the list!
However, if it is used as 'Reduce' function in the iterator, noone
will ever see the original element, so we don't care about side
effects.
"""
ret=mcl[0]
ret.putdata(vector_lengths(numpy.array(deepcopydatatolist(mcl)),axis=0))
return [ret]
class messagecontainer(object):
grbmsg=None
data=None
def __init__(self,grbmsg):
self.grbmsg=grbmsg
self.data=self.grbmsg.values.copy()
def __repr__(self):
return self.grbmsg.__repr__()
def getdata(self):
return self.data.copy()
def data2grbmsg(self):
self.grbmsg.values=self.data.copy()
def putdata(self,d):
self.data=d.copy()
self.data2grbmsg()
class CFSRwrapper(pygrib.open):
"""Iterator for traversing grbs file skipping spin-up timestamps"""
recpertimestep=None
spinup=None
nonspinup=None
instant=None
unaverage=None
fieldnametoparam={
'dswsfc':(1,1,6,False,True,None),
'tmp2m':(1,1,6,True,False,None),
'wnd10m':(2,1,6,True,False,ReduceToVecLengths),
}
_previousdata=None
_currentdata=None
_Reduce=None
def __init__(self, fname,recpertimestep=None,spinup=None,nonspinup=None,instant=None,unaverage=None,autoconf=True,Reduce=None):
if autoconf:
assert(recpertimestep==None)
assert(spinup==None)
assert(nonspinup==None)
assert(instant==None)
assert(unaverage==None)
assert(Reduce==None)
fieldname=re.match("(|.*/)([^\./]+)(\.l\.|\.)gdas\..*grb2",fname).groups()[1]
(self.recpertimestep,self.spinup,self.nonspinup,self.instant,self.unaverage,self._Reduce)=self.fieldnametoparam[fieldname]
else:
self.recpertimestep=recpertimestep
self.spinup=spinup
self.nonspinup=nonspinup
self.instant=instant
self.unaverage=unaverage
self._Reduce=Reduce
pygrib.open.__init__(self,fname)
def __new__(cls, fname, *args, **kwargs):
obj=pygrib.open.__new__(cls, fname)
return obj
def __iter__(self):
self.rewind()
return self
def messagestep(self):
return self.messagenumber/self.recpertimestep % ((self.spinup + self.nonspinup))
def isspinupstep(self):
return self.spinup and (self.messagestep() == 0)
def read(self,N=None):
return map(messagecontainer,pygrib.open.read(self,N))
def step(self):
"""read the next step of the raw data series
If 'unaverage' reverse the averaging, and return instantaneous
data. Unaveraging is done according to
$$ \tilde{x}_T = \frac{1}{T} \sum_{t=1}^{T} x_t \\
= \sum_{t=1}^{T-1} \frac{T-1}{T} \frac{x_t}{T-1} + \frac{x_T}{T} \\
= \frac{T-1}{T} \tilde{x}_{T-1} + \frac{x_T}{T}\\
x_{t} = T\tilde{x}_{T} -(T-1)\tilde{x}_{T-1} $$
where $x_t$ is the instantaneous data and $\tilde{x}_t$ is the
averaged data at time $t$.
"""
if self.unaverage:
T=self.messagestep()
if T <= 1: self._previousdata=None
currentdata=self.read(self.recpertimestep)
currentdata[0].stepType='instant'
self._currentdata=deepcopydatatolist(currentdata)
if self._previousdata:
ret=currentdata
for p,c,r in zip(self._previousdata, self._currentdata, ret):
r.putdata(T*c - (T-1)*p)
self._previousdata=self._currentdata
return ret
else:
#first step after spin up
self._previousdata=self._currentdata
return currentdata
else:
currentdata=self.read(self.recpertimestep)
return currentdata
def next(self):
if self.messagenumber + self.recpertimestep > self.messages:
raise StopIteration
if self.isspinupstep():
self.step()
if self._Reduce:
return self._Reduce(self.step())
else:
return self.step()
def openfields(infields,year,month,lowres=False):
import itertools
r=[]
for f in infields:
r.append(CFSRwrapper(filenamefromfield(f,year,month,lowres)))
return itertools.izip(*r)
def unpackandapply(i,conv,outf):
rnp=conv(map(lambda x:x[0].data,i))
rmc=i[0][0]
rmc.putdata(rnp)
# changing the units is not that easy, apparently the units have to match a 'concept'
#rmc.grbmsg.units='MW'
outf.write(rmc.grbmsg.tostring())
return rmc
def iterateandapply(it,conv,outf):
assert(outf.mode=='wb')
for i in it:
rmc=unpackandapply(i,conv,outf)
parser = OptionParser()
parser.add_option("--year", dest="year", default=2000, type=int,
help="year")
parser.add_option("--month", dest="month", default=1, type=int,
help="day of month")
parser.add_option("--lowres",
action="store_true", dest="lowres",
help="convert low resolution data")
parser.add_option("--debug",
action="store_true", dest="debug",
help="run in debug mode")