typed distribution_functions

binarycpython/utils/distribution_functions.py

@@ -3,33 +3,41 @@ Module containing the predefined distribution functions
 
 The user can use any of these distribution functions to
 generate probability distributions for sampling populations
+
+There are distributions for the following properties:
+- mass
+- period
+- mass ratio
+- binary fraction
+
+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 global constants stuff
+TODO: add eccentricity distribution: thermal
+TODO: Add SFH distributions depending on redshift
+TODO: Add metallicity distributions depending on redshift
+TODO: Add initial rotational velocity distributions
 """
 
 
 import math
 import numpy as np
-
+from typing import Optional, Union
 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: Add the stuff from the IMF file
-# TODO: call all of these functions to check whether they work
-# TODO: make global constants stuff
-# TODO: make description of module submodule
-
 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):
+def prepare_dict(global_dict: dict, list_of_sub_keys: list) -> None:
     """
-    Function that makes sure that the global dict is prepared to have a value set there
+    Function that makes sure that the global dict is prepared to have a value set there. This dictionary will store values and factors for the distribution functions, so that they dont have to be calculated each time.
+
+    Args:
+        global_dict: globablly acessible dictionary where factors are stored in
+        list_of_sub_keys: List of keys that must become be(come) present in the global_dict
     """
 
     internal_dict_value = global_dict
@@ -44,25 +52,87 @@ def prepare_dict(global_dict, list_of_sub_keys):
         internal_dict_value = internal_dict_value[k]
 
 
-def flat():
+def set_opts(opts: dict, newopts: dict) -> dict:
+    """
+    Function to take a default dict and override it with newer values.
+
+    # TODO: consider changing this to just a dict.update
+
+    Args:
+        opts: dictionary with default values
+        newopts: dictionary with new values
+
+    Returns:
+        returns an updated dictionary
+    """
+
+    if newopts:
+        for opt in newopts.keys():
+            if opt in opts.keys():
+                opts[opt] = newopts[opt]
+
+    return opts
+
+
+def flat() -> float:
     """
     Dummy distribution function that returns 1
+
+    Returns:
+        a flat uniform distribution: 1    
     """
 
-    return 1
+    return 1.0
 
 
-def number(value):
+def number(value: Union[int, float]) -> Union[int, float]:
     """
     Dummy distribution function that returns the input
+    
+    Args:
+        value: the value that will be returned by this function.
+
+    Returns: 
+        the value that was provided
     """
 
     return value
 
 
-def powerlaw_constant(min_val, max_val, k):
+def const(min_bound: Union[int, float], max_bound: Union[int, float], val: float=None) -> Union[int, float]:
+    """
+    a constant distribution function between min=min_bound and max=max_bound.
+
+    Args:
+        min_bound: lower bound of the range
+        max_bound: upper bound of the range
+    
+    Returns:
+            returns the value of 1/(max_bound-min_bound). If val is provided, it will check whether min_bound < val <= max_bound. if not: returns 0
+    """
+
+    if val:
+        if not min_bound < val <= max_bound:
+            print("out of bounds")
+            prob = 0
+            return prob
+    prob = 1.0 / (min_bound - max_bound)
+    return prob
+
+
+def powerlaw_constant(min_val: Union[int, float], max_val: Union[int, float], k: Union[int, float]) -> Union[int, float]:
     """
     Function that returns the constant to normalise a powerlaw
+    
+    TODO: what if k is -1?
+        
+    Args:
+        min_val: lower bound of the range
+        max_val: upper bound of the range
+        k: powerlaw slope
+
+    Returns:
+        constant to normalize the given powerlaw between the min_val and max_val range
     """
 
     k1 = k + 1.0
@@ -76,9 +146,18 @@ def powerlaw_constant(min_val, max_val, k):
     return powerlaw_const
 
 
-def powerlaw(min_val, max_val, k, x):
+def powerlaw(min_val: Union[int, float], max_val: Union[int, float], k: Union[int, float], x: Union[int, float]) -> Union[int, float]:
     """
     Single powerlaw with index k at x from min to max
