DISPATCH
derived_variables.py
1 # -*- coding: utf-8 -*-
2 
3 """
4  Derived, physical quantities for use in analysis and plotting of DISPATCH2 results.
5 """
6 
7 import numpy as np
8 import itertools
9 
10 def divergence(p, varindices, gf, lnorm=True, lcalcnorm=False, lsum=False, lzc=False):
11  """
12  Return the divergence of a vector field given the components.
13  Input must be 3-D numpy arrays (even if one or more of the dimensions
14  has length 1).
15 
16  Originally written by David Clarke.
17  Modified by JPR for zone-centred magnetic fields and re-written in Python.
18 
19  Note: To apply this function in curvilinear coords., the geometric factors
20  must be provided (which is not currently implemented)!
21 
22  Another note: This function is known to be relatively slow due to the nested do-loops.
23 
24  v1, v2, v3 => components of a vector field in the form of numpy arrays
25  lnorm = T => normalise divergence
26  = F => do not normalise divergence
27  lcalcnorm = F => do not calculate the normalisation factor
28  = T => calculate a normalisation constant for the divergence, but DO NOT
29  apply it, only return its value.
30  lsum = F => perform a reduction on the divergence, returning the
31  absolute sum ('sumd')
32  = T => no reduction
33  lzc = F => input vector components are *face*-centred (e.g., ZEUS, STAGGER)
34  = T => input vector components are *zone*-centred (e.g., RAMSES)
35 
36  """
37 
38  assert isinstance(varindices,list) or isinstance(varindices,tuple), "varindices must be a tuple or a list."
39  assert len(varindices) == 3, "you need to supply the indices of three components of a vector field."
40 
41  TINY = 1.0e-99
42 
43  nx, ny, nz, nv = np.shape(p.data)
44  vf1 = p.data[:,:,:,varindices[0]]
45  vf2 = p.data[:,:,:,varindices[1]]
46  vf3 = p.data[:,:,:,varindices[2]]
47 
48  divv = np.empty(vf1.shape)
49 
50  if nx > 1: ione = 1
51  else: ione = 0
52  if ny > 1: jone = 1
53  else: jone = 0
54  if nz > 1: kone = 1
55  else: kone = 0
56 
57  # calculate divergence
58  if not lzc:
59  # face-centred inputs
60  dvl1fi = 1.0 / ( gf['dvol1f'] + TINY )
61  dvl2fi = 1.0 / ( gf['dvol2f'] + TINY )
62  dvl3fi = 1.0 / ( gf['dvol3f'] + TINY )
63 
64  for i,j,k in itertools.product(xrange(nx-ione),xrange(ny-jone),xrange(nz-kone)):
65  kp1 = k + kone
66  jp1 = j + jone
67  ip1 = i + ione
68  divv[i,j,k] = ( ( gf['dar1c'][ip1] * vf1[ip1,j,k]
69  - gf['dar1c'][i ] * vf1[i ,j,k] ) * dvl1fi[i]
70  + ( gf['h32f'][jp1] * vf2[i,jp1,k]
71  - gf['h32f'][j ] * vf2[i,j ,k] )
72  * dvl2fi[j] * gf['dar2f'][i]
73  + ( vf3[i,j,kp1] - vf3[i,j,k] )
74  * dvl3fi[k] * gf['dar31f'][i] * gf['dar32f'][j] )
75  else:
76  # zone-centred inputs
77  # note that, to end up with a zone-centred divergence, differences are taken
78  # from -1 to +1, and thus skipping i,j,k entirely.
79  dvl1ci = 1.0 / ( gf['dvol1c'] + TINY )
80  dvl2ci = 1.0 / ( gf['dvol2c'] + TINY )
81  dvl3ci = 1.0 / ( gf['dvol3c'] + TINY )
82 
83  for i,j,k in itertools.product(xrange(ione,nx-ione),xrange(jone,ny-jone),xrange(kone,nz-kone)):
84  km1 = k - kone
85  kp1 = k + kone
86  jm1 = j - jone
87  jp1 = j + jone
88  im1 = i - ione
89  ip1 = i + ione
90  divv[i,j,k] = ( ( gf['dar1f'][ip1] * vf1[ip1,j,k]
91  - gf['dar1f'][im1] * vf1[im1,j,k] )
92  * ( dvl1ci[ip1] + dvl1ci[i] )
93  + ( gf['h32c'][jp1] * vf2[i,jp1,k]
94  - gf['h32c'][jm1] * vf2[i,jm1,k] )
95  * ( dvl2ci[jp1] + dvl2ci[j] ) * gf['dar2f'][i]
96  + ( vf3[i,j,kp1] - vf3[i,j,km1] )
97  * ( dvl3ci[kp1] + dvl3ci[k] )
98  * gf['dar31f'][i] * gf['dar32f'][j] )
99 
100  # calculate the reduction, if desired
101  sumd = None
102  if lsum: sumd = np.sum(np.abs(divv))
103 
104  # normalise the divergence, if desired
105  nmax = None
106  if lnorm or lcalcnorm:
107  norm = np.zeros(vf1.shape)
108 
109  # Evaluate two normalising constants:
110  # 1) sumn = sum over all grid zones the ratio of (average absolute
111  # vector field) / (sum of grid zone dimensions)
112  # 2) nmax = maximum over the grid of the above ratio.
113  if not lzc:
114  for i,j,k in itertools.product(xrange(nx-ione),xrange(ny-jone),xrange(nz-kone)):
115  kp1 = k + kone
116  jp1 = j + jone
117  ip1 = i + ione
118  norm[i,j,k] = ( ( abs( vf1[ip1,j,k] + vf1[i,j,k] )
119  + abs( vf2[i,jp1,k] + vf2[i,j,k] )
120  + abs( vf3[i,j,kp1] + vf3[i,j,k] ) ) * 0.5
121  / ( ione * gf['dx1f'][i]
122  + jone * gf['h2c'][i] * gf['dx2f'][j]
123  + kone * gf['h31c'][i] * gf['h32c'][j] * gf['dx3f'][k] ) )
124  else:
125  for i,j,k in itertools.product(xrange(nx-ione),xrange(ny-jone),xrange(nz-kone)):
126  norm[i,j,k] = ( ( abs( vf1[i,j,k] )
127  + abs( vf2[i,j,k] )
128  + abs( vf3[i,j,k] ) )
129  / ( ione * gf['dx1f'][i]
130  + jone * gf['h2c'][i] * gf['dx2f'][j]
131  + kone * gf['h31c'][i] * gf['h32c'][j] * gf['dx3f'][k] ) )
132 
133  sumn = np.sum(norm)
134  nmax = np.nanmax(norm)
135 
136  # Apply normalising constants 'sumn' to scalar 'sumd' and 'nmax' to 'divv'.
137  if lsum:
138  if sumn < TINY:
139  sumd = 0.0
140  else:
141  sumd = sumd / sumn
142  if lnorm:
143  if nmax < TINY:
144  nmaxi = 0.0
145  else:
146  nmaxi = 1.0 / nmax
147  divv = divv * nmaxi
148 
149  return divv, sumd, nmax
150 
151 def gas_temperature (p, dindex, eindex, gf, gamma=5.0/3.0):
152  """Calculate the gas temperature given the density and internal energy."""
153 
154  #import ipdb
155  #ipdb.set_trace()
156  e = p.data[:,:,:,eindex]
157  d = p.data[:,:,:,dindex]
158 
159  return ((gamma-1.0) * e / d)
160 
161 def speed (p, varidx, gf):
162  """Calculate the speed = magnitude of the velocity vector."""
163 
164  s1 = p.data[:,:,:,varidx['s1']]
165  s2 = p.data[:,:,:,varidx['s2']]
166  s3 = p.data[:,:,:,varidx['s3']]
167  d = p.data[:,:,:,varidx['d ']]
168 
169  v1 = s1 / d
170  v2 = s2 / d / gf['h2c'][:,np.newaxis,np.newaxis]
171  v3 = (s3 / d / gf['h31c'][:,np.newaxis,np.newaxis]
172  / gf['h32c'][np.newaxis,:,np.newaxis])
173 
174  # FIXME: technically incorrect for ZEUS and STAGGER data, which should be
175  # averaged to zone-centre first.
176  return np.sqrt(v1**2 + v2**2 + v3**2)
177 
178 def gradient_scalar (p, varidx, gf, idir=0):
179  """
180  Calculate the gradient of a scalar (zone-centred).
181  The result is a staggered, 2nd-order difference.
182 
183  """
184 
185  na = np.newaxis
186  n1, n2, n3 = p.gn
187  if n1 > 1: ione = 1
188  else: ione = 0
189  if n2 > 1: jone = 1
190  else: jone = 0
191  if n3 > 1: kone = 1
192  else: kone = 0
193 
194  cvar = p.data[:,:,:,varidx]
195  if idir == 1:
196  g = np.empty_like(cvar)
197  for i in xrange(1,n1):
198  g[i,:,:] = (cvar[i,:,:] - cvar[i-ione,:,:]) / gf['dx1c'][i]
199  return g
200  elif idir == 2:
201  g = np.empty_like(cvar)
202  for j in xrange(1,n2):
203  g[:,j,:] = (cvar[:,j,:] - cvar[:,j-jone,:]) / gf['dx2c'][j] / gf['h2c']
204  return g
205  elif idir == 3:
206  g = np.empty_like(cvar)
207  for k in xrange(1,n3):
208  g[:,:,k] = (cvar[:,:,k] - cvar[:,:,k-kone]) / gf['dx3c'][k] / gf['h31c'] / gf['h32c']
209  return g
210  else:
211  g1 = np.empty_like(cvar)
212  for i in xrange(1,n1):
213  g1[i,:,:] = (cvar[i,:,:] - cvar[i-ione,:,:]) / gf['dx1c'][i]
214  g2 = np.empty_like(cvar)
215  for j in xrange(1,n2):
216  g2[:,j,:] = (cvar[:,j,:] - cvar[:,j-jone,:]) / gf['dx2c'][j] / gf['h2c']
217  g3 = np.empty_like(cvar)
218  for k in xrange(1,n3):
219  g3[:,:,k] = (cvar[:,:,k] - cvar[:,:,k-kone]) / gf['dx3c'][k] / gf['h31c'] / gf['h32c']
220  return (g1,g2,g3)