From d947bdccb8b83256e9d651c32bf7c7c8b571285d Mon Sep 17 00:00:00 2001 From: Miikka Vaisala Date: Tue, 14 Jan 2020 14:23:24 +0800 Subject: [PATCH] General purpose Python tool improvements. --- analysis/python/astar/data/read.py | 115 ++++++++++++++++++++--- analysis/python/astar/visual/lineplot.py | 28 +++++- analysis/python/astar/visual/slices.py | 55 ++++++----- 3 files changed, 162 insertions(+), 36 deletions(-) diff --git a/analysis/python/astar/data/read.py b/analysis/python/astar/data/read.py index d97d449..9b2b820 100644 --- a/analysis/python/astar/data/read.py +++ b/analysis/python/astar/data/read.py @@ -21,6 +21,12 @@ import numpy as np +#Optional YT interface +try: + import yt + yt_present = True +except ImportError: + yt_present = False def set_dtype(endian, AcRealSize): if endian == 0: @@ -55,7 +61,8 @@ def read_bin(fname, fdir, fnum, minfo, numtype=np.longdouble): return array, timestamp, read_ok -def read_meshtxt(fdir, fname): +def read_meshtxt(fdir, fname, dbg_output): + with open(fdir+fname) as f: filetext = f.read().splitlines() @@ -65,30 +72,33 @@ def read_meshtxt(fdir, fname): line = line.split() if line[0] == 'int': contents[line[1]] = np.int(line[2]) - print(line[1], contents[line[1]]) + if dbg_output: + print(line[1], contents[line[1]]) elif line[0] == 'size_t': contents[line[1]] = np.int(line[2]) - print(line[1], contents[line[1]]) + if dbg_output: + print(line[1], contents[line[1]]) elif line[0] == 'int3': contents[line[1]] = [np.int(line[2]), np.int(line[3]), np.int(line[4])] - print(line[1], contents[line[1]]) + if dbg_output: + print(line[1], contents[line[1]]) elif line[0] == 'real': contents[line[1]] = np.float(line[2]) - print(line[1], contents[line[1]]) + if dbg_output: + print(line[1], contents[line[1]]) elif line[0] == 'real3': contents[line[1]] = [np.float(line[2]), np.float(line[3]), np.float(line[4])] - print(line[1], contents[line[1]]) + if dbg_output: + print(line[1], contents[line[1]]) else: print(line) print('ERROR: ' + line[0] +' not recognized!') return contents -#Now just 2nd order def DERX(array, dx): output = np.zeros_like(array) for i in range(3, array.shape[0]-3): #Keep boundary poits as 0 - #output[i,:,:] = (- array[i-1,:,:] + array[i+1,:,:])/(2.0*dx) output[i,:,:] =( -45.0*array[i-1,:,:] + 45.0*array[i+1,:,:] + 9.0*array[i-2,:,:] - 9.0*array[i+2,:,:] - array[i-3,:,:] + array[i+3,:,:] )/(60.0*dx) @@ -97,7 +107,6 @@ def DERX(array, dx): def DERY(array, dy): output = np.zeros_like(array) for i in range(3,array.shape[1]-3): - #output[:,i,:] = (- array[:,i-1,:] + array[:,i+1,:])/(2.0*dy) output[:,i,:] =( -45.0*array[:,i-1,:] + 45.0*array[:,i+1,:] + 9.0*array[:,i-2,:] - 9.0*array[:,i+2,:] - array[:,i-3,:] + array[:,i+3,:] )/(60.0*dy) @@ -106,7 +115,6 @@ def DERY(array, dy): def DERZ(array, dz): output = np.zeros_like(array) for i in range(3, array.shape[2]-3): - #output[:,:,i] = (- array[:,:,i-1] + array[:,:,i+1])/(2.0*dz) output[:,:,i] =( -45.0*array[:,:,i-1] + 45.0*array[:,:,i+1] + 9.0*array[:,:,i-2] - 9.0*array[:,:,i+2] - array[:,:,i-3] + array[:,:,i+3] )/(60.0*dz) @@ -124,8 +132,8 @@ def curl(aa, minfo): class MeshInfo(): '''Object that contains all mesh info''' - def __init__(self, fdir): - self.contents = read_meshtxt(fdir, 'mesh_info.list') + def __init__(self, fdir, dbg_output=False): + self.contents = read_meshtxt(fdir, 'mesh_info.list', dbg_output) class Mesh: '''Class tha contains all 3d mesh data''' @@ -176,9 +184,88 @@ 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): + def Bfield(self, get_jj = False): self.bb = curl(self.aa, self.minfo) - print(self.bb[2]) + if get_jj: + self.jj = curl(self.bb, self.minfo) + + def yt_conversion(self): + if yt_present: + self.ytdict = dict(density = (np.exp(self.lnrho)*self.minfo.contents['AC_unit_density'], "g/cm**3"), + uux = (self.uu[0]*self.minfo.contents['AC_unit_velocity'], "cm/s"), + uuy = (self.uu[1]*self.minfo.contents['AC_unit_velocity'], "cm/s"), + uuz = (self.uu[2]*self.minfo.contents['AC_unit_velocity'], "cm/s"), + bbx = (self.bb[0]*self.minfo.contents['AC_unit_magnetic'], "gauss"), + bby = (self.bb[1]*self.minfo.contents['AC_unit_magnetic'], "gauss"), + bbz = (self.bb[2]*self.minfo.contents['AC_unit_magnetic'], "gauss"), + ) + bbox = self.minfo.contents['AC_unit_length'] \ + *np.array([[self.xx.min(), self.xx.max()], [self.yy.min(), self.yy.max()], [self.zz.min(), self.zz.max()]]) + self.ytdata = yt.load_uniform_grid(self.ytdict, self.lnrho.shape, length_unit="cm", bbox=bbox) + else: + print("ERROR. No YT support found!") + + def export_csv(self): + csvfile = open("grid.csv.%s" % self.framenum, "w") + csvfile.write("xx, yy, zz, rho, uux, uuy, uuz, bbx, bby, bbz\n") + ul = self.minfo.contents['AC_unit_length'] + uv = self.minfo.contents['AC_unit_velocity'] + ud = self.minfo.contents['AC_unit_density'] + um = self.minfo.contents['AC_unit_magnetic'] + for kk in np.arange(3, self.zz.size-3): + for jj in np.arange(3, self.yy.size-3): + for ii in np.arange(3, self.xx.size-3): + #print(self.xx.size, self.yy.size, self.zz.size) + linestring = "%e, %e, %e, %e, %e, %e, %e, %e, %e, %e\n"% (self.xx[ii]*ul, self.yy[jj]*ul, self.zz[kk]*ul, + np.exp(self.lnrho[ii, jj, kk])*ud, + self.uu[0][ii, jj, kk]*uv, self.uu[1][ii, jj, kk]*uv, + self.uu[2][ii, jj, kk]*uv, + self.bb[0][ii, jj, kk]*um, self.bb[1][ii, jj, kk]*um, + self.bb[2][ii, jj, kk]*um) + csvfile.write(linestring) + csvfile.close() + + def export_raw(self): + uv = self.minfo.contents['AC_unit_velocity'] + ud = self.minfo.contents['AC_unit_density'] + um = self.minfo.contents['AC_unit_magnetic'] + print(self.lnrho.shape, set_dtype(self.minfo.contents['endian'], self.minfo.contents['AcRealSize'])) + + f = open("rho%s.raw" % self.framenum, 'w+b') + binary_format =(np.exp(self.lnrho)*ud).tobytes() + f.write(binary_format) + f.close() + + f = open("uux%s.raw" % self.framenum, 'w+b') + binary_format =(self.uu[0]*uv).tobytes() + f.write(binary_format) + f.close() + + f = open("uuy%s.raw" % self.framenum, 'w+b') + binary_format =(self.uu[1]*uv).tobytes() + f.write(binary_format) + f.close() + + f = open("uuz%s.raw" % self.framenum, 'w+b') + binary_format =(self.uu[2]*uv).tobytes() + f.write(binary_format) + f.close() + + f = open("bbx%s.raw" % self.framenum, 'w+b') + binary_format =(self.bb[0]*um).tobytes() + f.write(binary_format) + f.close() + + f = open("bby%s.raw" % self.framenum, 'w+b') + binary_format =(self.bb[1]*um).tobytes() + f.write(binary_format) + f.close() + + f = open("bbz%s.raw" % self.framenum, 'w+b') + binary_format =(self.bb[2]*um).tobytes() + f.write(binary_format) + f.close() + def parse_ts(fdir, fname): diff --git a/analysis/python/astar/visual/lineplot.py b/analysis/python/astar/visual/lineplot.py index d70acf4..e793a6f 100644 --- a/analysis/python/astar/visual/lineplot.py +++ b/analysis/python/astar/visual/lineplot.py @@ -37,7 +37,9 @@ def plot_min_man_rms(ts, xaxis, yaxis1, yaxis2, yaxis3): plt.close() def plot_ts(ts, isotherm=False, show_all=False, lnrho=False, uutot=False, - uux=False, uuy=False, uuz=False, ss=False, acc=False, sink=False): + uux=False, uuy=False, uuz=False, + aax=False, aay=False, aaz=False, + ss=False, acc=False, sink=False, rho=False): if show_all: lnrho=True @@ -108,6 +110,30 @@ def plot_ts(ts, isotherm=False, show_all=False, lnrho=False, uutot=False, yaxis2 = 'uuz_min' yaxis3 = 'uuz_max' plot_min_man_rms(ts, xaxis, yaxis1, yaxis2, yaxis3) + + if aax: + plt.figure() + xaxis = 't_step' + yaxis1 = 'aax_rms' + yaxis2 = 'aax_min' + yaxis3 = 'aax_max' + plot_min_man_rms(ts, xaxis, yaxis1, yaxis2, yaxis3) + + if aay: + plt.figure() + xaxis = 't_step' + yaxis1 = 'aay_rms' + yaxis2 = 'aay_min' + yaxis3 = 'aay_max' + plot_min_man_rms(ts, xaxis, yaxis1, yaxis2, yaxis3) + + if aaz: + plt.figure() + xaxis = 't_step' + yaxis1 = 'aaz_rms' + yaxis2 = 'aaz_min' + yaxis3 = 'aaz_max' + plot_min_man_rms(ts, xaxis, yaxis1, yaxis2, yaxis3) if ss: plt.figure() diff --git a/analysis/python/astar/visual/slices.py b/analysis/python/astar/visual/slices.py index 5b0aadc..4d95ce8 100644 --- a/analysis/python/astar/visual/slices.py +++ b/analysis/python/astar/visual/slices.py @@ -26,7 +26,8 @@ CM_INFERNO = plt.get_cmap('inferno') def plot_3(mesh, input_grid, title = '', fname = 'default', bitmap=False, slicetype = 'middle', colrange=None, colormap=CM_INFERNO , - contourplot=False, points_from_centre = -1, bfieldlines=False, velfieldlines=False): + contourplot=False, points_from_centre = -1, bfieldlines=False, velfieldlines=False, trimghost = 0): + fig = plt.figure(figsize=(8, 8)) grid = gridspec.GridSpec(2, 3, wspace=0.4, hspace=0.4, width_ratios=[1,1, 0.15]) ax00 = fig.add_subplot( grid[0,0] ) @@ -40,21 +41,30 @@ def plot_3(mesh, input_grid, title = '', fname = 'default', bitmap=False, yz_slice = input_grid[mesh.xmid, :, :] xz_slice = input_grid[:, mesh.ymid, :] xy_slice = input_grid[:, :, mesh.zmid] - if colrange==None: - plotnorm = colors.Normalize(vmin=input_grid.min(),vmax=input_grid.max()) - else: - plotnorm = colors.Normalize(vmin=colrange[0],vmax=colrange[1]) elif slicetype == 'sum': yz_slice = np.sum(input_grid, axis=0) xz_slice = np.sum(input_grid, axis=1) xy_slice = np.sum(input_grid, axis=2) - cmin = np.amin([yz_slice.min(), xz_slice.min(), xy_slice.min()]) - cmax = np.amax([yz_slice.max(), xz_slice.max(), xy_slice.max()]) - if colrange==None: - plotnorm = colors.Normalize(vmin=cmin,vmax=cmax) - else: - plotnorm = colors.Normalize(vmin=colrange[0],vmax=colrange[1]) - + + yz_slice = yz_slice[trimghost : yz_slice.shape[0]-trimghost, + trimghost : yz_slice.shape[1]-trimghost] + xz_slice = xz_slice[trimghost : xz_slice.shape[0]-trimghost, + trimghost : xz_slice.shape[1]-trimghost] + xy_slice = xy_slice[trimghost : xy_slice.shape[0]-trimghost, + trimghost : xy_slice.shape[1]-trimghost] + mesh_xx_tmp = mesh.xx[trimghost : mesh.xx.shape[0]-trimghost] + mesh_yy_tmp = mesh.yy[trimghost : mesh.yy.shape[0]-trimghost] + mesh_zz_tmp = mesh.zz[trimghost : mesh.zz.shape[0]-trimghost] + + #Set the coulourscale + cmin = np.amin([yz_slice.min(), xz_slice.min(), xy_slice.min()]) + cmax = np.amax([yz_slice.max(), xz_slice.max(), xy_slice.max()]) + if colrange==None: + plotnorm = colors.Normalize(vmin=cmin,vmax=cmax) + else: + plotnorm = colors.Normalize(vmin=colrange[0],vmax=colrange[1]) + + if points_from_centre > 0: yz_slice = yz_slice[int(yz_slice.shape[0]/2)-points_from_centre : int(yz_slice.shape[0]/2)+points_from_centre, int(yz_slice.shape[1]/2)-points_from_centre : int(yz_slice.shape[1]/2)+points_from_centre] @@ -62,11 +72,11 @@ def plot_3(mesh, input_grid, title = '', fname = 'default', bitmap=False, int(xz_slice.shape[1]/2)-points_from_centre : int(xz_slice.shape[1]/2)+points_from_centre] xy_slice = xy_slice[int(xy_slice.shape[0]/2)-points_from_centre : int(xy_slice.shape[0]/2)+points_from_centre, int(xy_slice.shape[1]/2)-points_from_centre : int(xy_slice.shape[1]/2)+points_from_centre] - mesh.xx = mesh.xx[int(mesh.xx.shape[0]/2)-points_from_centre : int(mesh.xx.shape[0]/2)+points_from_centre] - mesh.yy = mesh.yy[int(mesh.yy.shape[0]/2)-points_from_centre : int(mesh.yy.shape[0]/2)+points_from_centre] - mesh.zz = mesh.zz[int(mesh.zz.shape[0]/2)-points_from_centre : int(mesh.zz.shape[0]/2)+points_from_centre] + mesh_xx_tmp = mesh.xx[int(mesh.xx.shape[0]/2)-points_from_centre : int(mesh.xx.shape[0]/2)+points_from_centre] + mesh_yy_tmp = mesh.yy[int(mesh.yy.shape[0]/2)-points_from_centre : int(mesh.yy.shape[0]/2)+points_from_centre] + mesh_zz_tmp = mesh.zz[int(mesh.zz.shape[0]/2)-points_from_centre : int(mesh.zz.shape[0]/2)+points_from_centre] - yy, zz = np.meshgrid(mesh.yy, mesh.zz, indexing='ij') + yy, zz = np.meshgrid(mesh_xx_tmp, mesh_xx_tmp, indexing='ij') if contourplot: map1 = ax00.contourf(yy, zz, yz_slice, norm=plotnorm, cmap=colormap, nlev=10) else: @@ -76,9 +86,10 @@ def plot_3(mesh, input_grid, title = '', fname = 'default', bitmap=False, ax00.set_title('%s t = %.4e' % (title, mesh.timestamp) ) ax00.set_aspect('equal') - ax00.contour(yy, zz, np.sqrt((yy-yy.max()/2.0)**2.0 + (zz-zz.max()/2.0)**2.0), [mesh.minfo.contents["AC_accretion_range"]]) + if mesh.minfo.contents["AC_accretion_range"] > 0.0: + ax00.contour(yy, zz, np.sqrt((yy-yy.max()/2.0)**2.0 + (zz-zz.max()/2.0)**2.0), [mesh.minfo.contents["AC_accretion_range"]]) - xx, zz = np.meshgrid(mesh.xx, mesh.zz, indexing='ij') + xx, zz = np.meshgrid(mesh_xx_tmp, mesh_zz_tmp, indexing='ij') if contourplot: ax10.contourf(xx, zz, xz_slice, norm=plotnorm, cmap=colormap, nlev=10) else: @@ -87,9 +98,10 @@ def plot_3(mesh, input_grid, title = '', fname = 'default', bitmap=False, ax10.set_ylabel('z') ax10.set_aspect('equal') - ax10.contour(xx, zz, np.sqrt((xx-xx.max()/2.0)**2.0 + (zz-zz.max()/2.0)**2.0), [mesh.minfo.contents["AC_accretion_range"]]) + if mesh.minfo.contents["AC_accretion_range"] > 0.0: + ax10.contour(xx, zz, np.sqrt((xx-xx.max()/2.0)**2.0 + (zz-zz.max()/2.0)**2.0), [mesh.minfo.contents["AC_accretion_range"]]) - xx, yy = np.meshgrid(mesh.xx, mesh.yy, indexing='ij') + xx, yy = np.meshgrid(mesh_xx_tmp, mesh_yy_tmp, indexing='ij') if contourplot: ax11.contourf(xx, yy, xy_slice, norm=plotnorm, cmap=colormap, nlev=10) else: @@ -98,7 +110,8 @@ def plot_3(mesh, input_grid, title = '', fname = 'default', bitmap=False, ax11.set_ylabel('y') ax11.set_aspect('equal') - ax11.contour(xx, yy, np.sqrt((xx-xx.max()/2.0)**2.0 + (yy-yy.max()/2.0)**2.0), [mesh.minfo.contents["AC_accretion_range"]]) + if mesh.minfo.contents["AC_accretion_range"] > 0.0: + ax11.contour(xx, yy, np.sqrt((xx-xx.max()/2.0)**2.0 + (yy-yy.max()/2.0)**2.0), [mesh.minfo.contents["AC_accretion_range"]]) if bfieldlines: ax00.streamplot(mesh.yy, mesh.zz, np.mean(mesh.bb[1], axis=0), np.mean(mesh.bb[2], axis=0))