Merge 28063ae into f811ab3

Author: smcguire-cmu

src/hipscat/catalog/__init__.py

@@ -1,7 +1,7 @@
 """Catalog data wrappers"""
 
-from .catalog import Catalog
 from .association_catalog.association_catalog import AssociationCatalog
-from .dataset.dataset import Dataset
+from .catalog import Catalog
 from .catalog_type import CatalogType
+from .dataset.dataset import Dataset
 from .partition_info import PartitionInfo

src/hipscat/catalog/association_catalog/association_catalog.py

@@ -1,20 +1,19 @@
 from typing import Tuple, Union
+
 import pandas as pd
 from typing_extensions import TypeAlias
 
 from hipscat.catalog.association_catalog.association_catalog_info import AssociationCatalogInfo
 from hipscat.catalog.association_catalog.partition_join_info import PartitionJoinInfo
 from hipscat.catalog.catalog_type import CatalogType
-from hipscat.catalog.dataset.dataset import Dataset
-from hipscat.io import FilePointer, paths
+from hipscat.catalog.healpix_dataset.healpix_dataset import HealpixDataset, PixelInputTypes
+from hipscat.io import FilePointer, file_io, paths
 
 
-class AssociationCatalog(Dataset):
+class AssociationCatalog(HealpixDataset):
     """A HiPSCat Catalog for enabling fast joins between two HiPSCat catalogs
 
-    Catalogs of this type are partitioned based on the partitioning of both joining catalogs in the
-    form 'Norder=/Dir=/Npix=/join_Norder=/join_Dir=/join_Npix=.parquet'. Where each partition
-    contains the matching pair of hipscat indexes from each catalog's respective partitions to join.
+    Catalogs of this type are partitioned based on the partitioning of the left catalog.
     The `partition_join_info` metadata file specifies all pairs of pixels in the Association
     Catalog, corresponding to each pair of partitions in each catalog that contain rows to join.
     """
@@ -29,13 +28,14 @@ class AssociationCatalog(Dataset):
     def __init__(
         self,
         catalog_info: CatalogInfoClass,
+        pixels: PixelInputTypes,
         join_pixels: JoinPixelInputTypes,
         catalog_path=None,
-        storage_options: dict=None
+        storage_options: dict = None,
     ) -> None:
         if not catalog_info.catalog_type == CatalogType.ASSOCIATION:
             raise ValueError("Catalog info `catalog_type` must be 'association'")
-        super().__init__(catalog_info, catalog_path, storage_options=storage_options)
+        super().__init__(catalog_info, pixels, catalog_path, storage_options=storage_options)
         self.join_info = self._get_partition_join_info_from_pixels(join_pixels)
 
     def get_join_pixels(self) -> pd.DataFrame:
@@ -60,11 +60,21 @@ def _get_partition_join_info_from_pixels(
     @classmethod
     def _read_args(
         cls, catalog_base_dir: FilePointer, storage_options: dict = None
-    ) -> Tuple[CatalogInfoClass, JoinPixelInputTypes]:  # type: ignore[override]
+    ) -> Tuple[CatalogInfoClass, PixelInputTypes, JoinPixelInputTypes]:  # type: ignore[override]
         args = super()._read_args(catalog_base_dir, storage_options=storage_options)
         partition_join_info_file = paths.get_partition_join_info_pointer(catalog_base_dir)
         partition_join_info = PartitionJoinInfo.read_from_file(
-            partition_join_info_file,
-            storage_options=storage_options
+            partition_join_info_file, storage_options=storage_options
         )
         return args + (partition_join_info,)
+
+    @classmethod
+    def _check_files_exist(cls, catalog_base_dir: FilePointer, storage_options: dict = None):
+        super()._check_files_exist(catalog_base_dir, storage_options=storage_options)
+        partition_join_info_file = paths.get_partition_join_info_pointer(catalog_base_dir)
+        if not file_io.does_file_or_directory_exist(
+            partition_join_info_file, storage_options=storage_options
+        ):
+            raise FileNotFoundError(
+                f"No partition join info found where expected: {str(partition_join_info_file)}"
+            )

src/hipscat/catalog/catalog.py

@@ -1,35 +1,29 @@
 """Container class to hold catalog metadata and partition iteration"""
 from __future__ import annotations
