source: CMIP6dreqbuild/trunk/src/framework/dreqPy/volsum.py @ 839

Subversion URL: http://proj.badc.rl.ac.uk/svn/exarch/CMIP6dreqbuild/trunk/src/framework/dreqPy/volsum.py@839
Revision 839, 11.7 KB checked in by mjuckes, 3 years ago (diff)

more imports fixed

Line 
1
2import xlsxwriter
3from xlsxwriter.utility import xl_rowcol_to_cell
4import collections, os
5
6try:
7  import dreq
8
9  import scope
10  import makeTables
11  import overviewTabs
12except:
13  import dreqPy.dreq as dreq
14  import dreqPy.scope as scope
15  import dreqPy.makeTables as makeTables
16  import dreqPy.overviewTabs as overviewTabs
17
18###cell = xl_rowcol_to_cell(1, 2)  # C2
19 
20
21class xlsx(object):
22  def __init__(self,fn,txtOpts=None):
23    self.txtOpts = txtOpts
24    self.mcfgNote = 'Reference Volume (1 deg. atmosphere, 0.5 deg. ocean)'
25    self.wb = xlsxwriter.Workbook('%s.xlsx' % fn)
26    self.hdr_cell_format = self.wb.add_format({'text_wrap': True, 'font_size': 14, 'font_color':'#0000ff', 'bold':1, 'fg_color':'#aaaacc'})
27    self.hdr_cell_format.set_text_wrap()
28    self.sect_cell_format = self.wb.add_format({'text_wrap': True, 'font_size': 14, 'font_color':'#0000ff', 'bold':1, 'fg_color':'#ccccbb'})
29    self.sect_cell_format.set_text_wrap()
30    self.cell_format = self.wb.add_format({'text_wrap': True, 'font_size': 11})
31    self.cell_format.set_text_wrap()
32
33  def newSheet(self,name):
34    self.sht = self.wb.add_worksheet(name=name)
35    return self.sht
36
37  def tabrec(self,j,orec):
38        for i in range(len(orec)):
39          if orec[i] != '' and type( orec[i] ) == type( '' ) and orec[i][0] == '$':
40             self.sht.write_formula(j,i, '{=%s}' % orec[i][1:])
41          else:
42             if j == 0:
43               self.sht.write( j,i, orec[i], self.hdr_cell_format )
44             else:
45               self.sht.write( j,i, orec[i], self.cell_format )
46
47  def close(self):
48      self.wb.close()
49
50class vsum(object):
51  def __init__(self,sc,odsz,npy,exptFilter=None, odir='xls'):
52    idir = dreq.DOC_DIR
53    self.sc = sc
54    self.odsz=odsz
55    self.npy = npy
56    self.exptFilter = exptFilter
57    self.strSz = dict()
58    self.accum = False
59    self.odir = odir
60    self.efnsfx = ''
61    if sc.gridPolicyDefaultNative:
62      self.efnsfx = '_dn'
63    if not os.path.isdir( odir ):
64      print ( 'Creating new directory for xlsx output: %s' % odir )
65      os.mkdir( odir )
66    ii = open( '%s/sfheadings.csv' % idir, 'r' )
67    self.infoRows = []
68    for l in ii.readlines():
69      ll = l.strip().split( '\t' )
70      assert len(ll) == 2, 'Failed to parse info row: %s' % l
71      self.infoRows.append( ll )
72    ii.close()
73
74  def analAll(self,pmax):
75      volsmm={}
76      volsmmt={}
77      volsme={}
78      volsue={}
79      for m in ['TOTAL',] + self.sc.mips:
80        if m != 'TOTAL':
81          cmv1 = self.sc.cmvByInvMip(m,pmax=pmax,includeYears=True)
82          self.uniqueCmv = self.sc.differenceSelectedCmvDict(  cmv1, cmvTotal )
83        self.run( m, '%s/requestVol_%s_%s_%s' % (self.odir,m,self.sc.tierMax,pmax), pmax=pmax )
84
85        self.anal(olab=m,doUnique=True, makeTabs=True)
86        ttl = sum( [x for k,x in self.res['vu'].items()] )*2.*1.e-12
87        print ( '%s volume: %8.2fTb' % (m,ttl) )
88        volsmm[m] = self.res['vm']
89        volsmmt[m] = self.res['vmt']
90        volsme[m] = self.res['ve']
91        volsue[m] = self.res['uve']
92        if m == 'TOTAL':
93          cmvTotal = self.sc.selectedCmv.copy()
94          self.uniqueCmv =  {}
95      r1 = overviewTabs.r1( self.sc, pmax=pmax, vols=( volsmm, volsme, volsmmt,volsue ) )
96
97  def _analSelectedCmv(self,cmv):
98    lex = collections.defaultdict( list )
99    vet = collections.defaultdict( int )
100    vf = collections.defaultdict( int )
101    vu = collections.defaultdict( float )
102    mvol = collections.defaultdict( dict )
103
104    for u in cmv:
105      i = self.sc.dq.inx.uid[u]
106      if i._h.label != 'remarks':
107        npy = self.npy[ i.frequency ]
108        isClim = i.frequency.lower().find( 'clim' ) != -1
109        st = self.sc.dq.inx.uid[i.stid]
110        c1 = 0
111        for e,g in cmv[u]:
112          ee = self.sc.dq.inx.uid[e]
113          if ee.mip not in ['SolarMIP']:
114            lex[e].append( u )
115            t1, tt = self.sc.getStrSz( g, stid=i.stid, tt=True )
116            np = t1[1]*npy
117            if not isClim:
118              np = np*cmv[u][(e,g)]
119            c1 += cmv[u][(e,g)]
120            vet[(e,i.mipTable)] += np
121            vf[i.frequency] += np
122            vu[u] += np
123          else:
124            print ('ERROR.obsoleteMip.00001: %s,%s,%s' % (ee.mip,ee.label,ee.uid) )
125        if i.frequency == 'mon':
126            mvol[tt][u] = c1
127
128    return lex, vet, vf, vu, mvol
129
130  def anal(self,olab=None,doUnique=False,makeTabs=False):
131    vmt = collections.defaultdict( int )
132    vm = collections.defaultdict( int )
133    ve = collections.defaultdict( int )
134    uve = collections.defaultdict( int )
135    lm = collections.defaultdict( set )
136
137    lex, vet, vf, vu, mvol = self._analSelectedCmv(self.sc.selectedCmv )
138    if olab != 'TOTAL' and doUnique:
139      s_lex, s_vet, s_vf, s_vu, s_mvol = self._analSelectedCmv(self.uniqueCmv )
140      s_lm = set( self.uniqueCmv.keys() )
141      s_cc = collections.defaultdict( int )
142      for e,t in s_vet:
143        s_cc[t] += s_vet[(e,t)]
144        vm['Unique'] += s_vet[(e,t)]
145        vmt[('Unique',t)] += s_vet[(e,t)]
146        uve[e] += s_vet[(e,t)]
147
148    checkMvol = 1
149    for k in mvol:
150      sp = self.sc.dq.inx.uid[k[0]]
151      if k not in self.mvol:
152        print ( '%s missing from mvol: ' % str(k) )
153      else:
154        if checkMvol > 1:
155          for u in mvol[k]:
156            la = self.sc.dq.inx.uid[u].label
157            if self.mvol[k][u] != mvol[k][u]:
158              print ( 'MISMATCH IN %s (%s): %s:: %s,%s' % (str(k),sp.label,la,mvol[k][u],self.mvol[k][u]) )
159         
160    for e in lex:
161      ee = self.sc.dq.inx.uid[e]
162      for i in lex[e]:
163        lm[ee.mip].add( i )
164
165    for e,t in vet:
166      ee = self.sc.dq.inx.uid[e]
167      if ee.mip == 'SolarMIP':
168        print ('ERROR.solarmip: %s,%s,%s' % (ee.label, ee.title, ee.uid) )
169      vmt[(ee.mip,t)] += vet[(e,t)]
170      vm[ee.mip] += vet[(e,t)]
171      ve[e] += vet[(e,t)]
172
173##
174## makeTab needs: cc[m]: volume summary, by table,   lm[m]: list of CMOR variables
175##
176    cc = collections.defaultdict( dict )
177    cct = collections.defaultdict( int )
178    for m,t in vmt:
179      cc[m][t] = vmt[(m,t) ]
180    ss = set()
181    for m in sorted( vm.keys() ):
182      if olab != None:
183        for t in cc[m]:
184          cct[t] += cc[m][t]
185        ss = ss.union( lm[m] )
186        if makeTabs:
187          makeTables.makeTab(self.sc.dq, subset=lm[m], dest='%s/cmvmm_%s_%s_%s_%s%s' % (self.odir,olab,m,self.sc.tierMax,self.pmax,self.efnsfx), collected=cc[m])
188
189    if olab != None and makeTabs:
190        makeTables.makeTab(self.sc.dq, subset=ss, dest='%s/cmvmm_%s_%s_%s_%s%s' % (self.odir,olab,'TOTAL',self.sc.tierMax,self.pmax,self.efnsfx), collected=cct)
191        if olab != 'TOTAL' and doUnique:
192          makeTables.makeTab(self.sc.dq, subset=s_lm, dest='%s/cmvmm_%s_%s_%s_%s%s' % (self.odir,olab,'Unique',self.sc.tierMax,self.pmax,self.efnsfx), collected=s_cc)
193
194    cc = collections.defaultdict( dict )
195    ucc = collections.defaultdict( dict )
196    cct = collections.defaultdict( int )
197    for e,t in vet:
198      cc[e][t] = vet[(e,t) ]
199    for e in sorted( ve.keys() ):
200      if olab != None and makeTabs:
201        el = self.sc.dq.inx.uid[e].label
202        makeTables.makeTab(self.sc.dq, subset=lex[e], dest='%s/cmvme_%s_%s_%s_%s%s' % (self.odir,olab,el,self.sc.tierMax,self.pmax,self.efnsfx), collected=cc[e])
203
204    if olab != 'TOTAL' and doUnique:
205      for e,t in s_vet:
206        ucc[e][t] = s_vet[(e,t) ]
207      for e in sorted( uve.keys() ):
208        if olab != None and makeTabs:
209          el = self.sc.dq.inx.uid[e].label
210          makeTables.makeTab(self.sc.dq, subset=s_lex[e], dest='%s/cmvume_%s_%s_%s_%s%s' % (self.odir,olab,el,self.sc.tierMax,self.pmax,self.efnsfx), collected=ucc[e])
211
212    self.res = { 'vmt':vmt, 'vet':vet, 'vm':vm, 'uve':uve, 've':ve, 'lm':lm, 'lex':lex, 'vu':vu, 'cc':cc, 'cct':cct}
213    cc8 = collections.defaultdict( int )
214    for e,t in vet:
215      cc8[t] += vet[(e,t)]
216    for f in sorted( vf.keys() ):
217      print ( 'Frequency: %s: %s' % (f,vf[f]*2.*1.e-12 ) )
218       
219  def csvFreqStrSummary(self,mip,pmax=1):
220    sf, c3 = self.sc.getFreqStrSummary(mip,pmax=pmax)
221    self.c3 = c3
222    self.pmax = pmax
223    lf = sorted( list(sf) )
224    hdr0 = ['','','','']
225    hdr1 = ['','','','']
226    for f in lf:
227      hdr0 += [f,'','','']
228      hdr1 += ['','','',str( self.npy.get( f, '****') )]
229    orecs = [hdr0,hdr1,]
230    crecs = [None,None,]
231    self.mvol = collections.defaultdict( dict )
232    self.rvol = dict()
233
234    ix = 3
235    for tt in sorted( c3.keys() ):
236      s,o,g = tt
237      i = self.sc.dq.inx.uid[ s ]
238      if o != '':
239        msg = '%48.48s [%s]' % (i.title,o)
240      else:
241        msg = '%48.48s' % i.title
242      if g != 'native':
243        msg += '{%s}' % g
244      szg = self.sc.getStrSz( g, s=s, o=o )[1]
245      self.rvol[tt] = szg
246
247      rec = [msg,szg,2,'']
248      crec = ['','','','']
249      for f in lf:
250        if f in c3[tt]:
251            nn,ny,ne,labs,expts = c3[tt][f]
252            rec += [nn,ny,ne,'']
253            clabs = [self.sc.dq.inx.uid[x].label for x in labs.keys()]
254            crec += [sorted(clabs),'',expts,'']
255            if f.lower().find( 'clim' ) == -1:
256              assert abs( nn*ny - sum( [x for k,x in labs.items()] ) ) < .1, 'Inconsistency in year count: %s, %s, %s' % (str(tt),nn,ny)
257            if f == 'mon':
258              for k in labs:
259                self.mvol[tt][k] = labs[k]
260        else:
261            rec += ['','','','']
262            crec += ['','','','']
263      colr = xl_rowcol_to_cell(0, len(rec))
264      colr = colr[:-1]
265      eq = '$SUMPRODUCT(--(MOD(COLUMN(E%(ix)s:%(colr)s%(ix)s)-COLUMN(A%(ix)s)+1,4)=0),E%(ix)s:%(colr)s%(ix)s)' % {'ix':ix,'colr':colr}
266      ix += 1
267      rec[3] = eq
268      orecs.append( rec )
269      crecs.append( crec )
270   
271    return (orecs, crecs)
272
273  def byExpt(self):
274    for cmv in self.sc.selectedCmv.keys():
275      pass
276     
277  def run(self,mip='_all_',fn='test',pmax=1):
278    if mip == '_all_':
279      mip = set(self.sc.mips )
280    self.mip = mip
281    self.x = xlsx( fn )
282    self.sht = self.x.newSheet( 'Volume' )
283    orecs, crecs = self.csvFreqStrSummary(mip,pmax=pmax)
284    oh = orecs[0]
285    self.sht.set_column(0,0,60)
286    self.sht.set_column(1,1,15)
287    self.sht.set_column(2,2,4)
288    self.sht.set_column(3,3,15)
289    for k in range( int( (len(oh)-3)/4 ) ):
290      self.sht.set_column((k+1)*4,(k+1)*4,4)
291      self.sht.set_column((k+1)*4+1,(k+1)*4+1,8)
292      self.sht.set_column((k+1)*4+2,(k+1)*4+2,4)
293      self.sht.set_column((k+1)*4+3,(k+1)*4+3,12)
294     
295    oo = []
296    for i in range( len(oh) ):
297      oo.append( '' )
298    kk = 0
299    rr1 = 2
300    rr1p = rr1 + 1
301    for ix in range(len(orecs)):
302      o = orecs[ix]
303      kk += 1
304      if kk > 2:
305        for i in range( 7,len(o),4):
306          frq = oh[i-3]
307         
308          cell = xl_rowcol_to_cell(0, i)[:-1]
309          ca = xl_rowcol_to_cell(0, i-3)[:-1]
310          ##if frq.lower().find( 'clim' ) == -1:
311          cb = xl_rowcol_to_cell(0, i-2)[:-1]
312          ##else:
313          ##cb = xl_rowcol_to_cell(0, i-1)[:-1]
314          eq = '$%(cell)s$%(rr1)s*%(cb)s%(kk)s*%(ca)s%(kk)s*$B%(kk)s*$C%(kk)s*0.000000001' % locals()
315          o[i] = eq
316        self.x.tabrec(kk-1, o )
317        if crecs[ix] != None:
318          crec = crecs[ix]
319          for j in range(len(crec)):
320            if crec[j] != '':
321              self.sht.write_comment( kk-1, j, ' '.join( crec[j] ) )
322      else:
323        if kk == 1:
324          for i in range( 4,len(o),4):
325            cell = xl_rowcol_to_cell(0, i)[:-1]
326            cell2 = xl_rowcol_to_cell(0, i+3)[:-1]
327            self.sht.merge_range('%s1:%s1' % (cell,cell2), 'Merged Range')
328        self.x.tabrec(kk-1, o )
329
330    n = len(orecs)
331    for i in range( 3,len(oo),4):
332      cell = xl_rowcol_to_cell(0, i)[:-1]
333      oo[i] = '$SUM(%(cell)s%(rr1p)s:%(cell)s%(n)s)*0.001' % locals()
334    for i in range( 5,len(oo),4):
335      oo[i] = oh[i-1]
336    oo[0] = 'TOTAL VOLUME (Tb)'
337    self.x.tabrec(kk, oo )
338
339    n += 2
340    for a,b in self.infoRows:
341       self.sht.merge_range('B%s:H%s' % (n+1,n+1), 'Merged Range')
342       self.sht.write( n,0, a )
343       self.sht.write( n,1, b )
344       n += 1
345
346    self.x.close()
Note: See TracBrowser for help on using the repository browser.