<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.table import Table
from astropy.io import ascii
from astropy import log
from astropy.extern import six
import warnings

HAS_PYYAML = True
try:
    import yaml
except ImportError:
    HAS_PYYAML = False

__all__ = ["save_diagnostic_plots", "save_results_table"]

def save_diagnostic_plots(outname, sampler, modelidxs=None, pdf=False, sed=None,
        blob_labels=None, **kwargs):
    """
    Generate diagnostic plots.

    - A plot for each of the chain parameters showing walker progression, final
      sample distribution and several statistical measures of this distribution:
      ``outname_chain_parN.png`` (see `naima.plot_chain`).
    - A corner plot of sample density in the two dimensional parameter space of
      all parameter pairs of the run, with the Maximum Likelihood parameter
      vector indicated in blue: ``outname_corner.png`` (see `triangle.corner`).
    - A plot for each of the models returned as blobs by the model function. The
      maximum likelihood model is shown, as well as the 1 and 3 sigma confidence
      level contours. The first model will be compared with observational data
      and residuals shown. ``outname_fit_modelN.png`` (see `naima.plot_fit` and
      `naima.plot_blob`).

    Parameters
    ----------
    outname : str
        Name to be used to save diagnostic plot files.

    sampler : `emcee.EnsembleSampler` instance
        Sampler instance from which chains, blobs and data are read.

    modelidxs : iterable of integers, optional
        Model numbers to be plotted. Default: All returned in sampler.blobs

    blob_labels : list of strings, optional
        Label for each of the outputs of the model. They will be used as title
        for the corresponding plot.

    pdf : bool, optional
        Whether to save plots to multipage pdf.
    """

    from .plot import plot_chain, plot_blob

    if pdf:
        from matplotlib import pyplot as plt
        plt.rc('pdf', fonttype=42)
        log.info('Saving diagnostic plots in file '
                '{0}_plots.pdf'.format(outname))
        from matplotlib.backends.backend_pdf import PdfPages
        outpdf = PdfPages('{0}_plots.pdf'.format(outname))

    # Chains

    for par, label in six.moves.zip(six.moves.range(sampler.chain.shape[-1]),
                                    sampler.labels):
        try:
            log.info('Plotting chain of parameter {0}...'.format(label))
            f = plot_chain(sampler, par, **kwargs)
            if pdf:
                f.savefig(outpdf, format='pdf')
            else:
                if 'log(' in label or 'log10(' in label:
                    label = label.split('(')[-1].split(')')[0]
                f.savefig('{0}_chain_{1}.png'.format(outname, label))
            del f
        except Exception as e:
            log.warning('plot_chain failed for paramter'
                    ' {0} ({1}): {2}'.format(label,par,e))

    # Corner plot

    try:
        from triangle import corner
        from .plot import find_ML

        log.info('Plotting corner plot...')

        ML, MLp, MLvar, model_ML = find_ML(sampler, 0)
        f = corner(sampler.flatchain, labels=sampler.labels,
                   truths=MLp, quantiles=[0.16, 0.5, 0.84],
                   verbose=False, **kwargs)
        if pdf:
            f.savefig(outpdf, format='pdf')
        else:
            f.savefig('{0}_corner.png'.format(outname))
        del f
    except ImportError:
        log.warning('triangle_plot not installed, corner plot not available')

    # Fit

    if modelidxs is None:
        nmodels = len(sampler.blobs[-1][0])
        modelidxs = list(six.moves.range(nmodels))

    if sed is None:
        sed = [None for idx in modelidxs]
    elif isinstance(sed, bool):
        sed = [sed for idx in modelidxs]

    if blob_labels is None:
        blob_labels = ['Model output {0}'.format(idx) for idx in modelidxs]
    elif len(modelidxs)==1 and isinstance(blob_labels, str):
        blob_labels = [blob_labels,]
    elif len(blob_labels) < len(modelidxs):
        # Add labels
        n = len(blob_labels)
        blob_labels += ['Model output {0}'.format(idx) for idx in modelidxs[n:]]

    for modelidx, plot_sed, label in six.moves.zip(modelidxs, sed, blob_labels):

        try:
            log.info('Plotting {0}...'.format(label))
            f = plot_blob(sampler, blobidx=modelidx, label=label,
                          sed=plot_sed, n_samples=100, **kwargs)
            if pdf:
                f.savefig(outpdf, format='pdf')
            else:
                f.savefig('{0}_model{1}.png'.format(outname, modelidx))
            del f
        except Exception as e:
            log.warning('plot_blob failed for model output {0}: {1}'.format(par,e))

    if pdf:
        outpdf.close()


