Esempio 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)])
Esempio 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)] )
Esempio 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)])
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
Esempio n. 5
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
Esempio n. 6
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)])