Re-factor get_pixels

astronomy-commons · Oct 19, 2023 · fbedb31 · fbedb31
1 parent 0b5494a
commit fbedb31
Show file tree

Hide file tree

Showing 4 changed files with 10 additions and 10 deletions.
diff --git a/src/hipscat_import/margin_cache/margin_cache_arguments.py b/src/hipscat_import/margin_cache/margin_cache_arguments.py
@@ -40,8 +40,7 @@ def _check_arguments(self):
 
         self.catalog = Catalog.read_from_hipscat(self.input_catalog_path)
 
-        partition_stats = self.catalog.get_pixels()
-        highest_order = np.max(partition_stats["Norder"].values)
+        highest_order = self.catalog.partition_info.get_highest_order()
         margin_pixel_k = highest_order + 1
         if self.margin_order > -1:
             if self.margin_order < margin_pixel_k:

diff --git a/tests/hipscat_import/catalog/test_run_import.py b/tests/hipscat_import/catalog/test_run_import.py
@@ -82,7 +82,7 @@ def test_resume_dask_runner(
     assert catalog.catalog_info.ra_column == "ra"
     assert catalog.catalog_info.dec_column == "dec"
     assert catalog.catalog_info.total_rows == 131
-    assert len(catalog.get_pixels()) == 1
+    assert len(catalog.get_healpix_pixels()) == 1
 
     # Check that the catalog parquet file exists and contains correct object IDs
     output_file = os.path.join(args.catalog_path, "Norder=0", "Dir=0", "Npix=11.parquet")
@@ -122,7 +122,7 @@ def test_resume_dask_runner(
     assert catalog.catalog_info.ra_column == "ra"
     assert catalog.catalog_info.dec_column == "dec"
     assert catalog.catalog_info.total_rows == 131
-    assert len(catalog.get_pixels()) == 1
+    assert len(catalog.get_healpix_pixels()) == 1
     assert_parquet_file_ids(output_file, "id", expected_ids)
 
 
@@ -153,7 +153,7 @@ def test_dask_runner(
     assert catalog.catalog_info.ra_column == "ra"
     assert catalog.catalog_info.dec_column == "dec"
     assert catalog.catalog_info.total_rows == 131
-    assert len(catalog.get_pixels()) == 1
+    assert len(catalog.get_healpix_pixels()) == 1
 
     # Check that the catalog parquet file exists and contains correct object IDs
     output_file = os.path.join(args.catalog_path, "Norder=0", "Dir=0", "Npix=11.parquet")

diff --git a/tests/hipscat_import/catalog/test_run_round_trip.py b/tests/hipscat_import/catalog/test_run_round_trip.py
@@ -52,7 +52,7 @@ def test_import_source_table(
     assert catalog.catalog_path == args.catalog_path
     assert catalog.catalog_info.ra_column == "source_ra"
     assert catalog.catalog_info.dec_column == "source_dec"
-    assert len(catalog.get_pixels()) == 14
+    assert len(catalog.get_healpix_pixels()) == 14
 
 
 @pytest.mark.dask
@@ -347,7 +347,7 @@ def read(self, input_file):
     assert catalog.on_disk
     assert catalog.catalog_path == args.catalog_path
     assert catalog.catalog_info.total_rows == 131
-    assert len(catalog.get_pixels()) == 1
+    assert len(catalog.get_healpix_pixels()) == 1
 
     # Check that the catalog parquet file exists and contains correct object IDs
     output_file = os.path.join(args.catalog_path, "Norder=0", "Dir=0", "Npix=11.parquet")

diff --git a/tests/hipscat_import/margin_cache/test_margin_cache.py b/tests/hipscat_import/margin_cache/test_margin_cache.py
@@ -35,6 +35,7 @@ def test_margin_cache_gen(small_sky_source_catalog, tmp_path, dask_client):
 
     assert len(data) == 13
 
+
 @pytest.mark.dask(timeout=150)
 def test_margin_cache_gen_negative_pixels(small_sky_source_catalog, tmp_path, dask_client):
     """Test that margin cache generation can generate a file for a negative pixel."""
@@ -78,6 +79,7 @@ def test_partition_margin_pixel_pairs(small_sky_source_catalog, tmp_path):
     npt.assert_array_equal(margin_pairs.iloc[:10]["margin_pixel"], expected)
     assert len(margin_pairs) == 196
 
+
 def test_partition_margin_pixel_pairs_negative(small_sky_source_catalog, tmp_path):
     """Ensure partition_margin_pixel_pairs can generate negative tree pixels."""
     args = MarginCacheArguments(
@@ -91,9 +93,7 @@ def test_partition_margin_pixel_pairs_negative(small_sky_source_catalog, tmp_pat
     negative_pixels = args.catalog.generate_negative_tree_pixels()
     combined_pixels = partition_stats + negative_pixels
 
-    margin_pairs = mc._find_partition_margin_pixel_pairs(
-        combined_pixels, args.margin_order
-    )
+    margin_pairs = mc._find_partition_margin_pixel_pairs(combined_pixels, args.margin_order)
 
     expected_order = 0
     expected_pixel = 10
@@ -104,6 +104,7 @@ def test_partition_margin_pixel_pairs_negative(small_sky_source_catalog, tmp_pat
     npt.assert_array_equal(margin_pairs.iloc[-10:]["margin_pixel"], expected)
     assert len(margin_pairs) == 536
 
+
 def test_create_margin_directory(small_sky_source_catalog, tmp_path):
     """Ensure create_margin_directory works on main partition_pixels"""
     args = MarginCacheArguments(