Ejemplo n.º 1
0
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)] )
Ejemplo n.º 2
0
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)])
Ejemplo n.º 3
0
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)])
Ejemplo n.º 4
0
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)])
Ejemplo n.º 5
0
def plot_variables( x, y, figure=[] ):
   ''' Plots x and y variables from the input with pylab

       :param x:        Some variable to be plotted in the x-axis
       :param y:        Some variable to be plotted in the y-axis
       :param figure:   If one wants to plot into an existing figure then the matplotlib figure should be passed as an argument (OPTIONAL)
       :returns: a pylab figure with the plot

       .. code-block:: python

          # Example usage:
          plot_variables( distances, rho )
          # This would plot rho as a function of distance
   '''   
   # Get dimensions of x and y
   import numpy as np
   x_dim = len(np.shape(x))
   y_dim = len(np.shape(y))
   if x_dim != y_dim:
      if x_dim == y_dim - 1:
         from variable import get_data, get_name, get_units
         new_x = [get_data(x) for i in xrange(len(y)-1)]
         new_x.append(x)
         return plot_multiple_variables( new_x, y, figure, clean_xticks=True )
      else:
         print "ERROR; BAD X AND Y DIMENSIONS " + str(x_dim) + " " + str(y_dim)
         return []
   else:
      if x_dim == 0 and y_dim == 0:
         return plot_multiple_variables( [x], [y], figure )
      else:
         return plot_multiple_variables( x, y, figure )
Ejemplo n.º 6
0
def plot_variables( x, y, figure=[] ):
   ''' Plots x and y variables from the input with pylab

       :param x:        Some variable to be plotted in the x-axis
       :param y:        Some variable to be plotted in the y-axis
       :param figure:   If one wants to plot into an existing figure then the matplotlib figure should be passed as an argument (OPTIONAL)
       :returns: a pylab figure with the plot

       .. code-block:: python

          # Example usage:
          plot_variables( distances, rho )
          # This would plot rho as a function of distance
   '''   
   # Get dimensions of x and y
   import numpy as np
   x_dim = len(np.shape(x))
   y_dim = len(np.shape(y))
   if x_dim != y_dim:
      if x_dim == y_dim - 1:
         from variable import get_data, get_name, get_units
         new_x = [get_data(x) for i in xrange(len(y)-1)]
         new_x.append(x)
         return plot_multiple_variables( new_x, y, figure, clean_xticks=True )
      else:
         print "ERROR; BAD X AND Y DIMENSIONS " + str(x_dim) + " " + str(y_dim)
         return []
   else:
      if x_dim == 0 and y_dim == 0:
         return plot_multiple_variables( [x], [y], figure )
      else:
         return plot_multiple_variables( x, y, figure )
Ejemplo n.º 7
0
def plot_multiple_variables( variables_x_list, variables_y_list, figure=[], clean_xticks=False ):
   ''' Plots multiple variables from the input with pylab

       :param variables_x_list:        Some list of variables to be plotted in the x-axis
       :param variables_y_list:        Some list of variables to be plotted in the y-axis
       :param figure:                  If one wants to plot into an existing figure then the matplotlib figure should be passed as an argument (OPTIONAL)
       :returns: a pylab figure with the plot

       .. code-block:: python

          #Example usage:
          plot_multiple_variables( [distances, xcoordinates], [rho, B_x] )
          # This would plot rho_values as a function of distance and B_x_values as a function of xcoordinates

       .. note:: Multiplot expects variables to be saved in the VariableInfo class

       .. note:: If for some reason some variable list (x or y) is empty, e.g. variables_x_list = [B_x, [], B_z, rho], then the variable will not be plotted. This can be used if one wants to plot only into certain subplots.
   '''
   yticks = {}
   for i in xrange(18):
      tick = i+1
      yticks[tick] = 7 - (int)(i)/(int)(4)


   import numpy as np
   variables_x_list = np.ma.asarray(variables_x_list)
   variables_y_list = np.ma.asarray(variables_y_list)
   if len(variables_x_list) != len(variables_y_list):
      # Attempt to fix the lengths:
      if (len(variables_x_list) == 1):
         if (len(np.atleast_1d(variables_x_list[0])) == len(variables_y_list)):
            variables_y_list = [variables_y_list]
   
      if (len(variables_y_list) == 1):
         if (len(np.atleast_1d(variables_y_list[0])) == len(variables_x_list)):
            variables_x_list = [variables_x_list]

   if len(variables_x_list) != len(variables_y_list):
      print "BAD VARIABLE LENGTH: " + str(len(variables_x_list)) + " " + str(len(variables_y_list))
      return []
   if len(variables_y_list) > 18:
      print "TOO MANY VARIABLES: " + str(len(variables_y_list))
      return []
      
   length_of_list = len(variables_x_list)

   if figure != []:
      fig = pl.figure
      if len(fig.get_axes()) < length_of_list:
         for i in (np.arange(length_of_list-len(fig.get_axes())) + len(fig.get_axes())):
            fig.add_subplot(length_of_list,1,i)
   else:
      fig = pl.figure()
      for i in xrange(length_of_list):
         fig.add_subplot(length_of_list,1,i+1)

   axes = fig.get_axes()
   from variable import get_data, get_name, get_units
   for i in xrange(length_of_list):
      
      x = variables_x_list[i]
      y = variables_y_list[i]
      # Check the length of the list
      if (len(np.atleast_1d(x)) == 0) or (len(np.atleast_1d(y)) == 0):
         continue
      ax = axes[i]
      ax.plot(get_data(x), get_data(y), lw=2)

      if get_units(x) != "":
         ax.set_xlabel(get_name(x) + " [" + get_units(x) + "]")
      else:
         ax.set_xlabel(get_name(x))

      if get_units(y) != "":
         ax.set_ylabel(get_name(y) + " [" + get_units(y) + "]")
      else:
         ax.set_ylabel(get_name(y))

      # Set limits
      xlength = np.max(get_data(x)) - np.min(get_data(x))
      ylength = np.max(get_data(y)) - np.min(get_data(y))
      ax.set_xlim([np.min(get_data(x)) - 0.01*xlength, np.max(get_data(x)) + 0.01*xlength])
      ax.set_ylim([np.min(get_data(y)) - 0.05*ylength, np.max(get_data(y)) + 0.05*ylength])      
      # Set format
      ax.ticklabel_format(style='sci', axis='y', scilimits=(-3,3))

   if clean_xticks == True:
      for i in xrange(len(np.atleast_1d(axes))-1):
         axes[i].set_xticks([])

   # Set yticks:
   fig = set_yticks( fig, yticks[len(axes)] )
   return fig
