Source code for smefit.analyze.contours_2d

# -*- coding: utf-8 -*-
import matplotlib.pyplot as plt
import numpy as np
import scipy.stats
import seaborn as sns
from matplotlib import patches, transforms
from matplotlib.patches import Ellipse


[docs] def split_solution(full_solution): """Split the posterior solution""" min_val = min(full_solution) max_val = max(full_solution) mid = np.mean([max_val, min_val]) # solution 1 should be closer to 0 solution1 = full_solution[full_solution < mid] solution2 = full_solution[full_solution > mid] if min(abs(solution2)) < min(abs(solution1)): solution1, solution2 = solution2, solution1 return solution1, solution2
[docs] def confidence_ellipse( coeff1, coeff2, ax, facecolor="none", confidence_level=95, **kwargs ): """ Draws 95% C.L. ellipse for data points `x` and `y` Parameters ---------- coeff1: array_like ``(N,) ndarray`` containing ``N`` posterior samples for the first coefficient coeff2: array_like ``(N,) ndarray`` containing ``N`` posterior samples for the first coefficient ax: matplotlib.axes Axes object to plot on facecolor: str, optional Color of the ellipse **kwargs Additional plotting settings passed to matplotlib.patches.Ellipse Returns ------- matplotlib.patches.Ellipse Ellipse object """ if coeff1.size != coeff2.size: raise ValueError("coeff1 and coeff2 must be the same size") # construct covariance matrix of coefficients cov = np.cov(coeff1, coeff2) # diagonalise eig_val, eig_vec = np.linalg.eig(cov) # eigenvector with largest eigenvalue eig_vec_max = eig_vec[:, np.argmax(eig_val)] # angle of eigenvector with largest eigenvalue with the horizontal axis cos_th = eig_vec[0, np.argmax(eig_val)] / np.linalg.norm(eig_vec_max) if eig_vec_max[1] > 0: inclination = np.arccos(cos_th) else: # pay attention to range of arccos (extend to [0, -\pi] domain) inclination = -np.arccos(cos_th) eigval_sort = np.sort(eig_val) chi2_qnt = scipy.stats.chi2.ppf(confidence_level / 100.0, 2) ell_radius_x = np.sqrt(chi2_qnt * eigval_sort[-1]) ell_radius_y = np.sqrt(chi2_qnt * eigval_sort[-2]) ellipse = Ellipse( (0, 0), width=ell_radius_x * 2, height=ell_radius_y * 2, facecolor=facecolor, **kwargs ) mean_coeff1 = np.median(coeff1) mean_coeff2 = np.median(coeff2) transf = ( transforms.Affine2D().rotate(inclination).translate(mean_coeff1, mean_coeff2) ) ellipse.set_transform(transf + ax.transData) return ax.add_patch(ellipse)
[docs] def plot_contours( ax, posterior, ax_labels, coeff1, coeff2, kde, clr_idx, confidence_level=95, double_solution=None, ): """ Parameters ---------- ax: matplotlib.axes Axes object to plot on posterior: pandas.DataFrame Posterior samples per coefficient ax_labels: list Latex names coeff1: str Name of first coefficient coeff2: str Name of second coefficient kde: bool Performs kernel density estimate (kde) when quadratics are turned on clr_idx: int Color index that makes sure each fit gets associated a different color confidence_level: int, optional Confidence level interval, set to 95% by default double_solution: dict, optional Dictionary of operators with double (disjoint) solution per fit Returns ------- hndls: list List of Patch objects """ colors = plt.rcParams["axes.prop_cycle"].by_key()["color"] x_values = posterior[coeff2].values y_values = posterior[coeff1].values # perform kde for quadratic EFT fit if kde: solution1x = solution2x = x_values if coeff2 in double_solution: solution1x, solution2x = split_solution(posterior[coeff2]) solution1y = solution2y = y_values if coeff1 in double_solution: solution1y, solution2y = split_solution(posterior[coeff1]) sns.kdeplot( x=x_values, y=y_values, levels=[1 - confidence_level / 100.0, 1.0], bw_adjust=1.2, ax=ax, fill=True, alpha=0.3, color=colors[clr_idx], ) sns.kdeplot( x=x_values, y=y_values, levels=[1 - confidence_level / 100.0], bw_adjust=1.2, ax=ax, alpha=1, color=colors[clr_idx], label=confidence_level, ) ax.scatter( np.mean(solution1x), np.mean(solution1y), c=colors[clr_idx], s=50, marker="o", ) ax.scatter( np.mean(solution2x), np.mean(solution2y), c=colors[clr_idx], s=50, marker="o", ) hndls = ( patches.Patch(ec=colors[clr_idx], fc=colors[clr_idx], fill=True, alpha=0.3), patches.Patch( ec=colors[clr_idx], fc=colors[clr_idx], fill=False, alpha=1.0 ), ) else: # draw ellipses for linear EFT fit p1 = confidence_ellipse( x_values, y_values, ax, alpha=1, edgecolor=colors[clr_idx], confidence_level=confidence_level, ) p2 = confidence_ellipse( x_values, y_values, ax, alpha=0.3, facecolor=colors[clr_idx], edgecolor=None, confidence_level=confidence_level, ) ax.scatter( np.mean(x_values, axis=0), np.mean(y_values, axis=0), c=colors[clr_idx], s=50, marker="o", ) hndls = (p1, p2) plt.xlabel(ax_labels[1], fontsize=26) plt.ylabel(ax_labels[0], fontsize=26) plt.tick_params(which="both", direction="in", labelsize=22) return hndls