Source Code for Module parasol.CarpetPlot

  1   
  2  from pylab import * 
  3  import copy 
  4  from math import * 
  5   
  6  #def setAxes(): 
  7  #    ax = gca() 
  8  #    ax.apply_aspect() 
  9  #     
 10  #    print ax.get_xlim()  
 11  #    print ax.get_ylim() 
 12   
 13 -class Carpet( object ): 
 14       
 15 -    def __init__(self, func,  
 16          aName='A', aList=[1.,2.,3.,4.],  bName='B', bList=[1.,2.,3.,4.], 
 17          nStepsA=20, nStepsB=20, iLabelsX=0, iLabelsY=1000, alphaInLineLabels=0.5, 
 18          linewidth=2, inLineLabels=1, showGrid=1, showLabels=1, 
 19          logX=0, logY=0, titleStr='', xLabelStr='', yLabelStr=''): 
 20          '''func returns the x,y coordinates that are plotted''' 
 21   
 22          self.func = func 
 23          self.aName = aName 
 24          self.aList = aList 
 25          self.bName = bName 
 26          self.bList = bList 
 27           
 28          self.nStepsA=nStepsA 
 29          self.nStepsB=nStepsB 
 30          self.iLabelsX = iLabelsX 
 31          self.iLabelsY = iLabelsY 
 32          self.alphaInLineLabels = alphaInLineLabels 
 33           
 34          self.linewidth = linewidth 
 35          self.inLineLabels = inLineLabels 
 36          self.showLabels = showLabels 
 37          self.showGrid = showGrid 
 38          self.logX = logX 
 39          self.logY = logY 
 40          self.titleStr = titleStr 
 41          self.yLabelStr = yLabelStr 
 42          self.xLabelStr = xLabelStr 
 43           
 44          self.aValL = copy.deepcopy( aList ) 
 45          self.aValL.sort() 
 46          self.bValL = copy.deepcopy( bList ) 
 47          self.bValL.sort() 
 48           
 49          self.stepA = (self.aValL[-1]-self.aValL[0]) / float( self.nStepsA ) 
 50          self.stepB = (self.bValL[-1]-self.bValL[0]) / float( self.nStepsB ) 
 51           
 52          # make curves along const b direction 
 53          self.bCurveXYL = [] 
 54          for b in self.bValL: 
 55              aL = [] 
 56              xL = [] 
 57              yL = [] 
 58              for a in self.aValL: 
 59                  x,y = func(a,b) 
 60                  aL.append(a) 
 61                  xL.append(x) 
 62                  yL.append(y) 
 63              self.bCurveXYL.append( [b,aL,xL,yL] ) 
 64               
 65          # fill out curves 
 66          for b,aL,xL,yL in self.bCurveXYL: 
 67              i = len(aL)-1 
 68              while i>0: 
 69                  while aL[i]-aL[i-1]>self.stepA*1.2: 
 70                      a = aL[i]-self.stepA 
 71                      x,y = func(a,b) 
 72                      aL.insert(i,a) 
 73                      xL.insert(i,x) 
 74                      yL.insert(i,y) 
 75                  i = i -1 
 76           
 77          # make curves along const a direction 
 78          self.aCurveXYL = [] 
 79          for a in self.aValL: 
 80              bL = [] 
 81              xL = [] 
 82              yL = [] 
 83              for b in self.bValL: 
 84                  x,y = func(a,b) 
 85                  bL.append(b) 
 86                  xL.append(x) 
 87                  yL.append(y) 
 88              self.aCurveXYL.append( [a,bL,xL,yL] ) 
 89               
 90          # fill out curves 
 91          for a,bL,xL,yL in self.aCurveXYL: 
 92              i = len(bL)-1 
 93              while i>0: 
 94                  while bL[i]-bL[i-1]>self.stepB*1.2: 
 95                      b = bL[i]-self.stepB 
 96                      x,y = func(a,b) 
 97                      bL.insert(i,b) 
 98                      xL.insert(i,x) 
 99                      yL.insert(i,y) 
