diff options
Diffstat (limited to 'golemflavor')
| -rw-r--r-- | golemflavor/fr.py | 50 | ||||
| -rw-r--r-- | golemflavor/llh.py | 27 | ||||
| -rw-r--r-- | golemflavor/mcmc.py | 17 | ||||
| -rw-r--r-- | golemflavor/plot.py | 74 |
4 files changed, 105 insertions, 63 deletions
diff --git a/golemflavor/fr.py b/golemflavor/fr.py index 6945ce4..3b5d259 100644 --- a/golemflavor/fr.py +++ b/golemflavor/fr.py @@ -19,25 +19,25 @@ from golemflavor.misc import enum_parse, parse_bool import mpmath as mp mp.mp.dps = 100 # Computation precision -# DTYPE = np.float128 -# CDTYPE = np.complex256 -# PI = np.arccos(DTYPE(-1)) -# SQRT = np.sqrt -# COS = np.cos -# SIN = np.sin -# ACOS = np.arccos -# ASIN = np.arcsin -# EXP = np.exp - -DTYPE = mp.mpf -CDTYPE = mp.mpc -PI = mp.pi -SQRT = mp.sqrt -COS = mp.cos -SIN = mp.sin -ACOS = mp.acos -ASIN = mp.asin -EXP = mp.exp +DTYPE = np.float128 +CDTYPE = np.complex256 +PI = np.arccos(DTYPE(-1)) +SQRT = np.sqrt +COS = np.cos +SIN = np.sin +ACOS = np.arccos +ASIN = np.arcsin +EXP = np.exp + +# DTYPE = mp.mpf +# CDTYPE = mp.mpc +# PI = mp.pi +# SQRT = mp.sqrt +# COS = mp.cos +# SIN = mp.sin +# ACOS = mp.acos +# ASIN = mp.asin +# EXP = mp.exp MASS_EIGENVALUES = [7.40E-23, 2.515E-21] """SM mass eigenvalues.""" @@ -237,8 +237,8 @@ def cardano_eqn(ham): return mm -def normalise_fr(fr): - """Normalise an input flavor combination to a flavor ratio. +def normalize_fr(fr): + """Normalize an input flavor combination to a flavor ratio. Parameters ---------- @@ -251,8 +251,8 @@ def normalise_fr(fr): Examples ---------- - >>> from fr import normalise_fr - >>> print(normalise_fr((1, 2, 3))) + >>> from fr import normalize_fr + >>> print(normalize_fr((1, 2, 3))) array([ 0.16666667, 0.33333333, 0.5 ]) """ @@ -294,7 +294,7 @@ def fr_to_angles(ratios): ---------- TODO(shivesh) """ - fr0, fr1, fr2 = normalise_fr(ratios) + fr0, fr1, fr2 = normalize_fr(ratios) cphi2 = fr2 sphi2 = (1.0 - cphi2) @@ -307,7 +307,7 @@ def fr_to_angles(ratios): sphi4 = sphi2**2 c2psi = COS(ACOS(SQRT(cpsi2))*2) - return map(float, (sphi4, c2psi)) + return (sphi4, c2psi) NUFIT_U = angles_to_u((0.307, (1-0.02195)**2, 0.565, 3.97935)) diff --git a/golemflavor/llh.py b/golemflavor/llh.py index 14678b6..645ba6e 100644 --- a/golemflavor/llh.py +++ b/golemflavor/llh.py @@ -23,15 +23,34 @@ from golemflavor.misc import enum_parse, gen_identifier, parse_bool def GaussianBoundedRV(loc=0., sigma=1., lower=-np.inf, upper=np.inf): - """Normalised gaussian bounded between lower and upper values""" + """Normalized Gaussian bounded between lower and upper values""" low, up = (lower - loc) / sigma, (upper - loc) / sigma g = scipy.stats.truncnorm(loc=loc, scale=sigma, a=low, b=up) return g -def multi_gaussian(fr, fr_bf, sigma, offset=-320): - """Multivariate gaussian likelihood.""" - cov_fr = np.identity(3) * sigma +def multi_gaussian(fr, fr_bf, smearing, offset=-320): + """ + Multivariate Gaussian log likelihood. + + Parameters + ---------- + fr : List[float], length 3 + The flavour composition to evaluate at. + fr_bf : List[float], length 3 + The bestfit / injected flavour composition. + smearing : float + The amount of smearing. + offset : float, optional + An amount to offset the magnitude of the log likelihood. + + Returns + ---------- + llh : float + The log likelihood evaluated at `fr`. + + """ + cov_fr = np.identity(3) * pow(smearing, 2) return np.log(multivariate_normal.pdf(fr, mean=fr_bf, cov=cov_fr)) + offset diff --git a/golemflavor/mcmc.py b/golemflavor/mcmc.py index a1d3e27..c002620 100644 --- a/golemflavor/mcmc.py +++ b/golemflavor/mcmc.py @@ -9,10 +9,14 @@ Useful functions to use an MCMC for the BSM flavor ratio analysis from __future__ import absolute_import, division, print_function +import sys from functools import partial import emcee -import tqdm +if 'ipykernel' in sys.modules: + from tqdm import tqdm_notebook as tqdm +else: + from tqdm import tqdm import numpy as np @@ -20,21 +24,20 @@ from golemflavor.enums import MCMCSeedType from golemflavor.misc import enum_parse, make_dir, parse_bool -def mcmc(p0, ln_prob, ndim, nwalkers, burnin, nsteps, args, threads=1): +def mcmc(p0, ln_prob, ndim, nwalkers, burnin, nsteps, threads=1): """Run the MCMC.""" sampler = emcee.EnsembleSampler( nwalkers, ndim, ln_prob, threads=threads ) print("Running burn-in") - for result in tqdm.tqdm(sampler.sample(p0, iterations=burnin), total=burnin): + for result in tqdm(sampler.sample(p0, iterations=burnin), total=burnin): pos, prob, state = result sampler.reset() print("Finished burn-in") - args.burnin = False print("Running") - for _ in tqdm.tqdm(sampler.sample(pos, iterations=nsteps), total=nsteps): + for _ in tqdm(sampler.sample(pos, iterations=nsteps), total=nsteps): pass print("Finished") @@ -114,6 +117,10 @@ def save_chains(chains, outfile): Output file location of chains """ + if outfile[-4:] == '.npy': + of = outfile + else: + of = outfile + '.npy' make_dir(outfile) print('Saving chains to location {0}'.format(outfile+'.npy')) np.save(outfile+'.npy', chains) diff --git a/golemflavor/plot.py b/golemflavor/plot.py index abc4633..de645ee 100644 --- a/golemflavor/plot.py +++ b/golemflavor/plot.py @@ -10,18 +10,24 @@ Plotting functions for the BSM flavor ratio analysis from __future__ import absolute_import, division, print_function import os +import sys import socket from copy import deepcopy +import warnings +warnings.filterwarnings("ignore") + import numpy as np import numpy.ma as ma from scipy.interpolate import splev, splprep from scipy.ndimage.filters import gaussian_filter +import matplotlib import matplotlib as mpl import matplotlib.patches as patches import matplotlib.gridspec as gridspec -mpl.use('Agg') +if 'submitter' in socket.gethostname() or 'cobalt' in socket.gethostname(): + mpl.use('Agg', warn=False) from matplotlib import rc from matplotlib import pyplot as plt @@ -51,6 +57,24 @@ from golemflavor.misc import get_units, make_dir, solve_ratio, interval from golemflavor.fr import angles_to_u, flat_angles_to_u, angles_to_fr from golemflavor.fr import SCALE_BOUNDARIES +if os.path.isfile('./plot_llh/paper.mplstyle'): + plt.style.use('./plot_llh/paper.mplstyle') +elif os.path.isfile('./paper.mplstyle'): + plt.style.use('./paper.mplstyle') +elif os.environ.get('GOLEMSOURCEPATH') is not None: + plt.style.use(os.environ['GOLEMSOURCEPATH']+'/GolemFit/scripts/paper/paper.mplstyle') +if 'submitter' in socket.gethostname(): + rc('text', usetex=False) +else: + rc('text', usetex=True) + +mpl.rcParams['text.latex.preamble'] = [ + r'\usepackage{xcolor}', + r'\usepackage{amsmath}', + r'\usepackage{amssymb}'] +if sys.version_info < (3, 0): + mpl.rcParams['text.latex.unicode'] = True + BAYES_K = 1. # Strong degree of belief. # BAYES_K = 3/2. # Very strong degree of belief. @@ -76,20 +100,6 @@ PLANCK_SCALE = { } -if os.path.isfile('./plot_llh/paper.mplstyle'): - plt.style.use('./plot_llh/paper.mplstyle') -elif os.environ.get('GOLEMSOURCEPATH') is not None: - plt.style.use(os.environ['GOLEMSOURCEPATH']+'/GolemFit/scripts/paper/paper.mplstyle') -if 'submitter' in socket.gethostname(): - rc('text', usetex=False) - -mpl.rcParams['text.latex.preamble'] = [ - r'\usepackage{xcolor}', - r'\usepackage{amsmath}', - r'\usepackage{amssymb}'] -mpl.rcParams['text.latex.unicode'] = True - - def gen_figtext(args): """Generate the figure text.""" t = r'$' @@ -209,11 +219,11 @@ def heatmap(data, scale, vmin=None, vmax=None, style='triangular'): vertices = [] for v, value in vertices_values: - vertices.append(v) + vertices.append(list(v)) return vertices -def get_tax(ax, scale, ax_labels, rot_ax_labels=False, fontsize=23): +def get_tax(ax, scale, ax_labels=None, rot_ax_labels=False, fontsize=23): ax.set_aspect('equal') # Boundary and Gridlines @@ -227,6 +237,12 @@ def get_tax(ax, scale, ax_labels, rot_ax_labels=False, fontsize=23): # Set Axis labels and Title if rot_ax_labels: roty, rotz = (-60, 60) else: roty, rotz = (0, 0) + if ax_labels is None: + ax_labels = [ + r'$\nu_e\:\:{\rm fraction}\:\left( f_{e,\oplus}\right)$', + r'$\nu_\mu\:\:{\rm fraction}\:\left( f_{\mu,\oplus}\right)$', + r'$\nu_\tau\:\:{\rm fraction}\:\left( f_{\tau,\oplus}\right)$' + ] tax.bottom_axis_label( ax_labels[0], fontsize=fontsize+4, position=(0.55, -0.10/2, 0.5), rotation=0 @@ -340,7 +356,7 @@ def flavor_contour(frs, nbins, coverage, ax=None, smoothing=0.4, **kwargs): """Plot the flavor contour for a specified coverage.""" # Histogram in flavor space - os_nbins = nbins * oversample + os_nbins = int(nbins * oversample) H, b = np.histogramdd( (frs[:,0], frs[:,1], frs[:,2]), bins=(os_nbins+1, os_nbins+1, os_nbins+1), @@ -403,7 +419,7 @@ def flavor_contour(frs, nbins, coverage, ax=None, smoothing=0.4, ev_polygon = np.dstack((xi, yi))[0] # Remove points interpolated outside flavor triangle - f_ev_polygon = np.array(map(lambda x: project_toflavor(x, nbins), ev_polygon)) + f_ev_polygon = np.array(list(map(lambda x: project_toflavor(x, nbins), ev_polygon))) xf, yf, zf = f_ev_polygon.T mask = np.array((xf < 0) | (yf < 0) | (zf < 0) | (xf > nbins) | @@ -425,10 +441,10 @@ def flavor_contour(frs, nbins, coverage, ax=None, smoothing=0.4, else: return ev_polygon -def plot_Tchain(Tchain, axes_labels, ranges): +def plot_Tchain(Tchain, axes_labels, ranges, names=None): """Plot the Tchain using getdist.""" Tsample = mcsamples.MCSamples( - samples=Tchain, labels=axes_labels, ranges=ranges + samples=Tchain, names=names, labels=axes_labels, ranges=ranges ) Tsample.updateSettings({'contours': [0.90, 0.99]}) @@ -774,12 +790,12 @@ def plot_table_sens(data, outfile, outformat, args, show_lvatmo=True): fig.savefig(outfile+'.'+of, bbox_inches='tight', dpi=150) -def plot_x(data, outfile, outformat, args, normalise=False): +def plot_x(data, outfile, outformat, args, normalize=False): """Limit plot as a function of the source flavor ratio for each operator texture.""" print('Making X sensitivity plot') dim = args.dimension - if dim < 5: normalise = False + if dim < 5: normalize = False srcs = args.source_ratios x_arr = np.array([i[0] for i in srcs]) if args.texture is Texture.NONE: @@ -802,7 +818,7 @@ def plot_x(data, outfile, outformat, args, normalise=False): ax = fig.add_subplot(111) ylims = SCALE_BOUNDARIES[dim] - if normalise: + if normalize: if dim == 5: ylims = (-24, -8) if dim == 6: ylims = (-12, 8) if dim == 7: ylims = (0, 20) @@ -879,7 +895,7 @@ def plot_x(data, outfile, outformat, args, normalise=False): scales, statistic = ma.compress_rows(d).T lim = get_limit(deepcopy(scales), deepcopy(statistic), args, mask_initial=True) if lim is None: continue - if normalise: + if normalize: lim -= np.log10(PLANCK_SCALE[dim]) lims[isrc] = lim @@ -888,7 +904,7 @@ def plot_x(data, outfile, outformat, args, normalise=False): if size == 0: continue print('x_arr, lims', zip(x_arr, lims)) - if normalise: + if normalize: zeropoint = 100 else: zeropoint = 0 @@ -944,7 +960,7 @@ def plot_x(data, outfile, outformat, args, normalise=False): ) LV_lim = np.log10(LV_ATMO_90PC_LIMITS[dim]) - if normalise: + if normalize: LV_lim -= np.log10(PLANCK_SCALE[dim]) ax.add_patch(patches.Rectangle( (xlims[0], LV_lim[1]), np.diff(xlims), LV_lim[0]-LV_lim[1], @@ -953,7 +969,7 @@ def plot_x(data, outfile, outformat, args, normalise=False): if dim in PLANCK_SCALE: ps = np.log10(PLANCK_SCALE[dim]) - if normalise and dim == 6: + if normalize and dim == 6: ps -= np.log10(PLANCK_SCALE[dim]) ax.add_patch(Arrow( 0.24, -0.009, 0, -5, width=0.12, capstyle='butt', @@ -975,7 +991,7 @@ def plot_x(data, outfile, outformat, args, normalise=False): ax.axhline(y=ps, color='purple', alpha=1., linewidth=1.5) cpt = r'c' if dim % 2 == 0 else r'a' - if normalise: + if normalize: ft = r'${\rm New\:Physics\:Scale}\:[\:{\rm log}_{10} \left (\mathring{'+cpt+r'}^{(' + \ r'{0}'.format(args.dimension)+r')}\cdot{\rm E}_{\:\rm P}' if dim > 5: ft += r'^{\:'+ r'{0}'.format(args.dimension-4)+ r'}' |