added script that shows the different distributions that are available, and plots them

binarycpython/utils/distribution_functions.py

@@ -17,6 +17,7 @@ def flat(parameter):
     """
     Dummt distribution function that returns 1
     """
+
     return 1
 
 
@@ -70,8 +71,7 @@ def powerlaw(min_val, max_val, k, x):
         return y
 
 def calculate_constants_three_part_powerlaw(m0, m1, m2, m_max, p1, p2, p3):
-    print("Initialising constants for the three-part powerlaw: m0={} m1={} m2={} m_max={} p1={} p2={} p3={}\n".format(m0, m1, m2, m_max, p1, p2, p3))
-    # local $SIG{__DIE__} = sub { Carp::confess @_ }; ?
+    # print("Initialising constants for the three-part powerlaw: m0={} m1={} m2={} m_max={} p1={} p2={} p3={}\n".format(m0, m1, m2, m_max, p1, p2, p3))
 
     array_constants_three_part_powerlaw = [0, 0, 0]
 
@@ -83,7 +83,6 @@ def calculate_constants_three_part_powerlaw(m0, m1, m2, m_max, p1, p2, p3):
     array_constants_three_part_powerlaw[0] = array_constants_three_part_powerlaw[1] * ((m1**p2)*(m1**(-p1)))
     array_constants_three_part_powerlaw[2] = array_constants_three_part_powerlaw[1] * ((m2**p2)*(m2**(-p3)))
 
-    # print(array_constants_three_part_powerlaw)
     return array_constants_three_part_powerlaw
     # $$array[1]=(($m1**$p2)*($m1**(-$p1)))*
     # (1.0/(1.0+$p1))*
@@ -141,13 +140,14 @@ def set_opts(opts, newopts):
     """
 
     # DONE: put in check to make sure that the newopts keys are contained in opts
+    # TODO: change this to just a dict.update
 
     if newopts:
         for opt in newopts.keys():
             if opt in opts.keys():
                 opts[key] = newopts[key]
 
-    return newopts
+    return opt
 
 def gaussian(x, mean, sigma, gmin, gmax):
     """
@@ -200,19 +200,22 @@ def gaussian_func(x, mean, sigma):
 # Mass distributions
 ##### 
 
-def Kroupa2001(m, newopts):
+def Kroupa2001(m, newopts=None):
     """
-    probability distribution function for kroupa 2001 IMF 
+    Probability distribution function for kroupa 2001 IMF 
 
-    Input: Mass, and dict of new options
+    Input: Mass, (and optional: dict of new options. Input the  
+        default = {'m0':0.1, 'm1':0.5, 'm2':1, 'mmax':100, 'p1':-1.3, 'p2':-2.3, 'p3':-2.3}
     """
 
     # Default params and override them 
-    opts = set_opts({'m0':0.1, 'm1':0.5, 'm2':1, 'mmax':100, 'p1':-1.3, 'p2':-2.3, 'p3':-2.3}, newopts)
-
-    return three_part_powerlaw(m, opts['m0'], opts['m0'], opts['m2'], opts['m0'], opts['m0'], opts['m0'], opts['m0'])
+    default = {'m0':0.1, 'm1':0.5, 'm2':1, 'mmax':100, 'p1':-1.3, 'p2':-2.3, 'p3':-2.3}
+    value_dict = default.copy()
+    if newopts:
+        value_dict.update(newopts)
+    return three_part_powerlaw(m, value_dict['m0'], value_dict['m1'], value_dict['m2'], value_dict['mmax'], value_dict['p1'], value_dict['p2'], value_dict['p3'])
 
-def ktg93():
+def ktg93(m, newopts):
     """
     Wrapper for mass distribution of KTG93
     """
@@ -243,9 +246,12 @@ def ktg93():
     # }
 
     # set options
-    opts = set_opts({'m0':0.1, 'm1':0.5, 'm2':1.0, 'mmax':80, 'p1':-1.3, 'p2':-2.2, 'p3':-2.7}, newopts)
-
-    return three_part_powerlaw(m, opts['m0'], opts['m0'], opts['m2'], opts['m0'], opts['m0'], opts['m0'], opts['m0'])
+    # opts = set_opts({'m0':0.1, 'm1':0.5, 'm2':1.0, 'mmax':80, 'p1':-1.3, 'p2':-2.2, 'p3':-2.7}, newopts)
+    defaults = {'m0':0.1, 'm1':0.5, 'm2':1.0, 'mmax':80, 'p1':-1.3, 'p2':-2.2, 'p3':-2.7}
+    value_dict = default.copy()
+    if newopts:
+        value_dict.update(newopts)
+    return three_part_powerlaw(m, value_dict['m0'], value_dict['m0'], value_dict['m2'], value_dict['m0'], value_dict['m0'], value_dict['m0'], value_dict['m0'])
 
 
 # sub ktg93_lnspace
