From 2ea6db910d68f13b692ea501aaf2f301b598e772 Mon Sep 17 00:00:00 2001
From: David Hendriks <davidhendriks93@gmail.com>
Date: Mon, 21 Nov 2022 13:46:37 +0000
Subject: [PATCH] putting functions back into the main population class

---
 .../distribution_functions.py                 | 535 +++++++++++-------
 .../monte_carlo_sampling.py                   |  62 ++
 2 files changed, 390 insertions(+), 207 deletions(-)

diff --git a/binarycpython/utils/population_extensions/distribution_functions.py b/binarycpython/utils/population_extensions/distribution_functions.py
index b023dcb87..429924073 100644
--- a/binarycpython/utils/population_extensions/distribution_functions.py
+++ b/binarycpython/utils/population_extensions/distribution_functions.py
@@ -569,9 +569,180 @@ class distribution_functions:
         return prob
 
     ########################################################################
-    # Binary fractions
+    # Binary / Multiplicity fractions
     ########################################################################
 
+    def get_max_multiplicity(self, multiplicity_array):
+        """
+        Function to get the maximum multiplicity. Used in the M&S calculations
+        """
+
+        max_multiplicity = 0
+        for n in range(4):
+            if multiplicity_array[n] > 0:
+                max_multiplicity = n + 1
+        return max_multiplicity
+
+    def _get_multiplicity_dict(self, system_dict):
+        """
+        Function to get the multiplicity dict according to a prescription choice
+        TODO: combine with the M&S multiplicity function: this function can be used to get the multiplicity dict, and then that function wraps this one.
+        """
+
+        # Just return 1 if no option has been chosen
+        if self.grid_options["multiplicity_fraction_function"] in [0, "None"]:
+
+            # Raise an error if the multiplicity is not set
+            if not system_dict.get("multiplicity", None):
+                msg = "Multiplicity value has not been set. When not using a specific multiplicity fraction function please set the multiplicity"
+                raise ValueError(msg)
+
+            verbose_print(
+                "_calculate_multiplicity_fraction: Chosen not to use any multiplicity fraction.",
+                self.grid_options["verbosity"],
+                3,
+            )
+
+            multiplicity_fraction_dict = {
+                1: 0,
+                2: 0,
+                3: 0,
+                4: 0,
+            }
+            multiplicity_fraction_dict[system_dict["multiplicity"]] = 1
+
+        # Go over the chosen options
+        if self.grid_options["multiplicity_fraction_function"] in [1, "Arenou2010"]:
+            # Arenou 2010 will be used
+            verbose_print(
+                "_calculate_multiplicity_fraction: Using Arenou 2010 to calculate multiplicity fractions",
+                self.grid_options["verbosity"],
+                3,
+            )
+
+            binary_fraction = self.Arenou2010_binary_fraction(system_dict["M_1"])
+            multiplicity_fraction_dict = {
+                1: 1 - binary_fraction,
+                2: binary_fraction,
+                3: 0,
+                4: 0,
+            }
+
+        elif self.grid_options["multiplicity_fraction_function"] in [2, "Raghavan2010"]:
+            # Raghavan 2010 will be used
+            verbose_print(
+                "_calculate_multiplicity_fraction: Using Raghavan (2010) to calculate multiplicity fractions",
+                self.grid_options["verbosity"],
+                3,
+            )
+
+            binary_fraction = self.raghavan2010_binary_fraction(system_dict["M_1"])
+            multiplicity_fraction_dict = {
+                1: 1 - binary_fraction,
+                2: binary_fraction,
+                3: 0,
+                4: 0,
+            }
+
+        elif self.grid_options["multiplicity_fraction_function"] in [3, "Moe2017"]:
+            # We need to check several things now here:
+
+            # First, are the options for the MOE2017 grid set? On start it is filled with the default settings
+            if not self.grid_options["Moe2017_options"]:
+                msg = "The MOE2017 options do not seem to be set properly. The value is {}".format(
+                    self.grid_options["Moe2017_options"]
+                )
+                raise ValueError(msg)
+
+            # Second: is the Moecache filled.
+            if not Moecache:
+                verbose_print(
+                    "_calculate_multiplicity_fraction: Moecache is empty. It needs to be filled with the data for the interpolators. Loading the data now",
+                    self.grid_options["verbosity"],
+                    3,
+                )
+
+                # Load the data
+                self._load_moe_di_stefano_data()
+
+            # record the prev value
+            prev_M1_value_ms = self.grid_options["Moe2017_options"].get("M_1", None)
+
+            # Set value of M1 of the current system
+            self.grid_options["Moe2017_options"]["M_1"] = system_dict["M_1"]
+
+            # Calculate the multiplicity fraction
+            multiplicity_fraction_list = (
+                self.Moe_di_Stefano_2017_multiplicity_fractions(
+                    self.grid_options["Moe2017_options"], self.grid_options["verbosity"]
+                )
+            )
+
+            # Turn into dict
+            multiplicity_fraction_dict = {
+                el + 1: multiplicity_fraction_list[el]
+                for el in range(len(multiplicity_fraction_list))
+            }
+
+            # Set the prev value back
+            self.grid_options["Moe2017_options"]["M_1"] = prev_M1_value_ms
+
+        # we don't know what to do next
+        else:
+            msg = "Chosen value for the multiplicity fraction function is not known."
+            raise ValueError(msg)
+
+        # To make sure we normalize the dictionary
+        multiplicity_fraction_dict = normalize_dict(
+            multiplicity_fraction_dict
+        )
+
+        return multiplicity_fraction_dict
+
+    def merge_multiplicities(self, result_array, max_multiplicity, verbosity=0):
+        """
+        Function to fold the multiplicities higher than the max_multiplicity onto the max_multiplicity
+
+        if max_multiplicity == 1:
+            All the multiplicities are folded onto multiplicity == 1. This will always total to 1
+        if max_multiplicity == 2:
+            The multiplicity fractions of the triple and quadruples are folded onto that of the binary multiplicity fraction
+        if max_multiplicity == 3:
+            The multiplicity fractions of the quadruples are folded onto that of the triples
+        """
+
+        if max_multiplicity not in range(1, 5):
+            msg = "\tMoe and di Stefano 2017: merge_multiplicities: max_multiplicity has to be between 1 and 4. It is {} now".format(
+                max_multiplicity
+            )
+            verbose_print(
+                msg,
+                verbosity,
+                0,
+            )
+            raise ValueError(msg)
+
+        # Fold multiplicities:
+        verbose_print(
+            "\tMoe and di Stefano 2017: merge_multiplicities: Merging multiplicities with initial array {} and max multiplicity {}".format(
+                result_array, max_multiplicity
+            ),
+            verbosity,
+            _MOE2017_VERBOSITY_LEVEL,
+        )
+        for i in range(max_multiplicity, len(result_array))[::-1]:
+            result_array[i - 1] += result_array[i]
+            result_array[i] = 0
+        verbose_print(
+            "\tMoe and di Stefano 2017: merge_multiplicities: Merging multiplicities to new array {}".format(
+                result_array
+            ),
+            verbosity,
+            _MOE2017_VERBOSITY_LEVEL,
+        )
+
+        return result_array
+
     def Arenou2010_binary_fraction(self, m: Union[int, float]) -> Union[int, float]:
         """
         Arenou 2010 function for the binary fraction as f(M1)
@@ -615,6 +786,161 @@ class distribution_functions:
             ),
         )
 
+    def Moe_di_Stefano_2017_multiplicity_fractions(self, options, verbosity=0):
+        """
+        Function that creates a list of probability fractions and
+        normalises and merges them according to the users choice.
+
+        TODO: make an extrapolation functionality in this. log10(1.6e1)
+        is quite low.
+
+        The default result that is returned when sampling the mass outside
+        of the mass range is now the last known value
+
+        Returns a list of multiplicity fractions for a given input of mass
+        """
+
+        # Use the global Moecache
+        global Moecache
+
+        multiplicity_modulator_array = np.array(
+            options["multiplicity_modulator"]
+        )  # Modulator array
+
+        # Check for length
+        if len(multiplicity_modulator_array) != 4:
+            msg = "Multiplicity modulator has to have 4 elements. Now it is {}, len: {}".format(
+                multiplicity_modulator_array, len(multiplicity_modulator_array)
+            )
+            verbose_print(
+                msg,
+                verbosity,
+                0,
+            )
+            raise ValueError(msg)
+
+        # Set up some arrays
+        full_fractions_array = np.zeros(4)  # Meant to contain the real fractions
+        weighted_fractions_array = np.zeros(
+            4
+        )  # Meant to contain the fractions multiplied by the multiplicity modulator
+
+        # Get max multiplicity
+        max_multiplicity = self.get_max_multiplicity(multiplicity_modulator_array)
+
+        # ... it's better to interpolate the multiplicity and then
+        # use a Poisson distribution to calculate the fractions
+        # (this is more accurate)
+
+        # Set up the multiplicity interpolator
+        if not Moecache.get("rinterpolator_multiplicity", None):
+            Moecache["rinterpolator_multiplicity"] = py_rinterpolate.Rinterpolate(
+                table=Moecache["multiplicity_table"],  # Contains the table of data
+                nparams=1,  # logM1
+                ndata=4,  # The number of datapoints (the parameters that we want to interpolate)
+                verbosity=verbosity - (_MOE2017_VERBOSITY_INTERPOLATOR_LEVEL - 1),
+            )
+
+        if options["multiplicity_model"] == "Poisson":
+            multiplicity = Moecache["rinterpolator_multiplicity"].interpolate(
+                [np.log10(options["M_1"])]
+            )[0]
+
+            # Fill the multiplicity array
+            for n in range(4):
+                full_fractions_array[n] = self.poisson(multiplicity, n, 3)
+
+            # Normalize it so it fills to one when taking all the multiplicities:
+            full_fractions_array = full_fractions_array / np.sum(full_fractions_array)
+
+            verbose_print(
+                "\tMoe and di Stefano 2017: Moe_di_Stefano_2017_multiplicity_fractions: using model {}: full_fractions_array: {}".format(
+                    "Poisson", full_fractions_array
+                ),
+                verbosity,
+                _MOE2017_VERBOSITY_LEVEL,
+            )
+
+        elif options["multiplicity_model"] == "data":
+            # use the fractions calculated from Moe's data directly
+            #
+            # note that in this case, there are no quadruples: these
+            # are combined with triples
+
+            # Fill with the raw values
+            for n in range(3):
+                full_fractions_array[n] = Moecache[
+                    "rinterpolator_multiplicity"
+                ].interpolate([np.log10(options["M_1"])])[n + 1]
+
+            # Set last value
+            full_fractions_array[3] = 0.0  # no quadruples
+            verbose_print(
+                "\tMoe and di Stefano 2017: Moe_di_Stefano_2017_multiplicity_fractions: using model {}: full_fractions_array: {}".format(
+                    "data", full_fractions_array
+                ),
+                verbosity,
+                _MOE2017_VERBOSITY_LEVEL,
+            )
+
+        # Normalisation:
+        if options["normalize_multiplicities"] == "raw":
+            # Don't multiply by the multiplicity_array, but do give a fractions array
+            verbose_print(
+                "\tMoe and di Stefano 2017: Moe_di_Stefano_2017_multiplicity_fractions: Not normalising (using raw results): results: {}".format(
+                    full_fractions_array
+                ),
+                verbosity,
+                _MOE2017_VERBOSITY_LEVEL,
+            )
+            result = full_fractions_array
+
+        elif options["normalize_multiplicities"] == "norm":
+            # Multiply the full_multiplicity_fraction array by the multiplicity_multiplier_array, creating a weighted fractions array
+            weighted_fractions_array = (
+                full_fractions_array * multiplicity_modulator_array
+            )
+
+            # Normalise this so it is in total 1:
+            result = weighted_fractions_array / np.sum(weighted_fractions_array)
+
+            verbose_print(
+                "\tMoe and di Stefano 2017: Moe_di_Stefano_2017_multiplicity_fractions: Normalising with {}. result: {}".format(
+                    "norm", result
+                ),
+                verbosity,
+                _MOE2017_VERBOSITY_LEVEL,
+            )
+
+        elif options["normalize_multiplicities"] == "merge":
+            # We first take the full multiplicity array
+            # (i.e. not multiplied by multiplier) and do the merging
+            result = self.merge_multiplicities(
+                full_fractions_array, max_multiplicity, verbosity=verbosity
+            )
+
+            # Then normalise to be sure
+            result = result / np.sum(result)
+
+            verbose_print(
+                "\tMoe and di Stefano 2017: Moe_di_Stefano_2017_multiplicity_fractions: Normalising with {}, max_multiplicity={} result={}".format(
+                    "merge", max_multiplicity, result
+                ),
+                verbosity,
+                _MOE2017_VERBOSITY_LEVEL,
+            )
+
+        verbose_print(
+            "\tMoe and di Stefano 2017: Moe_di_Stefano_2017_multiplicity_fractions: {}".format(
+                str(result)
+            ),
+            verbosity,
+            _MOE2017_VERBOSITY_LEVEL,
+        )
+
+        # return array reference
+        return result
+
     ########################################################################
     # Period distributions
     ########################################################################
@@ -655,7 +981,7 @@ class distribution_functions:
 
         example args: 10, 5, sep(M1, M2, P), sep, ?, -2, 12, -0.55
 
-        # TODO: Fix this function! Half of the input here can be taken out and calculated within the function itself.
+        TODO: Fix this function! Half of the input here can be taken out and calculated within the function itself.
 
         Args:
             M1: Mass of primary
@@ -900,214 +1226,9 @@ class distribution_functions:
         """
         return (lambda_val**n) * np.exp(-lambda_val) / (1.0 * math.factorial(n))
 
