From 3c9dbb109fc1d47a52b26268a69f4f932698bf24 Mon Sep 17 00:00:00 2001
From: David Hendriks <davidhendriks93@gmail.com>
Date: Wed, 29 Jan 2020 22:55:24 +0000
Subject: [PATCH] expanded the grid generation code. still needs work but its
 more complete now. need to fix the phasevol correctly tho. Made a function to
 call the grid, and loop over it and write all the binary_c calls to a file

---
 binarycpython/utils/grid.py                  | 264 ++++++++++++++-----
 binarycpython/utils/grid_options_defaults.py |   9 +-
 tests/population/grid_tests.py               |  11 +-
 3 files changed, 207 insertions(+), 77 deletions(-)

diff --git a/binarycpython/utils/grid.py b/binarycpython/utils/grid.py
index d55aae0b5..384a9c924 100644
--- a/binarycpython/utils/grid.py
+++ b/binarycpython/utils/grid.py
@@ -711,15 +711,36 @@ class Population(object):
         print("\n\nBinary_c output:")
         print(output)
 
+
+
+
+
     ###################################################
     # Unordered functions
     ###################################################
 
-    def generate_grid_code(self):
+    def increment_probtot(self, prob):
+        """
+        Function to add to the total probability
+        """
+
+        self.grid_options['probtot'] += prob
+
+    def increment_count(self):
+        """
+        Function to add to the total amount of stars
+        """
+        self.grid_options['count'] += 1
+
+    ###################################################
+    # Gridcode functions
+    ###################################################
+
+    def generate_grid_code(self, dry_run=False):
         """
         Function that generates the code from which the population will be made.
 
-        # TODO: make a generator for this.  
+        # DONE: make a generator for this.  
         # TODO: Add correct logging everywhere
         # TODO: add different types of grid. 
         # TODO: add part to handle separation if orbital_period is added
@@ -741,7 +762,7 @@ class Population(object):
         # Import packages
         code_string += "import math\n"
         code_string += "import numpy as np\n"
-        code_string += "from binarycpython.utils.probability_distributions import *\n"
+        code_string += "from binarycpython.utils.distribution_functions import *\n"
         code_string += "from binarycpython.utils.spacing_functions import *\n"
         code_string += "\n\n"
 
@@ -781,6 +802,7 @@ class Population(object):
             )
         )
         code_string += indent * depth + "parameter_dict = {}\n"
+        code_string += indent * depth + "phasevol = 0\n"        
         code_string += indent * depth + "\n"
 
         code_string += indent * depth + "# setting probability lists\n"
@@ -801,10 +823,10 @@ class Population(object):
         code_string += indent * depth + "\n"
         # Generate code
         print("Generating grid code")
-        for el in sorted(
+        for loopnr, el in enumerate(sorted(
             self.grid_options["grid_variables"].items(),
             key=lambda x: x[1]["grid_variable_number"],
-        ):
+        )):
             print("Constructing/adding: {}".format(el[0]))
             grid_variable = el[1]
 
@@ -859,7 +881,7 @@ class Population(object):
             )
 
             #########################
-            # Setting up pre-code and
+            # Setting up pre-code and value in some cases
             # Add pre-code
             code_string += (
                 indent * (depth + 1)
@@ -869,6 +891,10 @@ class Population(object):
                 + "\n"
             )
 
+
+            # Set phasevol
+            code_string += (indent * (depth + 1) + "if {} > 0: phasevol *= {}\n".format(grid_variable["name"], grid_variable["name"]))
+
             #######################
             # Probabilities
             # Calculate probability
@@ -912,6 +938,7 @@ class Population(object):
             code_string += indent * (
                 depth + 1
             ) + "# Increment starcount for {}\n".format(grid_variable["parameter_name"])
+
             code_string += (
                 indent * (depth + 1)
                 + "starcounts[{}] += 1".format(grid_variable["grid_variable_number"],)
@@ -930,71 +957,110 @@ class Population(object):
             # Add some space
             code_string += "\n"
 
+            # The final parts of the code, where things are returned, are within the deepest loop, 
+            # but in some cases code from a higher loop needs to go under it again
+            # SO I think its better to put an ifstatement here that checks whether this is the last loop. 
+            if loopnr == len(self.grid_options["grid_variables"]) -1:
+
+                #################################################################################
+                # Here are the calls to the queuing or other solution. this part is for every system
+                # Add comment
+                code_string += (indent * (depth + 1) +  "#" * 40 + "\n")
+                code_string += (
+                    indent * (depth + 1)
+                    + "# Code below will get evaluated for every generated system\n"
+                )
+
+                # Calculate value
+                code_string += (
+                    indent * (depth + 1)
+                    + 'probability = self.grid_options["weight"] * probabilities_list[{}]'.format(
+                        grid_variable["grid_variable_number"]
+                    )
+                    + "\n"
+                )
+                code_string += (
+                    indent * (depth + 1)
+                    + 'repeat_probability = probability / self.grid_options["repeat"]'
+                    + "\n"
+                )
+                code_string += indent * (depth + 1) + "total_starcount += 1\n"
+
+                # set probability and phasevol values
+                code_string += (
+                    indent * (depth + 1)
+                    + 'parameter_dict["{}"] = {}'.format(
+                        'probability', 'probability'
+                    )
+                    + "\n"
+                )
+                code_string += (
+                    indent * (depth + 1)
+                    + 'parameter_dict["{}"] = {}'.format(
+                        'phasevol', 'phasevol'
+                    )
+                    + "\n"
+                )
+
+
+                # Some prints. will be removed
+                # code_string += indent * (depth + 1) + "print(probabilities)\n"
+                # code_string += (
+                #     indent * (depth + 1) + 'print("total_starcount: ", total_starcount)\n'
+                # )
+
+                # Increment total probability
+                code_string += indent * (depth + 1) + 'self.increment_probtot(probability)\n'
+
+                if not dry_run:
+                    # Handling of what is returned, or what is not.
+                    # TODO: think of whether this is a good method
+                    code_string += indent * (depth + 1) + "yield(parameter_dict)\n"
+
+                    # The below solution might be a good one to add things to specific queues
+                    # $self->queue_evolution_code_run($self->{_flexigrid}->{thread_q},
+                    # $system);
+
+                # If its a dry run, dont do anything with it
+                else:
+                    code_string += indent * (depth + 1) + "pass\n"
+
+                code_string += (indent * (depth + 1) +  "#" * 40 + "\n")
+
             # increment depth
             depth += 1
 
-        #################################################################################
-        # Here are the calls to the queuing or other solution. this part is for every system
-        # Add comment
-        code_string += (
-            indent * (depth)
-            + "# Code below will get evaluated for every generated system\n"
-        )
 
-        # Calculate value
-        code_string += (
-            indent * (depth)
-            + 'probability = self.grid_options["weight"] * probabilities_list[{}]'.format(
-                grid_variable["grid_variable_number"]
+        depth -= 1
+        code_string += "\n"
+        # Write parts to write below the part that yield the results. this has to go in a reverse order:
+        # Here comes the stuff that is put after the deepest nested part that calls returns stuff. 
+        for loopnr, el in enumerate(sorted(
+            self.grid_options["grid_variables"].items(),
+            key=lambda x: x[1]["grid_variable_number"], reverse=True
+        )):
+            code_string += (indent * (depth + 1) +  "#" * 40 + "\n")
+            code_string += (
+                indent * (depth + 1)
+                + "# Code below is for finalising the handling of this iteration of the parameter\n"
             )
-            + "\n"
-        )
-        code_string += (
-            indent * (depth)
-            + 'repeat_probability = probability / self.grid_options["repeat"]'
-            + "\n"
-        )
-        code_string += indent * (depth) + "total_starcount += 1\n"
-        code_string += indent * (depth) + "print(probabilities)\n"
-        code_string += (
-            indent * (depth) + 'print("total_starcount: ", total_starcount)\n'
-        )
-        code_string += indent * (depth) + "yield(parameter_dict)\n"
-
-        # {
-        #     {
-        #         $self->increment_probtot($prob);
-        #         {
-        #             my $system;
-        #             if($self->{_grid_options}->{binary})
-        #             {
-        #                 $system={
-        #                 M_1=>$m1,
-        #                 M_2=>$m2,
-        #                 metallicity=>$self->metallicity(),
-        #                 orbital_period=>$per,
-        #                 eccentricity=>$eccentricity,
-        #                 probability=>$repeat_prob,
-        #                 phasevol=>$phasevol
-        #             };
-        #         }
-        #         else
-        #         {
-        #             $system={
-        #             M_1=>$m1,
-        #             M_2=>0.01,
-        #             metallicity=>$self->metallicity(),
-        #             orbital_period=>$self->{_grid_options}->{single_star_period},
-        #             eccentricity=>0.0,
-        #             probability=>$repeat_prob,
-        #             phasevol=>$phasevol
-        #         };
-        #     }
-
-        #     $self->queue_evolution_code_run($self->{_flexigrid}->{thread_q},
-        #     $system);
-
-        # }
+
+            # Set phasevol
+            # TODO: fix. this isnt supposed to be the value that we give it here. discuss
+            code_string += (indent * (depth + 1) + "if {} > 0: phasevol /= {}\n".format(grid_variable["name"], grid_variable["name"]))
+
+            depth -= 1
+
+
+        ################
+        # Finalising print statements
+        #
+        code_string += (indent * (depth + 1) +  "\n")
+        code_string += (indent * (depth + 1) +  "#" * 40 + "\n")
+        code_string += (indent * (depth + 1) + "print('Grid has handled {} stars'.format(total_starcount))\n")
+        code_string += (indent * (depth + 1) + "print('with a total probability of {}'.format(2))\n")
+
+
 
         #################################################################################
         # Stop of code generation. Here the code is saved and written
@@ -1006,7 +1072,6 @@ class Population(object):
         self.grid_options["code_string"] = code_string
 
         # Write to file
-
         gridcode_filename = os.path.join(
             self.grid_options["tmp_dir"], "example_grid.py"
         )
@@ -1018,6 +1083,18 @@ class Population(object):
         with open(gridcode_filename, "w") as f:
             f.write(code_string)
 
+    def cleanup_grid(self):
+        """
+        Function that handles all the cleaning up after the grid has been generated and/or run
+    
+        - reset values to 0
+        - remove grid file
+        - unload grid function/module
+        """
+
+        pass
+
+
     def load_grid_function(self):
         """
         Test function to run grid stuff. mostly to test the import
