DISPATCH
dispatch_data.py
1 # -*- coding: utf-8 -*-
2 
3 """
4  Support functions for reading DISPATCH2 data.
5 """
6 
7 # Pythn 2/3 compatibility
8 from __future__ import print_function
9 import sys
10 if sys.version_info[0]!=2:
11  from importlib import reload
12 
13 import numpy as np
14 from abc import abstractproperty, ABCMeta
15 import os
16 from time import time
17 
18 import legacy_reader
19 reload(legacy_reader)
20 import stagger_utils as su
21 
22 class Task(object):
23  """An abstract class to describe tasks."""
24  __metaclass__ = ABCMeta
25  def __init__(self):
26  pass
27  @abstractproperty
28  def data(self):
29  """An iterable that yields the data associated with the object."""
30  pass
31 
32 class Patch(Task):
33  """A derived class for patches."""
34  __metaclass__ = ABCMeta
35  def __init__(self, filename, suffix='.txt', read_derivs=False, overlap=0.01, verbose=False, fd=None):
36  self.filename = filename
37  self.__data = None
38  self.offset=0
39  if fd:
40  fp=fd
41  else:
42  fp = open(filename+suffix,'r')
43  legacy_reader.single_patch (self, fp, verbose=False, fd=fd)
44  self.active_corners[1] = self.active_corners[1] + self.ds*overlap
45  self.corners[1] = self.corners[1] + self.ds*overlap
46  if self.kind == 'poisson':
47  self.nbytes = 8
48  else:
49  self.nbytes = 4
50  if read_derivs: self.nvar = 2 * self.nvar
51  if fd:
52  self.offset=(self.id-1)*np.product(self.nwrite)*self.nvar*4
53  else:
54  fp.close()
55  self.varidx = self.variable_indices(read_derivs)
56  self.varnames = []
57  self.__cached = 0
58  for k in self.varidx.keys():
59  self.varnames.append(k)
60  for k in self.varidx.keys():
61  v = self.varidx[k]
62  self.varnames[v] = k
63 
64  @property
65  def data(self):
66  """
67  The actual data for a patch is accessed through a numpy 'memory map',
68  which uses on-demand access rather than reading the whole file in.
69  Note that a file handle is opened when this is invoked for the first time
70  and, if there are a very large number of patches, python will raise
71  an exception when it reaches its file handler limit.
72  """
73  if self.__data is None:
74  self.__data = self._read_legacy_data(self.filename+'.dat', self.nwrite, self.nvar, self.nbytes)
75  return self.__data
76 
77  def _read_legacy_data(self, filename, dims, nvar, nbytes):
78  """Return a memory map pointing to MHD data."""
79  myshape = dims + (nvar,)
80  if nbytes == 4:
81  thistype = np.float32
82  elif nbytes == 8:
83  thistype = np.float64
84 
85  return np.memmap(filename, dtype=thistype, offset=self.offset, mode='r', order='F', shape=myshape)
86 
87  def _determine_corners(self, active_only=False):
88  """
89  For a given patch, calculate the corners of the cuboid.
90 
91  `corners` are defined as the absolute lowest coords. of a patch (i.e,
92  the first or last interface). `active_corners` are instead defined to
93  only include active cells.
94  """
95  c_lo = []
96  c_hi = []
97  if not active_only:
98  for o,n,d,ng in zip(self.origin, self.n, self.dx, self.nghost):
99  c_lo.append(o - ng * d)
100  c_hi.append(o + (n + ng) * d)
101  else:
102  for o,n,d in zip(self.origin, self.n, self.dx):
103  c_lo.append(o)
104  c_hi.append(o + n * d)
105  # we want the data to be immutable
106  c_lo = np.array(c_lo)
107  c_hi = np.array(c_hi)
108  return c_lo, c_hi
109 
110  def stripv (self, v):
111  l=np.array(self.nghost)
112  u=l+np.array(self.n)
113  #if self.ioformat==2:
114  # u=u-1
115  return v[l[0]:u[0],l[1]:u[1],l[2]:u[2]]
116  def trans(self,tr):
117  if self.kind[0:9]=='rt_solver':
118  self.t = self.t.transpose(tr)
119  self.Q = self.Q.transpose(tr)
120  elif self.kind[0:7]=='poisson':
121  self.d = self.d.transpose(tr)
122  self.phi = self.phi.transpose(tr)
123  else:
124  '''Remove guard cells from cached arrays'''
125  if not 'd' in dir(self):
126  self.cache()
127  self.d = self.d.transpose(tr)
128  self.e = self.e.transpose(tr)
129  self.s = self.s.transpose(tr)
130  self.sd = self.sd.transpose(tr)
131  self.temp = self.temp.transpose(tr)
132  self.pg = self.pg.transpose(tr)
133  self.u1 = self.u1.transpose(tr)
134  self.u2 = self.u2.transpose(tr)
135  self.u3 = self.u3.transpose(tr)
136  self.p1 = self.p1.transpose(tr)
137  self.p2 = self.p2.transpose(tr)
138  self.p3 = self.p3.transpose(tr)
139  if self.nvar>8:
140  self.phi= self.phi.transpose(tr)
141  if self.nvar>=8:
142  self.b1 = self.b1.transpose(tr)
143  self.b2 = self.b2.transpose(tr)
144  self.b3 = self.b3.transpose(tr)
145  t1=[]
146  t1.append(self.x); t1.append(self.y); t1.append(self.z)
147  self.x=t1[tr[0]]; self.y=t1[tr[1]]; self.z=t1[tr[2]]
148  ac=self.active_corners
149  self.active_corners[0]=[ac[0][tr[0]],ac[0][tr[1]],ac[0][tr[2]]]
150  self.active_corners[1]=[ac[1][tr[0]],ac[1][tr[1]],ac[1][tr[2]]]
151 
152  l=self.nghost[0]; u=l+self.n[0]; self.x = self.x[l:u]; self.xs = self.xs[l:u]
153  l=self.nghost[1]; u=l+self.n[1]; self.y = self.y[l:u]; self.ys = self.ys[l:u]
154  l=self.nghost[2]; u=l+self.n[2]; self.z = self.z[l:u]; self.zs = self.zs[l:u]
155 
156  def stripit(self):
157  if self.kind[0:9]=='rt_solver':
158  self.t = self.stripv(self.t)
159  self.Q = self.stripv(self.Q)
160  elif self.kind[0:7]=='poisson':
161  self.d = self.stripv(self.d)
162  self.phi = self.stripv(self.phi)
163  else:
164  '''Remove guard cells from cached arrays'''
165  if not 'd' in dir(self):
166  self.cache()
167  self.d = self.stripv (self.d)
168  self.e = self.stripv (self.e)
169  self.s = self.stripv (self.s)
170  self.sd = self.stripv (self.sd)
171  self.temp = self.stripv (self.temp)
172  self.pg = self.stripv (self.pg)
173  self.u1 = self.stripv (self.u1)
174  self.u2 = self.stripv (self.u2)
175  self.u3 = self.stripv (self.u3)
176  self.p1 = self.stripv (self.p1)
177  self.p2 = self.stripv (self.p2)
178  self.p3 = self.stripv (self.p3)
179  if self.nvar==6 or self.nvar>8:
180  self.phi= self.stripv(self.phi)
181  if self.nvar>=8:
182  self.b1 = self.stripv (self.b1)
183  self.b2 = self.stripv (self.b2)
184  self.b3 = self.stripv (self.b3)
185  l=self.nghost; u=l+self.n
186  #if self.ioformat==2:
187  # u=u-1
188  self.x = self.x[l[0]:u[0]]; self.xs = self.xs[l[0]:u[0]]
189  self.y = self.y[l[1]:u[1]]; self.ys = self.ys[l[1]:u[1]]
190  self.z = self.z[l[2]:u[2]]; self.zs = self.zs[l[2]:u[2]]
191 
192  def cache(self):
193  if self.__cached:
194  return
195  if self.kind[0:9]=='rt_solver':
196  self.Q = self.data[:,:,:,self.varidx['Q']]
197  self.S = self.data[:,:,:,self.varidx['S']]
198  self.rk = self.data[:,:,:,self.varidx['rk']]
199  elif self.kind[0:7]=='poisson':
200  self.d = self.data[:,:,:,self.varidx['d']]
201  self.phi = self.data[:,:,:,self.varidx['phi']]
202  else:
203  #if not 'd' in dir(self):
204  self.d =self.data[:,:,:,self.varidx['d']]
205  self.p1=self.data[:,:,:,self.varidx['p1']]
206  self.p2=self.data[:,:,:,self.varidx['p2']]
207  self.p3=self.data[:,:,:,self.varidx['p3']]
208  if self.nvar==6 or self.nvar>8:
209  self.phi=self.data[:,:,:,self.varidx['phi']]
210  if self.nvar>=8:
211  self.b1=self.data[:,:,:,self.varidx['b1']]
212  self.b2=self.data[:,:,:,self.varidx['b2']]
213  self.b3=self.data[:,:,:,self.varidx['b3']]
214  g1=self.gamma-1.0
215  if g1==0.0: g1=1e-4
216  if self.kind[0:7]=='stagger':
217  logd=np.log10(self.d)
218  ln10=np.log(10.)
219  self.u1=self.p1/np.exp(su.xdn(ln10*logd))
220  self.u2=self.p2/np.exp(su.ydn(ln10*logd))
221  self.u3=self.p3/np.exp(su.zdn(ln10*logd))
222  if self.kind[0:9]=='stagger2_':
223  self.s =self.data[:,:,:,self.varidx['s']]
224  self.e =np.exp((self.s/self.d+ln10*logd)*g1)/g1
225  self.pg=self.d*self.e*g1
226  elif self.kind[0:9]=='stagger2e':
227  self.e =self.data[:,:,:,self.varidx['e']]/self.d
228  self.pg=self.d*self.e*g1
229  self.s=np.log(self.pg/self.d**self.gamma)/g1
230  elif self.kind[0:6]=='ramses':
231  self.e =self.data[:,:,:,self.varidx['e']]/self.d
232  self.u1=self.p1/self.d
233  self.u2=self.p2/self.d
234  self.u3=self.p3/self.d
235  self.e=self.e-0.5*(self.u1**2+self.u2**2+self.u3**2)
236  self.pg=self.d*self.e*g1
237  self.s=np.log(self.pg/self.d**self.gamma)/g1
238  self.sd =self.s/self.d
239  self.temp = self.pg/self.d
240  self.__cached=1
241 
242  def vars(self):
243  '''Make arrays available in a .var dict'''
244  if not 'var' in dir(self):
245  self.cache()
246  if self.kind=='rt_solver':
247  self.var={'Q':self.Q,\
248  'S':self.S,\
249  'rk':self.rk}
250  else:
251  self.var={'d':self.d,\
252  'e':self.e,\
253  's':self.s,\
254  'sd':self.sd,\
255  'p1':self.p1,\
256  'p2':self.p2,\
257  'p3':self.p3,\
258  'u1':self.u1,\
259  'u2':self.u2,\
260  'u3':self.u3,\
261  'logd':np.log10(self.d),\
262  'loge':np.log10(self.e)}
263  if self.nvar>7:
264  self.var['b1']=self.b1
265  self.var['b2']=self.b2
266  self.var['b3']=self.b3
267  if self.nvar==6 or self.nvar>8:
268  self.var['phi']=self.phi
269 
270  def corner_radii(self):
271  ''' array of corner radiii '''
272  c=self.active_corners
273  pt=[]
274  for x in (c[0,0],c[1,0]):
275  for y in (c[0,1],c[1,1]):
276  for z in (c[0,2],c[1,2]):
277  pt.append(x**2+y**2+z**2)
278  return np.sqrt(np.array(pt))
279 
280  def variable_indices(self, read_derivs=False):
281  """
282  A simple function to provide the indices of the variable stored in the binary
283  dump.
284  """
285  assert isinstance(self.kind, str), "expecting a string for `self.kind`"
286  solver = self.kind.lower()
287  if solver[0:7]== 'stagger':
288  indices = {'d': 0,
289  'p1': 2, 'p2': 3, 'p3': 4}
290  if solver[7:9] == '2e':
291  indices['e'] = 1
292  else:
293  indices['s'] = 1
294  if self.nvar==6:
295  indices['phi']=5
296  if self.nvar>=8:
297  indices['b1']=5
298  indices['b2']=6
299  indices['b3']=7
300  if self.nvar==11:
301  indices['v1']=8
302  indices['v2']=9
303  indices['v3']=10
304  if self.nvar==9:
305  indices['phi']=8
306  if read_derivs:
307  indicesd = {'dd': 8,
308  'ds': 9,
309  'dp1': 10, 'dp2': 11, 'dp3': 12,
310  'db1': 13, 'db2': 14, 'db3': 15,}
311  indices.update(indicesd)
312  elif solver[0:6] == 'ramses':
313  indices = {'d': 0,
314  'e': 4,
315  'p1': 1, 'p2': 2, 'p3': 3,
316  'b1': 5, 'b2': 6, 'b3': 7,
317  'v1': 8, 'v2': 9, 'v3':10}
318  if read_derivs:
319  indicesd = {'dd': 1,
320  'de': 9,
321  'dp1': 10, 'dp2': 11, 'dp3': 12,
322  'db1': 13, 'db2': 14, 'db3': 15,}
323  indices.update(indicesd)
324  elif solver[0:4] == 'zeus':
325  indices = {'d': 0,
326  'e1': 1, 'et': 2,
327  'p1': 3, 'p2': 4, 'p3': 5,
328  'b1': 6, 'b2': 7, 'b3': 8,
329  'gp': 9}
330  elif solver.lower() == 'rt_solver':
331  indices = {'Q': 0,
332  'S': 1,
333  'rk': 2}
334  elif solver.lower() == 'poisson':
335  indices = {'d': 0, 'phi': 1}
336  if read_derivs:
337  indicesd = {'dd': 0, 'dphi': 1}
338  indices.updates(indicesd)
339  else:
340  raise Exception("Unrecognised solver type! => "+solver_type.lower())
341  return indices
342 
343 def Search(fd, ip, verbose=False):
344  line=fd.readline().split()
345  ioformat=int(line[0])
346  ioformat=2
347  status=False
348  while (True):
349  try:
350  line=fd.readline().split()
351  id=np.int(line[0])
352  if (verbose>1):
353  print(id,ip)
354  if id==ip:
355  status=True
356  if (verbose):
357  print(id,ip,status)
358  break
359  nskip=np.int(line[1])
360  for i in range(nskip):
361  void=fd.readline()
362  except:
363  status=False
364  break
365  return status
366 
367 def Patches(file, overlap=0.01, verbose=False):
368  with open(file+'.txt','r') as fd:
369  patches=[]
370  line=fd.readline().split()
371  ioformat=int(line[0])
372  ioformat=2
373  while (True):
374  try:
375  id,nskip=fd.readline().split()
376  except:
377  break
378  patch=Patch(file,fd=fd,overlap=overlap,verbose=verbose)
379  patches.append(patch)
380  return patches
381 
382 class Snapshot:
383  def __init__(self, iout=0, run='', data='../data', overlap=0.01, verbose=0):
384  start=time()
385  if run=='':
386  dir=data+'/{:05d}/'.format(iout)
387  else:
388  dir=data+'/'+run+'/{:05d}/'.format(iout)
389  self.directory=dir
390  self.rankfiles=[dir+f.replace('.txt','') for f in os.listdir(dir) if f.endswith('.rank.txt')]
391  if (verbose):
392  print(self.directory, self.rankfiles)
393  # If .rank.txt files were found, add patches listed there
394  if np.size(self.rankfiles)>0:
395  self.patches=[]
396  for file in self.rankfiles:
397  patches=Patches(file,overlap=overlap,verbose=verbose)
398  for patch in patches:
399  self.patches.append(patch)
400  self.files=[]
401  for p in self.patches:
402  file=self.directory+'{:05d}'.format(p.id)
403  self.files.append(file)
404  # If not, list .txt files, and add one patch from each file
405  else:
406  self.patches = []
407  # ioformat <= 2
408  suffix = '.info'
409  self.files=[dir+f.replace(suffix,'') for f in os.listdir(dir) if f.endswith(suffix)]
410  if len(self.files) == 0:
411  # ioformat > 2
412  suffix = '.txt'
413  self.files=[dir+f.replace(suffix,'') for f in os.listdir(dir) if f.endswith(suffix)]
414  for file in self.files:
415  self.patches.append (Patch(file,suffix=suffix,overlap=overlap))
416  if verbose:
417  print('{} snapshot.patches ({:.2f} s)'.format(np.size(self.patches),time()-start))
418 
def _read_legacy_data(self, filename, dims, nvar, nbytes)
def stripv(self, v)
def variable_indices(self, read_derivs=False)