<commit_before>from EXOSIMS.PlanetPopulation.KeplerLike2 import KeplerLike2
from EXOSIMS.util.InverseTransformSampler import InverseTransformSampler
import astropy.units as u
import astropy.constants as const
import numpy as np
import scipy.integrate as integrate

class SAG13(KeplerLike2):
    """Planet Population module based on SAG13 occurrence rates.
    
    This is the current working model based on averaging multiple studies. 
    These do not yet represent official scientific values.
    
    Attributes: 
        SAG13coeffs (float 4x2 ndarray):
            Coefficients used by the SAG13 broken power law. The 4 lines
            correspond to Gamma, alpha, beta, and the minimum radius.
        Gamma (float ndarray):
            Gamma coefficients used by SAG13 broken power law.
        alpha (float ndarray):
            Alpha coefficients used by SAG13 broken power law.
        beta (float ndarray):
            Beta coefficients used by SAG13 broken power law.
        Rplim (float ndarray):
            Minimum radius used by SAG13 broken power law.
        SAG13starMass (astropy Quantity):
            Assumed stellar mass corresponding to the given set of coefficients.
        mu (astropy Quantity):
            Gravitational parameter associated with SAG13starMass.
        Ca (float 2x1 ndarray):
            Constants used for sampling.
    
    """

    def __init__(self, SAG13coeffs=[[.38, -.19, .26, 0.],[.73, -1.18, .59, 3.4]],
            SAG13starMass=1., Rprange=[2/3., 17.0859375],
            arange=[0.09084645, 1.45354324], **specs):
        
        # first initialize with KeplerLike constructor
        specs['Rprange'] = Rprange
        specs['arange'] = arange
        # load SAG13 star mass in solMass: 1.3 (F), 1 (G), 0.70 (K), 0.35 (M)
        self.SAG13starMass = float(SAG13starMass)*u.solMass
        self.mu = const.G*self.SAG13starMass
        
        # load SAG13 coefficients (Gamma, alpha, beta, Rplim)
        self.SAG13coeffs = np.array(SAG13coeffs, dtype=float)
        assert self.SAG13coeffs.ndim <= 2, "SAG13coeffs array dimension must be <= 2."
        # if only one row of coefficients, make sure the forth element
        # (minimum radius) is set to zero
        if self.SAG13coeffs.ndim == 1:
            self.SAG13coeffs = np.array(np.append(self.SAG13coeffs[:3], 0.), ndmin=2)
        # make sure the array is of shape (4, n) where the forth row
        # contains the minimum radius values (broken power law)
        if len(self.SAG13coeffs) != 4:
            self.SAG13coeffs = self.SAG13coeffs.T
        assert len(self.SAG13coeffs) == 4, "SAG13coeffs array must have 4 rows."
        # sort by minimum radius
        self.SAG13coeffs = self.SAG13coeffs[:,np.argsort(self.SAG13coeffs[3,:])]
        
        # split out SAG13 coeffs
        self.Gamma = self.SAG13coeffs[0,:]
        self.alpha = self.SAG13coeffs[1,:]
        self.beta = self.SAG13coeffs[2,:]
        self.Rplim = np.append(self.SAG13coeffs[3,:], np.inf)
        
        KeplerLike2.__init__(self, **specs)
        
        # intermediate function
        m = self.mu.to('AU3/year2').value
        ftmp = lambda x,b,m=m,ak=self.smaknee: (2.*np.pi*np.sqrt(x**3/m))**(b-1.)*(3.*np.pi*np.sqrt(x/m))*np.exp(-(x/ak)**3)
        # intermediate constants used elsewhere
        self.Ca = np.zeros((2,))
        for i in xrange(2):
            self.Ca[i] = integrate.quad(ftmp, self.arange[0].to('AU').value, self.arange[1].to('AU').value, args=(self.beta[i],))[0]
        
        # set up samplers for sma and Rp
        # probability density function of sma given Rp < Rplim[1]
        f_sma_given_Rp1 = lambda a, beta=self.beta[0], m=m, C=self.Ca[0], smaknee=self.smaknee: self.dist_sma_given_radius(a,beta,m,C,smaknee)
        # sampler for Rp < Rplim:
        # unitless sma range
        ar = self.arange.to('AU').value
        self.sma_sampler1 = InverseTransformSampler(f_sma_given_Rp1, ar[0], ar[1])
        # probability density function of sma given Rp > Rplim[1]
        f_sma_given_Rp2 = lambda a, beta=self.beta[1], m=m, C=self.Ca[1], smaknee=self.smaknee: self.dist_sma_given_radius(a,beta,m,C,smaknee)
        self.sma_sampler2 = InverseTransformSampler(f_sma_given_Rp2, ar[0], ar[1])
        
        self.Rp_sampler = InverseTransformSampler(self.dist_radius, self.Rprange[0].to('earthRad').value, self.Rprange[1].to('earthRad').value)
        
        # determine eta
        if self.Rprange[1].to('earthRad').value < self.Rplim[1]:
            self.eta = self.Gamma[0]*(self.Rprange[1].to('earthRad').value**self.alpha[0]-self.Rprange[0].to('earthRad').value**self.alpha[0])/self.alpha[0]*self.Ca[0]
        elif self.Rprange[0].to('earthRad').value > self.Rplim[1]:
            self.eta = self.Gamma[1]*(self.Rprange[1].to('earthRad').value**self.alpha[1]-self.Rprange[1].to('earthRad').value**self.alpha[1])/self.alpha[1]*self.Ca[1]
        else:
            self.eta = self.Gamma[0]*(self.Rplim[1]**self.alpha[0]-self.Rprange[0].to('earthRad').value**self.alpha[0])/self.alpha[0]*self.Ca[0]
            self.eta+= self.Gamma[1]*(self.Rprange[1].to('earthRad').value**self.alpha[1]-self.Rplim[1]**self.alpha[1])/self.alpha[1]*self.Ca[1]

        self._outspec['eta'] = self.eta
        
        # populate _outspec with SAG13 specific attributes
        self._outspec['SAG13starMass'] = self.SAG13starMass.to('solMass').value
        self._outspec['SAG13coeffs'] = self.SAG13coeffs
        self._outspec['Gamma'] = self.Gamma
        self._outspec['alpha'] = self.alpha
        self._outspec['beta'] = self.beta
        self._outspec['Rplim'] = self.Rplim

    def gen_radius_sma(self, n):
        """Generate radius values in earth radius and semi-major axis values in AU.
        
        This method is called by gen_radius and gen_sma.
        <commit_msg>fixing trivial bug in eta calculation<commit_after>from EXOSIMS.PlanetPopulation.KeplerLike2 import KeplerLike2
from EXOSIMS.util.InverseTransformSampler import InverseTransformSampler
import astropy.units as u
import astropy.constants as const
import numpy as np
import scipy.integrate as integrate

class SAG13(KeplerLike2):
    """Planet Population module based on SAG13 occurrence rates.
    
    This is the current working model based on averaging multiple studies. 
    These do not yet represent official scientific values.
    
    Attributes: 
        SAG13coeffs (float 4x2 ndarray):
            Coefficients used by the SAG13 broken power law. The 4 lines
            correspond to Gamma, alpha, beta, and the minimum radius.
        Gamma (float ndarray):
            Gamma coefficients used by SAG13 broken power law.
        alpha (float ndarray):
            Alpha coefficients used by SAG13 broken power law.
        beta (float ndarray):
            Beta coefficients used by SAG13 broken power law.
        Rplim (float ndarray):
            Minimum radius used by SAG13 broken power law.
        SAG13starMass (astropy Quantity):
            Assumed stellar mass corresponding to the given set of coefficients.
        mu (astropy Quantity):
            Gravitational parameter associated with SAG13starMass.
        Ca (float 2x1 ndarray):
            Constants used for sampling.
    
    """

    def __init__(self, SAG13coeffs=[[.38, -.19, .26, 0.],[.73, -1.18, .59, 3.4]],
            SAG13starMass=1., Rprange=[2/3., 17.0859375],
            arange=[0.09084645, 1.45354324], **specs):
        
        # first initialize with KeplerLike constructor
        specs['Rprange'] = Rprange
        specs['arange'] = arange
        # load SAG13 star mass in solMass: 1.3 (F), 1 (G), 0.70 (K), 0.35 (M)
        self.SAG13starMass = float(SAG13starMass)*u.solMass
        self.mu = const.G*self.SAG13starMass
        
        # load SAG13 coefficients (Gamma, alpha, beta, Rplim)
        self.SAG13coeffs = np.array(SAG13coeffs, dtype=float)
        assert self.SAG13coeffs.ndim <= 2, "SAG13coeffs array dimension must be <= 2."
        # if only one row of coefficients, make sure the forth element
        # (minimum radius) is set to zero
        if self.SAG13coeffs.ndim == 1:
            self.SAG13coeffs = np.array(np.append(self.SAG13coeffs[:3], 0.), ndmin=2)
        # make sure the array is of shape (4, n) where the forth row
        # contains the minimum radius values (broken power law)
        if len(self.SAG13coeffs) != 4:
            self.SAG13coeffs = self.SAG13coeffs.T
        assert len(self.SAG13coeffs) == 4, "SAG13coeffs array must have 4 rows."
        # sort by minimum radius
        self.SAG13coeffs = self.SAG13coeffs[:,np.argsort(self.SAG13coeffs[3,:])]
        
        # split out SAG13 coeffs
        self.Gamma = self.SAG13coeffs[0,:]
        self.alpha = self.SAG13coeffs[1,:]
        self.beta = self.SAG13coeffs[2,:]
        self.Rplim = np.append(self.SAG13coeffs[3,:], np.inf)
        
        KeplerLike2.__init__(self, **specs)
        
        # intermediate function
        m = self.mu.to('AU3/year2').value
        ftmp = lambda x,b,m=m,ak=self.smaknee: (2.*np.pi*np.sqrt(x**3/m))**(b-1.)*(3.*np.pi*np.sqrt(x/m))*np.exp(-(x/ak)**3)
        # intermediate constants used elsewhere
        self.Ca = np.zeros((2,))
        for i in xrange(2):
            self.Ca[i] = integrate.quad(ftmp, self.arange[0].to('AU').value, self.arange[1].to('AU').value, args=(self.beta[i],))[0]
        
        # set up samplers for sma and Rp
        # probability density function of sma given Rp < Rplim[1]
        f_sma_given_Rp1 = lambda a, beta=self.beta[0], m=m, C=self.Ca[0], smaknee=self.smaknee: self.dist_sma_given_radius(a,beta,m,C,smaknee)
        # sampler for Rp < Rplim:
        # unitless sma range
        ar = self.arange.to('AU').value
        self.sma_sampler1 = InverseTransformSampler(f_sma_given_Rp1, ar[0], ar[1])
        # probability density function of sma given Rp > Rplim[1]
        f_sma_given_Rp2 = lambda a, beta=self.beta[1], m=m, C=self.Ca[1], smaknee=self.smaknee: self.dist_sma_given_radius(a,beta,m,C,smaknee)
        self.sma_sampler2 = InverseTransformSampler(f_sma_given_Rp2, ar[0], ar[1])
        
        self.Rp_sampler = InverseTransformSampler(self.dist_radius, self.Rprange[0].to('earthRad').value, self.Rprange[1].to('earthRad').value)
        
        # determine eta
        if self.Rprange[1].to('earthRad').value < self.Rplim[1]:
            self.eta = self.Gamma[0]*(self.Rprange[1].to('earthRad').value**self.alpha[0]-self.Rprange[0].to('earthRad').value**self.alpha[0])/self.alpha[0]*self.Ca[0]
        elif self.Rprange[0].to('earthRad').value > self.Rplim[1]:
            self.eta = self.Gamma[1]*(self.Rprange[1].to('earthRad').value**self.alpha[1]-self.Rprange[0].to('earthRad').value**self.alpha[1])/self.alpha[1]*self.Ca[1]
        else:
            self.eta = self.Gamma[0]*(self.Rplim[1]**self.alpha[0]-self.Rprange[0].to('earthRad').value**self.alpha[0])/self.alpha[0]*self.Ca[0]
            self.eta+= self.Gamma[1]*(self.Rprange[1].to('earthRad').value**self.alpha[1]-self.Rplim[1]**self.alpha[1])/self.alpha[1]*self.Ca[1]

        self._outspec['eta'] = self.eta
        
        # populate _outspec with SAG13 specific attributes
        self._outspec['SAG13starMass'] = self.SAG13starMass.to('solMass').value
        self._outspec['SAG13coeffs'] = self.SAG13coeffs
        self._outspec['Gamma'] = self.Gamma
        self._outspec['alpha'] = self.alpha
        self._outspec['beta'] = self.beta
        self._outspec['Rplim'] = self.Rplim

    def gen_radius_sma(self, n):
        """Generate radius values in earth radius and semi-major axis values in AU.
        
        This method is called by gen_radius and gen_sma.
        
{'fraction_non_alphanumeric': 0.1307032590051458, 'fraction_numerical': 0.038593481989708404, 'mean_word_length': 2.7113303628262253, 'pattern_counts': {'":': 0, '<': 13, '<?xml version=': 0, '>': 7, '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': '2742106', 'n_tokens_mistral': 4302, 'n_tokens_neox': 4017, 'n_words': 875}