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

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

updates to tikz

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, parentLabel=%(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      oo.write( '\\section{%s}\n' % self.sections['fluxDiagram'][mip][1] )
217      oo.write( '''\\begin{figure}[h]
218\definecolor{rcolor}{RGB}{200,255,200}
219\definecolor{rpcolor}{RGB}{220,255,200}
220\definecolor{rppcolor}{RGB}{220,220,200}
221\definecolor{rptxcolor}{RGB}{210,255,210}
222\definecolor{lcolor}{RGB}{00,00,128}
223\\begin{tikzpicture}[>=latex']
224\\tikzstyle gline=[color=green!50,text=black]
225\\tikzstyle ggline=[color=green!80,text=black]
226 \\tikzset{block/.style= {draw=blue!80, thick, fill=rcolor, rectangle, align=center,minimum width=2cm,minimum height=1cm},
227        every arrow/.style={ draw=green!80, line width=4pt, },
228        every path/.style={ draw, line width=4pt, },
229        reservoir/.style={draw=blue!80, thick, fill=rcolor, shape=rectangle,rounded corners=1.5em,align=center,minimum width=2cm,minimum height=1cm},
230        reservoirp/.style={draw=blue!80, thick, fill=rpcolor, shape=rectangle,rounded corners=.5em,align=center,minimum width=11cm,minimum height=2.5cm},
231        reservoirpp/.style={draw=blue!80, thick, fill=rppcolor, shape=rectangle,rounded corners=.5em,align=center,minimum width=14cm,minimum height=5.cm},
232        reservoirptx/.style={draw=blue!80, thick, fill=rptxcolor, dotted, shape=rectangle,rounded corners=.5em,align=center,minimum width=3cm,minimum height=2cm}
233     }
234''' )
235
236      self.writeMip( mip, oo )
237
238      oo.write( '''
239\\end{tikzpicture}
240\\caption{%s}
241\end{figure}\n''' % self.sections['fluxDiagram'][mip][2] )
242
243      vars = sorted( list( self.vars ) )
244      oo.write( '''\\begin{itemize}
245\\setlength\\itemsep{1pt}
246\\setlength{\\parskip}{0pt}
247\\setlength{\\parsep}{0pt}\n''' )
248      for v in vars:
249        if v in self.vlk:
250          i = self.dq.inx.uid[self.vlk[v]]
251          oo.write( '\\item   %s: %s\n' % (v,i.title) )
252      oo.write( '\\end{itemize}\n' )
253
254  def writeMip(self,mip,oo):
255 
256    vars = set()
257    oo.write( '\\begin{pgfonlayer}{child}\n' )
258    if mip in self.reservoirs:
259      for this in self.reservoirs[mip] + self.reservoirs['all']:
260        this.dec()
261        if this.isParent:
262          oo.write( nodeTmpl % this.__dict__ )
263          vars.add( this.var )
264
265      for this in self.reservoirs[mip] + self.reservoirs['all']:
266        if this.isParent:
267          oo.write( nodeTmpl % this.tnode.__dict__ )
268          vars.add( this.tnode.var )
269        else:
270          oo.write( nodeTmpl % this.__dict__ )
271          vars.add( this.var )
272    oo.write( '\\end{pgfonlayer}{child}\n' )
273         
274         
275
276    if mip in self.fluxes:
277      cll = collections.defaultdict( list )
278      ll = []
279      llb = []
280      cc = collections.defaultdict(list)
281   
282      for this in self.fluxes[mip]:
283        t = tuple( sorted( [this.tto, this.ffrom] ) )
284        vars.add( this.label )
285        cc[t].append( this )
286      for t in cc:
287        if len(cc[t]) > 3:
288          cc[t][0].pathDec = '[bend left=25]'
289          cc[t][1].pathDec = '[bend left=10]'
290          crv = {True:'right', False:'left'}[ cc[t][1].tto == cc[t][-2].tto ]
291          cc[t][-2].pathDec = '[bend %s=10]' % crv
292          crv = {True:'right', False:'left'}[ cc[t][0].tto == cc[t][-1].tto ]
293          cc[t][-1].pathDec = '[bend %s=25]' % crv
294        elif len(cc[t]) > 1:
295          cc[t][0].pathDec = '[bend left=10]'
296          crv = {True:'right', False:'left'}[ cc[t][0].tto == cc[t][-1].tto ]
297          cc[t][-1].pathDec = '[bend %s=10]' % crv
298        for this in cc[t]:
299          this.dec(budgets=self.budgets,mip=mip)
300          ss = flxTmpl % this.__dict__
301          cll[ (this.oneway,this.subf) ].append( ss )
302          ##if this.oneway:
303            ##ll.append(ss)
304          ##else:
305            ##llb.append(ss)
306
307      for ow,subf in cll:
308        ahd = { True:'->', False:'<->' }[ow]
309        sty = { True:'gline', False:'ggline' }[subf]
310        oo.write( '\path[draw,%s,sloped,style=%s]\n' % (ahd,sty) )
311        for l in cll[(ow,subf)]:
312          oo.write( l )
313        oo.write( ';' )
314
315    self.vars = vars
316
317r = run()
318oo = open( 'test.itex', 'w' )
319r.write( oo )
320oo.close()
Note: See TracBrowser for help on using the repository browser.