Improved Python file reading routines.

analysis/python/astar/data/read.py

@@ -28,25 +28,27 @@ try:
 except ImportError:
     yt_present = False
 
-def set_dtype(endian, AcRealSize):
+def set_dtype(endian, AcRealSize, print_type = True):
     if endian == 0:
         en = '>'
     elif endian == 1: 
         en = '<'
     type_instruction = en + 'f' + str(AcRealSize)
+    if print_type:
         print("type_instruction", type_instruction)
     my_dtype = np.dtype(type_instruction)
     return my_dtype
 
-def read_bin(fname, fdir, fnum, minfo, numtype=np.longdouble):
+def read_bin(fname, fdir, fnum, minfo, numtype=np.longdouble, getfilename=True):
     '''Read in a floating point array'''
     filename = fdir + fname + '_' + fnum + '.mesh'
     datas = np.DataSource()
     read_ok = datas.exists(filename)
 
-    my_dtype = set_dtype(minfo.contents['endian'], minfo.contents['AcRealSize'])
+    my_dtype = set_dtype(minfo.contents['endian'], minfo.contents['AcRealSize'], print_type=getfilename)
 
     if read_ok:
+        if getfilename:
             print(filename)
         array = np.fromfile(filename, dtype=my_dtype)
 
@@ -92,10 +94,58 @@ def read_meshtxt(fdir, fname, dbg_output):
                 print(line[1], contents[line[1]])
         else: 
             print(line)
+            if dbg_output:
                 print('ERROR: ' + line[0] +' not recognized!')
 
     return contents
 
+def parse_directory(meshdir):
+    dirlist = os.listdir(meshdir)
+    dirlist = [k for k in dirlist if 'LNRHO' in k]
+    for i, item in enumerate(dirlist):
+        tmp = item.strip('.mesh')
+        tmp = tmp.strip('VTXBUF_LNRHO')
+        dirlist[i] = int(tmp)
+    dirlist.sort()
+    return dirlist
+
+def apply_boundcond(array, btype):
+
+    if btype == "p":
+        be = 3
+        bi = 6 
+        # Edges
+        # xx
+        array[  : 3,  :  ,  :  ] = array[-6:-3,  :  ,  :  ]
+        array[-3:  ,  :  ,  :  ] = array[ 3: 6,  :  ,  :  ]
+        # yy
+        array[  :  ,  : 3,  :  ] = array[  :  ,-6:-3,  :  ]
+        array[  :  ,-3:  ,  :  ] = array[  :  , 3: 6,  :  ]
+        # zz 
+        array[  :  ,  :  ,  : 3] = array[  :  ,  :  ,-6:-3]
+        array[  :  ,  :  ,-3:  ] = array[  :  ,  :  , 3: 6]
+        # Corner parts
+        # xy
+        array[  : 3,  : 3,  :  ] = array[-6:-3,-6:-3,  :  ]
+        array[-3:  ,-3:  ,  :  ] = array[ 3: 6, 3: 6,  :  ]
+        array[-3:  ,  : 3,  :  ] = array[ 3: 6,-6:-3,  :  ]
+        array[  : 3,-3:  ,  :  ] = array[-6:-3, 3: 6,  :  ]
+        # xz
+        array[  : 3,  :  ,  : 3] = array[-6:-3,  :  ,-6:-3]
+        array[-3:  ,  :  ,-3:  ] = array[ 3: 6,  :  , 3: 6]
+        array[-3:  ,  :  ,  : 3] = array[ 3: 6,  :  ,-6:-3]
+        array[  : 3,  :  ,-3:  ] = array[-6:-3,  :  , 3: 6]
+        # yz
+        array[  :  ,  : 3,  : 3] = array[  :  ,-6:-3,-6:-3]
+        array[  :  ,-3:  ,-3:  ] = array[  :  , 3: 6, 3: 6]
+        array[  :  ,-3:  ,  : 3] = array[  :  , 3: 6,-6:-3]
+        array[  :  ,  : 3,-3:  ] = array[  :  ,-6:-3, 3: 6]
+    else:
+        print("Unknown btype", btype)
+
+    return array
+
+
 def DERX(array, dx):
     output = np.zeros_like(array)
     for i in range(3, array.shape[0]-3): #Keep boundary poits as 0
