def gyrophase_angles(bulk_velocity, B_unit, velocity_cell_data, velocity_coordinates, plasmaframe=True, cosine=False): ''' Calculates the gyrophase angle angle distribution for a given cell :param bulk_velocity: TODO :param B_unit: TODO :param velocity_coordinates: TODO :param cosine: True if returning the gyrophase angles as a cosine plot :param plasmaframe: True if the user wants to get the gyrophase angle distribution in the plasma frame, default True :returns: gyrophase angles and avgs [gyro_angles, avgs] .. code-block:: python # Example usage: vlsvReader = VlsvReader("fullf.0001.vlsv") result = gyrophase_angles_from_file( vlsvReader=vlsvReader, cellid=1924, cosine=True, plasmaframe=False ) # Plot the data import pylab as pl pl.hist(result[0].data, weights=result[1].data, bins=100, log=False) ''' # Get avgs data: avgs = velocity_cell_data.values() # Shift to plasma frame if plasmaframe == True: velocity_coordinates = velocity_coordinates - bulk_velocity # Get norms: v_norms = np.sum(np.abs(velocity_coordinates)**2, axis=-1)**(1. / 2) # Get the angles: v_rotated = rotateVectorToVector(velocity_coordinates, B_unit) v_rotated = np.asarray(v_rotated) if cosine == True: gyro_angles = np.cos(np.arctan2(v_rotated[:, 0], v_rotated[:, 1])) units = "" else: gyro_angles = np.arctan2(v_rotated[:, 0], v_rotated[:, 1]) / (2 * np.pi) * 360 units = "degree" # Return the gyrophase angles and avgs values: from output import output_1d return output_1d([gyro_angles, avgs], ["Gyrophase_angle", "avgs"], [units, ""])
def gyrophase_angles(bulk_velocity, B_unit, velocity_cell_data, velocity_coordinates, plasmaframe=True, cosine=False): ''' Calculates the gyrophase angle angle distribution for a given cell :param bulk_velocity: TODO :param B_unit: TODO :param velocity_coordinates: TODO :param cosine: True if returning the gyrophase angles as a cosine plot :param plasmaframe: True if the user wants to get the gyrophase angle distribution in the plasma frame, default True :returns: gyrophase angles and avgs [gyro_angles, avgs] .. code-block:: python # Example usage: vlsvReader = VlsvReader("fullf.0001.vlsv") result = gyrophase_angles_from_file( vlsvReader=vlsvReader, cellid=1924, cosine=True, plasmaframe=False ) # Plot the data import pylab as pl pl.hist(result[0].data, weights=result[1].data, bins=100, log=False) ''' # Get avgs data: avgs = velocity_cell_data.values() # Shift to plasma frame if plasmaframe == True: velocity_coordinates = velocity_coordinates - bulk_velocity # Get norms: v_norms = np.sum(np.abs(velocity_coordinates)**2,axis=-1)**(1./2) # Get the angles: v_rotated = rotateVectorToVector(velocity_coordinates, B_unit) v_rotated = np.asarray(v_rotated) if cosine == True: gyro_angles = np.cos(np.arctan2(v_rotated[:,0], v_rotated[:,1])) units = "" else: gyro_angles = np.arctan2(v_rotated[:,0], v_rotated[:,1]) / (2*np.pi) * 360 units = "degree" # Return the gyrophase angles and avgs values: from output import output_1d return output_1d([gyro_angles, avgs], ["Gyrophase_angle", "avgs"], [units, ""])
def vSpaceReducer(vlsvReader, cid, slicetype, normvect, pop="proton", center=None, setThreshold=None,normvectX=None): # check if velocity space exists in this cell if vlsvReader.check_variable('fSaved'): #restart files will not have this value if vlsvReader.read_variable('fSaved',cid) != 1.0: return (False,0,0,0) if vlsvReader.check_variable('vg_f_saved'): #restart files will not have this value if vlsvReader.read_variable('vg_f_saved',cid) != 1.0: return (False,0,0,0) # Assume velocity cells are cubes [vxsize, vysize, vzsize] = vlsvReader.get_velocity_mesh_size(pop=pop) vxsize = int(vxsize) vysize = int(vysize) vzsize = int(vzsize) # Account for 4x4x4 cells per block vxsize = 4*vxsize vysize = 4*vysize vzsize = 4*vzsize [vxmin, vymin, vzmin, vxmax, vymax, vzmax] = vlsvReader.get_velocity_mesh_extent(pop=pop) inputcellsize=(vxmax-vxmin)/vxsize print("Input velocity grid cell size "+str(inputcellsize)) velcells = vlsvReader.read_velocity_cells(cid, pop=pop) velcellslist = list(zip(*velcells.items())) # check that velocity space has cells if(len(velcellslist) <= 0): return (False,0,0,0) f = np.asarray(velcellslist[1]) V = vlsvReader.get_velocity_cell_coordinates(velcellslist[0], pop=pop) print("Found "+str(len(V))+" v-space cells") # center on highest f-value if center == "peak": peakindex = np.argmax(f) Vpeak = V[peakindex,:] V = V - Vpeak print(peakindex) print("Transforming to frame of peak f-value, travelling at speed "+str(Vpeak)) elif not center is None: if len(center)==3: # assumes it's a vector print("Transforming to frame travelling at speed "+str(center)) V = V - center else: print("Error in shape of center vector! Give in form (vx,vy,vz).") if setThreshold==None: # Drop all velocity cells which are below the sparsity threshold. Otherwise the plot will show buffer # cells as well. if vlsvReader.check_variable('MinValue') == True: # Sparsity threshold used to be saved as MinValue setThreshold = vlsvReader.read_variable('MinValue',cid) print("Found a vlsv file MinValue of "+str(setThreshold)) elif vlsvReader.check_variable(pop+"/EffectiveSparsityThreshold") == True: setThreshold = vlsvReader.read_variable(pop+"/EffectiveSparsityThreshold",cid) print("Found a vlsv file value "+pop+"/EffectiveSparsityThreshold"+" of "+str(setThreshold)) elif vlsvReader.check_variable(pop+"/vg_effectivesparsitythreshold") == True: setThreshold = vlsvReader.read_variable(pop+"/vg_effectivesparsitythreshold",cid) print("Found a vlsv file value "+pop+"/vg_effectivesparsitythreshold"+" of "+str(setThreshold)) else: print("Warning! Unable to find a MinValue or EffectiveSparsityThreshold value from the .vlsv file.") print("Using a default value of 1.e-16. Override with setThreshold=value.") setThreshold = 1.e-16 ii_f = np.where(f >= setThreshold) print("Dropping velocity cells under setThreshold value "+str(setThreshold)) if len(ii_f) < 1: return (False,0,0,0) f = f[ii_f] V = V[ii_f,:][0,:,:] # Geometric magic to widen the slice to assure that each cell has some velocity grid points inside it. samplebox=np.array([ [0.0,0.0,0.0], [0.0,0.0,1.0], [0.0,1.0,0.0], [0.0,1.0,1.0], [1.0,0.0,0.0], [1.0,0.0,1.0], [1.0,1.0,0.0], [1.0,1.0,1.0] ]) sbrot = rotateVectorToVector(samplebox,normvect) rotminx=np.amin(sbrot[:,0]) rotmaxx=np.amax(sbrot[:,0]) rotminy=np.amin(sbrot[:,1]) rotmaxy=np.amax(sbrot[:,1]) rotminz=np.amin(sbrot[:,2]) rotmaxz=np.amax(sbrot[:,2]) gridratio = np.amax([ rotmaxx-rotminx, rotmaxy-rotminy, rotmaxz-rotminz ]) # if gridratio > 1.0: # adds a 5% margin to slice thickness gridratio = 1.05*gridratio slicethick=inputcellsize*gridratio if slicetype=="xy": VX = V[:,0] VY = V[:,1] Voutofslice = V[:,2] elif slicetype=="yz": VX = V[:,1] VY = V[:,2] Voutofslice = V[:,0] elif slicetype=="xz": VX = V[:,0] VY = V[:,2] Voutofslice = -V[:,1] elif slicetype=="vecperp": N = np.array(normvect)/np.sqrt(normvect[0]**2 + normvect[1]**2 + normvect[2]**2) Vrot = rotateVectorToVector(V,N) # aligns the Z axis of V with normvect VX = Vrot[:,0] VY = Vrot[:,1] Voutofslice = Vrot[:,2] elif slicetype=="Bperp" or slicetype=="Bpara" or slicetype=="Bpara1": # Find velocity components in rotated frame where B is aligned with Z and BcrossV is aligned with X N = np.array(normvect)/np.sqrt(normvect[0]**2 + normvect[1]**2 + normvect[2]**2) NX = np.array(normvectX)/np.sqrt(normvectX[0]**2 + normvectX[1]**2 + normvectX[2]**2) Vrot = rotateVectorToVector(V,N) # transforms V to frame where z is aligned with N=B NXrot = rotateVectorToVector(NX,N) # transforms NX=BcrossV to frame where z is aligned with N=B (hence NXrot in XY plane) Vrot2 = rotateVectorToVector_X(Vrot,NXrot) # transforms Vrot to frame where x is aligned with NXrot (hence preserves z) # Choose the correct components for this plot if slicetype=="Bperp": VX = Vrot2[:,0] # the X axis of the slice is BcrossV=perp1 VY = Vrot2[:,1] # the Y axis of the slice is Bcross(BcrossV)=perp2 Voutofslice = Vrot2[:,2] # the Z axis of the slice is B elif slicetype=="Bpara": VX = Vrot2[:,2] # the X axis of the slice is B VY = Vrot2[:,1] # the Y axis of the slice is Bcross(BcrossV)=perp2 Voutofslice = Vrot2[:,0] # the Z axis of the slice is -BcrossV=perp1 elif slicetype=="Bpara1": VX = Vrot2[:,2] # the X axis of the slice is B VY = Vrot2[:,0] # the Y axis of the slice is BcrossV=perp1 Voutofslice = Vrot2[:,1] # the Z axis of the slice is Bcross(BcrossV)=perp2 # Calculations for verification of rotation: testvectors = np.array([N,NX,np.cross(N,NX)]) # verifies B, BcrossV, and Bcross(BcrossV) testrot = rotateVectorToVector(testvectors,N) # transforms testvectors to frame where z is aligned with N=B testrot2 = rotateVectorToVector_X(testrot,NXrot) # transforms testrot to frame where x is aligned with NXrot (hence preserves z) if abs(1.0-np.linalg.norm(NXrot))>1.e-3: print("Error in rotation: NXrot not a unit vector") if abs(NXrot[2]) > 1.e-3: print("Error in rotation: NXrot not in x-y-plane") for count,testvect in enumerate(testrot2): if abs(1.0-np.linalg.norm(testvect))>1.e-3: print("Error in rotation: testvector ",count,testvect," not a unit vector") if abs(1.0-np.amax(testvect))>1.e-3: print("Error in rotation: testvector ",count,testvect," largest component is not unity") else: print("Error finding rotation of v-space!") return (False,0,0,0) # create three 1-dimensional profiles across VDF (bins1,bins2,bins3,axis1,axis2,axis3) = doProfiles(f,VX,VY,Voutofslice,slicethick) return (True,bins1,bins2,bins3,axis1,axis2,axis3)
def plot_vdf( filename=None, vlsvobj=None, filedir=None, step=None, cellids=None, pop="proton", coordinates=None, coordre=None, outputdir=None, draw=None, unit=None, title=None, cbtitle=None, colormap=None, box=None, cbar=None, run=None, wmark=None, thick=1.0, fmin=None, fmax=None, slicethick=None, cellsize=None, xy=None, xz=None, yz=None, normal=None, bpara=None, bperp=None, coordswap=None, cbulk=None, center=None, wflux=None, keepfmin=None, legend=None, noborder=None, scale=1.0, biglabel=None, biglabloc=None, noxlabels=None, noylabels=None, ): ''' Plots a coloured plot with axes and a colour bar. :kword filename: path to .vlsv file to use for input. Assumes a bulk file. :kword vlsvobj: Optionally provide a python vlsvfile object instead :kword filedir: Optionally provide directory where files are located and use step for bulk file name :kword step: output step index, used for constructing output (and possibly input) filename :kword outputdir: path to directory where output files are created (default: $HOME/Plots/) If directory does not exist, it will be created. If the string does not end in a forward slash, the final parti will be used as a perfix for the files. :kword cellids: list of cell IDs to plot VDF for :kword coordinates: list of 3-element spatial coordinates to plot VDF for (given in metres) :kword coordre: list of 3-element spatial coordinates to plot VDF for (given in Earth radii) :kword pop: Population to plot, default proton :kword colormap: colour scale for plot, use e.g. hot_desaturated, jet, viridis, plasma, inferno, magma, parula, nipy_spectral, RdBu, bwr :kword run: run identifier, used for constructing output filename :kword title: string to use as plot title instead of time :kword cbtitle: string to use as colorbar title instead of phase space density of flux :kword fmin,fmax: min and max values for colour scale and colour bar. If no values are given, min and max values for whole plot are used. :kword box: extents of plotted velocity grid as [x0,x1,y0,y1] (in m/s) :kword unit: Plot v-axes using 10^{unit} m/s (default: km/s) :kword xy: Perform slice in x-y-direction :kword xz: Perform slice in x-z-direction :kword yz: Perform slice in y-z-direction :kword normal: Perform slice in plane perpendicular to given vector :kword bpara: Perform slice in B_para / B_perp2 plane :kword bperp: Perform slice in B_perp1 / B_perp2 plane If no plane is given, default is simulation plane (for 2D simulations) :kword coordswap: Swap the parallel and perpendicular coordinates :kword cbulk: Center plot on position of total bulk velocity (or if not available, bulk velocity for this population) :kword center: Center plot on provided 3-element velocity vector position (in m/s) :kword wflux: Plot flux instead of distribution function :kword slicethick: Thickness of slice as multiplier of cell size (default: 1 or minimum for good coverage) :kword cellsize: Plotting grid cell size as multiplier of input cell size (default: 1 or minimum for good coverage) :kword keepfmin: Also draw buffer cells with values below fMin :kword wmark: If set to non-zero, will plot a Vlasiator watermark in the top left corner. :kword draw: Draw image on-screen instead of saving to file (requires x-windowing) :kword cbar: Plot colourbar legend (default off). :kword biglabel: Plot large label (in top-left corner) :kword biglabloc: Move large label to: 0: NW 1: NE 2: SE 3: SW corner :kword noborder: Plot figure edge-to-edge without borders (default off) :kword noxlabels: Suppress x-axis labels and title :kword noylabels: Suppress y-axis labels and title :kword scale: Scale text size (default=1.0) :kword thick: line and axis thickness, default=1.0 :returns: Outputs an image to a file or to the screen. .. code-block:: python # Example usage: Note tilted slices: By default, the program samples the V-space with a slice where each cell is cube the dimensions of which are found by performing a rotation on a sample square and finding the maximum xyz-extent. This ensures adequate coverage and decreases sampling effects. This behaviour can be overridden with the slicethick and cellsize keywords. ''' # Verify the location of this watermark image watermarkimage = os.path.join(os.path.dirname(__file__), 'logo_color.png') # watermarkimage=os.path.expandvars('$HOME/appl_taito/analysator/pyPlot/logo_color.png') # watermarkimage='/homeappl/home/marbat/appl_taito/analysator/logo_color.png' outputprefix = '' if outputdir == None: outputdir = os.path.expandvars('$HOME/Plots/') outputprefixind = outputdir.rfind('/') if outputprefixind >= 0: outputprefix = outputdir[outputprefixind + 1:] outputdir = outputdir[:outputprefixind + 1] if not os.path.exists(outputdir): os.makedirs(outputdir) # Input file or object if filename != None: vlsvReader = pt.vlsvfile.VlsvReader(filename) elif vlsvobj != None: vlsvReader = vlsvobj elif ((filedir != None) and (step != None)): filename = filedir + 'bulk.' + str(step).rjust(7, '0') + '.vlsv' vlsvReader = pt.vlsvfile.VlsvReader(filename) else: print( "Error, needs a .vlsv file name, python object, or directory and step" ) return if colormap == None: colormap = "hot_desaturated" cmapuse = matplotlib.cm.get_cmap(name=colormap) fontsize = 8 * scale # Most text fontsize2 = 10 * scale # Time title fontsize3 = 5 * scale # Colour bar ticks fontsize4 = 16 * scale # Big label # Plot title with time timeval = vlsvReader.read_parameter("time") if timeval == None: timeval = vlsvReader.read_parameter("t") if title == None: if timeval == None: plot_title = '' print("Unknown time format encountered") else: plot_title = "t=" + str(np.int(timeval)) + ' s' else: plot_title = title # step, used for file name if step != None: stepstr = '_' + str(step).rjust(7, '0') else: if timeval != None: stepstr = '_t' + str(np.int(timeval)) else: stepstr = '' # If run name isn't given, just put "plot" in the output file name if run == None: run = 'plot' # If population isn't defined i.e. defaults to protons, check if # instead should use old version "avgs" if pop == "proton": if not vlsvReader.check_population(pop): if vlsvReader.check_population("avgs"): pop = "avgs" print("Auto-switched to population avgs") else: print("Unable to detect population " + pop + " in .vlsv file!") exit() else: if not vlsvReader.check_population(pop): print("Unable to detect population " + pop + " in .vlsv file!") exit() #read in mesh size and cells in ordinary space # xsize = vlsvReader.read_parameter("xcells_ini") # ysize = vlsvReader.read_parameter("ycells_ini") # zsize = vlsvReader.read_parameter("zcells_ini") [xsize, ysize, zsize] = vlsvReader.get_spatial_mesh_size() # cellids = vlsvReader.read_variable("CellID") # vxsize = vlsvReader.read_parameter("vxblocks_ini")*4 # vysize = vlsvReader.read_parameter("vyblocks_ini")*4 # vzsize = vlsvReader.read_parameter("vzblocks_ini")*4 # vxmin = vlsvReader.read_parameter("vxmin") # vxmax = vlsvReader.read_parameter("vxmax") # vymin = vlsvReader.read_parameter("vymin") # vymax = vlsvReader.read_parameter("vymax") # vzmin = vlsvReader.read_parameter("vzmin") # vzmax = vlsvReader.read_parameter("vzmax") # These apparently work with multipop at least for the default mesh of protons. # Also works with older vlsv versions. [vxsize, vysize, vzsize] = vlsvReader.get_velocity_mesh_size(pop=pop) [vxmin, vymin, vzmin, vxmax, vymax, vzmax] = vlsvReader.get_velocity_mesh_extent(pop=pop) inputcellsize = (vxmax - vxmin) / vxsize # account for 4x4x4 cells per block vxsize = 4 * vxsize vysize = 4 * vysize vzsize = 4 * vzsize Re = 6.371e+6 # Earth radius in m # unit of velocity velUnit = 1e3 velUnitStr = '[km/s]' if unit != None: velUnit = np.power(10, int(unit)) if unit == 1: velUnitStr = r'[m/s]' else: velUnitStr = r'[$10^{' + str(int(unit)) + '}$ m/s]' # Select ploitting back-end based on on-screen plotting or direct to file without requiring x-windowing if draw != None: plt.switch_backend('TkAgg') else: plt.switch_backend('Agg') if (cellids == None and coordinates == None and coordre == None): print("Error: must provide either cell id's or coordinates") return -1 if coordre != None: # Transform to metres coordinates = Re * np.asarray(coordre) if coordinates != None: if type(coordinates) is not list: coordinates = [coordinates] # Calculate cell IDs from given coordinates xReq = np.asarray(coordinates).T[0] yReq = np.asarray(coordinates).T[1] zReq = np.asarray(coordinates).T[2] if xReq.shape == yReq.shape == zReq.shape: print('Number of points: ' + str(xReq.shape[0])) else: print('ERROR: bad coordinate variables given') exit() cidsTemp = [] for ii in range(xReq.shape[0]): cidRequest = (np.int64)(vlsvReader.get_cellid( np.array([xReq[ii], yReq[ii], zReq[ii]]))) cidNearestVspace = -1 if cidRequest > 0: cidNearestVspace = getNearestCellWithVspace( vlsvReader, cidRequest) else: print('ERROR: cell not found') exit() if (cidNearestVspace <= 0): print('ERROR: cell with vspace not found') exit() xCid, yCid, zCid = vlsvReader.get_cell_coordinates(cidRequest) xVCid, yVCid, zVCid = vlsvReader.get_cell_coordinates( cidNearestVspace) print('Point: ' + str(ii + 1) + '/' + str(xReq.shape[0])) print('Requested coordinates : ' + str(xReq[ii] / Re) + ', ' + str(yReq[ii] / Re) + ', ' + str(zReq[ii] / Re)) print('Nearest spatial cell : ' + str(xCid / Re) + ', ' + str(yCid / Re) + ', ' + str(zCid / Re)) print('Nearest vspace : ' + str(xVCid / Re) + ', ' + str(yVCid / Re) + ', ' + str(zVCid / Re)) cidsTemp.append(cidNearestVspace) cellids = np.unique(cidsTemp) print('Unique cells with vspace found: ' + str(len(cidsTemp))) else: print('Using given cell ids and assuming vspace is stored in them') # Loop over all cell ids if type(cellids) is not list: cellids = [cellids] print(cellids) for cellid in cellids: # Initialise some values fminuse = None fmaxuse = None x, y, z = vlsvReader.get_cell_coordinates(cellid) print('cellid ' + str(cellid) + ', x = ' + str(x) + ', y = ' + str(y) + ', z = ' + str(z)) savefigname = outputdir + outputprefix + run + "_vdf_" + pop + stepstr + "_cellid_" + str( cellid) + ".png" # Check slice to perform (and possibly normal vector) normvect = None if xy == None and xz == None and yz == None and normal == None and bpara == None and bperp == None: # Use default slice for this simulation # Check if ecliptic or polar run if ysize == 1: # polar slicetype = "xz" pltxstr = r"$v_x$ " + velUnitStr pltystr = r"$v_z$ " + velUnitStr normvect = [0, 1, 0] # used just for cell size normalisation elif zsize == 1: # ecliptic slicetype = "xy" pltxstr = r"$v_x$ " + velUnitStr pltystr = r"$v_y$ " + velUnitStr normvect = [0, 0, 1] # used just for cell size normalisation else: print("Problem finding default slice direction") slicetype = "yz" pltxstr = r"$v_y$ " + velUnitStr pltystr = r"$v_z$ " + velUnitStr normvect = [1, 0, 0] # used just for cell size normalisation elif normal != None: if len(normal) == 3: slicetype = "vecperp" normvect = normal pltxstr = r"$v_1$ " + velUnitStr pltystr = r"$v_2$ " + velUnitStr else: print("Error parsing slice normal vector!") exit() elif xy != None: slicetype = "xy" pltxstr = r"$v_x$ " + velUnitStr pltystr = r"$v_y$ " + velUnitStr normvect = [0, 0, 1] # used just for cell size normalisation elif xz != None: slicetype = "xz" pltxstr = r"$v_x$ " + velUnitStr pltystr = r"$v_z$ " + velUnitStr normvect = [0, 1, 0] # used just for cell size normalisation elif yz != None: slicetype = "yz" pltxstr = r"$v_y$ " + velUnitStr pltystr = r"$v_z$ " + velUnitStr normvect = [1, 0, 0] # used just for cell size normalisation elif bpara != None or bperp != None: # Rotate based on B-vector if vlsvReader.check_variable("B"): Bvect = vlsvReader.read_variable("B", cellid) elif (vlsvReader.check_variable("background_B") and vlsvReader.check_variable("perturbed_B")): # used e.g. for restart files BGB = vlsvReader.read_variable("background_B", cellid) PERBB = vlsvReader.read_variable("perturbed_B", cellid) Bvect = BGB + PERBB else: print("Error finding B vector direction!") exit() if Bvect.shape == (1, 3): Bvect = Bvect[0] normvect = Bvect if bperp != None: # slice in b_perp1/b_perp2 slicetype = "vecperp" pltxstr = r"$v_{\perp 1}$ " + velUnitStr pltystr = r"$v_{\perp 2}$ " + velUnitStr else: # means bpara!=None, slice in b_parallel/b_perp2 plane slicetype = "vecpara" pltxstr = r"$v_{\parallel}$ " + velUnitStr pltystr = r"$v_{\perp}$ " + velUnitStr # Extend velocity space and each cell to account for slice directions oblique to axes normvect = np.array(normvect) normvect = normvect / np.linalg.norm(normvect) # Geometric magic to stretch the grid to assure that each cell has some velocity grid points inside it. # Might still be incorrect, erring on the side of caution. # norm_srt = sorted(abs(normvect)) # if cellsize==None: # if norm_srt[1] > 0: # temp = norm_srt[0]/norm_srt[1] # aratio = (1.+temp)/np.sqrt( 1+temp**2) # else: # aratio = 1. # gridratio = aratio * norm_srt[2] * (1. + aratio * np.sqrt(norm_srt[0]**2 + norm_srt[1]**2) / norm_srt[2] ) # if gridratio>1.0: # gridratio = gridratio*1.01 # account for numerical inaccuracies # else: # gridratio = 1. if cellsize == None: samplebox = np.array([[0.0, 0.0, 0.0], [0.0, 0.0, 1.0], [0.0, 1.0, 0.0], [0.0, 1.0, 1.0], [1.0, 0.0, 0.0], [1.0, 0.0, 1.0], [1.0, 1.0, 0.0], [1.0, 1.0, 1.0]]) sbrot = rotateVectorToVector(samplebox, normvect) rotminx = np.amin(sbrot[:, 0]) rotmaxx = np.amax(sbrot[:, 0]) rotminy = np.amin(sbrot[:, 1]) rotmaxy = np.amax(sbrot[:, 1]) rotminz = np.amin(sbrot[:, 2]) rotmaxz = np.amax(sbrot[:, 2]) gridratio = np.amax( [rotmaxx - rotminx, rotmaxy - rotminy, rotmaxz - rotminz]) else: gridratio = cellsize # num must be vxsize+1 or vysize+1 in order to do both edges for each cell VXBins = np.linspace(vxmin * gridratio, vxmax * gridratio, num=vxsize + 1) VYBins = np.linspace(vymin * gridratio, vymax * gridratio, num=vysize + 1) # Read velocity data into histogram (checkOk, binsXY, edgesX, edgesY) = vSpaceReducer(vlsvReader, cellid, slicetype, normvect, VXBins, VYBins, pop=pop, slicethick=slicethick, wflux=wflux, cbulk=cbulk, center=center, keepfmin=keepfmin) # Check that data is ok and not empty if checkOk == False: print('ERROR: error from velocity space reducer') continue # Perform swap of coordinate axes, if requested if coordswap != None: temp = edgesX edgesX = edgesY edgesY = temp temp = pltxstr pltxstr = pltystr pltystr = temp binsXY = binsXY.T # If no other fmin fmax values are given, take min and max of array if fmin != None: fminuse = fmin else: nzindex = np.where(binsXY > 0) if np.any(nzindex): fminuse = np.amin(binsXY[nzindex]) else: fminuse = 1e-15 if fmax != None: fmaxuse = fmax else: nzindex = np.where(binsXY > 0) if np.any(nzindex): fmaxuse = np.amax(binsXY[nzindex]) else: fmaxuse = 1e-12 if vlsvReader.check_variable('MinValue') == True: fMinFile = vlsvReader.read_variable('MinValue', cellid) print("Active f range is " + str(fminuse) + " to " + str(fmaxuse) + " with a vlsv file fMin value of " + str(fMinFile)) else: print("Active f range is " + str(fminuse) + " to " + str(fmaxuse)) norm = LogNorm(vmin=fminuse, vmax=fmaxuse) ticks = LogLocator(base=10, subs=range(10)) # where to show labels if box != None: # extents of plotted velocity grid as [x0,y0,x1,y1] xvalsrange = [box[0], box[1]] yvalsrange = [box[2], box[3]] else: # Find extent of nonzero data xindexrange = [vxsize, 0] yindexrange = [vysize, 0] for xi in range(len(edgesX) - 1): for yi in range(len(edgesY) - 1): if binsXY[xi, yi] > 0: xindexrange[0] = np.amin([xindexrange[0], xi]) xindexrange[1] = np.amax([xindexrange[1], xi]) yindexrange[0] = np.amin([yindexrange[0], yi]) yindexrange[1] = np.amax([yindexrange[1], yi]) # leave some buffer xindexrange[0] = np.max( [0, 4 * int(np.floor((xindexrange[0] - 2.) / 4.))]) xindexrange[1] = np.min([ len(edgesX) - 1, 4 * int(np.ceil((xindexrange[1] + 2.) / 4.)) ]) yindexrange[0] = np.max( [0, 4 * int((np.floor(yindexrange[0] - 2.) / 4.))]) yindexrange[1] = np.min([ len(edgesY) - 1, 4 * int(np.ceil((yindexrange[1] + 2.) / 4.)) ]) # If empty VDF: plot whole v-space if ((xindexrange == [vxsize, 0]) and (yindexrange == [vysize, 0])): xindexrange = [0, vxsize] yindexrange = [0, vysize] xvalsrange = [edgesX[xindexrange[0]], edgesX[xindexrange[1]]] yvalsrange = [edgesY[yindexrange[0]], edgesY[yindexrange[1]]] # TODO make plot area square if it's almost square? # Define figure size ratio = (yvalsrange[1] - yvalsrange[0]) / (xvalsrange[1] - xvalsrange[0]) figsize = [3.0, 3.0 * ratio] if noxlabels == None: figsize[1] = figsize[1] + 1.00 #0.25 if noylabels == None: figsize[0] = figsize[0] + 1.00 #0.25 if cbar != None: figsize[0] = figsize[0] + 0.70 #0.20 if title != None and cbtitle != None: # If either title exists, make room for them: if len(title) != 0 or len(cbtitle) != 0: figsize[1] = figsize[1] + 0.70 #0.20 # If both titles have been turned off, do nothing else: # Default titles are in place, make room for them. figsize[1] = figsize[1] + 0.70 #0.20 # Create 300 dpi image of suitable size fig = plt.figure(figsize=figsize, dpi=300) # Plot the slice [XmeshXY, YmeshXY] = scipy.meshgrid( edgesX / velUnit, edgesY / velUnit) # Generates the mesh to map the data to fig1 = plt.pcolormesh(XmeshXY, YmeshXY, binsXY, cmap=colormap, norm=norm) ax1 = plt.gca() # get current axes plt.xlim([val / velUnit for val in yvalsrange]) plt.ylim([val / velUnit for val in xvalsrange]) ax1.set_aspect('equal') # Grid plt.grid(color='grey', linestyle='-') plt.minorticks_on() for axiss in ['top', 'bottom', 'left', 'right']: ax1.spines[axiss].set_linewidth(thick) ax1.xaxis.set_tick_params(width=thick, length=4) ax1.yaxis.set_tick_params(width=thick, length=4) ax1.xaxis.set_tick_params(which='minor', width=thick * 0.8, length=2) ax1.yaxis.set_tick_params(which='minor', width=thick * 0.8, length=2) ax1.set_title(plot_title, fontsize=fontsize2, fontweight='bold') if noxlabels == None: plt.xlabel(pltxstr, fontsize=fontsize, weight='black') plt.xticks(fontsize=fontsize, fontweight='black') ax1.xaxis.offsetText.set_fontsize(fontsize) if noylabels == None: plt.ylabel(pltystr, fontsize=fontsize, weight='black') plt.yticks(fontsize=fontsize, fontweight='black') ax1.yaxis.offsetText.set_fontsize(fontsize) if biglabel != None: if biglabloc == None: biglabloc = 0 # default top-left corner if biglabloc == 0: BLcoords = [0.02, 0.98] BLha = "left" BLva = "top" elif biglabloc == 1: BLcoords = [0.98, 0.98] BLha = "right" BLva = "top" elif biglabloc == 2: BLcoords = [0.98, 0.02] BLha = "right" BLva = "bottom" elif biglabloc == 3: BLcoords = [0.02, 0.02] BLha = "left" BLva = "bottom" plt.text(BLcoords[0], BLcoords[1], biglabel, fontsize=fontsize4, weight='black', transform=ax1.transAxes, ha=BLha, va=BLva) if cbar != None: # Colourbar title if cbtitle != None: if len(cbtitle) != 0: cb_title_locy = 1.0 + 0.03 / ratio plt.text(1.0, cb_title_locy, cbtitle, fontsize=fontsize3, weight='black', transform=ax1.transAxes) else: if wflux == None: plt.text(1.05, 1.22, r"$f(v)$", fontsize=fontsize3, weight='black', transform=ax1.transAxes) plt.text(1.02, 1.1, r"$[\mathrm{m}^{-6} \,\mathrm{s}^{3}]$", fontsize=fontsize3, weight='black', transform=ax1.transAxes) else: plt.text(1.05, 1.22, r"flux F", fontsize=fontsize3, weight='black', transform=ax1.transAxes) plt.text( 1.02, 1.1, r"$[\mathrm{m}^{-2} \,\mathrm{s}^{-1} \,\mathrm{sr}^{-1}]$", fontsize=fontsize3, weight='black', transform=ax1.transAxes) # Witchcraft used to place colourbar divider = make_axes_locatable(ax1) cax = divider.append_axes("right", size="5%", pad=0.05) # First draw colorbar cb = plt.colorbar(fig1, ticks=ticks, cax=cax) cb.ax.tick_params(labelsize=fontsize3) #,width=1.5,length=3) cb.outline.set_linewidth(thick) # # if too many subticks: # # For non-square pictures, adjust tick count # nlabels = len(cb.ax.yaxis.get_ticklabels()) / ratio # if nlabels > 10: # valids = ['1','2','3','4','5','6','8'] # if nlabels > 19: # valids = ['1','2','5'] # if nlabels > 28: # valids = ['1'] # # for label in cb.ax.yaxis.get_ticklabels()[::labelincrement]: # if nlabels > 10: # for label in cb.ax.yaxis.get_ticklabels(): # # labels will be in format $x.0\times10^{y}$ # if not label.get_text()[1] in valids: # label.set_visible(False) if noxlabels != None: for label in ax1.xaxis.get_ticklabels(): label.set_visible(False) if noylabels != None: for label in ax1.yaxis.get_ticklabels(): label.set_visible(False) if noborder == None: # adjust layout plt.tight_layout() savefig_pad = 0.1 # The default is 0.1 bbox_inches = None else: # adjust layout plt.tight_layout(pad=0.01) savefig_pad = 0.01 bbox_inches = 'tight' # Add Vlasiator watermark if wmark != None: wm = plt.imread(get_sample_data(watermarkimage)) newax = fig.add_axes([0.01, 0.90, 0.3, 0.08], anchor='NW', zorder=-1) newax.imshow(wm) newax.axis('off') # Save output or draw on-screen if draw == None: print(savefigname + "\n") plt.savefig(savefigname, dpi=300, bbox_inches=bbox_inches, pad_inches=savefig_pad) else: plt.draw() plt.show() plt.close() plt.clf()
def vSpaceReducer(vlsvReader, cid, slicetype, normvect, VXBins, VYBins, pop="proton", slicethick=None, wflux=None, cbulk=None, center=None, keepfmin=None): # check if velocity space exists in this cell if vlsvReader.check_variable( 'fSaved'): #restart files will not have this value if vlsvReader.read_variable('fSaved', cid) != 1.0: return (False, 0, 0, 0) # Assume velocity cells are cubes [vxsize, vysize, vzsize] = vlsvReader.get_velocity_mesh_size(pop=pop) # Account for 4x4x4 cells per block vxsize = 4 * vxsize vysize = 4 * vysize vzsize = 4 * vzsize [vxmin, vymin, vzmin, vxmax, vymax, vzmax] = vlsvReader.get_velocity_mesh_extent(pop=pop) inputcellsize = (vxmax - vxmin) / vxsize print("Input velocity grid cell size " + str(inputcellsize)) velcells = vlsvReader.read_velocity_cells(cid, pop=pop) V = vlsvReader.get_velocity_cell_coordinates(velcells.keys(), pop=pop) print("Found " + str(len(V)) + " v-space cells") if cbulk != None: print("Transforming to plasma frame") if vlsvReader.check_variable('moments'): # This should be a restart file moments = np.array(vlsvReader.read_variable('moments', cid)) if moments == None: print("Error reading moments from assumed restart file!") exit() if len(moments.shape) == 2: moments = moments[0] if moments[0] > 0.0: bulkv = moments[1:4] / moments[0] else: bulkv = [0., 0., 0.] elif vlsvReader.check_variable('v'): # Multipop file with bulk v saved directly bulkv = vlsvReader.read_variable('v', cid) elif (vlsvReader.check_variable('rho') and vlsvReader.check_variable('rho_v')): # Older regular bulk file rhov = vlsvReader.read_variable('rho_v', cid) rho = vlsvReader.read_variable('rho', cid) bulkv = rhov / rho elif vlsvReader.check_variable(pop + '/V'): # Multipop file without V saved, use per-population bulk velocity bulkv = vlsvReader.read_variable(pop + '/V', cid) else: print("Error in finding plasma bulk velocity!") exit() # shift to velocities plasma frame V = V - bulkv elif center != None: if len(center) == 3: print("Transforming to frame travelling at speed " + str(center)) V - V - center else: print("Error in shape of center vector! Give in form (vx,vy,vz).") f = zip(*velcells.items()) # check that velocity space has cells if (len(f) > 0): f = np.asarray(zip(*velcells.items())[1]) else: return (False, 0, 0, 0) if keepfmin == None: # Drop all velocity cells which are below the sparsity threshold. Otherwise the plot will show buffer cells as well. fMin = 1e-16 # default if vlsvReader.check_variable('MinValue') == True: fMin = vlsvReader.read_variable('MinValue', cid) ii_f = np.where(f >= fMin) print("Dropping velocity cells under fMin value " + str(fMin)) if len(ii_f) < 1: return (False, 0, 0, 0) f = f[ii_f] V = V[ii_f, :][0, :, :] if slicethick == None: # Geometric magic to widen the slice to assure that each cell has some velocity grid points inside it. # Might still be incorrect, erring on the side of caution. # norm_srt = sorted(abs(normvect)) # if norm_srt[1] > 0: # temp = norm_srt[0]/norm_srt[1] # aratio = (1.+temp)/np.sqrt( 1+temp**2) # else: # aratio = 1. # gridratio = aratio * norm_srt[2] * (1. + aratio * np.sqrt(norm_srt[0]**2 + norm_srt[1]**2) / norm_srt[2] ) # if gridratio>1.0: # # Account for numerical inaccuracy # slicethick=inputcellsize*gridratio*1.01 # else: # slicethick=inputcellsize samplebox = np.array([[0.0, 0.0, 0.0], [0.0, 0.0, 1.0], [0.0, 1.0, 0.0], [0.0, 1.0, 1.0], [1.0, 0.0, 0.0], [1.0, 0.0, 1.0], [1.0, 1.0, 0.0], [1.0, 1.0, 1.0]]) sbrot = rotateVectorToVector(samplebox, normvect) rotminx = np.amin(sbrot[:, 0]) rotmaxx = np.amax(sbrot[:, 0]) rotminy = np.amin(sbrot[:, 1]) rotmaxy = np.amax(sbrot[:, 1]) rotminz = np.amin(sbrot[:, 2]) rotmaxz = np.amax(sbrot[:, 2]) gridratio = np.amax( [rotmaxx - rotminx, rotmaxy - rotminy, rotmaxz - rotminz]) slicethick = inputcellsize * gridratio else: slicethick = inputcellsize * slicethick print("Performing slice with a counting thickness of " + str(slicethick)) if slicetype == "xy": VX = V[:, 0] VY = V[:, 1] Vpara = V[:, 2] elif slicetype == "yz": VX = V[:, 1] VY = V[:, 2] Vpara = V[:, 0] elif slicetype == "xz": VX = V[:, 0] VY = V[:, 2] Vpara = V[:, 1] elif slicetype == "vecperp": # Find velocity components in give nframe (e.g. B frame: (vx,vy,vz) -> (vperp2,vperp1,vpar)) N = np.array(normvect) / np.sqrt(normvect[0]**2 + normvect[1]**2 + normvect[2]**2) Vrot = rotateVectorToVector(V, N) VX = Vrot[:, 0] VY = Vrot[:, 1] Vpara = Vrot[:, 2] elif slicetype == "vecpara": N = np.array(normvect) / np.sqrt(normvect[0]**2 + normvect[1]**2 + normvect[2]**2) Vrot = rotateVectorToVector(V, N) VX = Vrot[:, 2] VY = Vrot[:, 1] Vpara = Vrot[:, 0] else: print("Error finding rotation of v-space!") return (False, 0, 0, 0) # TODO: better rotation so perpendicular component can be defined? # create 2-dimensional histogram of velocity components perpendicular to slice-normal vector (binsXY, edgesX, edgesY) = doHistogram(f, VX, VY, Vpara, VXBins, VYBins, slicethick, wflux=wflux) return (True, binsXY, edgesX, edgesY)