DISPATCH
_graphics.py
1 # -*- coding: utf-8 -*-
2 
3 # Pythn 2/3 compatibility
4 from __future__ import print_function
5 
6 import numpy as np
7 import matplotlib
8 import matplotlib.pyplot as pl
9 
10 import dispatch as di
11 import dispatch.select as ds
12 
13 def _kw_extract(kw,dict):
14  """ if keys from dict occur in kw, pop them """
15  for key,value in dict.items():
16  dict[key]=kw.pop(key,value)
17  return kw,dict
18 
19 def imshow(f,colorbar=None,origin='lower',interpolation='nearest',
20  cmap='coolwarm',verbose=0,update=False,**kw):
21  """
22  imshow with default options and updateable colorbar. Example use:
23 
24  import matplotlib.pyplot as pl
25  import dispatch.select as ds
26  import dispatch.graphics as dg
27  ...
28  s=dispatch.snapshot(1)
29  dg.imshow(ds.unigrid_plane(s,iv=0))
30  cb=pl.colorbar()
31  dg.imshow(ds.unigrid_plane(s,iv=1),cb)
32 
33  The second call updates the colorbar from the 1st call
34  """
35  if verbose>0:
36  print ('min:',f.min(),' max:',f.max())
37  pl.imshow(np.transpose(f),origin=origin,interpolation=interpolation,cmap=cmap,**kw)
38  """ Check if there is an existing colorbar """
39  cb=None
40  ims=pl.gca().images
41  if len(ims)>0:
42  for im in ims:
43  if type(im.colorbar)==matplotlib.colorbar.Colorbar:
44  cb=im.colorbar
45  """ If a colorbar exists, update the limits - if not, make one """
46  if cb:
47  cb.set_clim(f.min(),f.max())
48  cb.draw_all()
49  else:
50  pl.colorbar()
51 
52 def plot(f,**kw):
53  """ plot(f) allows f to be (x,y) tuple """
54  if type(f)==tuple:
55  x,y=f
56  pl.plot(x,y,**kw)
57  else:
58  pl.plot(f,**kw)
59 
60 def plot_values_along(pp,pt=[0.5,0.5,0.5],**kw):
61  """ Plot values along direction dir={0,1,2}, through point pt=[x,y,z] """
62  kv = {'dir':0, 'verbose':0, 'all':False, 'iv':0, 'i4':0, 'var':None}
63  kw,kv = _kw_extract(kw,kv)
64  plot(ds.values_along(pp,pt,iv=kv['iv'],dir=kv['dir'],all=kv['all']),**kw)
65 
66 def plot_patch_values_along(pp_in,pt=[0.5,0.5,0.5],hold=False,**kw):
67  """ Plot values along direction dir={0,1,2}, through point pt=[x,y,z] """
68  kv = {'dir':0, 'verbose':0, 'all':False, 'iv':0, 'i4':0, 'var':None}
69  kw,kv = _kw_extract(kw,kv)
70  pp = ds.patches_along(pp_in,pt,dir=kv['dir'],verbose=kv['verbose'])
71  xmin,xmax = ds.minmax_patches(pp,dir=kv['dir'])
72  if not hold:
73  pl.xlim(xmin,xmax)
74  for p in pp:
75  plot(ds.values_in(p,pt,**kv),**kw)
76 
77 def _rotate(d,e):
78  '''Rotate data and extent to landscape format'''
79  d=d.transpose()
80  e=[e[2],e[3],e[0],e[1]]
81  return d,e
82 
83 def power2d(f,*kw):
84  """plot power spectrum of 2D array"""
85  ft2=np.abs(np.fft.fft2(f))**2
86  m=f.shape
87  k=np.meshgrid(range(m[0]),range(m[1]))
88  k=np.sqrt(k[0]**2+k[1]**2)
89  a=2
90  k0=1.0/a**0.5
91  k1=1.0*a**0.5
92  power=[]
93  kk=[]
94  while(k1 <= m[0]//2):
95  kk.append((k0*k1)**0.5)
96  w=np.where((k>k0) & (k <= k1))
97  power.append(ft2[w].sum())
98  k0=k1
99  k1=k1*a
100  pl.loglog(kk,power,*kw)
101 
102 def show_plot(figure_id=None):
103  """ raise a figure to the top """
104  if figure_id is None:
105  fig = pl.gcf()
106  else:
107  # do this even if figure_id == 0
108  fig = pl.figure(num=figure_id)
109  pl.show()
110  pl.pause(1e-9)
111  fig.canvas.manager.window.activateWindow()
112  fig.canvas.manager.window.raise_()
113 
114 def pause(time=1e-6):
115  """ replace the normal pause, w/o raising window() """
116  pl.draw()
117  pl.gcf().canvas.start_event_loop(time)
118 
119 def image_plane(snap,x=None,y=None,z=None,iv=0,verbose=1):
120  if x: i=0
121  if y: i=1
122  if z: i=2
123  xyz=(x,y,z)
124  sdir=['x','y','z']
125  labels=[('y','z'),('z','x'),('y','z')]
126  pp=snap.patches_in(x=x,y=y,z=z)
127  pl.clf()
128  ll={}
129  p=pp[0]
130  f=p.plane(x=x,y=y,z=z,iv=iv)
131  fmin=f.min()
132  fmax=f.max()
133  e=p.extent[i]
134  emin=np.array((e[0],e[2]))
135  emax=np.array((e[1],e[3]))
136  for p in pp:
137  im=p.plane(x=x,y=y,z=z,iv=iv)
138  fmin=min(fmin,im.min())
139  fmax=max(fmax,im.max())
140  e=p.extent[i]
141  emin=np.minimum(emin,np.array((e[0],e[2])))
142  emax=np.maximum(emax,np.array((e[1],e[3])))
143  ll[p.id]=(im,e,p.level)
144  for id in sorted(ll.keys()):
145  im,e,level=ll[id]
146  print(id,level)
147  imshow(im,extent=e,vmin=fmin,vmax=fmax,hold=True)
148  pl.colorbar()
149  pl.xlim(emin[0],emax[0])
150  pl.ylim(emin[1],emax[1])
151  pl.title('{}={:.3f}'.format(sdir[i],xyz[i]))
152  pl.xlabel(labels[i][0])
153  pl.ylabel(labels[i][1])
154 
155 def pdf_d(iout,run='',data='../data',iv='d',i4=0,nbin=100,xlim=[-4,3],lnd=False):
156  """ Plot the PDF of the density """
157  s = di.snapshot(iout,run=run,data=data)
158  n = nbin
159  bins = np.linspace(xlim[0],xlim[1],n+1)
160  htot = 0.0
161  i = 0
162  for p in s.patches:
163  i += 1
164  if i%1000==0:
165  print('{:.1f}%'.format(i/len(s.patches)*100.0))
166  d = p.var(iv,i4=i4)
167  if lnd:
168  logd = d/np.log(10.)
169  else:
170  logd = np.log10(d)
171  h,e = np.histogram(logd,bins=bins)
172  htot += h
173  pl.hist(bins[0:n],bins=bins,weights=htot,log=True,density=True)
174  return bins,htot
175 
176 def unigrid_plane(sn,**kw):
177  import dispatch.select as dse
178  imshow(dse.unigrid_plane(sn,**kw))
179 
180 
181 def amr_plane(sn,iv='d',z=0.5,log=False,mesh=False,ident=False,cmap=None,
182  title=None,verbose=0):
183  def contains(p,z):
184  return p.level >= 6 and abs(z-p.position[2]) <= p.size[2]/2
185  def outline(e):
186  x=[e[0],e[0],e[1],e[1],e[0]]
187  y=[e[2],e[3],e[3],e[2],e[2]]
188  pl.plot(x,y,color='grey')
189  def interpolate(p,iv,z,log=False):
190  n=p.n[2]
191  dz=p.size[2]/n
192  z0=p.position[2]-p.size[2]/2+dz/2
193  w=(z-z0)/dz
194  i=int(w)
195  i=max(0,min(i,n-2))
196  w=w-i
197  f=p.var(iv)
198  if iv==0 or iv==4 or iv=='d' or iv=='e':
199  log=True
200  if log:
201  f=np.log(f)
202  f=(1-w)*f[:,:,i]+w*f[:,:,i+1]
203  if log:
204  f=np.exp(f)
205  return f
206  pp=[p for p in sn.patches if contains(p,z)]
207  vmin=1e20
208  vmax=0.0
209  for p in pp:
210  d=interpolate(p,iv,z)
211  vmin=min(vmin,d.min())
212  vmax=max(vmax,d.max())
213  if verbose>0:
214  print(p.id,d.min(),d.max())
215  pl.xlim(0,1); pl.ylim(0,1)
216  if log:
217  vmin=np.log10(vmin)
218  vmax=np.log10(vmax)
219  for p in pp:
220  e=p.extent[2]
221  d=interpolate(p,iv,z)
222  if log:
223  d=np.log10(interpolate(p,iv,z))
224  pl.imshow(d.transpose(),extent=e,
225  vmin=vmin,vmax=vmax,origin='lower',cmap=cmap)
226  x0,x1,y0,y1=p.extent[2]
227  if mesh:
228  outline(p.extent[2])
229  if ident:
230  x0,x1,y0,y1=p.extent[2]
231  pl.text(x0,y0,'{}'.format(p.id))
232  pl.colorbar()
233  if title:
234  pl.title(title)
235  else:
236  pl.title('iv={}, z={:.2f}, t={:.4f}'.format(iv,z,sn.time))
237  pl.tight_layout()
238 
239 def plot_pdf(pdf,**kwargs):
240  """ Plot the output of the dispatch.pdf() procedure """
241  pl.hist(pdf.bins,bins=pdf.bins,weights=pdf.counts,**kwargs)
242  return pdf.time
243 
244 def pdf(io,run,data,log=True,**kwargs):
245  """ Shortcut of dispatch.graphics.plot_pdf(dispatch.pdf(io,run,data)) """
246  tmp=di.pdf(io,run,data,**kwargs)
247  plot_pdf(tmp,**kwargs)
248