Ejemplo n.º 8
0
def plot_multiple_variables( variables_x_list, variables_y_list, figure=[], clean_xticks=False ):
   ''' Plots multiple variables from the input with pylab

       :param variables_x_list:        Some list of variables to be plotted in the x-axis
       :param variables_y_list:        Some list of variables to be plotted in the y-axis
       :param figure:                  If one wants to plot into an existing figure then the matplotlib figure should be passed as an argument (OPTIONAL)
       :returns: a pylab figure with the plot

       .. code-block:: python

          #Example usage:
          plot_multiple_variables( [distances, xcoordinates], [rho, B_x] )
          # This would plot rho_values as a function of distance and B_x_values as a function of xcoordinates

       .. note:: Multiplot expects variables to be saved in the VariableInfo class

       .. note:: If for some reason some variable list (x or y) is empty, e.g. variables_x_list = [B_x, [], B_z, rho], then the variable will not be plotted. This can be used if one wants to plot only into certain subplots.
   '''
   yticks = {}
   for i in xrange(18):
      tick = i+1
      yticks[tick] = 7 - (int)(i)/(int)(4)


   import numpy as np
   variables_x_list = np.ma.asarray(variables_x_list)
   variables_y_list = np.ma.asarray(variables_y_list)
   if len(variables_x_list) != len(variables_y_list):
      # Attempt to fix the lengths:
      if (len(variables_x_list) == 1):
         if (len(np.atleast_1d(variables_x_list[0])) == len(variables_y_list)):
            variables_y_list = [variables_y_list]
   
      if (len(variables_y_list) == 1):
         if (len(np.atleast_1d(variables_y_list[0])) == len(variables_x_list)):
            variables_x_list = [variables_x_list]

   if len(variables_x_list) != len(variables_y_list):
      print "BAD VARIABLE LENGTH: " + str(len(variables_x_list)) + " " + str(len(variables_y_list))
      return []
   if len(variables_y_list) > 18:
      print "TOO MANY VARIABLES: " + str(len(variables_y_list))
      return []
      
   length_of_list = len(variables_x_list)

   if figure != []:
      fig = pl.figure
      if len(fig.get_axes()) < length_of_list:
         for i in (np.arange(length_of_list-len(fig.get_axes())) + len(fig.get_axes())):
            fig.add_subplot(length_of_list,1,i)
   else:
      fig = pl.figure()
      for i in xrange(length_of_list):
         fig.add_subplot(length_of_list,1,i+1)

   axes = fig.get_axes()
   from variable import get_data, get_name, get_units
   for i in xrange(length_of_list):
      
      x = variables_x_list[i]
      y = variables_y_list[i]
      # Check the length of the list
      if (len(np.atleast_1d(x)) == 0) or (len(np.atleast_1d(y)) == 0):
         continue
      ax = axes[i]
      ax.plot(get_data(x), get_data(y), lw=2)

      if get_units(x) != "":
         ax.set_xlabel(get_name(x) + " [" + get_units(x) + "]")
      else:
         ax.set_xlabel(get_name(x))

      if get_units(y) != "":
         ax.set_ylabel(get_name(y) + " [" + get_units(y) + "]")
      else:
         ax.set_ylabel(get_name(y))

      # Set limits
      xlength = np.max(get_data(x)) - np.min(get_data(x))
      ylength = np.max(get_data(y)) - np.min(get_data(y))
      ax.set_xlim([np.min(get_data(x)) - 0.01*xlength, np.max(get_data(x)) + 0.01*xlength])
      ax.set_ylim([np.min(get_data(y)) - 0.05*ylength, np.max(get_data(y)) + 0.05*ylength])      
      # Set format
      ax.ticklabel_format(style='sci', axis='y', scilimits=(-3,3))

   if clean_xticks == True:
      for i in xrange(len(np.atleast_1d(axes))-1):
         axes[i].set_xticks([])

   # Set yticks:
   fig = set_yticks( fig, yticks[len(axes)] )
   return fig
