<commit_before># Licensed under a 3-clause BSD style license - see LICENSE.rst
from __future__ import (absolute_import, division, print_function,
                        unicode_literals)

import numpy as np
import astropy.units as u
from astropy.extern import six
from astropy import log

from .core import lnprobmodel, _prefit
from .plot import color_cycle, _plot_data_to_ax
from .utils import sed_conversion, validate_data_table
from .extern.validator import validate_array

__all__ = ['InteractiveModelFitter']

def _process_model(model):
    if ((isinstance(model, tuple) or isinstance(model, list))
            and not isinstance(model,np.ndarray)):
        return model[0]
    else:
        return model


class InteractiveModelFitter(object):
    """
    Interactive model fitter using matplotlib widgets


    Parameters
    ----------
    modelfn : function
        A function that takes a vector in the parameter space and the data
        table, and returns the expected fluxes at the energies in the
        spectrum.
    p0 : array
        Initial position vector.
    data : `~astropy.table.Table` or list of `~astropy.table.Table`
        Table or tables with observed spectra. Must follow the format described
        in `naima.run_sampler`.
    e_range : list of `~astropy.units.Quantity`, length 2, optional
        Limits in energy for the computation of the model.
        Note that setting this parameter will mean that the model output is
        computed twice when `data` is provided: once for display using
        `e_range` and once for computation of the log-probability using the
        energy values of the spectra.
    e_npoints : int, optional
        How many points to compute for the model if `e_range` is set. 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.
    sed : bool, optional
        Whether to plot SED or differential spectrum.
    auto_update : bool, optional
        Whether to update the model plot when parameter sliders are changed.
        Default is True and can also be changed through the GUI.
    """
    def __init__(self, modelfn, p0, data=None, e_range=None, e_npoints=100,
            labels=None, sed=True, auto_update=True):

        import matplotlib.pyplot as plt
        from matplotlib.widgets import Button, Slider, CheckButtons

        self.pars = p0
        self.P0_IS_ML = False
        npars = len(p0)
        if labels is None:
            labels = ['par{0}'.format(i) for i in range(npars)]
        elif len(labels) < npars:
            labels += ['par{0}'.format(i) for i in range(len(labels),npars)]

        self.hasdata = data is not None

        self.modelfn = modelfn

        self.fig = plt.figure()
        modelax = plt.subplot2grid((10,4),(0,0),rowspan=4,colspan=4)

        if e_range:
            e_range = validate_array('e_range', u.Quantity(e_range),
                    physical_type='energy')
            e_unit = e_range.unit
            energy = np.logspace(np.log10(e_range[0].value),
                    np.log10(e_range[1].value), e_npoints) * e_unit
        else:
            energy = np.logspace(-4,2,e_npoints)*u.TeV
            e_unit = u.TeV

        if sed:
            flux = np.zeros(e_npoints) * u.Unit('erg/(cm2 s)')
        else:
            flux = np.zeros(e_npoints) * u.Unit('1/(TeV cm2 s)')

        self.data = None
        if self.hasdata:
            self.data = validate_data_table(data)
            e_unit = self.data['energy'].unit
            _plot_data_to_ax(self.data, modelax, sed=sed, e_unit=e_unit)
            if not e_range:
                energy = self.data['energy']
                flux = self.data['flux']

        self.data_for_model = {'energy': energy,
                'flux': flux}
        model = _process_model(self.modelfn(p0, self.data_for_model))

        if self.hasdata:
            if not np.all(self.data_for_model['energy'] == self.data['energy']):
                # this will be sloooow, maybe interpolate already computed model?
                model_for_lnprob = _process_model(
                        self.modelfn(self.pars, self.data))
            else:
                model_for_lnprob = model
            lnprob = lnprobmodel(model_for_lnprob, self.data)
            if isinstance(lnprob,u.Quantity):
                lnprob = lnprob.decompose().value
            self.lnprobtxt = modelax.text(0.05, 0.05, r'', ha='left', va='bottom',
                    transform=modelax.transAxes, size=20)
            self.lnprobtxt.set_text(r'$\ln\mathcal{{L}} = {0:.1f}$'.format(lnprob))


        self.f_unit, self.sedf = sed_conversion(energy, model.unit, sed)
        if self.hasdata:
            modelax.set_xlim(
                    (self.data['energy'][0] - self.data['energy_error_lo'][0]).to(e_unit).value / 3,
                    (self.data['energy'][-1] + self.data['energy_error_hi'][-1]).to(e_unit).value * 3)
        else:
            # plot_data_to_ax has not set ylabel
            unit = self.f_unit.to_string('latex_inline')
            if sed:
                modelax.set_ylabel(r'$E^2 dN/dE [{0}]$'.format(unit))
            else:
                modelax.set_ylabel(r'$dN/dE [{0}]$'.format(unit))
            modelax.set_xlim(energy[0].value/3,
                    energy[-1].value*3)

        self.line, = modelax.loglog(energy,
                (model*self.sedf).to(self.f_unit), lw=2,
                c='k', zorder=10)

        modelax.set_xlabel(
                'Energy [{0}]'.format(energy.unit.to_string('latex_inline')))


        paraxes = []
        for n in range(npars):
            paraxes.append(plt.subplot2grid((2*npars,10),(npars+n,0),colspan=7))
        self.parsliders = []
        slider_props = {'facecolor': color_cycle[-1], 'alpha': 0.5}
        for label,parax,valinit in six.moves.zip(labels, paraxes, p0):
            # Attempt to estimate reasonable parameter ranges from label
            pmin, pmax = valinit/10, valinit*3
            if 'log' in label:
                span = 2
                if 'norm' in label or 'amplitude' in label:
                    # give more range for normalization parameters
                    span *= 2
                pmin, pmax = valinit - span, valinit + span
            elif ('index' in label) or ('alpha' in label):
                if valinit > 0.:
                    pmin, pmax = 0, 5
                else:
                    pmin, pmax = -5, 0
            elif 'norm' in label or 'amplitude' in label:
                # norm without log, it will not be pretty because sliders are
                # only linear
                pmin, pmax = valinit / 100, valinit * 100
