import numpy as np
import matplotlib.pyplot as plt
import matplotlib
import uncertainties
import uncertainties.unumpy as unp
import sympy
import matplotlib.pylab as pylab
import itertools
# local imports
from smpl import interpolate
from smpl import io
from smpl import util
from smpl import wrap
from smpl import doc
from smpl import fit as ffit
def set_plot_style():
# fig_size = (8, 6)
# fig_legendsize = 14
# fig_labelsize = 12 # ‘xx-small’, ‘x-small’, ‘small’, ‘medium’, ‘large’, ‘x-large’, ‘xx-large’.
params = {'legend.fontsize': 'x-large',
'figure.figsize': (8, 6),
'axes.labelsize': 'x-large',
'axes.titlesize': 'x-large',
'xtick.labelsize': 'x-large',
'ytick.labelsize': 'x-large'}
pylab.rcParams.update(params)
matplotlib.rcParams.update(params)
# matplotlib.rcParams.update({'font.size': fig_labelsize})
# colors = dict(mcolors.BASE_COLORS, **mcolors.CSS4_COLORS)
unv = unp.nominal_values
usd = unp.std_devs
default = {
# 'params' :[None ,"Initial fit parameters",
'xlabel': ["", "X axis label", ],
'ylabel': ["", "Y axis label", ],
'label': [None, "Legend name of plotted ``data``", ],
'fmt': ['.', "Format for plotting fit function", ],
'units': [None, "Units of the fit parameters as strings. Displayed in the Legend", ],
'save': [None, " File to save the plot", ],
'lpos': [0, "Legend position", ],
'tight': [True, "tight_layout", ],
# 'frange' :[None ,"Limit the fit to given range. First integer is the lowest and second the highest index.",],
'prange': [None, "Limit the plot of the fit to given range", ],
'sigmas': [0, "Color the array of given ``sigma`` times uncertainty. Only works if the fit function is coded with ``unp``", ],
'data_sigmas': [1, "Color the array of given ``sigma`` times uncertainty. Only works if the data has uncertainties", ],
'init': [False, "Initialize a new plot"],
'ss': [True, "save, add legends and grid to the plot", ],
'also_data': [True, " also plot the data"],
'also_fit': [True, "also plot the fit", ],
'auto_fit': [False, "automatically fit", ],
'logy': [False, "logarithmic x axis", ],
'logx': [False, "logarithmic y axis", ],
'function_color': [None, "Color of the function plot", ],
'data_color': [None, "Color of the data plot", ],
'fit_color': [None, "Color of the fit plot", ],
'fit_fmt': ["-", "Format of the fit plot", ],
'residue': [False, "Display difference between fit and data in a second plot", ],
'residue_err': [True, "Differences between fit and data will have errorbars", ],
'show': [False, "Call plt.show()", ],
'size': [None, "Size of the plot as a tuple (x,y). Only has an effect if ``init`` is True", ],
'number_format': [io.gf(4), "Format to display numbers.", ],
# , 'selector' :[ None ,"Function that takes ``x`` and ``y`` as parameters and returns an array mask in order to limit the data points for fitting. Alternatively a mask for selecting elements from datax and datay.",],
# , 'fixed_params' :[ True ,"Enable fixing parameters by choosing the same-named variables from ``kwargs``.",],
# , 'sortbyx' :[ True , "Enable sorting the x and y data so that x is sorted.",],
'interpolate': [False, "Enable interpolation of the data."],
'interpolate_fmt': ["-", "Either format string or linestyle tuple.", ],
'interpolate_label': [None, "Label for the interpolation.", ],
'extrapolate': [True, "Enable extrapolation of whole data if fit range is limited by ``frange`` or ``fselector``.", ],
'extrapolate_min': [None, "Lower extrapolation bound", ],
'extrapolate_max': [None, "Higher extrapolation bound", ],
'extrapolate_fmt': ["--", "Format of the extrapolation line", ],
'extrapolate_hatch': [r"||", "Extrapolation shape/hatch for filled area in case of ``sigmas``>0. See https://matplotlib.org/stable/gallery/shapes_and_collections/hatch_style_reference.html", ],
'bbox_to_anchor': [None, "Position in a tuple (x,y),Shift position of the legend out of the main pane. ", ],
'ncol': [None, "Columns in the legend if used with ``bbox_to_anchor``.", ],
'steps': [1000, "resolution of the plotted function", ],
'fitinline': [False, "No newlines for each fit parameter", ],
'grid': [True, "Enable grid for the plot", ],
'hist': [False, "Enable histogram plot", ],
'stairs': [False, "Enable stair plot", ],
'capsize': [5, "size of cap on error bar plot"],
'axes': [None, "set current axis"],
'linestyle': [None, "linestyle, only active if `fmt`=None"],
'xspace': [np.linspace, "xspace gets called with xspace(xmin,xmax,steps) in :func:`function` to get the points of the function that will be drawn."]
}
[docs]@doc.append_doc(ffit.fit_kwargs)
@doc.append_str("\t")
@doc.append_str(doc.table(default, init=False))
@doc.append_str(doc.table({"plot_kwargs ": ["default", "description"]}, bottom=False))
def plot_kwargs(kwargs):
"""Set default plot_kwargs if not set.
"""
kwargs = ffit.fit_kwargs(kwargs)
for k, v in default.items():
if not k in kwargs:
kwargs[k] = v[0]
return kwargs
# TODO optimize to minimize number of calls to function
# @append_doc(default_kwargs)
[docs]def fit(func,*adata, **kwargs):
"""
Fit and plot function to datax and datay.
Parameters
----------
datax : array_like
X data either as ``unp.uarray`` or ``np.array`` or ``list``
datay : array_like
Y data either as ``unp.uarray`` or ``np.array`` or ``list``
function : func
Fit function with parameters: ``x``, ``params``
**kwargs : optional
see :func:`plot_kwargs`.
Fit parameters can be fixed via ``kwargs`` eg. ``a=5``.
Returns
-------
array_like
Optimized fit parameters of ``function`` to ``datax`` and ``datay``.
If ``datay`` is complex, both the real and imaginary part are returned.
Examples
--------
.. plot::
:include-source:
>>> from smpl import functions as f
>>> from smpl import plot
>>> param = plot.fit([0,1,2],[0,1,2],f.line)
>>> plot.unv(param).round()[0]
1.0
"""
function = func
if 'function' in kwargs:
function = kwargs['function']
del kwargs['function']
adata = [func,*adata]
# Fix parameter order if necessary
elif isinstance(function, (list, tuple, np.ndarray)):
adata=[adata[-1],function,*adata[:-1]]
function = adata[0]
adata = adata[1:]
if util.true("bins",kwargs):
# yvalue will be overwritten
ndata = [*adata,*adata]
for i,o in enumerate(adata):
ndata[2*i] =o
ndata[2*i+1] =o*0
adata= ndata
assert len(adata)%2==0, "data must be pairs of x and y data"
if len(adata) ==2:
datax, datay = adata
else:
rs = []
for i in range(0, len(adata), 2):
datax, datay = adata[i], adata[i+1]
if util.true("bins",kwargs):
rs.append(fit(function,datax, **kwargs))
else:
rs.append(fit(function,datax, datay, **kwargs))
return rs
kwargs = plot_kwargs(kwargs)
if np.any(np.iscomplex(datay)):
label = util.get("label",kwargs, "")
kwargs['label'] = label+ "(real)"
r= fit(datax, datay.real, function=function,**kwargs)
kwargs['label'] = label + "(imag)"
i= fit(datax, datay.imag, function=function,**kwargs)
return r,i
if kwargs['auto_fit']:
print("function:",function)
best_f, best_ff, lambda_f = ffit.auto(datax, datay, function, **kwargs)
if best_f is not None:
del kwargs['auto_fit']
fit(datax, datay, best_f, **kwargs)
return best_f, best_ff, lambda_f
if kwargs['also_fit'] == False and kwargs['label'] == None and kwargs['lpos'] == 0:
kwargs['lpos'] = -1
return _fit_impl(datax, datay, function, **kwargs)
def _fit_impl(datax,datay,function, **kwargs):
x = None
y = None
rfit = None
ifit = None
fig = None
fig = init_plot(kwargs)
ll = None
if kwargs['also_data']:
ll = plt_data(datax, datay, **kwargs).get_color()
if kwargs['interpolate']:
ifit, _, x, y = plt_interpolate(datax, datay, icolor=ll, **kwargs)
if kwargs['also_fit']:
assert function is not None, "function must be given"
rfit, kwargs['fit_color'], _, _ = plt_fit(
datax, datay, function, **kwargs)
if kwargs['ss']:
kwargs['oldshow'] = kwargs['show']
kwargs['show'] = kwargs['show'] and not kwargs['residue']
save_plot(**kwargs)
kwargs['show'] = kwargs['oldshow']
if kwargs['residue'] and fig is not None:
plt_residue(datax, datay, function, rfit, fig, **kwargs)
if not kwargs["also_fit"] and kwargs["interpolate"]:
return (ifit, x, y)
# return ifit
return rfit
# @append_doc(default_kwargs)
[docs]def data(*data, function=None, **kwargs):
"""
Plot datay against datax via :func:`fit`
Parameters
----------
datax : array_like
X data either as ``unp.uarray`` or ``np.array`` or ``list``
datay : array_like
Y data either as ``unp.uarray`` or ``np.array`` or ``list``
function : func,optional
Fit function with parameters: ``x``, ``params``
**kwargs : optional
see :func:`plot_kwargs`.
Returns
-------
array_like
Optimized fit parameters of ``function`` to ``datax`` and ``datay``
"""
if 'also_fit' not in kwargs:
kwargs['also_fit'] = False
kwargs = plot_kwargs(kwargs)
return fit(function=function,*data, **kwargs)
# @append_doc(default_kwargs)
[docs]def auto(*adata, funcs=None, **kwargs):
"""
Automatically loop over functions and fit the best one.
Parameters
----------
funcs : function array
functions to consider as fit. Default all ``smpl.functions``.
**kwargs : optional
see :func:`plot_kwargs`.
Returns
-------
The best fit function and it's parameters. Also a lambda function where the parameters are already applied.
"""
if 'auto_fit' not in kwargs:
kwargs['auto_fit'] = True
kwargs = plot_kwargs(kwargs)
return fit(function=funcs,*adata, **kwargs)
def _function(func, xfit, **kwargs):
kargs = {}
if util.has('fmt', kwargs):
kargs["fmt"] = kwargs["fmt"]
if util.has('label', kwargs) and kwargs['label'] != "":
kargs['label'] = kwargs['label']
if util.has('function_color', kwargs) and kwargs['function_color'] != "":
kargs['color'] = kwargs['function_color']
if util.has('sigmas', kwargs) and kwargs['sigmas'] != "":
kargs['sigmas'] = kwargs['sigmas']
__function(func, xfit, **kargs)
def plt_plt(x, y, fmt, color, label, linestyle):
if linestyle is None and fmt is not None:
return plt.plot(x, y, fmt, label=label, color=color)
elif linestyle is not None and fmt is None:
return plt.plot(x, y, label=label, color=color, linestyle=linestyle)
elif linestyle is None and fmt is None:
return plt.plot(x, y, label=label, color=color)
def __function(gfunc, xlinspace, fmt="-", label=None, color=None, hatch=None, sigmas=0., linestyle=None):
func = gfunc
x = xlinspace
l = label
if isinstance(func(x)[0], uncertainties.UFloat):
if sigmas > 0:
ll, = plt_plt(x, unv(func(x)), fmt, label=None,
color=color, linestyle=linestyle)
y = func(x)
plt.fill_between(x, unv(y)-sigmas*usd(y), unv(
y)+sigmas*usd(y), alpha=0.4, label=l, color=ll.get_color(), hatch=hatch)
else:
ll, = plt_plt(x, unv(func(x)), fmt, label=l,
color=color, linestyle=linestyle)
else:
ll, = plt_plt(x, func(x), fmt, label=l,
color=color, linestyle=linestyle)
return ll
[docs]def function(func, *args, **kwargs):
"""
Plot function ``func`` between ``xmin`` and ``xmax``
Parameters
----------
func : function
Function to be plotted between ``xmin`` and ``xmax``, only taking `array_like` ``x`` as parameter
*args : optional
arguments for ``func``
**kwargs : optional
see :func:`plot_kwargs`.
"""
if not util.has("xmin", kwargs) or not util.has("xmin", kwargs):
raise Exception("xmin or xmax missing.")
# if not util.has('lpos', kwargs) and not util.has('label', kwargs):
# kwargs['lpos'] = -1
if not util.has('fmt', kwargs):
kwargs['fmt'] = "-"
if "label" not in kwargs:
kwargs = plot_kwargs(kwargs)
kwargs['label'] = get_fnc_legend(func, args, **kwargs)
else:
kwargs = plot_kwargs(kwargs)
xlin = kwargs["xspace"](kwargs['xmin'], kwargs['xmax'], kwargs['steps'])
init_plot(kwargs)
# kwargs['lpos'] = 0
#_plot(xfit, func(xfit, *args), **kwargs)
_function(wrap.get_lambda_argd(
func, kwargs['xvar'], *args), xlin, **kwargs)
if kwargs['ss']:
save_plot(**kwargs)
# xaxis="",yaxis="",fit_color=None,save = None,residue_err=True,show=False):
def plt_residue(datax, datay, gfunction, rfit, fig, **kwargs):
function = wrap.get_lambda(gfunction, kwargs['xvar'])
fig.add_axes((.1, .1, .8, .2))
kwargs['yaxis'] = "$\\Delta$" + kwargs['yaxis']
kwargs['data_color'] = kwargs['fit_color']
if kwargs['residue_err']:
plt_data(datax, datay-function(datax, *rfit), **kwargs)
else:
plt_data(unv(datax), unv(datay-function(datax, *rfit)), **kwargs)
kwargs['lpos'] = -1
save_plot(**kwargs)
def data_split(datax, datay, **kwargs):
return ffit.data_split(datax, datay, **kwargs)
def _fit(datax, datay, function, **kwargs):
"""
Returns a fit like :func:`fit` but does no plotting.
"""
return ffit.fit(datax, datay, function, **kwargs)
def plt_data(datax, datay, **kwargs):
"""
Plot datay vs datax
"""
x, y, xerr, yerr = data_split(datax, datay, **kwargs)
if xerr is not None:
xerr = xerr * kwargs['data_sigmas']
if yerr is not None:
yerr = yerr * kwargs['data_sigmas']
ll = None
if xerr is None and yerr is None:
if kwargs['fmt'] is None:
if kwargs['linestyle'] is None:
ll, = plt.plot(
x, y, label=kwargs['label'], color=kwargs['data_color'])
else:
ll, = plt.plot(
x, y, label=kwargs['label'], color=kwargs['data_color'], linestyle=kwargs['linestyle'])
elif kwargs['fmt'] == "step":
ll, = plt.step(x, y, where='mid',
label=kwargs['label'], color=kwargs['data_color'])
elif kwargs['fmt'] == "hist":
ll, = plt.step(x, y, where='mid',
label=kwargs['label'], color=kwargs['data_color'])
plt.fill_between(x, y, step="mid", color=ll.get_color())
else:
ll, = plt.plot(x, y, kwargs['fmt'], label=kwargs['label'],
color=kwargs['data_color'])
else:
if kwargs['fmt'] is None:
if kwargs['linestyle'] is None:
ll, _, _, = plt.errorbar(x, y, yerr=yerr, xerr=xerr, fmt=" ", capsize=kwargs["capsize"],
label=kwargs['label'], color=kwargs['data_color'])
else:
ll, _, _, = plt.errorbar(x, y, yerr=yerr, xerr=xerr, fmt=" ", capsize=kwargs["capsize"],
label=kwargs['label'], color=kwargs['data_color'], linestyle=kwargs['linestyle'])
elif kwargs['fmt'] == "step":
ll, = plt.step(x, y, where='mid',
color=kwargs['data_color'])
if xerr is not None:
for ix, xv in enumerate(x):
dx = (xerr[ix])
tx = [xv-dx, xv+dx]
plt.fill_between(tx, y[ix]-yerr[ix], y[ix]+yerr[ix],
label=kwargs['label']if ix == 1 else None, alpha=0.2, step='pre', color=ll.get_color())
else:
plt.fill_between(x, y-yerr, y+yerr,
label=kwargs['label'], alpha=0.2, step='mid', color=ll.get_color())
elif kwargs['fmt'] == "hist":
ll, _, _, = plt.errorbar(x, y, yerr=yerr, xerr=xerr, fmt=" ", capsize=kwargs["capsize"],
color="black")
plt.fill_between(x, y, step="mid",
label=kwargs['label'], color=ll.get_color())
else:
ll, _, _, = plt.errorbar(x, y, yerr=yerr, xerr=xerr, fmt=kwargs['fmt'], capsize=kwargs["capsize"],
label=kwargs['label'], color=kwargs['data_color'])
return ll
def get_fnc_legend(function, rfit, **kwargs):
l = wrap.get_latex(function)
vnames = wrap.get_varnames(function, kwargs['xvar'])
for i in range(1, len(vnames)):
l = l + ("\n" if not kwargs["fitinline"] or i == 1 else " ")
l = l + "$" + sympy.latex(sympy.symbols(str(vnames[i]))) + "$="
if kwargs['units'] is not None and usd(rfit[i-1]) > 0:
l = l + "("
if 'number_format' in kwargs:
l = l + kwargs['number_format'].format(rfit[i-1])
else:
l = l + "%s" % (rfit[i-1])
if kwargs['units'] is not None and usd(rfit[i-1]) > 0:
l = l + ")"
if kwargs['units'] is not None:
l = l + " " + kwargs['units'][i-1]
return l
def plt_fit_or_interpolate(datax, datay, fitted, l=None, c=None, f=None, ls=None, **kwargs):
if kwargs['prange'] is None:
x, _, _, _ = ffit.fit_split(datax, datay, **kwargs)
xfit = kwargs['xspace'](
np.min(unv(x)), np.max(unv(x)), kwargs['steps'])
else:
xfit = kwargs['xspace'](kwargs['prange'][0],
kwargs['prange'][1], kwargs['steps'])
ll = __function(fitted, xfit, kwargs['fit_fmt'] if f is not None and ls is None else f, label=l,
color=kwargs['fit_color'] if c is None else c, sigmas=kwargs['sigmas'], linestyle=ls)
if (kwargs['frange'] is not None or kwargs['fselector'] is not None) and util.true('extrapolate', kwargs) or util.has("extrapolate_max", kwargs) or util.has("extrapolate_min", kwargs):
xxfit = kwargs['xspace'](util.get("extrapolate_min", kwargs, np.min(
unv(datax))), util.get("extrapolate_max", kwargs, np.max(unv(datax))), kwargs['steps'])
__function(fitted, kwargs['xspace'](np.min(xxfit), np.min(xfit), kwargs['steps']), util.get("extrapolate_fmt", kwargs, "--"),
color=ll.get_color(), hatch=util.get("extrapolate_hatch", kwargs, r"||"), sigmas=kwargs['sigmas'])
__function(fitted, kwargs['xspace'](np.max(xfit), np.max(xxfit), kwargs['steps']), util.get("extrapolate_fmt", kwargs, "--"),
color=ll.get_color(), hatch=util.get("extrapolate_hatch", kwargs, r"||"), sigmas=kwargs['sigmas'])
return ll.get_color(), xfit, fitted(xfit)
def plt_interpolate(datax, datay, icolor=None, **kwargs):
"""
Interpolate and Plot that Interpolation.
"""
inter = interpolate.interpolate(datax, datay, **kwargs)
kargs = {}
if isinstance(kwargs["interpolate_fmt"], tuple):
kargs["ls"] = kwargs["interpolate_fmt"]
else:
kargs["f"] = kwargs["interpolate_fmt"]
if kwargs['interpolate_label'] is not None:
kargs["l"] = kwargs['interpolate_label']
# l = None so that no label
return (inter, *plt_fit_or_interpolate(datax, datay, inter, c=icolor, **kargs, **kwargs))
def plt_fit(datax, datay, gfunction, **kwargs):
"""
Fit and Plot that Fit.
"""
func = wrap.get_lambda(gfunction, kwargs['xvar'])
rfit = _fit(datax, datay, gfunction, **kwargs)
def fitted(x): return func(x, *rfit)
l = get_fnc_legend(gfunction, rfit, **kwargs)
return (rfit, *plt_fit_or_interpolate(datax, datay, fitted, l, **kwargs))
def init_plot(kwargs):
fig = None
if util.has("axes", kwargs) and kwargs["axes"] is not None:
plt.sca(kwargs["axes"])
fig = kwargs["axes"].get_figure()
if kwargs['init'] or util.true("residue", kwargs):
if kwargs['size'] is None:
fig = plt.figure()
else:
fig = plt.figure(figsize=kwargs['size'])
if kwargs['residue']:
fig.add_axes((.1, .3, .8, .6))
if util.has("xlabel", kwargs) and kwargs['xlabel'] != "":
plt.xlabel(kwargs['xlabel'])
if util.has("ylabel", kwargs) and kwargs['ylabel'] != "":
plt.ylabel(kwargs['ylabel'])
if util.has("xaxis", kwargs) and kwargs['xaxis'] != "":
plt.xlabel(kwargs['xaxis'])
if util.has("yaxis", kwargs) and kwargs['yaxis'] != "":
plt.ylabel(kwargs['yaxis'])
if util.has("next_color", kwargs) and not kwargs['next_color']:
it1, it2 = itertools.tee(iter(plt.gca()._get_lines.prop_cycler))
plt.gca()._get_lines.prop_cycler = it2
tmp_color = next(it1)['color']
if kwargs['data_color'] is None:
kwargs['data_color'] = tmp_color
if kwargs['fit_color'] is None:
kwargs['fit_color'] = tmp_color
if kwargs['function_color'] is None:
kwargs['function_color'] = tmp_color
return fig
def save_plot(**kwargs):
"""
save plot
"""
if 'logy' in kwargs and kwargs['logy']:
plt.gca().set_yscale('log')
if 'logx' in kwargs and kwargs['logx']:
plt.gca().set_xscale('log')
if 'tight' in kwargs and kwargs['tight']:
plt.tight_layout()
if 'lpos' in kwargs and kwargs['lpos'] >= 0:
if(util.has('bbox_to_anchor', kwargs)):
if(util.has('ncol', kwargs)):
plt.legend(loc=kwargs['lpos'], bbox_to_anchor=kwargs['bbox_to_anchor'],
ncol=kwargs['ncol'], borderaxespad=0)
else:
plt.legend(loc=kwargs['lpos'],
bbox_to_anchor=kwargs['bbox_to_anchor'])
else:
plt.legend(loc=kwargs['lpos'])
# plt.gca().set_xlim([kwargs['xmin'],kwargs['xmax']])
# plt.gca().set_ylim([kwargs['ymin'],kwargs['ymax']])
if 'save' in kwargs and not kwargs['save'] == None:
io.mkdirs(kwargs['save'])
plt.savefig(kwargs['save'] + ".pdf")
plt.grid(b=kwargs["grid"])
if 'show' in kwargs and kwargs['show']:
show(**kwargs)
def show(**kwargs):
kwargs = plot_kwargs(kwargs)
plt.grid(b=kwargs["grid"])
plt.show()
if __name__ == "__main__":
import doctest
doctest.testmod()