def save_results_table(outname, sampler, format='ascii.ecsv',
        convert_log=True, last_step=True, include_blobs=False, **kwargs):
    """
    Save an ASCII table with the results stored in the `~emcee.EnsembleSampler`.

    The table contains the median, 16th and 84th percentile confidence region
    (~1sigma) for each parameter.

    Parameters
    ----------
    outname : str
        Root name to be used to save the table. ``_results.dat`` will be
        appended for the output filename.

    sampler : `emcee.EnsembleSampler` instance
        Sampler instance from which chains, blobs and data are read.

    format : str, optional
        Format of the saved table. Must be a format string accepted by
        `astropy.table.Table.write`, see the `astropy unified file read/write
        interface documentation
        <https://astropy.readthedocs.org/en/latest/io/unified.html>`_. Only the
        ``ascii.ecsv`` and ``ascii.ipac`` formats are able to preserve all the
        information stored in the ``run_info`` dictionary of the sampler.
        Defaults to ``ascii.ecsv`` if available (only in astropy > v1.0), else
        ``ascii.ipac``.

    convert_log : bool, optional
        Whether to convert natural or base-10 logarithms into original values in
        addition to saving the logarithm value.

    last_step : bool, optional
        Whether to only use the positions in the final step of the run (True,<commit_msg>really make include_blobs True by default<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.table import Table
from astropy.io import ascii
from astropy import log
from astropy.extern import six
import warnings

HAS_PYYAML = True
try:
    import yaml
except ImportError:
    HAS_PYYAML = False

__all__ = ["save_diagnostic_plots", "save_results_table"]

def save_diagnostic_plots(outname, sampler, modelidxs=None, pdf=False, sed=None,
        blob_labels=None, **kwargs):
    """
    Generate diagnostic plots.

    - A plot for each of the chain parameters showing walker progression, final
      sample distribution and several statistical measures of this distribution:
      ``outname_chain_parN.png`` (see `naima.plot_chain`).
    - A corner plot of sample density in the two dimensional parameter space of
      all parameter pairs of the run, with the Maximum Likelihood parameter
      vector indicated in blue: ``outname_corner.png`` (see `triangle.corner`).
    - A plot for each of the models returned as blobs by the model function. The
      maximum likelihood model is shown, as well as the 1 and 3 sigma confidence
      level contours. The first model will be compared with observational data
      and residuals shown. ``outname_fit_modelN.png`` (see `naima.plot_fit` and
      `naima.plot_blob`).

    Parameters
    ----------
    outname : str
        Name to be used to save diagnostic plot files.

    sampler : `emcee.EnsembleSampler` instance
        Sampler instance from which chains, blobs and data are read.

    modelidxs : iterable of integers, optional
        Model numbers to be plotted. Default: All returned in sampler.blobs

    blob_labels : list of strings, optional
        Label for each of the outputs of the model. They will be used as title
        for the corresponding plot.

    pdf : bool, optional
        Whether to save plots to multipage pdf.
    """

    from .plot import plot_chain, plot_blob

    if pdf:
        from matplotlib import pyplot as plt
        plt.rc('pdf', fonttype=42)
        log.info('Saving diagnostic plots in file '
                '{0}_plots.pdf'.format(outname))
        from matplotlib.backends.backend_pdf import PdfPages
        outpdf = PdfPages('{0}_plots.pdf'.format(outname))

    # Chains

    for par, label in six.moves.zip(six.moves.range(sampler.chain.shape[-1]),
                                    sampler.labels):
        try:
            log.info('Plotting chain of parameter {0}...'.format(label))
            f = plot_chain(sampler, par, **kwargs)
            if pdf:
                f.savefig(outpdf, format='pdf')
            else:
                if 'log(' in label or 'log10(' in label:
                    label = label.split('(')[-1].split(')')[0]
                f.savefig('{0}_chain_{1}.png'.format(outname, label))
            del f
        except Exception as e:
            log.warning('plot_chain failed for paramter'
                    ' {0} ({1}): {2}'.format(label,par,e))

    # Corner plot

    try:
        from triangle import corner
        from .plot import find_ML

        log.info('Plotting corner plot...')

        ML, MLp, MLvar, model_ML = find_ML(sampler, 0)
        f = corner(sampler.flatchain, labels=sampler.labels,
                   truths=MLp, quantiles=[0.16, 0.5, 0.84],
                   verbose=False, **kwargs)
        if pdf:
            f.savefig(outpdf, format='pdf')
        else:
            f.savefig('{0}_corner.png'.format(outname))
        del f
    except ImportError:
        log.warning('triangle_plot not installed, corner plot not available')

    # Fit

    if modelidxs is None:
        nmodels = len(sampler.blobs[-1][0])
        modelidxs = list(six.moves.range(nmodels))

    if sed is None:
        sed = [None for idx in modelidxs]
    elif isinstance(sed, bool):
        sed = [sed for idx in modelidxs]

    if blob_labels is None:
        blob_labels = ['Model output {0}'.format(idx) for idx in modelidxs]
    elif len(modelidxs)==1 and isinstance(blob_labels, str):
        blob_labels = [blob_labels,]
    elif len(blob_labels) < len(modelidxs):
        # Add labels
        n = len(blob_labels)
        blob_labels += ['Model output {0}'.format(idx) for idx in modelidxs[n:]]

    for modelidx, plot_sed, label in six.moves.zip(modelidxs, sed, blob_labels):

        try:
            log.info('Plotting {0}...'.format(label))
            f = plot_blob(sampler, blobidx=modelidx, label=label,
                          sed=plot_sed, n_samples=100, **kwargs)
            if pdf:
                f.savefig(outpdf, format='pdf')
            else:
                f.savefig('{0}_model{1}.png'.format(outname, modelidx))
            del f
        except Exception as e:
            log.warning('plot_blob failed for model output {0}: {1}'.format(par,e))

    if pdf:
        outpdf.close()


def save_results_table(outname, sampler, format='ascii.ecsv',
        convert_log=True, last_step=True, include_blobs=True, **kwargs):
    """
    Save an ASCII table with the results stored in the `~emcee.EnsembleSampler`.

    The table contains the median, 16th and 84th percentile confidence region
    (~1sigma) for each parameter.

    Parameters
    ----------
    outname : str
        Root name to be used to save the table. ``_results.dat`` will be
        appended for the output filename.

    sampler : `emcee.EnsembleSampler` instance
        Sampler instance from which chains, blobs and data are read.

    format : str, optional
        Format of the saved table. Must be a format string accepted by
        `astropy.table.Table.write`, see the `astropy unified file read/write
        interface documentation
        <https://astropy.readthedocs.org/en/latest/io/unified.html>`_. Only the
        ``ascii.ecsv`` and ``ascii.ipac`` formats are able to preserve all the
        information stored in the ``run_info`` dictionary of the sampler.
        Defaults to ``ascii.ecsv`` if available (only in astropy > v1.0), else
        ``ascii.ipac``.

    convert_log : bool, optional
        Whether to convert natural or base-10 logarithms into original values in
        addition to saving the logarithm value.

    last_step : bool, optional
        Whether to only use the positions in the final step of the run (True,
{'fraction_non_alphanumeric': 0.09490687070098153, 'fraction_numerical': 0.007728572532653219, 'mean_word_length': 2.403471856917412, 'pattern_counts': {'":': 0, '<': 7, '<?xml version=': 0, '>': 7, 'https://': 2, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 2, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '5585595', 'n_tokens_mistral': 3856, 'n_tokens_neox': 3652, 'n_words': 1258}