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 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