<commit_before># Import smorgasbord
from __future__ import print_function
import pdb
import sys
import os
import copy
import dill
import re
import gc
import warnings
warnings.filterwarnings('ignore')
from difflib import SequenceMatcher
import multiprocessing as mp
import numpy as np
import scipy.stats
import scipy.interpolate
import scipy.optimize
import scipy.ndimage
import pandas as pd
import matplotlib
#matplotlib.use('agg')
import matplotlib.pyplot as plt
import matplotlib.ticker
import seaborn as sns
import sklearn.neighbors
import progress.bar
import termcolor
import acor
import corner
import emcee

# Disable interactive plotting
plt.ioff()

# Define physical constants
c = 3E8
h = 6.64E-34
k = 1.38E-23





def Fit(gal_dict,
        bands_frame,
        beta_vary = True,
        beta = 2.0,
        components = 2,
        kappa_0 = 0.051,
        kappa_0_lambda = 500E-6,
        plot = True,
        correl_unc = None,
        priors = None,
        mcmc_n_walkers = 250,
        mcmc_n_steps = 2500,
        mcmc_n_threads = int(round(mp.cpu_count()*1.0)),
        simple_clean = False,
        full_posterior = True,
        danger = False,
        verbose = True,
        test = False):
        """
        Function that runs the ChrisFit dust SED fitting routine.

        Arguments:
            gal_dict:
                    A dictionary, containing entries called 'name', 'distance', and 'redshift', giving the
                    values for the target source in question (with distance in units of parsecs)
            bands_frame:
                    A dataframe, with columns called 'band', 'wavelength, 'flux', 'error', and 'limit' providing the
                    relevant values for each band for the target source in question

        Keyword arguments:
            beta_vary:
                    A boolean, stating whether or not beta (the emissivity slope) should a free parameter,
                    or fixed
            beta:
                    A float, or list of floats, stating the value(s) of beta to use. If only a single float
                    is given, this is used for all components. If beta_vary is set to true, beta will
                    provide starting position for MCMC
            components:
                    An integer, stating how many modified blackbody components should make up the model
                    being fit
            kappa_0:
                    The value of the dust mass absorption coefficient, kappa_d, to use to calculate dust mass
                    (uses Clark et al., 2016, value by default)
            kappa_0_lambda:
                    The reference wavelength for kappa_0; corresponding value of kappa_0 at other
                    wavelengths extrapolated via (kappa_0_lambda/lambda)**beta
            plot:
                    A boolean, stating whether to generate plots of the resulting SED fit; or,
                    alternatively, a string pointing to the desired plotting output directory
            correl_unc:
                    A list, each element of which (if any) is a dictionary describing band-covariant
                    uncertainties; for the 5% Hershcel-SPIRE band covariance, correl_unc would be:
                    [{'correl_bands':['SPIRE_250','SPIRE_350','SPIRE_500'],
                    'correl_scale':0.04,
                    'correl_distr':'flat'}],
                    where 'bands' describes the bands (as named in bands_frame) in question, 'correl_scale'
                    describes the size of the covariant component of the flux uncertainty (as a fraction of
                    measured source flux), and 'correl_distr' is the assumed distribution of the uncertainty
                    (currently accepting either 'flat', 'normal', or a defined function)
            priors:
                    A dictionary, of lists, of functions (yeah, I know); dictionary entries can be called
                    'temp', 'mass', and 'beta', each entry being an n-length list, where n is the number of
                    components, with the n-th list element being a function giving the ln-like prior for the
                    parameter in question (ie, temperature, mass, or beta) of the n-th model component; note
                    that the priors for any correlated uncertainty terms should be provided through the
                    correl_unc kwarg instead
            mcmc_n_walkers:
                    An int, stating how many emcee MCMC walkers should be used to explore the posterior; the higher the
                    number, the less likely you are to find yourself stuck in a local minima
            mcmc_n_steps:
                    An integer, stating how many steps each emcee MCMC walker should take whilst exploring the
                    posterior; the more steps, the better the likelihood of convergence and good sampling
            mcmc_n_threads:
                    An integer, stating how many CPU threads emcee should use for its MCMC sampling; if you have a
                    small number of walkers, setting this to 1 might be faster
            simple_clean:
                    A boolean or float, descrbing the type of chain cleaning to be performed. If False, a full cleaning
                    will be performed, with convergance diagnostics to remove burn-in, metastability analysis to exclude
                    bad chains. If a float in the range 0-1 is given, then all that is done is that fraction of the
                    all is removed as burn-in; if you have enough of walkers (ie, hundreds) and steps (ie, thousands),
                    then using this simple option and removing the first 50% all chains should
            full_posterior:
                    A boolean, stating whether the full posterior distribution of each parameter should be
                    returned, or just the summary of median, credible interval, etc
            danger:
                    A boolean, which if True will prioritise speed over caution (enabling the emcee live_dangerously
                    kwarg, and scatting the initial positions of the walkers a bit less)
            verbose:
                    A boolean, stating whether ChrisFit should provide verbose output whilst operating
            test:
                    A boolean, stating whether to read in a previously-saved posterior (only really needful for testing)
            """


        # State name of source being processed
        if sys.stdout.isatty():
            name_bracket_prefix = termcolor.colored('['+gal_dict['name']+']' + (' '*(25-len('['+gal_dict['name']+']'))), 'cyan', attrs=['bold'])
        else:
            name_bracket_prefix = '['+gal_dict['name']+']' + (' '*(25-len('['+gal_dict['name']+']')))
        if verbose:
            print(name_bracket_prefix  + 'Commencing processing')

        # Add column to bands_frame, to record which fluxes are larger than their uncertainty
        bands_frame['det'] = bands_frame.loc[:,'flux'] > bands_frame.loc[:,'error']

        # Parse beta argument, so that each model component is assigned its own value (even if they are all the same)
        if not hasattr(beta, '__iter__'):
            beta = np.array([beta])
        if len(beta) == 1 and components > 1:
            beta = np.array(beta)
        elif (len(beta) > 1) and (len(beta)<components) and (components > 1):
            Exception('Either provide a single value of beta, or a list of values of length the number of components')

        # Parse correl_unc argument, so that if no value provided, an empty list is used throughout the rest of the function
        if not hasattr(correl_unc, '__iter__'):
            correl_unc = []

        # Bundle various fitting arguments in to a dictionary
        fit_dict = {'bands_frame':bands_frame,
                    'gal_dict':gal_dict,
                    'components':components,
                    'beta_vary':beta_vary,
                    'beta':beta,
                    'correl_unc':correl_unc,
                    'bounds':False,
                    'priors':priors,
                    'mcmc_n_walkers':mcmc_n_walkers,
                    'mcmc_n_steps':mcmc_n_steps,
                    'distance':gal_dict['distance'],
                    'kappa_0':kappa_0,
                    'kappa_0_lambda':kappa_0_lambda,
                    'danger':danger}

        # Determine number of parameters
        n_params = (2 * int(components)) + (int(fit_dict['beta_vary']) * len(fit_dict['beta'])) + len(correl_unc)
        fit_dict['n_params'] = n_params

        # Read in colour-correction tables
        fit_dict['colour_corrections'] = ColourCorrections(fit_dict)

        # No custom priors provided, construct priors ahead of time, but warn that this is slower
        if isinstance(fit_dict['priors'], dict):
            pass
        else:
            fit_dict['priors'] = PriorsConstruct(fit_dict)
            if mcmc_n_threads > 1:
                print(name_bracket_prefix  + 'No custom priors provided; constructing default priors (note that the multithreaded MCMC is *MUCH FASTER* when working with custom priors, as functions passed from externally can be handled more efficiently)')


        # Generate initial guess values for maximum-likelihood estimation (which will then itself be used to initialise emcee's estimation)
        mle_fit_dict = copy.deepcopy(fit_dict)
        mle_fit_dict['bounds'] = True
        mle_fit_dict['correl_unc'] = False
        mle_initial = MaxLikeInitial(bands_frame, mle_fit_dict)



        # Find Maximum Likelihood Estimate (MLE)
        if test and os.path.exists(os.path.join(plot,gal_dict['name']+'_MCMC.dj')):
            mcmc_chains = dill.load(open(os.path.join(plot,gal_dict['name']+'_MCMC.dj'),'rb'))
        if not test:
            if verbose:
                print(name_bracket_prefix + 'Performing maximum likelihood estimation to initialise MCMC')
            NegLnLike = lambda *args: -LnLike(*args)
            mle_opt = scipy.optimize.minimize(NegLnLike, mle_initial, args=(mle_fit_dict), method='Powell', tol=5E-4, options={'maxiter':1000,'maxfev':1000})
            mle_params = mle_opt.x

            # Re-introduce any correlated uncertainty parameters that were excluded from maximum-likelihood fit<commit_msg>Neatened up verbose formatting<commit_after># Import smorgasbord
