From a6bfbffe6457a692eff179b71e5f99a8cd407639 Mon Sep 17 00:00:00 2001
From: Robert Izzard <r.izzard@surrey.ac.uk>
Date: Thu, 7 Oct 2021 12:34:46 +0100
Subject: [PATCH] add gaussian-zoomed spacing function, useful for zooming into
 part of the parameter sace

add option such that when resolution==0, the resolution is calculated automatically from the sample list.

add delete_grid_variable function

In my CEMP simulations, I use the above with something like this to zoom in around 0.8Msun in M1:

emp_population.update_grid_variable(
    name="lnm1",
    spacingfunc="gaussian_zoom(np.log(0.4),np.log(6.0),np.log(0.8),0.05,0.99,100)",
    resolution=0
)
---
 binarycpython/utils/grid.py                  | 63 +++++++++++++----
 binarycpython/utils/grid_options_defaults.py |  2 +
 binarycpython/utils/spacing_functions.py     | 72 +++++++++++++++++++-
 3 files changed, 122 insertions(+), 15 deletions(-)

diff --git a/binarycpython/utils/grid.py b/binarycpython/utils/grid.py
index 9cbb6872a..73e73a909 100644
--- a/binarycpython/utils/grid.py
+++ b/binarycpython/utils/grid.py
@@ -244,7 +244,7 @@ class Population:
                         key, kwargs[key]
                     ),
                     self.grid_options["verbosity"],
-                    0, # NOTE: setting this to be 0 prevents mistakes being overlooked. 
+                    0, # NOTE: setting this to be 0 prevents mistakes being overlooked.
                 )
                 self.custom_options[key] = kwargs[key]
 
@@ -401,6 +401,22 @@ class Population:
                 1,
             )
 