@@ -1041,12 +1118,59 @@ class Population(object):
         spec.loader.exec_module(grid_file)
         generator = grid_file.grid_code(self)
 
+        self.grid_options['system_generator'] = generator
+
         if self.grid_options["verbose"] > 0:
             print("Grid code loaded")
 
-        print(next(generator))
-        print(next(generator))
-        print(next(generator))
+        # for el in generator:
+        #     print(el)
+
+        # print(next(generator))
+        # print(next(generator))
+        # print(next(generator))
+
+
+    def write_binary_c_calls_to_file(self, output_dir=None, output_filename=None):
+        """
+        Function that loops over the gridcode and writes the generated parameters to a file. In the form of a commandline call
+         
+        On default this will write to the datadir, if it exists
+        """
+
+        if self.grid_options['system_generator']:
+            if self.custom_options.get('data_dir', None):
+                binary_c_calls_output_dir = self.custom_options['data_dir']
+                print('yo')
+            else:
+                if not output_dir:
+                    # if self.grid_options['verbose'] > 0:
+                    print("Error. No data_dir configured and you gave no output_dir. Aborting")
+                    raise ValueError
+                else:
+                    binary_c_calls_output_dir = output_dir
+
+            if output_filename:
+                binary_c_calls_filename = output_filename                 
+            else:
+                binary_c_calls_filename = 'binary_c_calls.txt'
+
+            print(binary_c_calls_output_dir, binary_c_calls_filename)
+
+            with open(os.path.join(binary_c_calls_output_dir, binary_c_calls_filename), 'w') as f:
+                for system in self.grid_options['system_generator']:
+                    full_system_dict = self.bse_options.copy()
+                    full_system_dict.update(system)
+
+                    binary_cmdline_string = self.return_argline(full_system_dict)
+                    f.write(binary_cmdline_string + '\n')
+        else:
+            if self.grid_options['verbose'] > 0: 
+                print("Error. No grid function found!")
+                raise KeyError
+
+
+
 
 
 ################################################################################################
