source: CMIP6dreq/trunk/dreqPy/volsum.py @ 1265

Subversion URL: http://proj.badc.rl.ac.uk/svn/exarch/CMIP6dreq/trunk/dreqPy/volsum.py@1265
Revision 1265, 15.7 KB checked in by mjuckes, 12 months ago (diff)

01.00.28

Line 
1import xlsxwriter
2from xlsxwriter.utility import xl_rowcol_to_cell
3import collections, os
4
5
6try:
7  import dreq
8  imm=1
9except:
10  import dreqPy.dreq  as dreq
11  imm=2
12
13if imm == 1:
14  import misc_utils
15  import table_utils
16  import overviewTabs
17  from extensions import collect as extCollect
18##
19## the double underscore causes some confusion, creating an error message _vsum__requestLink__expt not found.
20  from extensions.collect import _requestLink__expt as requestLink__expt
21else:
22  import dreqPy.misc_utils as misc_utils
23  import dreqPy.table_utils as table_utils
24  import dreqPy.overviewTabs as overviewTabs
25  from dreqPy.extensions import collect as extCollect
26  from dreqPy.extensions.collect import _requestLink__expt as requestLink__expt
27
28
29
30class xlsx(object):
31  def __init__(self,fn,txtOpts=None):
32    """Class to write spreadsheets of CMOR variables"""
33    self.txtOpts = txtOpts
34    self.mcfgNote = 'Reference Volume (1 deg. atmosphere, 0.5 deg. ocean)'
35    self.wb = xlsxwriter.Workbook('%s.xlsx' % fn)
36    self.hdr_cell_format = self.wb.add_format({'text_wrap': True, 'font_size': 14, 'font_color':'#0000ff', 'bold':1, 'fg_color':'#aaaacc'})
37    self.hdr_cell_format.set_text_wrap()
38    self.sect_cell_format = self.wb.add_format({'text_wrap': True, 'font_size': 14, 'font_color':'#0000ff', 'bold':1, 'fg_color':'#ccccbb'})
39    self.sect_cell_format.set_text_wrap()
40    self.cell_format = self.wb.add_format({'text_wrap': True, 'font_size': 11})
41    self.cell_format.set_text_wrap()
42
43  def newSheet(self,name):
44    self.sht = self.wb.add_worksheet(name=name)
45    return self.sht
46
47  def tabrec(self,j,orec):
48        for i in range(len(orec)):
49          if orec[i] != '' and type( orec[i] ) == type( '' ) and orec[i][0] == '$':
50             self.sht.write_formula(j,i, '{=%s}' % orec[i][1:])
51          else:
52             if j == 0:
53               self.sht.write( j,i, orec[i], self.hdr_cell_format )
54             else:
55               self.sht.write( j,i, orec[i], self.cell_format )
56
57  def close(self):
58      self.wb.close()
59
60class vsum(object):
61  def __init__(self,sc,odsz,npy,exptFilter=None, odir='xls', tabByFreqRealm=False,txt=False,txtOpts=None):
62    self.tabByFreqRealm = tabByFreqRealm
63    self.doTxt = txt
64    self.txtOpts = txtOpts
65    idir = dreq.DOC_DIR
66    if 'collect' not in sc.dq._extensions_:
67      extCollect.add(sc.dq)
68    self.sc = sc
69    self.exptMipRql = collections.defaultdict( set )
70    if exptFilter != None:
71      assert type( exptFilter ) == type( set() ), 'FATAl.vsum.001: exptFilter argument must be type set: %s' % type( exptFilter )
72    for i in self.sc.dq.coll['requestLink'].items:
73      expts = requestLink__expt(i,rql=[i.uid,])
74      if exptFilter != None:
75        expts = exptFilter.intersection( expts )
76      for e in expts:
77        self.exptMipRql[ (i.mip,self.sc.dq.inx.uid[e].label) ].add( i.uid)
78    self.odsz=odsz
79    self.npy = npy
80    self.exptFilter = exptFilter
81    self.strSz = dict()
82    self.accum = False
83    self.odir = odir
84    self.efnsfx = ''
85    if sc.gridPolicyForce == 'native':
86      self.efnsfx = '_fn'
87    elif sc.gridPolicyForce == '1deg':
88      self.efnsfx = '_f1'
89    elif sc.gridPolicyDefaultNative:
90      self.efnsfx = '_dn'
91    if not os.path.isdir( odir ):
92      print ( 'Creating new directory for xlsx output: %s' % odir )
93      os.mkdir( odir )
94
95    self.xlsPrefixM = 'cmvmm'
96    self.xlsPrefixE = 'cmvme'
97    self.xlsPrefixU = 'cmvume'
98    if tabByFreqRealm:
99      self.xlsPrefixM += 'fr'
100      self.xlsPrefixE += 'fr'
101      self.xlsPrefixU += 'fr'
102    ii = open( '%s/sfheadings.csv' % idir, 'r' )
103    self.infoRows = []
104    for l in ii.readlines():
105      ll = l.strip().split( '\t' )
106      assert len(ll) == 2, 'Failed to parse info row: %s' % l
107      self.infoRows.append( ll )
108    ii.close()
109
110  def analAll(self,pmax,mips=None,html=True,makeTabs=True):
111      volsmm={}
112      volsmmt={}
113      volsme={}
114      volsue={}
115      if mips == None:
116        theseMips =  ['TOTAL',] + self.sc.mips
117      else:
118        theseMips = list(mips)
119
120## move *TOTAL to end of list.
121      if '*TOTAL' in theseMips:
122         theseMips.remove( '*TOTAL' )
123         theseMips.append( '*TOTAL' )
124
125      self.rres = {}
126      self.rresu = {}
127      self.accPdict = collections.defaultdict( set )
128   
129      for m in theseMips:
130        olab = m
131        useAccPdict = False
132        if m == '*TOTAL':
133            thism = theseMips[:]
134            if type( thism ) == type( set() ):
135              thism.remove( '*TOTAL' )
136            else:
137              thism.remove( '*TOTAL' )
138            olab = misc_utils.setMlab( thism )
139            useAccPdict = True
140        elif type( theseMips ) == type( dict() ):
141            thism = {m:theseMips[m]}
142        else:
143            thism = m
144
145        if m != 'TOTAL' and 'TOTAL' in theseMips:
146          cmv1, cmvts = self.sc.cmvByInvMip(thism,pmax=pmax,includeYears=True)
147          self.uniqueCmv = self.sc.differenceSelectedCmvDict(  cmv1, cmvTotal )
148
149        self.run( thism, '%s/requestVol_%s_%s_%s' % (self.odir,olab,self.sc.tierMax,pmax), pmax=pmax,doxlsx=makeTabs )
150
151        self.anal(olab=olab,doUnique='TOTAL' in theseMips, makeTabs=makeTabs, useAccPdict=useAccPdict)
152        ttl = sum( [x for k,x in self.res['vu'].items()] )*2.*1.e-12
153        volsmm[m] = self.res['vm']
154        volsmmt[m] = self.res['vmt']
155        volsme[m] = self.res['ve']
156        volsue[m] = self.res['uve']
157        self.rres[m] = self.res['vf'].copy()
158        self.rresu[m] = self.res['vu'].copy()
159        if m == 'TOTAL':
160          cmvTotal = self.sc.selectedCmv.copy()
161          self.uniqueCmv =  {}
162      if html:
163        r1 = overviewTabs.r1( self.sc, table_utils.tables, pmax=pmax, vols=( volsmm, volsme, volsmmt,volsue ) )
164
165  def _analSelectedCmv(self,cmv):
166    lex = collections.defaultdict( list )
167    vet = collections.defaultdict( int )
168    vf = collections.defaultdict( int )
169    vu = collections.defaultdict( float )
170    mvol = collections.defaultdict( dict )
171
172    for u in cmv:
173      i = self.sc.dq.inx.uid[u]
174      if i._h.label != 'remarks':
175        npy = self.npy[ i.frequency ]
176        isClim = i.frequency.lower().find( 'clim' ) != -1 or i.frequency in ['monC', '1hrCM']
177        st = self.sc.dq.inx.uid[i.stid]
178        c1 = 0
179        for e,g in cmv[u]:
180          ee = self.sc.dq.inx.uid[e]
181          if ee.mip not in ['SolarMIP']:
182            lex[e].append( u )
183            t1, tt = self.sc.getStrSz( g, stid=i.stid, tt=True, cmv=u )
184            np = t1[1]*npy
185            if not isClim:
186              np = np*cmv[u][(e,g)]
187            c1 += cmv[u][(e,g)]
188            vet[(e,i.mipTable)] += np
189            vf[i.frequency] += np
190            vu[u] += np
191          else:
192            print ('ERROR.obsoleteMip.00001: %s,%s,%s' % (ee.mip,ee.label,ee.uid) )
193        if i.frequency in ['mon','monPt']:
194            mvol[tt][u] = c1
195
196    return dict(lex), dict(vet), dict(vf), dict(vu), dict(mvol)
197
198  def xlsDest(self,mode,olab,lab2):
199    if mode == 'e':
200      return '%s/%s_%s_%s_%s_%s%s' % (self.odir,self.xlsPrefixE,olab,lab2,self.sc.tierMax,self.pmax,self.efnsfx)
201    elif mode == 'u':
202      return '%s/%s_%s_%s_%s_%s%s' % (self.odir,self.xlsPrefixU,olab,lab2,self.sc.tierMax,self.pmax,self.efnsfx)
203    else:
204      return '%s/%s_%s_%s_%s_%s%s' % (self.odir,self.xlsPrefixM,olab,lab2,self.sc.tierMax,self.pmax,self.efnsfx)
205
206  def anal(self,olab=None,doUnique=False,makeTabs=False,mode='full',useAccPdict=False):
207    vmt = collections.defaultdict( int )
208    vm = collections.defaultdict( int )
209    ve = collections.defaultdict( int )
210    uve = collections.defaultdict( int )
211    lm = collections.defaultdict( set )
212
213    lex, vet, vf, vu, mvol = self._analSelectedCmv(self.sc.selectedCmv )
214    if mode == 'short':
215      self.res = { 'vet':vet,  'lex':lex, 'vu':vu, 'vf':vf}
216      return
217
218    if olab != 'TOTAL' and doUnique:
219      s_lex, s_vet, s_vf, s_vu, s_mvol = self._analSelectedCmv(self.uniqueCmv )
220      s_lm = set( self.uniqueCmv.keys() )
221      s_cc = collections.defaultdict( int )
222      for e,t in s_vet:
223        s_cc[t] += s_vet[(e,t)]
224        vm['Unique'] += s_vet[(e,t)]
225        vmt[('Unique',t)] += s_vet[(e,t)]
226        uve[e] += s_vet[(e,t)]
227
228    checkMvol = -1
229    if checkMvol > 0:
230      for k in mvol:
231        sp = self.sc.dq.inx.uid[k[0]]
232        if k not in self.mvol:
233          print ( 'INFO.volsum.01001: %s missing from mvol: ' % str(k) )
234        else:
235          if checkMvol > 1:
236            for u in mvol[k]:
237              la = self.sc.dq.inx.uid[u].label
238              if self.mvol[k][u] != mvol[k][u]:
239                print ( 'MISMATCH IN %s (%s): %s:: %s,%s' % (str(k),sp.label,la,mvol[k][u],self.mvol[k][u]) )
240         
241    for e in lex:
242      ee = self.sc.dq.inx.uid[e]
243      for i in lex[e]:
244        lm[ee.mip].add( i )
245
246    for e,t in vet:
247      ee = self.sc.dq.inx.uid[e]
248      vmt[(ee.mip,t)] += vet[(e,t)]
249      vm[ee.mip] += vet[(e,t)]
250      ve[e] += vet[(e,t)]
251##
252## makeTab needs: cc[m]: volume summary, by table,   lm[m]: list of CMOR variables
253##
254    cc = collections.defaultdict( dict )
255    cct = collections.defaultdict( int )
256    for m,t in vmt:
257      cc[m][t] = vmt[(m,t) ]
258    ss = set()
259    for m in sorted( vm.keys() ):
260      if olab != None:
261        for t in cc[m]:
262          cct[t] += cc[m][t]
263        ss = ss.union( lm[m] )
264        if makeTabs:
265          ##table_utils.makeTab(self.sc.dq, subset=lm[m], dest=self.xlsDest('m',olab,m), collected=cc[m],exptUid=self.sc.exptByLabel.get(m,m) )
266          table_utils.makeTab(self.sc, subset=lm[m], dest=self.xlsDest('m',olab,m), collected=cc[m], txt=self.doTxt, txtOpts=self.txtOpts )
267
268    if olab != None and makeTabs:
269        table_utils.makeTab(self.sc, subset=ss, dest=self.xlsDest('m',olab,'TOTAL'), collected=cct, txt=self.doTxt, txtOpts=self.txtOpts )
270        if olab != 'TOTAL' and doUnique:
271          table_utils.makeTab(self.sc, subset=s_lm, dest=self.xlsDest('m',olab,'Unique'), collected=s_cc, txt=self.doTxt, txtOpts=self.txtOpts )
272
273    cc = collections.defaultdict( dict )
274    ucc = collections.defaultdict( dict )
275    cct = collections.defaultdict( int )
276    for e,t in vet:
277      cc[e][t] = vet[(e,t) ]
278    for e in sorted( ve.keys() ):
279      if olab != None and makeTabs:
280        el = self.sc.dq.inx.uid[e].label
281
282        if useAccPdict:
283          pdict = collections.defaultdict( set )
284          for vid, p in self.accPdict[e]:
285               pdict[vid].add( p )
286             
287        elif olab in ['Total','TOTAL']:
288          pdict = None
289        elif (olab,el) in self.exptMipRql:
290          pdict = collections.defaultdict( set )
291          for rqlid in self.exptMipRql[(olab,el)]:
292             rql = self.sc.dq.inx.uid[rqlid]
293             for rqvid in self.sc.dq.inx.iref_by_sect[rql.refid].a['requestVar']:
294               rqv = self.sc.dq.inx.uid[rqvid]
295               pdict[rqv.vid].add( rqv.priority )
296               self.accPdict[e].add( (rqv.vid,rqv.priority) )
297        else:
298          print ( 'INFO.00201: olab,e not found:',olab,el )
299          pdict = None
300
301        tslice = {}
302        for v in self.sc.cmvts:
303          if e in self.sc.cmvts[v]:
304            tslice[v] = self.sc.cmvts[v][e]
305        dest = self.xlsDest('e',olab,el)
306        mode ='e'
307        table_utils.makeTab(self.sc, subset=lex[e], dest=self.xlsDest(mode,olab,el), collected=cc[e],byFreqRealm=self.tabByFreqRealm, tslice=tslice, exptUid=e, tabMode=mode, pdict=pdict, txt=self.doTxt, txtOpts=self.txtOpts )
308
309    if olab != 'TOTAL' and doUnique:
310      for e,t in s_vet:
311        ucc[e][t] = s_vet[(e,t) ]
312      for e in sorted( uve.keys() ):
313        if olab != None and makeTabs:
314          el = self.sc.dq.inx.uid[e].label
315          table_utils.makeTab(self.sc, subset=s_lex[e], dest=self.xlsDest('u',olab,el), collected=ucc[e], txt=self.doTxt, txtOpts=self.txtOpts)
316
317    self.res = { 'vmt':vmt, 'vet':vet, 'vm':vm, 'uve':uve, 've':ve, 'lm':lm, 'lex':lex, 'vu':vu, 'cc':cc, 'cct':cct, 'vf':vf}
318       
319  def csvFreqStrSummary(self,mip,pmax=1):
320    sf, c3 = self.sc.getFreqStrSummary(mip,pmax=pmax)
321    self.c3 = c3
322    self.pmax = pmax
323    lf = sorted( list(sf) )
324    hdr0 = ['','','','']
325    hdr1 = ['','','','']
326    for f in lf:
327      hdr0 += [f,'','','']
328      hdr1 += ['','','',str( self.npy.get( f, '****') )]
329    orecs = [hdr0,hdr1,]
330    crecs = [None,None,]
331    self.mvol = collections.defaultdict( dict )
332    self.rvol = dict()
333    ix = 3
334    for tt in sorted( c3.keys() ):
335      s,o,g = tt
336      i = self.sc.dq.inx.uid[ s ]
337      if o != '' and type(o) == type('x'):
338        msg = '%48.48s [%s]' % (i.title,o)
339      else:
340        msg = '%48.48s' % i.title
341      if g != 'native':
342        msg += '{%s}' % g
343      szg = self.sc.getStrSz( g, s=s, o=o )[1]
344      self.rvol[tt] = szg
345
346      rec = [msg,szg,2,'']
347      crec = ['','','','']
348      for f in lf:
349        if f in c3[tt]:
350            nn,ny,ne,labs,expts = c3[tt][f]
351            rec += [nn,ny,ne,'']
352            clabs = [self.sc.dq.inx.uid[x].label for x in labs.keys()]
353            crec += [sorted(clabs),'',expts,'']
354            if f.lower().find( 'clim' ) == -1:
355              assert abs( nn*ny - sum( [x for k,x in labs.items()] ) ) < .1, 'Inconsistency in year count: %s, %s, %s' % (str(tt),nn,ny)
356            if f == 'mon':
357              for k in labs:
358                self.mvol[tt][k] = labs[k]
359        else:
360            rec += ['','','','']
361            crec += ['','','','']
362      colr = xl_rowcol_to_cell(0, len(rec))
363      colr = colr[:-1]
364      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}
365      ix += 1
366      rec[3] = eq
367      orecs.append( rec )
368      crecs.append( crec )
369   
370    return (orecs, crecs)
371
372  def byExpt(self):
373    for cmv in self.sc.selectedCmv.keys():
374      pass
375     
376  def run(self,mip='_all_',fn='test',pmax=1,doxlsx=True):
377    if mip == '_all_':
378      mip = set(self.sc.mips )
379    self.mip = mip
380    orecs, crecs = self.csvFreqStrSummary(mip,pmax=pmax)
381    if not doxlsx:
382      return
383    print ('INFO.volsum.01002: Writing %s.xlsx' % fn )
384    self.x = xlsx( fn )
385    self.sht = self.x.newSheet( 'Volume' )
386    oh = orecs[0]
387    self.sht.set_column(0,0,60)
388    self.sht.set_column(1,1,15)
389    self.sht.set_column(2,2,4)
390    self.sht.set_column(3,3,15)
391    for k in range( int( (len(oh)-3)/4 ) ):
392      self.sht.set_column((k+1)*4,(k+1)*4,4)
393      self.sht.set_column((k+1)*4+1,(k+1)*4+1,8)
394      self.sht.set_column((k+1)*4+2,(k+1)*4+2,4)
395      self.sht.set_column((k+1)*4+3,(k+1)*4+3,12)
396     
397    oo = []
398    for i in range( len(oh) ):
399      oo.append( '' )
400    kk = 0
401    rr1 = 2
402    rr1p = rr1 + 1
403    for ix in range(len(orecs)):
404      o = orecs[ix]
405      kk += 1
406      if kk > 2:
407        for i in range( 7,len(o),4):
408          frq = oh[i-3]
409         
410          cell = xl_rowcol_to_cell(0, i)[:-1]
411          ca = xl_rowcol_to_cell(0, i-3)[:-1]
412          ##if frq.lower().find( 'clim' ) == -1:
413          cb = xl_rowcol_to_cell(0, i-2)[:-1]
414          ##else:
415          ##cb = xl_rowcol_to_cell(0, i-1)[:-1]
416          eq = '$%(cell)s$%(rr1)s*%(cb)s%(kk)s*%(ca)s%(kk)s*$B%(kk)s*$C%(kk)s*0.000000001' % locals()
417          o[i] = eq
418        self.x.tabrec(kk-1, o )
419        if crecs[ix] != None:
420          crec = crecs[ix]
421          for j in range(len(crec)):
422            if crec[j] != '':
423              self.sht.write_comment( kk-1, j, ' '.join( crec[j] ) )
424      else:
425        if kk == 1:
426          for i in range( 4,len(o),4):
427            cell = xl_rowcol_to_cell(0, i)[:-1]
428            cell2 = xl_rowcol_to_cell(0, i+3)[:-1]
429            self.sht.merge_range('%s1:%s1' % (cell,cell2), 'Merged Range')
430        self.x.tabrec(kk-1, o )
431
432    n = len(orecs)
433    for i in range( 3,len(oo),4):
434      cell = xl_rowcol_to_cell(0, i)[:-1]
435      oo[i] = '$SUM(%(cell)s%(rr1p)s:%(cell)s%(n)s)*0.001' % locals()
436    for i in range( 5,len(oo),4):
437      oo[i] = oh[i-1]
438    oo[0] = 'TOTAL VOLUME (Tb)'
439    self.x.tabrec(kk, oo )
440
441    n += 2
442    for a,b in self.infoRows:
443       self.sht.merge_range('B%s:H%s' % (n+1,n+1), 'Merged Range')
444       self.sht.write( n,0, a )
445       self.sht.write( n,1, b )
446       n += 1
447
448    self.x.close()
Note: See TracBrowser for help on using the repository browser.