@@ -255,9 +261,51 @@ def ktg93():
 #     return ktg93(@_) * $m;
 # }
 
-## 
+
+def imf_tinsley1980(m):
+    """
+    From Tinsley 1980 (defined up until 80Msol)
+    """
+
+    return three_part_powerlaw(m, 0.1, 2.0, 10.0, 80.0, -2.0, -2.3, -3.3)
+
+def imf_scalo1986(m):
+    """
+    From Scalo 1986 (defined up until 80Msol)
+    """
+    return three_part_powerlaw(m, 0.1, 1.0, 2.0, 80.0, -2.35, -2.35, -2.70)
+
+
+def imf_scalo1998(m):
+    """
+    From scalo 1998
+    """
+    return three_part_powerlaw(m, 0.1, 1.0, 10.0, 80.0, -1.2, -2.7, -2.3)
+
+def imf_chabrier2003(m):
+    """
+    # IMF of Chabrier 2003 PASP 115:763-795    
+    """    
+    Chabrier_logmc = math.log10(0.079)
+    Chabrier_sigma2 = (0.69*0.69)
+    Chabrier_A1 = 0.158
+    Chabrier_A2 = 4.43e-2
+    Chabrier_x = -1.3
+    if  m < 0:
+        print("below bounds")
+        raise ValueError
+    if (0 < m < 1.0):
+        A = 0.158
+        dm = math.log10(m) - Chabrier_logmc
+        p = Chabrier_A1 * math.exp(-(dm**2)/(2.0*Chabrier_sigma2))
+    else:
+        p = Chabrier_A2 * (m ** Chabrier_x)
+    p = p/(0.1202462 * m * math.log(10))
+    return p
+
+########################################################################
 # Binary fractions
-## 
+########################################################################
 def Arenou2010_binary_fraction(m):
     # Arenou 2010 function for the binary fraction as f(M1)
     #
@@ -288,9 +336,18 @@ def raghavan2010_binary_fraction(m):
         )
 # print(raghavan2010_binary_fraction(2))
 
-## 
+######################################################################## 
 # Period distributions
-## 
+######################################################################## 
+
+def duquennoy1991(x):
+    """
+    Period distribution from Duquennoy + Mayor 1991
+
+    Input:
+        x: logperiod
+    """
+    return gaussian(x, 4.8, 2.3, -2, 12)
 
 def sana12(M1, M2, a, P, amin, amax, x0, x1, p): # TODO: ? wtf. vague input
     """
@@ -302,10 +359,13 @@ def sana12(M1, M2, a, P, amin, amax, x0, x1, p): # TODO: ? wtf. vague input
     and is be given by dp/dlogP ~ (logP)^p for all other binary configurations (default p=-0.55)
 
     arguments are M1, M2, a, Period P, amin, amax, x0=log P0, x1=log P1, p
+
+    example args: 10, 5, ?, P, ?, ?, -2, 12, -0.55
     """
 
     res = 0
 
+
     if (M1 < 15) or (M2/M1 < 0.1):
         res = (1.0/(math.log(amax)-math.log(amin)))
     else:
@@ -334,6 +394,10 @@ def sana12(M1, M2, a, P, amin, amax, x0, x1, p): # TODO: ? wtf. vague input
     return res
 # print(sana12(10, 2, 10, 100, 1, 1000, math.log(10), math.log(1000), 6))
 
+########################################################################
+# Mass ratio distributions
+########################################################################
+
 def flatsections(x, opts):
     """
     

examples/example_plotting_distributions.py

+import math
+
+import numpy as np
+import matplotlib.pyplot as plt
+
+from binarycpython.utils.distribution_functions import three_part_powerlaw, Kroupa2001, Arenou2010_binary_fraction, raghavan2010_binary_fraction,imf_scalo1998,  imf_scalo1986, imf_tinsley1980, imf_scalo1998, imf_chabrier2003, flatsections, duquennoy1991, sana12
+from binarycpython.utils.useful_funcs import calc_sep_from_period
+################################################
+# Example script to plot the available probability distributions.
+# TODO: 
+
+################################################
+# 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()
+
+################################################
+# 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]
+
+# 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
+
+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
+
+sana12_distribution_q05 = [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 = 10
+m2 = 1
+sana12_distribution_q01 = [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 = 10
+m2 = 10
+sana12_distribution_q1 = [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]
+
+
+plt.plot(logperiod_values, duquennoy1991_distribution, label='Duquennoy & Mayor 1991')
+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.title('Period distributions')
+plt.ylabel(r'Probability')
+plt.xlabel(r'Log10(orbital period)')
+plt.grid()
+plt.legend()
+plt.show()
+
+
+################################################
+# 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.