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

Subversion URL: http://proj.badc.rl.ac.uk/svn/exarch/CMIP6dreq/trunk/dreqPy/volsum.py@880
Revision 880, 13.0 KB checked in by mjuckes, 3 years ago (diff)

Updated setup for tag 01.beta.43

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