Merge pull request #29 from SWIFTSIM/adaptive_bin_cache_tweak

Adaptive Bin Caching

Author: JBorrow

examples/auto_plotter_example.yml

@@ -1,3 +1,29 @@
+stellar_mass_gas_metallicity:
+  type: "2dhistogram"
+  x:
+    quantity: "apertures.mass_star_100_kpc"
+    units: Solar_Mass
+    start: 1e6
+    end: 1e12
+  y:
+    quantity: "apertures.zmet_gas_100_kpc"
+    units: ""
+    start: 1e-6
+    end: 1e-1
+  median:
+    plot: true
+    log: true
+    number_of_bins: 25
+    adaptive: true
+    start:
+      value: 1e6
+      units: Solar_Mass
+    end:
+      value: 1e12
+      units: Solar_Mass
+  metadata:
+    title: Stellar Mass-Gas Metallicity relation (100 kpc aperture)
+
 stellar_mass_halo_mass_adaptive_100:
   type: "2dhistogram"
   x:
@@ -147,6 +173,7 @@ stellar_mass_stellar_metallicity:
   median:
     plot: true
     log: true
+    adaptive: true
     number_of_bins: 25
     start:
       value: 1e6
@@ -157,27 +184,3 @@ stellar_mass_stellar_metallicity:
   metadata:
     title: Stellar Mass-Stellar Metallicity relation (100 kpc aperture)
 
-stellar_mass_gas_metallicity:
-  type: "2dhistogram"
-  x:
-    quantity: "apertures.mass_star_100_kpc"
-    units: Solar_Mass
-    start: 1e6
-    end: 1e12
-  y:
-    quantity: "apertures.zmet_gas_100_kpc"
-    units: ""
-    start: 1e-6
-    end: 1e-1
-  median:
-    plot: true
-    log: true
-    number_of_bins: 25
-    start:
-      value: 1e6
-      units: Solar_Mass
-    end:
-      value: 1e12
-      units: Solar_Mass
-  metadata:
-    title: Stellar Mass-Gas Metallicity relation (100 kpc aperture)

velociraptor/__version__.py

@@ -1 +1 @@
-__version__ = "0.10.3"
+__version__ = "0.10.4"

velociraptor/catalogue/catalogue.py

@@ -117,6 +117,7 @@ def getter(self):
                 try:
                     setattr(self, f"_{name}", unyt.unyt_array(handle[field][...], unit))
                     getattr(self, f"_{name}").name = full_name
+                    getattr(self, f"_{name}").file = filename
                 except KeyError:
                     print(f"Could not read {field}")
                     return None

velociraptor/tools/adaptive.py

@@ -6,6 +6,31 @@
 import unyt
 
 
+adaptive_bin_cache = {}
+
+
+def adaptive_bin_hash(
+    values,
+    lowest_value,
+    highest_value,
+    base_n_bins,
+    minimum_in_bin,
+    logarithmic,
+    stretch_final_bin,
+):
+    """
+    Hash for adaptive binning. Note that this can raise AttributeError
+    in the case where the array is unhashable.
+    """
+
+    this_hash = (
+        f"{values.size}{values.name}{lowest_value}{highest_value}"
+        f"{base_n_bins}{minimum_in_bin}{logarithmic}{stretch_final_bin}"
+    )
+
+    return this_hash
+
+
 def create_adaptive_bins(
     values: unyt.unyt_array,
     lowest_value: unyt.unyt_quantity,
@@ -58,7 +83,36 @@ def create_adaptive_bins(
     bin_edges: unyt.unyt_array, optional
         Bin edges that were used in the binning process.
 
+    Notes
+    -----
+
+    Caches the output as this procedure can be very expensive, and will be
+    repeated several times.
+
     """
+
+    # First we check in the cache to see if we have already performed
+    # the binning procedure.
+    try:
+        this_hash = adaptive_bin_hash(
+            values=values,
+            lowest_value=lowest_value,
+            highest_value=highest_value,
+            base_n_bins=base_n_bins,
+            minimum_in_bin=minimum_in_bin,
+            logarithmic=logarithmic,
+            stretch_final_bin=stretch_final_bin,
+        )
+    except AttributeError:
+        this_hash = False
+
+    if this_hash:
+        try:
+            return adaptive_bin_cache[this_hash]
+        except:
+            # Not created yet!
+            pass
+
     assert (
         values.units == lowest_value.units and lowest_value.units == highest_value.units
     ), "Please ensure that all value quantities have the same units."
@@ -162,5 +216,8 @@ def create_adaptive_bins(
             np.array([*bin_edges_left, bin_edges_right[-1]]), units=values.units
         )
 
+    if this_hash:
+        adaptive_bin_cache[this_hash] = (bin_centers, bin_edges)
+
     return bin_centers, bin_edges
 

velociraptor/tools/mass_functions.py

@@ -242,12 +242,21 @@ def create_adaptive_mass_function(
             bin_medians.append(
                 np.median(sorted_masses[current_lower_index : index + 1])
             )
-            bin_edges_right.append(mass)
-            bin_edges_left.append(mass)
+
+            # The new bin edge lives half way in between our current value and the
+            # next value, if it exists.
+            try:
+                new_edge = 0.5 * (sorted_masses[index + 1] + mass)
+            except IndexError:
+                # This case is where `value` is the last item in the array.
+                new_edge = mass
+
+            bin_edges_right.append(new_edge)
+            bin_edges_left.append(new_edge)
             number_in_bin.append(current_bin_count)
 
             # Reset for the lads
-            current_edge_left = mass
+            current_edge_left = new_edge
             current_lower_index = index
             current_bin_count = 0
         else: