From 46ff561e305dd09d8369cb3c3453d90b601a77f1 Mon Sep 17 00:00:00 2001
From: Robert Izzard <r.izzard@surrey.ac.uk>
Date: Sat, 27 Nov 2021 09:26:19 +0000
Subject: [PATCH] added option for dry_run_hook
---
binarycpython/utils/grid_options_defaults.py | 2 ++
binarycpython/utils/gridcode.py | 2 ++
binarycpython/utils/spacing_functions.py | 14 ++++++++------
3 files changed, 12 insertions(+), 6 deletions(-)
diff --git a/binarycpython/utils/grid_options_defaults.py b/binarycpython/utils/grid_options_defaults.py
index cb903637e..9b242f985 100644
--- a/binarycpython/utils/grid_options_defaults.py
+++ b/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
diff --git a/binarycpython/utils/gridcode.py b/binarycpython/utils/gridcode.py
index c6afbc3b3..2e14da981 100644
--- a/binarycpython/utils/gridcode.py
+++ b/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")
diff --git a/binarycpython/utils/spacing_functions.py b/binarycpython/utils/spacing_functions.py
index d18f25253..e07fcdeb7 100644
--- a/binarycpython/utils/spacing_functions.py
+++ b/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
--
GitLab