@@ -120,6 +170,34 @@ def DERZ(array, dz):
                          -     array[:,:,i-3] +      array[:,:,i+3] )/(60.0*dz)
     return output
 
+def DER2X(array, dx):
+    output = np.zeros_like(array)
+    for i in range(3, array.shape[0]-3): #Keep boundary poits as 0
+        output[i,:,:] =(    2.0*array[i-1,:,:] +   2.0*array[i+1,:,:]
+                         - 27.0*array[i-2,:,:] -  27.0*array[i+2,:,:]
+                         +270.0*array[i-3,:,:] + 270.0*array[i+3,:,:] 
+                         -490.0*array[i  ,:,:]        )/(180.0*dx*dx)
+    return output
+
+def DER2Y(array, dy):
+    output = np.zeros_like(array)
+    for i in range(3,array.shape[1]-3):
+        output[:,i,:] =(    2.0*array[:,i-1,:] +   2.0*array[:,i+1,:]
+                         - 27.0*array[:,i-2,:] -  27.0*array[:,i+2,:]
+                         +270.0*array[:,i-3,:] + 270.0*array[:,i+3,:] 
+                         -490.0*array[:,i  ,:]        )/(180.0*dy*dy)
+    return output
+
+def DER2Z(array, dz):
+    output = np.zeros_like(array)
+    for i in range(3, array.shape[2]-3):
+        output[:,:,i] =(    2.0*array[:,:,i-1] +   2.0*array[:,:,i+1]
+                         - 27.0*array[:,:,i-2] -  27.0*array[:,:,i+2]
+                         +270.0*array[:,:,i-3] + 270.0*array[:,:,i+3] 
+                         -490.0*array[:,:,i  ]        )/(180.0*dz*dz)
+    return output
+
+
 def curl(aa, minfo):
     dx = minfo.contents['AC_dsx']
     dy = minfo.contents['AC_dsy']
@@ -128,6 +206,44 @@ def curl(aa, minfo):
             DERZ(aa[0], dz)-DERX(aa[2], dx), 
             DERX(aa[1], dx)-DERY(aa[0], dy))
 
+def div(array, minfo):
+    dx = minfo.contents['AC_dsx']
+    dy = minfo.contents['AC_dsy']
+    dz = minfo.contents['AC_dsz']
+    return (  DERX(array[0], dx)  
+            + DERY(array[1], dy)  
+            + DERZ(array[2], dz))
+
+def grad(array, minfo):
+    dx = minfo.contents['AC_dsx']
+    dy = minfo.contents['AC_dsy']
+    dz = minfo.contents['AC_dsz']
+    return (DERX(array, dx), 
+            DERY(array, dy), 
+            DERZ(array, dz))
+
+def grad_div(array, minfo):
+    scalar = div(array, minfo)
+    scalar = apply_boundcond(scalar, "p") 
+    vec = grad(scalar, minfo)
+    return vec
+
+def laplace_scal(array, minfo):
+    dx = minfo.contents['AC_dsx']
+    dy = minfo.contents['AC_dsy']
+    dz = minfo.contents['AC_dsz']
+    return (DER2X(array, dx) + DER2Y(array, dy) + DER2Z(array, dz))
+
+def laplace_vec(array, minfo):
+    return (laplace_scal(array[0], minfo), 
+            laplace_scal(array[1], minfo), 
+            laplace_scal(array[2], minfo))
+
+def curl_of_curl(array, minfo):
+    array1 = curl(array, minfo) 
+    array2 = (apply_boundcond(array1[0], "p"), apply_boundcond(array1[1], "p"), apply_boundcond(array1[2], "p"))
+    return curl(array2, minfo)
+
 
 class MeshInfo():
     '''Object that contains all mesh info'''
@@ -138,45 +254,40 @@ class MeshInfo():
 class Mesh:
     '''Class tha contains all 3d mesh data'''
 
-    def __init__(self, fnum, fdir=""):
+    def __init__(self, fnum, fdir="", only_info = False, pdiag = True):
         fnum = str(fnum)
         self.framenum = fnum.zfill(10)
 
         self.minfo = MeshInfo(fdir)
 