-    def get_max_multiplicity(self, multiplicity_array):
-        """
-        Function to get the maximum multiplicity
-        """
-        max_multiplicity = 0
-        for n in range(4):
-            if multiplicity_array[n] > 0:
-                max_multiplicity = n + 1
-        return max_multiplicity
-
-    def merge_multiplicities(self, result_array, max_multiplicity, verbosity=0):
-        """
-        Function to fold the multiplicities higher than the max_multiplicity onto the max_multiplicity
-
-        if max_multiplicity == 1:
-            All the multiplicities are folded onto multiplicity == 1. This will always total to 1
-        if max_multiplicity == 2:
-            The multiplicity fractions of the triple and quadruples are folded onto that of the binary multiplicity fraction
-        if max_multiplicity == 3:
-            The multiplicity fractions of the quadruples are folded onto that of the triples
-        """
-
-        if max_multiplicity not in range(1, 5):
-            msg = "\tMoe and di Stefano 2017: merge_multiplicities: max_multiplicity has to be between 1 and 4. It is {} now".format(
-                max_multiplicity
-            )
-            verbose_print(
-                msg,
-                verbosity,
-                0,
-            )
-            raise ValueError(msg)
-
-        # Fold multiplicities:
-        verbose_print(
-            "\tMoe and di Stefano 2017: merge_multiplicities: Merging multiplicities with initial array {} and max multiplicity {}".format(
-                result_array, max_multiplicity
-            ),
-            verbosity,
-            _MOE2017_VERBOSITY_LEVEL,
-        )
-        for i in range(max_multiplicity, len(result_array))[::-1]:
-            result_array[i - 1] += result_array[i]
-            result_array[i] = 0
-        verbose_print(
-            "\tMoe and di Stefano 2017: merge_multiplicities: Merging multiplicities to new array {}".format(
-                result_array
-            ),
-            verbosity,
-            _MOE2017_VERBOSITY_LEVEL,
-        )
-
-        return result_array
-
-    def Moe_di_Stefano_2017_multiplicity_fractions(self, options, verbosity=0):
-        """
-        Function that creates a list of probability fractions and
-        normalises and merges them according to the users choice.
-
-        TODO: make an extrapolation functionality in this. log10(1.6e1)
-        is quite low.
-
-        The default result that is returned when sampling the mass outside
-        of the mass range is now the last known value
-
-        Returns a list of multiplicity fractions for a given input of mass
-        """
-
-        # Use the global Moecache
-        global Moecache
-
-        multiplicity_modulator_array = np.array(
-            options["multiplicity_modulator"]
-        )  # Modulator array
-
-        # Check for length
-        if len(multiplicity_modulator_array) != 4:
-            msg = "Multiplicity modulator has to have 4 elements. Now it is {}, len: {}".format(
-                multiplicity_modulator_array, len(multiplicity_modulator_array)
-            )
-            verbose_print(
-                msg,
-                verbosity,
-                0,
-            )
-            raise ValueError(msg)
-
-        # Set up some arrays
-        full_fractions_array = np.zeros(4)  # Meant to contain the real fractions
-        weighted_fractions_array = np.zeros(
-            4
-        )  # Meant to contain the fractions multiplied by the multiplicity modulator
-
-        # Get max multiplicity
-        max_multiplicity = self.get_max_multiplicity(multiplicity_modulator_array)
-
-        # ... it's better to interpolate the multiplicity and then
-        # use a Poisson distribution to calculate the fractions
-        # (this is more accurate)
-
-        # Set up the multiplicity interpolator
-        if not Moecache.get("rinterpolator_multiplicity", None):
-            Moecache["rinterpolator_multiplicity"] = py_rinterpolate.Rinterpolate(
-                table=Moecache["multiplicity_table"],  # Contains the table of data
-                nparams=1,  # logM1
-                ndata=4,  # The number of datapoints (the parameters that we want to interpolate)
-                verbosity=verbosity - (_MOE2017_VERBOSITY_INTERPOLATOR_LEVEL - 1),
-            )
-
-        if options["multiplicity_model"] == "Poisson":
-            multiplicity = Moecache["rinterpolator_multiplicity"].interpolate(
-                [np.log10(options["M_1"])]
-            )[0]
-
-            # Fill the multiplicity array
-            for n in range(4):
-                full_fractions_array[n] = self.poisson(multiplicity, n, 3)
-
-            # Normalize it so it fills to one when taking all the multiplicities:
-            full_fractions_array = full_fractions_array / np.sum(full_fractions_array)
-
-            verbose_print(
-                "\tMoe and di Stefano 2017: Moe_di_Stefano_2017_multiplicity_fractions: using model {}: full_fractions_array: {}".format(
-                    "Poisson", full_fractions_array
-                ),
-                verbosity,
-                _MOE2017_VERBOSITY_LEVEL,
-            )
-
-        elif options["multiplicity_model"] == "data":
-            # use the fractions calculated from Moe's data directly
-            #
-            # note that in this case, there are no quadruples: these
-            # are combined with triples
-
-            # Fill with the raw values
-            for n in range(3):
-                full_fractions_array[n] = Moecache[
-                    "rinterpolator_multiplicity"
-                ].interpolate([np.log10(options["M_1"])])[n + 1]
-
-            # Set last value
-            full_fractions_array[3] = 0.0  # no quadruples
-            verbose_print(
-                "\tMoe and di Stefano 2017: Moe_di_Stefano_2017_multiplicity_fractions: using model {}: full_fractions_array: {}".format(
-                    "data", full_fractions_array
-                ),
-                verbosity,
-                _MOE2017_VERBOSITY_LEVEL,
-            )
-
-        # Normalisation:
-        if options["normalize_multiplicities"] == "raw":
-            # Don't multiply by the multiplicity_array, but do give a fractions array
-            verbose_print(
-                "\tMoe and di Stefano 2017: Moe_di_Stefano_2017_multiplicity_fractions: Not normalising (using raw results): results: {}".format(
-                    full_fractions_array
-                ),
-                verbosity,
-                _MOE2017_VERBOSITY_LEVEL,
-            )
-            result = full_fractions_array
 