-from typing import List, Tuple, Union
 
 import dataclasses
+from typing import List
 
 import healpy as hp
 import numpy as np
-import pandas as pd
-
 from typing_extensions import TypeAlias
+
 from hipscat.catalog.catalog_info import CatalogInfo
 from hipscat.catalog.catalog_type import CatalogType
-from hipscat.catalog.dataset.dataset import Dataset
-from hipscat.catalog.partition_info import PartitionInfo
-from hipscat.io import FilePointer, file_io, paths
+from hipscat.catalog.healpix_dataset.healpix_dataset import HealpixDataset, PixelInputTypes
 from hipscat.pixel_math import HealpixPixel
 from hipscat.pixel_math.cone_filter import filter_pixels_by_cone
 from hipscat.pixel_tree.pixel_node_type import PixelNodeType
-from hipscat.pixel_tree.pixel_tree import PixelTree
-from hipscat.pixel_tree.pixel_tree_builder import PixelTreeBuilder
 
 
-class Catalog(Dataset):
+class Catalog(HealpixDataset):
     """A HiPSCat Catalog with data stored in a HEALPix Hive partitioned structure
 
     Catalogs of this type are partitioned spatially, contain `partition_info` metadata specifying
     the pixels in Catalog, and on disk conform to the parquet partitioning structure
     `Norder=/Dir=/Npix=.parquet`
     """
 
-    PixelInputTypes = Union[pd.DataFrame, PartitionInfo, PixelTree, List[HealpixPixel]]
     HIPS_CATALOG_TYPES = [CatalogType.OBJECT, CatalogType.SOURCE, CatalogType.MARGIN]
 
     # Update CatalogInfoClass, used to check if the catalog_info is the correct type, and
@@ -42,7 +36,7 @@ def __init__(
         catalog_info: CatalogInfoClass,
         pixels: PixelInputTypes,
         catalog_path: str = None,
-        storage_options: dict = None
+        storage_options: dict = None,
     ) -> None:
         """Initializes a Catalog
 
@@ -60,76 +54,7 @@ def __init__(
                 f"Catalog info `catalog_type` must be one of "
                 f"{', '.join([t.value for t in self.HIPS_CATALOG_TYPES])}"
             )
-        super().__init__(catalog_info, catalog_path, storage_options)
-        self.partition_info = self._get_partition_info_from_pixels(pixels)
-        self.pixel_tree = self._get_pixel_tree_from_pixels(pixels)
-
-    @staticmethod
-    def _get_partition_info_from_pixels(pixels: PixelInputTypes) -> PartitionInfo:
-        if isinstance(pixels, PartitionInfo):
-            return pixels
-        if isinstance(pixels, pd.DataFrame):
-            return PartitionInfo(pixels)
-        if isinstance(pixels, PixelTree):
-            return PartitionInfo.from_healpix(
-                [
-                    HealpixPixel(node.hp_order, node.hp_pixel)
-                    for node in pixels.root_pixel.get_all_leaf_descendants()
-                ]
-            )
-        if pd.api.types.is_list_like(pixels):
-            return PartitionInfo.from_healpix(pixels)
-        raise TypeError("Pixels must be of type PartitionInfo, Dataframe, PixelTree, or List[HealpixPixel]")
-
-    @staticmethod
-    def _get_pixel_tree_from_pixels(pixels: PixelInputTypes) -> PixelTree:
-        if isinstance(pixels, PartitionInfo):
-            return PixelTreeBuilder.from_partition_info_df(pixels.data_frame)
-        if isinstance(pixels, pd.DataFrame):
-            return PixelTreeBuilder.from_partition_info_df(pixels)
-        if isinstance(pixels, PixelTree):
-            return pixels
-        if pd.api.types.is_list_like(pixels):
-            return PixelTreeBuilder.from_healpix(pixels)
-        raise TypeError("Pixels must be of type PartitionInfo, Dataframe, PixelTree, or List[HealpixPixel]")
-
-    def get_pixels(self):
-        """Get all healpix pixels that are contained in the catalog
-
-        Returns:
-            data frame with per-pixel data.
-
-            The data frame contains the following columns:
-
-            - order: order of the destination pixel
-            - pixel: pixel number *at the above order*
-            - num_objects: the number of rows in the pixel's partition
-        """
-        return self.partition_info.data_frame
-
-    def get_healpix_pixels(self) -> List[HealpixPixel]:
-        """Get healpix pixel objects for all pixels contained in the catalog.
-
-        Returns:
-            List of HealpixPixel
-        """
-        return self.partition_info.get_healpix_pixels()
-
-    @classmethod
-    def _read_args(
-        cls, catalog_base_dir: FilePointer, storage_options: dict = None
-    ) -> Tuple[CatalogInfoClass, PartitionInfo]:
-        args = super()._read_args(catalog_base_dir, storage_options=storage_options)
-        partition_info_file = paths.get_partition_info_pointer(catalog_base_dir)
-        partition_info = PartitionInfo.read_from_file(partition_info_file, storage_options=storage_options)
-        return args + (partition_info,)
-
-    @classmethod
-    def _check_files_exist(cls, catalog_base_dir: FilePointer, storage_options: dict = None):
-        super()._check_files_exist(catalog_base_dir, storage_options=storage_options)
-        partition_info_file = paths.get_partition_info_pointer(catalog_base_dir)
-        if not file_io.does_file_or_directory_exist(partition_info_file, storage_options=storage_options):
-            raise FileNotFoundError(f"No partition info found where expected: {str(partition_info_file)}")
+        super().__init__(catalog_info, pixels, catalog_path, storage_options)
 
     def filter_by_cone(self, ra: float, dec: float, radius: float) -> Catalog:
         """Filter the pixels in the catalog to only include the pixels that overlap with a cone

