Merge pull request #87 from SWIFTSIM/box-size-correction

Box size correction

Author: bwvdnbro

velociraptor-compute-box-size-correction

@@ -0,0 +1,128 @@
+#!/usr/bin/env python
+
+"""
+Compute a box size correction file that can be used as the 'box_size_correction'
+argument for an autoplotter plot.
+
+Usage:
+  velociraptor-compute-box-size-correction \
+    smallbox largebox plotname plottype output
+
+with:
+  - smallbox/largebox: data*.yml output file from a pipeline run
+  - plotname: Name of a particular plot in the data*.yml files
+  - plottype: Type of plot (currently supported: mass_function)
+  - output: Name of an output .yml file. If the .yml extension is missing, it is
+    added.
+"""
+
+import argparse
+import os
+import yaml
+import numpy as np
+import scipy.interpolate as interpol
+
+argparser = argparse.ArgumentParser("Compute the box size correction for a plot.")
+argparser.add_argument(
+    "smallbox", help="Pipeline output for the small box that needs to be corrected."
+)
+argparser.add_argument(
+    "largebox", help="Pipeline output for the large box that we want to correct to."
+)
+argparser.add_argument("plotname", help="Name of the plot that we want to correct.")
+argparser.add_argument("plottype", help="Type of the plot we want to correct.")
+argparser.add_argument(
+    "output", help="Name of the output file that will store the correction."
+)
+args = argparser.parse_args()
+
+if not args.plottype in ["mass_function"]:
+    raise AttributeError(
+        f"Cannot compute box size correction for plot type {args.plottype}!"
+    )
+
+log_x = False
+log_y = False
+if args.plottype in ["mass_function"]:
+    log_x = True
+    log_y = True
+
+small_box = args.smallbox
+large_box = args.largebox
+for file in [args.smallbox, args.largebox]:
+    if not os.path.exists(file):
+        raise AttributeError(f"File {file} could not be found!")
+
+output_file = args.output
+if not output_file.endswith(".yml"):
+    output_file += ".yml"
+try:
+    open(output_file, "w").close()
+except:
+    raise AttributeError(f"Can not write to {output_file}!")
+
+with open(args.smallbox, "r") as handle:
+    small_box = yaml.safe_load(handle)
+with open(args.largebox, "r") as handle:
+    large_box = yaml.safe_load(handle)
+
+try:
+    small_box_data = small_box[args.plotname]["lines"]
+except:
+    raise AttributeError(f"Could not find {args.plotname} in {args.smallbox}!")
+try:
+    large_box_data = large_box[args.plotname]["lines"]
+except:
+    raise AttributeError(f"Could not find {args.plotname} in {args.largebox}!")
+
+try:
+    small_box_plot_data = small_box_data[args.plottype]
+except:
+    raise AttributeError(
+        f"{args.plottype} not found in plot {args.plotname} in {args.smallbox}!"
+    )
+try:
+    large_box_plot_data = large_box_data[args.plottype]
+except:
+    raise AttributeError(
+        f"{args.plottype} not found in plot {args.plotname} in {args.largebox}!"
+    )
+
+small_box_x = small_box_plot_data["centers"]
+small_box_y = small_box_plot_data["values"]
+large_box_x = large_box_plot_data["centers"]
+large_box_y = large_box_plot_data["values"]
+
+if log_x:
+    small_box_x = np.log10(small_box_x)
+    large_box_x = np.log10(large_box_x)
+
+if log_y:
+    small_box_y = np.log10(small_box_y)
+    large_box_y = np.log10(large_box_y)
+
+small_spline = interpol.InterpolatedUnivariateSpline(small_box_x, small_box_y)
+large_spline = interpol.InterpolatedUnivariateSpline(large_box_x, large_box_y)
+
+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)
+
+if log_y:
+    small_y_range = 10.0**small_y_range
+    large_y_range = 10.0**large_y_range
+
+correction = large_y_range / small_y_range
+
+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:
+    yaml.safe_dump(correction_data, handle)

velociraptor/autoplotter/box_size_correction.py

@@ -0,0 +1,46 @@
+"""
+Functionality to apply a mass dependent correction to quantities that have been
+binned in mass bins (e.g. a mass function).
+"""
+
+import numpy as np
+import yaml
+import os
+import scipy.interpolate as interpol
+import unyt
+from typing import Tuple
+
+
+class VelociraptorBoxSizeCorrection:
+    def __init__(self, filename: str, correction_directory: str):
+        correction_file = f"{correction_directory}/{filename}"
+        if not os.path.exists(correction_file):
+            raise FileNotFoundError(f"Could not find {correction_file}!")
+        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: 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:
+            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
+
+        return bin_centers, corrected_mass_function, error

velociraptor/autoplotter/lines.py

@@ -18,6 +18,7 @@
 )
 from velociraptor.tools.histogram import create_histogram_given_bins
 from velociraptor.tools.adaptive import create_adaptive_bins
