diff --git a/pyCalculations/calculations.py b/pyCalculations/calculations.py index 483f8bc7..f4f798f6 100644 --- a/pyCalculations/calculations.py +++ b/pyCalculations/calculations.py @@ -39,8 +39,9 @@ # List of functions and classes that should be imported into the interface from intpol_file import vlsv_intpol_file from intpol_points import vlsv_intpol_points -from cutthrough import cut_through, cut_through_step +from cutthrough import cut_through, cut_through_step, cut_through_curve, cut_through_swath from fourier import fourier +from spectra import get_spectrum_energy, get_spectrum_alongaxis_vel from variable import VariableInfo from timeevolution import cell_time_evolution from pitchangle import pitch_angles diff --git a/pyCalculations/cutthrough.py b/pyCalculations/cutthrough.py index 49d4e3d5..1f338cc3 100644 --- a/pyCalculations/cutthrough.py +++ b/pyCalculations/cutthrough.py @@ -32,7 +32,7 @@ def get_cellids_coordinates_distances( vlsvReader, xmax, xmin, xcells, ymax, ymi :type vlsvReader: :class:`vlsvfile.VlsvReader` :param point1: The starting point of a cut-through line :param point2: The ending point of a cut-through line - :returns: Coordinates of for the cell cut-through as well as cell ids and distances + :returns: Coordinates of for the cell cut-through as well as cell ids and distances. NB: Last CID is a duplicate. ''' point1 = np.array(point1, copy=False) point2 = np.array(point2, copy=False) @@ -47,6 +47,13 @@ def get_cellids_coordinates_distances( vlsvReader, xmax, xmin, xcells, ymax, ymi epsilon = sys.float_info.epsilon*2 iterator = point1 + + # Special case: both points in the same cell. + #if(vlsvReader.get_cellid(point1) == vlsvReader.get_cellid(point2)): + # cellids.append(vlsvReader.get_cellid(point2)) + + + unit_vector = (point2 - point1) / np.linalg.norm(point2 - point1 + epsilon) while True: # Get the cell id @@ -58,9 +65,8 @@ def get_cellids_coordinates_distances( vlsvReader, xmax, xmin, xcells, ymax, ymi min_bounds = vlsvReader.get_cell_coordinates(cellid) - 0.5 * cell_lengths max_bounds = np.array([min_bounds[i] + cell_lengths[i] for i in range(0,3)]) # Check which boundary we hit first when we move in the unit_vector direction: - - coefficients_min = np.divide((min_bounds - iterator), unit_vector) - coefficients_max = np.divide((max_bounds - iterator) , unit_vector) + coefficients_min = np.divide((min_bounds - iterator), unit_vector, out=np.full(min_bounds.shape, sys.float_info.max), where=np.logical_not(unit_vector==0.0)) + coefficients_max = np.divide((max_bounds - iterator), unit_vector, out=np.full(min_bounds.shape, sys.float_info.max), where=np.logical_not(unit_vector==0.0)) # Remove negative coefficients: for i in range(3): if coefficients_min[i] <= 0: @@ -71,10 +77,15 @@ def get_cellids_coordinates_distances( vlsvReader, xmax, xmin, xcells, ymax, ymi # Find the minimum coefficient for calculating the minimum distance from a boundary - coefficient = min([min(coefficients_min), min(coefficients_max)]) * 1.01 + coefficient = min([min(coefficients_min), min(coefficients_max)]) * 1.000001 # Get the cell id in the new coordinates newcoordinate = iterator + coefficient * unit_vector + + # Check to make sure the iterator is not moving past point2 - if it is, use point2 instead: + if min((point2 - newcoordinate)* unit_vector) < 0: + newcoordinate = point2 + newcellid = vlsvReader.get_cellid( newcoordinate ) # Check if valid cell id: if newcellid == 0: @@ -107,7 +118,7 @@ def cut_through( vlsvReader, point1, point2 ): :type vlsvReader: :class:`vlsvfile.VlsvReader` :param point1: The starting point of a cut-through line :param point2: The ending point of a cut-through line - :returns: an array containing cell ids, coordinates and distances in the following format: [cell ids, distances, coordinates] + :returns: an array containing cell ids, coordinates and distances in the following format: [cell ids, distances, coordinates]. NB. Last cellid is a duplicate. .. code-block:: python @@ -151,11 +162,33 @@ def cut_through( vlsvReader, point1, point2 ): return get_cellids_coordinates_distances( vlsvReader, xmax, xmin, xcells, ymax, ymin, ycells, zmax, zmin, zcells, cell_lengths, point1, point2 ) +def cut_through_swath(vlsvReader, point1, point2, width, normal): + init_cut = cut_through(vlsvReader, point1, point2) + init_cids = init_cut[0].data + + print('swath initial CellIds', init_cids) + + #find the other vector spanning the swath + s = np.array(point2)-np.array(point1) + w = np.cross(s, np.array(normal)) + w = width*w/np.linalg.norm(w) + + out_cids = [] + for cid in init_cids: + mid = vlsvReader.get_cell_coordinates(cid) + p1 = mid - w/2 + p2 = mid + w/2 + temp_cut = cut_through_step(vlsvReader, p1, p2) + out_cids.append(temp_cut[0].data) + + init_cut[0] = np.array(out_cids) + print(init_cut[0]) + return init_cut def cut_through_step( vlsvReader, point1, point2 ): ''' Returns cell ids and distances from point 1 to point 2, returning not every cell in a line - but rather the amount of cells which correspons with the largest axis-aligned component of the line. + but rather the amount of cells which corresponds with the largest axis-aligned component of the line. :param vlsvReader: Some open VlsvReader :type vlsvReader: :class:`vlsvfile.VlsvReader` @@ -228,4 +261,100 @@ def cut_through_step( vlsvReader, point1, point2 ): # Return the coordinates, cellids and distances for processing from output import output_1d return output_1d( [np.array(cellids, copy=False), np.array(distances, copy=False), np.array(coordinates, copy=False)], ["CellID", "distances", "coordinates"], ["", "m", "m"] ) - + + +# Take a curve (list of 3-coords), return cells and distances along curve as with cut_through. +# NB: calls get_cells_coordinates_distances for each curve segment. Feel free to optimize your curves beforehand. +def cut_through_curve(vlsvReader, curve): + ''' Returns cell ids and distances from point 1 for every cell in a line between given point1 and point2 + + :param vlsvReader: Some open VlsvReader + :type vlsvReader: :class:`vlsvfile.VlsvReader` + :param point1: The starting point of a cut-through line + :returns: an array containing cell ids, edge distances, and the coordinates of edges in the following format: [cell ids, distances, coordinates]. NB. Last cellid is a duplicate. + + .. code-block:: python + + Example: + vlsvReader = VlsvReader(\"testfile.vlsv\") + cut_through_curve = cut_through(vlsvReader, [[0,0,0], [2,5e6,0]]) + cellids = cut_through[0] + distances = cut_through[1] + print \"Cell ids: \" + str(cellids) + print \"Distance from point 1 for every cell: \" + str(distances) + ''' + # init cut_through static values, then do what cut_through does for each segment + # Get parameters from the file to determine a good length between points (step length): + # Get xmax, xmin and xcells_ini + mesh_limits = vlsvReader.get_spatial_mesh_extent() + mesh_size = vlsvReader.get_spatial_mesh_size() + xmax = mesh_limits[3] + xmin = mesh_limits[0] + xcells = mesh_size[0] + # Do the same for y + ymax = mesh_limits[4] + ymin = mesh_limits[1] + ycells = mesh_size[1] + # And for z + zmax = mesh_limits[5] + zmin = mesh_limits[2] + zcells = mesh_size[2] + + + #Calculate cell lengths: + cell_lengths = np.array([(xmax - xmin)/(float)(xcells), (ymax - ymin)/(float)(ycells), (zmax - zmin)/(float)(zcells)]) + + if(len(curve) < 2): + print("ERROR, less than 2 points in a curve") + return None + + cellIds = [] + edges = [] + coords = [] + for i in range(len(curve)-1): + point1 = np.array(curve[i]) + point2 = np.array(curve[i+1]) + # Make sure point1 and point2 are inside bounds + if vlsvReader.get_cellid(point1) == 0: + print("ERROR, POINT1 IN CUT-THROUGH-CURVE OUT OF BOUNDS!") + if vlsvReader.get_cellid(point2) == 0: + print("ERROR, POINT2 IN CUT-THROUGH-CURVE OUT OF BOUNDS!") + cut = get_cellids_coordinates_distances( vlsvReader, xmax, xmin, xcells, ymax, ymin, ycells, zmax, zmin, zcells, cell_lengths, point1, point2) + ccid = cut[0].data + cedges = cut[1].data + try: + edgestart = edges[-1] + except: + edgestart = 0 + for ci,citer in enumerate(ccid): + cellIds.append(citer) + edges.append(edgestart+cedges[ci]) + coords.append(cut[2].data[ci]) + + #print('sum of diffs', np.sum(np.diff(edges))) + #reduce the cellIDs and edges + reduct = True + if reduct: + rCellIds = [cellIds[0]] + rEdges = [edges[0]] + rCoords = [coords[0]] + cid0 = cellIds[0] + for i in range(1, len(cellIds)-1): + c1 = cellIds[i] + if(cid0 == c1): + rEdges[-1] = edges[i] + rCoords[-1] = coords[i] + continue + rCellIds.append(c1) + rEdges.append(edges[i]) + rCoords.append(coords[i]) + cid0 = c1 + rEdges.append(edges[-1]) + rCoords.append(coords[-1]) + #print('sum of r-diffs', np.sum(np.diff(rEdges))) + else: + rCellIds = cellIds + rEdges = edges + rCoords = coords + from output import output_1d + return output_1d( [np.array(rCellIds, copy=False), np.array(rEdges, copy=False), np.array(rCoords, copy=False)], ["CellID", "distances", "coordinates"], ["", "m", "m"] ) diff --git a/pyCalculations/fieldtracer.py b/pyCalculations/fieldtracer.py index 2b9450a7..39d66887 100644 --- a/pyCalculations/fieldtracer.py +++ b/pyCalculations/fieldtracer.py @@ -46,7 +46,7 @@ def dynamic_field_tracer( vlsvReader_list, x0, max_iterations, dx): x0 = x0 + v*dt iterations = iterations + 1 -def static_field_tracer( vlsvReader, x0, max_iterations, dx, direction='+', bvar='B' ): +def static_field_tracer( vlsvReader, x0, max_iterations, dx, direction='+', bvar='B', centering='default', boundary_inner=-1): ''' Field tracer in a static frame :param vlsvReader: An open vlsv file @@ -55,10 +55,12 @@ def static_field_tracer( vlsvReader, x0, max_iterations, dx, direction='+', bvar :param dx: One iteration step length :param direction: '+' or '-' or '+-' Follow field in the plus direction or minus direction :param bvar: String, variable name to trace [default 'B'] + :param centering: String, variable centering: 'face', 'volume', 'node' [defaults to 'face'] + :param boundary_inner: Float, stop propagation if closer to origin than this value [default -1] :returns: List of coordinates ''' - if(bvar is not 'B'): + if(bvar != 'B'): warnings.warn("User defined tracing variable detected. fg, volumetric variable results may not work as intended, use face-values instead.") if direction == '+-': @@ -93,31 +95,67 @@ def static_field_tracer( vlsvReader, x0, max_iterations, dx, direction='+', bvar if zsize <= 1: indices = [1,0] + if 'vol' in bvar and centering == 'default': + warnings.warn("Found 'vol' in variable name, assuming volumetric variable and adjusting centering") + centering = 'volume' + # Read face_B: - face_B = f.read_variable(bvar) - face_Bx = face_B[:,0] - face_By = face_B[:,1] - face_Bz = face_B[:,2] - - face_Bx = face_Bx[cellids.argsort()].reshape(sizes[indices]) - face_By = face_By[cellids.argsort()].reshape(sizes[indices]) - face_Bz = face_Bz[cellids.argsort()].reshape(sizes[indices]) - - face_B = np.array([face_Bx, face_By, face_Bz]) - - # Create x, y, and z coordinates: - x = np.arange(mins[0], maxs[0], dcell[0]) + 0.5*dcell[0] - y = np.arange(mins[1], maxs[1], dcell[1]) + 0.5*dcell[1] - z = np.arange(mins[2], maxs[2], dcell[2]) + 0.5*dcell[2] - coordinates = np.array([x,y,z]) - # Debug: - if( len(x) != sizes[0] ): - print("SIZE WRONG: " + str(len(x)) + " " + str(sizes[0])) - - # Create grid interpolation - interpolator_face_B_0 = interpolate.RectBivariateSpline(coordinates[indices[0]] - 0.5*dcell[indices[0]], coordinates[indices[1]], face_B[indices[0]], kx=2, ky=2, s=0) - interpolator_face_B_1 = interpolate.RectBivariateSpline(coordinates[indices[0]], coordinates[indices[1]] - 0.5*dcell[indices[1]], face_B[indices[1]], kx=2, ky=2, s=0) - interpolators = [interpolator_face_B_0, interpolator_face_B_1]#, interpolator_face_B_2] + if centering == 'face' or centering == 'default': + face_B = f.read_variable(bvar) + face_Bx = face_B[:,0] + face_By = face_B[:,1] + face_Bz = face_B[:,2] + + face_Bx = face_Bx[cellids.argsort()].reshape(sizes[indices]) + face_By = face_By[cellids.argsort()].reshape(sizes[indices]) + face_Bz = face_Bz[cellids.argsort()].reshape(sizes[indices]) + + face_B = np.array([face_Bx, face_By, face_Bz]) + + # Create x, y, and z coordinates: + x = np.arange(mins[0], maxs[0], dcell[0]) + 0.5*dcell[0] + y = np.arange(mins[1], maxs[1], dcell[1]) + 0.5*dcell[1] + z = np.arange(mins[2], maxs[2], dcell[2]) + 0.5*dcell[2] + coordinates = np.array([x,y,z]) + # Debug: + if( len(x) != sizes[0] ): + print("SIZE WRONG: " + str(len(x)) + " " + str(sizes[0])) + + # Create grid interpolation + interpolator_face_B_0 = interpolate.RectBivariateSpline(coordinates[indices[0]] - 0.5*dcell[indices[0]], coordinates[indices[1]], face_B[indices[0]], kx=2, ky=2, s=0) + interpolator_face_B_1 = interpolate.RectBivariateSpline(coordinates[indices[0]], coordinates[indices[1]] - 0.5*dcell[indices[1]], face_B[indices[1]], kx=2, ky=2, s=0) + interpolators = [interpolator_face_B_0, interpolator_face_B_1]#, interpolator_face_B_2] + elif centering == 'volume': + vol_B = f.read_variable(bvar) + vol_Bx = vol_B[:,0] + vol_By = vol_B[:,1] + vol_Bz = vol_B[:,2] + + vol_Bx = vol_Bx[cellids.argsort()].reshape(sizes[indices]) + vol_By = vol_By[cellids.argsort()].reshape(sizes[indices]) + vol_Bz = vol_Bz[cellids.argsort()].reshape(sizes[indices]) + + vol_B = np.array([vol_Bx, vol_By, vol_Bz]) + + # Create x, y, and z coordinates: + x = np.arange(mins[0], maxs[0], dcell[0]) + 0.5*dcell[0] + y = np.arange(mins[1], maxs[1], dcell[1]) + 0.5*dcell[1] + z = np.arange(mins[2], maxs[2], dcell[2]) + 0.5*dcell[2] + coordinates = np.array([x,y,z]) + # Debug: + if( len(x) != sizes[0] ): + print("SIZE WRONG: " + str(len(x)) + " " + str(sizes[0])) + + # Create grid interpolation + interpolator_vol_B_0 = interpolate.RectBivariateSpline(coordinates[indices[0]], coordinates[indices[1]], vol_B[indices[0]], kx=2, ky=2, s=0) + interpolator_vol_B_1 = interpolate.RectBivariateSpline(coordinates[indices[0]], coordinates[indices[1]], vol_B[indices[1]], kx=2, ky=2, s=0) + interpolators = [interpolator_vol_B_0, interpolator_vol_B_1]#, interpolator_face_B_2] + elif centering == 'node': + print("Nodal variables not implemented") + return + else: + print("Unrecognized centering:", centering) + return ####################################################### if direction == '-': @@ -125,14 +163,17 @@ def static_field_tracer( vlsvReader, x0, max_iterations, dx, direction='+', bvar else: multiplier = 1 - points = [x0] + points = [np.array(x0)] for i in range(max_iterations): - previous_point = points[len(points)-2] + previous_point = points[-1] B_unit = np.zeros(3) B_unit[indices[0]] = interpolators[0](previous_point[indices[0]], previous_point[indices[1]]) B_unit[indices[1]] = interpolators[1](previous_point[indices[0]], previous_point[indices[1]]) B_unit = B_unit / float(np.linalg.norm(B_unit)) - points.append( previous_point + multiplier*B_unit * dx ) + next_point = previous_point + multiplier*B_unit * dx + if(np.linalg.norm(next_point) < boundary_inner): + break + points.append(next_point) ####################################################### return points diff --git a/pyCalculations/spectra.py b/pyCalculations/spectra.py new file mode 100644 index 00000000..c7df2137 --- /dev/null +++ b/pyCalculations/spectra.py @@ -0,0 +1,162 @@ +# +# This file is part of Analysator. +# Copyright 2013-2016 Finnish Meteorological Institute +# Copyright 2017-2021 University of Helsinki +# +# For details of usage, see the COPYING file and read the "Rules of the Road" +# at http://www.physics.helsinki.fi/vlasiator/ +# +# This program is free software; you can redistribute it and/or modify +# it under the terms of the GNU General Public License as published by +# the Free Software Foundation; either version 2 of the License, or +# (at your option) any later version. +# +# This program is distributed in the hope that it will be useful, +# but WITHOUT ANY WARRANTY; without even the implied warranty of +# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +# GNU General Public License for more details. +# +# You should have received a copy of the GNU General Public License along +# with this program; if not, write to the Free Software Foundation, Inc., +# 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA. +# +import numpy as np +import pytools +# Function to reduce the velocity space in a spatial cell to an omnidirectional energy spectrum +# Weighted by particle flux/none +def get_spectrum_energy(vlsvReader, + cid, + population="proton", + fMin = 1e-21, + EMin=100, + EMax=80e3, + nBins=66, + mass=1.6726219e-27, # default: mp + weight='flux', + restart=True): + + EkinBinEdges = np.logspace(np.log10(EMin),np.log10(EMax),nBins) + vlsvReader = pytools.vlsvfile.VlsvReader(vlsvReader) + # check if velocity space exists in this cell + if not restart and vlsvReader.read_variable('fSaved',cid) != 1.0: + return (False,np.zeros(nBins), EkinBinEdges) + if vlsvReader.check_variable('MinValue') == True: + fMin = vlsvReader.read_variable('MinValue',cid) + #print('Cell ' + str(cid).zfill(9)) + velcells = vlsvReader.read_velocity_cells(cid, population) + V = vlsvReader.get_velocity_cell_coordinates(list(velcells.keys()), pop=population) + V2 = np.sum(np.square(V),1) + JtoeV = 1/1.60217662e-19 + Ekin = 0.5*mass*V2*JtoeV + f = list(zip(*velcells.items())) + + + # check that velocity space has cells - still return a zero histogram in the same shape + if(len(f) > 0): + f = np.asarray(f[1]) + else: + return (False,np.zeros(nBins), EkinBinEdges) + ii_f = np.where(np.logical_and(f >= fMin, Ekin > 0)) + if len(ii_f) < 1: + return (False,np.zeros(nBins), EkinBinEdges) + f = f[ii_f] + Ekin = Ekin[ii_f] + V2 = V2[ii_f] + + dV3 = np.prod(vlsvReader.get_velocity_mesh_dv(population)) + # normalization + if(weight == 'flux'): + fw = f * np.sqrt(V2) * dV3 / (4*np.pi * Ekin) # use particle flux as weighting + units = "1/(s m^2 sr eV)" + latexunits = '$\mathrm{s}^-1\,\mathrm{m}^{-2}\,\mathrm{sr}^{-1}\,\mathrm{eV}^{-1}$' + latex='$\mathcal{F}(\vec{r},E)$' + else: + fw = f * dV3 / Ekin + units = "1/(m^3 eV)" + latexunits = '$\mathrm{m}^{-3}\,\mathrm{eV}^{-1}$' + latex='$f(\vec{r},E)$' + weight = 'particles' + + #Ekin[Ekin < min(EkinBinEdges)] = min(EkinBinEdges) + #Ekin[Ekin > max(EkinBinEdges)] = max(EkinBinEdges) + + # compute histogram + (nhist,edges) = np.histogram(Ekin,bins=EkinBinEdges,weights=fw,normed=0) + + dE = EkinBinEdges[1:] - EkinBinEdges[0:-1] + #nhist = np.divide(nhist,dE) + vari = pytools.calculations.VariableInfo(nhist, + name="Omnidirectional differential energy spectrum "+population+' ('+weight+')', + units=units, + latex=latex, + latexunits=latexunits) + return (True,vari,edges) + + +# Function to reduce the velocity space in a spatial cell to a distribution in velocity along a vector +def get_spectrum_alongaxis_vel(vlsvReader, + cid, + population="proton", + vector=None, + vectorVar="vg_b_vol", + fMin=1e-21, + VMin=-2e6, + VMax=2e6, + nBins=200, + restart=True): + + vlsvReader = pytools.vlsvfile.VlsvReader(vlsvReader) + + if vectorVar is not None and vector is None: + vector=vlsvReader.read_variable(vectorVar, cid) + vector = vector/np.linalg.norm(vector) + VBinEdges = np.linspace(VMin,VMax,nBins) + # check if velocity space exists in this cell + if not restart and vlsvReader.read_variable('fSaved',cid) != 1.0: + return (False,np.zeros(nBins), VBinEdges) + if vlsvReader.check_variable('MinValue') == True: + fMin = vlsvReader.read_variable('MinValue',cid) + #print('Cell ' + str(cid).zfill(9)) + velcells = vlsvReader.read_velocity_cells(cid, population) + V = vlsvReader.get_velocity_cell_coordinates(list(velcells.keys()), pop=population) + V2 = np.sum(np.square(V),1) + Vproj = np.dot(V,vector) + f = list(velcells.values()) + + + # check that velocity space has cells - still return a zero histogram in the same shape + if(len(f) > 0): + f = np.asarray(f) + else: + return (False,np.zeros(nBins), VBinEdges) + ii_f = np.where(f >= fMin) + #ii_f = np.where(f >= fMin or True) + if len(ii_f) < 1: + return (False,np.zeros(nBins), VBinEdges) + #f = f[ii_f] + #Vproj = Vproj[ii_f] + #V2 = V2[ii_f] + + Vproj[Vproj < min(VBinEdges)] = min(VBinEdges) + Vproj[Vproj > max(VBinEdges)] = max(VBinEdges) + dV3 = np.prod(vlsvReader.get_velocity_mesh_dv(population)) + # normalization + units = "s/m^4" + latexunits = '$\mathrm{s}\mathrm{m}^{-4}$' + latex='$f(\vec{r})\,\Deltav^-1\$' + weight = 'particles' + fw = f*dV3 + + (nhist,edges) = np.histogram(Vproj,bins=VBinEdges,weights=fw,normed=0) + # normalization + dv = VBinEdges[1:] - VBinEdges[0:-1] + nhist = np.divide(nhist,dv) + + vari = pytools.calculations.VariableInfo(nhist, + name="Parallel distribution of "+population+' ('+weight+')', + units=units, + latex=latex, + latexunits=latexunits) + + return (True,vari,edges) + diff --git a/pyCalculations/variable.py b/pyCalculations/variable.py index 6694a8b0..3312f349 100644 --- a/pyCalculations/variable.py +++ b/pyCalculations/variable.py @@ -23,6 +23,8 @@ # This file has a class "Variable" that holds all the important data for variables e.g. variable's name, the units and the data on the variable import numpy as np +#import pytools +from numbers import Number class VariableInfo: ''' A class/struct for holding variable info. This includes the variable data in array form, the name and the units. @@ -51,6 +53,7 @@ def __init__(self, data_array, name="", units="", latex="", latexunits=""): self.units = units self.latex = latex self.latexunits = latexunits + self.scaleDict = {} def __repr__(self): return "VariableInfo(Name: \'" + str(self.name) + "\' Units: \'" + str(self.units) + "\')" @@ -68,7 +71,6 @@ def get_variable(self, index ): E.g. if the data is an array of 3d vectors, get_variable(0) would return the variable with data[:,0] as the data ''' - import numpy as np if len(self.data) <= 0: print("BAD DATA LENGTH") return [] @@ -78,12 +80,139 @@ def get_variable(self, index ): return VariableInfo( self.data[:,index], self.name, self.units, self.latex, self.latexunits ) + def get_scaling_metadata(self, env='EarthSpace', manualDict=None, vscale=None): + ''' Return scaling metadata + + :param env: A string to choose the scaling dictionary [default: EarthSpace] + :param manualDict: a dictionary of {units : {scalingparams}}; used to update the included dictionary + :param vscale: float, factor to scale the variable with + :returns: (norming factor, scaledUnits, scaledLatexUnits) + + .. note:: + + ''' + + # + if env=='EarthSpace': + self.scaleDict = { + 's' : {'defaultScale':1, + 1e6: {'scaledUnits':'us', 'scaledLatexUnit':r'$\mu\mathrm{s}$'}, + 1e3: {'scaledUnits':'ms', 'scaledLatexUnit':r'$\mathrm{ms}$'} + }, + 'T' : {'defaultScale':1e9, + 1e9: {'scaledUnits':'nT', 'scaledLatexUnit':r'$\mathrm{nT}$'} + }, + 'K' : {'defaultScale':1e-6, + 1e-6:{'scaledUnits':'MK', 'scaledLatexUnit':r'$\mathrm{MK}$'} + }, + 'Pa' : {'defaultScale':1e9, + 1e9:{'scaledUnits':'nPa', 'scaledLatexUnit':r'$\mathrm{nPa}$'} + }, + '1/m^3' : {'defaultScale':1e-6, + 1e-6:{'scaledUnits':'1/cm^3', 'scaledLatexUnit':r'$\mathrm{cm}^{-3}$'} + }, + '1/m3' : {'defaultScale':1e-6, + 1e-6:{'scaledUnits':'1/cm^3', 'scaledLatexUnit':r'$\mathrm{cm}^{-3}$'} + }, + 'm/s' : {'defaultScale':1e-3, + 1e-3:{'scaledUnits':'km/s', 'scaledLatexUnit':r'$\mathrm{km}\,\mathrm{s}^{-1}$'} + }, + 'V/m' : {'defaultScale':1e3, + 1e3:{'scaledUnits':'mV/m', 'scaledLatexUnit':r'$\mathrm{mV}\,\mathrm{m}^{-1}$'} + }, + 'eV/cm3' : {'defaultScale':1e-3, + 1e-3:{'scaledUnits':'keV/cm^3', 'scaledLatexUnit':r'$\mathrm{keV}\,\mathrm{cm}^{-3}$'} + }, + 'eV/cm^3': {'defaultScale':1e-3, + 1e-3:{'scaledUnits':'keV/cm^3', 'scaledLatexUnit':r'$\mathrm{keV}\,\mathrm{cm}^{-3}$'} + }, + } + + else: + self.scaleDict = {} + if manualDict is not None: + self.scaleDict.update(manualDict) + unitScale = 1.0 + scaledUnits = self.units + scaledLatexUnits = self.latexunits + + if self.units != '': + dictKey = self.units + try: + udict = self.scaleDict[dictKey] + except: + print('No entry in specialist dict for unit "' + self.units+ '"') + return 1.0, self.units, self.latexunits + if vscale is None: + try: + unitScale = udict['defaultScale'] + except: + print('No vscale or defaultScale in specialist dict for unit "' + self.units +'"') + return 1.0, self.units, self.latexunits + elif np.isclose(vscale, 1.0): + return 1.0, self.units, self.latexunits + else: + unitScale = vscale + + if not any([np.isclose(unitScale, tryScale) for tryScale in udict.keys() if isinstance(tryScale, Number)]): + # + return vscale, self.units+" x{vscale:e}".format(vscale=vscale), self.latexunits+r"{\times}"+pytools.plot.cbfmtsci(vscale,None) + try: + #above guarantees the list comprehension does not give an empty list + unitScale = [scale for scale in udict.keys() if isinstance(scale, Number) and np.isclose(scale,unitScale)][0] + scaledUnits = udict[unitScale]['scaledUnits'] + except KeyError: + print('Missing scaledUnits in specialist dict for' + self.units + ' for unitScale='+str(unitScale)) + return 1.0, self.units, self.latexunits + try: + scaledLatexUnits = udict[unitScale]['scaledLatexUnit'] + except: + print('Missing scaledLatexUnits in specialist dict for ' + self.units+ ' for unitScale='+str(unitScale)) + return 1.0, self.units, self.latexunits + else: + if vscale is None or np.isclose(vscale, 1.0): + return 1.0, self.units, self.latexunits + else: + return vscale, self.units+"x{vscale:e}".format(vscale=vscale), self.latexunits+r"{\times}"+pytools.plot.cbfmtsci(vscale,None) + + return unitScale, scaledUnits, scaledLatexUnits + + # A utility to get variableinfo with corresponding units for simple plotting. Add "canonical" scalings as + # necessary, for default/other environments. + def get_scaled_var(self, data=None, env='EarthSpace', manualDict=None, vscale=None): + ''' Automatically scales the variableinfo data and adjusts the units correspondingly with the + default dictionaries. + + :param data: in case you wish to provide new data array (why, though?) + :param env: A string to choose the scaling dictionary [default: EarthSpace] + :param manualDict: a dictionary of {units : {scalingparams}}; used to update the included dictionary + :returns: self, with scaled units with pre-formatted units included in the varinfo. + + .. note:: + + ''' + + if data is None: + data = self.data + else: + self.data = data + + unitScale, scaledUnits, scaledLatexUnits = self.get_scaling_metadata(env, manualDict, vscale) + if unitScale == 1: # no change, optimize out the calculation + return self + + self.data = data*unitScale + self.units = scaledUnits + self.latexunits = scaledLatexUnits + return self + + def get_data( variable ): ''' Function to use when not sure if variable is in raw form ( simply a list with data ), or a VariableInfo instance - :param variable: The variable as a VariableInfo instance or a list - :returns: data of the variable + :param variable: The variable as a VariableInfo instance or a list + :returns: data of the variable ''' if isinstance(variable, VariableInfo): return variable.data @@ -93,8 +222,8 @@ def get_data( variable ): def get_name( variable ): ''' Function to use when not sure if variable is in raw form ( simply a list with data ), or a VariableInfo instance - :param variable: The variable as a VariableInfo instance or a list - :returns: the name of the variable or \"\" if not a VariableInfo instance + :param variable: The variable as a VariableInfo instance or a list + :returns: the name of the variable or \"\" if not a VariableInfo instance ''' if isinstance(variable, VariableInfo): return variable.name @@ -104,8 +233,8 @@ def get_name( variable ): def get_units( variable ): ''' Function to use when not sure if variable is in raw form ( simply a list with data ), or a VariableInfo instance - :param variable: The variable as a VariableInfo instance or a list - :returns: the units of the variable or \"\" if not a VariableInfo instance + :param variable: The variable as a VariableInfo instance or a list + :returns: the units of the variable or \"\" if not a VariableInfo instance ''' if isinstance(variable, VariableInfo): return variable.units @@ -115,8 +244,8 @@ def get_units( variable ): def get_latex( variable ): ''' Function to use when not sure if variable is in raw form ( simply a list with data ), or a VariableInfo instance - :param variable: The variable as a VariableInfo instance or a list - :returns: the variable in LaTeX format or \"\" if not a VariableInfo instance + :param variable: The variable as a VariableInfo instance or a list + :returns: the variable in LaTeX format or \"\" if not a VariableInfo instance ''' if isinstance(variable, VariableInfo): return variable.latex @@ -126,8 +255,8 @@ def get_latex( variable ): def get_latexunits( variable ): ''' Function to use when not sure if variable is in raw form ( simply a list with data ), or a VariableInfo instance - :param variable: The variable as a VariableInfo instance or a list - :returns: the units of the variable in LaTeX format or \"\" if not a VariableInfo instance + :param variable: The variable as a VariableInfo instance or a list + :returns: the units of the variable in LaTeX format or \"\" if not a VariableInfo instance ''' if isinstance(variable, VariableInfo): return variable.latexunits diff --git a/pyPlots/plot.py b/pyPlots/plot.py index 698758ea..afe4458c 100644 --- a/pyPlots/plot.py +++ b/pyPlots/plot.py @@ -182,34 +182,3 @@ def textbfstring(string): return r'\textbf{'+string+'}' # LaTex output off return string - -# Helper routine for allowing specialist units for known vscale and unit combinations -def scaleunits(datamap_info, vscale): - # Check if vscale is in use? - if np.isclose(vscale,1.): - return datamap_info.latexunits - # Check for known variables? - if datamap_info.units=="s" and np.isclose(vscale,1.e6): - return r"\mu"+rmstring("s") - if datamap_info.units=="s" and np.isclose(vscale,1.e3): - return rmstring("ms") - if datamap_info.units=="T" and np.isclose(vscale,1.e9): - return rmstring("nT") - if datamap_info.units=="K" and np.isclose(vscale,1.e-6): - return rmstring("MK") - if datamap_info.units=="Pa" and np.isclose(vscale,1.e9): - return rmstring("nPa") - if datamap_info.units=="1/m3" and np.isclose(vscale,1.e-6): - return rmstring("cm")+"^{-3}" - if datamap_info.units=="1/m^3" and np.isclose(vscale,1.e-6): - return rmstring("cm")+"^{-3}" - if datamap_info.units=="m/s" and np.isclose(vscale,1.e-3): - return rmstring("km")+"\,"+rmstring("s")+"^{-1}" - if datamap_info.units=="V/m" and np.isclose(vscale,1.e3): - return rmstring("mV")+"\,"+rmstring("m")+"^{-1}" - if datamap_info.units=="eV/cm3" and np.isclose(vscale,1.e-3): - return rmstring("keV")+"\,"+rmstring("cm")+"^{-3}" - if datamap_info.units=="eV/cm^3" and np.isclose(vscale,1.e-3): - return rmstring("keV")+"\,"+rmstring("cm")+"^{-3}" - # fallthrough - return datamap_info.latexunits+r"{\times}"+cbfmtsci(vscale,None) diff --git a/pyPlots/plot_colormap.py b/pyPlots/plot_colormap.py index 819627aa..61863feb 100644 --- a/pyPlots/plot_colormap.py +++ b/pyPlots/plot_colormap.py @@ -61,7 +61,7 @@ def plot_colormap(filename=None, absolute=False, symmetric=False, pass_vars=None, pass_times=None, pass_full=False, - fluxfile=None, fluxdir=None, + fluxfile=None, fluxdir=None, flux_levels=None, fluxthick=1.0, fluxlines=1, fsaved=None, Earth=None, @@ -151,7 +151,7 @@ def plot_colormap(filename=None, :kword vscale: Scale all values with this before plotting. Useful for going from e.g. m^-3 to cm^-3 or from tesla to nanotesla. Guesses correct units for colourbar for some known - variables. + variables. Set to None to search for a default scaling settings. :kword absolute: Plot the absolute of the evaluated variable :kword pass_vars: Optional list of map names to pass to the external/expression functions @@ -169,6 +169,8 @@ def plot_colormap(filename=None, :kword pass_full: Set to anything but None in order to pass the full arrays instead of a zoomed-in section :kword fluxfile: Filename to plot fluxfunction from + :kword flux_levels: A list of flux function values to plot as the contours (default: None, a set of constant + intervals: np.linspace(-10,10,fluxlines*60)) :kword fluxdir: Directory in which fluxfunction files can be found :kword fluxthick: Scale fluxfunction line thickness :kword fluxlines: Relative density of fluxfunction contours @@ -478,7 +480,7 @@ def exprMA_cust(exprmaps, requestvariables=False): cb_title_use = datamap_info.latex datamap_unit = datamap_info.latexunits # Check if vscale results in standard unit - datamap_unit = pt.plot.scaleunits(datamap_info, vscale) + vscale, datamap_unit_plain, datamap_unit = datamap_info.get_scaling_metadata(vscale=vscale) # Add unit to colorbar title if datamap_unit: @@ -912,7 +914,7 @@ def exprMA_cust(exprmaps, requestvariables=False): norm=norm, interpolation=imshowinterp, origin='lower', - extent=(boxcoords[0], boxcoords[1], boxcoords[2], boxcoords[3]), + extent=(np.min(XmeshPass), np.max(XmeshPass), np.min(YmeshPass), np.max(YmeshPass)) ) @@ -929,8 +931,6 @@ def exprMA_cust(exprmaps, requestvariables=False): ax1.spines[axis].set_linewidth(thick) ax1.xaxis.set_tick_params(width=thick,length=3) ax1.yaxis.set_tick_params(width=thick,length=3) - #ax1.xaxis.set_tick_params(which='minor',width=3,length=5) - #ax1.yaxis.set_tick_params(which='minor',width=3,length=5) if not noxlabels: xlabelstr = pt.plot.mathmode(pt.plot.bfstring('x\,['+axisunitstr+']')) @@ -993,7 +993,10 @@ def exprMA_cust(exprmaps, requestvariables=False): flux_function = np.ma.array(flux_function, mask=XYmask) # The flux level contours must be fixed instead of scaled based on min/max values in order # to properly account for flux freeze-in and advection with plasma - flux_levels = np.linspace(-10,10,fluxlines*60) + if flux_levels is None: + flux_levels = np.linspace(-10,10,fluxlines*60) + else: + pass #This was given, do nothing fluxcont = ax1.contour(XmeshCentres,YmeshCentres,flux_function,flux_levels,colors='k',linestyles='solid',linewidths=0.5*fluxthick,zorder=2) # add fSaved identifiers diff --git a/pyPlots/plot_colormap3dslice.py b/pyPlots/plot_colormap3dslice.py index 7e273271..5c1cd595 100644 --- a/pyPlots/plot_colormap3dslice.py +++ b/pyPlots/plot_colormap3dslice.py @@ -149,7 +149,7 @@ def plot_colormap3dslice(filename=None, :kword vscale: Scale all values with this before plotting. Useful for going from e.g. m^-3 to cm^-3 or from tesla to nanotesla. Guesses correct units for colourbar for some known - variables. + variables. Set to None to seek for a default scaling. :kword absolute: Plot the absolute of the evaluated variable :kword pass_vars: Optional list of map names to pass to the external/expression functions @@ -547,8 +547,8 @@ def exprMA_cust(exprmaps, requestvariables=False): cb_title_use = datamap_info.latex datamap_unit = datamap_info.latexunits # Check if vscale results in standard unit - datamap_unit = pt.plot.scaleunits(datamap_info, vscale) - + vscale, datamap_unit_plain, datamap_unit = datamap_info.get_scaling_metadata(vscale=vscale) + # Add unit to colorbar title if datamap_unit: cb_title_use = cb_title_use + "\,["+datamap_unit+"]" diff --git a/pyPlots/plot_threeslice.py b/pyPlots/plot_threeslice.py index 899df1dc..020a0ac4 100644 --- a/pyPlots/plot_threeslice.py +++ b/pyPlots/plot_threeslice.py @@ -668,7 +668,7 @@ def plot_threeslice(filename=None, :kword vscale: Scale all values with this before plotting. Useful for going from e.g. m^-3 to cm^-3 or from tesla to nanotesla. Guesses correct units for colourbar for some known - variables. + variables. Set to None to search for a default scaling. :kword viewangle: Azimuth and elevation angles giving the point of view on the 3D axes, in degrees. (default=(-60.,30.); corresponds to dayside, morningside, northern hemisphere) @@ -1041,7 +1041,7 @@ def plot_threeslice(filename=None, cb_title_use = datamap_info.latex datamap_unit = datamap_info.latexunits # Check if vscale results in standard unit - datamap_unit = pt.plot.scaleunits(datamap_info, vscale) + vscale, datamap_unit_plain, datamap_unit = datamap_info.get_scaling_metadata(vscale=vscale) # Add unit to colorbar title if datamap_unit: diff --git a/pyPlots/plot_vdf.py b/pyPlots/plot_vdf.py index 3bf36cea..88214df6 100644 --- a/pyPlots/plot_vdf.py +++ b/pyPlots/plot_vdf.py @@ -22,6 +22,7 @@ # import matplotlib +import warnings import pytools as pt import numpy as np import matplotlib.pyplot as plt @@ -40,6 +41,9 @@ from rotation import rotateVectorToVector,rotateVectorToVector_X +# Resample reducer flag - set to True to perform resampling in log-scaled values +# Retained for reference, if the large differences in VDF values come back to haunt +logspaceResample = False # find nearest spatial cell with vspace to cid def getNearestCellWithVspace(vlsvReader,cid): @@ -65,10 +69,10 @@ def verifyCellWithVspace(vlsvReader,cid): return found # create a 2-dimensional histogram -def doHistogram(f,VX,VY,Voutofslice,vxBinEdges,vyBinEdges,vthick,wflux=None): +def doHistogram(f,VX,VY,Voutofslice,vxBinEdges,vyBinEdges,vthick,reducer="integrate", wflux=None, initial_dV=1.0): # Flux weighting? if wflux is not None: - fw = f*np.linalg.norm([VX,VY,Voutofslice]) # use particle flux as weighting in the histogram + fw = f*np.linalg.norm([VX,VY,Voutofslice], axis=0)/(4*np.pi) # use particle flux as weighting in the histogram else: fw = f # use particle phase-space density as weighting in the histogram @@ -83,14 +87,19 @@ def doHistogram(f,VX,VY,Voutofslice,vxBinEdges,vyBinEdges,vthick,wflux=None): # Gather histogram of values (nVhist,VXEdges,VYEdges) = np.histogram2d(VX[tuple(indexes)],VY[tuple(indexes)],bins=(vxBinEdges,vyBinEdges),weights=fw[tuple(indexes)],normed=0) - # Gather histogram of how many cells were summed for the histogram - (Chist,VXEdges,VYEdges) = np.histogram2d(VX[tuple(indexes)],VY[tuple(indexes)],bins=(vxBinEdges,vyBinEdges),normed=0) - # Correct for summing multiple cells into one histogram output cell - nonzero = np.where(Chist != 0) - nVhist[nonzero] = np.divide(nVhist[nonzero],Chist[nonzero]) + # Correct for summing multiple cells into one histogram output cell with the averaging reducer + if reducer == "average": + # Gather histogram of how many cells were summed for the histogram + (Chist,VXEdges,VYEdges) = np.histogram2d(VX[tuple(indexes)],VY[tuple(indexes)],bins=(vxBinEdges,vyBinEdges),normed=0) + nonzero = np.where(Chist != 0) + nonzero = Chist > 0 + nVhist[nonzero] = np.divide(nVhist[nonzero],Chist[nonzero]) + + dV = np.abs(vxBinEdges[-1] - vxBinEdges[-2]) # assumes constant bin size + if vthick==0: - # slickethick=0, perform a projection. This is done by taking averages for each sampled stack of cells + # slickethick=0, perform a projection. This is done by taking averages for each sampled stack of cells (above) # (in order to deal with rotated sampling issues) and then rescaling the resultant 2D VDF with the unsampled # VDF in order to get the particle counts to agree. Here we gather the total particle counts for the # unprojected VDF. @@ -98,23 +107,126 @@ def doHistogram(f,VX,VY,Voutofslice,vxBinEdges,vyBinEdges,vthick,wflux=None): weights_total_proj = nVhist.sum() # Now rescale the projection back up in order to retain particle counts. nVhist = nVhist * weights_total_all/weights_total_proj + else: + pass # Please note that the histogram does not follow the Cartesian convention where x values are on the abscissa # and y values on the ordinate axis. Rather, x is histogrammed along the first dimension of the array (vertical), # and y along the second dimension of the array (horizontal). This ensures compatibility with histogramdd. nVhist = nVhist.transpose() - - # Flux weighting - if wflux is not None: - dV = np.abs(vxBinEdges[-1] - vxBinEdges[-2]) # assumes constant bin size - nVhist = np.divide(nVhist,(dV*4*np.pi)) # normalization + + # Rotation contributions are already normalized, this is just to rescaling to correct units + # nb: maybe not great for finite slices...? + if reducer == "integrate": + nVhist = nVhist*initial_dV**3/dV**2 # normalization return (nVhist,VXEdges,VYEdges) +def resampleReducer(V,f, inputcellsize, setThreshold, normvect, normvectX, slicetype, slicethick, reducer="integrate", wflux=None): + + if wflux is not None: + fw = f*np.linalg.norm(V, axis=-1)/(4*np.pi) # use particle flux as weighting in the histogram + else: + fw = f # use particle phase-space density as weighting in the histogram + + NX = np.array(normvect)/np.linalg.norm(normvect) + NY = np.array(normvectX)/np.linalg.norm(normvectX) + NZ = np.cross(NX,NY) + if slicetype=="Bpara": + R = np.stack((NX, NY, NZ)).T + elif slicetype=="Bpara1": + R = np.stack((NX, NZ, NY)).T + elif slicetype=="Bperp": + R = np.stack((NY, NX, NZ)).T + elif slicetype=="vecperp": + R = np.stack((NY, NX, NZ)).T + #print(R) + vmins = np.amin(V,axis=0) + vmaxs = np.amax(V,axis=0) + + # let's see where the bounding box corners end up and pad the array accordingly + vexts =[[vmins[0], vmins[1], vmins[2]], + [vmins[0], vmins[1], vmaxs[2]], + [vmins[0], vmaxs[1], vmaxs[2]], + [vmaxs[0], vmaxs[1], vmaxs[2]], + [vmaxs[0], vmaxs[1], vmins[2]], + [vmaxs[0], vmins[1], vmins[2]]] + + vextsR = np.array([np.matmul(R, v) for v in vexts]) + vexts = np.array(vexts) # also needs the initial extent, so a thin box isn't rotated out of the initial box + + vminsR,vmaxsR = np.minimum(np.amin(vextsR, axis=0),np.amin(vexts, axis=0)), np.maximum(np.amax(vextsR, axis=0),np.amax(vexts, axis=0)) + + leftpads = (np.abs(vmins/inputcellsize - vminsR/inputcellsize)).astype(int) + rightpads = (np.abs(vmaxs/inputcellsize - vmaxsR/inputcellsize)).astype(int) + + szs = np.trunc((vmaxs-vmins)/inputcellsize+1).astype(int) + + pivot = -vmins/inputcellsize+leftpads #assumes a previously centered V-space wrt. rotations, just need this in cellwidth units + + #Form a dense array of the initial data + basearray, edges = np.histogramdd(V, bins=tuple(szs), + range=list(zip(vmins-0.5*inputcellsize,vmaxs+0.5*inputcellsize)), + density=False,weights=fw*inputcellsize**3) + + newedges = [np.linspace(vminsR[d],vmaxsR[d], num=(szs+leftpads+rightpads)[d]+1) for d in [0,1,2]] + #need to expand the array a bit to not chop off anything + basearray = np.pad(basearray,tuple(zip(leftpads,rightpads))) + + + nonzero = basearray > 0 + + + if logspaceResample: + newarray = np.ones_like(basearray) + newarray = newarray*np.log(sys.float_info.min) + np.log(basearray, where=nonzero, out=newarray) + else: + newarray = np.zeros_like(basearray) + newarray[nonzero] = basearray[nonzero] + + newarray = scipy.ndimage.affine_transform(newarray,R, + mode='constant',cval=np.log(setThreshold/10), + order=1, + offset=pivot-np.matmul(R,pivot), + ) + if logspaceResample: + newarray = np.exp(newarray) + newarray[np.logical_not(np.isfinite(newarray))] = 0 + + newarray[newarray 0.1*min(min(vxsize,vysize),vzsize): + warnings.warn("You seem to be analysing a thick slice of the VDF.\nPlease check this is your intent - slicethick is amount of cells, not SI velocity space length!") + if slicethick > 0 and reducer=="average": + warnings.warn("You seem to be averaging across some width of the VDF. Are you sure you don't want to integrate instead?") + # check that velocity space has cells if(len(velcellslist) <= 0): return (False,0,0,0) @@ -220,58 +343,67 @@ def vSpaceReducer(vlsvReader, cid, slicetype, normvect, VXBins, VYBins, pop="pro elif slicetype=="vecperp": # Find velocity components in rotated frame where normavect is outofslice and optional # normvectX is in VX direction - N = np.array(normvect)/np.sqrt(normvect[0]**2 + normvect[1]**2 + normvect[2]**2) - Vrot = rotateVectorToVector(V,N) - if normvectX is not None: - NX = np.array(normvectX)/np.sqrt(normvectX[0]**2 + normvectX[1]**2 + normvectX[2]**2) - NXrot = rotateVectorToVector(NX,N) - Vrot2 = rotateVectorToVector_X(Vrot,NXrot) - Vrot = Vrot2 - VX = Vrot[:,0] - VY = Vrot[:,1] - Voutofslice = Vrot[:,2] + if resampler is False: + N = np.array(normvect)/np.sqrt(normvect[0]**2 + normvect[1]**2 + normvect[2]**2) + Vrot = rotateVectorToVector(V,N) + if normvectX is not None: + NX = np.array(normvectX)/np.sqrt(normvectX[0]**2 + normvectX[1]**2 + normvectX[2]**2) + NXrot = rotateVectorToVector(NX,N) + Vrot2 = rotateVectorToVector_X(Vrot,NXrot) + Vrot = Vrot2 + VX = Vrot[:,0] + VY = Vrot[:,1] + Voutofslice = Vrot[:,2] + else: + if normvectX is None: + warnings.warn("Please provide a normvectX for the resampler!") + return(False,0,0,0) + return resampleReducer(V,f, inputcellsize,setThreshold, normvect, normvectX, slicetype, slicethick, reducer=reducer, wflux=wflux) + 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 - # intuition says this should be -Vrot2[:,0], but testing of profiles across the VDF prove otherwise - 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") - + if resampler is False: + # 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 + # intuition says this should be -Vrot2[:,0], but testing of profiles across the VDF prove otherwise + 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: + return resampleReducer(V,f, inputcellsize, setThreshold, normvect, normvectX, slicetype, slicethick, reducer=reducer, wflux=wflux) else: print("Error finding rotation of v-space!") return (False,0,0,0) # create 2-dimensional histogram of velocity components perpendicular to slice-normal vector - (binsXY,edgesX,edgesY) = doHistogram(f,VX,VY,Voutofslice,VXBins,VYBins, slicethick, wflux=wflux) + (binsXY,edgesX,edgesY) = doHistogram(f,VX,VY,Voutofslice,VXBins,VYBins, slicethick, reducer=reducer, wflux=wflux, initial_dV=inputcellsize) return (True,binsXY,edgesX,edgesY) @@ -282,26 +414,27 @@ def plot_vdf(filename=None, coordinates=None, coordre=None, outputdir=None, outputfile=None, nooverwrite=None, - draw=None,axisunit=None,title=None, cbtitle=None, + draw=None,axisunit=None,axiskmps=None,title=None, cbtitle=None, tickinterval=None, colormap=None, box=None, nocb=None, internalcb=None, run=None, thick=1.0, wmark=None, wmarkb=None, - fmin=None, fmax=None, slicethick=None, cellsize=None, + fmin=None, fmax=None, slicethick=None, reducer='integrate', resampler=True, + cellsize=None, xy=None, xz=None, yz=None, normal=None, normalx=None, bpara=None, bpara1=None, bperp=None, coordswap=None, - bvector=None, + bvector=None,bvectorscale=0.2, cbulk=None, center=None, wflux=None, setThreshold=None, - noborder=None, scale=1.0, + noborder=None, scale=1.0, scale_text=8.0, scale_title=10.0,scale_cb=5.0,scale_label=12.0, biglabel=None, biglabloc=None, noxlabels=None, noylabels=None, axes=None, cbaxes=None, contours=None ): - ''' Plots a coloured plot with axes and a colour bar. + ''' Plots a coloured 2D plot of a VDF (a slice of given thickness or fully projected) 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 @@ -332,6 +465,7 @@ def plot_vdf(filename=None, :kword box: extents of plotted velocity grid as [x0,x1,y0,y1] (in m/s) :kword axisunit: Plot v-axes using 10^{axisunit} m/s (default: km/s) + :kword axiskmps: Plot v-axes using 10^{axiskmps} km/s (default: km/s, when the kword has a value) :kword tickinterval: Interval at which to have ticks on axes :kword xy: Perform slice in x-y-direction @@ -347,6 +481,7 @@ def plot_vdf(filename=None, :kword coordswap: Swap the parallel and perpendicular coordinates :kword bvector: Plot a magnetic field vector projection in-plane + :kword bvectorscale: Scale of bvector (default: 0.2 in units of axis lengths) :kword cbulk: Center plot on position of total bulk velocity (or if not available, bulk velocity for this population) @@ -356,6 +491,9 @@ def plot_vdf(filename=None, :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). This can be set to zero in order to project the whole VDF to a plane. + :kword reducer: How to reduce to 2D - default 'integrate' for LOS integration + and reduced units, 'average' for old (slightly questionable) behaviour + :kword resampler: Resample onto a regular grid? Default: yes, use False to disable. :kword cellsize: Plotting grid cell size as multiplier of input cell size (default: 1 or minimum for good coverage) :kword setThreshold: Use given setThreshold value instead of EffectiveSparsityThreshold or MinValue value read from file Useful if EffectiveSparsityThreshold wasn't saved, or user wants to draw buffer cells @@ -379,7 +517,11 @@ def plot_vdf(filename=None, :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 scale: Scale text size everywhere (default=1.0) + :kword scale_text: Most text additional scale factor (default=8.0) + :kword scale_title: Title additional scale factor (default=10.0) + :kword scale_cb: Colour bar text additional scale factor (default=5.0) + :kword scale_label: Big label text additional scale factor (default=12.0) :kword thick: line and axis thickness, default=1.0 @@ -429,10 +571,10 @@ def plot_vdf(filename=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=12*scale # Big label + fontsize=scale_text*scale # Most text + fontsize2=scale_title*scale # Time title + fontsize3=scale_cb*scale # Colour bar ticks + fontsize4=scale_label*scale # Big label # Plot title with time timeval=vlsvReader.read_parameter("time") @@ -450,6 +592,19 @@ def plot_vdf(filename=None, else: plot_title = title + if reducer == "average": + print("V-space reduction via averages") + warnings.warn('Average-reduction is kept for backward-compatibility for now; consider using "integrate"!') + pass + elif reducer == "integrate": + print("V-space reduction via integration") + pass + else: + raise ValueError("Unknown reducer ("+reducer+'), accepted values are "average", "integrate"') + + if wflux is not None: + warnings.warn("Does flux weighting make sense? Tread carefully.") + if draw is None and axes is None: # step, used for file name if step is not None: @@ -534,14 +689,16 @@ def plot_vdf(filename=None, vxsize = int(vxsize) vysize = int(vysize) vzsize = int(vzsize) - [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 + # Account for WID3 cells per block + widval=4 #default WID=4 + if vlsvReader.check_parameter("velocity_block_width"): + widval = vlsvReader.read_parameter("velocity_block_width") + vxsize = widval*vxsize + vysize = widval*vysize + vzsize = widval*vzsize Re = 6.371e+6 # Earth radius in m # unit of velocity @@ -553,6 +710,12 @@ def plot_vdf(filename=None, velUnitStr = r'[m s$^{-1}$]' else: velUnitStr = r'[$10^{'+str(int(axisunit))+'}$ m s$^{-1}$]' + if axiskmps is not None: + velUnit = np.power(10,int(axiskmps)+3) + if np.isclose(axiskmps,0): + velUnitStr = r'[km s$^{-1}$]' + else: + velUnitStr = r'[$10^{'+str(int(axiskmps))+'}$ km s$^{-1}$]' # Select plotting back-end based on on-screen plotting or direct to file without requiring x-windowing if axes is None: # If axes are provided, leave backend as-is. @@ -859,7 +1022,7 @@ def plot_vdf(filename=None, # Read velocity data into histogram (checkOk,binsXY,edgesX,edgesY) = vSpaceReducer(vlsvReader,cellid,slicetype,normvect,VXBins, VYBins,pop=pop, - slicethick=slicethick, wflux=wflux, + slicethick=slicethick, reducer=reducer, resampler=resampler, wflux=wflux, center=center,setThreshold=setThreshold,normvectX=normvectX) # Check that data is ok and not empty @@ -900,9 +1063,11 @@ def plot_vdf(filename=None, fmaxuse = 1e-10 # No valid values! use extreme default. print("Active f range is "+str(fminuse)+" to "+str(fmaxuse)) - norm = LogNorm(vmin=fminuse,vmax=fmaxuse) - ticks = LogLocator(base=10,subs=list(range(10))) # where to show labels + + ticks = LogLocator(base=10,subs=list(range(0,10)))#, + #numticks=max(2,np.rint(np.log10(fmaxuse/fminuse))) ) # where to show labels + # tries to force at least 2 labels if box is not None: # extents of plotted velocity grid as [x0,y0,x1,y1] xvalsrange=[box[0],box[1]] @@ -921,10 +1086,10 @@ def plot_vdf(filename=None, 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.)) ]) + xindexrange[0] = np.max([0, widval * int(np.floor((xindexrange[0]-0.5*widval)/widval)) ]) + xindexrange[1] = np.min([len(edgesX)-1, widval * int(np.ceil((xindexrange[1]+0.5*widval)/widval)) ]) + yindexrange[0] = np.max([0, widval * int((np.floor(yindexrange[0]-0.5*widval)/widval)) ]) + yindexrange[1] = np.min([len(edgesY)-1, widval * int(np.ceil((yindexrange[1]+0.5*widval)/widval)) ]) # If empty VDF: plot whole v-space if ((xindexrange==[vxsize,0]) and (yindexrange==[vysize,0])): @@ -1063,8 +1228,8 @@ def plot_vdf(filename=None, # normalize bvector = binplane/np.linalg.norm(binplane) # Length default is 1/5 of axis length - bvectormultiplier = np.amin([yvalsrange[1]-yvalsrange[0],xvalsrange[1]-xvalsrange[0]])/(velUnit*5.) - bvector *= bvectormultiplier + bvectormultiplier = np.amin([yvalsrange[1]-yvalsrange[0],xvalsrange[1]-xvalsrange[0]])/(velUnit) + bvector *= bvectormultiplier*bvectorscale # ax1.arrow(0,0,bvector[0],bvector[1],width=0.02*thick,head_width=0.1*thick, # head_length=0.2*thick,zorder=10,color='k') #ax1.arrow(origin[0],origin[1],bvector[0],bvector[1],zorder=10,color='k') @@ -1078,9 +1243,15 @@ def plot_vdf(filename=None, cb_title_use = cbtitle if cb_title_use is None: if wflux is None: - cb_title_use=r"$f(v)\,["+pt.plot.rmstring('m')+"^{-6} \,"+pt.plot.rmstring('s')+"^{3}]$" + if reducer == 'average': + cb_title_use=r"$f(v)\,["+pt.plot.rmstring('m')+"^{-6} \,"+pt.plot.rmstring('s')+"^{3}]$" + elif reducer == 'integrate': + cb_title_use=r"$f(v)\,["+pt.plot.rmstring('m')+"^{-5} \,"+pt.plot.rmstring('s')+"^{2}]$" else: - cb_title_use=r"flux $F\,["+pt.plot.rmstring('m')+"^{-2} \,"+pt.plot.rmstring('s')+"^{-1} \,"+pt.plot.rmstring('sr')+"^{-1}]$" + if reducer == 'average': + cb_title_use=r"flux $F\,["+pt.plot.rmstring('m')+"^{-2} \,"+pt.plot.rmstring('s')+"^{-1} \,"+pt.plot.rmstring('sr')+"^{-1}]$" + elif reducer == 'integrate': + cb_title_use=r"flux $F\,["+pt.plot.rmstring('m')+"^{-1} \,"+pt.plot.rmstring('s')+"^{-2} \,"+pt.plot.rmstring('sr')+"^{-1}]$" if cbaxes is not None: cax = cbaxes @@ -1178,6 +1349,7 @@ def plot_vdf(filename=None, if draw is None and axes is None: try: plt.savefig(savefigname,dpi=300, bbox_inches=bbox_inches, pad_inches=savefig_pad) + plt.close() except: print("Error with attempting to save figure due to matplotlib LaTeX integration.") print(savefigname+"\n") diff --git a/pyPlots/plot_vdf_profiles.py b/pyPlots/plot_vdf_profiles.py index f130fd1a..ef5b75df 100644 --- a/pyPlots/plot_vdf_profiles.py +++ b/pyPlots/plot_vdf_profiles.py @@ -119,10 +119,13 @@ def vSpaceReducer(vlsvReader, cid, slicetype, normvect, pop="proton", vxsize = int(vxsize) vysize = int(vysize) vzsize = int(vzsize) - # Account for 4x4x4 cells per block - vxsize = 4*vxsize - vysize = 4*vysize - vzsize = 4*vzsize + # Account for WID3 cells per block + widval=4 #default WID=4 + if vlsvReader.check_parameter("velocity_block_width"): + widval = vlsvReader.read_parameter("velocity_block_width") + vxsize = widval*vxsize + vysize = widval*vysize + vzsize = widval*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)) @@ -271,7 +274,7 @@ def plot_vdf_profiles(filename=None, axes=None ): - ''' Plots a coloured plot with axes and a colour bar. + ''' Plots vdf values along axis-aligned lines (see axis definitions). :kword filename: path to .vlsv file to use for input. Assumes a bulk file. :kword vlsvobj: Optionally provide a python vlsvfile object instead @@ -475,10 +478,13 @@ def plot_vdf_profiles(filename=None, [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 + # Account for WID3 cells per block + widval=4 #default WID=4 + if vlsvReader.check_parameter("velocity_block_width"): + widval = vlsvReader.read_parameter("velocity_block_width") + vxsize = widval*vxsize + vysize = widval*vysize + vzsize = widval*vzsize Re = 6.371e+6 # Earth radius in m # unit of velocity diff --git a/pyVlsv/reduction.py b/pyVlsv/reduction.py index 80ff5b0d..da394ed6 100644 --- a/pyVlsv/reduction.py +++ b/pyVlsv/reduction.py @@ -364,6 +364,8 @@ def ParallelVectorComponent( variables ): ''' Data reducer function for vector component parallel to the magnetic field (or another vector) ''' inputvector = variables[0] + if inputvector is None: + raise ValueError("Missing inputvector") bgvector = variables[1] if( np.ndim(inputvector)==1 ): if( np.linalg.norm(bgvector) != 0): @@ -379,6 +381,8 @@ def PerpendicularVectorComponent( variables ): ''' Data reducer function for vector component perpendicular to the magnetic field (or another vector) ''' inputvector = variables[0] + if inputvector is None: + raise ValueError("Missing inputvector") bgvector = variables[1] if( np.ndim(inputvector)==1 ): if( np.linalg.norm(bgvector) != 0): diff --git a/pyVlsv/vlsvreader.py b/pyVlsv/vlsvreader.py index 48328414..e665d75b 100644 --- a/pyVlsv/vlsvreader.py +++ b/pyVlsv/vlsvreader.py @@ -153,7 +153,7 @@ def __init__(self, file_name): pop.__vxblocks = (int)(self.read_parameter("vxblocks_ini")) pop.__vyblocks = (int)(self.read_parameter("vyblocks_ini")) pop.__vzblocks = (int)(self.read_parameter("vzblocks_ini")) - pop.__vxblock_size = 4 + pop.__vxblock_size = 4 # Old files will always have WID=4, newer files read it from bbox pop.__vyblock_size = 4 pop.__vzblock_size = 4 pop.__vxmin = self.read_parameter("vxmin") @@ -168,7 +168,7 @@ def __init__(self, file_name): pop.__dvz = ((pop.__vzmax - pop.__vzmin) / (float)(pop.__vzblocks)) / (float)(pop.__vzblock_size) else: - #no velocity space in this file, e.g., file n ot written by Vlasiator + #no velocity space in this file, e.g., file not written by Vlasiator pop.__vxblocks = 0 pop.__vyblocks = 0 pop.__vzblocks = 0 @@ -209,7 +209,7 @@ def __init__(self, file_name): pop.__dvz = ((pop.__vzmax - pop.__vzmin) / (float)(pop.__vzblocks)) / (float)(pop.__vzblock_size) self.__meshes[popname]=pop - if os.getenv('PTNONINTERACTIVE') == None: + if not os.getenv('PTNONINTERACTIVE'): print("Found population " + popname) # Precipitation energy bins @@ -1462,6 +1462,13 @@ def get_cell_neighbor(self, cellid, offset, periodic): cellid_neighbour = self.get_cellid(coord_neighbour) return cellid_neighbour + def get_WID(self): + # default WID=4 + widval=4 + if self.check_parameter("velocity_block_width"): + widval = self.read_parameter("velocity_block_width") + return widval + def get_velocity_cell_ids(self, vcellcoord, pop="proton"): ''' Returns velocity cell ids of given coordinate @@ -1471,17 +1478,20 @@ def get_velocity_cell_ids(self, vcellcoord, pop="proton"): .. seealso:: :func:`get_velocity_cell_coordinates` ''' + WID=self.get_WID() + WID2=WID*WID + WID3=WID2*WID vmin = np.array([self.__meshes[pop].__vxmin, self.__meshes[pop].__vymin, self.__meshes[pop].__vzmin]) dv = np.array([self.__meshes[pop].__dvx, self.__meshes[pop].__dvy, self.__meshes[pop].__dvz]) - block_index = np.floor((vcellcoord - vmin) / (4 * dv)) - cell_index = np.floor(np.remainder(vcellcoord - vmin, 4 * dv) / dv) + block_index = np.floor((vcellcoord - vmin) / (WID * dv)) + cell_index = np.floor(np.remainder(vcellcoord - vmin, WID * dv) / dv) vcellid = int(block_index[0]) vcellid += int(block_index[1] * self.__meshes[pop].__vxblocks) vcellid += int(block_index[2] * self.__meshes[pop].__vxblocks * self.__meshes[pop].__vyblocks) - vcellid *= 64 + vcellid *= WID3 vcellid += int(cell_index[0]) - vcellid += int(cell_index[1] * 4) - vcellid += int(cell_index[2] * 16) + vcellid += int(cell_index[1] * WID) + vcellid += int(cell_index[2] * WID2) return vcellid def get_velocity_cell_coordinates(self, vcellids, pop="proton"): @@ -1494,20 +1504,23 @@ def get_velocity_cell_coordinates(self, vcellids, pop="proton"): .. seealso:: :func:`get_cell_coordinates` :func:`get_velocity_block_coordinates` ''' vcellids = np.atleast_1d(vcellids) + WID=self.get_WID() + WID2=WID*WID + WID3=WID2*WID # Get block ids: - blocks = vcellids.astype(int) // 64 + blocks = vcellids.astype(int) // WID3 # Get block coordinates: blockIndicesX = np.remainder(blocks.astype(int), (int)(self.__meshes[pop].__vxblocks)) blockIndicesY = np.remainder(blocks.astype(int)//(int)(self.__meshes[pop].__vxblocks), (int)(self.__meshes[pop].__vyblocks)) blockIndicesZ = blocks.astype(int)//(int)(self.__meshes[pop].__vxblocks*self.__meshes[pop].__vyblocks) - blockCoordinatesX = blockIndicesX.astype(float) * self.__meshes[pop].__dvx * 4 + self.__meshes[pop].__vxmin - blockCoordinatesY = blockIndicesY.astype(float) * self.__meshes[pop].__dvy * 4 + self.__meshes[pop].__vymin - blockCoordinatesZ = blockIndicesZ.astype(float) * self.__meshes[pop].__dvz * 4 + self.__meshes[pop].__vzmin + blockCoordinatesX = blockIndicesX.astype(float) * self.__meshes[pop].__dvx * WID + self.__meshes[pop].__vxmin + blockCoordinatesY = blockIndicesY.astype(float) * self.__meshes[pop].__dvy * WID + self.__meshes[pop].__vymin + blockCoordinatesZ = blockIndicesZ.astype(float) * self.__meshes[pop].__dvz * WID + self.__meshes[pop].__vzmin # Get cell indices: - cellids = np.remainder(vcellids.astype(int), (int)(64)) - cellIndicesX = np.remainder(cellids.astype(int), (int)(4)) - cellIndicesY = np.remainder((cellids.astype(int)//(int)(4)).astype(int), (int)(4)) - cellIndicesZ = cellids.astype(int)//(int)(16) + cellids = np.remainder(vcellids.astype(int), (int)(WID3)) + cellIndicesX = np.remainder(cellids.astype(int), (int)(WID)) + cellIndicesY = np.remainder((cellids.astype(int)//(int)(WID)).astype(int), (int)(WID)) + cellIndicesZ = cellids.astype(int)//(int)(WID2) # Get cell coordinates: cellCoordinates = np.array([blockCoordinatesX.astype(float) + (cellIndicesX.astype(float) + 0.5) * self.__meshes[pop].__dvx, blockCoordinatesY.astype(float) + (cellIndicesY.astype(float) + 0.5) * self.__meshes[pop].__dvy, @@ -1523,12 +1536,13 @@ def get_velocity_block_coordinates( self, blocks, pop="proton"): .. seealso:: :func:`get_velocity_cell_coordinates` ''' + WID=self.get_WID() blockIndicesX = np.remainder(blocks.astype(int), (int)(self.__meshes[pop].__vxblocks)) blockIndicesY = np.remainder(blocks.astype(int)//(int)(self.__meshes[pop].__vxblocks), (int)(self.__meshes[pop].__vyblocks)) blockIndicesZ = blocks.astype(int)//(int)(self.__meshes[pop].__vxblocks*self.__meshes[pop].__vyblocks) - blockCoordinatesX = blockIndicesX.astype(float) * self.__meshes[pop].__dvx * 4 + self.__meshes[pop].__vxmin - blockCoordinatesY = blockIndicesY.astype(float) * self.__meshes[pop].__dvy * 4 + self.__meshes[pop].__vymin - blockCoordinatesZ = blockIndicesZ.astype(float) * self.__meshes[pop].__dvz * 4 + self.__meshes[pop].__vzmin + blockCoordinatesX = blockIndicesX.astype(float) * self.__meshes[pop].__dvx * WID + self.__meshes[pop].__vxmin + blockCoordinatesY = blockIndicesY.astype(float) * self.__meshes[pop].__dvy * WID + self.__meshes[pop].__vymin + blockCoordinatesZ = blockIndicesZ.astype(float) * self.__meshes[pop].__dvz * WID + self.__meshes[pop].__vzmin # Return the coordinates: return np.array([blockCoordinatesX.astype(float), blockCoordinatesY.astype(float), @@ -1542,8 +1556,9 @@ def get_velocity_blocks( self, blockcoordinates, pop="proton" ): .. seealso:: :func:`get_velocity_block_coordinates` ''' + WID=self.get_WID() mins = np.array([self.__meshes[pop].__vxmin, self.__meshes[pop].__vymin, self.__meshes[pop].__vzmin]).astype(float) - dvs = np.array([4*self.__meshes[pop].__dvx, 4*self.__meshes[pop].__dvy, 4*self.__meshes[pop].__dvz]).astype(float) + dvs = np.array([WID*self.__meshes[pop].__dvx, WID*self.__meshes[pop].__dvy, WID*self.__meshes[pop].__dvz]).astype(float) multiplier = np.array([1, self.__meshes[pop].__vxblocks, self.__meshes[pop].__vxblocks * self.__meshes[pop].__vyblocks]).astype(float) velocity_block_ids = np.sum(np.floor(((blockCoordinates.astype(float) - mins) / dvs)) * multiplier, axis=-1) return velocity_block_ids @@ -1554,7 +1569,9 @@ def construct_velocity_cells( self, blocks ): :param blocks: list of block ids :returns: a numpy array containing the velocity cell ids e.g. np.array([4,2,56,44,522, ..]) ''' - return np.ravel(np.outer(np.array(blocks), np.ones(64)) + np.arange(64)) + WID=self.get_WID() + WID3=WID*WID*WID + return np.ravel(np.outer(np.array(blocks), np.ones(WID3)) + np.arange(WID3)) def construct_velocity_cell_coordinates( self, blocks ): ''' Returns velocity cell coordinates in given blocks @@ -1577,13 +1594,16 @@ def construct_velocity_cell_nodes( self, blocks, pop="proton" ): .. seealso:: :mod:`grid` ''' blocks = np.array(blocks) + WID=self.get_WID() + WID2=WID*WID + WID3=WID2*WID # Get block coordinates: blockIndicesX = np.remainder(blocks.astype(int), (int)(self.__meshes[pop].__vxblocks)).astype(np.uint16) blockIndicesY = np.remainder(blocks.astype(int)//(int)(self.__meshes[pop].__vxblocks), (int)(self.__meshes[pop].__vyblocks)).astype(np.uint16) blockIndicesZ = (blocks.astype(np.uint64)//(int)(self.__meshes[pop].__vxblocks*self.__meshes[pop].__vyblocks)).astype(np.uint16) - cellsPerDirection = 4 - cellsPerBlock = 64 + cellsPerDirection = WID + cellsPerBlock = WID3 # Get velocity cell min coordinates (per velocity block) vcellids = np.arange(cellsPerBlock).astype(np.uint32) @@ -1835,18 +1855,21 @@ def read_velocity_cells(self, cellid, pop="proton"): array_size = len(data_avgs) # Construct velocity cells: + WID=self.get_WID() + WID2=WID*WID + WID3=WID2*WID velocity_cell_ids = [] - for kv in range(4): - for jv in range(4): - for iv in range(4): - velocity_cell_ids.append(kv*16 + jv*4 + iv) + for kv in range(WID): + for jv in range(WID): + for iv in range(WID): + velocity_cell_ids.append(kv*WID2 + jv*WID + iv) for i in range(array_size): velocity_block_id = data_block_ids[i] avgIndex = 0 avgs = data_avgs[i] - for j in velocity_cell_ids + 64*velocity_block_id: + for j in velocity_cell_ids + WID3*velocity_block_id: velocity_cells[(int)(j)] = avgs[avgIndex] avgIndex = avgIndex + 1 return velocity_cells @@ -1909,6 +1932,13 @@ def get_velocity_mesh_extent(self, pop="proton"): ''' return np.array([self.__meshes[pop].__vxmin, self.__meshes[pop].__vymin, self.__meshes[pop].__vzmin, self.__meshes[pop].__vxmax, self.__meshes[pop].__vymax, self.__meshes[pop].__vzmax]) + def get_velocity_mesh_dv(self, pop="proton"): + ''' Read velocity mesh extent + + :returns: Velocity mesh grid size, array with three elements [dvx, dvy, dvz] + ''' + return np.array([self.__meshes[pop].__dvx, self.__meshes[pop].__dvy, self.__meshes[pop].__dvz]) + def read_blocks(self, cellid, pop="proton"): ''' Read raw block data from the open file and return the data along with block ids diff --git a/testpackage/testpackage_colormap.py b/testpackage/testpackage_colormap.py index e8b1fe93..46033dde 100644 --- a/testpackage/testpackage_colormap.py +++ b/testpackage/testpackage_colormap.py @@ -822,7 +822,7 @@ def extcontour(ax, XmeshXY,YmeshXY, extmaps, requestvariables=False): callrunindex.append(callindex) callindex += 1 for pop in run['pops']: - if pop is not 'avgs': + if pop != 'avgs': for call in v5multipopcalls: # Skip flux function calls if no flux files if "flux" in call and fluxLocation is None: @@ -854,7 +854,7 @@ def extcontour(ax, XmeshXY,YmeshXY, extmaps, requestvariables=False): callrunindex.append(callindex) callindex += 1 for pop in run['pops']: - if pop is not 'avgs': + if pop != 'avgs': for call in multipopcalls: # Skip flux function calls if no flux files if "flux" in call and fluxLocation is None: