Merge pull request #95 from SWIFTSIM/soap-catalogue

SOAP catalogue

Author: JBorrow

requirements.txt

@@ -7,3 +7,4 @@ tqdm
 pyyaml
 black
 matplotlib
+swiftsimio>=6.1.0

tests/helper.py

@@ -0,0 +1,56 @@
+"""
+Contains helper functions for the test routines.
+"""
+
+import subprocess
+import os
+
+webstorage_location = "http://virgodb.cosma.dur.ac.uk/swift-webstorage/IOExamples/"
+test_data_location = "test_data/"
+
+
+def requires(filename):
+    """
+    Use this as a decorator around tests that require data.
+    """
+
+    # First check if the test data directory exists
+    if not os.path.exists(test_data_location):
+        os.mkdir(test_data_location)
+
+    file_location = f"{test_data_location}{filename}"
+
+    if os.path.exists(file_location):
+        ret = 0
+    else:
+        # Download it!
+        ret = subprocess.call(
+            ["wget", f"{webstorage_location}{filename}", "-O", file_location]
+        )
+
+    if ret != 0:
+        Warning(f"Unable to download file at {filename}")
+        # It wrote an empty file, kill it.
+        subprocess.call(["rm", file_location])
+
+        def dont_call_test(func):
+            def empty(*args, **kwargs):
+                return True
+
+            return empty
+
+        return dont_call_test
+
+    else:
+        # Woo, we can do the test!
+
+        def do_call_test(func):
+            def final_test():
+                # Whack the path on there for good measure.
+                return func(f"{test_data_location}{filename}")
+
+            return final_test
+
+        return do_call_test
+
+    raise Exception("You should never have got here.")

tests/test_load_catalogue.py

@@ -3,11 +3,11 @@
 """
 
 from velociraptor import load
+from helper import requires
 
 
-def test_basic_load_catalogue_no_crash(
-    filename="/Users/mphf18/Desktop/halo_027_z00p101.properties",
-):
+@requires("cosmo_0000.properties")
+def test_basic_load_catalogue_no_crash(filename="test_data/cosmo_0000.properties",):
     catalogue = load(filename)
 
     return

velociraptor/__init__.py

@@ -1,7 +1,7 @@
 """
-The velociraptor module. 
+The velociraptor module.
 
-More information is available in the documnetation.
+More information is available in the documentation.
 """
 
 # First things first, we need to upgrade msun and mh from a symbol to a
@@ -27,7 +27,9 @@
     pass
 
 
-from velociraptor.catalogue.catalogue import VelociraptorCatalogue
+from velociraptor.catalogue.catalogue import Catalogue, CatalogueTypeError
+from velociraptor.catalogue.velociraptor_catalogue import VelociraptorCatalogue
+from velociraptor.catalogue.soap_catalogue import SOAPCatalogue
 from velociraptor.__version__ import __version__
 
 from typing import Union, List
@@ -38,10 +40,9 @@ def load(
     disregard_units: bool = False,
     registration_file_path: Union[List[str], str, None] = None,
     mask: slice = Ellipsis,
-) -> VelociraptorCatalogue:
+) -> Catalogue:
     """
-    Loads a velociraptor catalogue, producing a VelociraptorCatalogue
-    object.
+    Loads a velociraptor catalogue, producing a Catalogue object.
 
     Parameters
     ----------
@@ -73,14 +74,16 @@ def load(
     Returns
     -------
 
-    VelociraptorCatalogue
-        The VelociraptorCatalogue object that describes your
-        .properties file.
+    Catalogue
+        The Catalogue object that describes your .properties file.
     """
 
-    catalogue = VelociraptorCatalogue(
-        filename, disregard_units=disregard_units, mask=mask
-    )
+    try:
+        catalogue = VelociraptorCatalogue(
+            filename, disregard_units=disregard_units, mask=mask
+        )
+    except CatalogueTypeError:
+        catalogue = SOAPCatalogue(filename)
 
     if registration_file_path is not None:
         catalogue.register_derived_quantities(registration_file_path)

velociraptor/autoplotter/compare.py

@@ -21,11 +21,12 @@
 from velociraptor.autoplotter.objects import AutoPlotter
 from velociraptor.autoplotter.metadata import AutoPlotterMetadata
 from velociraptor.observations import load_observations
+from velociraptor.catalogue.catalogue import Catalogue
 
 
-class FakeCatalogue(object):
+class FakeCatalogue(Catalogue):
     """
-    Fake VelociraptorCatalogue used to store redshift and
+    Fake Catalogue used to store redshift and
     scale factor information if available.
     """
 

velociraptor/autoplotter/objects.py

@@ -2,7 +2,7 @@
 Main objects for holding information relating to the autoplotter.
 """
 
-from velociraptor import VelociraptorCatalogue
+from velociraptor import Catalogue
 from velociraptor.autoplotter.lines import VelociraptorLine, valid_line_types
 from velociraptor.autoplotter.box_size_correction import VelociraptorBoxSizeCorrection
 from velociraptor.exceptions import AutoPlotterError
@@ -19,7 +19,6 @@
 from pathlib import Path
 
 from os import path, mkdir
-from functools import reduce
 from collections import OrderedDict
 
 valid_plot_types = [
@@ -767,13 +766,13 @@ 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: Catalogue
     ) -> unyt_array:
         """
         Get a quantity from the catalogue using the mask.
         """
 
-        x = reduce(getattr, quantity.split("."), catalogue)
+        x = catalogue.get_quantity(quantity)
         # We give each dataset a custom name, that gets ruined when masking
         # in versions of unyt less than 2.6.0
         name = x.name
@@ -790,9 +789,9 @@ def get_quantity_from_catalogue_with_mask(
 
         if self.selection_mask is not None:
             # Create mask
-            self.structure_mask = reduce(
-                getattr, self.selection_mask.split("."), catalogue
-            ).astype(bool)
+            self.structure_mask = catalogue.get_quantity(self.selection_mask).astype(
+                bool
+            )
         if self.select_structure_type is not None:
             if self.select_structure_type == self.exclude_structure_type:
                 raise AutoPlotterError(
@@ -801,12 +800,14 @@ def get_quantity_from_catalogue_with_mask(
                 )
             self.structure_mask = logical_and(
                 self.structure_mask,
-                catalogue.structure_type.structuretype == self.select_structure_type,
+                catalogue.get_quantity("structure_type.structuretype")
+                == self.select_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,
+                catalogue.get_quantity("structure_type.structuretype")
+                != self.exclude_structure_type,
             )
 
         # combine global and structure masks
@@ -817,9 +818,7 @@ def get_quantity_from_catalogue_with_mask(
         x.name = name
         return x
 
-    def _make_plot_scatter(
-        self, catalogue: VelociraptorCatalogue
-    ) -> Tuple[Figure, Axes]:
+    def _make_plot_scatter(self, catalogue: Catalogue) -> Tuple[Figure, Axes]:
         """
         Makes a scatter plot and returns the figure and axes.
         """
@@ -835,9 +834,7 @@ def _make_plot_scatter(
 
         return fig, ax
 
-    def _make_plot_2dhistogram(
-        self, catalogue: VelociraptorCatalogue
-    ) -> Tuple[Figure, Axes]:
+    def _make_plot_2dhistogram(self, catalogue: Catalogue) -> Tuple[Figure, Axes]:
         """
         Makes a 2d histogram plot and returns the figure and axes.
         """
@@ -855,9 +852,7 @@ def _make_plot_2dhistogram(
 
         return fig, ax
 
-    def _make_plot_massfunction(
-        self, catalogue: VelociraptorCatalogue
-    ) -> Tuple[Figure, Axes]:
+    def _make_plot_massfunction(self, catalogue: Catalogue) -> Tuple[Figure, Axes]:
         """
         Makes a mass function plot and returns the figure and axes.
         """
@@ -888,15 +883,15 @@ def _make_plot_massfunction(
         return fig, ax
 
     def _make_plot_adaptivemassfunction(
-        self, catalogue: VelociraptorCatalogue
+        self, catalogue: Catalogue
     ) -> Tuple[Figure, Axes]:
         """
         Makes the _adaptive_ mass function plot. Same as mass function.
         """
         return self._make_plot_massfunction(catalogue=catalogue)
 
     def _make_plot_luminosityfunction(
-        self, catalogue: VelociraptorCatalogue
+        self, catalogue: Catalogue
     ) -> Tuple[Figure, Axes]:
         """
         Makes a luminosity function plot and returns the figure and axes.
@@ -927,9 +922,7 @@ def _make_plot_luminosityfunction(
 
         return fig, ax
 
-    def _make_plot_histogram(
-        self, catalogue: VelociraptorCatalogue
-    ) -> Tuple[Figure, Axes]:
+    def _make_plot_histogram(self, catalogue: Catalogue) -> Tuple[Figure, Axes]:
         """
         Make histogram plot and return the figure and axes.
         """
@@ -950,7 +943,7 @@ def _make_plot_histogram(
         return fig, ax
 
     def _make_plot_cumulative_histogram(
-        self, catalogue: VelociraptorCatalogue
+        self, catalogue: Catalogue
     ) -> Tuple[Figure, Axes]:
         """
         Make cumulative histogram plot and return the figure and axes.
@@ -984,7 +977,7 @@ def _make_plot_cumulative_histogram(
 
     def make_plot(
         self,
-        catalogue: VelociraptorCatalogue,
+        catalogue: Catalogue,
         directory: str,
         file_extension: str,
         no_plot: bool = False,
@@ -1065,7 +1058,7 @@ class AutoPlotter(object):
     # Forward declarations
     filename: Union[str, List[str]]
     multiple_yaml_files: bool
-    catalogue: VelociraptorCatalogue
+    catalogue: Catalogue
     yaml: Dict[str, Union[Dict, str]]
     plots: List[VelociraptorPlot]
     # Directory containing the observational data.
@@ -1150,9 +1143,7 @@ def parse_yaml(self):
 
         return
 
-    def link_catalogue(
-        self, catalogue: VelociraptorCatalogue, global_mask_tag: Union[None, str]
-    ):
+    def link_catalogue(self, catalogue: Catalogue, global_mask_tag: Union[None, str]):
         """
         Links a catalogue with this object so that the plots
         can actually be created.
@@ -1161,7 +1152,7 @@ def link_catalogue(
         self.catalogue = catalogue
 
         if global_mask_tag is not None:
-            self.global_mask = reduce(getattr, global_mask_tag.split("."), catalogue)
+            self.global_mask = catalogue.get_quantity(global_mask_tag)
         else:
             self.global_mask = True
         return
@@ -1202,8 +1193,8 @@ def create_plots(
                     import sys, traceback
 
                     _, _, exc_traceback = sys.exc_info()
-                    print("Traceback:")
-                    traceback.print_tb(exc_traceback, limit=10, file=sys.stdout)
+                    print("Traceback:", file=sys.stderr)
+                    traceback.print_tb(exc_traceback, limit=10, file=sys.stderr)
             except UnitConversionError as e:
                 print(
                     f"Unable to create plot {plot.filename} due to an error when "
@@ -1218,7 +1209,19 @@ def create_plots(
                     import sys, traceback
 
                     _, _, exc_traceback = sys.exc_info()
-                    print("Traceback:")
-                    traceback.print_tb(exc_traceback, limit=10, file=sys.stdout)
+                    print("Traceback:", file=sys.stderr)
+                    traceback.print_tb(exc_traceback, limit=10, file=sys.stderr)
+            except Exception as e:
+                print(
+                    f"Unable to create plot {plot.filename} due to an unkown error: {e}!",
+                    file=sys.stderr,
+                )
+                self.created_successfully.append(False)
+                if debug:
+                    import sys, traceback
+
+                    _, _, exc_traceback = sys.exc_info()
+                    print("Traceback:", file=sys.stderr)
+                    traceback.print_tb(exc_traceback, limit=10, file=sys.stderr)
 
         return

velociraptor/autoplotter/plot.py

@@ -9,7 +9,7 @@
 from matplotlib import rcParams
 
 from matplotlib.colors import LogNorm
-from velociraptor import VelociraptorCatalogue
+from velociraptor import Catalogue
 from velociraptor.autoplotter.objects import VelociraptorLine
 from typing import Tuple, Union
 
@@ -124,7 +124,7 @@ def histogram(
 
 def decorate_axes(
     ax: plt.Axes,
-    catalogue: VelociraptorCatalogue,
+    catalogue: Catalogue,
     comment: Union[str, None] = None,
     legend_loc: str = "upper left",
     redshift_loc: str = "lower right",