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

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

candidate 01.00.12

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