from __future__ import print_function
import pdb
import sys
import os
import copy
import dill
import re
import gc
import warnings
warnings.filterwarnings('ignore')
from difflib import SequenceMatcher
import multiprocessing as mp
import numpy as np
import scipy.stats
import scipy.interpolate
import scipy.optimize
import scipy.ndimage
import pandas as pd
import matplotlib
#matplotlib.use('agg')
import matplotlib.pyplot as plt
import matplotlib.ticker
import seaborn as sns
import sklearn.neighbors
import progress.bar
import termcolor
import acor
import corner
import emcee

# Disable interactive plotting
plt.ioff()

# Define physical constants
c = 3E8
h = 6.64E-34
k = 1.38E-23





def Fit(gal_dict,
        bands_frame,
        beta_vary = True,
        beta = 2.0,
        components = 2,
        kappa_0 = 0.051,
        kappa_0_lambda = 500E-6,
        plot = True,
        correl_unc = None,
        priors = None,
        mcmc_n_walkers = 250,
        mcmc_n_steps = 2500,
        mcmc_n_threads = int(round(mp.cpu_count()*1.0)),
        simple_clean = False,
        full_posterior = True,
        danger = False,
        verbose = True,
        test = False):
        """
        Function that runs the ChrisFit dust SED fitting routine.

        Arguments:
            gal_dict:
                    A dictionary, containing entries called 'name', 'distance', and 'redshift', giving the
                    values for the target source in question (with distance in units of parsecs)
            bands_frame:
                    A dataframe, with columns called 'band', 'wavelength, 'flux', 'error', and 'limit' providing the
                    relevant values for each band for the target source in question

        Keyword arguments:
            beta_vary:
                    A boolean, stating whether or not beta (the emissivity slope) should a free parameter,
                    or fixed
            beta:
                    A float, or list of floats, stating the value(s) of beta to use. If only a single float
                    is given, this is used for all components. If beta_vary is set to true, beta will
                    provide starting position for MCMC
            components:
                    An integer, stating how many modified blackbody components should make up the model
                    being fit
            kappa_0:
                    The value of the dust mass absorption coefficient, kappa_d, to use to calculate dust mass
                    (uses Clark et al., 2016, value by default)
            kappa_0_lambda:
                    The reference wavelength for kappa_0; corresponding value of kappa_0 at other
                    wavelengths extrapolated via (kappa_0_lambda/lambda)**beta
            plot:
                    A boolean, stating whether to generate plots of the resulting SED fit; or,
                    alternatively, a string pointing to the desired plotting output directory
            correl_unc:
                    A list, each element of which (if any) is a dictionary describing band-covariant
                    uncertainties; for the 5% Hershcel-SPIRE band covariance, correl_unc would be:
                    [{'correl_bands':['SPIRE_250','SPIRE_350','SPIRE_500'],
                    'correl_scale':0.04,
                    'correl_distr':'flat'}],
                    where 'bands' describes the bands (as named in bands_frame) in question, 'correl_scale'
                    describes the size of the covariant component of the flux uncertainty (as a fraction of
                    measured source flux), and 'correl_distr' is the assumed distribution of the uncertainty
                    (currently accepting either 'flat', 'normal', or a defined function)
            priors:
                    A dictionary, of lists, of functions (yeah, I know); dictionary entries can be called
                    'temp', 'mass', and 'beta', each entry being an n-length list, where n is the number of
                    components, with the n-th list element being a function giving the ln-like prior for the
                    parameter in question (ie, temperature, mass, or beta) of the n-th model component; note
                    that the priors for any correlated uncertainty terms should be provided through the
                    correl_unc kwarg instead
            mcmc_n_walkers:
                    An int, stating how many emcee MCMC walkers should be used to explore the posterior; the higher the
                    number, the less likely you are to find yourself stuck in a local minima
            mcmc_n_steps:
                    An integer, stating how many steps each emcee MCMC walker should take whilst exploring the
                    posterior; the more steps, the better the likelihood of convergence and good sampling
            mcmc_n_threads:
                    An integer, stating how many CPU threads emcee should use for its MCMC sampling; if you have a
                    small number of walkers, setting this to 1 might be faster
            simple_clean:
                    A boolean or float, descrbing the type of chain cleaning to be performed. If False, a full cleaning
                    will be performed, with convergance diagnostics to remove burn-in, metastability analysis to exclude
                    bad chains. If a float in the range 0-1 is given, then all that is done is that fraction of the
                    all is removed as burn-in; if you have enough of walkers (ie, hundreds) and steps (ie, thousands),
                    then using this simple option and removing the first 50% all chains should
            full_posterior:
                    A boolean, stating whether the full posterior distribution of each parameter should be
                    returned, or just the summary of median, credible interval, etc
            danger:
                    A boolean, which if True will prioritise speed over caution (enabling the emcee live_dangerously
                    kwarg, and scatting the initial positions of the walkers a bit less)
            verbose:
                    A boolean, stating whether ChrisFit should provide verbose output whilst operating
            test:
                    A boolean, stating whether to read in a previously-saved posterior (only really needful for testing)
            """


        # State name of source being processed
        if sys.stdout.isatty():
            name_bracket_prefix = termcolor.colored('['+gal_dict['name']+']' + (' '*(23-len('['+gal_dict['name']+']'))), 'cyan', attrs=['bold'])
        else:
            name_bracket_prefix = '['+gal_dict['name']+']' + (' '*(25-len('['+gal_dict['name']+']')))
        if verbose:
            print(name_bracket_prefix  + 'Commencing processing')

        # Add column to bands_frame, to record which fluxes are larger than their uncertainty
        bands_frame['det'] = bands_frame.loc[:,'flux'] > bands_frame.loc[:,'error']

        # Parse beta argument, so that each model component is assigned its own value (even if they are all the same)
        if not hasattr(beta, '__iter__'):
            beta = np.array([beta])
        if len(beta) == 1 and components > 1:
            beta = np.array(beta)
        elif (len(beta) > 1) and (len(beta)<components) and (components > 1):
            Exception('Either provide a single value of beta, or a list of values of length the number of components')

        # Parse correl_unc argument, so that if no value provided, an empty list is used throughout the rest of the function
        if not hasattr(correl_unc, '__iter__'):
            correl_unc = []

        # Bundle various fitting arguments in to a dictionary
        fit_dict = {'bands_frame':bands_frame,
                    'gal_dict':gal_dict,
                    'components':components,
                    'beta_vary':beta_vary,
                    'beta':beta,
                    'correl_unc':correl_unc,
                    'bounds':False,
                    'priors':priors,
                    'mcmc_n_walkers':mcmc_n_walkers,
                    'mcmc_n_steps':mcmc_n_steps,
                    'distance':gal_dict['distance'],
                    'kappa_0':kappa_0,
                    'kappa_0_lambda':kappa_0_lambda,
                    'danger':danger}

        # Determine number of parameters
        n_params = (2 * int(components)) + (int(fit_dict['beta_vary']) * len(fit_dict['beta'])) + len(correl_unc)
        fit_dict['n_params'] = n_params

        # Read in colour-correction tables
        fit_dict['colour_corrections'] = ColourCorrections(fit_dict)

        # No custom priors provided, construct priors ahead of time, but warn that this is slower
        if isinstance(fit_dict['priors'], dict):
            pass
        elif mcmc_n_threads == 1:
            fit_dict['priors'] = PriorsConstruct(fit_dict)
        else:
            print(name_bracket_prefix  + 'No custom priors provided; constructing default priors (note that the multithreaded MCMC is *MUCH FASTER* when working with custom priors, as functions defined outsite the fitter can be handled more efficiently)')


        # Generate initial guess values for maximum-likelihood estimation (which will then itself be used to initialise emcee's estimation)
        mle_fit_dict = copy.deepcopy(fit_dict)
        mle_fit_dict['bounds'] = True
        mle_fit_dict['correl_unc'] = False
        mle_initial = MaxLikeInitial(bands_frame, mle_fit_dict)



        # Find Maximum Likelihood Estimate (MLE)
        if test and os.path.exists(os.path.join(plot,gal_dict['name']+'_MCMC.dj')):
            mcmc_chains = dill.load(open(os.path.join(plot,gal_dict['name']+'_MCMC.dj'),'rb'))
        if not test:
            if verbose:
                print(name_bracket_prefix + 'Performing maximum likelihood estimation to initialise MCMC')
            NegLnLike = lambda *args: -LnLike(*args)
            mle_opt = scipy.optimize.minimize(NegLnLike, mle_initial, args=(mle_fit_dict), method='Powell', tol=5E-4, options={'maxiter':1000,'maxfev':1000})
            mle_params = mle_opt.x

            # Re-introduce any correlated uncertainty parameters that were excluded from maximum-likelihood fit
{'fraction_non_alphanumeric': 0.06948215671062839, 'fraction_numerical': 0.008242823894491854, 'mean_word_length': 2.1445342277786246, 'pattern_counts': {'":': 0, '<': 5, '<?xml version=': 0, '>': 12, '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': '592110', 'n_tokens_mistral': 5575, 'n_tokens_neox': 5164, 'n_words': 2041}