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