+    def delete_grid_variable(
+        self,
+        name: str,
+    ) -> None:
+        try:
+            del self.grid_options["_grid_variables"][name]
+            verbose_print(
+                "Deleted grid variable: {}".format(name),
+                self.grid_options["verbosity"],
+                1,
+            )
+        except:
+            msg = "Failed to remove grid variable {} : please check it exists.".format(name)
+            raise ValueError(msg)
+
+
     def rename_grid_variable(
         self,
         oldname: str,
@@ -425,6 +441,11 @@ class Population:
         try:
             self.grid_options["_grid_variables"][newname] = self.grid_options["_grid_variables"].pop(oldname)
             self.grid_options["_grid_variables"][newname]["name"] = newname
+            verbose_print(
+                "Rename grid variable: {} to {}".format(oldname,newname),
+                self.grid_options["verbosity"],
+                1,
+            )
         except:
             msg = "Failed to rename grid variable {} to {}.".format(oldname,newname)
             raise ValueError(msg)
@@ -471,7 +492,7 @@ class Population:
                 Examples:
                     longname = 'Primary mass'
             range:
-                Range of values to take. Does not get used really, the spacingfunction is used to
+                Range of values to take. Does not get used really, the spacingfunc is used to
                 get the values from
 
                 Examples:
@@ -482,13 +503,13 @@ class Population:
 
                 Examples:
                     resolution = resolution["M_1"]
-            spacingfunction:
+            spacingfunc:
                 Function determining how the range is sampled. You can either use a real function,
                 or a string representation of a function call. Will get written to a file and
                 then evaluated.
 
                 Examples:
-                    spacingfunction = "const(math.log(m_min), math.log(m_max), {})".format(resolution['M_1'])
+                    spacingfunc = "const(math.log(m_min), math.log(m_max), {})".format(resolution['M_1'])
 
             precode:
                 Extra room for some code. This code will be evaluated within the loop of the
@@ -831,7 +852,7 @@ class Population:
 
         # set previous logging system number to 0
         self.shared_memory["prev_log_system_number"] = multiprocessing.Value('i', 0)
-        
+
         # array to store memory use per-thread
         mem = mem_use()
         self.shared_memory["memory_use_per_thread"] = multiprocessing.Array('d', [1.0 * mem] * self.grid_options["amt_cores"])
@@ -864,7 +885,7 @@ class Population:
 
         Returns an dictionary containing the analytics of the run
 
-        TODO: change the way this is done. Slurm & CONDOR should probably do this different
+        TODO: change the way this is done. Slurm & CONDOR should probably do this differently
         NOTE: SLURM and CONDOR options are not working properly yet
         """
 
@@ -1807,6 +1828,9 @@ class Population:
                     )
                 )
 
+                if self.grid_options['exit_after_dry_run']:
+                    exit()
+
             #######################
             # Reset values and prepare the grid function
             self.grid_options[
@@ -2083,13 +2107,26 @@ class Population:
             #     + "\n"
             # )
 
-            code_string += (
-                indent * depth
-                + "for {}_sample_number in range({}):".format(
-                    grid_variable["name"], grid_variable["resolution"]
+
+            if grid_variable["resolution"] == 0:
+                # generate resolution automatically
+                code_string += (
+                    indent * depth
+                    + "for {}_sample_number in range(len(sampled_values_{})):".format(
+                        grid_variable["name"], grid_variable["name"]
+                    )
+                    + "\n"
                 )
-                + "\n"
-            )
+            else:
+                # use resolution passed in
+                code_string += (
+                    indent * depth
+                    + "for {}_sample_number in range({}):".format(
+                        grid_variable["name"], grid_variable["resolution"]
+                    )
+                    + "\n"
+                )
+
             if grid_variable['gridtype'] == 'edge':
                 code_string += (
                     indent * (depth+1)
@@ -2548,7 +2585,7 @@ class Population:
 
         Requires the grid to be built as a dry run grid
         """
-
+        verbose_print("Dry run of the grid", self.grid_options["verbosity"], 1)
         system_generator = self.grid_options["_system_generator"]
         total_starcount = system_generator(self)
         self.grid_options["_total_starcount"] = total_starcount
diff --git a/binarycpython/utils/grid_options_defaults.py b/binarycpython/utils/grid_options_defaults.py
index 00b3e0d07..67c845abc 100644
--- a/binarycpython/utils/grid_options_defaults.py
+++ b/binarycpython/utils/grid_options_defaults.py
@@ -43,6 +43,8 @@ grid_options_defaults_dict = {
     "_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
+    "exit_after_dry_run": False, # Exit after dry run?
+
     ##########################
     # Execution log:
     ##########################
diff --git a/binarycpython/utils/spacing_functions.py b/binarycpython/utils/spacing_functions.py
index 228426773..0f96a1c9e 100644
--- a/binarycpython/utils/spacing_functions.py
+++ b/binarycpython/utils/spacing_functions.py
@@ -7,7 +7,7 @@ Tasks:
 
 from typing import Union
 import numpy as np
-
+import math
 
 def const(
     min_bound: Union[int, float], max_bound: Union[int, float], steps: int
@@ -18,10 +18,78 @@ def const(
     Args:
         min_bound: lower bound of range
         max_bound: upper bound of range
-        steps: amount of segments between min_bound and max_bound
+        steps: number of segments between min_bound and max_bound
 
     Returns:
         np.linspace(min_bound, max_bound, steps+1)
     """
 
     return np.linspace(min_bound, max_bound, steps+1)
+
+############################################################
+def peak_normalized_gaussian_func(
+    x: Union[int, float], mean: Union[int, float], sigma: Union[int, float]
+) -> Union[int, float]:
+    """
+    Function to evaluate a Gaussian at a given point, note
+    that the normalization is such that the peak is always 1.0,
+    not that the integral is 1.0
+
+    Args:
+        x: location at which to evaluate the distribution
+        mean: mean of the Gaussian
+        sigma: standard deviation of the Gaussian
+
+    Returns:
+        value of the Gaussian at x
+    """
+    gaussian_prefactor = 1.0 #/ math.sqrt(2.0 * math.pi)
+
+    r = 1.0 / sigma
+    y = (x - mean) * r
+    return math.exp(-0.5 * y ** 2)
+
+############################################################
+def gaussian_zoom(
+        min_bound: Union[int, float],
+        max_bound: Union[int, float],
+        zoom_mean: Union[int, float],
+        zoom_dispersion: Union[int,float],
+        zoom_magnitude: Union[int,float],
+        steps: int
+) -> list:
+    """
+    Samples such that a region is zoomed in according to a 1-Gaussian function
+
+    Args:
+        min_bound: lower bound of range
+        max_bound: upper bound of range
+        zoom_mean: mean of the Gaussian zoom location
+        zoom_dispersion: dispersion of the Gaussian
+        zoom_magnitude: depth of the Gaussian (should be 0<= zoom_magntiude <1)
+        steps: number of segments between min_bound and max_bound assuming a linear step
+               this is what you'd normally call "resolution"
+
+    Returns:
+        np.linspace(min_bound, max_bound, steps+1)
+    """
+
+    # linear spacing: this is what we'd have
+    # in the absence of a Gaussian zoom
+    linear_spacing = (max_bound - min_bound)/(steps-1)
+
+    # make the list of values
+    x = min_bound
+    array = np.array([])
+    while x <= max_bound:
+        array = np.append(array,x)
+        g = peak_normalized_gaussian_func(x,zoom_mean,zoom_dispersion)
+        f = 1.0 - zoom_magnitude * g
+        dx = linear_spacing * f
+        x = x + dx
+
+    # force the last array member to be max_bound if it's not
+    if array[-1] != max_bound:
+        array[-1] = max_bound
+
+    return array
-- 
GitLab