diff --git a/binarycpython/utils/grid.py b/binarycpython/utils/grid.py
index fb24d0661b9496b7f611be79e72956f1e8c30371..f762a75fb26bd3db6ae93c0f29f8db1b94dce4b8 100644
--- a/binarycpython/utils/grid.py
+++ b/binarycpython/utils/grid.py
@@ -685,6 +685,130 @@ class Population(object):
# Testing functions
###################################################
+ def test_evolve_population(self):
+ """
+ The to test running a population
+ """
+
+ import time
+ import multiprocessing as mp
+ from pathos.multiprocessing import ProcessingPool as Pool
+ import random
+
+ #######################
+ ### Custom logging code:
+ self.set_custom_logging()
+
+ ### Load store
+ self.grid_options["store_memaddr"] = binary_c_python_api.return_store("")
+
+
+ #######################
+ # Dry run and getting starcount
+ self.generate_grid_code(dry_run=True)
+
+ self.load_grid_function()
+
+ self.dry_run()
+
+ print(self.grid_options['total_starcount'])
+
+
+ #######################
+ # Linear run
+ self.generate_grid_code(dry_run=False)
+
+ self.load_grid_function()
+
+ for i, system in enumerate(self.grid_options["system_generator"](self)):
+ full_system_dict = self.bse_options.copy()
+ full_system_dict.update(system)
+
+ binary_cmdline_string = self.return_argline(full_system_dict)
+ print(i+1, binary_cmdline_string)
+
+ #######################
+ # MP run
+ start_mp = time.time()
+
+ self.generate_grid_code(dry_run=False)
+
+ self.load_grid_function()
+
+ def evolve_system(system):
+ full_system_dict = self.bse_options.copy()
+ full_system_dict.update(system)
+
+ binary_cmdline_string = self.return_argline(full_system_dict)
+ print(binary_cmdline_string)
+
+ # self.set_bse_option("M_1", mass)
+ # out = binary_c_python_api.run_population(
+ # self.return_argline(),
+ # self.grid_options["custom_logging_func_memaddr"],
+ # self.grid_options["store_memaddr"],
+ # )
+ # # parse_function(self, out)
+
+ def yield_system():
+ for i, system in enumerate(self.grid_options["system_generator"](self)):
+ print(i+1)
+ yield system
+ print("generator done")
+
+ # Create pool
+ p = Pool(nodes=self.grid_options["amt_cores"])
+
+ # Execute
+ r = list(p.imap(evolve_system, yield_system()))
+
+ stop_mp = time.time()
+
+ # Give feedback
+
+ print(
+ "with mp: {} systems took {}s using {} cores".format(
+ self.grid_options['total_starcount'],
+ stop_mp - start_mp,
+ self.grid_options["amt_cores"],
+ )
+ )
+
+ quit()
+ ########################
+
+ # evolve with mp
+ start_mp = time.time()
+
+ def evolve_mp(mass):
+ self.set_bse_option("M_1", mass)
+ out = binary_c_python_api.run_population(
+ self.return_argline(),
+ self.grid_options["custom_logging_func_memaddr"],
+ self.grid_options["store_memaddr"],
+ )
+ parse_function(self, out)
+
+ p = Pool(nodes=self.grid_options["amt_cores"])
+
+ def g(mass_distribution):
+ masses = mass_distribution
+ for mass in masses:
+ yield mass
+ print("generator done")
+
+ r = list(p.imap(evolve_mp, g(mass_distribution)))
+ stop_mp = time.time()
+
+ print(
+ "with mp: {} systems took {}s using {} cores".format(
+ len(mass_distribution),
+ stop_mp - start_mp,
+ self.grid_options["amt_cores"],
+ )
+ )
+
+
def test_evolve_single(self):
"""
Function to test the evolution of a system. Calls the api binding directly.
@@ -740,6 +864,8 @@ class Population(object):
"""
Function that generates the code from which the population will be made.
+ dry_run: when True, it will return the starcount at the end so that we know what the total amount of systems is.
+
The phasevol values are handled by generating a second array
@@ -1076,6 +1202,12 @@ class Population(object):
indent * (depth + 1) + "print('with a total probability of {}'.format(2))\n"
)
+ if dry_run:
+ code_string += (
+ indent * (depth + 1) + "return total_starcount\n"
+ )
+
+
#################################################################################
# Stop of code generation. Here the code is saved and written
@@ -1104,6 +1236,8 @@ class Population(object):
- reset values to 0
- remove grid file
- unload grid function/module
+ - remove dry grid file
+ - unload dry grid function/module
"""
pass
@@ -1125,23 +1259,28 @@ class Population(object):
spec = importlib.util.spec_from_file_location(
"binary_c_python_grid",
- os.path.join(self.grid_options["tmp_dir"], "example_grid.py"),
+ os.path.join(self.grid_options["gridcode_filename"]),
)
grid_file = importlib.util.module_from_spec(spec)
spec.loader.exec_module(grid_file)
- generator = grid_file.grid_code(self)
+ generator = grid_file.grid_code
self.grid_options["system_generator"] = generator
if self.grid_options["verbose"] > 0:
print("Grid code loaded")
- # for el in generator:
- # print(el)
+ def dry_run(self):
+ """
+ Function to dry run the grid and know how many stars it will run
+
+ Requires the grid to be built as a dry run grid
+ """
+
+ system_generator = self.grid_options["system_generator"]
+ total_starcount = system_generator(self)
+ self.grid_options['total_starcount'] = total_starcount
- # print(next(generator))
- # print(next(generator))
- # print(next(generator))
def write_binary_c_calls_to_file(self, output_dir=None, output_filename=None):
"""
diff --git a/tests/population/grid_tests.py b/tests/population/grid_tests.py
index 986fdca4dd0d51f3dec4cafd821d49dc37f23a6f..a0f4e976138fabddc82c211f8eddc688727b5304 100644
--- a/tests/population/grid_tests.py
+++ b/tests/population/grid_tests.py
@@ -213,6 +213,84 @@ test_pop.set(
### Grid generating testing
+# test_pop.add_grid_variable(
+# name="lnm1",
+# longname="log primary mass",
+# valuerange=[10, 20],
+# resolution="10",
+# spacingfunc="np.linspace(0.213, 10.2, 10)",
+# precode="M_1=math.exp(lnm1)",
+# probdist="flat(M_1)",
+# # probdist='self.custom_options["extra_prob_function"](M_1)',
+# dphasevol="",
+# parameter_name="M_1",
+# condition="",
+# )
+
+# ### Grid generating test.
+# test_pop.add_grid_variable(
+# name="period",
+# longname="period",
+# valuerange=["M_1", 20],
+# resolution="10",
+# spacingfunc="np.linspace(1, 10, 10)",
+# precode="orbital_period = period**2",
+# probdist="flat(orbital_period)",
+# parameter_name="orbital_period",
+# dphasevol="",
+# condition='self.grid_options["binary"]==0',
+# )
+
+# 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"])
+
+####
+# Dry run:
+# test_pop.set(verbose=1)
+# test_pop.add_grid_variable(
+# name="lnm1",
+# longname="log primary mass",
+# valuerange=[10, 20],
+# resolution="10",
+# spacingfunc="np.linspace(0.213, 10.2, 10)",
+# precode="M_1=math.exp(lnm1)",
+# probdist="flat(M_1)",
+# # probdist='self.custom_options["extra_prob_function"](M_1)',
+# dphasevol="",
+# parameter_name="M_1",
+# condition="",
+# )
+
+# ### Grid generating test.
+# test_pop.add_grid_variable(
+# name="period",
+# longname="period",
+# valuerange=["M_1", 20],
+# resolution="10",
+# spacingfunc="np.linspace(1, 10, 10)",
+# precode="orbital_period = period**2",
+# probdist="flat(orbital_period)",
+# parameter_name="orbital_period",
+# dphasevol="",
+# condition='self.grid_options["binary"]==0',
+# )
+
+# test_pop.generate_grid_code(dry_run=True)
+
+# test_pop.load_grid_function()
+
+# test_pop.dry_run()
+
+####
+# testing population:
+
+test_pop.set(verbose=1, amt_cores=4)
test_pop.add_grid_variable(
name="lnm1",
longname="log primary mass",
@@ -241,11 +319,5 @@ test_pop.add_grid_variable(
condition='self.grid_options["binary"]==0',
)
-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"])
+test_pop.test_evolve_population()
+# test_pop.generate_grid_code(dry_run=True)