DISPATCH
dispatch_grid.py
1 # -*- coding: utf-8 -*-
2 
3 """
4  Support functions for reading or calculating DISPATCH2 grid and geometry information.
5 """
6 
7 import numpy as np
8 from abc import abstractproperty, abstractmethod, ABCMeta
9 from curvilinear_transforms import *
10 from dispatch_data import Patch
11 from legacy_reader import mesh_types
12 
13 class Grid(dict):
14  """An abstract grid class based on `dict()`."""
15 
16  __metaclass__= ABCMeta
17 
19  """Calculate and store the geometric factors used by curvilinear grids."""
20 
21  def __init__(self):
22  """Constructor."""
23 
24  # Define geometric factors with the same notation as in `mesh_mod` ("c"
25  # for zone-centred and "f" for face-centred).
26  self['h2c'] = None
27  self['h2f'] = None
28  self['h31c'] = None
29  self['h31f'] = None
30  self['h32c'] = None
31  self['h32f'] = None
32  self['dx1c'] = None
33  self['dx1f'] = None
34  self['dx2c'] = None
35  self['dx2f'] = None
36  self['dx3c'] = None
37  self['dx3f'] = None
38  self['dvol1c'] = None
39  self['dvol1f'] = None
40  self['dvol2c'] = None
41  self['dvol2f'] = None
42  self['dvol3c'] = None
43  self['dvol3f'] = None
44  self['dar1c'] = None
45  self['dar1f'] = None
46  self['dar2c'] = None
47  self['dar2f'] = None
48  self['dar31c'] = None
49  self['dar31f'] = None
50  self['dar32c'] = None
51  self['dar32f'] = None
52 
53  def init_grid(self, p, x_i):
54  """Initialise geometric factors based on coord. type."""
55 
56  if p.mesh_type == 'Cartesian': self.init_Cartesian(p, x_i)
57  elif p.mesh_type == 'cylindrical': self.init_cylindrical(p, x_i)
58  elif p.mesh_type == 'spherical': self.init_spherical(p, x_i)
59 
60  def init_Cartesian(self, p, x_i):
61  """Initialise geometric factors for a Cartesian coord. system."""
62 
63  n1, n2, n3 = p.nwrite
64 
65  # 1-direction
66  self['h2c'] = np.ones(n1)
67  self['h2f'] = np.ones(n1)
68  self['h31c'] = self['h2c'].view()
69  self['h31f'] = self['h2f'].view()
70 
71  # 2-direction
72  self['h32c'] = np.ones(n2)
73  self['h32f'] = self['h32c'].view()
74 
75  # linear size elements
76  self['dx1c'] = np.ones(n1) * p.dx[0]
77  self['dx1f'] = np.ones(n1) * p.dx[0]
78  self['dx2c'] = np.ones(n2) * p.dx[1]
79  self['dx2f'] = np.ones(n2) * p.dx[1]
80  self['dx3c'] = np.ones(n3) * p.dx[2]
81  self['dx3f'] = np.ones(n3) * p.dx[2]
82 
83  # volume elements
84  self['dvol1c'] = np.ones(n1) * p.dx[0]
85  self['dvol1f'] = np.ones(n1) * p.dx[0]
86  self['dvol2c'] = np.ones(n2) * p.dx[1]
87  self['dvol2f'] = np.ones(n2) * p.dx[1]
88  self['dvol3c'] = np.ones(n3) * p.dx[2]
89  self['dvol3f'] = np.ones(n3) * p.dx[2]
90 
91  # area elements
92  self['dar1c'] = self['h31c'] * self['h2c']
93  self['dar1f'] = self['h31f'] * self['h2f']
94  self['dar2c'] = self['h31f'] * p.dx[0] / self['dvol1c']
95  self['dar2f'] = self['h31c'] * p.dx[0] / self['dvol1f']
96  self['dar31c'] = self['h2f'] * p.dx[0] / self['dvol1c']
97  self['dar31f'] = self['h2c'] * p.dx[0] / self['dvol1f']
98  self['dar32c'] = p.dx[1] / self['dvol2c']
99  self['dar32f'] = p.dx[1] / self['dvol2f']
100 
101  def init_cylindrical(self, p, x_i):
102  """Initialise geometric factors for a cylindrical coord. system."""
103 
104  n1, n2, n3 = p.nwrite
105 
106  # 1-direction
107  self['h2c'] = np.ones(n1)
108  self['h2f'] = np.ones(n1)
109  self['h31c'] = self['h2c'].view()
110  self['h31f'] = self['h2f'].view()
111 
112  # 2-direction
113  pos_c = np.array(x_i[1])
114  pos_f = np.array(x_i[1] - 0.5 * p.dx[1])
115  self['h32c'] = abs(pos_c)
116  self['h32f'] = abs(pos_f)
117 
118  # linear size elements
119  self['dx1c'] = np.ones(n1) * p.dx[0]
120  self['dx1f'] = np.ones(n1) * p.dx[0]
121  self['dx2c'] = np.ones(n2) * p.dx[1]
122  self['dx2f'] = np.ones(n2) * p.dx[1]
123  self['dx3c'] = np.ones(n3) * p.dx[2]
124  self['dx3f'] = np.ones(n3) * p.dx[2]
125 
126  # volume elements
127  self['dvol1c'] = np.ones(n1) * p.dx[0]
128  self['dvol1f'] = np.ones(n1) * p.dx[0]
129  self['dvol2c'] = np.empty_like(pos_c)
130  self['dvol2f'] = np.empty_like(pos_f)
131  for j in xrange(len(pos_c)-1):
132  self['dvol2c'][j] = 0.5 * abs( self['h32f'][j+1] * pos_f[j+1]
133  - self['h32f'][j ] * pos_f[j ] )
134  self['dvol2f'][j] = 0.5 * abs( self['h32c'][j ] * pos_c[j ]
135  - self['h32c'][j-1] * pos_c[j-1] )
136  self['dvol3c'] = np.ones(n3) * p.dx[2]
137  self['dvol3f'] = np.ones(n3) * p.dx[2]
138 
139  # area elements
140  self['dar1c'] = self['h31c'] * self['h2c']
141  self['dar1f'] = self['h31f'] * self['h2f']
142  self['dar2c'] = self['h31f'] * p.dx[0] / self['dvol1c']
143  self['dar2f'] = self['h31c'] * p.dx[0] / self['dvol1f']
144  self['dar31c'] = self['h2f'] * p.dx[0] / self['dvol1c']
145  self['dar31f'] = self['h2c'] * p.dx[0] / self['dvol1f']
146  self['dar32c'] = p.dx[1] / self['dvol2c']
147  self['dar32f'] = p.dx[1] / self['dvol2f']
148 
149  def init_spherical(self, p, x_i):
150  """Initialise geometric factors for a spherical coord. system."""
151 
152  n1, n2, n3 = p.nwrite
153 
154  # 1-direction
155  rpos_c = np.array(x_i[0])
156  rpos_f = np.array(x_i[0] - 0.5 * p.dx[0])
157  self['h2c'] = abs(rpos_c)
158  self['h2f'] = abs(rpos_f)
159  self['h31c'] = self['h2c'].view()
160  self['h31f'] = self['h2f'].view()
161 
162  # 2-direction
163  tpos_c = np.array(x_i[1])
164  tpos_f = np.array(x_i[1] - 0.5 * p.dx[1])
165  self['h32c'] = abs(np.sin(tpos_c))
166  self['h32f'] = abs(np.sin(tpos_f))
167 
168  # linear size elements
169  self['dx1c'] = np.ones(n1) * p.dx[0]
170  self['dx1f'] = np.ones(n1) * p.dx[0]
171  self['dx2c'] = np.ones(n2) * p.dx[1]
172  self['dx2f'] = np.ones(n2) * p.dx[1]
173  self['dx3c'] = np.ones(n3) * p.dx[2]
174  self['dx3f'] = np.ones(n3) * p.dx[2]
175 
176  # volume elements
177  self['dvol1c'] = np.empty_like(rpos_c)
178  self['dvol1f'] = np.empty_like(rpos_f)
179  for i in xrange(len(rpos_c)-1):
180  self['dvol1c'] = (self['h2f'][i+1] * self['h31f'][i+1] * rpos_f[i+1] / 3.0
181  - self['h2f'][i ] * self['h31f'][i ] * rpos_f[i ] / 3.0)
182  self['dvol1f'] = (self['h2c'][i ] * self['h31c'][i ] * rpos_c[i ] / 3.0
183  - self['h2c'][i-1] * self['h31c'][i-1] * rpos_c[i-1] / 3.0)
184  self['dvol2c'] = np.empty_like(tpos_c)
185  self['dvol2f'] = np.empty_like(tpos_f)
186  for j in xrange(len(tpos_c)-1):
187  self['dvol2c'] = -np.cos(tpos_c[j+1]) + np.cos(tpos_c[j ])
188  self['dvol2f'] = -np.cos(tpos_f[j ]) + np.cos(tpos_f[j-1])
189  self['dvol3c'] = np.ones(n3) * p.dx[2]
190  self['dvol3f'] = np.ones(n3) * p.dx[2]
191 
192  # area elements
193  self['dar1c'] = self['h31c'] * self['h2c']
194  self['dar1f'] = self['h31f'] * self['h2f']
195  self['dar2c'] = self['h31f'] * p.dx[0] / self['dvol1c']
196  self['dar2f'] = self['h31c'] * p.dx[0] / self['dvol1f']
197  self['dar31c'] = self['h2f'] * p.dx[0] / self['dvol1c']
198  self['dar31f'] = self['h2c'] * p.dx[0] / self['dvol1f']
199  self['dar32c'] = p.dx[1] / self['dvol2c']
200  self['dar32f'] = p.dx[1] / self['dvol2f']
201 
202 def dispatch_grid(p, truecoords=True):
203  """
204  Given grid information for a patch, calculate the three *cell-centred*
205  coordinate arrays in the *coord. system of the patch*.
206  To convert to Cartesian coords., see `plotting_grid`.
207 
208  """
209 
210  if np.any(p.n != p.nwrite):
211  if truecoords:
212  x1 = p.corners[0][0] + np.arange(p.nwrite[0]) * p.dx[0] + 0.5 * p.dx[0]
213  x2 = p.corners[0][1] + np.arange(p.nwrite[1]) * p.dx[1] + 0.5 * p.dx[1]
214  x3 = p.corners[0][2] + np.arange(p.nwrite[2]) * p.dx[2] + 0.5 * p.dx[2]
215  else:
216  x1 = p.llc[0] - p.nghost[0] * p.dx[0] + np.arange(p.nwrite[0]) * p.dx[0] + 0.5 * p.dx[0]
217  x2 = p.llc[1] - p.nghost[1] * p.dx[1] + np.arange(p.nwrite[1]) * p.dx[1] + 0.5 * p.dx[1]
218  x3 = p.llc[2] - p.nghost[2] * p.dx[2] + np.arange(p.nwrite[2]) * p.dx[2] + 0.5 * p.dx[2]
219  else:
220  if truecoords:
221  x1 = p.active_corners[0][0] + np.arange(p.nwrite[0]) * p.dx[0] + 0.5 * p.dx[0]
222  x2 = p.active_corners[0][1] + np.arange(p.nwrite[1]) * p.dx[1] + 0.5 * p.dx[1]
223  x3 = p.active_corners[0][2] + np.arange(p.nwrite[2]) * p.dx[2] + 0.5 * p.dx[2]
224  else:
225  x1 = p.llc[0] + np.arange(p.nwrite[0]) * p.dx[0] + 0.5 * p.dx[0]
226  x2 = p.llc[1] + np.arange(p.nwrite[1]) * p.dx[1] + 0.5 * p.dx[1]
227  x3 = p.llc[2] + np.arange(p.nwrite[2]) * p.dx[2] + 0.5 * p.dx[2]
228 
229  if p.ioformat == 2:
230  x1 = x1 - p.dx[0] * 0.5
231  x2 = x2 - p.dx[1] * 0.5
232  x3 = x3 - p.dx[2] * 0.5
233 
234  return (x1, x2, x3)
235 
236 def plotting_grid(p, x_i, slicedir, sliceloc, sliceeps=1.0e-2, truecoords=False,
237  phi_rotate=0.0, theta_rotate=0.0):
238  """
239  Given slice information, return the appropriate coord. vectors for plotting,
240  including the conversion to Cartesian coords. (required for plotting with
241  matplotlib).
242 
243  If the `truecoords` keyword is set True, then the conversion to Cartesian
244  coords. is skipped.
245 
246  The following curvilinear coord. systems are known:
247  (x,y,z), (Z,R,Phi), (r,theta,phi)
248  The `phi_rotate` and `theta_rotate` arguments allow for patch coords. which
249  are rotated relative to the convention of theta = phi = 0. They are silently
250  passed to the coord. transformation functions.
251 
252  """
253 
254  assert (slicedir > 1) or (slicedir < 3), "slice direction must be 1, 2, or 3."
255 
256  if p.mesh_type == 'Cartesian' or truecoords:
257  x, y, z = np.meshgrid(x_i[0], x_i[1], x_i[2], indexing='ij')
258  elif p.mesh_type == 'cylindrical':
259  x, y, z = zrp_to_xyz(x_i[0], x_i[1], x_i[2], phi_rotate=phi_rotate)
260  elif p.mesh_type == 'spherical':
261  x, y, z = rtp_to_xyz(x_i[0], x_i[1], x_i[2], phi_rotate=phi_rotate, theta_rotate=theta_rotate)
262 
263  # adjust coords. using the patch position (which is in Cartesian coords.).
264  #x += p.pos[0]
265  #y += p.pos[1]
266  #z += p.pos[2]
267 
268  # assign coords. actually needed for plotting; cyclical.
269  if slicedir == 1:
270  xp, yp = y,z
271  elif slicedir == 2:
272  xp, yp = z,x
273  elif slicedir == 3:
274  xp, yp = x,y
275 
276  # account for symmetry (i.e., less than 3-D).
277  sliceloc_act = sliceloc
278  if p.mesh_type == 'Cartesian':
279  if slicedir == 1 and len(x_i[0]) == 1: sliceloc_act = x_i[0]
280  if slicedir == 2 and len(x_i[1]) == 1: sliceloc_act = x_i[1]
281  if slicedir == 3 and len(x_i[2]) == 1: sliceloc_act = x_i[2]
282  if p.mesh_type == 'cylindrical':
283  if slicedir == 3 and len(x_i[0]) == 1: sliceloc_act = x_i[0]
284  if p.mesh_type == 'spherical':
285  if slicedir == 3 and len(x_i[1]) == 1: sliceloc_act = np.cos(x_i[1])
286 
287  # finally, determine the (nearest) slice location for the current patch.
288  if slicedir == 1:
289  diff = x - sliceloc_act
290  elif slicedir == 2:
291  diff = y - sliceloc_act
292  elif slicedir == 3:
293  diff = z - sliceloc_act
294  # create mask
295  m = np.ma.masked_outside(diff,-sliceeps,sliceeps)
296  xpcut = np.ma.array(xp,mask=m.mask)
297  ypcut = np.ma.array(yp,mask=m.mask)
298 
299  return xpcut, ypcut, m.mask
300 
301 def patch_grid_edges(p, slicedir, scaling=1.0, truecoords=True):
302  """
303  Using the corners of the patch, create a matplotlib patch that draws
304  the grid edges.
305 
306  """
307 
308  from matplotlib.patches import Rectangle, Wedge
309 
310  assert isinstance(p, Patch), "Input must be an instance of a DISPATCH `Patch`!"
311 
312  if truecoords:
313  if slicedir == 1:
314  llx, urx = p.active_corners[0][1], p.active_corners[1][1]
315  lly, ury = p.active_corners[0][2], p.active_corners[1][2]
316  elif slicedir == 2:
317  llx, urx = p.active_corners[0][0], p.active_corners[1][0]
318  lly, ury = p.active_corners[0][2], p.active_corners[1][2]
319  elif slicedir == 3:
320  llx, urx = p.active_corners[0][0], p.active_corners[1][0]
321  lly, ury = p.active_corners[0][1], p.active_corners[1][1]
322  else:
323  if slicedir == 1:
324  llx, urx = p.llc[1], p.llc[1] + p.size[1]
325  lly, ury = p.llc[2], p.llc[2] + p.size[2]
326  elif slicedir == 2:
327  llx, urx = p.llc[0], p.llc[0] + p.size[0]
328  lly, ury = p.llc[2], p.llc[2] + p.size[2]
329  elif slicedir == 3:
330  llx, urx = p.llc[0], p.llc[0] + p.size[0]
331  lly, ury = p.llc[1], p.llc[1] + p.size[1]
332 
333  if p.mesh_type == mesh_types[1] or truecoords: # Cartesian
334  llx *= scaling
335  lly *= scaling
336  urx *= scaling
337  ury *= scaling
338  edges = Rectangle((llx,lly),width=(urx-llx),height=(ury-lly),fc='none',lw=1.0,
339  ec='k',linestyle='dashed',zorder=0.01/min(p.dx)+0.1)
340  elif p.mesh_type == mesh_types[2]: # spherical
341  thetamin = round(p.active_corners[0][2]*180.0/np.pi,3)
342  thetamax = round(p.active_corners[1][2]*180.0/np.pi,3)
343  llx, urx = p.active_corners[0][0], p.active_corners[1][0]
344  llx *= scaling
345  lly *= scaling
346  urx *= scaling
347  ury *= scaling
348  edges = Wedge((p.pos[0]*scaling,p.pos[1]*scaling),urx,thetamin,thetamax,width=(urx-llx),fc='none',
349  lw=1.0,ec='k',linestyle='dotted',zorder=0.01/min(p.dx)+0.1)
350  elif p.mesh_type == mesh_types[3]: # cylindrical
351  phimin = round(p.active_corners[0][2]*180.0/np.pi,3)
352  phimax = round(p.active_corners[1][2]*180.0/np.pi,3)
353  llx, urx = p.active_corners[0][1], p.active_corners[1][1]
354  llx *= scaling
355  lly *= scaling
356  urx *= scaling
357  ury *= scaling
358  edges = Wedge((p.pos[0]*scaling,p.pos[1]*scaling),urx,phimin,phimax,width=(urx-llx),fc='none',
359  lw=1.0,ec='k',linestyle='dotted',zorder=0.01/min(p.dx)+0.1)
360 
361  return edges
def init_Cartesian(self, p, x_i)
def init_cylindrical(self, p, x_i)
def init_spherical(self, p, x_i)