added distribution. will add others later. added method to store globally normalisation constants

import math

import numpy as np

from binarycpython.utils.useful_funcs import (

    calc_period_from_sep,

# File containing probability distributions

# Mostly copied from the perl modules

# TODO: make some things globally present? rob does this in his module. i guess it saves calculations but not sure if im gonna do that now

# TODO: make some things globally present? rob does this in his module. 

    # i guess it saves calculations but not sure if im gonna do that now

# TODO: Add the stuff from the IMF file

# TODO: call all of these functions to check whether they work

# TODO: make global constants stuff

LOG_LN_CONVERTER = 1.0 / math.log(10.0)

distribution_constants = {} # To store the constants in

def prepare_dict(global_dict, list_of_sub_keys):

    """

    Function that makes sure that the global dict is prepared to have a value set there

    """

    internal_dict_value = global_dict

    # This loop almost mimics a recursive loop into the dictionary.

    # It checks whether the first key of the list is present, if not; set it with an empty dict.

    # Then it overrides itself to be that (new) item, and goes on to do that again, until the list exhausted

    for k in list_of_sub_keys:

        # If the sub key doesnt exist then make an empty dict

        if not internal_dict_value.get(k, None):

            internal_dict_value[k] = {}

        internal_dict_value = internal_dict_value[k]

def flat():

    """

    Dummy distribution function that returns 1

########################################################################

def duquennoy1991(x):

def duquennoy1991(logper):

    """

    Period distribution from Duquennoy + Mayor 1991

    Input:

        x: logperiod

        logper: logperiod

    """

    return gaussian(x, 4.8, 2.3, -2, 12)

    return gaussian(logper, 4.8, 2.3, -2, 12)

def sana12(M1, M2, a, P, amin, amax, x0, x1, p): 

    """

    res = 0

    if (M1 < 15) or (M2 / M1 < 0.1):

        res = 1.0 / (math.log(amax) - math.log(amin))

    else:

        # For more details see the LyX document of binary_c for this distribution

        # where the variables and normalizations are given

        # we use the notation x=log(P), xmin=log(Pmin), x0=log(P0), ... to determine the

        x = LOG_LN_CONVERTER * math.log(p)

        x = LOG_LN_CONVERTER * math.log(P)

        xmin = LOG_LN_CONVERTER * math.log(calc_period_from_sep(M1, M2, amin))

        xmax = LOG_LN_CONVERTER * math.log(calc_period_from_sep(M1, M2, amax))

        print("M1 M2 amin amax P x xmin xmax")

        print(M1, M2, amin, amax, P, x, xmin, xmax)

        # my $x0 = 0.15;

        # my $x1 = 3.5;

    return res

def Izzard2012_period_distribution(P, M1, log10Pmin=1):

    """

    period distribution which interpolates between 

    Duquennoy and Mayor 1991 at low mass (G/K spectral type <~1.15Msun)

    and Sana et al 2012 at high mass (O spectral type >~16.3Msun)

    This gives dN/dlogP, i.e. DM/Raghavan's Gaussian in log10P at low mass

    and Sana's power law (as a function of logP) at high mass

    input::

        P (float): period

        M1 (float): Primary star mass

        log10Pmin (float): Minimum period in base log10 (optional)

    """

    # Check if there is input and force it to be at least 1

    log10Pmin //= -1.0

    log10Pmin = max(-1.0, log10Pmin)

    # save mass input and limit mass used (M1 from now on) to fitted range

    Mwas = M1

    M1 = max(1.15, min(16.3, M1))

    print("Izzard2012 called for M={} (trunc'd to {}), P={}\n".format(Mwas, M1, P))

    # Calculate the normalisations

    # need to normalize the distribution for this mass 

    # (and perhaps secondary mass)

    prepare_dict(distribution_constants, ['Izzard2012', M1])

    if not distribution_constants['Izzard2012'][M1].get(log10Pmin):

        distribution_constants['Izzard2012'][M1][log10Pmin] = 1 # To prevent this loop from going recursive

        N = 200.0 # Resolution for normalisation. I hope 1000 is enough

        dlP = (10.0 - log10Pmin)/N

        C = 0 # normalisation const.

        for lP in np.arange(log10Pmin, 10, dlP):

            C += dlP * Izzard2012_period_distribution(10**lP, M1, log10Pmin)

        distribution_constants['Izzard2012'][M1][log10Pmin] = 1.0/C;

    print("Normalization constant for Izzard2012 M={} (log10Pmin={}) is\

        {}\n".format(M1, log10Pmin, distribution_constants['Izzard2012'][M1][log10Pmin]))

    lP = math.log10(P); # log period

    # # fits

    mu = interpolate_in_mass_izzard2012(M1, -17.8, 5.03)

    sigma = interpolate_in_mass_izzard2012(M1, 9.18, 2.28)

    K = interpolate_in_mass_izzard2012(M1, 6.93e-2, 0.0)

    nu = interpolate_in_mass_izzard2012(M1, 0.3, -1)

    g = 1.0/(1.0+1e-30**(lP-nu))

    lPmu = lP - mu

    print("M={} ({}) P={} : mu={} sigma={} K={} nu={} norm=%g\n".format(

        Mwas, M1, P, mu, sigma, K, nu))

    # print "FUNC $distdata{Izzard2012}{$M}{$log10Pmin} * (exp(- (x-$mu)**2/(2.0*$sigma*$sigma) ) + $K/MAX(0.1,$lP)) * $g;\n";

    if ((lP < log10Pmin) or (lP > 10.0)):

        return 0

    else:

        return distribution_constants['Izzard2012'][M1][log10Pmin] * (math.exp(- lPmu * lPmu / (2.0 * sigma * sigma)) + K/max(0.1, lP)) * g;

def interpolate_in_mass_izzard2012(M, high, low):

    """

    Function to interpolate in mass

    high: at M=16.3

    low: at 1.15

    """

    log_interpolation = False

    if log_interpolation:

        return (high-low)/(math.log10(16.3)-math.log10(1.15)) * (math.log10(M)-math.log10(1.15)) + low

    else:

        return (high-low)/(16.3-1.15) * (M-1.15) + low

# print(sana12(10, 2, 10, 100, 1, 1000, math.log(10), math.log(1000), 6))

    Kroupa2001,

    Arenou2010_binary_fraction,

    raghavan2010_binary_fraction,

    imf_scalo1998,

    imf_scalo1986,

    imf_tinsley1980,

    imf_scalo1998,

    imf_chabrier2003,

    flatsections,

    duquennoy1991,

    sana12,

    Izzard2012_period_distribution,

)

from binarycpython.utils.useful_funcs import calc_sep_from_period

################################################

# mass distribution plots

################################################

# mass_values = np.arange(0.11, 80, .1)

# kroupa_probability = [Kroupa2001(mass) for mass in mass_values]

# scalo1986 = [imf_scalo1986(mass) for mass in mass_values]

# tinsley1980 = [imf_tinsley1980(mass) for mass in mass_values]

# scalo1998 = [imf_scalo1998(mass) for mass in mass_values]

# chabrier2003 = [imf_chabrier2003(mass) for mass in mass_values]

# plt.plot(mass_values, kroupa_probability, label='Kroupa')

# plt.plot(mass_values, scalo1986, label='scalo1986')

# plt.plot(mass_values, tinsley1980, label='tinsley1980')

# plt.plot(mass_values, scalo1998, label='scalo1998')

# plt.plot(mass_values, chabrier2003, label='chabrier2003')

# plt.title('Probability distribution for mass of primary')

# plt.ylabel(r'Probability')

# plt.xlabel(r'Mass (M$_{\odot}$)')

# plt.yscale('log')

# plt.xscale('log')

# plt.grid()

# plt.legend()

# plt.show()

def plot_mass_distributions():

    mass_values = np.arange(0.11, 80, .1)

    kroupa_probability = [Kroupa2001(mass) for mass in mass_values]

    scalo1986 = [imf_scalo1986(mass) for mass in mass_values]

    tinsley1980 = [imf_tinsley1980(mass) for mass in mass_values]

    scalo1998 = [imf_scalo1998(mass) for mass in mass_values]

    chabrier2003 = [imf_chabrier2003(mass) for mass in mass_values]

    plt.plot(mass_values, kroupa_probability, label='Kroupa')

    plt.plot(mass_values, scalo1986, label='scalo1986')

    plt.plot(mass_values, tinsley1980, label='tinsley1980')

    plt.plot(mass_values, scalo1998, label='scalo1998')

    plt.plot(mass_values, chabrier2003, label='chabrier2003')

    plt.title('Probability distribution for mass of primary')

    plt.ylabel(r'Probability')

    plt.xlabel(r'Mass (M$_{\odot}$)')

    plt.yscale('log')

    plt.xscale('log')

    plt.grid()

    plt.legend()

    plt.show()

################################################

# Binary fraction distributions

################################################

# arenou_binary_distibution = [Arenou2010_binary_fraction(mass) for mass in mass_values]

# raghavan2010_binary_distribution = [raghavan2010_binary_fraction(mass) for mass in mass_values ]

# plt.plot(mass_values, arenou_binary_distibution, label='arenou 2010')

# plt.plot(mass_values, raghavan2010_binary_distribution, label='Raghavan 2010')

# plt.title('Binary fractions distributions')

# plt.ylabel(r'Binary fraction')

# plt.xlabel(r'Mass (M$_{\odot}$)')

# # plt.yscale('log')

# plt.xscale('log')

# plt.grid()

# plt.legend()

# plt.show()

def plot_binary_fraction_distributions():

    arenou_binary_distibution = [Arenou2010_binary_fraction(mass) for mass in mass_values]

    raghavan2010_binary_distribution = [raghavan2010_binary_fraction(mass) for mass in mass_values ]

    plt.plot(mass_values, arenou_binary_distibution, label='arenou 2010')

    plt.plot(mass_values, raghavan2010_binary_distribution, label='Raghavan 2010')

    plt.title('Binary fractions distributions')

    plt.ylabel(r'Binary fraction')

    plt.xlabel(r'Mass (M$_{\odot}$)')

    # plt.yscale('log')

    plt.xscale('log')

    plt.grid()

    plt.legend()

    plt.show()

################################################

# Mass ratio distributions

################################################

# mass_ratios = np.arange(0, 1, .01)

# example_mass = 2

# flat_dist = [flatsections(q, opts=[{'min':0.1/example_mass, 'max':0.8, 'height':1}, {'min': 0.8, 'max':1.0, 'height': 1.0}]) for q in mass_ratios]

def plot_mass_ratio_distributions():

    mass_ratios = np.arange(0, 1, .01)

    example_mass = 2

    flat_dist = [

        flatsections(

            q,

            opts=[

                    {'min':0.1/example_mass, 'max':0.8, 'height':1},

                    {'min': 0.8, 'max':1.0, 'height': 1.0}

                ]

            ) for q in mass_ratios]

# plt.plot(mass_ratios, flat_dist, label='Flat')

# plt.title('Mass ratio distributions')

# plt.ylabel(r'Probability')

# plt.xlabel(r'Mass ratio (q = $\frac{M1}{M2}$) ')

# plt.grid()

# plt.legend()

# plt.show()

    plt.plot(mass_ratios, flat_dist, label='Flat')

    plt.title('Mass ratio distributions')

    plt.ylabel(r'Probability')

    plt.xlabel(r'Mass ratio (q = $\frac{M1}{M2}$) ')

    plt.grid()

    plt.legend()

    plt.show()

################################################

# Period distributions

################################################

# TODO: fix this

def plot_period_distributions():

    logperiod_values = np.arange(-2, 12, 0.1)

    duquennoy1991_distribution = [duquennoy1991(logper) for logper in logperiod_values]

    # Sana12 distributions

m1 = 10

m2 = 5

    period_min = 10 ** 0.15

    period_max = 10 ** 5.5

    m1 = 20

    m2 = 15

    sana12_distribution_q05 = [

        sana12(

            m1,

        for logper in logperiod_values

    ]

m1 = 10

    m1 = 30

    m2 = 1

    sana12_distribution_q0033 = [

        sana12(

            m1,

            m2,

            calc_sep_from_period(m1, m2, 10 ** logper),

            10 ** logper,

            calc_sep_from_period(m1, m2, period_min),

            calc_sep_from_period(m1, m2, period_max),

            math.log10(period_min),

            math.log10(period_max),

            -0.55,

        )

        for logper in logperiod_values

    ]

    m1 = 30

    m2 = 3

    sana12_distribution_q01 = [

        sana12(

            m1,

        for logper in logperiod_values

    ]

m1 = 10

m2 = 10

    m1 = 30

    m2 = 30

    sana12_distribution_q1 = [

        sana12(

            m1,

        for logper in logperiod_values

    ]

    Izzard2012_period_distribution_10 = [Izzard2012_period_distribution(10**logperiod, 10) for logperiod in logperiod_values]

    Izzard2012_period_distribution_20 = [Izzard2012_period_distribution(10**logperiod, 20) for logperiod in logperiod_values]

    plt.plot(logperiod_values, duquennoy1991_distribution, label="Duquennoy & Mayor 1991")

    plt.plot(logperiod_values, sana12_distribution_q0033, label="Sana 12 (q=0.033)")

    plt.plot(logperiod_values, sana12_distribution_q05, label="Sana 12 (q=0.5)")

    plt.plot(logperiod_values, sana12_distribution_q01, label="Sana 12 (q=0.1)")

    plt.plot(logperiod_values, sana12_distribution_q1, label="Sana 12 (q=1)")

    plt.plot(logperiod_values, Izzard2012_period_distribution_10, label='Izzard2012 (M=10)')

    plt.plot(logperiod_values, Izzard2012_period_distribution_20, label='Izzard2012 (M=20)')

    plt.title("Period distributions")

    plt.ylabel(r"Probability")

    plt.xlabel(r"Log10(orbital period)")

    plt.legend()

    plt.show()

plot_period_distributions()

################################################

# Sampling part of distribution and calculating probability ratio

################################################

# TODO show the difference between sampling over the full range, or taking a smaller range initially and compensating for it.

# val = Izzard2012_period_distribution(1000, 10)

# print(val)

# val2 = Izzard2012_period_distribution(100, 10)

# print(val2)

# from binarycpython.utils.distribution_functions import (interpolate_in_mass_izzard2012)

# # print(interpolate_in_mass_izzard2012(15, 0.3, -1))