+from velociraptor.autoplotter.box_size_correction import VelociraptorBoxSizeCorrection
 
 valid_line_types = [
     "median",
@@ -61,6 +62,8 @@ class VelociraptorLine(object):
     bins: unyt_array = None
     # Scatter can be: "none", "errorbar", or "shaded"
     scatter: str
+    # Box size correction?
+    box_size_correction: Union[None, VelociraptorBoxSizeCorrection]
     # Output: centers, values, scatter, additional_x, additional_y - initialised here
     # to prevent crashes in other code.
     output: Tuple[unyt_array] = (
@@ -71,7 +74,12 @@ class VelociraptorLine(object):
         unyt_array([]),
     )
 
-    def __init__(self, line_type: str, line_data: Dict[str, Union[Dict, str]]):
+    def __init__(
+        self,
+        line_type: str,
+        line_data: Dict[str, Union[Dict, str]],
+        box_size_correction: Union[None, VelociraptorBoxSizeCorrection] = None,
+    ):
         """
         Initialise a line with data from the yaml file.
         """
@@ -82,6 +90,8 @@ def __init__(self, line_type: str, line_data: Dict[str, Union[Dict, str]]):
         self.data = line_data
         self._parse_data()
 
+        self.box_size_correction = box_size_correction
+
         return
 
     def _parse_line_type(self):
@@ -265,6 +275,10 @@ def create_line(
             mass_function_output = create_mass_function_given_bins(
                 masked_x, self.bins, box_volume=box_volume
             )
+            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),
@@ -309,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

@@ -4,6 +4,7 @@
 
 from velociraptor import VelociraptorCatalogue
 from velociraptor.autoplotter.lines import VelociraptorLine, valid_line_types
+from velociraptor.autoplotter.box_size_correction import VelociraptorBoxSizeCorrection
 from velociraptor.exceptions import AutoPlotterError
 from velociraptor.observations import load_observations
 
@@ -87,6 +88,9 @@ class VelociraptorPlot(object):
     exclude_structure_type: Union[None, int]
     structure_mask: Union[None, array]
     selection_mask: Union[None, array]
+    # Apply a box size correction to the plot?
+    correction_directory: str
+    box_size_correction: Union[None, VelociraptorBoxSizeCorrection]
     # Where should the legend and z, a information be placed?
     legend_loc: str
     redshift_loc: str
@@ -103,13 +107,15 @@ def __init__(
         filename: str,
         data: Dict[str, Union[Dict, str]],
         observational_data_directory: str,
+        correction_directory: str,
     ):
         """
         Initialise the plot object variables.
         """
         self.filename = filename
         self.data = data
         self.observational_data_directory = observational_data_directory
+        self.correction_directory = correction_directory
 
         self._parse_data()
 
@@ -523,6 +529,13 @@ def _parse_massfunction(self) -> None:
 
         self._parse_common_histogramtype()
 
+        try:
+            box_size_correction = str(self.data["box_size_correction"])
+            self.box_size_correction = VelociraptorBoxSizeCorrection(
+                box_size_correction, self.correction_directory
+            )
+        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
         # VelociraptorLine methods.
@@ -535,6 +548,7 @@ def _parse_massfunction(self) -> None:
                 start=dict(value=self.x_lim[0].value, units=self.x_lim[0].units),
                 end=dict(value=self.x_lim[1].value, units=self.x_lim[1].units),
             ),
+            box_size_correction=self.box_size_correction,
         )
 
         return
@@ -550,6 +564,13 @@ def _parse_adaptivemassfunction(self) -> None:
 
         self._parse_common_histogramtype()
 
+        try:
+            box_size_correction = str(self.data["box_size_correction"])
+            self.box_size_correction = VelociraptorBoxSizeCorrection(
+                box_size_correction, self.correction_directory
+            )
+        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
         # VelociraptorLine methods.
@@ -563,6 +584,7 @@ def _parse_adaptivemassfunction(self) -> None:
                 end=dict(value=self.x_lim[1].value, units=self.x_lim[1].units),
                 adaptive=True,
             ),
+            box_size_correction=self.box_size_correction,
         )
 
         return
@@ -1042,6 +1064,8 @@ class AutoPlotter(object):
     plots: List[VelociraptorPlot]
     # Directory containing the observational data.
     observational_data_directory: str
+    # Directory containing box size correction data
+    correction_directory: str
     # Whether or not the plots were created successfully.
     created_successfully: List[bool]
     # global mask
@@ -1051,6 +1075,7 @@ def __init__(
         self,
         filename: Union[str, List[str]],
         observational_data_directory: Union[None, str] = None,
+        correction_directory: Union[None, str] = None,
     ) -> None:
         """
         Initialises the AutoPlotter object with the yaml filename(s).
@@ -1075,6 +1100,9 @@ def __init__(
             if observational_data_directory is not None
             else ""
         )
+        self.correction_directory = Path(
+            correction_directory if correction_directory is not None else ""
+        )
 
         self.load_yaml()
         self.parse_yaml()
@@ -1105,7 +1133,12 @@ def parse_yaml(self):
         """
 
         self.plots = [
-            VelociraptorPlot(filename, plot, self.observational_data_directory)
+            VelociraptorPlot(
+                filename,
+                plot,
+                self.observational_data_directory,
+                self.correction_directory,
+            )
             for filename, plot in self.yaml.items()
         ]