typed distribution_functions (af2bd4e8) · Commits · binary_c / binary_c-python

binarycpython/utils/distribution_functions.py

+330 −130

Original line number	Diff line number	Diff line
		@@ -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}

		Args:
		m: mass to evaluate the distribution at
		newopts: optional dict to override the default values.

		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}
		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

		opts: list of dicts with settings for the flat sections
		x: location to calculate the y value
		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

		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
		########################################################################