Merge pull request #68 from james-trayford/global_mask_bg_parts

Global masking

Author: EvgeniiChaikin

velociraptor-plot

@@ -149,7 +149,7 @@ if __name__ == "__main__":
         registration_file_path=args.registration,
     )
     print_if_debug(f"Linking catalogue and AutoPlotter instance.")
-    auto_plotter.link_catalogue(catalogue=catalogue)
+    auto_plotter.link_catalogue(catalogue=catalogue, global_mask_tag=None)
 
     print_if_debug(
         f"Creating figures with extension .{args.file_type} in {args.output}."

velociraptor/__init__.py

@@ -76,5 +76,5 @@ def load(
 
     if registration_file_path is not None:
         catalogue.register_derived_quantities(registration_file_path)
-
+        
     return catalogue

velociraptor/autoplotter/objects.py

@@ -11,7 +11,7 @@
 
 from unyt import unyt_quantity, unyt_array, matplotlib_support
 from unyt.exceptions import UnitConversionError
-from numpy import log10, linspace, logspace, array, logical_and
+from numpy import log10, linspace, logspace, array, logical_and, ones
 from matplotlib.pyplot import Axes, Figure, close
 from yaml import safe_load
 from typing import Union, List, Dict, Tuple
@@ -95,6 +95,8 @@ class VelociraptorPlot(object):
     observational_data_filenames: List[str]
     observational_data_bracket_width: float
     observational_data_directory: str
+    # global mask
+    global_mask: Union[None, array]
 
     def __init__(
         self,
@@ -743,7 +745,7 @@ def _add_lines_to_axes(self, ax: Axes, x: unyt_array, y: unyt_array) -> None:
         return
 
     def get_quantity_from_catalogue_with_mask(
-        self, quantity: str, catalogue: VelociraptorCatalogue
+            self, quantity: str, catalogue: VelociraptorCatalogue,
     ) -> unyt_array:
         """
         Get a quantity from the catalogue using the mask.
@@ -753,62 +755,48 @@ def get_quantity_from_catalogue_with_mask(
         # We give each dataset a custom name, that gets ruined when masking
         # in versions of unyt less than 2.6.0
         name = x.name
-
+        
         if self.structure_mask is not None:
-            x = x[self.structure_mask]
+            # if structure_mask already set, mask and return
+            x_mask = logical_and(self.global_mask, self.structure_mask)
+            x = x[x_mask]
             x.name = name
-        elif self.selection_mask is not None:
+            return x
+        
+        # allow all entries by default
+        self.structure_mask = ones(x.shape).astype(bool)
+        
+        if self.selection_mask is not None:
             # Create mask
             self.structure_mask = reduce(
                 getattr, self.selection_mask.split("."), catalogue
             ).astype(bool)
-
-            if self.select_structure_type is not None:
-                if self.select_structure_type == self.exclude_structure_type:
-                    raise AutoPlotterError(
-                        f"Cannot simultaneously select and exclude structure"
-                        " type {self.select_structure_type}"
-                    )
-                self.structure_mask = logical_and(
-                    self.structure_mask,
-                    catalogue.structure_type.structuretype
-                    == self.select_structure_type,
-                )
-
-            elif self.exclude_structure_type is not None:
-                self.structure_mask = logical_and(
-                    self.structure_mask,
-                    catalogue.structure_type.structuretype
-                    != self.exclude_structure_type,
-                )
-                
-            x = x[self.structure_mask]
-            x.name = name
-        elif self.select_structure_type is not None:
+        if self.select_structure_type is not None:
             if self.select_structure_type == self.exclude_structure_type:
                 raise AutoPlotterError(
                     f"Cannot simultaneously select and exclude structure"
                     " type {self.select_structure_type}"
                 )
-            
-            # Need to create mask
-            self.structure_mask = (
-                catalogue.structure_type.structuretype == self.select_structure_type
+            self.structure_mask = logical_and(
+                self.structure_mask,
+                catalogue.structure_type.structuretype
+                == self.select_structure_type,
             )
-
-            x = x[self.structure_mask]
-            x.name = name
-        elif self.exclude_structure_type is not None:
-            # Need to create mask
-            self.structure_mask = (
-                catalogue.structure_type.structuretype != self.exclude_structure_type
+        if self.exclude_structure_type is not None:
+            self.structure_mask = logical_and(
+                self.structure_mask,
+                catalogue.structure_type.structuretype
+                != self.exclude_structure_type,
             )
+            
+        # combine global and structure masks
+        x_mask = logical_and(self.global_mask, self.structure_mask)
 
-            x = x[self.structure_mask]
-            x.name = name
-
+        # apply to the unyt array of values
+        x = x[x_mask]
+        x.name = name
         return x
-
+                
     def _make_plot_scatter(
         self, catalogue: VelociraptorCatalogue
     ) -> Tuple[Figure, Axes]:
@@ -974,7 +962,7 @@ def _make_plot_cumulative_histogram(
         return fig, ax
 
     def make_plot(
-        self, catalogue: VelociraptorCatalogue, directory: str, file_extension: str
+        self, catalogue: VelociraptorCatalogue, directory: str, file_extension: str,
     ):
         """
         Federates out data parsing to individual functions based on the
@@ -1058,7 +1046,9 @@ class AutoPlotter(object):
     observational_data_directory: str
     # Whether or not the plots were created successfully.
     created_successfully: List[bool]
-
+    # global mask
+    global_mask: Union[None, array]
+    
     def __init__(
         self,
         filename: Union[str, List[str]],
@@ -1123,14 +1113,18 @@ def parse_yaml(self):
 
         return
 
-    def link_catalogue(self, catalogue: VelociraptorCatalogue):
+    def link_catalogue(self, catalogue: VelociraptorCatalogue, global_mask_tag: Union[None, str]):
         """
         Links a catalogue with this object so that the plots
         can actually be created.
         """
 
         self.catalogue = catalogue
 
+        if global_mask_tag is not None:
+            self.global_mask = reduce(getattr, global_mask_tag.split("."), catalogue)
+        else:
+            self.global_mask = True
         return
 
     def create_plots(
@@ -1150,6 +1144,7 @@ def create_plots(
 
         for plot in self.plots:
             try:
+                plot.global_mask = self.global_mask
                 plot.make_plot(
                     catalogue=self.catalogue,
                     directory=directory,

velociraptor/catalogue/catalogue.py

@@ -412,4 +412,3 @@ def register_derived_quantities(self, registration_file_path: str) -> None:
         self.derived_quantities = DerivedQuantities(registration_file_path, self)
 
         return
-

velociraptor/catalogue/registration.py

@@ -37,7 +37,8 @@ def registration_fail_all(
     + name: A fancy (possibly LaTeX'd) name for the field.
     + snake_case: A correct snake_case name for the field.
     """
-
+    
+    
     if field_path == "ThisFieldPathWouldNeverExist":
         return (
             unit_system.length,
@@ -307,6 +308,8 @@ def registration_masses(
         full_name = "$M_{\\rm FOF}$"
     elif field_path == "Mass_tot":
         full_name = r"$M$"
+    elif field_path == "Mass_interloper":
+        full_name = "$M_{\\rm BG}$"
 
     # General regex matching case.
 
@@ -321,7 +324,6 @@ def registration_masses(
     )
     regex = cached_regex(match_string)
     match = regex.match(field_path)
-
     if match and not full_name:
         mass = match.group(1)
         radius = match.group(2)
@@ -1441,8 +1443,7 @@ def registration_spherical_overdensities(
         return unit, full_name, snake_case
     else:
         raise RegistrationDoesNotMatchError
-
-
+    
 # TODO
 # lambda_B
 # q