+
+    Args:
+        min_val: lower bound of the powerlaw
+        max_val: upper bound of the powerlaw
+        k: slope of the power law
+        x: position at which we want to evaluate
+
+    Returns:
+        `probability` at the given position(x)
     """
 
     # Handle faulty value
@@ -102,9 +181,23 @@ def powerlaw(min_val, max_val, k, x):
     return prob
 
 
-def calculate_constants_three_part_powerlaw(m0, m1, m2, m_max, p1, p2, p3):
+def calculate_constants_three_part_powerlaw(m0: Union[int, float], m1: Union[int, float], m2: Union[int, float], m_max: Union[int, float], p1: Union[int, float], p2: Union[int, float], p3: Union[int, float]) -> Union[int, float]:
     """
     Function to calculate the constants for a three-part powerlaw
+
+    TODO: use the powerlaw_constant function to calculate all these values
+
+    Args:
+        m0: lower bound mass
+        m1: second boundary, between the first slope and the second slope
+        m2: third boundary, between the second slope and the third slope
+        m_max: upper bound mass
+        p1: first slope
+        p2: second slope
+        p3: third slope
+
+    Returns:
+        array of normalisation constants
     """
 
     # print("Initialising constants for the three-part powerlaw: m0={} m1={} m2={}\
@@ -152,68 +245,106 @@ def calculate_constants_three_part_powerlaw(m0, m1, m2, m_max, p1, p2, p3):
     # $threepart_powerlaw_consts{"@_"}=[@$array];
 
 
-def three_part_powerlaw(M, M0, M1, M2, M_MAX, P1, P2, P3):
+def three_part_powerlaw(m: Union[int, float], m0: Union[int, float], m1: Union[int, float], m2: Union[int, float], m_max: Union[int, float], p1: Union[int, float], p2: Union[int, float], p3: Union[int, float]) -> Union[int, float]:
     """
     Generalized three-part power law, usually used for mass distributions
+
+    Args:
+        m: mass at which we want to evaluate the distribution.
+        m0: lower bound mass
+        m1: second boundary, between the first slope and the second slope
+        m2: third boundary, between the second slope and the third slope
+        m_max: upper bound mass
+        p1: first slope
+        p2: second slope
+        p3: third slope
+
+    Returns:
+        'probability' at given mass m
     """
 
     # TODO: add check on whether the values exist
 
     three_part_powerlaw_constants = calculate_constants_three_part_powerlaw(
-        M0, M1, M2, M_MAX, P1, P2, P3
+        m0, m1, m2, m_max, p1, p2, p3
     )
 
     #
-    if M < M0:
+    if m < m0:
         prob = 0  # Below lower bound
-    elif M0 < M <= M1:
-        prob = three_part_powerlaw_constants[0] * (M ** P1)  # Between M0 and M1
-    elif M1 < M <= M2:
-        prob = three_part_powerlaw_constants[1] * (M ** P2)  # Between M1 and M2
-    elif M2 < M <= M_MAX:
-        prob = three_part_powerlaw_constants[2] * (M ** P3)  # Between M2 and M_MAX
+    elif m0 < m <= m1:
+        prob = three_part_powerlaw_constants[0] * (m ** p1)  # Between M0 and M1
+    elif m1 < m <= m2:
+        prob = three_part_powerlaw_constants[1] * (m ** p2)  # Between M1 and M2
+    elif m2 < m <= m_max:
+        prob = three_part_powerlaw_constants[2] * (m ** p3)  # Between M2 and M_MAX
     else:
         prob = 0  # Above M_MAX
 
     return prob
 
 
-def const(min_bound, max_bound, val=None):
+def gaussian_normalizing_const(mean: Union[int, float], sigma: Union[int, float], gmin: Union[int, float], gmax: Union[int, float]) -> Union[int, float]:
     """
-    a constant distribution function between min=$_[0] and max=$_[1]
+    Function to calculate the normalisation constant for the gaussian
+
+    Args:
+        mean: mean of the gaussian
+        sigma: standard deviation of the gaussian
+        gmin: lower bound of the range to calculate the probabilities in
+        gmax: upper bound of the range to calculate the probabilities in
+
+    Returns:
+        normalisation constant for the gaussian distribution(mean, sigma) between gmin and gmax
     """
 
-    if val:
-        if not min_bound < val <= max_bound:
-            print("out of bounds")
-            prob = 0
-            return prob
-    prob = 1.0 / (min_bound - max_bound)
-    return prob
+    # First time; calculate multipllier for given mean and sigma
+    ptot = 0
+    resolution = 1000
+    d = (gmax - gmin) / resolution
 
+    for i in range(resolution):
+        y = gmin + i * d
+        ptot += d * gaussian_func(y, mean, sigma)
 
-def set_opts(opts, newopts):
-    """
-    Function to take a default dict and override it with newer values.
+    # TODO: Set value in global
+    return ptot
+
+
+def gaussian_func(x: Union[int, float], mean: Union[int, float], sigma: Union[int, float]) -> Union[int, float]:
     """
+    Function to evaluate a gaussian at a given point, but this time without any boundaries. 
 
-    # DONE: put in check to make sure that the newopts keys are contained in opts
-    # TODO: change this to just a dict.update
+    Args:
+        x: location at which to evaluate the distribution
+        mean: mean of the gaussian
+        sigma: standard deviation of the gaussian
 
-    if newopts:
-        for opt in newopts.keys():
-            if opt in opts.keys():
-                opts[opt] = newopts[opt]
+    Returns:
+        value of the gaussian at x
+    """
+    gaussian_prefactor = 1.0 / math.sqrt(2.0 * math.pi)
 