Ejemplo n.º 9
0
    def __picker_callback(self, picker):
        """ This gets called when clicking on a cell
      """
        if (self.picker != "Cut_through"):
            # Make sure the last pick is null (used in cut_through)
            self.__last_pick = []

        coordinates = picker.pick_position
        coordinates = np.array(
            [coordinates[0], coordinates[1], coordinates[2]])
        # For numerical inaccuracy
        epsilon = 80
        # Check for numerical inaccuracy
        for i in xrange(3):
            if (coordinates[i] < self.__mins[i]) and (coordinates[i] + epsilon
                                                      > self.__mins[i]):
                # Correct the numberical inaccuracy
                coordinates[i] = self.__mins[i] + 1
            if (coordinates[i] > self.__maxs[i]) and (coordinates[i] - epsilon
                                                      < self.__maxs[i]):
                # Correct the values
                coordinates[i] = self.__maxs[i] - 1
        print "COORDINATES:" + str(coordinates)
        cellid = self.__vlsvReader.get_cellid(coordinates)
        print "CELL ID: " + str(cellid)
        # Check for an invalid cell id
        if cellid == 0:
            print "Invalid cell id"
            return

        if (self.picker == "Velocity_space"):
            # Set label to give out the location of the cell:
            self.__add_label(cellid)
            # Generate velocity space
            self.__generate_velocity_grid(cellid)
        elif (self.picker == "Velocity_space_nearest_cellid"):
            # Find the nearest cell id with distribution:
            # Read cell ids with velocity distribution in:
            cell_candidates = self.__vlsvReader.read("SpatialGrid",
                                                     "CELLSWITHBLOCKS")
            # Read in the coordinates of the cells:
            cell_candidate_coordinates = [
                self.__vlsvReader.get_cell_coordinates(cell_candidate)
                for cell_candidate in cell_candidates
            ]
            # Read in the cell's coordinates:
            pick_cell_coordinates = self.__vlsvReader.get_cell_coordinates(
                cellid)
            # Find the nearest:
            from operator import itemgetter
            norms = np.sum(
                (cell_candidate_coordinates - pick_cell_coordinates)**2,
                axis=-1)**(1. / 2)
            norm, i = min((norm, idx) for (idx, norm) in enumerate(norms))
            # Get the cell id:
            cellid = cell_candidates[i]
            # Set label to give out the location of the cell:
            self.__add_label(cellid)
            # Generate velocity grid
            self.__generate_velocity_grid(cellid)
        elif (self.picker == "Velocity_space_iso_surface"):
            # Set label to give out the location of the cell:
            self.__add_label(cellid)
            self.__generate_velocity_grid(cellid, True)
        elif (self.picker == "Velocity_space_nearest_cellid_iso_surface"):
            # Find the nearest cell id with distribution:
            # Read cell ids with velocity distribution in:
            cell_candidates = self.__vlsvReader.read("SpatialGrid",
                                                     "CELLSWITHBLOCKS")
            # Read in the coordinates of the cells:
            cell_candidate_coordinates = [
                self.__vlsvReader.get_cell_coordinates(cell_candidate)
                for cell_candidate in cell_candidates
            ]
            # Read in the cell's coordinates:
            pick_cell_coordinates = self.__vlsvReader.get_cell_coordinates(
                cellid)
            # Find the nearest:
            from operator import itemgetter
            norms = np.sum(
                (cell_candidate_coordinates - pick_cell_coordinates)**2,
                axis=-1)**(1. / 2)
            norm, i = min((norm, idx) for (idx, norm) in enumerate(norms))
            # Get the cell id:
            cellid = cell_candidates[i]
            # Set label to give out the location of the cell:
            self.__add_label(cellid)
            # Generate velocity grid
            self.__generate_velocity_grid(cellid, True)
        elif (self.picker == "Pitch_angle"):
            # Set label to give out the location of the cell:
            self.__add_label(cellid)
            # Plot pitch angle distribution:
            from pitchangle import pitch_angles
            result = pitch_angles(vlsvReader=self.__vlsvReader,
                                  cellid=cellid,
                                  cosine=True,
                                  plasmaframe=True)
            # plot:
            pl.hist(result[0].data, weights=result[1].data, bins=50, log=False)
            pl.show()
        elif (self.picker == "Gyrophase_angle"):
            # Plot gyrophase angle distribution:
            from gyrophaseangle import gyrophase_angles_from_file
            result = gyrophase_angles_from_file(vlsvReader=self.__vlsvReader,
                                                cellid=cellid)
            # plot:
            pl.hist(result[0].data,
                    weights=result[1].data,
                    bins=36,
                    range=[-180.0, 180.0],
                    log=True,
                    normed=1)
            pl.show()
        elif (self.picker == "Cut_through"):
            if len(self.__last_pick) == 3:
                from cutthrough import cut_through
                # Get a cut-through
                self.cut_through = cut_through(self.__vlsvReader,
                                               point1=self.__last_pick,
                                               point2=coordinates)
                # Get cell ids and distances separately
                cellids = self.cut_through[0].data
                distances = self.cut_through[1]
                # Get any arguments from the user:
                args = self.args.split()
                if len(args) == 0:
                    #Do nothing
                    print "Bad args"
                    self.__last_pick = []
                    return
                plotCut = False
                plotRankine = False
                # Optimize file read:
                self.__vlsvReader.optimize_open_file()
                variables = []
                # Save variables
                for i in xrange(len(args)):
                    # Check if the user has given the plot argument
                    if args[i] == "plot":
                        plotCut = True
                    elif args[i] == "rankine":
                        # set labels:
                        self.__add_normal_labels(point1=self.__last_pick,
                                                 point2=coordinates)
                        fig = plot_rankine(self.__vlsvReader,
                                           point1=self.__last_pick,
                                           point2=coordinates)
                        #pl.show()
                        self.__last_pick = []
                        self.plot = fig
                        return
                    else:
                        if args[i].find(",") != -1:
                            _variable = args[i].split(',')[0]
                            _operator = args[i].split(',')[1]
                            variable_info = self.__vlsvReader.read_variable_info(
                                name=_variable,
                                cellids=cellids,
                                operator=_operator)
                            variables.append(variable_info)
                            self.cut_through.append(variable_info)
                        else:
                            variable_info = self.__vlsvReader.read_variable_info(
                                name=args[i], cellids=cellids)
                            variables.append(variable_info)
                            self.cut_through.append(variable_info)
                if plotCut == True:
                    # Set label to give out the location of the cell:
                    self.__add_label(cellids[0])
                    self.__add_label(cellids[len(cellids) - 1])
                    if plotRankine == True:
                        # Plot Rankine-Hugoniot jump conditions:
                        normal_vector = (coordinates - self.__last_pick
                                         ) / np.linalg.norm(coordinates -
                                                            self.__last_pick)
                        # Read V, B, T and rho
                        V = self.__vlsvReader.read_variable("v",
                                                            cellids=cellids[0])
                        B = self.__vlsvReader.read_variable("B",
                                                            cellids=cellids[0])
                        T = self.__vlsvReader.read_variable("Temperature",
                                                            cellids=cellids[0])
                        rho = self.__vlsvReader.read_variable(
                            "rho", cellids=cellids[0])
                        # Get parallel and perpendicular components:
                        Vx = np.dot(V, normal_vector)
                        Vy = np.linalg.norm(V - Vx * normal_vector)
                        Bx = np.dot(B, normal_vector)
                        By = np.linalg.norm(B - Bx * normal_vector)
                        # Calculate jump conditions
                        conditions = oblique_shock(Vx, Vy, Bx, By, T, rho)
                        rankine_variables = []
                        for i in xrange(len(get_data(distances))):
                            if i < len(get_data(distances)) * 0.5:
                                rankine_variables.append(rho)
                            else:
                                rankine_variables.append(conditions[5])
                        variables.append(rankine_variables)
                    from plot import plot_multiple_variables
                    fig = plot_multiple_variables(
                        [distances for i in xrange(len(args) - 1)],
                        variables,
                        figure=[])
                    pl.show()
                # Close the optimized file read:
                self.__vlsvReader.optimize_close_file()
                # Read in the necessary variables:
                self.__last_pick = []
            else:
                self.__last_pick = coordinates