Source code for naima.core

# -*- coding: utf-8 -*-
# Licensed under a 3-clause BSD style license - see LICENSE.rst
import multiprocessing
import warnings
from import Iterable

import astropy.units as u
import emcee
import numpy as np
from astropy import log

from .utils import sed_conversion, validate_data_table

__all__ = [

# Define phsyical types used in plot and utils.validate_data_table
u.def_physical_type(u.erg / ** 2 / u.s, "flux")
u.def_physical_type(u.Unit("1/(s cm2 erg)"), "differential flux")
u.def_physical_type(u.Unit("1/(s erg)"), "differential power")
u.def_physical_type(u.Unit("1/TeV"), "differential energy")
u.def_physical_type(u.Unit("1/cm3"), "number density")
u.def_physical_type(u.Unit("1/(eV cm3)"), "differential number density")

# Prior functions

[docs]def uniform_prior(value, umin, umax): """Uniform prior distribution.""" if umin <= value <= umax: return 0.0 else: return -np.inf
[docs]def normal_prior(value, mean, sigma): """Normal prior distribution.""" return -0.5 * (2 * np.pi * sigma) - (value - mean) ** 2 / (2.0 * sigma)
def log_uniform_prior(value, umin=0, umax=None): """Log-uniform prior distribution.""" if value > 0 and value >= umin: if umax is not None: if value <= umax: return 1 / value else: return -np.inf else: return 1 / value else: return -np.inf # Probability function def lnprobmodel(model, data): # Check if conversion is required model_is_sed = model.unit.physical_type in ["power", "flux"] data_is_sed = data["flux"].unit.physical_type in ["power", "flux"] if model_is_sed != data_is_sed: unit, sed_factor = sed_conversion( data["energy"], model.unit, data_is_sed ) model = (model * sed_factor).to(data["flux"].unit) ul = data["ul"] notul = ~ul difference = model[notul] - data["flux"][notul] # use different errors for model above or below data sign = difference > 0 loerr, hierr = 1 * ~sign, 1 * sign logprob = -(difference ** 2) / ( 2.0 * ( loerr * data["flux_error_lo"][notul] + hierr * data["flux_error_hi"][notul] ) ** 2 ) totallogprob = np.sum(logprob) if np.sum(ul) > 0: # deal with upper limits at CL set by data['cl'] violated_uls = np.sum(model[ul] > data["flux"][ul]) totallogprob += violated_uls * np.log(1.0 - data["cl"][violated_uls]) return totallogprob def lnprob(pars, data, modelfunc, priorfunc): if priorfunc is None: lnprob_priors = 0.0 else: lnprob_priors = priorfunc(pars) modelout = modelfunc(pars, data) # Save blobs or save model if no blobs given if isinstance(modelout, Iterable) and not isinstance(modelout, np.ndarray): model = modelout[0] blob = modelout else: model = modelout # Avoid squeezing of the model output in order to retain units by # adding a second dummy object blob blob = (modelout, np.nan) if not np.isinf(lnprob_priors): lnprob_model = lnprobmodel(model, data) total_lnprob = lnprob_model + lnprob_priors else: total_lnprob = lnprob_priors return (total_lnprob, *blob) # Sampler funcs def _run_mcmc(sampler, pos, nrun): for i, state in enumerate( sampler.sample(pos, iterations=nrun, store=True) ): progress = 100.0 * float(i) / float(nrun) if progress % 5 < (5.0 / float(nrun)): print( "\nProgress of the run: {0:.0f} percent" " ({1} of {2} steps)".format(int(progress), i, nrun) ) paravg, parstd = [], [] npars = sampler.get_chain().shape[-1] for npar in range(npars): paravg.append(np.median(state.coords[:, npar])) parstd.append(np.std(state.coords[:, npar])) print( " " + (" ".join(["{%i:-^15}" % i for i in range(npars)])).format( *sampler.labels ) ) print( " Last ensemble median : " + (" ".join(["{%i:^15.3g}" % i for i in range(npars)])).format( *paravg ) ) print( " Last ensemble std : " + (" ".join(["{%i:^15.3g}" % i for i in range(npars)])).format( *parstd ) ) print( " Last ensemble lnprob : avg: {0:.3f}, max: {1:.3f}".format( np.average(state.log_prob), np.max(state.log_prob) ) ) return sampler, state def _prefit(p0, data, model, prior): P0_IS_ML = False from .extern.minimize import minimize def flat_prior(*args): return 0.0 if prior is None: prior = flat_prior def nll(*args): return -lnprob(*args)[0] "Finding Maximum Likelihood parameters through Nelder-Mead fitting..." )" Initial parameters: {0}".format(p0)) " Initial lnprob(p0): {0:.3f}".format(-nll(p0, data, model, prior)) ) with warnings.catch_warnings(): warnings.simplefilter("ignore") result = minimize( nll, p0, args=(data, model, flat_prior), method="Nelder-Mead", options={"maxfev": 500, "xtol": 1e-1, "ftol": 1e-3}, ) ll_prior = lnprob(result["x"], data, model, prior)[0] if (result["success"] or result["status"] == 1) and not np.isinf(ll_prior): # also keep result if we have reached maxiter, it is likely # better than p0 if result["status"] == 1:" Maximum number of function evaluations reached!") if result["status"] == 1:" New parameters : {0}".format(result["x"])) else:" New ML parameters : {0}".format(result["x"])) P0_IS_ML = True if -result["fun"] == ll_prior:" Maximum lnprob(p0): {0:.3f}".format(-result["fun"])) else:"flat prior lnprob(p0): {0:.3f}".format(-result["fun"]))"full prior lnprob(p0): {0:.3f}".format(ll_prior)) p0 = result["x"] elif np.isinf(ll_prior): log.warning( "Maximum Likelihood procedure converged on a parameter" " vector forbidden by prior," " using original parameters for MCMC" ) else: log.warning( "Maximum Likelihood procedure failed to converge," " using original parameters for MCMC" ) return p0, P0_IS_ML
[docs]def get_sampler( data_table=None, p0=None, model=None, prior=None, nwalkers=500, nburn=100, guess=True, interactive=False, prefit=False, labels=None, threads=None, data_sed=None, ): """Generate a new MCMC sampler. Parameters ---------- data_table : `~astropy.table.Table` or list of `~astropy.table.Table` Table containing the observed spectrum. If multiple tables are passed as a string, they will be concatenated in the order given. Each table needs at least these columns, with the appropriate associated units (with the physical type indicated in brackets below) as either a `~astropy.units.Unit` instance or parseable string: - ``energy``: Observed photon energy [``energy``] - ``flux``: Observed fluxes [``flux`` or ``differential flux``] - ``flux_error``: 68% CL gaussian uncertainty of the flux [``flux`` or ``differential flux``]. It can also be provided as ``flux_error_lo`` and ``flux_error_hi`` (see below). Optional columns: - ``energy_width``: Width of the energy bin [``energy``], or - ``energy_error``: Half-width of the energy bin [``energy``], or - ``energy_error_lo`` and ``energy_error_hi``: Distance from bin center to lower and upper bin edges [``energy``], or - ``energy_lo`` and ``energy_hi``: Energy edges of the corresponding energy bin [``energy``] - ``flux_error_lo`` and ``flux_error_hi``: 68% CL gaussian lower and upper uncertainties of the flux. - ``ul``: Flag to indicate that a flux measurement is an upper limit. - ``flux_ul``: Upper limit to the flux. If not present, the ``flux`` column will be taken as an upper limit for those measurements with the ``ul`` flag set to True or 1. The ``keywords`` metadata field of the table can be used to provide the confidence level of the upper limits with the keyword ``cl``, which defaults to 90%. The `` reader can recover all the needed information from ASCII tables in the :class:`` and :class:`` formats, and everything except the ``cl`` keyword from tables in the :class:``. For the latter format, the cl keyword can be added after reading the table with:: data.meta['keywords']['cl']=0.99 p0 : array Initial position vector. The distribution for the ``nwalkers`` walkers will be computed as a multidimensional gaussian of width 5% around the initial position vector ``p0``. model : function A function that takes a vector in the parameter space and the data dictionary, and returns the expected fluxes at the energies in the spectrum. Additional return objects will be saved as blobs in the sampler chain, see `the emcee documentation for the format <>`_. prior : function, optional A function that takes a vector in the parameter space and returns the log-likelihood of the Bayesian prior. Parameter limits can be specified through a uniform prior, returning 0. if the vector is within the parameter bounds and ``-np.inf`` otherwise. nwalkers : int, optional The number of Goodman & Weare “walkers”. Default is 500. nburn : int, optional Number of burn-in steps. After ``nburn`` steps, the sampler is reset and chain history discarded. It is necessary to settle the sampler into the maximum of the parameter space density. Default is 100. labels : iterable of strings, optional Labels for the parameters included in the position vector ``p0``. If not provided ``['par1','par2', ... ,'parN']`` will be used. threads : int, optional Number of parallel processes to use for sampling. Default is the number of CPU cores. guess : bool, optional Whether to attempt to guess the normalization (first) parameter of the model. Default is True. interactive : bool, optional Whether to launch the interactive fitting window to set the initial values for the prefitting or the MCMC run. Requires matplotlib. Default is False. prefit : bool, optional Whether to attempt to find the maximum likelihood parameters with a Nelder-Mead algorithm and use them as starting point of the MCMC run. The parameter values in `p0` will be used as starting points for the minimization. Note that the initial optimization is done without taking the prior function into account to avoid the possibility of infinite values in the objective function. If the best-fit parameter vector without prior is forbidden by the prior given, it will be discarded. data_sed : bool, optional When providing more than one data table, whether to convert them to SED format. If unset or None, all tables will be converted to the format of the first table. Returns ------- sampler : :class:`~emcee.EnsembleSampler` instance Ensemble sampler with walker positions after ``nburn`` burn-in steps. pos : :class:`~numpy.ndarray` Final position vector array. See also -------- emcee.EnsembleSampler """ if data_table is None: raise TypeError("Data table is missing!") else: data = validate_data_table(data_table, sed=data_sed) if model is None: raise TypeError("Model function is missing!") # Add parameter labels if not provided or too short if labels is None: # First is normalization labels = ["norm"] + ["par{0}".format(i) for i in range(1, len(p0))] elif len(labels) < len(p0): labels += ["par{0}".format(i) for i in range(len(labels), len(p0))] # Check that the model returns fluxes in same physical type as data modelout = model(p0, data) if (type(modelout) == tuple or type(modelout) == list) and ( type(modelout) != np.ndarray ): spec = modelout[0] else: spec = modelout # check whether both can be converted to same physical type through # sed_conversion try: # If both can be converted to differential flux, they can be compared # Otherwise, sed_conversion will raise a u.UnitsError sed_conversion(data["energy"], spec.unit, False) sed_conversion(data["energy"], data["flux"].unit, False) except u.UnitsError: raise u.UnitsError( "The physical type of the model and data units are not compatible," " please modify your model or data so they match:\n" " Model units: {0} [{1}]\n Data units: {2} [{3}]\n".format( spec.unit, spec.unit.physical_type, data["flux"].unit, data["flux"].unit.physical_type, ) ) if guess: normNames = ["norm", "ampl", "we", "wp"] normNameslog = ["log({0}".format(name) for name in normNames] normNameslog10 = ["log10({0}".format(name) for name in normNames] normNames += normNameslog + normNameslog10 idxs = [] for normName in normNames: for l2 in labels: if l2.lower().startswith(normName): # check with startswith to include normalization, # amplitude, etc. idxs.append(labels.index(l2)) if len(idxs) == 1: nunit, sedf = sed_conversion(data["energy"], spec.unit, False) currFlux = np.trapz( data["energy"] * (spec * sedf).to(nunit), data["energy"] ) nunit, sedf = sed_conversion( data["energy"], data["flux"].unit, False ) dataFlux = np.trapz( data["energy"] * (data["flux"] * sedf).to(nunit), data["energy"], ) ratio = dataFlux / currFlux if labels[idxs[0]].startswith("log("): p0[idxs[0]] += np.log(ratio) elif labels[idxs[0]].startswith("log10("): p0[idxs[0]] += np.log10(ratio) else: p0[idxs[0]] *= ratio elif len(idxs) == 0: log.warning( "No label starting with [{0}] found: not applying" " normalization guess.".format(",".join(normNames)) ) elif len(idxs) > 1: log.warning( "More than one label starting with [{0}] found:" " not applying normalization guess.".format( ",".join(normNames) ) ) P0_IS_ML = False if interactive: try:"Launching interactive model fitter, close when finished") import matplotlib.pyplot as plt from .model_fitter import InteractiveModelFitter iprev = plt.rcParams["interactive"] plt.rcParams["interactive"] = False imf = InteractiveModelFitter( model, p0, data, labels=labels, sed=True ) p0 = P0_IS_ML = imf.P0_IS_ML plt.rcParams["interactive"] = iprev except ImportError as e: log.warning( "Interactive fitting is not available because" " matplotlib is not installed: {0}".format(e) ) # If we already did the prefit call in ModelWidget (and didn't modify the # parameters afterwards), avoid doing it here if prefit and not P0_IS_ML: p0, P0_IS_ML = _prefit(p0, data, model, prior) sampler = emcee.EnsembleSampler( nwalkers, len(p0), lnprob, args=[data, model, prior], pool=multiprocessing.Pool(threads), blobs_dtype=np.dtype(object), ) # Add data and parameters properties to sampler sampler.data_table = data_table = data sampler.labels = labels # Add model function to sampler sampler.modelfn = model # Add run_info dict sampler.run_info = { "n_walkers": nwalkers, "n_burn": nburn, # convert from np.float to regular float "p0": [float(p) for p in p0], "guess": guess, } # Initialize walkers in a ball of relative size 0.5% in all dimensions if # the parameters have been fit to their ML values, or to 10% otherwise spread = 0.005 if P0_IS_ML else 0.1 p0var = np.array([spread * pp for pp in p0]) p0 = np.vstack( [p0 + p0var * np.random.normal(size=len(p0)) for i in range(nwalkers)] ) if nburn > 0: print( "Burning in the {0} walkers with {1} steps...".format( nwalkers, nburn ) ) sampler, state = _run_mcmc(sampler, p0, nburn) else: state = emcee.State(p0) sampler.run_info["p0_burn_median"] = [ float(p) for p in np.median(state.coords, axis=0) ] return sampler, state
[docs]def run_sampler(nrun=100, sampler=None, pos=None, **kwargs): """Run an MCMC sampler. If no sampler or initial position vector is provided, extra ``kwargs`` are passed to `get_sampler` to generate a new sampler. Parameters ---------- nrun : int, optional Number of steps to run sampler : :class:`~emcee.EnsembleSampler` instance, optional Sampler. pos : :class:`~numpy.ndarray`, optional A list of initial position vectors for the walkers. It should have dimensions of ``(nwalkers,dim)``, where ``dim`` is the number of free parameters. `emcee.utils.sample_ball` can be used to generate a multidimensional gaussian distribution around a single initial position. Returns ------- sampler : :class:`~emcee.EnsembleSampler` instance Sampler containing the paths of the walkers during the ``nrun`` steps. pos : array List of final position vectors after the run. """ if sampler is None or pos is None: sampler, pos = get_sampler(**kwargs) sampler.run_info["n_run"] = nrun print("\nWalker burn in finished, running {0} steps...".format(nrun)) sampler.reset() sampler, pos = _run_mcmc(sampler, pos, nrun) return sampler, pos