added option for dry_run_hook

binarycpython/utils/grid_options_defaults.py

@@ -55,6 +55,7 @@ class grid_options_defaults():
             "_zero_prob_stars_skipped": 0,
             "ensemble_factor_in_probability_weighted_mass": False,  # Whether to multiply the ensemble results by 1/probability_weighted_mass
             "do_dry_run": True,  # Whether to do a dry run to calculate the total probability for this run
+            "dry_run_hook" : None, # Function hook for the dry run: this function is called, if not None, for every star in the dry run. Useful for checking initial distributions.
             "custom_generator": None,  # Place for the custom system generator
             "exit_after_dry_run": False,  # Exit after dry run?
             "print_stack_on_exit" : False, # print the stack trace on exit calls?
@@ -392,6 +393,7 @@ class grid_options_defaults():
             "m&s_options": "Internal variable that holds the Moe and di Stefano (2017) options. Don't write to this your self",
             "_loaded_Moe2017_data": "Internal variable storing whether the Moe and di Stefano (2017) data has been loaded into memory",
             "do_dry_run": "Whether to do a dry run to calculate the total probability for this run",
+            "dry_run_hook" : "Function hook to be called for every system in a dry run. The function is passed a dict of the system parameters. Does nothing if None (the default).",
             "exit_after_dry_run": "If True, exits after a dry run. Default is False.",
             "print_stack_on_exit" : "If True, prints a stack trace when the population's exit method is called.",
             "_Moe2017_JSON_data": "Location to store the loaded Moe&diStefano2017 dataset",  # Stores the data

binarycpython/utils/gridcode.py

@@ -701,6 +701,8 @@ class gridcode():
 
         # If its a dry run, dont do anything with it
         else:
+            if self.grid_options['dry_run_hook']:
+                self._add_code("self.grid_options['dry_run_hook'](parameter_dict)\n",indent=1)
             self._add_code("pass\n", indent=1)
 
         self._add_code("#" * 40 + "\n")

binarycpython/utils/spacing_functions.py

@@ -5,16 +5,16 @@ Tasks:
     TODO: add more spacing functions to this module.
 """
 
-from typing import Union
-from diskcache import Cache
 import cachetools
 import datetime
+import diskcache
 import functools
 import json
 import math
 import numpy as np
 import py_rinterpolate
 import sys
+from typing import Union
 
 class spacing_functions():
 
@@ -225,8 +225,9 @@ class spacing_functions():
         print("Cache dir {}".format(self.grid_options['cache_dir']))
         if cachedir == None:
             cachedir = self.grid_options['cache_dir'] + '/const_dt_cache'
-        cache = Cache(cachedir)
-
+            cache = diskcache.Cache(cachedir)
+        else:
+            cache = None
         def _const_dt_wrapper(cachedir=None,
                               num_cores=None,
                               bse_options=None,
@@ -261,7 +262,7 @@ class spacing_functions():
                     del bse_stripped[x]
                 except:
                     pass
-                
+
             # make a JSON string of the options (this can be
             # used to check the cache)
             bse_options_json = json.dumps(bse_stripped,
@@ -289,7 +290,8 @@ class spacing_functions():
                              showtable=showtable,
                              usecache=usecache)
 
-        @cache.memoize() # memoize to disc
+        if cache:
+            eval('@cache.memoize()') # memoize to disc
         def _const_dt(cachedir=None,
                       num_cores=None,
                       bse_options_json=None, # JSON string