<commit_msg>various fixes for IMF<commit_after># Licensed under a 3-clause BSD style license - see LICENSE.rst
from __future__ import (absolute_import, division, print_function,
                        unicode_literals)

import numpy as np
import astropy.units as u
from astropy.extern import six
from astropy import log

from .core import lnprobmodel, _prefit
from .plot import color_cycle, _plot_data_to_ax
from .utils import sed_conversion, validate_data_table
from .extern.validator import validate_array

__all__ = ['InteractiveModelFitter']

def _process_model(model):
    if ((isinstance(model, tuple) or isinstance(model, list))
            and not isinstance(model,np.ndarray)):
        return model[0]
    else:
        return model


class InteractiveModelFitter(object):
    """
    Interactive model fitter using matplotlib widgets


    Parameters
    ----------
    modelfn : function
        A function that takes a vector in the parameter space and the data
        table, and returns the expected fluxes at the energies in the
        spectrum.
    p0 : array
        Initial position vector.
    data : `~astropy.table.Table` or list of `~astropy.table.Table`
        Table or tables with observed spectra. Must follow the format described
        in `naima.run_sampler`.
    e_range : list of `~astropy.units.Quantity`, length 2, optional
        Limits in energy for the computation of the model.
        Note that setting this parameter will mean that the model output is
        computed twice when `data` is provided: once for display using
        `e_range` and once for computation of the log-probability using the
        energy values of the spectra.
    e_npoints : int, optional
        How many points to compute for the model if `e_range` is set. 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.
    sed : bool, optional
        Whether to plot SED or differential spectrum.
    auto_update : bool, optional
        Whether to update the model plot when parameter sliders are changed.
        Default is True and can also be changed through the GUI.
    """
    def __init__(self, modelfn, p0, data=None, e_range=None, e_npoints=100,
            labels=None, sed=True, auto_update=True):

        import matplotlib.pyplot as plt
        from matplotlib.widgets import Button, Slider, CheckButtons

        self.pars = p0
        self.P0_IS_ML = False
        npars = len(p0)
        if labels is None:
            labels = ['par{0}'.format(i) for i in range(npars)]
        elif len(labels) < npars:
            labels += ['par{0}'.format(i) for i in range(len(labels),npars)]

        self.hasdata = data is not None
        self.data = None
        if self.hasdata:
            self.data = validate_data_table(data, sed=sed)

        self.modelfn = modelfn
        self.fig = plt.figure()
        modelax = plt.subplot2grid((10,4),(0,0),rowspan=4,colspan=4)

        if e_range:
            e_range = validate_array('e_range', u.Quantity(e_range),
                    physical_type='energy')
            energy = np.logspace(np.log10(e_range[0].value),
                    np.log10(e_range[1].value), e_npoints) * e_range.unit
            if self.hasdata:
                energy = energy.to(self.data['energy'].unit)
            else:
                e_unit = e_range.unit
        else:
            energy = np.logspace(-4,2,e_npoints)*u.TeV
            e_unit = u.TeV

        # Bogus flux array to send to model if not using data
        if sed:
            flux = np.zeros(e_npoints) * u.Unit('erg/(cm2 s)')
        else:
            flux = np.zeros(e_npoints) * u.Unit('1/(TeV cm2 s)')

        if self.hasdata:
            e_unit = self.data['energy'].unit
            _plot_data_to_ax(self.data, modelax, sed=sed, e_unit=e_unit)
            if not e_range:
                # use data for model
                energy = self.data['energy']
                flux = self.data['flux']

        self.data_for_model = {'energy': energy, 'flux': flux}

        model = _process_model(self.modelfn(p0, self.data_for_model))

        if self.hasdata:
            if not np.all(self.data_for_model['energy'] == self.data['energy']):
                # this will be sloooow, maybe interpolate already computed model?
                model_for_lnprob = _process_model(
                        self.modelfn(self.pars, self.data))
            else:
                model_for_lnprob = model
            lnprob = lnprobmodel(model_for_lnprob, self.data)
            if isinstance(lnprob,u.Quantity):
                lnprob = lnprob.decompose().value
            self.lnprobtxt = modelax.text(0.05, 0.05, r'', ha='left', va='bottom',
                    transform=modelax.transAxes, size=20)
            self.lnprobtxt.set_text(r'$\ln\mathcal{{L}} = {0:.1f}$'.format(lnprob))


        self.f_unit, self.sedf = sed_conversion(energy, model.unit, sed)

        if self.hasdata:
            datamin = (self.data['energy'][0] - self.data['energy_error_lo'][0]).to(e_unit).value / 3
            datamax = (self.data['energy'][-1] + self.data['energy_error_hi'][-1]).to(e_unit).value * 3
            xmin = min(datamin, energy[0].to(e_unit).value)
            xmax = max(datamin, energy[-1].to(e_unit).value)
            modelax.set_xlim(xmin, xmax)
        else:
            # plot_data_to_ax has not set ylabel
            unit = self.f_unit.to_string('latex_inline')
            if sed:
                modelax.set_ylabel(r'$E^2 dN/dE$ [{0}]'.format(unit))
            else:
                modelax.set_ylabel(r'$dN/dE$ [{0}]'.format(unit))
            modelax.set_xlim(energy[0].value, energy[-1].value)

        self.line, = modelax.loglog(energy.to(e_unit),
                (model*self.sedf).to(self.f_unit), lw=2,
                c='k', zorder=10)

        modelax.set_xlabel(
                'Energy [{0}]'.format(energy.unit.to_string('latex_inline')))


        paraxes = []
        for n in range(npars):
            paraxes.append(plt.subplot2grid((2*npars,10),(npars+n,0),colspan=7))
        self.parsliders = []
        slider_props = {'facecolor': color_cycle[-1], 'alpha': 0.5}
        for label,parax,valinit in six.moves.zip(labels, paraxes, p0):
            # Attempt to estimate reasonable parameter ranges from label
            pmin, pmax = valinit/10, valinit*3
            if 'log' in label:
                span = 2
                if 'norm' in label or 'amplitude' in label:
                    # give more range for normalization parameters
                    span *= 2
                pmin, pmax = valinit - span, valinit + span
            elif ('index' in label) or ('alpha' in label):
                if valinit > 0.:
                    pmin, pmax = 0, 5
                else:
                    pmin, pmax = -5, 0
            elif 'norm' in label or 'amplitude' in label:
                # norm without log, it will not be pretty because sliders are
                # only linear
                pmin, pmax = valinit / 100, valinit * 100

{'fraction_non_alphanumeric': 0.10578512396694215, 'fraction_numerical': 0.013223140495867768, 'mean_word_length': 2.0311478980614246, 'pattern_counts': {'":': 0, '<': 5, '<?xml version=': 0, '>': 5, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '2603337', 'n_tokens_mistral': 4348, 'n_tokens_neox': 4121, 'n_words': 1149}