diff --git a/binarycpython/utils/grid_options_defaults.py b/binarycpython/utils/grid_options_defaults.py
index 11253d260..5a747233d 100644
--- a/binarycpython/utils/grid_options_defaults.py
+++ b/binarycpython/utils/grid_options_defaults.py
@@ -20,9 +20,12 @@ grid_options_defaults_dict = {
     "log_args": 0,  #
     "log_args_dir": "/tmp/",
     # Grid variables: instructions to generate the values of the parameters
-    "grid_variables": {},
-    "grid_code": None,
-    "gridcode_filename": None,
+    "grid_variables": {}, # grid variables
+    "grid_code": None,  # literal grid code
+    "gridcode_filename": None, # filename of gridcode
+    "count": 0,  # total count of systems
+    "probtot": 0, # total probabilit
+    "system_generator": None, # value that holds the function that generates the system (result of building the grid script)
     # binary
     "binary": 0,
     # Locations:
diff --git a/tests/population/grid_tests.py b/tests/population/grid_tests.py
index 9ff0f4759..e439dc338 100644
--- a/tests/population/grid_tests.py
+++ b/tests/population/grid_tests.py
@@ -209,9 +209,9 @@ test_pop.set(
 
 ### Cleaning up custom logging code
 
-test_pop.set(C_logging_code='Printf("MY_STELLAR_DATA time=%g mass=%g radius=%g\\n", stardata->model.time, stardata->star[0].mass, stardata->star[0].radius);')
-print(test_pop.evolve_single())
-quit()
+# test_pop.set(C_logging_code='Printf("MY_STELLAR_DATA time=%g mass=%g radius=%g\\n", stardata->model.time, stardata->star[0].mass, stardata->star[0].radius);')
+# print(test_pop.evolve_single())
+# quit()
 
 
 
@@ -257,7 +257,10 @@ test_pop.add_grid_variable(
 )
 
 test_pop.generate_grid_code()
+
 test_pop.load_grid_function()
 
+test_pop.write_binary_c_calls_to_file(output_dir='/home/david/Desktop')
+
 
-print(test_pop.grid_options["code_string"])
+# print(test_pop.grid_options["code_string"])
-- 
GitLab