From 6ce11ff5604fcc99d2da2b15d7cc500c9f71c1ef Mon Sep 17 00:00:00 2001
From: Robert Izzard <r.izzard@surrey.ac.uk>
Date: Wed, 13 Oct 2021 20:01:10 +0100
Subject: [PATCH] add multiplicity functions by string as well as integer

---
 binarycpython/utils/grid.py                  | 38 ++++++++++++++++----
 binarycpython/utils/grid_options_defaults.py |  1 +
 2 files changed, 33 insertions(+), 6 deletions(-)

diff --git a/binarycpython/utils/grid.py b/binarycpython/utils/grid.py
index a28022334..00ce11cbe 100644
--- a/binarycpython/utils/grid.py
+++ b/binarycpython/utils/grid.py
@@ -484,6 +484,7 @@ class Population:
             parameter_name: str,
             gridtype: str = "centred",
             branchpoint: int = 0,
+            branchcode: Union[str, None] = None,
             precode: Union[str, None] = None,
             postcode: Union[str, None] = None,
             condition: Union[str, None] = None,
@@ -577,6 +578,7 @@ class Population:
             "condition": condition,
             "gridtype": gridtype,
             "branchpoint": branchpoint,
+            "branchcode": branchcode,
             "grid_variable_number": len(self.grid_options["_grid_variables"]),
         }
 
@@ -2562,6 +2564,7 @@ class Population:
                     grid_variable,
                     dry_run,
                     grid_variable["branchpoint"],
+                    grid_variable["branchcode"],
                 )
 
             # increment indent_depth
@@ -2620,6 +2623,7 @@ class Population:
                     reverse_sorted_grid_variables[loopnr + 1][1],
                     dry_run,
                     grid_variable["branchpoint"],
+                    grid_variable["branchcode"],
                 )
                 self._increment_indent_depth(-1)
                 self._add_code("\n")
@@ -2670,13 +2674,35 @@ class Population:
         with open(gridcode_filename, "w") as file:
             file.write(self.code_string)
 
-    def _write_gridcode_system_call(self, grid_variable, dry_run, branchpoint):
+        # perhaps create symlink to latest gridcode: note, we use os.replace()
+        # in case the symlink already exists
+        if self.grid_options['symlink latest gridcode']:
+            ftmp = os.path.join(self.grid_options['tmp_dir'],
+                                'binary_c_grid-latest-tmp-symlink')
+            fsym = os.path.join(self.grid_options['tmp_dir'],
+                                'binary_c_grid-latest')
+            os.symlink(gridcode_filename,ftmp)
+            os.replace(ftmp,'binary_c_grid-latest-tmp')
+            verbose_print(
+                "{}Symlinked grid code to {} {}".format(self.ANSI_colours['blue'],
+                                                        ftmp,
+                                                        self.ANSI_colours['reset']),
+                self.grid_options["verbosity"],
+                1,
+            )
+
+    def _write_gridcode_system_call(self, grid_variable, dry_run, branchpoint, branchcode):
         #################################################################################
         # Here are the calls to the queuing or other solution. this part is for every system
         # Add comment
         self._increment_indent_depth(+1)
         self._add_code("#" * 40 + "\n")
 
+        if branchcode:
+            self._add_code(
+                "# Branch code\nif {}:\n",
+                branchcode)
+
         if branchpoint:
             self._add_code(
                 "# Code below will get evaluated for every system at this level of multiplicity (last one of that being {})\n".format(
@@ -4552,7 +4578,7 @@ eccentricity3=0
         """
 
         # Just return 1 if no option has been chosen
-        if self.grid_options["multiplicity_fraction_function"] == 0:
+        if self.grid_options["multiplicity_fraction_function"] in [0, "None"]:
             verbose_print(
                 "_calculate_multiplicity_fraction: Chosen not to use any multiplicity fraction.",
                 self.grid_options["verbosity"],
@@ -4567,7 +4593,7 @@ eccentricity3=0
             raise ValueError(msg)
 
         # Go over the chosen options
-        if self.grid_options["multiplicity_fraction_function"] == 1:
+        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",
@@ -4583,10 +4609,10 @@ eccentricity3=0
                 4: 0,
             }
 
-        elif self.grid_options["multiplicity_fraction_function"] == 2:
+        elif self.grid_options["multiplicity_fraction_function"] in [2, 'Raghavan2010']:
             # Raghavan 2010 will be used
             verbose_print(
-                "_calculate_multiplicity_fraction: Using Rhagavan 2010 to calculate multiplicity fractions",
+                "_calculate_multiplicity_fraction: Using Raghavan (2010) to calculate multiplicity fractions",
                 self.grid_options["verbosity"],
                 3,
             )
@@ -4599,7 +4625,7 @@ eccentricity3=0
                 4: 0,
             }
 
-        elif self.grid_options["multiplicity_fraction_function"] == 3:
+        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
diff --git a/binarycpython/utils/grid_options_defaults.py b/binarycpython/utils/grid_options_defaults.py
index 44bb9ade4..a1d05e9dd 100644
--- a/binarycpython/utils/grid_options_defaults.py
+++ b/binarycpython/utils/grid_options_defaults.py
@@ -123,6 +123,7 @@ grid_options_defaults_dict = {
     ## Grid type evolution
     "_grid_variables": {},  # grid variables
     "gridcode_filename": None,  # filename of gridcode
+    "symlink latest gridcode":True, # symlink to latest gridcode
     ## Monte carlo type evolution
     # TODO: make MC options
     ## Evolution from source file
-- 
GitLab