-        self.lnrho, self.timestamp, self.ok = read_bin('VTXBUF_LNRHO', fdir, fnum, self.minfo)
+        if only_info == False:
+            self.lnrho, self.timestamp, self.ok = read_bin('VTXBUF_LNRHO', fdir, fnum, self.minfo, getfilename=pdiag)
+        else:
+            self.ok = False
 
         if self.ok:
 
-            self.ss, timestamp, ok = read_bin('VTXBUF_ENTROPY', fdir, fnum, self.minfo)
+            self.ss, timestamp, ok = read_bin('VTXBUF_ENTROPY', fdir, fnum, self.minfo, getfilename=pdiag)
 
-            self.accretion, timestamp, ok = read_bin('VTXBUF_ACCRETION', fdir, fnum, self.minfo)
+            self.accretion, timestamp, ok = read_bin('VTXBUF_ACCRETION', fdir, fnum, self.minfo, getfilename=pdiag)
  
             #TODO Generalize is a dict. Do not hardcode!  
-            uux, timestamp, ok = read_bin('VTXBUF_UUX', fdir, fnum, self.minfo)
+            uux, timestamp, ok = read_bin('VTXBUF_UUX', fdir, fnum, self.minfo, getfilename=pdiag)
-            uuy, timestamp, ok = read_bin('VTXBUF_UUY', fdir, fnum, self.minfo) 
+            uuy, timestamp, ok = read_bin('VTXBUF_UUY', fdir, fnum, self.minfo, getfilename=pdiag) 
-            uuz, timestamp, ok = read_bin('VTXBUF_UUZ', fdir, fnum, self.minfo)
+            uuz, timestamp, ok = read_bin('VTXBUF_UUZ', fdir, fnum, self.minfo, getfilename=pdiag)
             self.uu = (uux, uuy, uuz)
             uux = []
             uuy = [] 
             uuz = []
  
-            aax, timestamp, ok = read_bin('VTXBUF_AX', fdir, fnum, self.minfo)
+            aax, timestamp, ok = read_bin('VTXBUF_AX', fdir, fnum, self.minfo, getfilename=pdiag)
-            aay, timestamp, ok = read_bin('VTXBUF_AY', fdir, fnum, self.minfo) 
+            aay, timestamp, ok = read_bin('VTXBUF_AY', fdir, fnum, self.minfo, getfilename=pdiag) 
-            aaz, timestamp, ok = read_bin('VTXBUF_AZ', fdir, fnum, self.minfo)
+            aaz, timestamp, ok = read_bin('VTXBUF_AZ', fdir, fnum, self.minfo, getfilename=pdiag)
             self.aa = (aax, aay, aaz)
             aax = []
             aay = [] 
             aaz = []
 
-             
-            #self.aa[0][:,:,:] = 0.0
-            #self.aa[1][:,:,:] = 0.0 
-            #self.aa[2][:,:,:] = 0.0 
-            #for i in range(0, self.aa[0].shape[0]):
-            #    self.aa[0][:,i,:] = float(i)
-   
-
             self.xx =  np.arange(self.minfo.contents['AC_mx']) * self.minfo.contents['AC_dsx']
             self.yy =  np.arange(self.minfo.contents['AC_my']) * self.minfo.contents['AC_dsy']
             self.zz =  np.arange(self.minfo.contents['AC_mz']) * self.minfo.contents['AC_dsz']
@@ -184,10 +295,44 @@ class Mesh:
             self.xmid = int(self.minfo.contents['AC_mx']/2)
             self.ymid = int(self.minfo.contents['AC_my']/2)
             self.zmid = int(self.minfo.contents['AC_mz']/2)
-    def Bfield(self, get_jj = False):
+    def Bfield(self, get_jj = False, trim=False):
         self.bb = curl(self.aa, self.minfo) 
         if get_jj:
