<commit_before>from __future__ import print_function, division
import os,os.path
import numpy as np
import logging

try:
    import pandas as pd
except ImportError:
    pd = None
    
import numpy.random as rand
import scipy.optimize

try:
    import emcee
except ImportError:
    emcee = None

try:
    from plotutils.plotutils import setfig
except ImportError:
    setfig = None
    
import matplotlib.pyplot as plt

try:
    import triangle
except ImportError:
    triangle = None


from .extinction import EXTINCTION

class StarModel(object):
    """An object to represent a star, with observed properties, modeled by an Isochrone

    :param ic: 
        :class:`Isochrone` object used to model star.

    :param maxAV: (optional)
        Maximum allowed extinction (i.e. the extinction @ infinity in direction of star).  Default is 1.

    :param max_distance: (optional)
        Maximum allowed distance (pc).  Default is 1000.
    
    :param **kwargs:
        Keyword arguments must be properties of given isochrone, e.g., logg,
        feh, Teff, and/or magnitudes.  The values represent measurements of
        the star, and must be in (value,error) format. All such keyword
        arguments will be held in ``self.properties``.
        
    """
    def __init__(self,ic,maxAV=1,max_distance=1000,**kwargs):
        self.ic = ic
        self.properties = kwargs
        self.max_distance = max_distance
        self.maxAV = maxAV
        self._samples = None

    def add_props(self,**kwargs):
        """
        Adds observable properties to ``self.properties``.
        
        """
        for kw,val in kwargs.iteritems():
            self.properties[kw] = val

    def remove_props(self,*args):
        """
        Removes desired properties from ``self.properties``.
        
        """
        for arg in args:
            if arg in self.properties:
                del self.properties[arg]
    
    @property
    def fit_for_distance(self):
        """
        ``True`` if any of the properties are apparent magnitudes.
        
        """
        for prop in self.properties.keys():
            if prop in self.ic.bands:
                return True
        return False
            
    
    def loglike(self,p):
        """Log-likelihood of model at given parameters

        
        :param p: 
            mass, log10(age), feh, [distance, A_V (extinction)].
            Final two should only be provided if ``self.fit_for_distance``
            is ``True``; that is, apparent magnitudes are provided.
            

        :return:
           log-likelihood.  Will be -np.inf if values out of range.
        
        """
        if len(p)==5:
            fit_for_distance = True
            mass,age,feh,dist,AV = p
        elif len(p)==3:
            fit_for_distance = False
            mass,age,feh = p
                        
        if mass < self.ic.minmass or mass > self.ic.maxmass \
           or age < self.ic.minage or age > self.ic.maxage \
           or feh < self.ic.minfeh or feh > self.ic.maxfeh:
            return -np.inf
        if self.fit_for_distance:
            if dist < 0 or AV < 0 or dist > self.max_distance:
                return -np.inf
            if AV > self.maxAV:
                return -np.inf

        logl = 0
        for prop in self.properties.keys():
            val,err = self.properties[prop]
            if prop in self.ic.bands:
                if not fit_for_distance:
                    raise ValueError('must fit for mass, age, feh, dist, A_V if apparent magnitudes provided.')
                mod = self.ic.mag[prop](mass,age,feh) + 5*np.log10(dist) - 5
                A = AV*EXTINCTION[prop]
                mod += A
            elif prop=='feh':
                mod = feh
            else:
                mod = getattr(self.ic,prop)(mass,age,feh)
            logl += -(val-mod)**2/err**2

        if np.isnan(logl):
            logl = -np.inf

        #IMF prior
        logl += np.log(salpeter_prior(mass))
        
        #distance prior?
        

        ##prior to sample ages with linear prior
        #a0 = 10**self.ic.minage
        #a1 = 10**self.ic.maxage
        #da = a1-a0
        #a = 10**age
        #logl += np.log(a/(a1-a0))

        return logl

    def maxlike(self,nseeds=50):
        """Returns the best-fit parameters, choosing the best of multiple starting guesses

        :param nseeds: (optional)
            Number of starting guesses, uniformly distributed throughout
            allowed ranges.  Default=50.

        :return:
            list of best-fit parameters: ``[m,age,feh,[distance,A_V]]``.
            Note that distance and A_V values will be meaningless unless
            magnitudes are present in ``self.properties``.
        
        """
        m0,age0,feh0 = self.ic.random_points(nseeds)
        d0 = 10**(rand.uniform(0,np.log10(self.max_distance),size=nseeds))
        AV0 = rand.uniform(0,self.maxAV,size=nseeds)

        

        costs = np.zeros(nseeds)

        if self.fit_for_distance:
            pfits = np.zeros((nseeds,5))
        else:
            pfits = np.zeros((nseeds,3))
            
        def fn(p): #fmin is a function *minimizer*
            return -1*self.loglike(p)
        
        for i,m,age,feh,d,AV in zip(range(nseeds),
                                    m0,age0,feh0,d0,AV0):
                if self.fit_for_distance:
                    pfit = scipy.optimize.fmin(fn,[m,age,feh,d,AV],disp=False)
                else:
                    pfit = scipy.optimize.fmin(fn,[m,age,feh],disp=False)
                pfits[i,:] = pfit
                costs[i] = self.loglike(pfit)

        return pfits[np.argmax(costs),:]

            
    def fit_mcmc(self,nwalkers=200,nburn=100,niter=200,
                 p0=None,initial_burn=None,
                 ninitial=100, **kwargs):
        """Fits stellar model using MCMC.

        :param nwalkers: (optional)
            Number of walkers to pass to :class:`emcee.EnsembleSampler`.
            Default is 200.

        :param nburn: (optional)
            Number of iterations for "burn-in."  Default is 100.

        :param niter: (optional)
            Number of for-keeps iterations for MCMC chain.
            Default is 200.

        :param p0: (optional)
            Initial parameters for emcee.  If not provided, then chains
            will behave according to whether inital_burn is set.

        :param initial_burn: (optional)
            If `True`, then initialize walkers first with a random initialization,
            then cull the walkers, keeping only those with > 15% acceptance
            rate, then reinitialize sampling.  If `False`, then just do
            normal burn-in.  Default is `None`, which will be set to `True` if
            fitting for distance (i.e., if there are apparent magnitudes as
            properties of the model), and `False` if not.
            
        :param ninitial: (optional)
            Number of iterations to test walkers for acceptance rate before
            re-initializing.
        
        :return:
            :class:`emcee.EnsembleSampler` object.
            
        """

        #clear any saved _samples
        if self._samples is not None:
            self._samples = None
            

        if self.fit_for_distance:
            npars = 5
            if initial_burn is None:
                initial_burn = True
        else:
            if initial_burn is None:
                initial_burn = False
            npars = 3

        if p0 is None:
            m0,age0,feh0 = self.ic.random_points(nwalkers)
            d0 = 10**(rand.uniform(0,np.log10(self.max_distance),size=nwalkers))
            AV0 = rand.uniform(0,self.maxAV,size=nwalkers)
            if self.fit_for_distance:
                p0 = np.array([m0,age0,feh0,d0,AV0]).T
            else:
                p0 = np.array([m0,age0,feh0]).T
            if initial_burn:
                sampler = emcee.EnsembleSampler(nwalkers,npars,self.loglike,
                                                **kwargs)
                #ninitial = 300 #should this be parameter?
                pos, prob, state = sampler.run_mcmc(p0, ninitial) 
                wokinds = np.where((sampler.naccepted/ninitial > 0.15) &
                                   (sampler.naccepted/ninitial < 0.4))[0]
                i=1
                while len(wokinds)==0:
                    thresh = 0.15 - i*0.02
                    if thresh < 0:
                        raise RuntimeError('Initial burn has no acceptance?')
                    wokinds = np.where((sampler.naccepted/ninitial > thresh) &
                                       (sampler.naccepted/ninitial < 0.4))[0]
                    i += 1
                inds = rand.randint(len(wokinds),size=nwalkers)
                p0 = sampler.chain[wokinds[inds],:,:].mean(axis=1) #reset p0
                p0 *= (1 + rand.normal(size=p0.shape)*0.01)
        else:
            p0 = np.array(p0)
            p0 = rand.normal(size=(nwalkers,npars))*0.01 + p0.T[None,:]
            if self.fit_for_distance:
                p0[:,3] *= (1 + rand.normal(size=nwalkers)*0.5)
        
        sampler = emcee.EnsembleSampler(nwalkers,npars,self.loglike)
        pos, prob, state = sampler.run_mcmc(p0, nburn)
        sampler.reset()
        sampler.run_mcmc(pos, niter, rstate0=state)
        
        self._sampler = sampler
        return sampler

    def triangle_plots(self, basename=None, format='png',
                       **kwargs):
        """Returns two triangle plots, one with physical params, one observational

        :param basename:
            If basename is provided, then plots will be saved as
            "[basename]_physical.[format]" and "[basename]_observed.[format]"

        :param format:
            Format in which to save figures (e.g., 'png' or 'pdf')

        :param **kwargs:
            Additional keyword arguments passed to :func:`StarModel.triangle`
            and :func:`StarModel.prop_triangle`

        :return:
             * Physical parameters triangle plot (mass, radius, Teff, feh, age)
             * Observed properties triangle plot.
             
        """
        fig1 = self.triangle(plot_datapoints=False,
                            params=['mass','radius','Teff','feh','age'],
                            **kwargs)
        if basename is not None:
            plt.savefig('{}_physical.{}'.format(basename,format))
            plt.close()
        fig2 = self.prop_triangle(**kwargs)
        if basename is not None:
            plt.savefig('{}_observed.{}'.format(basename,format))
            plt.close()
        return fig1, fig2

    def triangle(self, params=None, query=None, extent=0.999,
                 **kwargs):
        """
        Makes a nifty corner plot.

        Uses :func:`triangle.corner`.

        :param params: (optional)
            Names of columns (from :attr:`StarModel.samples`)
            to plot.  If ``None``, then it will plot samples
            of the parameters used in the MCMC fit-- that is,
            mass, age, [Fe/H], and optionally distance and A_V.

        :param query: (optional)<commit_msg>added distance to physical triangle plot<commit_after>from __future__ import print_function, division
import os,os.path
import numpy as np
import logging

try:
    import pandas as pd
except ImportError:
    pd = None
    
import numpy.random as rand
import scipy.optimize

try:
    import emcee
except ImportError:
    emcee = None

try:
    from plotutils.plotutils import setfig
except ImportError:
    setfig = None
    
import matplotlib.pyplot as plt

try:
    import triangle
except ImportError:
    triangle = None


from .extinction import EXTINCTION

class StarModel(object):
    """An object to represent a star, with observed properties, modeled by an Isochrone

    :param ic: 
        :class:`Isochrone` object used to model star.

    :param maxAV: (optional)
        Maximum allowed extinction (i.e. the extinction @ infinity in direction of star).  Default is 1.

    :param max_distance: (optional)
        Maximum allowed distance (pc).  Default is 1000.
    
    :param **kwargs:
        Keyword arguments must be properties of given isochrone, e.g., logg,
        feh, Teff, and/or magnitudes.  The values represent measurements of
        the star, and must be in (value,error) format. All such keyword
        arguments will be held in ``self.properties``.
        
    """
    def __init__(self,ic,maxAV=1,max_distance=1000,**kwargs):
        self.ic = ic
        self.properties = kwargs
        self.max_distance = max_distance
        self.maxAV = maxAV
        self._samples = None

    def add_props(self,**kwargs):
        """
        Adds observable properties to ``self.properties``.
        
        """
        for kw,val in kwargs.iteritems():
            self.properties[kw] = val

    def remove_props(self,*args):
        """
        Removes desired properties from ``self.properties``.
        
        """
        for arg in args:
            if arg in self.properties:
                del self.properties[arg]
    
    @property
    def fit_for_distance(self):
        """
        ``True`` if any of the properties are apparent magnitudes.
        
        """
        for prop in self.properties.keys():
            if prop in self.ic.bands:
                return True
        return False
            
    
    def loglike(self,p):
        """Log-likelihood of model at given parameters

        
        :param p: 
            mass, log10(age), feh, [distance, A_V (extinction)].
            Final two should only be provided if ``self.fit_for_distance``
            is ``True``; that is, apparent magnitudes are provided.
            

        :return:
           log-likelihood.  Will be -np.inf if values out of range.
        
        """
        if len(p)==5:
            fit_for_distance = True
            mass,age,feh,dist,AV = p
        elif len(p)==3:
            fit_for_distance = False
            mass,age,feh = p
                        
        if mass < self.ic.minmass or mass > self.ic.maxmass \
           or age < self.ic.minage or age > self.ic.maxage \
           or feh < self.ic.minfeh or feh > self.ic.maxfeh:
            return -np.inf
        if self.fit_for_distance:
            if dist < 0 or AV < 0 or dist > self.max_distance:
                return -np.inf
            if AV > self.maxAV:
                return -np.inf

        logl = 0
        for prop in self.properties.keys():
            val,err = self.properties[prop]
            if prop in self.ic.bands:
                if not fit_for_distance:
                    raise ValueError('must fit for mass, age, feh, dist, A_V if apparent magnitudes provided.')
                mod = self.ic.mag[prop](mass,age,feh) + 5*np.log10(dist) - 5
                A = AV*EXTINCTION[prop]
                mod += A
            elif prop=='feh':
                mod = feh
            else:
                mod = getattr(self.ic,prop)(mass,age,feh)
            logl += -(val-mod)**2/err**2

        if np.isnan(logl):
            logl = -np.inf

        #IMF prior
        logl += np.log(salpeter_prior(mass))
        
        #distance prior?
        

        ##prior to sample ages with linear prior
        #a0 = 10**self.ic.minage
        #a1 = 10**self.ic.maxage
        #da = a1-a0
        #a = 10**age
        #logl += np.log(a/(a1-a0))

        return logl

    def maxlike(self,nseeds=50):
        """Returns the best-fit parameters, choosing the best of multiple starting guesses

        :param nseeds: (optional)
            Number of starting guesses, uniformly distributed throughout
            allowed ranges.  Default=50.

        :return:
            list of best-fit parameters: ``[m,age,feh,[distance,A_V]]``.
            Note that distance and A_V values will be meaningless unless
            magnitudes are present in ``self.properties``.
        
        """
        m0,age0,feh0 = self.ic.random_points(nseeds)
        d0 = 10**(rand.uniform(0,np.log10(self.max_distance),size=nseeds))
        AV0 = rand.uniform(0,self.maxAV,size=nseeds)

        

        costs = np.zeros(nseeds)

        if self.fit_for_distance:
            pfits = np.zeros((nseeds,5))
        else:
            pfits = np.zeros((nseeds,3))
            
        def fn(p): #fmin is a function *minimizer*
            return -1*self.loglike(p)
        
        for i,m,age,feh,d,AV in zip(range(nseeds),
                                    m0,age0,feh0,d0,AV0):
                if self.fit_for_distance:
                    pfit = scipy.optimize.fmin(fn,[m,age,feh,d,AV],disp=False)
                else:
                    pfit = scipy.optimize.fmin(fn,[m,age,feh],disp=False)
                pfits[i,:] = pfit
                costs[i] = self.loglike(pfit)

        return pfits[np.argmax(costs),:]

            
    def fit_mcmc(self,nwalkers=200,nburn=100,niter=200,
                 p0=None,initial_burn=None,
                 ninitial=100, **kwargs):
        """Fits stellar model using MCMC.

        :param nwalkers: (optional)
            Number of walkers to pass to :class:`emcee.EnsembleSampler`.
            Default is 200.

        :param nburn: (optional)
            Number of iterations for "burn-in."  Default is 100.

        :param niter: (optional)
            Number of for-keeps iterations for MCMC chain.
            Default is 200.

        :param p0: (optional)
            Initial parameters for emcee.  If not provided, then chains
            will behave according to whether inital_burn is set.

        :param initial_burn: (optional)
            If `True`, then initialize walkers first with a random initialization,
            then cull the walkers, keeping only those with > 15% acceptance
            rate, then reinitialize sampling.  If `False`, then just do
            normal burn-in.  Default is `None`, which will be set to `True` if
            fitting for distance (i.e., if there are apparent magnitudes as
            properties of the model), and `False` if not.
            
        :param ninitial: (optional)
            Number of iterations to test walkers for acceptance rate before
            re-initializing.
        
        :return:
            :class:`emcee.EnsembleSampler` object.
            
        """

        #clear any saved _samples
        if self._samples is not None:
            self._samples = None
            

        if self.fit_for_distance:
            npars = 5
            if initial_burn is None:
                initial_burn = True
        else:
            if initial_burn is None:
                initial_burn = False
            npars = 3

        if p0 is None:
            m0,age0,feh0 = self.ic.random_points(nwalkers)
            d0 = 10**(rand.uniform(0,np.log10(self.max_distance),size=nwalkers))
            AV0 = rand.uniform(0,self.maxAV,size=nwalkers)
            if self.fit_for_distance:
                p0 = np.array([m0,age0,feh0,d0,AV0]).T
            else:
                p0 = np.array([m0,age0,feh0]).T
            if initial_burn:
                sampler = emcee.EnsembleSampler(nwalkers,npars,self.loglike,
                                                **kwargs)
                #ninitial = 300 #should this be parameter?
                pos, prob, state = sampler.run_mcmc(p0, ninitial) 
                wokinds = np.where((sampler.naccepted/ninitial > 0.15) &
                                   (sampler.naccepted/ninitial < 0.4))[0]
                i=1
                while len(wokinds)==0:
                    thresh = 0.15 - i*0.02
                    if thresh < 0:
                        raise RuntimeError('Initial burn has no acceptance?')
                    wokinds = np.where((sampler.naccepted/ninitial > thresh) &
                                       (sampler.naccepted/ninitial < 0.4))[0]
                    i += 1
                inds = rand.randint(len(wokinds),size=nwalkers)
                p0 = sampler.chain[wokinds[inds],:,:].mean(axis=1) #reset p0
                p0 *= (1 + rand.normal(size=p0.shape)*0.01)
        else:
            p0 = np.array(p0)
            p0 = rand.normal(size=(nwalkers,npars))*0.01 + p0.T[None,:]
            if self.fit_for_distance:
                p0[:,3] *= (1 + rand.normal(size=nwalkers)*0.5)
        
        sampler = emcee.EnsembleSampler(nwalkers,npars,self.loglike)
        pos, prob, state = sampler.run_mcmc(p0, nburn)
        sampler.reset()
        sampler.run_mcmc(pos, niter, rstate0=state)
        
        self._sampler = sampler
        return sampler

    def triangle_plots(self, basename=None, format='png',
                       **kwargs):
        """Returns two triangle plots, one with physical params, one observational

        :param basename:
            If basename is provided, then plots will be saved as
            "[basename]_physical.[format]" and "[basename]_observed.[format]"

        :param format:
            Format in which to save figures (e.g., 'png' or 'pdf')

        :param **kwargs:
            Additional keyword arguments passed to :func:`StarModel.triangle`
            and :func:`StarModel.prop_triangle`

        :return:
             * Physical parameters triangle plot (mass, radius, Teff, feh, age, distance)
             * Observed properties triangle plot.
             
        """
        fig1 = self.triangle(plot_datapoints=False,
                            params=['mass','radius','Teff','feh','age','distance'],
                            **kwargs)
        if basename is not None:
            plt.savefig('{}_physical.{}'.format(basename,format))
            plt.close()
        fig2 = self.prop_triangle(**kwargs)
        if basename is not None:
            plt.savefig('{}_observed.{}'.format(basename,format))
            plt.close()
        return fig1, fig2

    def triangle(self, params=None, query=None, extent=0.999,
                 **kwargs):
        """
        Makes a nifty corner plot.

        Uses :func:`triangle.corner`.

        :param params: (optional)
            Names of columns (from :attr:`StarModel.samples`)
            to plot.  If ``None``, then it will plot samples
            of the parameters used in the MCMC fit-- that is,
            mass, age, [Fe/H], and optionally distance and A_V.

        :param query: (optional)
{'fraction_non_alphanumeric': 0.09958877167888432, 'fraction_numerical': 0.014392991239048811, 'mean_word_length': 1.6536591151702051, 'pattern_counts': {'":': 0, '<': 19, '<?xml version=': 0, '>': 19, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 10, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '2723451', 'n_tokens_mistral': 6834, 'n_tokens_neox': 6430, 'n_words': 1763}