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

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

cleaned examples and import structure

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