Applied more review suggestions. Properly deal with out-of-bounds values in spline interpolation by setting the correction to 1 outside the interpolation range.

Author: Bert Vandenbroucke

velociraptor-compute-box-size-correction

@@ -104,8 +104,8 @@ if log_y:
 small_spline = interpol.InterpolatedUnivariateSpline(small_box_x, small_box_y)
 large_spline = interpol.InterpolatedUnivariateSpline(large_box_x, large_box_y)
 
-xmin = min(small_box_x.min(), large_box_x.min())
-xmax = max(small_box_x.max(), large_box_x.max())
+xmin = max(small_box_x.min(), large_box_x.min())
+xmax = min(small_box_x.max(), large_box_x.max())
 x_range = np.linspace(xmin, xmax, 100)
 small_y_range = small_spline(x_range)
 large_y_range = large_spline(x_range)
@@ -120,6 +120,8 @@ correction_data = {}
 correction_data["plot_name"] = args.plotname
 correction_data["plot_type"] = args.plottype
 correction_data["is_log_x"] = True
+correction_data["x_units"] = small_box_plot_data["centers_units"]
+correction_data["x_limits"] = np.array([xmin, xmax]).tolist()
 correction_data["x"] = x_range.tolist()
 correction_data["y"] = correction.tolist()
 with open(output_file, "w") as handle:

velociraptor/autoplotter/box_size_correction.py

@@ -7,6 +7,8 @@
 import yaml
 import os
 import scipy.interpolate as interpol
+import unyt
+from typing import Tuple
 
 
 class VelociraptorBoxSizeCorrection:
@@ -17,18 +19,26 @@ def __init__(self, filename: str, correction_directory: str):
         with open(correction_file, "r") as handle:
             correction_data = yaml.safe_load(handle)
         self.is_log_x = correction_data["is_log_x"]
+        self.x_min, self.x_max = correction_data["x_limits"]
         x = np.array(correction_data["x"])
         y = np.array(correction_data["y"])
         self.correction_spline = interpol.InterpolatedUnivariateSpline(x, y)
 
-    def apply_mass_function_correction(self, mass_function_output):
+    def apply_mass_function_correction(
+        self,
+        mass_function_output: Tuple[unyt.unyt_array, unyt.unyt_array, unyt.unyt_array],
+    ) -> Tuple[unyt.unyt_array, unyt.unyt_array, unyt.unyt_array]:
 
         bin_centers, mass_function, error = mass_function_output
 
+        x_vals = bin_centers
+        correction = np.ones(x_vals.shape)
         if self.is_log_x:
-            correction = self.correction_spline(np.log10(bin_centers))
-        else:
-            correction = self.correction_spline(bin_centers)
+            x_vals = np.log10(x_vals)
+        # only apply the correction to bins that are within the range for which
+        # the correction is valid
+        x_mask = (self.x_min <= x_vals) & (x_vals <= self.x_max)
+        correction[x_mask] = self.correction_spline(x_vals[x_mask])
 
         corrected_mass_function = mass_function * correction
         corrected_mass_function.name = mass_function.name

velociraptor/autoplotter/lines.py

@@ -323,6 +323,10 @@ def create_line(
                 box_volume=box_volume,
                 return_bin_edges=True,
             )
+            if self.box_size_correction is not None:
+                mass_function_output = self.box_size_correction.apply_mass_function_correction(
+                    mass_function_output
+                )
             self.output = (
                 *mass_function_output,
                 unyt_array([], units=mass_function_output[0].units),

velociraptor/autoplotter/objects.py

@@ -534,7 +534,7 @@ def _parse_massfunction(self) -> None:
             self.box_size_correction = VelociraptorBoxSizeCorrection(
                 box_size_correction, self.correction_directory
             )
-        except:
+        except KeyError:
             self.box_size_correction = None
         # A bit of a hacky workaround - improve this in the future
         # by combining this functionality properly into the
@@ -569,7 +569,7 @@ def _parse_adaptivemassfunction(self) -> None:
             self.box_size_correction = VelociraptorBoxSizeCorrection(
                 box_size_correction, self.correction_directory
             )
-        except:
+        except KeyError:
             self.box_size_correction = None
         # A bit of a hacky workaround - improve this in the future
         # by combining this functionality properly into the