-    return opts
+    r = 1.0 / (sigma)
+    y = (x - mean) * r
+    return gaussian_prefactor * r * math.exp(-0.5 * y ** 2)
 
 
-def gaussian(x, mean, sigma, gmin, gmax):
+def gaussian(x: Union[int, float], mean: Union[int, float], sigma: Union[int, float], gmin: Union[int, float], gmax: Union[int, float]) -> Union[int, float]:
     """
     Gaussian distribution function. used for e..g Duquennoy + Mayor 1991
 
-    Input: location, mean, sigma, min and max:
+    Args:
+        x: location at which to evaluate the distribution
+        mean: mean of the gaussian
+        sigma: standard deviation of the gaussian
+        gmin: lower bound of the range to calculate the probabilities in
+        gmax: upper bound of the range to calculate the probabilities in
+
+    Returns:
+        'probability' of the gaussian distribution between the boundaries, evaluated at x
     """
+
     # # location (X value), mean and sigma, min and max range
     # my ($x,$mean,$sigma,$gmin,$gmax) = @_;
 
@@ -228,46 +359,21 @@ def gaussian(x, mean, sigma, gmin, gmax):
     return prob
 
 
-def gaussian_normalizing_const(mean, sigma, gmin, gmax):
-    """
-    Function to calculate the normalisation constant for the gaussian
-    """
-
-    # First time; calculate multipllier for given mean and sigma
-    ptot = 0
-    resolution = 1000
-    d = (gmax - gmin) / resolution
-
-    for i in range(resolution):
-        y = gmin + i * d
-        ptot += d * gaussian_func(y, mean, sigma)
-
-    # TODO: Set value in global
-    return ptot
-
-
-def gaussian_func(x, mean, sigma):
-    """
-    Function to evaluate a gaussian at a given point
-    """
-    gaussian_prefactor = 1.0 / math.sqrt(2.0 * math.pi)
-
-    r = 1.0 / (sigma)
-    y = (x - mean) * r
-    return gaussian_prefactor * r * math.exp(-0.5 * y ** 2)
-
-
 #####
 # Mass distributions
 #####
 
 
-def Kroupa2001(m, newopts=None):
+def Kroupa2001(m: Union[int, float], newopts: dict=None) -> Union[int, float]:
     """
-    Probability distribution function for kroupa 2001 IMF
+    Probability distribution function for kroupa 2001 IMF, where the default values to the three_part_powerlaw are: default = {"m0": 0.1, "m1": 0.5, "m2": 1, "mmax": 100, "p1": -1.3, "p2": -2.3,"p3": -2.3}
 
-    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}
+    Args:
+        m: mass to evaluate the distribution at
+        newopts: optional dict to override the default values.
+    
+    Returns:
+        'probability' of distribution function evaluated at m
     """
 
     # Default params and override them
@@ -297,10 +403,16 @@ def Kroupa2001(m, newopts=None):
         value_dict["p3"],
     )
 
-
-def ktg93(m, newopts=None):
+def ktg93(m: Union[int, float], newopts: dict=None) -> Union[int, float]:
     """
-    Wrapper for mass distribution of KTG93
+    Probability distribution function for KTG93 IMF, where the default values to the three_part_powerlaw are: default = {"m0": 0.1, "m1": 0.5, "m2": 1, "mmax": 80, "p1": -1.3, "p2": -2.2,"p3": -2.7}
+
+    Args:
+        m: mass to evaluate the distribution at
+        newopts: optional dict to override the default values.
+    
+    Returns:
+        'probability' of distribution function evaluated at m
     """
     # TODO: ask rob what this means
 