src/hipscat/catalog/dataset/dataset.py

@@ -18,10 +18,7 @@ class Dataset:
     CatalogInfoClass: TypeAlias = BaseCatalogInfo
 
     def __init__(
-        self,
-        catalog_info: CatalogInfoClass,
-        catalog_path=None,
-        storage_options: dict = None
+        self, catalog_info: CatalogInfoClass, catalog_path=None, storage_options: dict = None
     ) -> None:
         """Initializes a Dataset
 

src/hipscat/catalog/healpix_dataset/__init__.py

@@ -0,0 +1,113 @@
+from typing import List, Tuple, Union
+
+import pandas as pd
+from typing_extensions import TypeAlias
+
+from hipscat.catalog.dataset import BaseCatalogInfo, Dataset
+from hipscat.catalog.partition_info import PartitionInfo
+from hipscat.io import FilePointer, file_io, paths
+from hipscat.pixel_math import HealpixPixel
+from hipscat.pixel_tree.pixel_tree import PixelTree
+from hipscat.pixel_tree.pixel_tree_builder import PixelTreeBuilder
+
+PixelInputTypes = Union[pd.DataFrame, PartitionInfo, PixelTree, List[HealpixPixel]]
+
+
+class HealpixDataset(Dataset):
+    """A HiPSCat dataset partitioned with a HEALPix partitioning structure.
+
+    Catalogs of this type are partitioned based on the ra and dec of the points with each partition
+    containing points within a given HEALPix pixel. The files are in the form 'Norder=/Dir=/Npix=.parquet'.
+    """
+
+    CatalogInfoClass: TypeAlias = BaseCatalogInfo
+    catalog_info: CatalogInfoClass
+
+    def __init__(
+        self,
+        catalog_info: CatalogInfoClass,
+        pixels: PixelInputTypes,
+        catalog_path: str = None,
+        storage_options: dict = None,
+    ) -> None:
+        """Initializes a Catalog
+
+        Args:
+            catalog_info: CatalogInfo object with catalog metadata
+            pixels: Specifies the pixels contained in the catalog. Can be either a Dataframe with
+                columns `Norder`, `Dir`, and `Npix` matching a `partition_info.csv` file, a
+                `PartitionInfo object`, or a `PixelTree` object
+            catalog_path: If the catalog is stored on disk, specify the location of the catalog
+                Does not load the catalog from this path, only store as metadata
+            storage_options: dictionary that contains abstract filesystem credentials
+        """
+        super().__init__(catalog_info, catalog_path, storage_options)
+        self.partition_info = self._get_partition_info_from_pixels(pixels)
+        self.pixel_tree = self._get_pixel_tree_from_pixels(pixels)
+
+    def get_pixels(self):
+        """Get all healpix pixels that are contained in the catalog
+
+        Returns:
+            data frame with per-pixel data.
+
+            The data frame contains the following columns:
+
+            - order: order of the destination pixel
+            - pixel: pixel number *at the above order*
+            - num_objects: the number of rows in the pixel's partition
+        """
+        return self.partition_info.data_frame
+
+    def get_healpix_pixels(self) -> List[HealpixPixel]:
+        """Get healpix pixel objects for all pixels contained in the catalog.
+
+        Returns:
+            List of HealpixPixel
+        """
+        return self.partition_info.get_healpix_pixels()
+
+    @staticmethod
+    def _get_partition_info_from_pixels(pixels: PixelInputTypes) -> PartitionInfo:
+        if isinstance(pixels, PartitionInfo):
+            return pixels
+        if isinstance(pixels, pd.DataFrame):
+            return PartitionInfo(pixels)
+        if isinstance(pixels, PixelTree):
+            return PartitionInfo.from_healpix(
+                [
+                    HealpixPixel(node.hp_order, node.hp_pixel)
+                    for node in pixels.root_pixel.get_all_leaf_descendants()
+                ]
+            )
+        if pd.api.types.is_list_like(pixels):
+            return PartitionInfo.from_healpix(pixels)
+        raise TypeError("Pixels must be of type PartitionInfo, Dataframe, PixelTree, or List[HealpixPixel]")
+
+    @staticmethod
+    def _get_pixel_tree_from_pixels(pixels: PixelInputTypes) -> PixelTree:
+        if isinstance(pixels, PartitionInfo):
+            return PixelTreeBuilder.from_partition_info_df(pixels.data_frame)
+        if isinstance(pixels, pd.DataFrame):
+            return PixelTreeBuilder.from_partition_info_df(pixels)
+        if isinstance(pixels, PixelTree):
+            return pixels
+        if pd.api.types.is_list_like(pixels):
+            return PixelTreeBuilder.from_healpix(pixels)
+        raise TypeError("Pixels must be of type PartitionInfo, Dataframe, PixelTree, or List[HealpixPixel]")
+
+    @classmethod
+    def _read_args(
+        cls, catalog_base_dir: FilePointer, storage_options: dict = None
+    ) -> Tuple[CatalogInfoClass, PartitionInfo]:
+        args = super()._read_args(catalog_base_dir, storage_options=storage_options)
+        partition_info_file = paths.get_partition_info_pointer(catalog_base_dir)
+        partition_info = PartitionInfo.read_from_file(partition_info_file, storage_options=storage_options)
+        return args + (partition_info,)
+
+    @classmethod
+    def _check_files_exist(cls, catalog_base_dir: FilePointer, storage_options: dict = None):
+        super()._check_files_exist(catalog_base_dir, storage_options=storage_options)
+        partition_info_file = paths.get_partition_info_pointer(catalog_base_dir)
+        if not file_io.does_file_or_directory_exist(partition_info_file, storage_options=storage_options):
+            raise FileNotFoundError(f"No partition info found where expected: {str(partition_info_file)}")

src/hipscat/catalog/healpix_dataset/healpix_dataset.py

@@ -1,15 +1,17 @@
 from __future__ import annotations
-from typing import List
 
 import os
 import warnings
+from typing import List
+
 import pandas as pd
 
 from hipscat.catalog.catalog import CatalogType
 from hipscat.catalog.dataset.dataset import Dataset
 from hipscat.inspection.almanac_info import AlmanacInfo
 from hipscat.io.file_io import file_pointer
 
+
 class Almanac:
     """Single instance of an almanac, and available catalogs within namespaces
 

src/hipscat/inspection/almanac.py

@@ -1,11 +1,11 @@
 from __future__ import annotations
-from typing import List
 
 import dataclasses
 import os
 from dataclasses import dataclass, field
-import yaml
+from typing import List
 
+import yaml
 from typing_extensions import Self
 
 from hipscat.catalog.dataset import catalog_info_factory

src/hipscat/inspection/almanac_info.py

@@ -10,8 +10,6 @@
     get_partition_info_pointer,
     get_point_map_file_pointer,
     get_provenance_pointer,
-    pixel_association_directory,
-    pixel_association_file,
     pixel_catalog_file,
     pixel_directory,
 )