source: CMIP6dreqbuild/trunk/src/framework/tikz.py @ 1178

Subversion URL: http://proj.badc.rl.ac.uk/svn/exarch/CMIP6dreqbuild/trunk/src/framework/tikz.py@1207
Revision 1178, 11.3 KB checked in by mjuckes, 2 years ago (diff)

flux diagram updates

Line 
1import collections, string
2from utils_wb import workbook
3from dreqPy import dreq
4
5fluxStyles = {'SIMIP':{'sndmasswindrif':[('pathDec','[bend right=40]'),] },
6              'ISMIP6':{'mrro':[('pathDec','[bend left=25]'),] } }
7
8
9def floatify( x ):
10  if x == '':
11    return 0.
12  elif type(x) == type(''):
13    return float(x)
14  return c
15
16class node(object):
17  def __init__(self,label,title,var=None):
18    self.nclass = 'reservoir'
19    self.label = label
20    self.title = title
21    self.var=var
22    self.ncx = ''
23    self.decVar = ''
24    self.decTitle = title
25    self.isChild = False
26    self.isParent = False
27    self.child = []
28
29  def setPos(self,x,y):
30    self.x = x
31    self.y = y
32
33  def set(self,k,v,mode=None):
34    self.__dict__[k] = v
35    if mode == 'float':
36      self.__dict__[k] = floatify(v)
37
38  def addChildren(self,nlist):
39    '''Add a list of child nodes and create a node list string for use in tikz'''
40    self.child = nlist
41    self.childList = ''.join( [ '(%s)' % x.label for x in nlist] )
42    self.isParent = True
43
44  def dec(self):
45    if self.isParent:
46      self.nclass = 'reservoirp'
47      self.tnode = node( 'p-' + self.label, self.title, self.var )
48      ocx = self.__dict__.get( 'ocx', -3 )
49      ocy = self.__dict__.get( 'ocy', 0 )
50      if ocx == '':
51        ocx = -3.
52      if ocy == '':
53        ocy = 0.
54      print ocx, ocy, type(ocx), type(ocy)
55      self.tnode.setPos( str(float(self.x)+ocx), str( float(self.y) + ocy ) )
56      self.tnode.dec()
57      self.tnode.nclass = 'reservoirptx'
58      self.title = ''
59      ##self.decTitle = ''
60      self.var = None
61    if 'dx' in self.__dict__ and self.dx not in [None,'',0.]:
62        print '**********',self.dx,'*********', type(self.dx)
63        self.ncx += ', minimum width=%scm' % self.dx
64    if 'dy' in self.__dict__ and self.dy not in [None,'',0.]:
65        print '**********',self.dy,'*********', type(self.dy)
66        self.ncx += ', minimum height=%scm' % self.dy
67   
68    if len(self.child) > 0 and any( [x.isParent for x in self.child ] ):
69        self.nclass = 'reservoirpp'
70        print 'GRANDPARENT !!!!!!!!!!!!!!!!'
71     
72    if self.var != None:
73      self.decVar = '[%s]' % self.var
74      self.decTitle = '%s \\\\ {\\bf %s}' % (self.title, self.var) 
75
76class flux(object):
77  def __init__(self,label,title,tto,ffrom,oneway,isSubf):
78    self.label = label
79    self.title = title
80    self.tto = tto
81    self.ffrom = ffrom
82    self.oneway = oneway
83    self.subf = isSubf
84    self.flxlabel = ''
85    self.pathDec = ''
86
87  def dec(self,budgets=None,mip=None):
88    lab = self.label
89    if budgets != None:
90      if self.label[0] == '@':
91        print '###### ',self.label
92        if self.label[1:] in budgets:
93          self.label = budgets[ self.label[1:] ][5]
94          print '########## ',self.label, lab
95        else:
96          print sorted( budgets.keys() )
97
98    if mip in fluxStyles and lab in fluxStyles[mip]:
99      self.pathDec = fluxStyles[mip][lab][0][1]
100
101    if self.label[0] == '$':
102      self.flxLabel = flxLabelTmplAlt % self.label[1:]
103    else:
104      self.flxLabel = flxLabelTmpl % self.label
105 
106
107nodeTmpl = '''   \\node at (%(x)s,%(y)s) [%(nclass)s%(ncx)s] (%(label)s) {%(decTitle)s} ;\n'''
108parentNodeTmpl = '''   \\node [%(nclass)s%(ncx)s, header=%(decTitle)s, fit={%(childList)s}] (%(label)s) {} ;\n'''
109flxTmpl = '''  (%(ffrom)s) edge %(pathDec)s %(flxLabel)s (%(tto)s) \n'''
110flxLabelTmpl = ''' node[below,align=center]{\\bf %s} '''
111flxLabelTmplAlt = ''' node[below,align=center]{\\it %s} '''
112
113class run(object):
114  def __init__(self):
115    self.wb = workbook( 'budgets.xls' )
116    assert 'reservoirs' in self.wb.sns, 'reservoirs sheet not found'
117    assert 'fluxes' in self.wb.sns, 'fluxes sheet not found'
118    self.dq = dreq.loadDreq()
119    self.vlk = {}
120    for i in self.dq.coll['var'].items:
121      self.vlk[i.label] = i.uid
122
123    self.loadRes()
124    self.loadFl()
125    self.loadBudg()
126    self.loadSections()
127
128  def loadRes(self):
129    self.reservoirs = collections.defaultdict( list )
130    self.reservoirsIx = {}
131    parents = collections.defaultdict( list )
132    atmos = node( 'atmos', 'Atmosphere' )
133    atmos.setPos( 0., 4. )
134    ocean = node( 'ocean', 'Ocean' )
135    ocean.setPos( 8., -8. )
136    self.reservoirs['all'].append( atmos )
137    self.reservoirs['all'].append( ocean )
138    sht = self.wb.book.sheet_by_name( 'reservoirs' )
139    for k in range( 1, sht.nrows):
140      r = [str(x.value) for x in sht.row(k)]
141      label, title, stuff, mip = (r[0],r[1],r[2],r[3])
142      if len(r) > 5 and r[5] != '':
143        var = r[5]
144      else:
145        var = None
146      thisnode = node( label, title, var=var )
147      thisnode.setPos( r[8], r[9] )
148
149      if len(r) > 10:
150        thisnode.set( 'dx', r[10], mode='float' )
151      if len(r) > 11:
152        thisnode.set( 'dy', r[11], mode='float' )
153      if len(r) > 12:
154        thisnode.set( 'ocx', r[12], mode='float' )
155      if len(r) > 13:
156        thisnode.set( 'ocy', r[13], mode='float' )
157
158      if r[7] != '':
159        parent = r[7]
160        parents[(mip,parent)].append( thisnode )
161        thisnode.parent = parent
162        thisnode.isChild = True
163      self.reservoirsIx[ (mip,label) ] = thisnode
164      self.reservoirs[mip].append( thisnode )
165    for mip,p in parents:
166      try:
167        self.reservoirsIx[ (mip,p) ].addChildren( parents[(mip,p)] )
168        ##self.reservoirsIx[ (mip,p) ].child = parents[(mip,p)]
169        ##self.reservoirsIx[ (mip,p) ].child = thisnode
170        ##self.reservoirsIx[ (mip,p) ].isParent = True
171      except:
172        print 'Failed to set parent: ', mip,p
173
174  def loadFl(self):
175    self.fluxes = collections.defaultdict( list )
176    sht = self.wb.book.sheet_by_name( 'fluxes' )
177
178    for k in range( 1, sht.nrows):
179      r = [str(x.value) for x in sht.row(k)]
180      label, title, ffrom, tto, ow, subf, sgnc, mip, budg = (r[0],r[1],r[2],r[3],r[4],r[5],r[6],r[7],r[8])
181      if len(tto) > 0 and len(ffrom) > 0:
182        onew = ow == '1'
183        isSubf = subf == '1'
184        print r, onew, isSubf
185        if tto[0] == '@' or ffrom[0] == '@':
186          print 'Not supported yet: ',r
187        else:
188          thisnode = flux( label, title, tto, ffrom, onew, isSubf )
189          self.fluxes[mip].append( thisnode )
190
191  def loadBudg(self):
192    self.budgets = {}
193    sht = self.wb.book.sheet_by_name( 'budgets' )
194
195    for k in range( 1, sht.nrows):
196      r = [string.strip(str(x.value)) for x in sht.row(k)]
197      if len(r) > 0:
198        label, title, uid, mip, description, lhs, rhs = (r[0],r[1],r[2],r[3],r[4],r[5],r[6])
199        self.budgets[label] = (title, uid, mip, description, lhs, rhs)
200
201  def loadSections(self):
202    self.sections = collections.defaultdict( dict )
203    self.mips = set()
204    sht = self.wb.book.sheet_by_name( 'sections' )
205
206    for k in range( 1, sht.nrows):
207      r = [str(x.value) for x in sht.row(k)]
208      label, title, description, type, mip = (r[0],r[1],r[2],r[3],r[4])
209      self.sections[type][mip] = (label, title, description )
210      self.mips.add( mip )
211
212
213  def write(self,oo):
214    for mip in sorted( list( self.mips ) ):
215     ##if mip in ['ISMIP6','SIMIP','LS3MIP','C4MIP']:
216     if mip in ['SIMIP','C4MIP']:
217      oo.write( '\\section{%s}\n' % self.sections['fluxDiagram'][mip][1] )
218      oo.write( '''\\begin{figure}[h]
219\definecolor{rcolor}{RGB}{200,255,200}
220\definecolor{rpcolor}{RGB}{220,255,200}
221\definecolor{rppcolor}{RGB}{220,220,200}
222\definecolor{rptxcolor}{RGB}{210,255,210}
223\definecolor{lcolor}{RGB}{00,00,128}
224\\begin{tikzpicture}[>=latex']
225\\tikzstyle gline=[color=green!50,text=black]
226\\tikzstyle ggline=[color=green!80,text=black]
227 \\tikzset{block/.style= {draw=blue!80, thick, fill=rcolor, rectangle, align=center,minimum width=2cm,minimum height=1cm},
228        every arrow/.style={ draw=green!80, line width=4pt, },
229        every path/.style={ draw, line width=4pt, },
230        reservoir/.style={draw=blue!80, thick, fill=rcolor, shape=rectangle,rounded corners=1.5em,align=center,minimum width=2cm,minimum height=1cm},
231        reservoirp/.style={draw=blue!80, thick, fill=rpcolor, shape=rectangle,rounded corners=.5em,align=center,minimum height=2.5cm},
232        reservoirpp/.style={draw=blue!80, thick, fill=rppcolor, shape=rectangle,rounded corners=.5em,align=center,minimum width=14cm,minimum height=5.cm},
233        reservoirptx/.style={draw=blue!80, thick, fill=rptxcolor, dotted, shape=rectangle,rounded corners=.5em,align=center,minimum width=3cm,minimum height=2cm}
234     }
235''' )
236
237      self.writeMip( mip, oo )
238
239      oo.write( '''
240\\end{tikzpicture}
241\\caption{%s}
242\end{figure}\n''' % self.sections['fluxDiagram'][mip][2] )
243
244      vars = sorted( list( self.vars ) )
245      oo.write( '''\\begin{itemize}
246\\setlength\\itemsep{1pt}
247\\setlength{\\parskip}{0pt}
248\\setlength{\\parsep}{0pt}\n''' )
249      for v in vars:
250        if v in self.vlk:
251          i = self.dq.inx.uid[self.vlk[v]]
252          oo.write( '\\item   %s: %s\n' % (v,i.title) )
253      oo.write( '\\end{itemize}\n' )
254
255  def writeMip(self,mip,oo):
256 
257    vars = set()
258    if mip in self.reservoirs:
259      oo.write( '\\begin{pgfonlayer}{child}\n' )
260      #for this in self.reservoirs[mip] + self.reservoirs['all']:
261      #  if this.isParent:
262      #    oo.write( nodeTmpl % this.__dict__ )
263      #    vars.add( this.var )
264
265      np = 0
266      for this in self.reservoirs[mip] + self.reservoirs['all']:
267        this.dec()
268        if this.isParent:
269          np += 1
270          pass
271          ##oo.write( parentNodeTmpl % this.__dict__ )
272          ##vars.add( this.var )
273        else:
274          oo.write( nodeTmpl % this.__dict__ )
275          vars.add( this.var )
276      oo.write( '\\end{pgfonlayer}{child}\n' )
277      if np > 0:
278        oo.write( '\\begin{pgfonlayer}{parent}\n' )
279
280        for this in self.reservoirs[mip] + self.reservoirs['all']:
281          this.dec()
282          if this.isParent:
283            oo.write( parentNodeTmpl % this.__dict__ )
284            vars.add( this.var )
285        oo.write( '\\end{pgfonlayer}{parent}\n' )
286         
287         
288
289    if mip in self.fluxes:
290      oo.write( '\\begin{pgfonlayer}{child}\n' )
291      cll = collections.defaultdict( list )
292      ll = []
293      llb = []
294      cc = collections.defaultdict(list)
295   
296      for this in self.fluxes[mip]:
297        t = tuple( sorted( [this.tto, this.ffrom] ) )
298        vars.add( this.label )
299        cc[t].append( this )
300      for t in cc:
301        if len(cc[t]) > 3:
302          cc[t][0].pathDec = '[bend left=25]'
303          cc[t][1].pathDec = '[bend left=10]'
304          crv = {True:'right', False:'left'}[ cc[t][1].tto == cc[t][-2].tto ]
305          cc[t][-2].pathDec = '[bend %s=10]' % crv
306          crv = {True:'right', False:'left'}[ cc[t][0].tto == cc[t][-1].tto ]
307          cc[t][-1].pathDec = '[bend %s=25]' % crv
308        elif len(cc[t]) > 1:
309          cc[t][0].pathDec = '[bend left=10]'
310          crv = {True:'right', False:'left'}[ cc[t][0].tto == cc[t][-1].tto ]
311          cc[t][-1].pathDec = '[bend %s=10]' % crv
312        for this in cc[t]:
313          this.dec(budgets=self.budgets,mip=mip)
314          ss = flxTmpl % this.__dict__
315          cll[ (this.oneway,this.subf) ].append( ss )
316          ##if this.oneway:
317            ##ll.append(ss)
318          ##else:
319            ##llb.append(ss)
320
321      for ow,subf in cll:
322        ahd = { True:'->', False:'<->' }[ow]
323        sty = { True:'gline', False:'ggline' }[subf]
324        oo.write( '\path[draw,%s,sloped,style=%s]\n' % (ahd,sty) )
325        for l in cll[(ow,subf)]:
326          oo.write( l )
327        oo.write( ';' )
328      oo.write( '\\end{pgfonlayer}{child}\n' )
329
330    self.vars = vars
331
332r = run()
333oo = open( 'test.itex', 'w' )
334r.write( oo )
335oo.close()
Note: See TracBrowser for help on using the repository browser.