cwpearson/astaroth
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Improved Python file reading routines.

Browse Source
This commit is contained in:
Miikka Vaisala
2020-09-14 15:40:29 +08:00
parent 948b16aaaa
commit 66afc21162
1 changed files with 296 additions and 50 deletions
Download Patch File Download Diff File Expand all files Collapse all files

346
analysis/python/astar/data/read.py
View File

@@ -28,26 +28,28 @@ 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)
print("type_instruction", type_instruction)
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:
print(filename)
if getfilename:
print(filename)
array = np.fromfile(filename, dtype=my_dtype)
timestamp = array[0]
@@ -92,10 +94,58 @@ def read_meshtxt(fdir, fname, dbg_output):
print(line[1], contents[line[1]])
else:
print(line)
print('ERROR: ' + line[0] +' not recognized!')
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)
uuy, timestamp, ok = read_bin('VTXBUF_UUY', fdir, fnum, self.minfo)
uuz, timestamp, ok = read_bin('VTXBUF_UUZ', 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, getfilename=pdiag)
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)
aay, timestamp, ok = read_bin('VTXBUF_AY', fdir, fnum, self.minfo)
aaz, timestamp, ok = read_bin('VTXBUF_AZ', 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, getfilename=pdiag)
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,42 +411,143 @@ 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):
with open(fdir+fname) as f:
filetext = f.read().splitlines()
def parse_ts(fdir, fname, debug = False):
var = {}
line = filetext[0].split()
for i in range(len(line)):
line[i] = line[i].replace('VTXBUF_', "")
line[i] = line[i].replace('UU', "uu")
line[i] = line[i].replace('_total', "tot")
line[i] = line[i].replace('ACCRETION', "acc")
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('X', "x")
line[i] = line[i].replace('Y', "y")
line[i] = line[i].replace('Z', "z")
tsfile = fdir+fname
tsdata = np.loadtxt(fdir+fname,skiprows=1)
if os.path.exists(tsfile):
for i in range(len(line)):
var[line[i]] = tsdata[:,i]
with open(tsfile) as f:
filetext = f.read().splitlines()
var['step'] = np.int64(var['step'])
line = filetext[0].split()
for i in range(len(line)):
line[i] = line[i].replace('VTXBUF_', "")
line[i] = line[i].replace('UU', "uu")
line[i] = line[i].replace('_total', "tot")
line[i] = line[i].replace('ACCRETION', "acc")
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")
print("HERE ARE ALL KEYS FOR TS DATA:")
print(var.keys())
#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)