diff --git a/src/hipscat_import/margin_cache/margin_cache_arguments.py b/src/hipscat_import/margin_cache/margin_cache_arguments.py index 3f3cadca..429974ab 100644 --- 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 index bdccde4e..18e90dac 100644 --- 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 index b3b55532..5f71c570 100644 --- 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 index 79443c02..42fd8ee9 100644 --- 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(