DISPATCH
plot_utils.py
1 # Pythn 2/3 compatibility
2 from __future__ import print_function
3 import sys
4 if sys.version_info[0]!=2:
5  from importlib import reload
6 
7 import numpy as np
8 import matplotlib.pyplot as pl
9 import matplotlib.animation as an
10 import dispatch_utils as du
11 reload(du)
12 
13 def plot_utils():
14  '''
15  Methods: image(), yz_plane()
16 
17  See http://matplotlib.org/users/colormaps.html for color maps
18  '''
19 
20 def simple_slice(iout=0,run='',data='../data',x=None,y=None,z=None,iv=0,verbose=0,pp=None):
21  if not pp:
22  pp=du.patches(iout=iout,run=run,data=data,x=x,y=y,z=z)
23  if (np.size(pp)==0):
24  print('no patches found')
25  return
26  if verbose:
27  for p in pp:
28  print('id={:04d} pos={}'.format(p.id,p.pos))
29  ix=None; iy=None; iz=None
30  if x:
31  p = pp[0]
32  ix=np.int((x-p.x[0])/p.ds[0]+0.5)
33  print('ix =',ix,' xr=',p.active_corners[0:2,0])
34  if y:
35  p = pp[1]
36  iy=np.int((y-p.y[0])/p.ds[1]+0.5)
37  print('iy =',iy,' yr=',p.active_corners[0:2,1])
38  if z:
39  p = pp[0]
40  iz=np.int((z-p.z[0])/p.ds[2]+0.5)
41  print('iz =',iz,' zr=',p.active_corners[0:2,2])
42  ll=du.limits(pp,ix=ix,iy=iy,iz=iz,iv=iv,verbose=verbose)
43  pl.clf()
44  p=pp[0]
45  l=p.nghost
46  u=l+p.n-1
47  if x:
48  pl.xlim(ll.ylim)
49  pl.ylim(ll.zlim)
50  pl.title('iout={} iv={} time={} x={}'.format(iout,iv,p.time,x))
51  for p in pp:
52  image(p.data[ix,l[1]:u[1],l[2]:u[2],iv],vmin=ll.vlim[0],vmax=ll.vlim[1],extent=p.extent[0])
53  elif y:
54  pl.xlim(ll.xlim)
55  pl.ylim(ll.zlim)
56  pl.title('iout={} iv={} time={} y={}'.format(iout,iv,p.time,y))
57  for p in pp:
58  image(p.data[l[0]:u[0],iy,l[2]:u[2],iv],vmin=ll.vlim[0],vmax=ll.vlim[1],extent=p.extent[1])
59  elif z:
60  pl.xlim(ll.xlim)
61  pl.ylim(ll.ylim)
62  pl.title('iout={} iv={} time={} z={}'.format(iout,iv,p.time,z))
63  for p in pp:
64  image(p.data[l[0]:u[0],l[1]:u[1],iz,iv],vmin=ll.vlim[0],vmax=ll.vlim[1],extent=p.extent[2])
65  pl.colorbar()
66 
67 def image(f,verbose=0,vlim=None,keep=False,**kwargs):
68  '''
69  Turn images the Fortran way; 1st index increases to the right,
70  2nd index upwards
71  '''
72  if not keep:
73  pl.clf()
74  if np.size(vlim)==2:
75  im=pl.imshow(f.squeeze().transpose(),origin='lower',interpolation='nearest',\
76  vmin=vlim[0],vmax=vlim[0],**kwargs)
77  else:
78  im=pl.imshow(f.squeeze().transpose(),origin='lower',interpolation='nearest',**kwargs)
79  if not keep:
80  pl.colorbar()
81  return im
82 
83 def in_extent(e,p):
84  ok=True
85  if p[0]<e[0] or p[0]>e[1]: ok=False
86  if p[1]<e[2] or p[1]>e[3]: ok=False
87  return ok
88 
89 def yx_plane(iout=-1,run='',data='../data',z=0.0,var='logd',cmap='coolwarm',\
90  keep=0, extent=None,vlim=None,fraction=0.1,labels=None,verbose=0,\
91  pp=None,**kwargs):
92  '''
93  Display xy-image for given z
94  '''
95  pl.show()
96  if type(pp)==type(None):
97  pp=du.patches(run=run,data=data,iout=iout)
98  if np.size(pp)==0:
99  print("no files found")
100  return
101  def inside(p, z):
102  return z<=p.active_corners[1][2] and z>=p.active_corners[0][2]
103  select=[p for p in pp if inside(p,z)]
104  if verbose: print(np.size(pp),'patches in plane' )
105  vmin=1e9; vmax=-vmin
106  xlim=[vmin,vmax]
107  ylim=[vmin,vmax]
108  l=np.zeros(3,dtype=np.int32)
109  u=l+np.array(p.n)
110  def interpolate(p,x,x0,dx,n,f):
111  s=(x-x0)/dx
112  i=max(0,min(n-2,int(s)))
113  s=s-i
114  if verbose>1:
115  fmt='id:{:6d} lev:{:2d} i:{:3d} p:{:6.2f}'
116  print (fmt.format(p.id,p.level,i,s))
117  return (1.0-s)*f[l[0]:u[0],l[1]:u[1],i]+s*f[l[0]:u[0],l[1]:u[1],i+1]
118  for p in select:
119  p.cache()
120  p.stripit()
121  p.vars()
122  p.f=p.var[var]
123  pz=p.z
124  if var=='u3' or var=='p3': pz=p.zs
125  p.f=interpolate(p,z,pz[0],p.ds[2],p.n[2],p.f)
126  p.f,e=du.rotate(p.f,p.extent[2])
127  p.extent[2]=e
128  xlim=[min(xlim[0],e[0]),max(xlim[1],e[1])]
129  ylim=[min(ylim[0],e[2]),max(ylim[1],e[3])]
130  if extent != None:
131  if np.size(extent)==1:
132  w=extent
133  extent=[-w,w,-w,w]
134  xlim=extent[0:2]
135  ylim=extent[2:4]
136  def overlaps(p):
137  does = ((p.extent[0][0] > xlim[0] and p.extent[0][0] < xlim[1]) \
138  or (p.extent[0][1] > xlim[0] and p.extent[0][1] < xlim[1])) \
139  and ((p.extent[0][2] > ylim[0] and p.extent[0][2] < ylim[1]) \
140  or (p.extent[0][3] > ylim[0] and p.extent[0][3] < ylim[1]))
141  return does
142  select=[p for p in select if overlaps(p)]
143  if verbose: print(np.size(select),'patches remain')
144  else:
145  extent=[xlim[0],xlim[1],ylim[0],ylim[1]]
146  if np.size(vlim)==2:
147  vmin=vlim[0]
148  vmax=vlim[1]
149  else:
150  for p in select:
151  vmin=min(vmin,p.f.min())
152  vmax=max(vmax,p.f.max())
153  if verbose:
154  print(p.id,p.f.min(),p.f.max(),vmin,vmax)
155  if not keep:
156  pl.clf()
157  for p in select:
158  image(p.f,extent=p.extent[2],vmin=vmin,vmax=vmax,cmap=cmap,keep=1)
159  yx=[p.active_corners[0][1],p.active_corners[0][0]]
160  if labels and in_extent(extent,yx):
161  pl.text(yx[0],yx[1],p.id)
162  pl.title('var='+var+' t={:3.2f} z={:2.1f}'.format(p.time,z))
163  if not keep:
164  pl.colorbar(shrink=0.8,aspect=15,fraction=fraction,pad=0.02)
165  pl.xlim(xlim); pl.ylim(ylim); pl.xlabel('y'); pl.ylabel('x')
166  pl.tight_layout()
167  pl.draw()
168 
169 def yz_plane(iout=-1,run='',data='../data',var='logd',x=0.0,cmap='coolwarm',\
170  keep=0, extent=None,vlim=None,fraction=0.1,labels=None,verbose=0,overlap=0.01,**kwargs):
171  '''
172  Display yz-image for given x
173  '''
174  pl.show()
175  pp=du.patches(run=run,data=data,iout=iout,overlap=overlap)
176  if np.size(pp)==0:
177  print("no files found")
178  return
179  def inside(p,x):
180  return x<=p.active_corners[1][0] and x>=p.active_corners[0][0]
181  pp=[p for p in pp if inside(p,x)]
182  if verbose: print(np.size(pp),'patches in plane')
183  vmin=1e9; vmax=-vmin
184  ylim=[vmin,vmax]
185  zlim=[vmin,vmax]
186  l=np.zeros(3,dtype=np.int32)
187  u=l+np.array(p.n)
188  def interpolate(p,x,x0,dx,n,f):
189  s=(x-x0)/dx
190  i=max(0,min(n-2,int(s)))
191  s=s-i
192  if verbose>1:
193  fmt='id:{:6d} lev:{:2d} i:{:3d} p:{:6.2f}'
194  print (fmt.format(p.id,p.level,i,s))
195  return (1.0-s)*f[i,l[1]:u[1],l[2]:u[2]]+s*f[i+1,l[1]:u[1],l[2]:u[2]]
196  for p in pp:
197  p.cache()
198  p.stripit()
199  p.vars()
200  p.f=p.var[var]
201  px=p.x
202  if var=='u1' or var=='p1': px=p.xs
203  p.f=interpolate(p,x,px[0],p.dx[0],p.n[0],p.f)
204  e=p.extent[0]
205  ylim=[min(ylim[0],e[0]),max(ylim[1],e[1])]
206  zlim=[min(zlim[0],e[2]),max(zlim[1],e[3])]
207  if extent:
208  if np.size(extent)==1:
209  w=extent
210  extent=[-w,w,-w,w]
211  ylim=extent[0:2]
212  zlim=extent[2:4]
213  def overlaps(p):
214  does = ((p.extent[0][0] > ylim[0] and p.extent[0][0] < ylim[1]) \
215  or (p.extent[0][1] > ylim[0] and p.extent[0][1] < ylim[1])) \
216  and ((p.extent[0][2] > zlim[0] and p.extent[0][2] < zlim[1]) \
217  or (p.extent[0][3] > zlim[0] and p.extent[0][3] < zlim[1]))
218  return does
219  pp=[p for p in pp if overlaps(p)]
220  if verbose: print(np.size(pp),'patches remain')
221  else:
222  extent=[ylim[0],ylim[1],zlim[0],zlim[1]]
223  if np.size(vlim)==2:
224  vmin=vlim[0]
225  vmax=vlim[1]
226  else:
227  for p in pp:
228  vmin=min(vmin,p.f.min())
229  vmax=max(vmax,p.f.max())
230  if not keep:
231  pl.clf()
232  for p in pp:
233  image(p.f,extent=p.extent[0],vmin=vmin,vmax=vmax,cmap=cmap,keep=1)
234  yz=[p.active_corners[0][1],p.active_corners[0][2]]
235  if labels and in_extent(extent,yz):
236  pl.text(yz[0],yz[1],p.id)
237  pl.title('var='+var+' t={:3.2f} x={:2.1f}'.format(p.time,x))
238  if not keep:
239  pl.colorbar(shrink=0.8,aspect=15,fraction=fraction,pad=0.02)
240  pl.xlim(ylim); pl.ylim(zlim); pl.xlabel('y'); pl.ylabel('z')
241  pl.tight_layout()
242  pl.show()
243 
244 def xz_plane(iout=-1,run='',data='../data',var='logd',y=0.0,cmap='coolwarm',\
245  keep=0, extent=None,vlim=None,fraction=0.1,labels=None,verbose=0,**kwargs):
246  '''
247  Display xz-image for given y
248  '''
249  pl.show()
250  #print('data:',data,' run:',run,' iout:',iout)
251  pp=du.patches(run=run,data=data,iout=iout)
252  if np.size(pp)==0:
253  print("no files found")
254  return
255  def inside(p,y):
256  return y<=p.active_corners[1][1] and y>=p.active_corners[0][1]
257  pp=[p for p in pp if inside(p,y)]
258  if verbose: print(np.size(pp),'patches in plane')
259  vmin=1e9; vmax=-vmin
260  xlim=[vmin,vmax]
261  zlim=[vmin,vmax]
262  l=np.zeros(3,dtype=np.int32)
263  u=l+np.array(p.n)
264  def interpolate(p,y,y0,dy,n,f):
265  s=(y-y0)/dy
266  i=max(0,min(n-2,int(s)))
267  s=s-i
268  if verbose>1:
269  fmt='id:{:6d} lev:{:2d} i:{:3d} p:{:6.2f}'
270  print (fmt.format(p.id,p.level,i,s))
271  return (1.0-s)*f[l[0]:u[0],i,l[2]:u[2]]+s*f[l[0]:u[0],i+1,l[2]:u[2]]
272  for p in pp:
273  p.cache()
274  p.stripit()
275  p.vars()
276  p.f=p.var[var]
277  py=p.y
278  if var=='u2' or var=='p2': py=p.ys
279  p.f=interpolate(p,y,py[0],p.dx[0],p.n[0],p.f)
280  e=p.extent[1]
281  xlim=[min(xlim[0],e[0]),max(xlim[1],e[1])]
282  zlim=[min(zlim[0],e[2]),max(zlim[1],e[3])]
283  if extent:
284  if np.size(extent)==1:
285  w=extent
286  extent=[-w,w,-w,w]
287  xlim=extent[0:2]
288  zlim=extent[2:4]
289  def overlaps(p):
290  does = ((p.extent[1][0] > xlim[0] and p.extent[1][0] < xlim[1]) \
291  or (p.extent[1][1] > xlim[0] and p.extent[1][1] < xlim[1])) \
292  and ((p.extent[1][2] > zlim[0] and p.extent[1][2] < zlim[1]) \
293  or (p.extent[1][3] > zlim[0] and p.extent[1][3] < zlim[1]))
294  return does
295  #pp=[p for p in pp if overlaps(p)]
296  if verbose:(print(np.size(pp),'patches remain'))
297  else:
298  extent=[xlim[0],xlim[1],zlim[0],zlim[1]]
299  if np.size(vlim)==2:
300  vmin=vlim[0]
301  vmax=vlim[1]
302  else:
303  for p in pp:
304  vmin=min(vmin,p.f.min())
305  vmax=max(vmax,p.f.max())
306  if not keep:
307  pl.clf()
308  for p in pp:
309  #print (np.size(p.x),p.x[0],p.x[-1],p.extent[1])
310  image(p.f,extent=p.extent[1],vmin=vmin,vmax=vmax,cmap=cmap,keep=1)
311  xz=[p.active_corners[0][0],p.active_corners[0][2]]
312  if labels and in_extent(extent,xz):
313  pl.text(xz[0],xz[1],p.id)
314  pl.title('var='+var+' t={:3.2f} y={:2.1f}'.format(p.time,y))
315  if not keep:
316  pl.colorbar(shrink=0.8,aspect=15,fraction=fraction,pad=0.02)
317  pl.xlim(xlim); pl.ylim(zlim); pl.xlabel('x'); pl.ylabel('z')
318  pl.tight_layout()
319  pl.show()
320 
321 def yx_scan(fig,run='',data='../data',iout=0,dz=50,var='logd',cmap='coolwarm',\
322  verbose=0,fraction=0.1,extent=None,repeat=0,vlim=None,interval=200,\
323  labels=None,**kwargs):
324  '''
325  Display yx-image in steps of dz in z
326  '''
327  pl.show()
328  pp=du.patches(run=run,data=data,iout=iout)
329  co=du.corners(pp)
330  zmin=co[0,2]
331  zmax=co[1,2]
332  zmin=(np.floor(zmin/dz+0.5)+0.5)*dz
333  zmax=(np.floor(zmax/dz+0.5)-0.5)*dz
334  if extent != None:
335  if np.size(extent)==6:
336  zmin=extent[4]
337  zmax=extent[5]
338  #zmin=np.floor(zmin/dz+0.5)*dz
339  #zmax=np.floor(zmax/dz+0.5)*dz
340  print(zmin,zmax)
341  def anim(z):
342  def inside(p, z):
343  ac=p.active_corners
344  ok = z<=ac[1][2] and z>=ac[0][2]
345  if extent != None:
346  ok = ok and ac[1][0]>extent[0] and ac[0][0]<extent[1]
347  ok = ok and ac[1][1]>extent[2] and ac[0][1]<extent[3]
348  return ok
349  select=[p for p in pp if inside(p,z)]
350  p=select[0]
351  title='var='+var+' t={:d} (:,:,{:d})'.format(np.int(p.time),np.int(z))
352  vmin=1e9
353  vmax=-vmin
354  xlim=[vmin,vmax]
355  ylim=[vmin,vmax]
356  l=np.zeros(3,dtype=np.int32)
357  u=l+np.array(p.n)
358  def interpolate(p,x,x0,dx,n,f):
359  s=(x-x0)/dx
360  i=max(0,min(n-2,int(s)))
361  s=s-i
362  return (1.0-s)*f[l[0]:u[0],l[1]:u[1],i]+s*f[l[0]:u[0],l[1]:u[1],i+1]
363  for p in select:
364  p.cache()
365  p.stripit()
366  p.vars()
367  p.f=p.var[var]
368  if p.ioformat==1:
369  z0=p.active_corners[0][2]+0.5*p.dx[2]
370  else:
371  z0=p.active_corners[0][2]
372  if var=='u3' or var=='p3': z0=z0-0.5*p.dx[2]
373  p.f=interpolate(p,z,z0,p.dx[2],p.n[2],p.f)
374  e=p.extent[2]
375  p.f,e=du.rotate(p.f,e)
376  p.ext=e
377  vmin=min(vmin,p.f.min())
378  vmax=max(vmax,p.f.max())
379  xlim=[min(xlim[0],e[0]),max(xlim[1],e[1])]
380  ylim=[min(ylim[0],e[2]),max(ylim[1],e[3])]
381  pl.clf()
382  if extent==None:
383  pl.xlim(co[:,1]) # y-axis is horizontal
384  pl.ylim(co[:,0]) # x-axis is vertical
385  else:
386  pl.xlim(extent[2:4]) # y-axis is horizontal
387  pl.ylim(extent[0:2]) # x-axis is vertical
388  pl.xlabel('y')
389  pl.ylabel('x')
390  pl.tight_layout()
391  pl.title(title)
392  if np.size(vlim)==2:
393  vmin=vlim[0]
394  vmax=vlim[1]
395  for p in select:
396  im=image(p.f,extent=p.ext,vmin=vmin,vmax=vmax,cmap=cmap,keep=1,**kwargs)
397  if labels:
398  pl.text(p.active_corners[0][1],p.active_corners[0][0],p.id)
399  pl.colorbar(shrink=0.9,aspect=15,fraction=fraction,pad=0.02)
400  pl.draw()
401  pl.tight_layout()
402  return im
403  a=an.FuncAnimation(fig,anim,np.arange(zmin,zmax+dz,dz),\
404  repeat=repeat,interval=interval,**kwargs)
405  return a
406 
407 def yz_scan(fig,run='',data='../data',iout=0,dx=50,var='logd',cmap='coolwarm',\
408  verbose=0,fraction=0.1,extent=None,repeat=0,vlim=None,interval=200,\
409  labels=None,**kwargs):
410  '''
411  Display yz-image in steps of dx in x
412  '''
413  pl.show()
414  pp=du.patches(run=run,data=data,iout=iout)
415  co=du.corners(pp)
416  xmin=co[0,0]
417  xmax=co[1,0]
418  xmin=(np.floor(xmin/dx+0.5)+0.5)*dx
419  xmax=(np.floor(xmax/dx+0.5)-0.5)*dx
420  if extent != None:
421  if np.size(extent)==6:
422  xmin=extent[0]
423  xmax=extent[1]
424  xmin=np.floor(xmin/dx+0.5)*dx
425  xmax=np.floor(xmax/dx+0.5)*dx
426  def anim(x):
427  def inside(p, x):
428  ac=p.active_corners
429  ok = x<=ac[1][0] and x>=ac[0][0]
430  if extent != None:
431  ok = ok and ac[1][1]>extent[2] and ac[0][1]<extent[3] # y-axis is horizontal
432  ok = ok and ac[1][2]>extent[4] and ac[0][2]<extent[5] # z-axis is vertical
433  return ok
434  select=[p for p in pp if inside(p,x)] # and p.level<8]
435  p=select[0]
436  title='var='+var+' t={:d} ({:d},:,:)'.format(np.int(p.time),np.int(x))
437  vmin=1e9
438  vmax=-vmin
439  zlim=[vmin,vmax]
440  ylim=[vmin,vmax]
441  l=np.zeros(3,dtype=np.int32)
442  u=l+np.array(p.n)
443  def interpolate(p,x,x0,dx,n,f):
444  s=(x-x0)/dx
445  i=max(0,min(n-2,int(s)))
446  s=s-i
447  return (1.0-s)*f[i,l[1]:u[1],l[2]:u[2]]+s*f[i+1,l[1]:u[1],l[2]:u[2]]
448  for p in select:
449  p.cache()
450  p.stripit()
451  p.vars()
452  f=p.var[var]
453  p.f=interpolate(p,x,p.active_corners[0][0]+0.5*p.dx[0],p.dx[0],p.n[0],f)
454  fmin=np.float(p.f.min()); fmax=np.float(p.f.max())
455  vmin=min(vmin,fmin)
456  vmax=max(vmax,fmax)
457  e=p.extent[0]
458  zlim=[min(zlim[0],e[0]),max(zlim[1],e[1])]
459  ylim=[min(ylim[0],e[2]),max(ylim[1],e[3])]
460  pl.clf()
461  if extent==None:
462  pl.xlim(co[:,1])
463  pl.ylim(co[:,2])
464  else:
465  pl.xlim(extent[2:4])
466  pl.ylim(extent[4:6])
467  pl.xlabel('y')
468  pl.ylabel('z')
469  pl.tight_layout()
470  pl.title(title)
471  if np.size(vlim)==2:
472  vmin=vlim[0]
473  vmax=vlim[1]
474  for p in select:
475  im=image(p.f,extent=p.extent[0],vmin=vmin,vmax=vmax,cmap=cmap,keep=1,**kwargs)
476  if labels:
477  pl.text(p.active_corners[0][1],p.active_corners[0][2],p.id)
478  pl.colorbar(shrink=0.75,aspect=15,fraction=fraction,pad=0.02)
479  pl.draw()
480  pl.tight_layout()
481  return im
482  a=an.FuncAnimation(fig,anim,np.arange(xmin,xmax+dx,dx),repeat=repeat,**kwargs)
483  return a
Definition: image.py:1