@@ -366,32 +478,60 @@ def ktg93(m, newopts=None):
 # }
 
 
-def imf_tinsley1980(m):
+def imf_tinsley1980(m: Union[int, float]) -> Union[int, float]:
     """
-    From Tinsley 1980 (defined up until 80Msol)
+    Probability distribution function for tinsley 1980 IMF (defined up until 80Msol): three_part_powerlaw(m, 0.1, 2.0, 10.0, 80.0, -2.0, -2.3, -3.3)
+
+    Args:
+        m: mass to evaluate the distribution at
+    
+    Returns:
+        'probability' of distribution function evaluated at m
     """
 
     return three_part_powerlaw(m, 0.1, 2.0, 10.0, 80.0, -2.0, -2.3, -3.3)
 
 
-def imf_scalo1986(m):
+def imf_scalo1986(m: Union[int, float]) -> Union[int, float]:
     """
-    From Scalo 1986 (defined up until 80Msol)
+    Probability distribution function for Scalo 1986 IMF (defined up until 80Msol): three_part_powerlaw(m, 0.1, 1.0, 2.0, 80.0, -2.35, -2.35, -2.70)
+
+    Args:
+        m: mass to evaluate the distribution at
+    
+    Returns:
+        'probability' of distribution function evaluated at m
     """
     return three_part_powerlaw(m, 0.1, 1.0, 2.0, 80.0, -2.35, -2.35, -2.70)
 
 
-def imf_scalo1998(m):
+def imf_scalo1998(m: Union[int, float]) -> Union[int, float]:
     """
     From scalo 1998
+
+    Probability distribution function for Scalo 1998 IMF (defined up until 80Msol): three_part_powerlaw(m, 0.1, 1.0, 10.0, 80.0, -1.2, -2.7, -2.3)
+
+    Args:
+        m: mass to evaluate the distribution at
+    
+    Returns:
+        'probability' of distribution function evaluated at m
     """
+
     return three_part_powerlaw(m, 0.1, 1.0, 10.0, 80.0, -1.2, -2.7, -2.3)
 
 
-def imf_chabrier2003(m):
+def imf_chabrier2003(m: Union[int, float]) -> Union[int, float]:
     """
-    IMF of Chabrier 2003 PASP 115:763-795
+    Probability distribution function for IMF of Chabrier 2003 PASP 115:763-795
+
+    Args:
+        m: mass to evaluate the distribution at
+    
+    Returns:
+        'probability' of distribution function evaluated at m
     """
+
     chabrier_logmc = math.log10(0.079)
     chabrier_sigma2 = 0.69 * 0.69
     chabrier_a1 = 0.158
@@ -413,12 +553,19 @@ def imf_chabrier2003(m):
 ########################################################################
 # Binary fractions
 ########################################################################
-def Arenou2010_binary_fraction(m):
+
+def Arenou2010_binary_fraction(m: Union[int, float]) -> Union[int, float]:
     """
     Arenou 2010 function for the binary fraction as f(M1)
 
     GAIA-C2-SP-OPM-FA-054
     www.rssd.esa.int/doc_fetch.php?id=2969346
+
+    Args:
+        m: mass to evaluate the distribution at
+    
+    Returns:
+        binary fraction at m
     """
 
     return 0.8388 * math.tanh(0.688 * m + 0.079)
@@ -427,7 +574,7 @@ def Arenou2010_binary_fraction(m):
 # print(Arenou2010_binary_fraction(0.4))
 
 
