def pitch_angles(vlsvReader, cellid, cosine=True, plasmaframe=False): """ Calculates the pitch angle distribution for a given cell :param vlsvReader: Some VlsvReader class with a file open :type vlsvReader: :class:`vlsvfile.VlsvReader` :param cellid: The cell id whose pitch angle the user wants NOTE: The cell id must have a velocity distribution! :param cosine: True if returning the pitch angles as a cosine plot :param plasmaframe: True if the user wants to get the pitch angle distribution in the plasma frame :returns: pitch angles and avgs [pitch_angles, avgs] .. code-block:: python # Example usage: vlsvReader = VlsvReader("fullf.0001.vlsv") result = pitch_angles( 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) """ # Read the velocity cells: velocity_cell_data = vlsvReader.read_velocity_cells(cellid) # Read bulk velocity: if vlsvReader.read_variable("rho", cellid) != 0.0: bulk_velocity = np.array( vlsvReader.read_variable("rho_v", cellid) / vlsvReader.read_variable("rho", cellid), copy=False ) else: bulk_velocity = 0.0 # Calculate the pitch angles for the data: B = vlsvReader.read_variable("B", cellid) B_unit = B / np.linalg.norm(B) # Get cells: vcellids = velocity_cell_data.keys() # Get avgs data: avgs = velocity_cell_data.values() # Get a list of velocity coordinates: if plasmaframe == True: v = vlsvReader.get_velocity_cell_coordinates(vcellids) - bulk_velocity else: v = vlsvReader.get_velocity_cell_coordinates(vcellids) # Get norms: v_norms = np.sum(np.abs(v) ** 2, axis=-1) ** (1.0 / 2) # Get the angles: if cosine == True: pitch_angles = v.dot(B_unit) / v_norms units = "radian" else: pitch_angles = np.arccos(v.dot(B_unit) / v_norms) / (2 * np.pi) * 360 units = "degree" # Return the pitch angles and avgs values: from output import output_1d if vlsvReader.read_variable("rho", cellid) != 0.0: return output_1d([pitch_angles, avgs], ["Pitch_angle", "avgs"], [units, ""]) else: return output_1d([[0], [1e-9]], ["Pitch_angle", "avgs"], [units, ""])
def pitch_angles(vlsvReader, cellid, cosine=True, plasmaframe=False): ''' Calculates the pitch angle distribution for a given cell :param vlsvReader: Some VlsvReader class with a file open :type vlsvReader: :class:`vlsvfile.VlsvReader` :param cellid: The cell id whose pitch angle the user wants NOTE: The cell id must have a velocity distribution! :param cosine: True if returning the pitch angles as a cosine plot :param plasmaframe: True if the user wants to get the pitch angle distribution in the plasma frame :returns: pitch angles and avgs [pitch_angles, avgs] .. code-block:: python # Example usage: vlsvReader = VlsvReader("fullf.0001.vlsv") result = pitch_angles( 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) ''' # Read the velocity cells: velocity_cell_data = vlsvReader.read_velocity_cells(cellid) # Read bulk velocity: if vlsvReader.read_variable("rho", cellid) != 0.0: bulk_velocity = np.array(vlsvReader.read_variable("rho_v", cellid) / vlsvReader.read_variable("rho", cellid), copy=False) else: bulk_velocity = 0.0 # Calculate the pitch angles for the data: B = vlsvReader.read_variable("B", cellid) B_unit = B / np.linalg.norm(B) # Get cells: vcellids = velocity_cell_data.keys() # Get avgs data: avgs = velocity_cell_data.values() # Get a list of velocity coordinates: if plasmaframe == True: v = vlsvReader.get_velocity_cell_coordinates(vcellids) - bulk_velocity else: v = vlsvReader.get_velocity_cell_coordinates(vcellids) # Get norms: v_norms = np.sum(np.abs(v)**2, axis=-1)**(1. / 2) # Get the angles: if cosine == True: pitch_angles = v.dot(B_unit) / v_norms units = "radian" else: pitch_angles = np.arccos(v.dot(B_unit) / v_norms) / (2 * np.pi) * 360 units = "degree" # Return the pitch angles and avgs values: from output import output_1d if vlsvReader.read_variable("rho", cellid) != 0.0: return output_1d([pitch_angles, avgs], ["Pitch_angle", "avgs"], [units, ""]) else: return output_1d([[0], [1e-9]], ["Pitch_angle", "avgs"], [units, ""])
def gyrophase_angles_from_file( vlsvReader, cellid): ''' Calculates the gyrophase angle angle distribution for a given cell with a given file :param vlsvReader: Some VlsvReader class with a file open :type vlsvReader: :class:`vlsvfile.VlsvReader` :param cellid: The cell id whose gyrophase angle the user wants NOTE: The cell id must have a velocity distribution! :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) # Plot the data import pylab as pl pl.hist(result[0].data, weights=result[1].data, bins=100, log=False) ''' # Read the velocity cells: velocity_cell_data = vlsvReader.read_velocity_cells(cellid) if velocity_cell_data == []: from output import output_1d return output_1d([[0.0, 1.0], [1.0, 1.0]], ["Gyrophase_angle", "avgs"], ["", ""]) # Read bulk velocity: if vlsvReader.check_variable( "rho" ): if vlsvReader.read_variable("rho", cellid) != 0.0: bulk_velocity = np.array(vlsvReader.read_variable("rho_v", cellid) / vlsvReader.read_variable("rho", cellid), copy=False) else: from output import output_1d return output_1d([[0.0, 1.0], [1.0, 1.0]], ["Gyrophase_angle", "avgs"], ["", ""]) else: moments = np.array(vlsvReader.read_variable("moments", cellid)) if moments[0] != 0.0: bulk_velocity = np.array(np.divide(np.array(moments[1:]), moments[0])) else: from output import output_1d return output_1d([[0.0, 1.0], [1.0, 1.0]], ["Gyrophase_angle", "avgs"], ["", ""]) # Calculate the gyrophase angles for the data: if vlsvReader.check_variable( "B" ): B = vlsvReader.read_variable("B", cellid) else: B = vlsvReader.read_variable("background_B", cellid) + vlsvReader.read_variable("perturbed_B", cellid) if np.linalg.norm(B) != 0.0: B_unit = B / np.linalg.norm(B) else: from output import output_1d return output_1d([[0.0, 1.0], [1.0, 1.0]], ["Gyrophase_angle", "avgs"], ["", ""]) # Get cells: vcellids = velocity_cell_data.keys() # Get a list of velocity coordinates: velocity_coordinates = vlsvReader.get_velocity_cell_coordinates(vcellids) return gyrophase_angles(bulk_velocity, B_unit, velocity_cell_data, velocity_coordinates)
def subtract_1d_polynomial_fit(x, y): ''' Fits 1d polynomial into the data, subtracts it from the given data and returns the subtracted data. :param x: The x-axis :param y: Data to manipulate in either variable or raw form :returns: data from which a 1d polynomial fit has been subtracted .. note:: This may be useful for fourier transforms ''' # Fit a polynomial into the data from variable import get_data, get_name, get_units import numpy as np parameters = 2 def function(args, x, y): x = np.array(x) y = np.array(y) value = args[0] + args[1] * x return y - value from scipy import optimize fit = optimize.leastsq(function, np.ones(parameters), args=(get_data(x), get_data(y))) y_fitted = (-1) * function(fit[0], x, 0) # Create a new array y2 which has a forced constant amplitude for the (possible) waves: y2 = y - y_fitted # Return the data from output import output_1d return output_1d([y2], [get_name(y)], [get_units(y)])
def subtract_1d_polynomial_fit( x, y ): ''' Fits 1d polynomial into the data, subtracts it from the given data and returns the subtracted data. :param x: The x-axis :param y: Data to manipulate in either variable or raw form :returns: data from which a 1d polynomial fit has been subtracted .. note:: This may be useful for fourier transforms ''' # Fit a polynomial into the data from variable import get_data, get_name, get_units import numpy as np parameters = 2 def function(args, x, y): x = np.array(x) y = np.array(y) value = args[0] + args[1]*x return y - value from scipy import optimize fit = optimize.leastsq(function, np.ones(parameters), args=(get_data(x), get_data(y))) y_fitted = (-1)*function(fit[0], x, 0) # Create a new array y2 which has a forced constant amplitude for the (possible) waves: y2 = y - y_fitted # Return the data from output import output_1d return output_1d( [y2], [get_name(y)], [get_units(y)] )
def fourier(dt, y, kaiserwindowparameter=0): ''' Function for returning fourier series and frequencies of some given arrays t and y :param dt: Time :param y: Some variable data :returns: the frequencies, new time variables and frequencies .. note:: return format: [FFT, frequencies, t, y] .. note:: t must have a constant time stepping .. code-block:: python Example usage: fourier_data = fourier( t=np.arange(0,500,0.5), y=rho_data, kaiserwindowparameter=14 ) ''' # Get the data from variable import get_data #t_data = get_data(t) y_data = get_data(y) # First check the t array whether it has a constant dt #dt = get_data(t)[1] - get_data(t)[0] #for i in xrange(len(get_data(t))-1): # if dt != get_data(t)[i+1] - get_data(t)[i]: # print "Gave bad timestep to plot_fourier, the time step in array t must be constant (for now)" # Use kaiser window on y import numpy as np y_tmp = get_data(y) * np.kaiser(len(get_data(y)), kaiserwindowparameter) # Do FFT on the data fourier = np.fft.fft(y_tmp) * (1 / (float)(len(y_tmp))) # Get frequencies of the fourier freq = np.fft.fftfreq(len(fourier), d=dt) # Declare t2 (Note: This is the same as t but we want the steps to be thicker so the data looks smoother dt2 = dt * 0.01 t2 = np.arange(len(y_tmp) * 100) * dt2 # Declare y2 y2 = np.array([ np.sum(fourier * np.exp(complex(0, 1) * 2 * np.pi * freq * T)) for T in t2 ]) from output import output_1d from variable import get_name # Get the indexes: toIndex = (int)((len(freq) / 2) / 2.0 + 1) return output_1d( [2 * np.abs(fourier[1:toIndex]), freq[1:toIndex], t2, y2], ["FFT", "frequency", "time", get_name(y)])
def fourier( t, y, kaiserwindowparameter=0 ): ''' Function for returning fourier series and frequencies of some given arrays t and y :param t: Time :param y: Some variable data :returns: the frequencies, new time variables and frequencies .. note:: return format: [FFT, frequencies, t, y] .. note:: t must have a constant time stepping .. code-block:: python Example usage: fourier_data = fourier( t=np.arange(0,500,0.5), y=rho_data, kaiserwindowparameter=14 ) ''' # Get the data from variable import get_data #t_data = get_data(t) y_data = get_data(y) # First check the t array whether it has a constant t #t = get_data(t)[1] - get_data(t)[0] #for i in xrange(len(get_data(t))-1): # if t != get_data(t)[i+1] - get_data(t)[i]: # print "Gave bad timestep to plot_fourier, the time step in array t must be constant (for now)" # Use kaiser window on y import numpy as np y_tmp = get_data(y) * np.kaiser(len(get_data(y)), kaiserwindowparameter) # Do FFT on the data fourier=np.fft.fft(y_tmp) * (1/(float)(len(y_tmp))) # Get frequencies of the fourier freq=np.fft.fftfreq(len(fourier), d=t) # Declare t2 (Note: This is the same as t but we want the steps to be thicker so the data looks smoother t2=t*0.01 t2=np.arange(len(y_tmp)*100)*t2 # Declare y2 y2=np.array([np.sum(fourier*np.exp(complex(0,1)*2*np.pi*freq*T)) for T in t2]) from output import output_1d from variable import get_name # Get the indexes: toIndex = (int)((len(freq)/2)/2.0 + 1) return output_1d([2*np.abs(fourier[1:toIndex]), freq[1:toIndex], t2, y2], ["FFT", "frequency", "time", get_name(y)])
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 gyrophase_angles_from_file(vlsvReader, cellid): ''' Calculates the gyrophase angle angle distribution for a given cell with a given file :param vlsvReader: Some VlsvReader class with a file open :type vlsvReader: :class:`vlsvfile.VlsvReader` :param cellid: The cell id whose gyrophase angle the user wants NOTE: The cell id must have a velocity distribution! :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) # Plot the data import pylab as pl pl.hist(result[0].data, weights=result[1].data, bins=100, log=False) ''' # Read the velocity cells: velocity_cell_data = vlsvReader.read_velocity_cells(cellid) if velocity_cell_data == []: from output import output_1d return output_1d([[0.0, 1.0], [1.0, 1.0]], ["Gyrophase_angle", "avgs"], ["", ""]) # Read bulk velocity: if vlsvReader.check_variable("rho"): if vlsvReader.read_variable("rho", cellid) != 0.0: bulk_velocity = np.array( vlsvReader.read_variable("rho_v", cellid) / vlsvReader.read_variable("rho", cellid), copy=False) else: from output import output_1d return output_1d([[0.0, 1.0], [1.0, 1.0]], ["Gyrophase_angle", "avgs"], ["", ""]) else: moments = np.array(vlsvReader.read_variable("moments", cellid)) if moments[0] != 0.0: bulk_velocity = np.array( np.divide(np.array(moments[1:]), moments[0])) else: from output import output_1d return output_1d([[0.0, 1.0], [1.0, 1.0]], ["Gyrophase_angle", "avgs"], ["", ""]) # Calculate the gyrophase angles for the data: if vlsvReader.check_variable("B"): B = vlsvReader.read_variable("B", cellid) else: B = vlsvReader.read_variable("background_B", cellid) + vlsvReader.read_variable( "perturbed_B", cellid) if np.linalg.norm(B) != 0.0: B_unit = B / np.linalg.norm(B) else: from output import output_1d return output_1d([[0.0, 1.0], [1.0, 1.0]], ["Gyrophase_angle", "avgs"], ["", ""]) # Get cells: vcellids = velocity_cell_data.keys() # Get a list of velocity coordinates: velocity_coordinates = vlsvReader.get_velocity_cell_coordinates(vcellids) return gyrophase_angles(bulk_velocity, B_unit, velocity_cell_data, velocity_coordinates)
def pitch_angles(vlsvReader, cellid, nbins=10, filename=None, filedir=None, step=None, outputdir=None, outputfile=None, cosine=False, plasmaframe=False, vcut=None, pop="proton"): ''' Calculates the pitch angle distribution for a given cell :param vlsvReader: Some VlsvReader class with a file open. Can be overriden by keywords. :param cellid: The cell id whose pitch angle the user wants NOTE: The cell id must have a velocity distribution! :kword filename: path to .vlsv file to use for input. :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 nbins: How many bins to use for the distribution :kword cosine: True if returning the pitch angles as a cosine(alpha) plot [-1,1]. If false, returns as pitch angle in degrees [0,180]. :kword plasmaframe: True if the user wants to get the pitch angle distribution in the plasma frame (for this population). If set to a string, will try to use the string as a variable for the frame to transform into. If set to a 3-element vector, will use that frame instead. :kword vcut: Set to True to ignore velocity cells below 2x the thermal speed. If set to a number, will use that velocity in m/s instead. :kword outputdir: Optional (recommended) method to save results to a file in the given directory. If directory does not exist, it will be created. :kword outputfile: Provide exact output file name :kword pop: Active population, defaults to proton (avgs) :returns: pitch angles and avgs [pitch_angles, avgs] .. code-block:: python # Example usage: vlsvReader = VlsvReader("restart.0000798.vlsv") result = pitch_angles( vlsvReader=vlsvReader, 1924, cosine=True, plasmaframe=True, outputdir="/wrk/username/pitchangledirectory/" ) ''' # Input file or object if filename != None: vlsvReader = pt.vlsvfile.VlsvReader(filename) elif ((filedir != None) and (step != None)): filename = filedir + 'bulk.' + str(step).rjust(7, '0') + '.vlsv' vlsvReader = pt.vlsvfile.VlsvReader(filename) # Transform to a different frame? frame = [0., 0., 0.] if plasmaframe is not False: if isinstance(plasmaframe, str): frame = vlsvReader.read_variable(plasmaframe, cellid) elif plasmaframe is True: if vlsvReader.check_variable("moments"): # restart frame = vlsvReader.read_variable('restart_V', cellid) else: frame = vlsvReader.read_variable('V', cellid) elif len(plasmaframe) == 3: # Assume it's a vector frame = plasmaframe # Find the magnetic field direction B = vlsvReader.read_variable("B", cellid) Bmag = np.linalg.norm(B) B_unit = B / Bmag # verify population 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!") sys.exit() else: if not vlsvReader.check_population(pop): print("Unable to detect population " + pop + " in .vlsv file!") sys.exit() # Read temperature for thermal speed if vcut is True: if vlsvReader.check_variable( "moments"): # restart, use overall rho/rhom and pressure rhom = vlsvReader.read_variable("restart_rhom", cellid) PDiagonal = vlsvReader.read_variable("pressure", cellid) Pressure = (PDiagonal[0] + PDiagonal[1] + PDiagonal[2]) * (1. / 3.) vth = np.sqrt(Pressure * 8. / (rhom * np.pi)) else: if vlsvReader.check_variable("rhom"): # multipop vth = vlsvReader.read_variable(pop + "/vThermal", cellid) else: vth = vlsvReader.read_variable("vThermal", cellid) vcutoff = 2 * vth else: # assume it's a number to use as the speed in m/s vcutoff = vcut # Read the velocity cells: velocity_cell_data = vlsvReader.read_velocity_cells(cellid, pop=pop) vcellids = velocity_cell_data.keys() avgs = velocity_cell_data.values() # Transform to a frame v = vlsvReader.get_velocity_cell_coordinates(vcellids, pop=pop) - frame # Get the angles: v_norms = np.linalg.norm(v, axis=-1) v_unit = v / v_norms[:, np.newaxis] if cosine == True: pitch_angles = (v_unit * B_unit).sum(-1) units = "cos alpha" pitchrange = (-1, 1) else: pitch_angles = np.arccos((v_unit * B_unit).sum(-1)) * (180. / np.pi) units = "degree" pitchrange = (0, 180) # Calculate velocity cell sizes [vxsize, vysize, vzsize] = vlsvReader.get_velocity_mesh_size(pop=pop) [vxmin, vymin, vzmin, vxmax, vymax, vzmax] = vlsvReader.get_velocity_mesh_extent(pop=pop) dvx = (vxmax - vxmin) / (4 * vxsize) dvy = (vymax - vymin) / (4 * vysize) dvz = (vzmax - vzmin) / (4 * vzsize) dv3 = dvx * dvy * dvz # Clip negative avgs to zero # (Some elements may be negative due to ghost cell propagation effects) avgs = np.array(avgs).clip(min=0) * dv3 # Mask off cells below threshold condition = (v_norms > vcutoff) # Get the velocity cells above cutoff speed #vcellids_nonsphere = np.extract(condition, vcellids) # Get the avgs avgs_nonsphere = np.extract(condition, avgs) # Get the pitch-angles (or cosines) pitch_nonsphere = np.extract(condition, pitch_angles) # Generate a histogram weights, angles = np.histogram(pitch_nonsphere, nbins, range=pitchrange, weights=avgs_nonsphere) # Wrap the data into a custon format result = output_1d([angles, weights], ["Pitch_angle", "sum_avgs"], [units, "1/m3"]) rho_summed = np.sum(avgs) rho_nonsphere = np.sum(avgs_nonsphere) print("rho", rho_summed, rho_nonsphere) if outputfile != None or outputdir != None: # Save output to file # Generate filename timestr = '{:4.1f}'.format(vlsvReader.read_parameter("time")) if outputfile == None: outputfile = outputdir + "/pitchangle_weights_cellid_" + str( cellid).rjust(7, '0') + "_time_" + timestr + ".txt" # Check to find actual target sub-directory outputprefixind = outputfile.rfind('/') if outputprefixind >= 0: outputdir = outputfile[:outputprefixind + 1] # Ensure output directory exists and is writeable if not os.path.exists(outputdir): try: os.makedirs(outputdir) except: pass if not os.access(outputdir, os.W_OK): print("No write access for directory " + outputdir + "! Exiting.") return outfilewrite = open(outputfile, 'w') outfilewrite.write("# cellid time rho rho_nonsphere Bx By Bz pop\n") outfilewrite.write(str(int(cellid)) + " " + timestr) outfilewrite.write(' {:E} {:E} {:E} {:E} {:E}'.format( rho_summed, rho_nonsphere, B[0], B[1], B[2]) + " " + pop + "\n") outfilewrite.write("# nbins, bin_edges\n") for angle in angles: outfilewrite.write('{:E} '.format(angle)) outfilewrite.write("\n") outfilewrite.write("# bin_values\n") for weight in weights: outfilewrite.write('{:E} '.format(weight)) outfilewrite.write("\n") outfilewrite.close() # Return the histogram return result
def get_cellids_coordinates_distances(vlsvReader, xmax, xmin, xcells, ymax, ymin, ycells, zmax, zmin, zcells, cell_lengths, point1, point2): ''' Calculates coordinates to be used in the cut_through. The coordinates are calculated so that every cell gets picked in the coordinates. :param vlsvReader: Some open VlsvReader :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 ''' point1 = np.array(point1, copy=False) point2 = np.array(point2, copy=False) distances = [0] cellids = [vlsvReader.get_cellid(point1)] coordinates = [point1] epsilon = sys.float_info.epsilon * 2 iterator = point1 unit_vector = (point2 - point1) / np.linalg.norm(point2 - point1 + epsilon) while True: # Get the cell id cellid = vlsvReader.get_cellid(iterator) if cellid == 0: print("ERROR, invalid cell id!") return # Get the max and min boundaries: 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) # Remove negative coefficients: for i in xrange(3): if coefficients_min[i] <= 0: coefficients_min[i] = sys.float_info.max if coefficients_max[i] <= 0: coefficients_max[i] = sys.float_info.max # Find the minimum coefficient for calculating the minimum distance from a boundary coefficient = min([min(coefficients_min), min(coefficients_max)]) * 1.01 # Get the cell id in the new coordinates newcoordinate = iterator + coefficient * unit_vector newcellid = vlsvReader.get_cellid(newcoordinate) # Check if valid cell id: if newcellid == 0: break # Append the cell id: cellids.append(newcellid) # Append the coordinate: coordinates.append(newcoordinate) # Append the distance: distances.append(np.linalg.norm(newcoordinate - point1)) # Move the iterator to the next cell. Note: Epsilon is here to ensure there are no bugs with float rounding iterator = iterator + coefficient * unit_vector # Check to make sure the iterator is not moving past point2: if min((point2 - iterator) * unit_vector) < 0: break # 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"])
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. :param vlsvReader: Some open VlsvReader :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] .. code-block:: python Example: vlsvReader = VlsvReader(\"testfile.vlsv\") cut_through = cut_through_step(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) ''' # Transform point1 and point2 into numpy array: point1 = np.array(point1) point2 = np.array(point2) # Make sure point1 and point2 are inside bounds if vlsvReader.get_cellid(point1) == 0: print("ERROR, POINT1 IN CUT-THROUGH OUT OF BOUNDS!") if vlsvReader.get_cellid(point2) == 0: print("ERROR, POINT2 IN CUT-THROUGH OUT OF BOUNDS!") # Find path distances = point2 - point1 largestdistance = np.linalg.norm(distances) largestindex = np.argmax(abs(distances)) derivative = distances / abs(distances[largestindex]) # Re=6371000. # print("distances",distances/Re) # print("largestindex",largestindex) # print("derivative",derivative) # Get parameters from the file to determine a good length between points (step length): # Get xmax, xmin and xcells_ini [xcells, ycells, zcells] = vlsvReader.get_spatial_mesh_size() [xmin, ymin, zmin, xmax, ymax, zmax] = vlsvReader.get_spatial_mesh_extent() # Calculate cell lengths: cell_lengths = np.array([(xmax - xmin) / (float)(xcells), (ymax - ymin) / (float)(ycells), (zmax - zmin) / (float)(zcells)]) # Initialize lists distances = [0] cellids = [vlsvReader.get_cellid(point1)] coordinates = [point1] finalcellid = vlsvReader.get_cellid(point2) #print(" cellids init ",cellids,finalcellid) # Loop until final cellid is reached while True: newcoordinate = coordinates[-1] + cell_lengths * derivative newcellid = vlsvReader.get_cellid(newcoordinate) distances.append(np.linalg.norm(newcoordinate - point1)) coordinates.append(newcoordinate) cellids.append(newcellid) #print(distances[-1]/Re,np.array(coordinates[-1])/Re,cellids[-1]) if newcellid == finalcellid: break if distances[-1] > largestdistance: break # 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"])
def get_cellids_coordinates_distances( vlsvReader, xmax, xmin, xcells, ymax, ymin, ycells, zmax, zmin, zcells, cell_lengths, point1, point2 ): ''' Calculates coordinates to be used in the cut_through. The coordinates are calculated so that every cell gets picked in the coordinates. :param vlsvReader: Some open VlsvReader :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 ''' point1 = np.array(point1, copy=False) point2 = np.array(point2, copy=False) distances = [0] cellids = [vlsvReader.get_cellid(point1)] coordinates = [point1] epsilon = sys.float_info.epsilon*2 iterator = point1 unit_vector = (point2 - point1) / np.linalg.norm(point2 - point1 + epsilon) while True: # Get the cell id cellid = vlsvReader.get_cellid(iterator) if cellid == 0: print "ERROR, invalid cell id!" return # Get the max and min boundaries: 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) # Remove negative coefficients: for i in xrange(3): if coefficients_min[i] <= 0: coefficients_min[i] = sys.float_info.max if coefficients_max[i] <= 0: coefficients_max[i] = sys.float_info.max # Find the minimum coefficient for calculating the minimum distance from a boundary coefficient = min([min(coefficients_min), min(coefficients_max)]) * 1.01 # Get the cell id in the new coordinates newcoordinate = iterator + coefficient * unit_vector newcellid = vlsvReader.get_cellid( newcoordinate ) # Check if valid cell id: if newcellid == 0: break # Append the cell id: cellids.append( newcellid ) # Append the coordinate: coordinates.append( newcoordinate ) # Append the distance: distances.append( np.linalg.norm( newcoordinate - point1 ) ) # Move the iterator to the next cell. Note: Epsilon is here to ensure there are no bugs with float rounding iterator = iterator + coefficient * unit_vector # Check to make sure the iterator is not moving past point2: if min((point2 - iterator)* unit_vector) < 0: break # 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"] )
def cell_time_evolution(vlsvReader_list, variables, cellids, units=""): ''' Returns variable data from a time evolution of some certain cell ids :param vlsvReader_list: List containing VlsvReaders with a file open :type vlsvReader_list: :class:`vlsvfile.VlsvReader` :param variables: Name of the variables :param cellids: List of cell ids :param units: List of units for the variables (OPTIONAL) :returns: an array containing the data for the time evolution for every cell id .. code-block:: python import pytools as pt; import pylab as pl # Example of usage: time_data = pt.calculations.cell_time_evolution( vlsvReader_list=[VlsvReader("bulk.000.vlsv"), VlsvReader("bulk.001.vlsv"), VlsvReader("bulk.002.vlsv")], variables=["rho", "Pressure", "B"], cellids=[2,4], units=["N", "Pascal", "T"] ) # Check output print time_data # Now plot the results: time = time_data[0] rho = time_data[3] pt.plot.plot_variables(time, rho) pl.show() # Do post processing: rho_data = rho.data non_existing_example_function(rho_data) ''' vlsvReader_list = np.atleast_1d(vlsvReader_list) variables = np.atleast_1d(variables) cellids = np.atleast_1d(cellids) parameters = ["t", "tstep", "fileIndex"] parameter_units = ["s", "", ""] #construct empty units, if none are given if (units == "") or (len(units) != len(variables)): units = ["" for i in range(len(variables))] #construct data data = [[] for i in range(len(parameters) + len(cellids) * len(variables))] for t in range(len(vlsvReader_list)): # Get the vlsv reader vlsvReader = vlsvReader_list[t] # Open the vlsv reader's file: vlsvReader.optimize_open_file() #go through parameters for j in range(len(parameters)): # Read the parameter # Save the data into the right slot in the data array: data[j].append(vlsvReader.read_parameter(parameters[j])) # Go through variables: for j in range(len(variables)): variable = variables[j] # Read the variable for all cell ids #variables_for_cellids = vlsvReader.read_variables_for_cellids( variable, cellids ) # Save the data into the right slot in the data array: for i in range(len(cellids)): data[len(parameters) + i * len(variables) + j].append( vlsvReader.read_variable(variable, cellids[i])) # For optimization purposes we are now freeing vlsvReader's memory # Note: Upon reading data vlsvReader created an internal hash map that takes a lot of memory vlsvReader.optimize_clear_fileindex_for_cellid() # Close the vlsv reader's file: vlsvReader.optimize_close_file() from output import output_1d return output_1d( data, parameters + [ variables[(int)(i) % (int)(len(variables))] for i in range(len(data) - len(parameters)) ], parameter_units + [ units[(int)(i) % (int)(len(units))] for i in range(len(data) - len(parameters)) ])
def cell_time_evolution( vlsvReader_list, variables, cellids, units="" ): ''' Returns variable data from a time evolution of some certain cell ids :param vlsvReader_list: List containing VlsvReaders with a file open :type vlsvReader_list: :class:`vlsvfile.VlsvReader` :param variables: Name of the variables :param cellids: List of cell ids :param units: List of units for the variables (OPTIONAL) :returns: an array containing the data for the time evolution for every cell id .. code-block:: python import pytools as pt; import pylab as pl # Example of usage: time_data = pt.calculations.cell_time_evolution( vlsvReader_list=[VlsvReader("bulk.000.vlsv"), VlsvReader("bulk.001.vlsv"), VlsvReader("bulk.002.vlsv")], variables=["rho", "Pressure", "B"], cellids=[2,4], units=["N", "Pascal", "T"] ) # Check output print time_data # Now plot the results: time = time_data[0] rho = time_data[3] pt.plot.plot_variables(time, rho) pl.show() # Do post processing: rho_data = rho.data non_existing_example_function(rho_data) ''' vlsvReader_list = np.atleast_1d(vlsvReader_list) variables = np.atleast_1d(variables) cellids = np.atleast_1d(cellids) parameters = ["t","tstep","fileIndex"] parameter_units=["s","",""] #construct empty units, if none are given if (units == "") or (len(units) != len(variables)): units=[ "" for i in xrange(len(variables))] #construct data data = [[] for i in xrange(len(parameters)+len(cellids)*len(variables))] for t in xrange(len(vlsvReader_list)): # Get the vlsv reader vlsvReader = vlsvReader_list[t] # Open the vlsv reader's file: vlsvReader.optimize_open_file() #go through parameters for j in xrange(len(parameters)): # Read the parameter # Save the data into the right slot in the data array: data[j].append(vlsvReader.read_parameter( parameters[j])) # Go through variables: for j in xrange(len(variables)): variable = variables[j] # Read the variable for all cell ids #variables_for_cellids = vlsvReader.read_variables_for_cellids( variable, cellids ) # Save the data into the right slot in the data array: for i in xrange(len(cellids)): data[len(parameters)+i*len(variables)+j].append(vlsvReader.read_variable( variable, cellids[i] )) # For optimization purposes we are now freeing vlsvReader's memory # Note: Upon reading data vlsvReader created an internal hash map that takes a lot of memory vlsvReader.optimize_clear_fileindex_for_cellid() # Close the vlsv reader's file: vlsvReader.optimize_close_file() from output import output_1d return output_1d( data, parameters + [variables[(int)(i)%(int)(len(variables))] for i in xrange(len(data)-len(parameters))], parameter_units + [units[(int)(i)%(int)(len(units))] for i in xrange(len(data)-len(parameters))] )
def pitch_angles( vlsvReader, cellid, nbins=10, filename=None, filedir=None, step=None, outputdir=None, outputfile=None, cosine=False, plasmaframe=False, vcut=None, pop="proton"): ''' Calculates the pitch angle distribution for a given cell :param vlsvReader: Some VlsvReader class with a file open. Can be overriden by keywords. :param cellid: The cell id whose pitch angle the user wants NOTE: The cell id must have a velocity distribution! :kword filename: path to .vlsv file to use for input. :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 nbins: How many bins to use for the distribution :kword cosine: True if returning the pitch angles as a cosine(alpha) plot [-1,1]. If false, returns as pitch angle in degrees [0,180]. :kword plasmaframe: True if the user wants to get the pitch angle distribution in the plasma frame (for this population). If set to a string, will try to use the string as a variable for the frame to transform into. If set to a 3-element vector, will use that frame instead. :kword vcut: Set to True to ignore velocity cells below 2x the thermal speed. If set to a number, will use that velocity in m/s instead. :kword outputdir: Optional (recommended) method to save results to a file in the given directory. If directory does not exist, it will be created. :kword outputfile: Provide exact output file name :kword pop: Active population, defaults to proton (avgs) :returns: pitch angles and avgs [pitch_angles, avgs] .. code-block:: python # Example usage: vlsvReader = VlsvReader("restart.0000798.vlsv") result = pitch_angles( vlsvReader=vlsvReader, 1924, cosine=True, plasmaframe=True, outputdir="/wrk/username/pitchangledirectory/" ) ''' # Input file or object if filename!=None: vlsvReader=pt.vlsvfile.VlsvReader(filename) elif ((filedir!=None) and (step!=None)): filename = filedir+'bulk.'+str(step).rjust(7,'0')+'.vlsv' vlsvReader=pt.vlsvfile.VlsvReader(filename) # Transform to a different frame? frame = [0.,0.,0.] if plasmaframe is not False: if isinstance(plasmaframe, str): frame = vlsvReader.read_variable(plasmaframe, cellid) elif plasmaframe is True: if vlsvReader.check_variable("moments"): # restart frame = vlsvReader.read_variable('restart_V', cellid) else: frame = vlsvReader.read_variable('V', cellid) elif len(plasmaframe)==3: # Assume it's a vector frame = plasmaframe # Find the magnetic field direction B = vlsvReader.read_variable("B", cellid) Bmag = np.linalg.norm(B) B_unit = B / Bmag # verify population 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!") sys.exit() else: if not vlsvReader.check_population(pop): print("Unable to detect population "+pop+" in .vlsv file!") sys.exit() # Read temperature for thermal speed if vcut is True: if vlsvReader.check_variable("moments"): # restart, use overall rho/rhom and pressure rhom = vlsvReader.read_variable("restart_rhom", cellid) PDiagonal = vlsvReader.read_variable("pressure", cellid) Pressure = (PDiagonal[0] + PDiagonal[1] + PDiagonal[2])*(1./3.) vth = np.sqrt(Pressure*8./(rhom*np.pi)) else: if vlsvReader.check_variable("rhom"): # multipop vth = vlsvReader.read_variable(pop+"/vThermal", cellid) else: vth = vlsvReader.read_variable("vThermal", cellid) vcutoff=2*vth else: # assume it's a number to use as the speed in m/s vcutoff=vcut # Read the velocity cells: velocity_cell_data = vlsvReader.read_velocity_cells(cellid, pop=pop) vcellids = velocity_cell_data.keys() avgs = velocity_cell_data.values() # Transform to a frame v = vlsvReader.get_velocity_cell_coordinates(vcellids, pop=pop) - frame # Get the angles: v_norms = np.linalg.norm(v,axis=-1) v_unit = v/v_norms[:,np.newaxis] if cosine == True: pitch_angles = (v_unit*B_unit).sum(-1) units = "cos alpha" pitchrange = (-1,1) else: pitch_angles = np.arccos((v_unit*B_unit).sum(-1)) * (180./np.pi) units = "degree" pitchrange = (0,180) # Calculate velocity cell sizes [vxsize, vysize, vzsize] = vlsvReader.get_velocity_mesh_size(pop=pop) [vxmin, vymin, vzmin, vxmax, vymax, vzmax] = vlsvReader.get_velocity_mesh_extent(pop=pop) dvx=(vxmax-vxmin)/(4*vxsize) dvy=(vymax-vymin)/(4*vysize) dvz=(vzmax-vzmin)/(4*vzsize) dv3=dvx*dvy*dvz # Clip negative avgs to zero # (Some elements may be negative due to ghost cell propagation effects) avgs = np.array(avgs).clip(min=0) * dv3 # Mask off cells below threshold condition = (v_norms > vcutoff) # Get the velocity cells above cutoff speed #vcellids_nonsphere = np.extract(condition, vcellids) # Get the avgs avgs_nonsphere = np.extract(condition, avgs) # Get the pitch-angles (or cosines) pitch_nonsphere = np.extract(condition, pitch_angles) # Generate a histogram weights, angles = np.histogram(pitch_nonsphere, nbins, range=pitchrange, weights=avgs_nonsphere ) # Wrap the data into a custon format result = output_1d([angles, weights], ["Pitch_angle", "sum_avgs"], [units, "1/m3"]) rho_summed = np.sum(avgs) rho_nonsphere = np.sum(avgs_nonsphere) print("rho",rho_summed, rho_nonsphere) if outputfile!=None or outputdir!=None: # Save output to file # Generate filename timestr='{:4.1f}'.format(vlsvReader.read_parameter("time")) if outputfile==None: outputfile = outputdir+"/pitchangle_weights_cellid_"+str(cellid).rjust(7,'0')+"_time_"+timestr+".txt" # Check to find actual target sub-directory outputprefixind = outputfile.rfind('/') if outputprefixind >= 0: outputdir = outputfile[:outputprefixind+1] # Ensure output directory exists and is writeable if not os.path.exists(outputdir): try: os.makedirs(outputdir) except: pass if not os.access(outputdir, os.W_OK): print("No write access for directory "+outputdir+"! Exiting.") return outfilewrite=open(outputfile,'w') outfilewrite.write("# cellid time rho rho_nonsphere Bx By Bz pop\n") outfilewrite.write(str(int(cellid))+" "+timestr) outfilewrite.write(' {:E} {:E} {:E} {:E} {:E}'.format(rho_summed,rho_nonsphere,B[0],B[1],B[2])+" "+pop+"\n") outfilewrite.write("# nbins, bin_edges\n") for angle in angles: outfilewrite.write('{:E} '.format(angle)) outfilewrite.write("\n") outfilewrite.write("# bin_values\n") for weight in weights: outfilewrite.write('{:E} '.format(weight)) outfilewrite.write("\n") outfilewrite.close() # Return the histogram return result
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. :param vlsvReader: Some open VlsvReader :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] .. code-block:: python Example: vlsvReader = VlsvReader(\"testfile.vlsv\") cut_through = cut_through_step(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) ''' # Transform point1 and point2 into numpy array: point1 = np.array(point1) point2 = np.array(point2) # Make sure point1 and point2 are inside bounds if vlsvReader.get_cellid(point1) == 0: print "ERROR, POINT1 IN CUT-THROUGH OUT OF BOUNDS!" if vlsvReader.get_cellid(point2) == 0: print "ERROR, POINT2 IN CUT-THROUGH OUT OF BOUNDS!" # Find path distances = point2-point1 largestdistance = np.linalg.norm(distances) largestindex = np.argmax(abs(distances)) derivative = distances/abs(distances[largestindex]) # Re=6371000. # print("distances",distances/Re) # print("largestindex",largestindex) # print("derivative",derivative) # Get parameters from the file to determine a good length between points (step length): # Get xmax, xmin and xcells_ini [xcells, ycells, zcells] = vlsvReader.get_spatial_mesh_size() [xmin, ymin, zmin, xmax, ymax, zmax] = vlsvReader.get_spatial_mesh_extent() # Calculate cell lengths: cell_lengths = np.array([(xmax - xmin)/(float)(xcells), (ymax - ymin)/(float)(ycells), (zmax - zmin)/(float)(zcells)]) # Initialize lists distances = [0] cellids = [vlsvReader.get_cellid(point1)] coordinates = [point1] finalcellid = vlsvReader.get_cellid(point2) #print(" cellids init ",cellids,finalcellid) # Loop until final cellid is reached while True: newcoordinate = coordinates[-1]+cell_lengths*derivative newcellid = vlsvReader.get_cellid( newcoordinate ) distances.append( np.linalg.norm( newcoordinate - point1 ) ) coordinates.append(newcoordinate) cellids.append(newcellid) #print(distances[-1]/Re,np.array(coordinates[-1])/Re,cellids[-1]) if newcellid==finalcellid: break if distances[-1]>largestdistance: break # 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"] )