DISPATCH
_select.py
1 # -*- coding: utf-8 -*-
2 
3 # Pythn 2/3 compatibility
4 from __future__ import print_function
5 
6 import numpy as np
7 from time import time
8 import dispatch
9 
10 def minloc(a):
11  ''' Return the 3-D index with the minimum of array a'''
12  i = a.argmin()
13  return np.unravel_index(i,np.shape(a))
14 
15 def maxloc(a):
16  ''' Return the 3-D index with the maximum of array a'''
17  i = a.argmax()
18  return np.unravel_index(i,np.shape(a))
19 
20 def patch_at(pp,point=[0.5,0.5,0.5],verbose=0):
21  ''' Find the patch that contains a given point
22  '''
23  level=-1
24  p1=None
25  for p in pp:
26  if is_inside(p,point,verbose):
27  if p.level>level:
28  p1=p
29  level=p.level
30  return p1
31 
32 def is_inside(p,point=[0.5,0.5,0.5],verbose=0):
33  ''' Return True/False depending on if the point is inside the patch
34  '''
35  left=point>=p.llc_cart
36  right=point<=p.llc_cart+p.size
37  return left.all() and right.all()
38 
39 def count_inside(p,point=[0.5,0.5,0.5],verbose=0):
40  ''' Return the number of dimension in which point is inside the patch
41  '''
42  left=point>=p.llc_cart
43  right=point<=p.llc_cart+p.size
44  count=0
45  for i in range(3):
46  if left[i] and right[i]:
47  count+=1
48  return count
49 
50 def patches_along(pp,point=[0.5,0.5,0.5],dir=0,verbose=0):
51  ''' Get the patches along a given direction through a point
52  '''
53  if isinstance(pp,dispatch._dispatch.snapshot):
54  pp=pp.patches
55  pt=np.copy(point)
56  if verbose>0:
57  np.set_printoptions(formatter={'float_kind': '{:6.3f}'.format})
58 
59  p=patch_at(pp,pt)
60  if p:
61  p1=p
62  while p:
63  if verbose>1:
64  print('fwd {:4d} {} {}'.format(p.id,p.position,p.size))
65  pt[dir]=p.position[dir]-p.size[dir]*0.6
66  p=patch_at(pp,pt)
67  if p:
68  p1=p
69  p=p1
70  out=[]
71  while p:
72  out.append(p)
73  if verbose>1:
74  print('rev {:4d} {} {}'.format(p.id,p.position,p.size))
75  pt[dir]=p.position[dir]+p.size[dir]*0.6
76  p=patch_at(pp,pt)
77  if verbose>0:
78  for p in out:
79  print('id: {:4d} position: {}'.format(p.id,p.position))
80  else:
81  out=None
82  np.set_printoptions(formatter=None)
83  return out
84 
85 def patches_along_x(pp,point=[0.5,0.5,0.5],verbose=0):
86  ''' Get the patches along the x direction through a point '''
87  patches_along(pp,point,dir=0,verbose=verbose)
88 
89 def patches_along_y(pp,point=[0.5,0.5,0.5],verbose=0):
90  ''' Get the patches along the y direction through a point '''
91  patches_along(pp,point,dir=1,verbose=verbose)
92 
93 def patches_along_z(pp,point=[0.5,0.5,0.5],verbose=0):
94  ''' Get the patches along the z direction through a point '''
95  patches_along(pp,point,dir=2,verbose=verbose)
96 
97 def indices_and_weights(p,point=[0.5,0.5,0.5],iv=0):
98  ''' Return indices and interpolation weights for a point in a patch p
99  '''
100  x0=p.x[0]
101  y0=p.y[0]
102  z0=p.z[0]
103  if 'p1' in p.idx.dict.keys():
104  if iv==p.idx.p1: x0=p.xs[0]
105  if iv==p.idx.p2: y0=p.ys[0]
106  if iv==p.idx.p3: z0=p.zs[0]
107  if 'b1' in p.idx.dict.keys():
108  if iv==p.idx.b1: x0=p.xs[0]
109  if iv==p.idx.b2: y0=p.ys[0]
110  if iv==p.idx.b3: z0=p.zs[0]
111  corner=np.array([x0,y0,z0])
112  weights=(point-corner)/p.ds
113  indices=np.array(weights,dtype=np.int32)
114  #print('before:',indices,weights)
115  for i in range(3):
116  indices[i]=max(0,min(p.n[i]-2,indices[i]))
117  weights=weights-indices
118  #print(' after:',indices,weights)
119  return indices, weights
120 
121 def values_in(p,point=[0.5,0.5,0.5],dir=0,iv=0,i4=0,var=None,verbose=0,all=0):
122  ''' Return s,f(s) with s the coordinates and f the values in the iv
123  slot of data, taken along the direction v -- so far restricted
124  to axis values
125  '''
126 
127  ss=[]
128  ff=[]
129  ii,w=indices_and_weights(p,point,iv)
130  data=p.var(iv,i4=i4)
131  ione = (0,1)[p.gn[0] > 1]
132  jone = (0,1)[p.gn[1] > 1]
133  kone = (0,1)[p.gn[2] > 1]
134 
135  if verbose:
136  print('{:3d} {} {} {:5.1f} {:4.1f} {:4.1f}'.\
137  format(p.id,iv,ii,w[0],w[1],w[2]))
138  if not p.guard_zones:
139  m=np.zeros(3,dtype=np.int32)
140  n=p.n
141  elif all:
142  m=np.zeros(3,dtype=np.int32)
143  n=p.gn
144  else:
145  m=p.ng
146  n=p.n
147  if type(var)==str:
148  iv=p.idx.dict[var]
149  if dir==0:
150  j=ii[1]
151  k=ii[2]
152  for i in range(m[0],m[0]+n[0]):
153  f1=data[i,j,k ]*(1-w[1])+data[i,j+jone,k, ]*w[1]
154  f2=data[i,j,k+kone]*(1-w[1])+data[i,j+jone,k+kone]*w[1]
155  f=f1*(1-w[2])+f2*w[2]
156  if hasattr(p,'p1'):
157  if iv==p.idx.p1:
158  ss.append(p.xs[i])
159  else:
160  ss.append(p.x[i])
161  else:
162  ss.append(p.x[i])
163  ff.append(f)
164  elif dir==1:
165  i=ii[0]
166  k=ii[2]
167  for j in range(m[1],m[1]+n[1]):
168  f1=data[i,j,k ]*(1-w[0])+data[i+ione,j,k ]*w[0]
169  f2=data[i,j,k+kone]*(1-w[0])+data[i+ione,j,k+kone]*w[0]
170  f=f1*(1-w[2])+f2*w[2]
171  if hasattr(p,'p2'):
172  if iv==p.idx.p2:
173  ss.append(p.ys[j])
174  else:
175  ss.append(p.y[j])
176  else:
177  ss.append(p.y[j])
178  ff.append(f)
179  else:
180  i=ii[0]
181  j=ii[1]
182  for k in range(m[2],m[2]+n[2]):
183  f1=data[i,j ,k]*(1-w[0])+data[i+ione,j ,k]*w[0]
184  f2=data[i,j+jone,k]*(1-w[0])+data[i+ione,j+jone,k]*w[0]
185  f=f1*(1-w[1])+f2*w[1]
186  if hasattr(p,'p3'):
187  if iv==p.idx.p3:
188  ss.append(p.zs[k])
189  else:
190  ss.append(p.z[k])
191  else:
192  ss.append(p.z[k])
193  ff.append(f)
194  ss=np.array(ss)
195  ff=np.array(ff)
196  return ss,ff
197 
198 def values_along(pp,point=[0.5,0.5,0.5],dir=0,iv=0,var=None,verbose=0,all=0):
199  ''' Return s,f(s) with s the coordinates and f the values in the iv
200  slot of data, taken along the direction v -- so far restricted
201  to axis values
202  '''
203  patches=patches_along(pp,point,dir=dir,verbose=verbose)
204  ss=[]
205  ff=[]
206  for p in patches:
207  ss1,ff1=values_in(p,point,dir=dir,iv=iv,var=var,all=all,verbose=verbose)
208  for s,f in zip(ss1,ff1):
209  ss.append(s)
210  ff.append(f)
211  ss=np.array(ss)
212  ff=np.array(ff)
213  return ss,ff
214 
215 def minmax_patch(p,dir=0):
216  smin = p.llc_cart[dir] - p.ng[dir]*p.ds[dir]
217  smax = smin + p.size[dir] + p.ng[dir]*p.ds[dir]
218  return (smin,smax)
219 
220 def minmax_patches(pp,dir=0):
221  smin,smax=minmax_patch(pp[0],dir=dir)
222  for p in pp:
223  s=minmax_patch(p,dir=dir)
224  smin=min(smin,s[0])
225  smax=max(smax,s[1])
226  return (smin,smax)
227 
228 def _corners(p,active=1):
229  c=np.array([p.llc_cart,p.llc_cart+p.size])
230  if not active:
231  c=c+np.array([-p.ng*p.ds,+p.ng*p.ds])
232  return c
233 
234 def corners(pp,active=1):
235  '''Corners of patch collection'''
236  if type(pp)==list:
237  c=_corners(pp[0],active=active)
238  for p in pp[1:]:
239  cp=_corners(p,active=active)
240  for i in range(3): c[0,i]=min(c[0,i],cp[0,i])
241  for i in range(3): c[1,i]=max(c[1,i],cp[1,i])
242  else:
243  c=_corners(pp,active=active)
244  return c
245 
246 def values_along_x(p,point=[0.5,0.5,0.5],iv=0,verbose=0):
247  ''' Get the iv values along the x direction through a point '''
248  s,f=values_along(p,point,dir=0,iv=iv,verbose=verbose)
249  return s,f
250 
251 def values_along_y(p,point=[0.5,0.5,0.5],iv=0,verbose=0):
252  ''' Get the iv values along the y direction through a point '''
253  s,f=values_along(p,point,dir=1,iv=iv,verbose=verbose)
254  return s,f
255 
256 def values_along_z(p,point=[0.5,0.5,0.5],iv=0,verbose=0):
257  ''' Get the iv values along the z direction through a point '''
258  s,f=values_along(p,point,dir=2,iv=iv,verbose=verbose)
259  return s,f
260 
261 def _corner_radii(p,origin=[0.,0.,0.]):
262  ''' array of corner radiii '''
263  c=_corners(p,active=True)
264  pt=[]
265  for x in (c[0,0]-origin[0],c[1,0]-origin[0]):
266  for y in (c[0,1]-origin[1],c[1,1]-origin[1]):
267  for z in (c[0,2]-origin[2],c[1,2]-origin[2]):
268  pt.append(x**2+y**2+z**2)
269  return np.sqrt(np.array(pt))
270 
272  def __init__(self,pp,radius=0.25,nds=1,dr=0.05,origin=[0.5,0.5,0.5],iv=0,verbose=0):
273  ''' return a dict with shell data '''
274  if isinstance(pp,dispatch._dispatch.snapshot):
275  pp=pp.patches
276  if verbose:
277  start=time()
278  self.var={}
279  self.pp=pp
280  p=pp[0]
281  keys=[]
282  # Add keys corresponding to patch variables
283  for key,value in p.idx.dict.items():
284  if value >= 0:
285  self.var[key]=[]
286  keys.append(key)
287  # Add extra variables tot the dict
288  for key in ('x','y','z','r','dv'):
289  self.var[key]=[]
290  p.idx.dict[key]=-1
291  npatch=0
292  for p in pp:
293  if not hasattr(p,'var'):
294  p.var={}
295  c2=_corner_radii(p,origin=origin)
296  if verbose>2:
297  print (p.id,p.position-origin,c2.min(),c2.max())
298  if is_inside(origin,p) or (c2.min()<radius and c2.max()>radius):
299  # make relevant variables available in dict d
300  if verbose:
301  if verbose>1:
302  print('using patch id',p.id)
303  npatch+=1
304  o=origin
305  rr=np.meshgrid(p.x-o[0],p.y-o[1],p.z-o[2])
306  rr=np.array(rr)
307  r2=sum(rr**2,0)
308  d={}
309  for key in keys:
310  d[key]=p.var(iv)
311  d['x']=rr[0]
312  d['y']=rr[1]
313  d['z']=rr[2]
314  d['r']=np.sqrt(r2)
315  d['dv']=np.ones(r2.shape)*np.product(p.ds)
316  r2l=(radius-nds*p.ds[0])**2
317  r2u=(radius+nds*p.ds[0])**2
318  cmp=(r2-r2l)*(r2-r2u)
319  w=np.where(cmp<0.0)
320  for key,value in d.items():
321  vv=value[w]
322  for v in vv:
323  self.var[key].append(v)
324  if verbose:
325  if npatch==1: s='patch'
326  else: s='patches'
327  print ('{} {} used, {:.3f} sec processing time'.format(npatch,s,time()-start))
328  if verbose>1:
329  print('variable min max (in shell)')
330  for key in self.var.keys():
331  v=np.array(self.var[key])
332  self.var[key]=v
333  if verbose>1:
334  print('{:>8} {:12.3e} {:12.3e}'.format(key,v.min(),v.max()))
335 
336  def radial_components(self):
337  ''' compute radial components of mass flux and velocity '''
338  self.var['pr']=(self.var['px']*self.var['x']+self.var['py']*self.var['y']+self.var['pz']*self.var['z'])/self.var['r']
339  ux=self.var['px']/self.var['d']
340  uy=self.var['py']/self.var['d']
341  uz=self.var['pz']/self.var['d']
342  self.var['ur']=(ux*self.var['x']+uy*self.var['y']+uz*self.var['z'])/self.var['r']
343 
344  def angles(self):
345  ''' compute latitude and longitude '''
346  self.var['mu']=self.var['z']/self.var['r']
347  self.var['lat']=np.arcsin(self.var['mu'])
348  self.var['lon']=np.arctan2(self.var['x'],self.var['y'])
349 
350  def mass_flux(self):
351  ''' return the net average mass flux, weighting each cell by volume '''
352  dv=self.var['dv']
353  return np.sum(self.var['pr']*dv)/np.sum(dv)
354 
def __init__(self, pp, radius=0.25, nds=1, dr=0.05, origin=[0.5, iv=0, verbose=0)
Definition: _select.py:272