Created
December 24, 2019 14:23
-
-
Save TomNicholas/c8a0d4d7269df44bf36fdb7ca8b45c02 to your computer and use it in GitHub Desktop.
Fabio's implementation of 3D plotting with mayavi
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| from netCDF4 import Dataset | |
| import numpy as np | |
| import scipy.interpolate | |
| import scipy.io | |
| import matplotlib.pyplot as plt | |
| import math | |
| def extract_region(var, regions, ir, yguards=True): | |
| """ | |
| Extract the data from an N-dimensional (with N=2 or N=4) variable corresponding to a given region | |
| Inputs : | |
| - the N-dimensional variable | |
| - the list of all possible regions | |
| - the index of the region of interest | |
| - if it has to return y guard cells | |
| Output : | |
| - an N-dimensional array | |
| """ | |
| ixi = regions[ir].ix[0] | |
| iyi = regions[ir].iy[0] | |
| nx = regions[ir].ix.shape[0] | |
| ny = regions[ir].iy.shape[0] | |
| if(var.ndim == 2): | |
| # 2D (x-y) array | |
| if(yguards): | |
| vartmp = np.full((nx,ny+2),np.nan) | |
| vartmp[:,1:-1] = var[ixi:ixi+nx,iyi:iyi+ny] | |
| if(regions[ir].yl == None): | |
| # If no yl data: extrapolate | |
| vartmp[:,0] = 2.*vartmp[:,1] - vartmp[:,2] | |
| else: | |
| vartmp[:,0] = var[ixi:ixi+nx,regions[regions[ir].yl].iy[-1]] | |
| if(regions[ir].yu == None): | |
| # If no yu data: extrapolate | |
| vartmp[:,-1] = 2.*vartmp[:,-2] - vartmp[:,-3] | |
| else: | |
| vartmp[:,-1] = var[ixi:ixi+nx,regions[regions[ir].yu].iy[0]] | |
| return vartmp | |
| else: | |
| return var[ixi:ixi+nx,iyi:iyi+ny] | |
| elif(var.ndim == 4): | |
| # 4D (t-x-y-z) array | |
| if(yguards): | |
| nt = var.shape[0] | |
| nz = var.shape[3] | |
| vartmp = np.full((nt,nx,ny+2,nz),np.nan) | |
| vartmp[:,:,1:-1,:] = var[:,ixi:ixi+nx,iyi:iyi+ny,:] | |
| if(regions[ir].yl == None): | |
| # If no yl data: extrapolate | |
| vartmp[:,:,0,:] = 2.*vartmp[:,:,1,:] - vartmp[:,:,2,:] | |
| else: | |
| ixi = regions[ir].ix[0] | |
| vartmp[:,:,0,:] = var[:,ixi:ixi+nx,regions[regions[ir].yl].iy[-1],:] | |
| if(regions[ir].yu == None): | |
| # If no yl data: extrapolate | |
| vartmp[:,:,-1,:] = 2.*vartmp[:,:,-2,:] - vartmp[:,:,-3,:] | |
| else: | |
| ixi = regions[ir].ix[0] | |
| vartmp[:,:,-1,:] = var[:,ixi:ixi+nx,regions[regions[ir].yu].iy[0],:] | |
| return vartmp | |
| else: | |
| return var[:,ixi:ixi+nx,iyi:iyi+ny,:] | |
| def interpy(var, regions, ir, N): | |
| """ | |
| Return var interpolated along y on a grid N times more refined in region region | |
| Inputs : | |
| - the N-dimensional variable to be interpolated (for 4D variables these must be field aligned!) | |
| - the list of all regions | |
| - the index of the region of interest | |
| - the refining factor N | |
| Output : | |
| - the interpolated N-dimensional variable | |
| """ | |
| nx = regions[ir].ix.shape[0] | |
| ny = regions[ir].iy.shape[0] | |
| nyi = ny*N | |
| y = (np.arange(ny+2)-0.5)/ny | |
| yi = (np.arange(nyi)+0.5)/nyi | |
| vartmp = extract_region(var, regions, ir) | |
| if(var.ndim == 2): | |
| f = scipy.interpolate.interp1d(y, vartmp, axis=1) | |
| elif(var.ndim == 4): | |
| f = scipy.interpolate.interp1d(y, vartmp, axis=2) | |
| return f(yi) | |
| def shiftZ(var, zperiod, toaligned, zShift): | |
| """ | |
| Shift in z to aligned/orthogonal coordinates | |
| """ | |
| if(toaligned): | |
| sign = 1 | |
| else: | |
| sign = -1 | |
| nz = var.shape[3] | |
| var_fft = np.fft.fft(var,axis=3) | |
| kz = np.arange(0,nz*zperiod,zperiod) | |
| var_aligned_fft = var_fft*np.exp(sign*1j*np.transpose(np.tile(zShift,(nz,1,1)),(1,2,0))*kz) | |
| if(nz % 2 == 0): | |
| nfft = int(nz/2) | |
| var_aligned_fft[:,:,:,nfft] = var_aligned_fft[:,:,:,nfft].real | |
| var_aligned_fft[:,:,:,nfft+1:] = np.conj(var_aligned_fft[:,:,:,nfft-1:0:-1]) | |
| else: | |
| nfft = int((nz-1)/2) | |
| var_aligned_fft[:,:,:,nfft+1:] = np.conj(var_aligned_fft[:,:,:,nfft:0:-1]) | |
| return np.fft.ifft(var_aligned_fft).real | |
| class region: | |
| """ | |
| x-y subregion of the poloidal plane | |
| ix and iy are the indexes defining the region (must be a sub domain of 0:nx-1 and 0:ny-1, respectively) | |
| xl and xr point to the regions on the left/right of the subregion | |
| yl and yu point to the regions below/above the subregion | |
| """ | |
| def __init__(self, ix=[], iy=[], xl=None, xr=None, yl=None, yu=None): | |
| self.ix = ix | |
| self.iy = iy | |
| self.xl = xl | |
| self.xr = xr | |
| self.yl = yl | |
| self.yu = yu | |
| class field3D(object): | |
| """ | |
| 3D field containing: | |
| - the 4D (t,x,y,z) variable | |
| - the Rxy and Zxy grid | |
| - the zShift | |
| - the different regions in the poloidal plane | |
| - zperiod | |
| """ | |
| def __init__(self,filename=None, varname=None, yguards=True, isaligned=False, gridname=None, zperiod=None): | |
| if filename is not None and varname is not None and gridname is not None: | |
| self.filename = filename | |
| self.varname = varname | |
| self.isaligned = isaligned | |
| # Load the variable | |
| if(filename[-2:] == "nc"): | |
| dataset = Dataset(filename) | |
| self.var = np.array(dataset.variables[varname]) | |
| self.t = np.array(dataset.variables["t_array"]) | |
| elif(filename[-3:] == "mat"): | |
| #data = scipy.io.loadmat(filename, variable_names={'t',varname}) | |
| data = scipy.io.loadmat(filename, variable_names=varname) | |
| self.var = data[varname] | |
| #self.t = data['t'] | |
| if(zperiod == None): | |
| raise NameError("Please specify zperiod") | |
| # Load the grid | |
| self.griddata = Dataset(gridname) | |
| self.gridname = gridname | |
| self.Rxy = np.array(self.griddata.variables['Rxy']) | |
| self.Zxy = np.array(self.griddata.variables['Zxy']) | |
| self.zShift = np.array(self.griddata.variables['zShift']) | |
| self.ShiftAngle = np.array(self.griddata.variables['ShiftAngle']) | |
| self.zperiod = zperiod | |
| # Interpolated variable | |
| self.Rxy_i = None | |
| self.Zxy_i = None | |
| self.zShift_i = None | |
| self.var_i = None | |
| self.isaligned_i = None | |
| self.regions_i = None | |
| self.t_i = None | |
| ixseps1 = self.griddata.variables['ixseps1'][0] | |
| ixseps2 = self.griddata.variables['ixseps2'][0] | |
| jyseps1_1 = self.griddata.variables['jyseps1_1'][0] | |
| jyseps2_1 = self.griddata.variables['jyseps2_1'][0] | |
| jyseps1_2 = self.griddata.variables['jyseps1_2'][0] | |
| jyseps2_2 = self.griddata.variables['jyseps2_2'][0] | |
| ny_inner = self.griddata.variables['ny_inner'][0] | |
| nx = self.griddata.variables['nx'][0] | |
| ny = self.griddata.variables['ny'][0] | |
| if(self.var.shape[2] == ny and self.Rxy.shape[1] == ny): | |
| print("The variable and the grid have the same number of y points.") | |
| elif(self.var.shape[2] == ny+4 and self.Rxy.shape[1] == ny): | |
| print("The variable has y guard cells, the grid not, removing guard cells from variables.") | |
| self.var = self.var[:,:,2:-2,:] | |
| elif(self.var.shape[2] == ny+4 and self.Rxy.shape[1] == ny+8): | |
| print("The variable and the grid have y guard cells and the configuration has 4 targets, removing guard cells.") | |
| self.var = self.var[:,:,2:-2,:] | |
| self.Rxy = self.Rxy[:,np.r_[2:ny_inner+2,ny_inner+6:ny+6]] | |
| self.Zxy = self.Zxy[:,np.r_[2:ny_inner+2,ny_inner+6:ny+6]] | |
| self.zShift = self.zShift[:,np.r_[2:ny_inner+2,ny_inner+6:ny+6]] | |
| elif(self.var.shape[2] == ny+4 and self.Rxy.shape[1] == ny+4): | |
| print("The variable and the grid have y guard cells and the configuration has 2 targets, removing guard cells.") | |
| self.var = self.var[:,:,2:-2,:] | |
| self.Rxy = self.Rxy[:,2:-2] | |
| self.Zxy = self.Zxy[:,2:-2] | |
| self.zShift = self.zShift[:,2:-2] | |
| else: | |
| raise NameError("Not implemented for this number of y guard cells.") | |
| if(nx != self.var.shape[1] or ny != self.var.shape[2]): | |
| raise NameError("Variable size not consistent with grid size.") | |
| regions = [] | |
| # Create the regions depending on the topology considered. | |
| # Possibilities : | |
| # - Core only | |
| # - Target to target. | |
| # - s-alpha | |
| # - SN | |
| # - X-point | |
| # - Connected DN | |
| # - Disconnected DN | |
| if(jyseps2_1 == jyseps1_2): | |
| if(ixseps1 != ixseps2): | |
| raise NameError("Unknown topology.") | |
| elif(ixseps1 > nx-1): | |
| if(jyseps1_1 < 0 and jyseps2_2 >= ny-1): | |
| print("Core only topology") | |
| regions.append(region(ix=np.arange(nx), iy=np.arange(ny), yl=0, yu=0)) | |
| else: | |
| raise NameError("Unknown topology.") | |
| elif(ixseps1 <= 0): | |
| print("Target to target topology") | |
| regions.append(region(ix=np.arange(nx), iy=np.arange(ny))) | |
| elif(jyseps1_1 < 0 and jyseps2_2 >= ny-1): | |
| print("s-alpha topology") | |
| # Core region (0) | |
| regions.append(region(ix=np.arange(ixseps1), iy=np.arange(ny), xr=1, yl=0, yu=0)) | |
| # SOL region (1) | |
| regions.append(region(ix=np.arange(nx-ixseps1)+ixseps1, iy=np.arange(ny), xl=0)) | |
| elif(jyseps1_1 >= 0 and jyseps2_2 < ny-1): | |
| print("Single null topology") | |
| # PFR inner (0) | |
| regions.append(region(ix=np.arange(ixseps1), iy=np.arange(jyseps1_1+1), xr=4, yu=3)) | |
| # Core inner (1) | |
| regions.append(region(ix=np.arange(ixseps1), iy=np.arange(jyseps2_1-jyseps1_1)+jyseps1_1+1, xr=5, yl=2, yu=2)) | |
| # Core outer (2) | |
| regions.append(region(ix=np.arange(ixseps1), iy=np.arange(jyseps2_2-jyseps1_2)+jyseps1_2+1, xr=6, yl=1, yu=1)) | |
| # PFR outer (3) | |
| regions.append(region(ix=np.arange(ixseps1), iy=np.arange(ny-jyseps2_2-1)+jyseps2_2+1, xr=7, yl=0)) | |
| # SOL (4-7, clockwise) | |
| regions.append(region(ix=np.arange(nx-ixseps1)+ixseps1, iy=np.arange(jyseps1_1+1), xl=0, yu=5)) | |
| regions.append(region(ix=np.arange(nx-ixseps1)+ixseps1, iy=np.arange(jyseps2_1-jyseps1_1)+jyseps1_1+1, xl=1, yl=4, yu=6)) | |
| regions.append(region(ix=np.arange(nx-ixseps1)+ixseps1, iy=np.arange(jyseps2_2-jyseps1_2)+jyseps1_2+1, xl=2, yl=5, yu=7)) | |
| regions.append(region(ix=np.arange(nx-ixseps1)+ixseps1, iy= np.arange(ny-jyseps2_2-1)+jyseps2_2+1, xl=3, yl=6)) | |
| else: | |
| raise NameError("Unknown topology.") | |
| else: | |
| if(ixseps1 == ixseps2): | |
| if(jyseps1_1 == jyseps2_1 and jyseps1_2 == jyseps2_2): | |
| print("X-point topology") | |
| # Lower PFR, inner (0) | |
| regions.append(region(ix=np.arange(ixseps1), iy=np.arange(jyseps1_1+1), xr=4, yu=3)) | |
| # Upper PFR, inner (1) | |
| regions.append(region(ix=np.arange(ixseps1), iy=np.arange(ny_inner-jyseps2_1-1)+jyseps2_1+1, xr=5, yl=2)) | |
| # Upper PFR, outer (2) | |
| regions.append(region(ix=np.arange(ixseps1), iy=np.arange(jyseps1_2-ny_inner+1)+ny_inner, xr=6, yu=1)) | |
| # Lower PFR, outer (3) | |
| regions.append(region(ix=np.arange(ixseps1), iy=np.arange(ny-jyseps1_2-1)+jyseps1_2+1, xr=7, yl=0)) | |
| # Inner SOL (4-5) | |
| regions.append(region(ix=np.arange(nx-ixseps1)+ixseps1, iy=np.arange(jyseps1_1+1), xl=0, yu=5)) | |
| regions.append(region(ix=np.arange(nx-ixseps1)+ixseps1, iy=np.arange(ny_inner-jyseps2_1-1)+jyseps2_1+1, xl=1, yl=4)) | |
| # Outer SOL (6-7) | |
| regions.append(region(ix=np.arange(nx-ixseps1)+ixseps1, iy=np.arange(jyseps1_2-ny_inner+1)+ny_inner, xl=2, yu=7)) | |
| regions.append(region(ix=np.arange(nx-ixseps1)+ixseps1, iy=np.arange(ny-jyseps1_2-1)+jyseps1_2+1, xl=3, yl=6)) | |
| else: | |
| print("Connected double null topology") | |
| # Lower PFR, inner (0) | |
| regions.append(region(ix=np.arange(ixseps1), iy=np.arange(jyseps1_1+1), xr=6, yu=5)) | |
| # Core inner (1) | |
| regions.append(region(ix=np.arange(ixseps1), iy=np.arange(jyseps2_1-jyseps1_1)+jyseps1_1+1, xr=7, yl=4, yu=4)) | |
| # Upper PFR, inner (2) | |
| regions.append(region(ix=np.arange(ixseps1), iy=np.arange(ny_inner-jyseps2_1-1)+jyseps2_1+1, xr=8, yl=3)) | |
| # Upper PFR, outer (3) | |
| regions.append(region(ix=np.arange(ixseps1), iy=np.arange(jyseps1_2-ny_inner+1)+ny_inner, xr=9, yu=2)) | |
| # Core outer (4) | |
| regions.append(region(ix=np.arange(ixseps1), iy=np.arange(jyseps2_2-jyseps1_2)+jyseps1_2+1, xr=10, yl=1, yu=1)) | |
| # Lower PFR, outer (5) | |
| regions.append(region(ix=np.arange(ixseps1), iy=np.arange(ny-jyseps2_2-1)+jyseps2_2+1, xr=11, yl=0)) | |
| # Inner SOL (6-8) | |
| regions.append(region(ix=np.arange(nx-ixseps1)+ixseps1, iy=np.arange(jyseps1_1+1), xl=0, yu=7)) | |
| regions.append(region(ix=np.arange(nx-ixseps1)+ixseps1, iy=np.arange(jyseps2_1-jyseps1_1)+jyseps1_1+1, xl=1, yl=6, yu=8)) | |
| regions.append(region(ix=np.arange(nx-ixseps1)+ixseps1, iy=np.arange(ny_inner-jyseps2_1-1)+jyseps2_1+1, xl=2, yl=7)) | |
| # Outer SOL (9-11) | |
| regions.append(region(ix=np.arange(nx-ixseps1)+ixseps1, iy=np.arange(jyseps1_2-ny_inner+1)+ny_inner, xl=3, yu=10)) | |
| regions.append(region(ix=np.arange(nx-ixseps1)+ixseps1, iy=np.arange(jyseps2_2-jyseps1_2)+jyseps1_2+1, xl=4, yl=9, yu=11)) | |
| regions.append(region(ix=np.arange(nx-ixseps1)+ixseps1, iy=np.arange(ny-jyseps2_2-1)+jyseps2_2+1, xl=5, yl=10)) | |
| elif(ixseps1 > 0 and ixseps2 <= nx-1 and ixseps1 < ixseps2): | |
| print("Disconnected double null topology") | |
| # Lower PFR, inner (0) | |
| regions.append(region(ix=np.arange(ixseps1), iy=np.arange(jyseps1_1+1), xr=6, yu=5)) | |
| # Core inner (1) | |
| regions.append(region(ix=np.arange(ixseps1), iy=np.arange(jyseps2_1-jyseps1_1)+jyseps1_1+1, xr=7, yl=4, yu=4)) | |
| # Upper PFR, inner (2) | |
| regions.append(region(ix=np.arange(ixseps2), iy=np.arange(ny_inner-jyseps2_1-1)+jyseps2_1+1, xr=12, yl=3)) | |
| # Upper PFR, outer (3) | |
| regions.append(region(ix=np.arange(ixseps2), iy=np.arange(jyseps1_2-ny_inner+1)+ny_inner, xr=13, yu=2)) | |
| # Core outer (4) | |
| regions.append(region(ix=np.arange(ixseps1), iy=np.arange(jyseps2_2-jyseps1_2)+jyseps1_2+1, xr=8, yl=1, yu=1)) | |
| # Lower PFR, outer (5) | |
| regions.append(region(ix=np.arange(ixseps1), iy=np.arange(ny-jyseps2_2-1)+jyseps2_2+1, xr=9, yl=0)) | |
| # SOL between separatrices, inner (6-7) | |
| regions.append(region(ix=np.arange(ixseps2-ixseps1)+ixseps1, iy=np.arange(jyseps1_1+1), xl=0, xr=10, yu=7)) | |
| regions.append(region(ix=np.arange(ixseps2-ixseps1)+ixseps1, iy=np.arange(jyseps2_1-jyseps1_1)+jyseps1_1+1, xl=1, xr=11, yl=6, yu=8)) | |
| # SOL between separatrices, outer (8-9) | |
| regions.append(region(ix=np.arange(ixseps2-ixseps1)+ixseps1, iy=np.arange(jyseps2_2-jyseps1_2)+jyseps1_2+1, xl=4, xr=14, yl=7, yu=9)) | |
| regions.append(region(ix=np.arange(ixseps2-ixseps1)+ixseps1, iy=np.arange(ny-jyseps2_2-1)+jyseps2_2+1, xl=5, xr=15, yl=8)) | |
| # Inner SOL (10-12) | |
| regions.append(region(ix=np.arange(nx-ixseps2)+ixseps2, iy=np.arange(jyseps1_1+1), xl=6, yu=11)) | |
| regions.append(region(ix=np.arange(nx-ixseps2)+ixseps2, iy=np.arange(jyseps2_1-jyseps1_1)+jyseps1_1+1, xl=7, yl=10, yu=12)) | |
| regions.append(region(ix=np.arange(nx-ixseps2)+ixseps2, iy=np.arange(ny_inner-jyseps2_1-1)+jyseps2_1+1, xl=2, yl=11)) | |
| # Outer SOL (13-15) | |
| regions.append(region(ix=np.arange(nx-ixseps2)+ixseps2, iy=np.arange(jyseps1_2-ny_inner+1)+ny_inner, xl=3, yu=14)) | |
| regions.append(region(ix=np.arange(nx-ixseps2)+ixseps2, iy=np.arange(jyseps2_2-jyseps1_2)+jyseps1_2+1, xl=8, yl=13, yu=15)) | |
| regions.append(region(ix=np.arange(nx-ixseps2)+ixseps2, iy=np.arange(ny-jyseps2_2-1)+jyseps2_2+1, xl=9, yl=14)) | |
| else: | |
| raise NameError("Unknown topology.") | |
| self.regions = regions | |
| else: | |
| raise NameError("Need a filename, a varname, and a gridfile to create a Field3D.") | |
| def interp_var(self,N,it=None,aligned=False,log_interp=False): | |
| """ | |
| 3D interpolation in y of field. | |
| Inputs : | |
| - The refinement factor N | |
| - the timestep (None for all times) | |
| - the interpolated variable is field aligned? | |
| - Performing the interpolation in logarithmic? | |
| """ | |
| if(it == None): | |
| nt = self.var.shape[0] | |
| #self.t_i = self.t | |
| else: | |
| nt = 1 | |
| #self.t_i = self.t[it] | |
| nx = self.var.shape[1] | |
| ny = N*self.var.shape[2] | |
| nz = self.var.shape[3] | |
| zperiod = self.zperiod | |
| self.Rxy_i = np.full((nx,ny),np.nan) | |
| self.Zxy_i = np.full((nx,ny),np.nan) | |
| self.zShift_i = np.full((nx,ny),np.nan) | |
| self.var_i = np.full((nt,nx,ny,nz),np.nan) | |
| self.isaligned_i = aligned | |
| #Extract the relevant data | |
| if(it == None): | |
| var_tmp = self.var | |
| else: | |
| var_tmp = np.tile(self.var[it,:,:,:],(1,1,1,1)) | |
| if(log_interp): | |
| var_tmp = np.log(var_tmp+1e-12) | |
| if(not self.isaligned): | |
| #Go to aligned coordinates | |
| print("Shifting in z to field aligned...") | |
| var_aligned = shiftZ(var_tmp, zperiod, True, self.zShift) | |
| else: | |
| var_aligned = var_tmp | |
| # Interpolate in y | |
| print("Interpolate in y...") | |
| regions = [] | |
| for ir in range(len(self.regions)): | |
| r = self.regions[ir] | |
| nx = r.ix.shape[0] | |
| nyi = r.iy.shape[0]*N | |
| regions.append(region(ix=r.ix, iy=np.arange(nyi)+N*r.iy[0], xl=r.xl, xr=r.xr, yl=r.yl, yu=r.yu)) | |
| ixi = r.ix[0] | |
| iyi = regions[-1].iy[0] | |
| # Interpolate Rxy, Zxy, and zShift | |
| self.Rxy_i[ixi:ixi+nx,iyi:iyi+nyi] = interpy(self.Rxy, self.regions, ir, N) | |
| self.Zxy_i[ixi:ixi+nx,iyi:iyi+nyi] = interpy(self.Zxy, self.regions, ir, N) | |
| if(r.yl == r.yu and r.yl != None): | |
| # Core region, need to correct zShift | |
| zShift_tmp = extract_region(self.zShift, self.regions, ir) | |
| # Correct the guard cells | |
| if(r.yl == ir): | |
| # Only one core region, r | |
| zShift_tmp[:,0] = zShift_tmp[:,-2] - self.ShiftAngle[ixi:ixi+nx] | |
| zShift_tmp[:,-1] = zShift_tmp[:,1] + self.ShiftAngle[ixi:ixi+nx] | |
| elif(r.yl > ir): | |
| # Two core regions, I'm the lower one | |
| zShift_tmp[:,0] = self.zShift[ixi:ixi+nx,self.regions[r.yl].iy[-1]] - self.ShiftAngle[ixi:ixi+nx] | |
| else: | |
| # Two core regions, I'm the upper one | |
| zShift_tmp[:,-1] = self.zShift[ixi:ixi+nx,self.regions[r.yu].iy[0]] + self.ShiftAngle[ixi:ixi+nx] | |
| # Interpolate in y | |
| y = (np.arange(r.iy.shape[0]+2)-0.5)/r.iy.shape[0] | |
| yi = (np.arange(nyi)+0.5)/nyi | |
| f = scipy.interpolate.interp1d(y, zShift_tmp, axis=1) | |
| self.zShift_i[ixi:ixi+nx,iyi:iyi+nyi] = f(yi) | |
| else: | |
| self.zShift_i[ixi:ixi+nx,iyi:iyi+nyi] = interpy(self.zShift, self.regions, ir, N) | |
| # Interpolate in y var_aligned | |
| self.var_i[:,ixi:ixi+nx,iyi:iyi+nyi,:] = interpy(var_aligned, self.regions, ir, N) | |
| self.regions_i = regions | |
| if(not aligned): | |
| print("Shifting in z to orthogonal...") | |
| self.var_i = shiftZ(self.var_i, zperiod, False, self.zShift_i) | |
| if(log_interp): | |
| self.var_i = np.exp(self.var_i) | |
| def plot(self,method,interpolated=False,it=None,ix=None,iy=None,iz=None,remove_toroidal=0): | |
| """ | |
| Plot methods: | |
| - method = 'xy',... | |
| - plot for interpolated quantity? | |
| - it->which snapshot | |
| - ix = at which radial location? | |
| - iy = at which poloidal location? | |
| - iz = at which toroidal location? | |
| - remove_toroidal = 0 -> no modifications | |
| remove_toroidal = 1 -> A - A_DC (z-fluctuations) | |
| remove_toroidal = 2 -> (A-A_DC)/A_DC (relative z.fluctuations) | |
| """ | |
| if(interpolated): | |
| if(not isinstance(self.Rxy_i,np.ndarray)): | |
| raise NameError("Execute field3D.interp_var before trying to plot the interpolated variable.") | |
| var = self.var_i | |
| Rxy = self.Rxy_i | |
| Zxy = self.Zxy_i | |
| regions = self.regions_i | |
| isaligned = self.isaligned_i | |
| else: | |
| var = self.var | |
| Rxy = self.Rxy | |
| Zxy = self.Zxy | |
| regions = self.regions | |
| isaligned = self.isaligned | |
| if(var.shape[0] > 1 and it == None): | |
| raise NameError("Plot error: variable has several timesteps, but it not defined") | |
| elif(var.shape[0] > 1): | |
| var = var[it,:,:,:] | |
| else: | |
| var = var[0,:,:,:] | |
| nz = var.shape[2] | |
| if(remove_toroidal == 1): | |
| var = var - np.transpose(np.tile(np.mean(var,axis=2),(nz,1,1)),(1,2,0)) | |
| elif(remove_toroidal == 2): | |
| var_DC = np.transpose(np.tile(np.mean(var,axis=2),(nz,1,1)),(1,2,0)) + 1e-12 | |
| var = (var - var_DC)/var_DC | |
| if(method == 'xy'): | |
| if(isaligned): | |
| # For the moment 'xy' only for plotting a poloidal plane | |
| raise NameError("Plot xy error: variable is not orthogonal") | |
| fig = plt.figure() | |
| ax_list = fig.axes | |
| if(iz == None): | |
| raise NameError("Plot method xy, iz not defined.") | |
| else: | |
| var = var[:,:,iz] | |
| m = var[2:-2,:].min() | |
| M = var[2:-2,:].max() | |
| for ir in range(len(regions)): | |
| r = regions[ir] | |
| nx = r.ix.shape[0] | |
| ny = r.iy.shape[0] | |
| ixi = r.ix[0] | |
| iyi = r.iy[0] | |
| # Remove x guard cells | |
| if(r.xl == None): | |
| ixi = 2 | |
| nx = nx - 2 | |
| if(r.xr == None): | |
| nx = nx-2 | |
| var_tmp = var[ixi:ixi+nx,iyi:iyi+ny] | |
| Rxy_tmp = Rxy[ixi:ixi+nx,iyi:iyi+ny] | |
| Zxy_tmp = Zxy[ixi:ixi+nx,iyi:iyi+ny] | |
| # Add guard cells to remove gaps in plotting between different regions | |
| if(r.xr != None): | |
| var_tmp = np.vstack([var_tmp, var[regions[r.xr].ix[0],iyi:iyi+ny]]) | |
| Rxy_tmp = np.vstack([Rxy_tmp, Rxy[regions[r.xr].ix[0],iyi:iyi+ny] ]) | |
| Zxy_tmp = np.vstack([Zxy_tmp, Zxy[regions[r.xr].ix[0],iyi:iyi+ny] ]) | |
| nx = nx+1 | |
| if(r.yu != None): | |
| var_tmp = np.vstack([var_tmp.T, var[ixi:ixi+nx,regions[r.yu].iy[0]] ]).T | |
| Rxy_tmp = np.vstack([Rxy_tmp.T, Rxy[ixi:ixi+nx,regions[r.yu].iy[0]] ]).T | |
| Zxy_tmp = np.vstack([Zxy_tmp.T, Zxy[ixi:ixi+nx,regions[r.yu].iy[0]] ]).T | |
| plt.pcolormesh(Rxy_tmp,Zxy_tmp,var_tmp, vmin=m, vmax=M, shading='gouraud',cmap='jet') | |
| plt.colorbar() | |
| plt.gca().set_aspect('equal', adjustable='box') | |
| elif(method == 'xyz'): | |
| if(isaligned): | |
| raise NameError("Plot xyz error: variable is not orthogonal") | |
| try: | |
| from mayavi import mlab | |
| havemayavi = True | |
| except: | |
| from mpl_toolkits import mplot3d | |
| from matplotlib import cm | |
| ax = plt.axes(projection='3d') | |
| ax.view_init(elev=0., azim=-150) | |
| havemayavi = false | |
| print("Don't have mayavi installed, using mplot3d (bad performances...).") | |
| dz = 2*math.pi/(self.zperiod*nz) | |
| z = np.arange(nz)*dz | |
| nx = var.shape[0] | |
| ny = var.shape[1] | |
| # Create the coordinates | |
| print("Computing coordinates for 3D plot...") | |
| X = np.transpose(np.tile(Rxy,(nz,1,1)),(1,2,0))*np.tile(np.cos(z),(nx,ny,1)) | |
| Y = np.transpose(np.tile(Rxy,(nz,1,1)),(1,2,0))*np.tile(np.sin(z),(nx,ny,1)) | |
| Z = np.transpose(np.tile(Zxy,(nz,1,1)),(1,2,0)) | |
| m = var[2:-2,:,:].min() | |
| M = var[2:-2,:,:].max() | |
| print("Plotting...") | |
| for ir in range(len(regions)): | |
| r = regions[ir] | |
| nx = r.ix.shape[0] | |
| ny = r.iy.shape[0] | |
| ixi = r.ix[0] | |
| iyi = r.iy[0] | |
| # Remove x guard cells | |
| if(r.xl == None): | |
| ixi = 2 | |
| nx = nx - 2 | |
| if(r.xr == None): | |
| nx = nx-2 | |
| var_tmp = var[ixi:ixi+nx,iyi:iyi+ny,:] | |
| X_tmp = X[ixi:ixi+nx,iyi:iyi+ny,:] | |
| Y_tmp = Y[ixi:ixi+nx,iyi:iyi+ny,:] | |
| Z_tmp = Z[ixi:ixi+nx,iyi:iyi+ny,:] | |
| # Add guard cells to remove gaps in plotting between different regions | |
| if(r.xr != None): | |
| var_tmp = np.vstack([var_tmp, np.tile(var[regions[r.xr].ix[0],iyi:iyi+ny,:],(1,1,1))]) | |
| X_tmp = np.vstack([X_tmp, np.tile(X[regions[r.xr].ix[0],iyi:iyi+ny,:],(1,1,1))]) | |
| Y_tmp = np.vstack([Y_tmp, np.tile(Y[regions[r.xr].ix[0],iyi:iyi+ny,:],(1,1,1))]) | |
| Z_tmp = np.vstack([Z_tmp, np.tile(Z[regions[r.xr].ix[0],iyi:iyi+ny,:],(1,1,1))]) | |
| nx = nx+1 | |
| if(r.yu != None): | |
| var_tmp = np.vstack([np.transpose(var_tmp,(1,0,2)),np.tile(var[ixi:ixi+nx,regions[r.yu].iy[0],:],(1,1,1))]) | |
| X_tmp = np.vstack([np.transpose(X_tmp,(1,0,2)),np.tile(X[ixi:ixi+nx,regions[r.yu].iy[0],:],(1,1,1))]) | |
| Y_tmp = np.vstack([np.transpose(Y_tmp,(1,0,2)),np.tile(Y[ixi:ixi+nx,regions[r.yu].iy[0],:],(1,1,1))]) | |
| Z_tmp = np.vstack([np.transpose(Z_tmp,(1,0,2)),np.tile(Z[ixi:ixi+nx,regions[r.yu].iy[0],:],(1,1,1))]) | |
| var_tmp = np.transpose(var_tmp,(1,0,2)) | |
| X_tmp = np.transpose(X_tmp,(1,0,2)) | |
| Y_tmp = np.transpose(Y_tmp,(1,0,2)) | |
| Z_tmp = np.transpose(Z_tmp,(1,0,2)) | |
| if(not havemayavi): | |
| # Rescale the variable between 0 and 1 for color plotting | |
| var_tmp = (var_tmp-m)/(M-m) | |
| shade = False | |
| rstride = 1 | |
| cstride = 1 | |
| antialiased = False | |
| # Plot the inner and outer surfaces | |
| if(r.xl == None): | |
| surf = ax.plot_surface(X_tmp[0,:,:], Y_tmp[0,:,:], Z_tmp[0,:,:], \ | |
| rstride=rstride, cstride=cstride, facecolors=cm.jet(var_tmp[0,:,:]), \ | |
| linewidth=0, antialiased=antialiased, shade=shade) | |
| if(r.xr == None): | |
| surf = ax.plot_surface(X_tmp[-1,:,:], Y_tmp[-1,:,:], Z_tmp[-1,:,:], \ | |
| rstride=rstride, cstride=cstride, facecolors=cm.jet(var_tmp[-1,:,:]), \ | |
| linewidth=0, antialiased=antialiased, shade=shade) | |
| # Plot the target surfaces | |
| if(r.yl == None): | |
| surf = ax.plot_surface(X_tmp[:,0,:], Y_tmp[:,0,:], Z_tmp[:,0,:], \ | |
| rstride=rstride, cstride=cstride, facecolors=cm.jet(var_tmp[:,0,:]), \ | |
| linewidth=0, antialiased=antialiased, shade=shade) | |
| if(r.yu == None): | |
| surf = ax.plot_surface(X_tmp[:,-1,:], Y_tmp[:,-1,:], Z_tmp[:,-1,:], \ | |
| rstride=rstride, cstride=cstride, facecolors=cm.jet(var_tmp[:,-1,:]), \ | |
| linewidth=0, antialiased=antialiased, shade=shade) | |
| # Plot the first and last poloidal planes | |
| iz = 0 | |
| surf = ax.plot_surface(X_tmp[:,:,iz], Y_tmp[:,:,iz], Z_tmp[:,:,iz], \ | |
| rstride=rstride, cstride=cstride, facecolors=cm.jet(var_tmp[:,:,iz]), \ | |
| linewidth=0, antialiased=antialiased, shade=shade) | |
| iz = nz-1 | |
| surf = ax.plot_surface(X_tmp[:,:,iz], Y_tmp[:,:,iz], Z_tmp[:,:,iz], \ | |
| rstride=rstride, cstride=cstride, facecolors=cm.jet(var_tmp[:,:,iz]), \ | |
| linewidth=0, antialiased=antialiased, shade=shade) | |
| else: | |
| # Plot the inner and outer surfaces | |
| if(r.xl == None): | |
| mlab.mesh(X_tmp[0,:,:], Y_tmp[0,:,:], Z_tmp[0,:,:], scalars=var_tmp[0,:,:], vmin=m, vmax=M) | |
| if(r.xr == None): | |
| mlab.mesh(X_tmp[-1,:,:], Y_tmp[-1,:,:], Z_tmp[-1,:,:], scalars=var_tmp[-1,:,:], vmin=m, vmax=M) | |
| # Plot the target surfaces | |
| if(r.yl == None): | |
| mlab.mesh(X_tmp[:,0,:], Y_tmp[:,0,:], Z_tmp[:,0,:], scalars=var_tmp[:,0,:], vmin=m, vmax=M) | |
| if(r.yu == None): | |
| mlab.mesh(X_tmp[:,-1,:], Y_tmp[:,-1,:], Z_tmp[:,-1,:], scalars=var_tmp[:,-1,:], vmin=m, vmax=M) | |
| # Plot the first and last poloidal planes | |
| iz = 0 | |
| mlab.mesh(X_tmp[:,:,iz], Y_tmp[:,:,iz], Z_tmp[:,:,iz], scalars=var_tmp[:,:,iz], vmin=m, vmax=M) | |
| iz = nz-1 | |
| mlab.mesh(X_tmp[:,:,iz], Y_tmp[:,:,iz], Z_tmp[:,:,iz], scalars=var_tmp[:,:,iz], vmin=m, vmax=M) | |
| plt.show() | |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Produces plots like this:
Would be great to use to plot 3D s-alpha results.