From 3cd57516f6fe3283294c57d7d15dc8b32edd84e8 Mon Sep 17 00:00:00 2001 From: delucchi-cmu Date: Mon, 29 Jan 2024 12:12:32 -0500 Subject: [PATCH] Add progress bar to early stages of margin import --- .../margin_cache/margin_cache.py | 18 ++++++++++++------ .../margin_cache/test_margin_cache.py | 2 +- 2 files changed, 13 insertions(+), 7 deletions(-) diff --git a/src/hipscat_import/margin_cache/margin_cache.py b/src/hipscat_import/margin_cache/margin_cache.py index 6f30551e..75544089 100644 --- a/src/hipscat_import/margin_cache/margin_cache.py +++ b/src/hipscat_import/margin_cache/margin_cache.py @@ -106,17 +106,23 @@ def generate_margin_cache(args, client): args (MarginCacheArguments): A valid `MarginCacheArguments` object. client (dask.distributed.Client): A dask distributed client object. """ + with tqdm(total=5, desc="Planning ", disable=not args.progress_bar) as step_progress: # determine which order to generate margin pixels for - partition_stats = args.catalog.partition_info.get_healpix_pixels() + partition_stats = args.catalog.partition_info.get_healpix_pixels() + step_progress.update(1) - # get the negative tree pixels - negative_pixels = args.catalog.generate_negative_tree_pixels() + # get the negative tree pixels + negative_pixels = args.catalog.generate_negative_tree_pixels() + step_progress.update(1) - combined_pixels = partition_stats + negative_pixels + combined_pixels = partition_stats + negative_pixels + step_progress.update(1) - margin_pairs = _find_partition_margin_pixel_pairs(combined_pixels, args.margin_order) + margin_pairs = _find_partition_margin_pixel_pairs(combined_pixels, args.margin_order) + step_progress.update(1) - _create_margin_directory(combined_pixels, args.catalog_path) + _create_margin_directory(combined_pixels, args.catalog_path) + step_progress.update(1) _map_to_margin_shards( client=client, diff --git a/tests/hipscat_import/margin_cache/test_margin_cache.py b/tests/hipscat_import/margin_cache/test_margin_cache.py index b489375e..efe472a0 100644 --- a/tests/hipscat_import/margin_cache/test_margin_cache.py +++ b/tests/hipscat_import/margin_cache/test_margin_cache.py @@ -40,7 +40,7 @@ def test_margin_cache_gen(small_sky_source_catalog, tmp_path, dask_client): assert all(data[PartitionInfo.METADATA_ORDER_COLUMN_NAME] == norder) assert all(data[PartitionInfo.METADATA_PIXEL_COLUMN_NAME] == npix) - assert all(data[PartitionInfo.METADATA_DIR_COLUMN_NAME] == int(npix / 10000) * 10000) + assert all(data[PartitionInfo.METADATA_DIR_COLUMN_NAME] == 0) assert data.dtypes[PartitionInfo.METADATA_ORDER_COLUMN_NAME] == np.uint8 assert data.dtypes[PartitionInfo.METADATA_DIR_COLUMN_NAME] == np.uint64