100                  i = i -1 
101                       
102               
103          #for ac in self.aCurveXYL: 
104          #    print ac 
105          #    print len( ac[-1] ) 
106               
107 -    def plotCarpet(self, afmt='red', bfmt='blue', figObj=None): 
108           
109          if figObj: 
110              f = figObj 
111          else: 
112              f = figure() 
113           
114          if self.showGrid: grid() 
115       
116          if self.yLabelStr: 
117              ylabel( self.yLabelStr ) 
118       
119          if self.xLabelStr: 
120              xlabel( self.xLabelStr ) 
121               
122          if self.titleStr: 
123              title(self.titleStr) 
124       
125          majorFormatter = FormatStrFormatter('%g') 
126          gca().yaxis.set_major_formatter(majorFormatter) 
127          majorFormatter = FormatStrFormatter('%g') 
128          gca().xaxis.set_major_formatter(majorFormatter) 
129           
130          for a,bL,xL,yL in self.aCurveXYL: 
131              #plot(xL, yL, label='%s=%g'%(self.aName,a), linewidth=self.linewidth , color=afmt) 
132              lblStr = '%s=%g'%(self.aName,a) 
133   
134              if self.logX and self.logY: 
135                  loglog(xL, yL, label=lblStr, linewidth=self.linewidth , color=afmt) 
136              elif self.logX: 
137                  semilogx(xL, yL, label=lblStr, linewidth=self.linewidth , color=afmt) 
138              elif self.logY: 
139                  semilogy(xL, yL, label=lblStr, linewidth=self.linewidth , color=afmt) 
140              else: 
141                  plot(xL, yL, label=lblStr, linewidth=self.linewidth , color=afmt) 
142           
143   
144          for b,aL,xL,yL in self.bCurveXYL: 
145              #plot(xL, yL, label='%s=%g'%(self.bName,b), linewidth=self.linewidth , color=bfmt) 
146              lblStr = '%s=%g'%(self.bName,b) 
147   
148              if self.logX and self.logY: 
149                  loglog(xL, yL, label=lblStr, linewidth=self.linewidth , color=bfmt) 
150              elif self.logX: 
151                  semilogx(xL, yL, label=lblStr, linewidth=self.linewidth , color=bfmt) 
152              elif self.logY: 
153                  semilogy(xL, yL, label=lblStr, linewidth=self.linewidth , color=bfmt) 
154              else: 
155                  plot(xL, yL, label=lblStr, linewidth=self.linewidth , color=bfmt) 
156               
157          ax = gca() 
158          #ax.apply_aspect() 
159           
160          xmin,xmax = ax.get_xlim()  
161          #print 'xmin,xmax',xmin,xmax 
162          ymin,ymax = ax.get_ylim() 
163          #print 'ymin,ymax',ymin,ymax 
164           
165          # save plot limits as attributes 
166          self.xmin, self.xmax, self.ymin, self.ymax = xmin, xmax, ymin, ymax 
167           
168          # make enough room around the carpet to make labels 
169           
170           
171          bbox = { 'pad':0, 'facecolor':'w', 'edgecolor':'w', \ 
172                  'clip_on':False, 'lw':0, 'alpha':0.01}  
173          if self.inLineLabels: 
174              bbox['alpha'] = self.alphaInLineLabels 
175          props = {'ha':'center', 'va':'center', 'bbox':bbox} 
176           
177          dx = (xmax-xmin)/50.0 
178          dy = (ymax-ymin)/50.0 
179          label = '' 
180          if self.iLabelsX > len(self.bValL)-2: 
181              self.iLabelsX = len(self.bValL)-2 
182          for a,bL,xL,yL in self.aCurveXYL: 
183              #props = {'ha':'%s'%haLabel, 'va':'%s'%vaLabel, 'color':afmt, 'rotation':0., 
184              #    'bbox':dict(facecolor='white', alpha=0.1, lw=0, clip_on=False)} 
185              if label: 
186                  label = '%g '%a 
187              else: 
188                  label = '%s=%g '%(self.aName,a) 
189               
190              i1 = bL.index( self.bValL[self.iLabelsX+1] ) 
191              i0 = bL.index( self.bValL[self.iLabelsX] ) 
192              i = (i1+i0)/2 
193               
194              self.placeLabel(i, xL, yL, afmt, label,  xmin, xmax, ymin, ymax, props) 
195   
196          label = '' 
197          if self.iLabelsY > len(self.aValL)-2: 
198              self.iLabelsY = len(self.aValL)-2 
199          for b,aL,xL,yL in self.bCurveXYL: 
200              #props = {'ha':'%s'%haLabel, 'va':'%s'%vaLabel, 'color':bfmt, 'rotation':0., 
201              #    'bbox':dict(facecolor='white', alpha=0.1, lw=0, clip_on=False)} 
202               
203              if label: 
204                  label = '%g '%b 
205              else: 
206                  label = '%s=%g '%(self.bName,b) 
207                   
208              ilast = aL.index( self.aValL[self.iLabelsY+1] ) 
209              ipen = aL.index( self.aValL[self.iLabelsY] ) 
210              i = (ilast+ipen)/2 
211               
212              self.placeLabel(i, xL, yL, bfmt, label,  xmin, xmax, ymin, ymax, props) 
213               
214 -    def placeLabel(self, i, xL, yL, fmt, label,  xmin, xmax, ymin, ymax, props): 
215           
216          if not self.showLabels: 
217              return 
218           
219          if len(xL)<2: 
220              return  
221          else: 
222              if i>= len(xL): 
223                  i = len(xL)-2 
224              dx = xL[i+1] - xL[i] 
225              dy = yL[i+1] - yL[i] 
226              fdx = 8. * dx / (xmax-xmin) 
227              fdy = 6. * dy / (ymax-ymin) 
228               
229              if fdy==0.0: 
230                  if fdx>0.0: 
231                      ang = 0 
232                  else: 
233                      ang = 0 
234              else: 
235                  try: 
236                      ang = int( atan( fdy / fdx ) * 180.0 / pi ) 
237                  except: 
238                      ang = 90 
239                   
240              xave = (xL[i+1] + xL[i])/2. 
241              yave = (yL[i+1] + yL[i])/2. 
242              offx = 0.0 
243              offy = 0.0 
244              props['va'] = 'center' 
245              props['ha'] = 'center' 
246               
247              if ang > 45.0: 
248                  offx = (xmax-xmin)/50.0 
249              elif ang < -45.0: 
250                  offx = (xmax-xmin)/50.0 
251              else: 
252                  offy = (ymax-ymin)/50.0 
253               
254              props['color']=fmt 
255              props['rotation'] = ang 
256               
257              if self.inLineLabels: 
258                  text(xave, yave, label , props) 
259              else: 
260                  text(xave+offx, yave+offy, label , props) 
261           
262   
263  if __name__ == "__main__":  #self test 
264       
265      from math import * 
266 -    def func3(a,b): 
267          x = a**2 + b 
268          y = 2*b**2 + 3*a 
269          return x,y 
270           
271      C = Carpet(func3, linewidth=1, titleStr='Simple Function Test', xLabelStr='Horizontal Axis', yLabelStr='Vertical Axis') 
272      C.plotCarpet(figObj = figure(), afmt='green', bfmt='purple') 
273      show() 
274