-            self.jj = curl(self.bb, self.minfo) 
+            self.jj = curl_of_curl(self.aa, self.minfo)
+        if trim:
+            self.bb = (    self.bb[0][3:-3, 3:-3, 3:-3],self.bb[1][3:-3, 3:-3, 3:-3],self.bb[2][3:-3, 3:-3, 3:-3])
+            if get_jj:
+                self.jj = (self.jj[0][3:-3, 3:-3, 3:-3],self.jj[1][3:-3, 3:-3, 3:-3],self.jj[2][3:-3, 3:-3, 3:-3])
+
+    def get_jj(self, trim=False):
+        self.jj = curl_of_curl(self.aa, minfo, trim=False)
+        if trim:
+            self.jj =     (self.jj[0][3:-3, 3:-3, 3:-3],self.jj[1][3:-3, 3:-3, 3:-3],self.jj[2][3:-3, 3:-3, 3:-3])
+
+    def vorticity(self, trim=False):
+        self.oo = curl(self.uu, self.minfo) 
+        if trim:
+            self.oo =     (self.oo[0][3:-3, 3:-3, 3:-3],self.oo[1][3:-3, 3:-3, 3:-3],self.oo[2][3:-3, 3:-3, 3:-3])
+ 
+    def rad_vel(self):
+        print("Calculating spherical velocity components")
+        self.uu_pherical = np.zeros_like(self.uu)
+        xx, yy, zz = np.meshgrid(self.xx - self.xmid, self.yy - self.ymid, self.zz - self.zmid)       
+        rr                         = np.sqrt(xx**2.0 + yy**2.0 + zz**2.0)
+        theta                      = np.arccos(zz/rr)
+        phi                        = np.arctan2(yy,xx) 
+        sin_theta_sin_phi          = np.sin(theta)*np.sin(phi)
+        cos_theta_cos_phi          = np.cos(theta)*np.cos(phi)
+        sin_theta_cos_phi          = np.sin(theta)*np.cos(phi)
+        cos_theta_sin_phi          = np.cos(theta)*np.sin(phi)
+        ux = self.uu[0]; uy = self.uu[1]; uz = self.uu[2];
+        vr     = sin_theta_cos_phi*ux + sin_theta_sin_phi*uy + np.cos(theta)*uz
+        vtheta = cos_theta_cos_phi*ux + cos_theta_sin_phi*uy - np.sin(theta)*uz 
+        vphi   =      -np.sin(phi)*ux +       np.cos(phi)*uy 
+        self.uu_pherical[0] = vr
+        self.uu_pherical[1] = vtheta
+        self.uu_pherical[2] = vphi
+
 
     def yt_conversion(self):
         if yt_present:
@@ -266,14 +411,105 @@ class Mesh:
         f.write(binary_format)
         f.close()
 