-def raghavan2010_binary_fraction(m):
+def raghavan2010_binary_fraction(m: Union[int, float]) -> Union[int, float]:
     """
     Fit to the Raghavan 2010 binary fraction as a function of
     spectral type (Fig 12). Valid for local stars (Z=Zsolar).
@@ -438,6 +585,12 @@ def raghavan2010_binary_fraction(m):
     (based on Jaschek+Jaschek's Teff-spectral type table).
 
     Rob then fitted the result
+
+    Args:
+        m: mass to evaluate the distribution at
+    
+    Returns:
+        binary fraction at m
     """
 
     return min(
@@ -456,17 +609,21 @@ def raghavan2010_binary_fraction(m):
 ########################################################################
 
 
-def duquennoy1991(logper):
+def duquennoy1991(logper: Union[int, float]) -> Union[int, float]:
     """
-    Period distribution from Duquennoy + Mayor 1991
+    Period distribution from Duquennoy + Mayor 1991. Evaluated the function gaussian(logper, 4.8, 2.3, -2, 12)
 
-    Input:
-        logper: logperiod
+    Args:
+        logper: logarithm of period to evaluate the distribution at
+    
+    Returns:
+        'probability' at gaussian(logper, 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):
+
+def sana12(M1: Union[int, float], M2: Union[int, float], a: Union[int, float], P: Union[int, float], amin: Union[int, float], amax: Union[int, float], x0: Union[int, float], x1: Union[int, float], p: Union[int, float]) -> Union[int, float]:
     """
     distribution of initial orbital periods as found by Sana et al. (2012)
     which is a flat distribution in ln(a) and ln(P) respectively for stars
@@ -480,6 +637,20 @@ def sana12(M1, M2, a, P, amin, amax, x0, x1, p):
     example args: 10, 5, sep(M1, M2, P), sep, ?, -2, 12, -0.55
 
     # TODO: Fix this function!
+    
+    Args:
+        M1: Mass of primary
+        M2: Mass of secondary
+        a: separation of binary
+        P: period of binary
+        amin: minimum separation of the distribution (lower bound of the range)
+        amax: maximum separation of the distribution (upper bound of the range)
+        x0: log of minimum period of the distribution (lower bound of the range)
+        x1: log of maximum period of the distribution (upper bound of the range)
+        p: slope of the distributoon
+    
+    Returns:
+        'probability' of orbital period P given the other parameters
     """
 
     res = 0
@@ -515,8 +686,38 @@ def sana12(M1, M2, a, P, amin, amax, x0, x1, p):
 
     return res
 
+# print(sana12(10, 2, 10, 100, 1, 1000, math.log(10), math.log(1000), 6))
+
+
+def interpolate_in_mass_izzard2012(M: Union[int, float], high: Union[int, float], low: Union[int, float]) -> Union[int, float]:
+    """
+    Function to interpolate in mass
+
+    TODO: fix this function. 
+    TODO: describe the args
+    high: at M=16.3
+    low: at 1.15
+
+    Args:
+        M: mass
+        high:
+        low:
+
+    Returns:
+
+    """
+
+    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
+
 
-def Izzard2012_period_distribution(P, M1, log10Pmin=1):
+def Izzard2012_period_distribution(P: Union[int, float], M1: Union[int, float], log10Pmin: Union[int, float]=1) ->Union[int, float]:
     """
     period distribution which interpolates between
     Duquennoy and Mayor 1991 at low mass (G/K spectral type <~1.15Msun)
@@ -525,10 +726,15 @@ def Izzard2012_period_distribution(P, M1, log10Pmin=1):
     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)
+    TODO: fix this function
+
+    Args:
+        P: period
+        M1: Primary star mass
+        log10Pmin: minimum period in base log10 (optional)
+
+    Returns:
+        'probability' of interpolated distribution function at P and M1
 
     """
 
@@ -539,7 +745,6 @@ def Izzard2012_period_distribution(P, M1, log10Pmin=1):
     # 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
@@ -593,38 +798,20 @@ def Izzard2012_period_distribution(P, M1, log10Pmin=1):
             * 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))
-
 ########################################################################
 # Mass ratio distributions
 ########################################################################
 
-
-def flatsections(x, opts):
+def flatsections(x: float , opts: dict) -> Union[float, int]:
     """
     Function to generate flat distributions, possibly in multiple sections
+    
+    Args:
+        x: mass ratio value
+        opts: list containing the flat sections. Which are themselves dictionaries, with keys "max": upper bound, "min": lower bound and "height": value
 
-    opts: list of dicts with settings for the flat sections
-    x: location to calculate the y value
+    Returns: 
+        probability of that mass ratio.
     """
 
     c = 0
@@ -638,10 +825,23 @@ def flatsections(x, opts):
         if opt["min"] <= x <= opt["max"]:
             y = opt["height"]
             # print("Use this\n")
+
     c = 1.0 / c
     y = y * c
+
     # print("flatsections gives C={}: y={}\n",c,y)
     return y
 
-
 # print(flatsections(1, [{'min': 0, 'max': 2, 'height': 3}]))
+
+########################################################################
+# Eccentricity distributions
+########################################################################
+
+########################################################################
+# Star formation histories
+########################################################################
+
+########################################################################
+# Metallicity distributions
+########################################################################