# -*- coding: utf-8 -*-
"""
qspec.analyze
=============
Created on 07.05.2020
@author: Patrick Mueller
Module for analyzing/evaluating/fitting data.
Linear regression algorithms:
- york(); [York et al., Am. J. Phys. 72, 367 (2004)]
- linear_monte_carlo(); based on [Gebert et al., Phys. Rev. Lett. 115, 053003 (2015), Suppl.]
- linear_monte_carlo_nd(); based on [Gebert et al., Phys. Rev. Lett. 115, 053003 (2015), Suppl.]
Curve fitting methods:
- curve_fit(); Reimplements the scipy.optimize.curve_fit method to allow fixing parameters
and having parameter-dependent y-uncertainties.
- odr_fit(); Encapsulates the scipy.odr.odr method to accept inputs similarly to curve_fit().
Classes:
- King; Creates a King plot with isotope shifts or nuclear charge radii.
LICENSE NOTES:
The method curve_fit is a modified version of scipy.optimize.curve_fit.
Therefore, it is licensed under the 'BSD 3-Clause "New" or "Revised" License' provided with scipy.
"""
import inspect
import warnings
import numpy as np
from scipy.optimize._minpack_py import Bounds, LinAlgError, OptimizeWarning, _getfullargspec, _initialize_feasible, \
_wrap_jac, cholesky, least_squares, leastsq, prepare_bounds, solve_triangular, svd
from scipy.optimize import minimize
from scipy.stats import norm, multivariate_normal
from scipy import odr
import matplotlib.pyplot as plt
from qspec._types import *
from qspec import tools
from qspec.physics import me_u, me_u_d, mass_factor
__all__ = ['straight', 'straight_slope', 'straight_std', 'straight_x_std', 'draw_straight_unc_area',
'ellipse2d', 'draw_sigma2d', 'weight', 'york', 'linear_monte_carlo', 'linear_monte_carlo_nd', 'linear_alpha',
'odr_fit', 'curve_fit', 'King']
def _get_rtype2(ab_dict: dict, a: Union[Iterable, any], b: Union[Iterable, any], rtype: type = ndarray):
"""
:param ab_dict: The nested dictionary of depth 2, where values are extracted from.
:param a: The list of keys of the first level of 'ab_dict'.
:param b: The list of keys of the second level of 'ab_dict'.
:param rtype: The type in which the results will be returned. Currently supported special types are [dict, ].
The standard 'rtype' is a numpy array.
:returns: A subset of the dictionary 'ab_dict' as a dictionary or a numpy array.
:raises: ValueError . 'a' and 'b' must be Non-Iterables or an Iterable of the corresponding types.
:raises: KeyError . ab_dict[a][b] must exist for all a and b.
"""
a_arr, b_arr = np.asarray(a), np.asarray(b)
a_dim, line_dim = len(a_arr.shape), len(b_arr.shape)
if a_dim == 0 and line_dim == 0:
if rtype == dict:
return {a: {b: ab_dict[a][b]}}
else:
return ab_dict[a][b]
elif a_dim == 0 and line_dim == 1:
if rtype == dict:
return {b_i: ab_dict[a][b_i] for b_i in b}
else:
return np.array([ab_dict[a][b_i] for b_i in b])
elif a_dim == 1 and line_dim == 0:
if rtype == dict:
return {a_i: ab_dict[a_i][b] for a_i in a}
else:
return np.array([ab_dict[a_i][b] for a_i in a])
elif a_dim == 1 and line_dim == 1:
if rtype == dict:
return {a_i: {b_i: ab_dict[a_i][b_i] for b_i in b} for a_i in a}
else:
return np.array([[ab_dict[a_i][b_i] for b_i in b] for a_i in a])
else:
raise ValueError('\'a\' and \'b\' must be Non-Iterables or an Iterable of the corresponding types.')
def _wrap_func_pars(func, p0, p0_fixed):
def func_wrapped(x, *params):
_params = p0
_params[~p0_fixed] = np.asarray(params)
return func(x, *_params)
return func_wrapped
def _wrap_func_sigma(sigma, p0, p0_fixed):
def transform(xdata, ydata, yfunc, *params):
_params = p0
_params[~p0_fixed] = np.asarray(params)
return 1 / sigma(xdata, ydata, yfunc, *_params)
return transform
def _wrap_func(func, xdata, ydata, transform):
# Copied from scipy, extended for callables.
if transform is None:
def func_wrapped(params):
return func(xdata, *params) - ydata
elif callable(transform):
def func_wrapped(params):
yfunc = func(xdata, *params)
return transform(xdata, ydata, yfunc, *params) * (yfunc - ydata)
elif transform.ndim == 1:
def func_wrapped(params):
return transform * (func(xdata, *params) - ydata)
else:
def func_wrapped(params):
return solve_triangular(transform, func(xdata, *params) - ydata, lower=True)
return func_wrapped
[docs]def straight(x: array_like, a: array_like, b: array_like) -> ndarray:
"""
:param x: The x values.
:param a: The y-intercept.
:param b: The slope.
:returns: The y values resulting from the 'x' values via the given linear relation.
"""
x, a, b = np.asarray(x), np.asarray(a), np.asarray(b)
return a + b * x
[docs]def straight_slope(p_0: array_iter, p_1: array_iter, axis=-1) -> ndarray:
"""
:param p_0: The first point(s).
:param p_1: The second point(s).
:param axis: The axis along which the vector components are aligned.
:returns: The parametrization of a straight line in n dimensions.
"""
p_0, p_1 = np.asarray(p_0), np.asarray(p_1)
dr = p_1 - p_0
return dr / np.expand_dims(tools.absolute(dr, axis=axis), axis=axis)
[docs]def straight_std(x: array_like, sigma_a: array_like, sigma_b: array_like, corr_ab: array_like) -> ndarray:
"""
:param x: The x values.
:param sigma_a: The standard deviation of the y-intercept.
:param sigma_b: The standard deviation of the slope.
:param corr_ab: The correlation coefficient between the slope and y-intercept.
:returns: The standard deviation of a straight line where the x values do not have uncertainties.
"""
x = np.asarray(x)
sigma_a, sigma_b, corr_ab = np.asarray(sigma_a), np.asarray(sigma_b), np.asarray(corr_ab)
return np.sqrt(sigma_a ** 2 + (x * sigma_b) ** 2 + 2 * x * sigma_a * sigma_b * corr_ab)
[docs]def straight_x_std(x: array_like, b: array_like, sigma_x: array_like,
sigma_a: array_like, sigma_b: array_like, corr_ab: array_like) -> ndarray:
"""
:param x: The x values.
:param b: The slope.
:param sigma_x: The standard deviation of the x values.
:param sigma_a: The standard deviation of the y-intercept.
:param sigma_b: The standard deviation of the slope.
:param corr_ab: The correlation coefficient between the slope and y-intercept.
:returns: The standard deviation of a straight line where all input values have uncertainties.
"""
x, sigma_x = np.asarray(x), np.asarray(sigma_x)
b, sigma_a, sigma_b, corr_ab = np.asarray(b), np.asarray(sigma_a), np.asarray(sigma_b), np.asarray(corr_ab)
return np.sqrt(sigma_a ** 2 + (x * sigma_b) ** 2 + 2 * x * sigma_a * sigma_b * corr_ab
+ (b * sigma_x) ** 2 + (sigma_b * sigma_x) ** 2)
[docs]def draw_straight_unc_area(x: array_like, y: array_like, sigma_a: array_like, sigma_b: array_like,
corr_ab: array_like, **kwargs):
"""
:param x: The x values.
:param y: The y values.
:param sigma_a: The standard deviation of the y-intercept.
:param sigma_b: The standard deviation of the slope.
:param corr_ab: The correlation coefficient between the slope and y-intercept.
:param kwargs: The keyword arguments for the fill_between function.
:returns: The standard deviation of a straight line where the x values do not have uncertainties.
"""
unc = straight_std(x, sigma_a, sigma_b, corr_ab)
plt.fill_between(x, y - unc, y + unc, **kwargs)
[docs]def ellipse2d(x: array_like, y: array_like, scale_x: array_like, scale_y: array_like,
phi: array_like, corr: array_like) -> (ndarray, ndarray):
"""
:param x: The x-component of the position of the ellipse.
:param y: The y-component of the position of the ellipse.
:param scale_x: The amplitude of the x-component.
:param scale_y: The amplitude of the y-component.
:param phi: The angle between the vector to the point on the ellipse and the x-axis.
:param corr: The correlation coefficient between the x and y data.
:returns: A point on an ellipse in 2d-space with amplitudes 'x', 'y'
and correlation 'corr' between x- and y-component.
"""
x, y, scale_x, scale_y, corr = np.asarray(x), np.asarray(y), \
np.asarray(scale_x), np.asarray(scale_y), np.asarray(corr)
return x + scale_x * np.cos(phi), y + scale_y * (corr * np.cos(phi) + np.sqrt(1 - corr ** 2) * np.sin(phi))
[docs]def draw_sigma2d(x: array_iter, y: array_iter, sigma_x: array_iter, sigma_y: array_iter,
corr: array_iter, n: int = 1, **kwargs):
"""
:param x: The x data.
:param y: The y data.
:param sigma_x: The 1-sigma uncertainties of the x data.
:param sigma_y: The 1-sigma uncertainties of the y data.
:param corr: The correlation coefficients between the x and y data.
:param n: The maximum sigma region to draw
:param kwargs: Additional keyword arguments are passed to plt.plot().
Use key 'fmt' to specify the third argument of plt.plot().
:returns: None. Draws the sigma-bounds of the given data points (x, y) until the n-sigma region.
"""
fmt = '-k'
if 'fmt' in list(kwargs.keys()):
fmt = kwargs.pop('fmt')
phi = np.arange(0., 2 * np.pi, 0.001)
for x_i, y_i, s_x, s_y, r in zip(x, y, sigma_x, sigma_y, corr):
for i in range(1, n + 1, 1):
plt.plot(*ellipse2d(x_i, y_i, i * s_x, i * s_y, phi, r), fmt, **kwargs)
[docs]def weight(sigma):
"""
:param sigma: The 1-sigma uncertainty.
:returns: The weight corresponding to the 1-sigma uncertainty 'sigma'
"""
return 1. / sigma ** 2
[docs]def york(x: array_iter, y: array_iter, sigma_x: array_iter = None, sigma_y: array_iter = None,
corr: array_iter = None, iter_max: int = 200, report: bool = False, show: bool = False):
"""
A linear regression algorithm to find the best straight line, given normally distributed errors for x and y
and correlation coefficients between errors in x and y. The algorithm is described in
['Unified equations for the slope, intercept, and standard errors of the best straight line',
York et al., American Journal of Physics 72, 367 (2004)]. See the comments to compare the individual steps.
:param x: The x data.
:param y: The y data.
:param sigma_x: The 1-sigma uncertainties of the x data.
:param sigma_y: The 1-sigma uncertainties of the y data.
:param corr: The correlation coefficients between the x and y data.
:param iter_max: The maximum number of iterations to find the best slope.
:param report: Whether to print the result of the fit.
:param show: Whether to plot the fit result.
:returns: a, b, sigma_a, sigma_b, corr_ab. The best y-intercept and slope,
their respective 1-sigma uncertainties and their correlation coefficient.
"""
x, y = np.asarray(x), np.asarray(y)
if sigma_x is None:
sigma_x = np.full_like(x, 1.)
if report:
print('\nNo uncertainties for \'x\' were given. Assuming \'sigma_x\'=1.')
if sigma_y is None:
sigma_y = np.full_like(y, 1.)
if report:
print('\nNo uncertainties for \'y\' were given. Assuming \'sigma_y\'=1.')
sigma_2d = True
if corr is None:
sigma_2d = False
corr = 0.
sigma_x, sigma_y, corr = np.asarray(sigma_x), np.asarray(sigma_y), np.asarray(corr)
p_opt, _ = curve_fit(straight, x, y, p0=[0., 1.], sigma=sigma_y) # (1)
b_init = p_opt[1]
b = b_init
w_x, w_y = weight(sigma_x), weight(sigma_y) # w(X_i), w(Y_i), (2)
alpha = np.sqrt(w_x * w_y)
n = 0
r_tol = 1e-15
tol = 1.
mod_w = None
x_bar, y_bar = None, None
beta = None
while tol > r_tol and n < iter_max: # (6)
b_init = b
mod_w = w_x * w_y / (w_x + b_init ** 2 * w_y - 2. * b_init * corr * alpha) # W_i, (3)
x_bar, y_bar = np.average(x, weights=mod_w), np.average(y, weights=mod_w) # (4)
u = x - x_bar
v = y - y_bar
beta = mod_w * (u / w_y + v * b_init / w_x - straight(u, v, b_init) * corr / alpha)
b = np.sum(mod_w * beta * v) / np.sum(mod_w * beta * u) # (5)
tol = abs((b - b_init) / b)
n += 1 # (6)
a = y_bar - b * x_bar # (7)
x_i = x_bar + beta # (8)
x_bar = np.average(x_i, weights=mod_w) # (9)
u = x_i - x_bar
sigma_b = np.sqrt(1. / np.sum(mod_w * u ** 2)) # (10)
sigma_a = np.sqrt(1. / np.sum(mod_w) + x_bar ** 2 * sigma_b ** 2)
corr_ab = -x_bar * sigma_b / sigma_a
chi2 = np.sum(mod_w * (y - b * x - a) ** 2) / (x.size - 2)
if report:
if n == iter_max:
print('\nMaximum number of iterations ({}) was reached.'.format(iter_max))
tools.printh('York-fit result:')
print('a: {} +/- {}\nb: {} +/- {}\ncorr_ab: {}\nchi2: {}'
.format(a, sigma_a, b, sigma_b, corr_ab, chi2))
if show:
x_cont = np.linspace(np.min(x) - 0.1 * (np.max(x) - np.min(x)), np.max(x) + 0.1 * (np.max(x) - np.min(x)), 1001)
plt.xlabel('x')
plt.ylabel('y')
if sigma_2d:
plt.plot(x, y, 'k.', label='Data')
draw_sigma2d(x, y, sigma_x, sigma_y, corr, n=1)
else:
plt.errorbar(x, y, xerr=sigma_x, yerr=sigma_y, fmt='k.', label='Data')
plt.plot(x_cont, straight(x_cont, a, b), 'b-', label='Fit')
y_min = straight(x_cont, a, b) - straight_std(x_cont, sigma_a, sigma_b, corr_ab)
y_max = straight(x_cont, a, b) + straight_std(x_cont, sigma_a, sigma_b, corr_ab)
plt.fill_between(x_cont, y_min, y_max, color='b', alpha=0.3, antialiased=True)
plt.legend(loc='best', numpoints=1)
plt.show()
return a, b, sigma_a, sigma_b, corr_ab
[docs]def linear_monte_carlo(x: array_iter, y: array_iter, sigma_x: array_iter = None, sigma_y: array_iter = None,
corr: array_iter = None, n: int = 1000000, report: bool = True, show: bool = False):
"""
:param x: The x data.
:param y: The y data.
:param sigma_x: The 1-sigma uncertainties of the x data.
:param sigma_y: The 1-sigma uncertainties of the y data.
:param corr: The correlation coefficients between the x and y data.
:param n: The number of samples generated for each data point.
:param report: Whether to print the result of the fit.
:param show: Whether to plot the fit result.
:returns: a, b, sigma_a, sigma_b, corr_ab. The best y-intercept and slope,
their respective 1-sigma uncertainties and their correlation coefficient.
"""
if sigma_x is None:
sigma_x = [1., ] * len(x)
if report:
print('\nNo uncertainties for \'x\' were given. Assuming \'sigma_x\'=1.')
if sigma_y is None:
sigma_y = [1., ] * len(y)
if report:
print('\nNo uncertainties for \'y\' were given. Assuming \'sigma_y\'=1.')
x, y, sigma_x, sigma_y, corr = np.asarray(x), np.asarray(y),\
np.asarray(sigma_x), np.asarray(sigma_y), np.asarray(corr)
tools.check_shape_like(x, y, sigma_x, sigma_y, corr, allow_scalar=False)
tools.check_dimension(x.size, 0, x, y, sigma_x, sigma_y, corr)
mean = np.concatenate((np.expand_dims(x, axis=1), np.expand_dims(y, axis=1)), axis=1)
cov = np.array([[[s_x ** 2, r * s_x * s_y],
[r * s_y * s_x, s_y ** 2]] for s_x, s_y, r in zip(sigma_x, sigma_y, corr)])
a, b, sigma_a, sigma_b, corr_ab = linear_monte_carlo_nd(mean, cov=cov, axis=0, n=n, report=report, show=show)
return a[0], b[0], sigma_a[0], sigma_b[0], corr_ab[0, 1]
def _test_order_linear_monte_carlo_nd(x: ndarray, cov: ndarray, n: int):
"""
:param x: The x data. 'x' is a numpy array of vectors with arbitrary but fixed dimension.
:param cov: The covariance matrices of the x data vectors. Must have shape (x.shape[0], x.shape[1], x.shape[1]).
:param n: The number of samples generated for each data point.
:returns: The order of axis 0 of 'x' which yields the most accepted samples.
"""
indices = [(i, j) for i in range(x.shape[0]) for j in range(x.shape[0]) if j > i]
best_n = 0
best_order = np.arange(x.shape[0], dtype=int)
for (i, j) in indices:
order = np.array([i, ] + [k for k in range(x.shape[0]) if k != i and k != j] + [j, ])
n_accepted = np.sum(_generate_collinear_points(x[order], cov[order], n)[2])
# print(n_accepted)
if n_accepted > best_n:
best_n = n_accepted
best_order = order
# print(best_n, best_order)
return best_order
def _generate_collinear_points(x: ndarray, cov: ndarray, n: int):
"""
:param x: The x data. 'x' is a numpy array of vectors with arbitrary but fixed dimension.
:param cov: The covariance matrices of the x data vectors. Must have shape (x.shape[0], x.shape[1], x.shape[1]).
:param n: The number of samples generated for each data point.
:returns: The randomly generated data vectors p [shape=(n, )+ x.shape] aligned along a straight line,
the unit vectors dr [shape=(n, x.shape[1])] in the directions of the straight lines
and a mask of the accepted samples. The accepted samples are p[accepted] and dr[accepted].
"""
p_0 = multivariate_normal.rvs(mean=x[0], cov=cov[0], size=n)
p_1 = multivariate_normal.rvs(mean=x[-1], cov=cov[-1], size=n)
dr = straight_slope(p_0, p_1, axis=-1)
t_0 = [np.einsum('ij,ij->i', np.expand_dims(x_i, axis=0) - p_0, dr) for x_i in x[1:-1]]
t = [norm.rvs(loc=t_0_i, scale=np.sqrt(np.min(np.diag(cov_i)))) for t_0_i, cov_i in zip(t_0, cov[1:-1])]
p_new = [p_0 + np.expand_dims(t_i, axis=-1) * dr for t_i in t]
f = np.prod([multivariate_normal.pdf(p_i, mean=x_i, cov=cov_i)
for p_i, x_i, cov_i in zip(p_new, x[1:-1], cov[1:-1])], axis=0)
g = np.prod([norm.pdf(t_i, loc=t_0_i, scale=np.sqrt(np.max(np.diag(cov_i))))
for t_i, t_0_i, cov_i in zip(t, t_0, cov[1:-1])], axis=0)
if p_new:
u = f / g
u /= np.max(u)
accepted = np.random.random(size=n) < u
p = np.concatenate((np.expand_dims(p_0, 0), p_new, np.expand_dims(p_1, 0)), axis=0)
else:
accepted = np.full(n, True, dtype=bool)
p = np.concatenate((np.expand_dims(p_0, 0), np.expand_dims(p_1, 0)), axis=0)
p = np.transpose(p, axes=[1, 0, 2])
return p, dr, accepted
[docs]def linear_monte_carlo_nd(x: array_iter, cov: array_iter = None, axis: int = None,
n: int = 1000000, full_output: bool = False, report: bool = True, show: bool = False):
"""
A Monte-Carlo fitter that finds a straight line in n-dimensional space.
:param x: The x data. 'x' is an iterable of vectors with arbitrary but fixed dimension.
:param cov: The covariance matrices of the x data vectors. Must have shape (x.shape[0], x.shape[1], x.shape[1]).
:param axis: The axis to use for the parametrization. If None, the directional vector is normalized.
:param n: The number of samples generated for each data point.
:param full_output: Whether to also return the generated points 'p' and a mask for the 'accepted' points.
The accepted points are p[accepted]. p has shape (n, x.shape).
:param report: Whether to print the result of the fit.
:param show: Whether to plot the fit result.
:returns: A list of results as defined by the routine 'linear_monte_carlo_nd':
A tuple (a, b, sigma_a, sigma_b, corr_ab) of arrays. The best y-intercepts and slopes,
their respective 1-sigma uncertainties and their correlation matrix.
The y-intercepts and slopes are defined through 'axis' by a(x[axis] == 0)
and b(dx[i != axis] / dx[axis]), respectively.
Additionally, a second tuple with the accepted points, offsets, slopes and a mask
for the accepted points is returned if full_output == True.
"""
x = np.asarray(x)
if len(x.shape) != 2:
raise ValueError('x must be an iterable of vectors with arbitrary but fixed dimension.'
' Hence, len(x.shape) =: 2, but is {}.'.format(len(x.shape)))
if cov is None:
cov = [np.identity(x.shape[1]), ] * x.shape[0]
cov = np.asarray(cov)
tools.check_shape((x.shape[0], x.shape[1], x.shape[1]), cov, allow_scalar=False)
order = _test_order_linear_monte_carlo_nd(x, cov, 100000)
inverse_order = np.array([int(np.where(order == i)[0])
for i in range(order.size)]) # Invert the order for later.
p, dr, accepted = _generate_collinear_points(x[order], cov[order], n)
p = p[:, inverse_order, :]
n_accepted = np.sum(accepted)
p_0 = p[accepted, 0, :]
dr = dr[accepted]
if show:
for sample in range(0, int(n_accepted), int(n_accepted / 10)):
x_plot = p[accepted][sample, :, 0]
y_plot = p[accepted][sample, :, 1]
plot_order = np.argsort(x_plot)
x_plot = x_plot[plot_order]
y_plot = y_plot[plot_order]
plt.plot([x_plot[0], x_plot[-1]], [y_plot[0], y_plot[-1]], 'C0-')
plt.plot(x_plot, y_plot, 'C1.')
# plt.errorbar(x[:, 0], x[:, 1], xerr=np.sqrt(cov[:, 0, 0]), yerr=np.sqrt(cov[:, 1, 1]), fmt='k.')
draw_sigma2d(x[:, 0], x[:, 1], np.sqrt(cov[:, 0, 0]), np.sqrt(cov[:, 1, 1]),
cov[:, 0, 1] / (cov[:, 0, 0] * cov[:, 1, 1]), n=2)
# plt.gca().set_aspect('equal')
plt.show()
mask = [True, ] * x.shape[1]
if axis is not None:
if axis == -1:
axis = x.shape[1] - 1
mask = np.arange(x.shape[1]) != axis
t = -p_0[:, axis] / dr[:, axis]
p_0 = p_0 + np.expand_dims(t, axis=-1) * dr
dr = dr / dr[:, [axis]]
a = np.mean(p_0, axis=0)[mask]
b = np.mean(dr, axis=0)[mask]
sigma_a = np.std(p_0, axis=0, ddof=1)[mask]
sigma_b = np.std(dr, axis=0, ddof=1)[mask]
corr_ab = np.corrcoef(p_0.T[mask], dr.T[mask])
# print(dr.shape)
# plt.plot(p[accepted, :, 0], dr[:, 1], '.')
# plt.show()
if report:
tools.printh('Linear Monte-Carlo-fit result:')
print('Accepted samples: {} / {} ({}%)\na: {} +- {}\nb: {} +- {}\ncorr_ab:'
.format(n_accepted, n, np.around(n_accepted / n * 100, decimals=2),
a, sigma_a, b, sigma_b))
tools.print_cov(corr_ab, normalize=True)
if full_output:
return (a, b, sigma_a, sigma_b, corr_ab), (p[accepted], p_0, dr, accepted)
return a, b, sigma_a, sigma_b, corr_ab
[docs]def linear_alpha(x: array_iter, y: array_iter, sigma_x: array_like = None, sigma_y: array_like = None,
corr: array_iter = None, func: Callable = york, alpha: scalar = 0, find_alpha: bool = True,
report: bool = True, show: bool = False, **kwargs):
"""
:param x: The x data.
:param y: The y data.
:param sigma_x: The 1-sigma uncertainty of the x data.
:param sigma_y: The 1-sigma uncertainty of the y data.
:param corr: The correlation coefficients between the x and y data.
:param func: The fitting routine.
:param alpha: An x-axis offset to reduce the correlation coefficient between the y-intercept and the slope.
:param find_alpha: Whether to search for the best 'alpha'. Uses the given 'alpha' as a starting point.
May not give the desired result if 'alpha' was initialized to far from its optimal value.
:param report: Whether to print the result of the fit.
:param show: Whether to plot the fit result.
:param kwargs: Additional keyword arguments are passed to the fitting routine.
:returns: a, b, sigma_a, sigma_b, corr_ab, alpha. The best y-intercept and slope,
their respective 1-sigma uncertainties, their correlation coefficient and the used alpha.
"""
n = tools.floor_log10(alpha)
def cost(x0):
_, _, _, _, c = func(x - x0[0], y, sigma_x=sigma_x, sigma_y=sigma_y, corr=corr,
report=False, show=False, **kwargs)
print('alpha: {}(1.), \tcost: {}({})'.format(x0[0], c ** 2 * 10. ** (n + 1), 10. ** (n - 3)))
return c ** 2 * 10. ** (n + 1)
if find_alpha:
alpha = minimize(cost, np.array([alpha]), method='Nelder-Mead',
options={'xatol': 1., 'fatol': 0.01 ** 2 * 10. ** (n + 1)}).x[0]
a, b, sigma_a, sigma_b, corr_ab =\
func(x - alpha, y, sigma_x=sigma_x, sigma_y=sigma_y, corr=corr,
report=report, show=show, **kwargs)
if report:
print('alpha: {}'.format(alpha))
return a, b, sigma_a, sigma_b, corr_ab, alpha
[docs]def odr_fit(f: Callable, x: array_iter, y: array_iter, sigma_x: array_iter = None, sigma_y: array_iter = None,
p0: array_iter = None, p0_d: array_iter = None, p0_fixed: array_iter = None,
report: bool = False, **kwargs) -> (ndarray, ndarray):
"""
:param f: The model function to fit to the data.
:param x: The x data.
:param y: The y data.
:param sigma_x: The 1-sigma uncertainty of the x data.
:param sigma_y: The 1-sigma uncertainty of the y data.
:param p0: A numpy array or an Iterable of the initial guesses for the parameters.
Must have at least the same length as the minimum number of parameters required by the function 'f'.
If 'p0' is None, 1 is taken as an initial guess for all non-keyword parameters.
:param p0_d: A numpy array or an Iterable of the uncertainties of the initial guesses for the parameters.
Must have the same length as p0.
:param p0_fixed: A numpy array or an Iterable of bool values specifying, whether to fix a parameter.
Must have the same length as p0.
:param report: Whether to print the result of the fit.
:param kwargs: Keyword arguments passed to odr.ODR.
:returns: popt, pcov. The optimal parameters and their covariance matrix.
"""
x, y = np.asarray(x), np.asarray(y)
sx, sy, covx, covy = None, None, None, None
if sigma_x is not None:
sx = np.asarray(sigma_x)
# if sigma_x.shape == (x.size, x.size):
# covx = sigma_x
# sy = None
if sigma_y is not None:
sy = np.asarray(sigma_y)
# if sigma_y.shape == (y.size, y.size):
# covy = sigma_y
# sy = None
if p0 is None:
arg_spec = inspect.getfullargspec(f)
n = len(arg_spec.args) - 1
if (n < 1 and arg_spec.defaults is None) or (arg_spec.defaults is not None and n - len(arg_spec.defaults) < 1):
raise ValueError('\'f\' must have at least one parameter that can be optimized.')
p0 = [1., ] * n
if arg_spec.defaults is not None:
p0 = [1., ] * (n - len(arg_spec.defaults))
p0 = np.asarray(p0)
ifixb = None
if p0_fixed is not None:
ifixb = (~np.asarray(p0_fixed)).astype(int)
if ifixb.shape != p0.shape:
raise ValueError('\'p0_fixed\' must have the same shape as \'p0\'.')
data = odr.RealData(x, y, sx=sx, sy=sy, covx=covx, covy=covy)
# noinspection PyTypeChecker
model = odr.Model(lambda beta, x_i: f(x_i, *beta))
odr_i = odr.ODR(data, model, beta0=p0, delta0=p0_d, ifixb=ifixb, **kwargs)
out = odr_i.run()
if report:
tools.printh('odr_fit result:')
for i, (val, err) in enumerate(zip(out.beta, np.sqrt(np.diag(out.cov_beta)))):
print('{}: {} +/- {}'.format(i, val, err))
return out.beta, out.cov_beta
[docs]def curve_fit(f: Callable, x: Union[array_like, object], y: array_like, p0: array_iter = None,
p0_fixed: array_iter = None, sigma: Union[array_iter, Callable] = None, absolute_sigma: bool = False,
check_finite: bool = True, bounds: (ndarray, ndarray) = (-np.inf, np.inf), method: str = None,
jac: Union[Callable, str] = None, full_output: bool = False, report: bool = False, **kwargs):
"""
:param f: The model function to fit to the data.
:param x: The x data.
:param y: The y data.
:param p0: A numpy array or an Iterable of the initial guesses for the parameters.
Must have at least the same length as the minimum number of parameters required by the function 'f'.
If 'p0' is None, 1 is taken as an initial guess for all non-keyword parameters.
:param p0_fixed: A numpy array or an Iterable of bool values specifying, whether to fix a parameter.
Must have the same length as p0.
:param sigma: The 1-sigma uncertainty of the y data.
This can also be a function g such that 'g(x, y, f(x, \*params), \*params) -> sigma'.
:param absolute_sigma: See scipy.optimize.curve_fit.
:param check_finite: See scipy.optimize.curve_fit.
:param bounds: See scipy.optimize.curve_fit.
:param method: See scipy.optimize.curve_fit.
:param jac: See scipy.optimize.curve_fit. Must not be callable if 'sigma' is callable.
:param full_output: See scipy.optimize.curve_fit.
:param report: Whether to print the result of the fit.
:param kwargs: See scipy.optimize.curve_fit.
:returns: popt, pcov. The optimal parameters and their covariance matrix. Additional output if full_output is True.
See scipy.optimize.curve_fit.
"""
if p0 is None:
# determine number of parameters by inspecting the function
sig = _getfullargspec(f)
args = sig.args
if len(args) < 2:
raise ValueError("Unable to determine number of fit parameters.")
n = len(args) - 1
else:
p0 = np.atleast_1d(p0)
n = p0.size
if p0_fixed is None:
p0_fixed = np.zeros(n, dtype=bool)
else:
p0_fixed = np.atleast_1d(p0_fixed)
if isinstance(bounds, Bounds):
lb, ub = bounds.lb, bounds.ub
else:
lb, ub = prepare_bounds(bounds, n)
if p0 is None:
p0 = _initialize_feasible(lb, ub)
# Truncate p0, bounds and func to use free parameters only.
p0_free = p0[~p0_fixed]
lb, ub = lb[~p0_fixed], ub[~p0_fixed]
func = _wrap_func_pars(f, p0, p0_fixed)
bounds = (lb, ub)
bounded_problem = np.any((lb > -np.inf) | (ub < np.inf))
if method is None:
if bounded_problem:
method = 'trf'
else:
method = 'lm'
if method == 'lm' and bounded_problem:
raise ValueError("Method 'lm' only works for unconstrained problems. "
"Use 'trf' or 'dogbox' instead.")
# optimization may produce garbage for float32 inputs, cast them to float64
# NaNs cannot be handled
if check_finite:
y = np.asarray_chkfinite(y, float)
else:
y = np.asarray(y, float)
if isinstance(x, (list, tuple, np.ndarray)):
# 'x' is passed straight to the user-defined 'f', so allow
# non-array_like 'x'.
if check_finite:
x = np.asarray_chkfinite(x, float)
else:
x = np.asarray(x, float)
if y.size == 0:
raise ValueError("'y' must not be empty!")
# Determine type of sigma
if sigma is not None:
if callable(sigma):
transform = _wrap_func_sigma(sigma, p0, p0_fixed)
else:
sigma = np.asarray(sigma)
# if 1-D, sigma are errors, define transform = 1/sigma
if sigma.shape == (y.size,):
transform = 1.0 / sigma
# if 2-D, sigma is the covariance matrix,
# define transform = L such that L L^T = C
elif sigma.shape == (y.size, y.size):
try:
# scipy.linalg.cholesky requires lower=True to return L L^T = A
transform = cholesky(sigma, lower=True)
except LinAlgError as e:
raise ValueError("'sigma' must be positive definite.") from e
else:
raise ValueError("'sigma' has incorrect shape.")
else:
transform = None
func = _wrap_func(func, x, y, transform)
if callable(jac):
if callable(sigma):
raise ValueError("'jac' must not be callable if 'sigma' is callable.")
jac = _wrap_jac(jac, x, transform)
elif jac is None and method != 'lm':
jac = '2-point'
if 'args' in kwargs:
# The specification for the model function 'f' does not support
# additional arguments. Refer to the 'curve_fit' docstring for
# acceptable call signatures of 'f'.
raise ValueError("'args' is not a supported keyword argument.")
if method == 'lm':
# if y.size == 1, this might be used for broadcast.
if y.size != 1 and n > y.size:
raise TypeError(f"The number of func parameters={n} must not exceed the number of data points={y.size}")
res = leastsq(func, p0_free, Dfun=jac, full_output=True, **kwargs)
popt, pcov, infodict, errmsg, ier = res
ysize = len(infodict['fvec'])
cost = np.sum(infodict['fvec'] ** 2)
if ier not in [1, 2, 3, 4]:
raise RuntimeError("Optimal parameters not found: " + errmsg)
else:
# Rename maxfev (leastsq) to max_nfev (least_squares), if specified.
if 'max_nfev' not in kwargs:
kwargs['max_nfev'] = kwargs.pop('maxfev', None)
res = least_squares(func, p0_free, jac=jac, bounds=bounds, method=method, **kwargs)
if not res.success:
raise RuntimeError("Optimal parameters not found: " + res.message)
infodict = dict(nfev=res.nfev, fvec=res.fun)
ier = res.status
errmsg = res.message
ysize = len(res.fun)
cost = 2 * res.cost # res.cost is half sum of squares!
popt = res.x
# Do Moore-Penrose inverse discarding zero singular values.
# noinspection PyTupleAssignmentBalance
_, s, vt = svd(res.jac, full_matrices=False)
threshold = np.finfo(float).eps * max(res.jac.shape) * s[0]
s = s[s > threshold]
vt = vt[:s.size]
pcov = np.dot(vt.T / s ** 2, vt)
warn_cov = False
if pcov is None:
# indeterminate covariance
pcov = np.zeros((len(popt), len(popt)), dtype=float)
pcov.fill(np.inf)
warn_cov = True
elif not absolute_sigma:
if ysize > p0.size:
s_sq = cost / (ysize - p0.size)
pcov = pcov * s_sq
else:
pcov.fill(np.inf)
warn_cov = True
popt_ret, pcov_ret = p0.copy(), np.zeros((p0.size, p0.size), dtype=float)
popt_ret[~p0_fixed] = popt
pcov_mask = ~(p0_fixed[np.newaxis, :] + p0_fixed[:, np.newaxis])
pcov_ret[pcov_mask] = pcov.flatten()
if report:
digits = int(np.floor(np.log10(np.abs(p0.size)))) + 1
tools.printh("curve_fit result:")
for i, (val, err, fix) in enumerate(zip(popt_ret, np.sqrt(np.diag(pcov_ret)), p0_fixed)):
print(f"{str(i).zfill(digits)}: {val} +/- {err} ({'fixed' if fix else 'free'})")
if warn_cov:
warnings.warn("Covariance of the parameters could not be estimated", category=OptimizeWarning)
if full_output:
return popt_ret, pcov_ret, infodict, errmsg, ier
else:
return popt_ret, pcov_ret
[docs]class King:
def __init__(self, a: array_iter, m: array_iter, x_abs: array_iter = None,
subtract_electrons: scalar = 0., element_label: str = None):
"""
:param a: An Iterable of the mass numbers of the used isotopes.
:param m: The masses and their uncertainties of the isotopes 'a' (amu). 'm' must have shape (len(a), 2).
:param x_abs: Absolute values of the x-axis corresponding to the mass numbers a.
If given, the 'x' parameter can be omitted when fitting (in 'fit' and 'fit_nd')
and is determined automatically as the difference between the mass numbers 'a' and 'a_ref'.
Must have shape (len(a), 2) or (len(a), n, 2) where n is the number of dimensions of the king plot without 'y'.
:param subtract_electrons: The number of electron masses
that should be subtracted from the specified isotope masses. If the ionization energy must be considered,
less electrons can be subtracted. 'subtract_electrons' does not have to be an integer.
:param element_label: The label of the element to enhance the printed and plotted information.
"""
self.fontsize = 12.
self.scale_y = 1e-3
self.n = 1000000
self.subtract_electrons = subtract_electrons
self.element_label = element_label
if self.element_label is None:
self.element_label = ''
self.a: Union[object, list] = np.asarray(a, dtype=int).tolist()
self.x_abs = x_abs
if self.x_abs is not None:
self.x_abs = np.asarray(x_abs)
if len(self.x_abs.shape) < 3:
self.x_abs = np.expand_dims(self.x_abs, axis=1)
self.m: Union[object, list, ndarray] = np.asarray(m, dtype=float).tolist()
self.m = np.array([[m_i[0] - self.subtract_electrons * me_u,
np.sqrt(m_i[1] ** 2 + (self.subtract_electrons * me_u_d) ** 2)] for m_i in self.m])
tools.check_shape((len(a), 2), self.m, allow_scalar=False)
self.m_mod = np.array([[list(mass_factor(_m[0], _m_ref[0], m_d=_m[1], m_ref_d=_m_ref[1], k_inf=True))
if _i != _j else [0., 0.]
for _j, _m in enumerate(self.m)] for _i, _m_ref in enumerate(self.m)])
self.m_sample = None
self.a_fit, self.a_ref = None, None
self.x, self.y = None, None
self.x_mod, self.y_mod = None, None
self.corr = None
self.results: Union[None, tuple] = None
self.x_nd = None
self.x_nd_mod = None
self.cov = None
self.results_nd: Union[None, tuple] = None
self.x_results: Union[None, tuple] = None
self.axis = 0
def _correlation(self):
"""
:returns: None. Sets the correlation coefficients of the modified x- and y-data points by sampling the data.
"""
i = np.array([self.a.index(a_i) for a_i in self.a_fit])
i_ref = np.array([self.a.index(a_i) for a_i in self.a_ref])
self.m_sample = norm.rvs(loc=self.m[:, 0], scale=self.m[:, 1], size=(self.n, self.m.shape[0])).T
x = norm.rvs(loc=self.x[:, 0], scale=self.x[:, 1], size=(self.n, self.x.shape[0])).T
y = norm.rvs(loc=self.y[:, 0], scale=self.y[:, 1], size=(self.n, self.y.shape[0])).T
m_mod_sample = mass_factor(self.m_sample[i], self.m_sample[i_ref], k_inf=True)[0]
x_mod = m_mod_sample * x
y_mod = m_mod_sample * y
self.corr = np.diag(np.corrcoef(x_mod, y_mod)[:x_mod.shape[0], x_mod.shape[0]:])
def _correlation_nd(self):
"""
:returns: None. Sets the correlation coefficients of the modified data by sampling it.
"""
i = np.array([self.a.index(a_i) for a_i in self.a_fit])
i_ref = np.array([self.a.index(a_i) for a_i in self.a_ref])
self.m_sample = norm.rvs(loc=self.m[:, 0], scale=self.m[:, 1], size=(self.n, self.m.shape[0])).T
x_nd = norm.rvs(loc=self.x_nd[:, :, 0], scale=self.x_nd[:, :, 1],
size=(self.n, self.x_nd.shape[0], self.x_nd.shape[1]))
x_nd = np.transpose(x_nd, axes=[1, 2, 0])
m_mod_sample = mass_factor(self.m_sample[i], self.m_sample[i_ref], k_inf=True)[0]
x_mod = np.expand_dims(m_mod_sample, axis=1) * x_nd
self.cov = np.array([np.cov(x_mod_i) for x_mod_i in x_mod])
[docs] def fit(self, a: array_iter, a_ref: array_iter, x: array_iter = None, y: array_iter = None,
xy: Iterable[int] = None, func: Callable = york, alpha: scalar = 0, find_alpha: bool = False,
show: bool = True, **kwargs):
"""
:param a: An Iterable of the mass numbers of the used isotopes.
:param a_ref: An Iterable of the mass numbers of the used reference isotopes.
:param x: The x-values and their uncertainties of the King-fit to be performed. If 'mode' is "radii",
this must contain the differences of mean square nuclear charge radii or the Lambda-factor.
'x' must have shape (len(a), 2). Units: (fm ** 2) or (MHz).
:param y: The isotope shifts and their uncertainties of the King-fit to be performed.
'y' must have shape (len(a), 2). If None, 'y' is tried to be inherited from 'self.f'. Units: (MHz).
:param xy: A 2-tuple of indices (ix, iy) which select the axes to use for the King-fit from 'King.x_abs'
if x or y is not specified. The default value is (0, 1), fitting the second against the first axis.
:param func: The fitting routine.
:param alpha: An x-axis offset to reduce the correlation coefficient between the y-intercept and the slope.
Unit: (u fm ** 2) or (u MHz).
:param find_alpha: Whether to search for the best 'alpha'. Uses the given 'alpha' as a starting point.
May not give the desired result if 'alpha' was initialized to far from its optimal value.
:param show: Whether to plot the fit result.
:param kwargs: Additional keyword arguments are passed to 'self.plot'.
:returns: A list of results as defined by the routine 'linear_alpha':
a, b, sigma_a, sigma_b, corr_ab, alpha. The best y-intercept and slope,
their respective 1-sigma uncertainties, their correlation coefficient and the used alpha.
"""
self.a_fit, self.a_ref = np.asarray(a), np.asarray(a_ref)
i = np.array([self.a.index(a_i) for a_i in self.a_fit])
i_ref = np.array([self.a.index(a_i) for a_i in self.a_ref])
if xy is None:
xy = (0, 1)
if x is None:
self.x = np.zeros((self.a_fit.size, 2), dtype=float)
self.x[:, 0] = self.x_abs[i, xy[0], 0] - self.x_abs[i_ref, xy[0], 0]
self.x[:, 1] = np.sqrt(self.x_abs[i, xy[0], 1] ** 2 + self.x_abs[i_ref, xy[0], 1] ** 2)
else:
self.x = np.asarray(x)
if y is None:
self.y = np.zeros((self.a_fit.size, 2), dtype=float)
self.y[:, 0] = self.x_abs[i, xy[1], 0] - self.x_abs[i_ref, xy[1], 0]
self.y[:, 1] = np.sqrt(self.x_abs[i, xy[1], 1] ** 2 + self.x_abs[i_ref, xy[1], 1] ** 2)
else:
self.y = np.asarray(y)
m_mod = self.m_mod[i_ref, i, :]
self.x_mod = np.array([self.x[:, 0] * m_mod[:, 0],
straight_x_std(self.x[:, 0], m_mod[:, 0], self.x[:, 1], 0., m_mod[:, 1], 0.)]).T
self.y_mod = np.array([self.y[:, 0] * m_mod[:, 0],
straight_x_std(self.y[:, 0], m_mod[:, 0], self.y[:, 1], 0., m_mod[:, 1], 0.)]).T
self._correlation()
self.results = linear_alpha(self.x_mod[:, 0], self.y_mod[:, 0],
sigma_x=self.x_mod[:, 1], sigma_y=self.y_mod[:, 1],
corr=self.corr, func=func, alpha=alpha, find_alpha=find_alpha, show=False)
if show:
self.plot(**kwargs)
return self.results
[docs] def fit_nd(self, a: array_iter, a_ref: array_iter, x: array_iter = None,
axis: int = 0, show: bool = True, **kwargs):
"""
:param a: An Iterable of the mass numbers of the used isotopes.
:param a_ref: An Iterable of the mass numbers of the used reference isotopes.
:param x: The x data. 'x' is an iterable of vectors with uncertainties with arbitrary but fixed dimension n_vec.
For a set of n_data data points, the shape of x must be (n_data, n_vec, 2).
:param axis: The axis to use for the parametrization. For example, a King plot with the isotope shifts
of two transitions ['D1', 'D2'], yields the slope F_D2 / F_D1 if 'axis' == 0.
:param show: Whether to plot the fit result.
:param kwargs: Additional keyword arguments are passed to 'self.plot_nd'.
:returns: A list of results as defined by the routine 'linear_monte_carlo_nd':
A tuple (a, b, sigma_a, sigma_b, corr_ab) of arrays. The best y-intercepts and slopes,
their respective 1-sigma uncertainties and their correlation matrix.
Additionally, a second tuple with the accepted points, offsets, slopes and a mask
for the accepted points is returned if full_output == True.
"""
self.axis = axis
if axis is None:
self.axis = 0
if axis == -1:
self.axis = x.shape[1] - 1
self.a_fit, self.a_ref = np.asarray(a), np.asarray(a_ref)
i = np.array([self.a.index(a_i) for a_i in self.a_fit])
i_ref = np.array([self.a.index(a_i) for a_i in self.a_ref])
if x is None:
self.x_nd = np.zeros((self.a_fit.size, self.x_abs.shape[1], 2), dtype=float)
self.x_nd[:, :, 0] = self.x_abs[i, :, 0] - self.x_abs[i_ref, :, 0]
self.x_nd[:, :, 1] = np.sqrt(self.x_abs[i, :, 1] ** 2 + self.x_abs[i_ref, :, 1] ** 2)
else:
self.x_nd = np.asarray(x)
m_mod = np.expand_dims(self.m_mod[i_ref, i, 0], axis=1)
self.x_nd_mod = m_mod * self.x_nd[:, :, 0]
self._correlation_nd()
self.results_nd, self.x_results = linear_monte_carlo_nd(self.x_nd_mod, cov=self.cov,
axis=axis, full_output=True, show=False)
# plt.plot(self.x_results[0][self.x_results[1], :, 0], self.x_results[1][self.x_results[1], :, 1], '.')
# plt.show()
if show:
self.plot_nd(axis=self.axis, **kwargs)
return self.results_nd, self.x_results
[docs] def get_popt(self) -> (ndarray, ndarray):
"""
:returns: The y-intercept and slope of the performed fit and their correlation,
if alpha was zero (Only changes the y-intercept).
This returns the mass- and field-shift factors K and F if the x-axis are radii
and the y-intercept Ky - Fy/Fx * Kx and field-shift ratio Fy/Fx if the x-axis are isotope shifts.
"""
if self.results is None:
print('There are no results yet. Please use one of the fitting options.')
return None, None
(a, b, sigma_a, sigma_b, corr, alpha) = self.results
f = [b, sigma_b]
k = [a - b * alpha, np.sqrt(sigma_a ** 2 + (alpha * sigma_b) ** 2)]
tools.printh('\npopt:')
print('f = {} +/- {}\nk = {} +/- {}\ncorr = {}'.format(*f, *k, corr))
return np.array(k), np.array(f), corr
[docs] def get_popt_nd(self) -> (ndarray, ndarray):
"""
:returns: The slopes and y-intercepts of the performed fit.
This returns the mass- and field-shift factors K and F if the x-axis are radii
and the y-intercept Ky - Fy/Fx * Kx and field-shift ratio Fy/Fx if the x-axis are isotope shifts.
The returned arrays k and f have shape (n, 2) were n is number of dimensions of the King-fit.
"""
if self.results_nd is None:
print('There are no results yet. Please use one of the fitting options.')
return None, None
(a, b, sigma_a, sigma_b, corr) = self.results_nd
f = [b, sigma_b]
k = [a, sigma_a]
r = np.diag(corr[a.size:, :a.size])
tools.printh('\npopt_nd:')
for i in range(b.size):
print('{}{}:{}'.format(tools.COLORS.OKCYAN, i, tools.COLORS.ENDC))
print('f = {}\t{}\nk = {}\t{}\ncorr = {}'.format(f[0][i], f[1][i], k[0][i], k[1][i], r[i]))
return np.array(k).T, np.array(f).T, r
[docs] def get_unmodified_x(self, a: array_iter, a_ref: array_iter, y: array_iter, results: tuple = None,
show: bool = False, **kwargs) -> Union[ndarray, None]:
"""
:param a: An Iterable of mass numbers corresponding to the isotopes of the given 'y' values.
:param a_ref: An Iterable of mass numbers corresponding to reference isotopes of the given 'y' values.
:param y: The unmodified y values of the given mass numbers. Must have shape (len(a), 2).
:param results: The results of a King plot. Must be a tuple
of the shape (a, b, sigma_a, sigma_b, corr_ab, alpha), compare the module-level method 'linear_alpha'.
If None, the results are derived from 'self.results'.
:param show: Whether to draw the King plot with the specified values.
:param kwargs: Additional keyword arguments are passed to 'self.plot'.
:returns: The unmodified x values calculated with the fit results and the given 'y' values.
"""
if results is None and self.results is None:
print('There are no results yet. Please use one of the fitting options.')
return
if results is not None:
(_a, b, sigma_a, sigma_b, corr, alpha) = results
else:
(_a, b, sigma_a, sigma_b, corr, alpha) = self.results
a, a_ref, y = np.asarray(a), np.asarray(a_ref), np.asarray(y)
i = np.array([self.a.index(a_i) for a_i in a])
i_ref = np.array([self.a.index(a_i) for a_i in a_ref])
m_mod = self.m_mod[i_ref, i, :]
y_mod = y[:, 0] * m_mod[:, 0]
y_mod_d = y[:, 1] * m_mod[:, 0]
x_mod = straight(y_mod, -_a / b, 1. / b) + alpha
x_mod_d = np.sqrt((y_mod_d / b) ** 2 + (sigma_a / b) ** 2 + (sigma_b * (x_mod - alpha) / b) ** 2
+ 2. * (x_mod - alpha) / (b ** 2) * sigma_a * sigma_b * corr)
if show:
add_xy = np.array([x_mod, x_mod_d, y_mod, y_mod_d]).T
add_a = np.array([a, a_ref]).T
self.plot(add_xy=add_xy, add_a=add_a, **kwargs)
return np.array([x_mod / m_mod[:, 0], np.abs(x_mod_d / m_mod[:, 0])]).T
[docs] def get_unmodified_y(self, a: array_iter, a_ref: array_iter, x: array_iter, results: tuple = None) \
-> Union[ndarray, None]:
"""
:param a: An Iterable of mass numbers corresponding to the isotopes of the given 'x' values.
:param a_ref: An Iterable of mass numbers corresponding to reference isotopes of the given 'x' values.
:param x: The unmodified x values of the given mass numbers. Must have shape (len(a), 2).
:param results: The results of a King plot. Must be a tuple
of the shape (a, b, sigma_a, sigma_b, corr_ab, alpha), compare the module-level method 'linear_alpha'.
If None, the results are derived from 'self.results'.
:returns: The unmodified y values calculated with the fit results and the given 'x' values.
"""
if results is None and self.results is None:
print('There are no results yet. Please use one of the fitting options.')
return
if results is not None:
(_a, b, sigma_a, sigma_b, corr, alpha) = results
else:
(_a, b, sigma_a, sigma_b, corr, alpha) = self.results
a, a_ref, x = np.asarray(a), np.asarray(a_ref), np.asarray(x)
i = np.array([self.a.index(a_i) for a_i in a])
i_ref = np.array([self.a.index(a_i) for a_i in a_ref])
m_mod = self.m_mod[i_ref, i, :]
x_mod = x[:, 0] * m_mod[:, 0] - alpha
x_mod_d = x[:, 1] * m_mod[:, 0]
y_mod = straight(x_mod, _a, b)
y_mod_d = straight_x_std(x_mod, b, x_mod_d, sigma_a, sigma_b, corr)
return np.array([y_mod / m_mod[:, 0], y_mod_d / m_mod[:, 0]]).T
[docs] def get_unmodified_x_nd(self, a: array_iter, a_ref: array_iter, y: array_iter, axis: int) -> Union[ndarray, None]:
"""
:param a: An Iterable of mass numbers corresponding to the isotopes of the given 'y' values.
:param a_ref: An Iterable of mass numbers corresponding to reference isotopes of the given 'y' values.
:param y: The unmodified y values of the given mass numbers. Must have shape (len(a), 2).
:param axis: The axis of the input values.
:returns: The unmodified values improved by 'fit_nd'.
"""
if self.x_results is None:
print('There are no results yet. Please use one of the fitting options.')
return None
_a = np.expand_dims(self.x_results[1], axis=1)
b = np.expand_dims(self.x_results[2], axis=1)
a, a_ref, y = np.asarray(a), np.asarray(a_ref), np.asarray(y)
i = np.array([self.a.index(a_i) for a_i in a])
i_ref = np.array([self.a.index(a_i) for a_i in a_ref])
m = self.m_sample[i, :]
m = m[:, self.x_results[-1]]
m_ref = self.m_sample[i_ref, :]
m_ref = m_ref[:, self.x_results[-1]]
m_mod_sample = mass_factor(m, m_ref, k_inf=True)[0].T
y_sample = norm.rvs(loc=y[:, 0], scale=y[:, 1], size=(m_mod_sample.shape[0], len(a)))
y_sample = m_mod_sample * y_sample
x_sample = straight(y_sample, -_a[:, :, axis] / b[:, :, axis], 1. / b[:, :, axis])
x_sample = straight(np.expand_dims(x_sample, axis=-1), _a, b)
x_sample /= np.expand_dims(m_mod_sample, axis=-1)
return np.array([np.mean(x_sample, axis=0), np.std(x_sample, axis=0, ddof=1)]).T
[docs] def plot(self, mode: str = '', sigma2d: bool = True, show: bool = True, add_xy: array_like = None,
add_a: array_like = None):
"""
:param mode: The mode of the King-fit. If mode='radii', the :math:'x'-axis must contain the differences of
mean square nuclear charge radii or the Lambda-factor. For every other value,
the :math:'x'-axis is assumed to be an isotope shift such that the slope corresponds to
a field-shift ratio :math:'F(y_i) / F(x)'.
:param sigma2d: Whether to draw the actual two-dimensional uncertainty bounds or the classical errorbars.
:param show: Whether to show the plot.
:param add_xy: Additional :math:'x' and :math:'y' data to plot. Must have shape (:, 4).
:param add_a: Additional mass numbers for the additional data- Must have shape (:, 2).
The reference is the second column.
:returns: None. Generates a King-Plot based on the modified axes 'self.x_mod' and 'self.y_mod'
as well as the fit results 'self.results'.
"""
if self.results is None:
print('There are no results yet. Please use one of the fitting options.')
return
(a, b, sigma_a, sigma_b, corr, alpha) = self.results
scale_x = 1e-3
offset_str = ' - {}'.format(round(alpha * scale_x, 1)) if alpha != 0 else ''
xlabel = 'x{} (mode=\'shifts\' or =\'radii\' for right x-label)'.format(offset_str)
if mode == 'radii':
scale_x = 1.
offset_str = ' - {}'.format(round(alpha * scale_x, 1)) if alpha != 0 else ''
xlabel = r'$\mu\,\delta\langle r^2\rangle^{A, A^\prime}' + offset_str + r'\quad(\mathrm{u}\,\mathrm{fm}^2)$'
elif mode == 'shifts':
scale_x = 1e-3
offset_str = ' - {}'.format(round(alpha * scale_x, 1)) if alpha != 0 else ''
xlabel = r'$\mu\,\delta\nu_x^{A, A^\prime}' + offset_str + r'\quad(\mathrm{u\,GHz})$'
if sigma2d:
plt.plot((self.x_mod[:, 0] - alpha) * scale_x,
self.y_mod[:, 0] * self.scale_y, 'k.', label='Data')
draw_sigma2d((self.x_mod[:, 0] - alpha) * scale_x, self.y_mod[:, 0] * self.scale_y,
self.x_mod[:, 1] * scale_x, self.y_mod[:, 1] * self.scale_y, self.corr, n=2)
if add_xy is not None:
plt.plot((add_xy[:, 0] - alpha) * scale_x,
add_xy[:, 2] * self.scale_y, 'r.', label='Data')
draw_sigma2d((add_xy[:, 0] - alpha) * scale_x, add_xy[:, 2] * self.scale_y,
add_xy[:, 1] * scale_x, add_xy[:, 3] * self.scale_y, self.corr, n=2, **{'fmt': '-r'})
else:
plt.errorbar((self.x_mod[:, 0] - alpha) * scale_x, self.y_mod[:, 0] * self.scale_y,
xerr=self.x_mod[:, 1] * scale_x, yerr=self.y_mod[:, 1] * self.scale_y,
fmt='k.', label='Data')
if add_xy is not None:
plt.errorbar((add_xy[:, 0] - alpha) * scale_x, add_xy[:, 2] * self.scale_y,
xerr=add_xy[:, 1] * scale_x, yerr=add_xy[:, 3] * self.scale_y,
fmt='r.', label='Add. Data')
for j in range(len(self.a_fit)):
pm = -1. if b > 0 else 1.
plt.text((self.x_mod[j, 0] + 1. * self.x_mod[j, 1] - alpha) * scale_x,
(self.y_mod[j, 0] + pm * 5. * self.y_mod[j, 1]) * self.scale_y,
'{} - {}'.format(self.a_fit[j], self.a_ref[j]), fontsize=self.fontsize,
horizontalalignment='left', verticalalignment='top' if b > 0 else 'bottom')
if add_a is not None and add_xy is not None:
for j in range(add_xy.shape[0]):
pm = 1. if b > 0 else -1.
plt.text((add_xy[j, 0] - 1. * add_xy[j, 1] - alpha) * scale_x,
(add_xy[j, 2] + pm * 5. * add_xy[j, 3]) * self.scale_y,
'{} - {}'.format(*add_a[j]), fontsize=self.fontsize,
horizontalalignment='right', verticalalignment='bottom' if b > 0 else 'top')
min_x, max_x = np.min(self.x_mod[:, 0]), np.max(self.x_mod[:, 0])
if add_xy is not None:
min_x, max_x = np.min([min_x, np.min(add_xy[:, 0])]), np.max([max_x, np.max(add_xy[:, 0])])
x_cont = np.linspace(min_x - 0.1 * (max_x - min_x), max_x + 0.1 * (max_x - min_x), 1001) - alpha
plt.plot(x_cont * scale_x, straight(x_cont, a, b) * self.scale_y, 'b-', label='Fit')
y_min = straight(x_cont, a, b) - straight_std(x_cont, sigma_a, sigma_b, corr)
y_max = straight(x_cont, a, b) + straight_std(x_cont, sigma_a, sigma_b, corr)
plt.fill_between(x_cont * scale_x, y_min * self.scale_y, y_max * self.scale_y,
color='b', alpha=0.3, antialiased=True)
plt.xlabel(xlabel, fontsize=self.fontsize)
plt.ylabel(r'$\mu\,\delta\nu_y^{A, A^\prime}\quad(\mathrm{u\,GHz})$', fontsize=self.fontsize)
plt.xticks(fontsize=self.fontsize)
plt.yticks(fontsize=self.fontsize)
plt.legend(numpoints=1, loc='best', fontsize=self.fontsize)
plt.margins(0.1)
plt.tight_layout()
if show:
plt.show()
plt.gcf().clear()
[docs] def plot_nd(self, axis: int = 0, mode: str = '', sigma2d: bool = True):
"""
:param axis: The axis which is used as the x-axis throughout the plots.
:param mode: The mode of the King-plot. If mode='radii', the :math:'x'-axis must contain the differences of
mean square nuclear charge radii or the Lambda-factor. For every other value,
the :math:'x'-axis is assumed to be an isotope shift such that the slope corresponds to
a field-shift ratio :math:'F(y_i) / F(x)'.
:param sigma2d: Whether to draw the actual two-dimensional uncertainty bounds or the classical errorbars.
:returns: None. Generates a King-Plot based on the modified axes 'self.x_mod_nd' and 'self.y_mod_nd'
as well as the fit results 'self.results_nd'.
"""
if self.results_nd is None:
print('There are no results yet. Please use one of the fitting options.')
return
(a, b, sigma_a, sigma_b, corr) = self.results_nd
alpha = 0
scale_x = 1e-3
offset_str = ' - {}'.format(round(alpha * scale_x, 1)) if alpha != 0 else ''
xlabel = 'x{} (mode=\'shifts\' or =\'radii\' for right x-label)'.format(offset_str)
if mode == 'radii':
scale_x = 1.
offset_str = ' - {}'.format(round(alpha * scale_x, 1)) if alpha != 0 else ''
xlabel = r'$\mu\,\delta\langle r^2\rangle^{A, A^\prime}' + offset_str + r'\quad(\mathrm{u}\,\mathrm{fm}^2)$'
elif mode == 'shifts':
scale_x = 1e-3
offset_str = ' - {}'.format(round(alpha * scale_x, 1)) if alpha != 0 else ''
xlabel = r'$\mu\,\delta\nu_x^{A, A^\prime}' + offset_str + r'\quad(\mathrm{u\,GHz})$'
n_rows = self.x_nd_mod.shape[1] - 1
rows = [i for i in range(n_rows + 1) if i != axis]
for row, i in enumerate(rows):
plt.subplot(n_rows, 1, row + 1)
if sigma2d:
plt.plot((self.x_nd_mod[:, axis] - alpha) * scale_x,
self.x_nd_mod[:, i] * self.scale_y, 'k.', label='Data')
draw_sigma2d((self.x_nd_mod[:, axis] - alpha) * scale_x, self.x_nd_mod[:, i] * self.scale_y,
np.sqrt(self.cov[:, axis, axis]) * scale_x, np.sqrt(self.cov[:, i, i]) * self.scale_y,
self.cov[:, axis, i] / np.sqrt(self.cov[:, axis, axis] * self.cov[:, i, i]), n=2)
else:
plt.errorbar((self.x_nd_mod[:, axis] - alpha) * scale_x, self.x_nd_mod[:, i] * self.scale_y,
xerr=np.sqrt(self.cov[:, axis, axis]) * scale_x,
yerr=np.sqrt(self.cov[:, i, i]) * self.scale_y, fmt='k.', label='Data')
for j in range(len(self.a_fit)):
pm = -1. if b[row] > 0 else 1.
plt.text((self.x_nd_mod[j, axis] + 1. * np.sqrt(self.cov[j, axis, axis]) - alpha) * scale_x,
(self.x_nd_mod[j, i] + pm * 5. * np.sqrt(self.cov[j, i, i])) * self.scale_y,
'{} - {}'.format(self.a_fit[j], self.a_ref[j]), fontsize=self.fontsize,
horizontalalignment='left', verticalalignment='top' if b[row] > 0 else 'bottom')
min_x, max_x = np.min(self.x_nd_mod[:, axis]), np.max(self.x_nd_mod[:, axis])
x_cont = np.linspace(min_x - 0.1 * (max_x - min_x), max_x + 0.1 * (max_x - min_x), 1001) - alpha
plt.plot(x_cont * scale_x, straight(x_cont, a[row], b[row]) * self.scale_y, 'b-', label='Fit')
y_min = straight(x_cont, a[row], b[row]) \
- straight_std(x_cont, sigma_a[row], sigma_b[row], corr[row, row + n_rows])
y_max = straight(x_cont, a[row], b[row]) \
+ straight_std(x_cont, sigma_a[row], sigma_b[row], corr[row, row + n_rows])
plt.fill_between(x_cont * scale_x, y_min * self.scale_y, y_max * self.scale_y,
color='b', alpha=0.3, antialiased=True)
plt.xlabel(xlabel, fontsize=self.fontsize)
plt.ylabel(r'$\mu\,\delta\nu_{y' + str(i) + r'}^{A, A^\prime}\quad(\mathrm{u\,GHz})$',
fontsize=self.fontsize)
plt.xticks(fontsize=self.fontsize)
plt.yticks(fontsize=self.fontsize)
plt.legend(numpoints=1, loc='best', fontsize=self.fontsize)
plt.margins(0.1)
# plt.tight_layout()
plt.show()
plt.gcf().clear()