+    def export_vtk_ascii(self, Beq = 1.0):
+        #BASED ON https://lorensen.github.io/VTKExamples/site/VTKFileFormats/#dataset-attribute-format 
+        self.Bfield()
+
+        f = open("GRID%s.vtk" % self.framenum, 'w')
+
+        Ntot = self.minfo.contents['AC_mx']*self.minfo.contents['AC_my']*self.minfo.contents['AC_mz']
+
+        mx = self.minfo.contents['AC_mx']
+        my = self.minfo.contents['AC_my']
+        mz = self.minfo.contents['AC_mz']
+
+        print("Writing GRID%s.vtk" % self.framenum)
+
+        f.write("# vtk DataFile Version 2.0\n")
+        f.write("Astaroth grid for visualization\n")
+        f.write("ASCII\n")
+        #f.write("DATASET STRUCTURED_GRID\n")
+        #f.write("DIMENSIONS %i %i %i \n" % (mx, my, mz))
+        #f.write("POINTS %i float \n" % (Ntot))
+        #for i in range(mx):
+        #    for j in range(my):
+        #        for k in range(mz):
+        #             f.write("%e %e %e \n" % (i, j, k))
+        f.write("DATASET RECTILINEAR_GRID\n")
+        f.write("DIMENSIONS %i %i %i \n" % (mx, my, mz))
+        f.write("X_COORDINATES %i float \n" % mx)
+        for i in range(mx):  
+            f.write("%e " % (i))
+        f.write("\n")
+        f.write("Y_COORDINATES %i float \n" % my)
+        for j in range(my):  
+            f.write("%e " % (j))
+        f.write("\n")
+        f.write("Z_COORDINATES %i float \n" % mz)
+        for k in range(mz):  
+            f.write("%e " % (k))
+        f.write("\n")
+        f.write("POINT_DATA %i \n" % (Ntot))
+        f.write("VECTORS velocity float \n" )
+        for i in range(mx):
+            if (i % 8) == 0:
+                print("i = %i / %i" %(i, mx-1))
+            for j in range(my):
+                for k in range(mz):
+                     f.write("%e %e %e \n" % ( self.uu[0][i, j, k], self.uu[1][i, j, k], self.uu[2][i, j, k]))
+        #f.write("POINT_DATA %i \n" % (Ntot))
+        f.write("VECTORS bfield float \n")
+        eqprint = True
+        for i in range(mx):
+            if (i % 8) == 0:
+                print("i = %i / %i" %(i, mx-1))
+                eqprint = True
+            for j in range(my):
+                for k in range(mz):
+                     #Beq is the equipartition magnetic field.
+                     while(eqprint):
+                         print("normal        B %e %e %e \n" % (self.bb[0][i, j, k],     self.bb[1][i, j, k],     self.bb[2][i, j, k]    ))
+                         print("equipartition B %e %e %e \n" % (self.bb[0][i, j, k]/Beq, self.bb[1][i, j, k]/Beq, self.bb[2][i, j, k]/Beq))
+                         eqprint = False
+                     f.write("%e %e %e \n" % ( self.bb[0][i, j, k]/Beq, self.bb[1][i, j, k]/Beq, self.bb[2][i, j, k]/Beq))
+        #ADD DENSITY SCALAR
+        f.write("\n")
+
+        print("Done.")
+
+        f.close()
+
+def find_explosion_index(array, criterion = 1e5):
+    for i, ar in enumerate(array):
+        if (np.abs(array[i])-np.abs(array[i-1])) > criterion:
+            return i
+    return -1
+
+def mask_bad_values_ts(ts, criterion = 1e5):
+    indexmask = np.zeros_like(ts)
+    
+    ts = np.ma.masked_invalid(ts)
+    
+    index  = find_explosion_index(ts, criterion = criterion)
+    
+    if index >= 0:
+        indexmask[index:] = 1
+    
+    ts = np.ma.array(ts, mask=indexmask)
+    
+    return ts
 
 
-def parse_ts(fdir, fname):
+def parse_ts(fdir, fname, debug = False):
-    with open(fdir+fname) as f:
-        filetext = f.read().splitlines()
-
     var = {}  
 
+    tsfile = fdir+fname
+
+    if os.path.exists(tsfile): 
+
+        with open(tsfile) as f:
+            filetext = f.read().splitlines()
+
         line = filetext[0].split()
         for i in range(len(line)):
             line[i] = line[i].replace('VTXBUF_', "")
@@ -283,25 +519,35 @@ def parse_ts(fdir, fname):
             line[i] = line[i].replace('A', "aa")
             line[i] = line[i].replace('LNRHO', "lnrho")
             line[i] = line[i].replace('ENTROPY', "ss")
+            line[i] = line[i].replace('BFIELD', "bb")
             line[i] = line[i].replace('X', "x")
             line[i] = line[i].replace('Y', "y")
             line[i] = line[i].replace('Z', "z")
+            line[i] = line[i].replace('vaa', "vA")
 
-    tsdata = np.loadtxt(fdir+fname,skiprows=1)
+        #tsdata = np.loadtxt(fdir+fname,skiprows=1)
+        tsdata = np.genfromtxt(fdir+fname,skip_header=1, skip_footer=1)
 
         for i in range(len(line)):
             var[line[i]] = tsdata[:,i]
 
         var['step'] = np.int64(var['step'])
 
+        var['exist'] = True
+
+    else:
+        var['exist'] = False
+
+    if debug:
         print("HERE ARE ALL KEYS FOR TS DATA:")
         print(var.keys())
    
     return var
 
+
 class TimeSeries:
     '''Class for time series data'''
 
-    def __init__(self, fdir="", fname="timeseries.ts"):
+    def __init__(self, fdir="", fname="timeseries.ts", debug = False):
 
-        self.var = parse_ts(fdir, fname)
+        self.var = parse_ts(fdir, fname, debug = debug)