-        elif options["normalize_multiplicities"] == "norm":
-            # Multiply the full_multiplicity_fraction array by the multiplicity_multiplier_array, creating a weighted fractions array
-            weighted_fractions_array = (
-                full_fractions_array * multiplicity_modulator_array
-            )
-
-            # Normalise this so it is in total 1:
-            result = weighted_fractions_array / np.sum(weighted_fractions_array)
 
-            verbose_print(
-                "\tMoe and di Stefano 2017: Moe_di_Stefano_2017_multiplicity_fractions: Normalising with {}. result: {}".format(
-                    "norm", result
-                ),
-                verbosity,
-                _MOE2017_VERBOSITY_LEVEL,
-            )
 
-        elif options["normalize_multiplicities"] == "merge":
-            # We first take the full multiplicity array
-            # (i.e. not multiplied by multiplier) and do the merging
-            result = self.merge_multiplicities(
-                full_fractions_array, max_multiplicity, verbosity=verbosity
-            )
-
-            # Then normalise to be sure
-            result = result / np.sum(result)
-
-            verbose_print(
-                "\tMoe and di Stefano 2017: Moe_di_Stefano_2017_multiplicity_fractions: Normalising with {}, max_multiplicity={} result={}".format(
-                    "merge", max_multiplicity, result
-                ),
-                verbosity,
-                _MOE2017_VERBOSITY_LEVEL,
-            )
-
-        verbose_print(
-            "\tMoe and di Stefano 2017: Moe_di_Stefano_2017_multiplicity_fractions: {}".format(
-                str(result)
-            ),
-            verbosity,
-            _MOE2017_VERBOSITY_LEVEL,
-        )
-
-        # return array reference
-        return result
 
     def build_q_table(self, options, m, p, verbosity=0):
         """
diff --git a/binarycpython/utils/population_extensions/monte_carlo_sampling.py b/binarycpython/utils/population_extensions/monte_carlo_sampling.py
index 7a91275af..2fa22da09 100644
--- a/binarycpython/utils/population_extensions/monte_carlo_sampling.py
+++ b/binarycpython/utils/population_extensions/monte_carlo_sampling.py
@@ -253,3 +253,65 @@ class monte_carlo_sampling:
 
         # raise NotImplementedError("Functionality in not implemented")
         pass
+
+
+
+
+
+    ##############
+    # Writing to file:
+    def monte_carlo_sampling_write_binary_c_calls_to_file(
+        self,
+        output_file,
+        include_defaults
+    ):
+        """
+        Function to write the generated grid to a file.
+        TODO: this can be merged with the original write_binary_c_calls to file ... Have the get_generator handle the setup
+        """
+
+        #######
+        # Start up the generator
+        generator = self._get_generator()
+
+        if not self.grid_options["monte_carlo_mass_threshold"] > 0:
+            raise ValueError("Cant write MC sampling to file without a mass threshold")
+
+        # then if the _system_generator is present, we go through it
+        if generator:
+            # Write to file
+            with self.open(
+                output_file, "w", encoding="utf-8"
+            ) as file:
+                # Get defaults and clean them, then overwrite them with the set values.
+                if include_defaults:
+                    # TODO: make sure that the defaults here are cleaned up properly
+                    cleaned_up_defaults = self.cleaned_up_defaults
+                    full_system_dict = cleaned_up_defaults.copy()
+                    full_system_dict.update(self.bse_options.copy())
+                else:
+                    full_system_dict = self.bse_options.copy()
+
+                for system_number, system_dict in enumerate(generator):
+                    # update values with current system values
+                    full_system_dict.update(system_dict)
+
+                    ######
+                    # Handle monte-carlo threshold based on evolved mass
+                    if self.grid_options["evolution_type"] == "monte_carlo":
+                        # Check based on mass threshold
+                        self._monte_carlo_sampling_check_mass_threshold(full_system_dict)
+
+                        ######
+                        # Check if evolution threshold is reached.
+                        if self.grid_options["_monte_carlo_threshold_reached"]:
+                            break
+
+                    #
+                    binary_cmdline_string = self._return_argline(full_system_dict)
+                    file.write(binary_cmdline_string + "\n")
+        else:
+            print("Error. No grid function found!")
+            raise ValueError
+
+        return output_file
\ No newline at end of file
-- 
GitLab