From 9d2ea2c765b504effd53a6289f90d1d436b0a4bd Mon Sep 17 00:00:00 2001
From: David Hendriks <davidhendriks93@gmail.com>
Date: Tue, 8 Jun 2021 14:50:06 +0100
Subject: [PATCH] Working on the multiplicity_fraction calculation

---
 binarycpython/utils/distribution_functions.py | 131 +++++++++---------
 1 file changed, 68 insertions(+), 63 deletions(-)

diff --git a/binarycpython/utils/distribution_functions.py b/binarycpython/utils/distribution_functions.py
index 93003faf4..f37efc8ca 100644
--- a/binarycpython/utils/distribution_functions.py
+++ b/binarycpython/utils/distribution_functions.py
@@ -1017,6 +1017,55 @@ def get_max_multiplicity(multiplicity_array):
             max_multiplicity = n + 1
     return max_multiplicity
 
+def merge_multiplicities(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 not max_multiplicity in range(1, 5):
+        msg = "\tM&S: 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(
+        "\tM&S: merge_multiplicities: Merging multiplicities with initial array {} and max multiplicity {}".format(result_array, max_multiplicity),
+        verbosity,
+        0,
+    )
+    for i in range(max_multiplicity, len(result_array))[::-1]:
+        result_array[i-1] += result_array[i]
+        result_array[i] = 0
+    verbose_print(
+        "\tM&S: merge_multiplicities: Merging multiplicities to new array {}".format(result_array),
+        verbosity,
+        0,
+    )
+
+    return result_array
+
+def normalize_dict(result_dict, verbosity=0):
+    """
+    Function to normalize a dictionary
+    """
+
+    sum_result = sum([result_dict[key] for key in result_dict.keys()])
+    for key in result_dict.keys():
+        result_dict[key] = result_dict[key] / sum_result
+    return result_dict
 
 def Moe_de_Stefano_2017_multiplicity_fractions(options, verbosity=0):
     """
@@ -1027,31 +1076,20 @@ def Moe_de_Stefano_2017_multiplicity_fractions(options, verbosity=0):
 
     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 list of masses
+    Returns a list of multiplicity fractions for a given input of mass
+    """
 
     # Use the global Moecache
     global Moecache
 
     result = {}
-
-    # if(0)
-    # {
-    #     # we have these in a table, so could interpolate, but...
-    #     $Moecache{'rinterpolator_fractions_M1_table'} //=
-    #         rinterpolate->new(
-    #             table => $Moecache{'fractions_M1_table'},
-    #             nparams => 1, # logM1
-    #             ndata => 4, # single/binary/triple/quadruple fractions
-    #         );
-
-    #     $r = $Moecache{'rinterpolator_fractions_M1_table'}->interpolate([log10($opts->{'M1'})]);
-    # }
+    new_result = {}
 
     # ... 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(
@@ -1067,12 +1105,21 @@ def Moe_de_Stefano_2017_multiplicity_fractions(options, verbosity=0):
         )[0]
 
         for n in range(4):
+            new_result[n] = poisson(multiplicity, n, 3, verbosity)
             result[n] = (
                 options["multiplicity_modulator"][n]
                 * poisson(multiplicity, n, 3, verbosity)
                 if options["multiplicity_modulator"][n] > 0
                 else 0
             )
+        print("For mass M1={}".format(options['M1']))
+        print("Multiplicty modulator: {}".format(options["multiplicity_modulator"]))
+        print("Old results: {}".format(result))
+
+        # Normalize dict
+        new_result = normalize_dict(new_result, verbosity=verbosity)
+        sum_result = sum([new_result[key] for key in new_result])
+        print("New result: {} (sum = {})".format(new_result, sum_result))
 
         verbose_print(
             "\tM&S: Moe_de_Stefano_2017_multiplicity_fractions: using model {}".format("Poisson"),
@@ -1115,9 +1162,7 @@ def Moe_de_Stefano_2017_multiplicity_fractions(options, verbosity=0):
 
     elif options["normalize_multiplicities"] == "norm":
         # simply normalize so all multiplicities add to 1
-        sum_result = sum([result[key] for key in result.keys()])
-        for key in result.keys():
-            result[key] = result[key] / sum_result
+        result = normalize_dict(result, verbosity=verbosity)
 
         verbose_print(
             "\tM&S: Moe_de_Stefano_2017_multiplicity_fractions: Normalizing with {}".format("norm"),
@@ -1126,55 +1171,15 @@ def Moe_de_Stefano_2017_multiplicity_fractions(options, verbosity=0):
         )
 
     elif options["normalize_multiplicities"] == "merge":
-        # if multiplicity == 1, merge binaries, triples and quads into singles
-        #    (de facto this is the same as "norm")
-        # if multiplicity == 2, merge triples and quads into binaries
-        # if multiplicity == 3, merge quads into triples
-        # if multiplicity == 4, do nothing, equivalent to 'raw'.
-
-        # TODO: ask rob about this part. Not sure if i understand it.
-
         max_multiplicity = get_max_multiplicity(options["multiplicity_modulator"])
 
-        if max_multiplicity == 1:
-            result[0] = 1.0
-            for n in range(1, 4):
-                result[0] += result[n]
-                result[n] = 0
-
-        elif max_multiplicity == 2:
-            # we have only singles and binaries:
-            # triples and quads are treated as binaries
-
-            for n in [2, 3]:
-                result[1] += result[n]
-                result[n] = 0
-
-        elif max_multiplicity == 3:
-            # we have singles, binaries and triples:
-            # quads are treated as triples
-            result[2] += result[3]
-            result[3] = 0
-
-        elif max_multiplicity == 4:
-            # we have singles, binaries, triples and quads,
-            # so do nothing
-            pass
-        else:
-            msg = "\tM&S: Error: in the Moe distribution, we seem to want no stars at all (multiplicity == {})".format(
-                max_multiplicity
-            )
-            verbose_print(
-                msg,
-                verbosity,
-                0,
-            )
-            raise ValueError(msg)
+        # Merge the multiplicities
+        result = merge_multiplicities(result, max_multiplicity, verbosity=verbosity)
+
+        new_result = merge_multiplicities(new_result, max_multiplicity, verbosity=verbosity)
 
         # Normalising results again
-        sum_result = sum([result[key] for key in result.keys()])
-        for key in result.keys():
-            result[key] = result[key] / sum_result
+        result = normalize_dict(result, verbosity=verbosity)
 
         verbose_print(
             "\tM&S: Moe_de_Stefano_2017_multiplicity_fractions: Normalizing with {}, max_multiplicity={}".format("merge", max_multiplicity),
-- 
GitLab