Merge 9ac5477 into 53a04bc

Author: olivialynn

src/hats/pixel_math/partition_stats.py

@@ -72,7 +72,12 @@ def generate_histogram(
 
 
 def generate_alignment(
-    histogram, highest_order=10, lowest_order=0, threshold=1_000_000, drop_empty_siblings=False
+    row_count_histogram,
+    highest_order=10,
+    lowest_order=0,
+    threshold=1_000_000,
+    drop_empty_siblings=False,
+    mem_size_histogram=None,
 ):
     """Generate alignment from high order pixels to those of equal or lower order
 
@@ -83,9 +88,9 @@ def generate_alignment(
 
     Parameters
     ----------
-    histogram : np.array
-        one-dimensional numpy array of long integers where the
-        value at each index corresponds to the number of objects found at the healpix pixel.
+    row_count_histogram : np.array
+        one-dimensional numpy array of long integers where the value at each index corresponds to
+        the number of objects found at the healpix pixel.
     highest_order : int
         the highest healpix order (e.g. 5-10) (Default value = 10)
     lowest_order : int
@@ -95,6 +100,10 @@ def generate_alignment(
         the maximum number of objects allowed in a single pixel (Default value = 1_000_000)
     drop_empty_siblings : bool
         if 3 of 4 pixels are empty, keep only the non-empty pixel (Default value = False)
+    mem_size_histogram : np.array or None
+        one-dimensional numpy array of long integers where the value at each index corresponds to
+        the memory size (in bytes) of objects found at the healpix pixel. If provided, this will be
+        used to determine the thresholding instead of the param `histogram`. (Default value = None)
 
     Returns
     -------
@@ -113,23 +122,41 @@ def generate_alignment(
         if the histogram is the wrong size, or some initial histogram bins
         exceed threshold.
     """
-    _validate_alignment_arguments(histogram, highest_order, lowest_order, threshold)
+    # Validate inputs.
+    _validate_alignment_arguments(
+        row_count_histogram, mem_size_histogram, highest_order, lowest_order, threshold
+    )
 
-    nested_sums = _get_nested_sums(histogram, highest_order, lowest_order)
+    # Generate nested sums.
+    nested_sums_row_count = _get_nested_sums(row_count_histogram, highest_order, lowest_order)
+    if mem_size_histogram is not None:
+        nested_sums_mem_size = _get_nested_sums(mem_size_histogram, highest_order, lowest_order)
+    else:
+        nested_sums_mem_size = None
 
+    # Generate alignment.
     if drop_empty_siblings:
-        return _get_alignment_dropping_siblings(nested_sums, highest_order, lowest_order, threshold)
-    return _get_alignment(nested_sums, highest_order, lowest_order, threshold)
+        return _get_alignment_dropping_siblings(
+            nested_sums_row_count, highest_order, lowest_order, threshold, nested_sums_mem_size
+        )
+    return _get_alignment(nested_sums_row_count, highest_order, lowest_order, threshold, nested_sums_mem_size)
 
 
-def _validate_alignment_arguments(histogram, highest_order, lowest_order, threshold):
-    if len(histogram) != hp.order2npix(highest_order):
+def _validate_alignment_arguments(
+    row_count_histogram, mem_size_histogram, highest_order, lowest_order, threshold
+):
+    if len(row_count_histogram) != hp.order2npix(highest_order):
         raise ValueError("histogram is not the right size")
     if lowest_order > highest_order:
         raise ValueError("lowest_order should be less than highest_order")
-    max_bin = np.amax(histogram)
-    if max_bin > threshold:
-        raise ValueError(f"single pixel count {max_bin} exceeds threshold {threshold}")
+    if mem_size_histogram is not None:
+        max_bin = np.amax(mem_size_histogram)
+        if max_bin > threshold:
+            raise ValueError(f"single pixel mem_size {max_bin} exceeds threshold {threshold}")
+    else:
+        max_bin = np.amax(row_count_histogram)
+        if max_bin > threshold:
+            raise ValueError(f"single pixel row count {max_bin} exceeds threshold {threshold}")
 
 
 def _get_nested_sums(histogram, highest_order, lowest_order):
@@ -150,16 +177,22 @@ def _get_nested_sums(histogram, highest_order, lowest_order):
     return nested_sums
 
 
-def _get_alignment(nested_sums, highest_order, lowest_order, threshold):
+def _get_alignment(nested_sums_row_count, highest_order, lowest_order, threshold, nested_sums_mem_size):
     """Method to aggregate pixels up to the threshold.
 
     Checks from low order (large areas), drilling down into higher orders (smaller areas) to
     find the appropriate order for an area of sky."""
+    if nested_sums_mem_size is not None:
+        nested_sums = nested_sums_mem_size
+    else:
+        nested_sums = nested_sums_row_count
+
+    # Initialize our alignment structure.
     nested_alignment = []
     for i in range(0, highest_order + 1):
         nested_alignment.append(np.full(hp.order2npix(i), None))
 
-    # work forward - determine if we should map to a lower order pixel, this pixel, or keep looking.
+    # Work forward - determine if we should map to a lower order pixel, this pixel, or keep looking.
     for read_order in range(lowest_order, highest_order + 1):
         parent_order = read_order - 1
         for index in range(0, len(nested_sums[read_order])):
@@ -173,16 +206,40 @@ def _get_alignment(nested_sums, highest_order, lowest_order, threshold):
             elif nested_sums[read_order][index] == 0:
                 continue
             elif nested_sums[read_order][index] <= threshold:
-                nested_alignment[read_order][index] = (
-                    read_order,
-                    index,
-                    nested_sums[read_order][index],
-                )
+                # For row_count mode, use tuple of (order, pixel, row_count)
+                if not nested_sums_mem_size:
+                    nested_alignment[read_order][index] = (
+                        read_order,
+                        index,
+                        nested_sums[read_order][index],
+                    )
+                # For mem_size mode, use tuple of (order, pixel, row_count, mem_size)
+                else:
+                    mem_size = nested_sums_mem_size[read_order][index]
+                    row_count = nested_sums_row_count[read_order][index]
+                    nested_alignment[read_order][index] = (
+                        read_order,
+                        index,
+                        row_count,
+                        mem_size,
+                    )
+
+    # We no longer need to store mem_size sums once the alignment has been constructed.
+    if nested_sums_mem_size is not None:
+        nested_alignment[highest_order] = np.array(
+            [
+                pixel_alignment[:3] if pixel_alignment else None
+                for pixel_alignment in nested_alignment[highest_order]
+            ],
+            dtype="object",
+        )
 
     return nested_alignment[highest_order]
 
 
-def _get_alignment_dropping_siblings(nested_sums, highest_order, lowest_order, threshold):
+def _get_alignment_dropping_siblings(
+    nested_sum_row_count, highest_order, lowest_order, threshold, nested_sums_mem_size
+):
     """Method to aggregate pixels up to the threshold that collapses completely empty pixels away.
 
     Checks from higher order (smaller areas) out to lower order (large areas). In this way, we are able to
@@ -209,6 +266,15 @@ def _get_alignment_dropping_siblings(nested_sums, highest_order, lowest_order, t
     - total number in cell is greater than the threshold
     - only one subcell contains values
     """
+    # If nested_sums_mem_size is provided, we're in mem_size mode (and thresholding by memory size).
+    # This means we'll want to use the mem_size sums to generate our alignment, but still keep track
+    # of the row counts for the output.
+    if nested_sums_mem_size is not None:
+        nested_sums = nested_sums_mem_size
+    else:
+        nested_sums = nested_sum_row_count
+
+    # Initialize our order map to the highest order.
     order_map = np.array(
         [highest_order if count > 0 else -1 for count in nested_sums[highest_order]], dtype=np.int32
     )
@@ -237,19 +303,22 @@ def _get_alignment_dropping_siblings(nested_sums, highest_order, lowest_order, t
         )
         for pixel_high_index, intended_order in enumerate(order_map)
     ]
+    # In both row_count and mem_size mode, use tuple of (order, pixel, row_count)
+    # as mem_size is no longer needed now that alignment has been constructed.
     nested_alignment = [
-        (tup[0], tup[1], nested_sums[tup[0]][tup[1]]) if tup else None for tup in nested_alignment
+        (tup[0], tup[1], nested_sum_row_count[tup[0]][tup[1]]) if tup else None for tup in nested_alignment
     ]
 
     return np.array(nested_alignment, dtype="object")
 
 
 def generate_incremental_alignment(
-    histogram: np.ndarray,
+    row_count_histogram: np.ndarray,
     existing_pixels: Sequence[tuple[int, int]],
     highest_order: int = 10,
     lowest_order: int = 0,
     threshold: int = 1_000_000,
+    mem_size_histogram: np.ndarray | None = None,
 ):
     """Generate alignment for an incremental catalog.
 
@@ -262,7 +331,7 @@ def generate_incremental_alignment(
 
     Parameters
     ----------
-    histogram : np.ndarray
+    row_count_histogram : np.ndarray
         one-dimensional numpy array of long integers where the
         value at each index corresponds to the number of objects
         found at the healpix pixel.
@@ -275,6 +344,10 @@ def generate_incremental_alignment(
         constrains the partitioning to prevent spatially large pixels. (Default value = 0)
     threshold : int
         the maximum number of objects allowed in a single pixel (Default value = 1_000_000)
+    mem_size_histogram : np.ndarray or None
+        one-dimensional numpy array of long integers where the value at each index corresponds to
+        the memory size (in bytes) of objects found at the healpix pixel. If provided, this will be
+        used to determine the thresholding instead of the param `histogram`. (Default value = None)
 
     Returns
     -------
@@ -287,9 +360,11 @@ def generate_incremental_alignment(
         - pixel number *at the above order*
         - the number of objects in the pixel
     """
-    _validate_alignment_arguments(histogram, highest_order, lowest_order, threshold)
+    _validate_alignment_arguments(
+        row_count_histogram, mem_size_histogram, highest_order, lowest_order, threshold
+    )
 
-    nested_sums = _get_nested_sums(histogram, highest_order, lowest_order)
+    nested_sums = _get_nested_sums(row_count_histogram, highest_order, lowest_order)
 
     tree = PixelTree.from_healpix(existing_pixels)
     if tree.tree_order > highest_order:

tests/hats/pixel_math/test_partition_stats.py

@@ -255,3 +255,159 @@ def test_incremental_alignment_highest_order_invalid():
         hist.generate_incremental_alignment(
             hist.empty_histogram(0), existing_pixels=[(1, 45)], highest_order=0
         )
+
+
+def test_generate_alignment_mem_size():
+    """Create alignment based on memory size histogram"""
+    initial_row_count_histogram = hist.empty_histogram(2)
+    filled_pixels = [11_000, 11, 14, 13, 5, 7, 8, 9, 11, 23, 4, 4, 17, 0, 1, 0]
+    initial_row_count_histogram[176:] = filled_pixels[:]
+
+    initial_mem_size_histogram = hist.empty_histogram(2)
+    filled_mem_sizes = [
+        4_000,
+        1_000,
+        1_000,
+        1_000,
+        5_000,
+        7_000,
+        8_000,
+        9_000,
+        1_000,
+        2_000,
+        4_000,
+        4_000,
+        7_000,
+        0,
+        1_000,
+        0,
+    ]
+    initial_mem_size_histogram[176:] = filled_mem_sizes[:]
+
+    # Generate alignment based on memory size thresholding.
+    result = hist.generate_alignment(
+        initial_row_count_histogram,
+        highest_order=2,
+        threshold=10_000,
+        mem_size_histogram=initial_mem_size_histogram,
+    )
+
+    expected = np.full(hp.order2npix(2), None)
+    tuples = [
+        (1, 44, 11038),
+        (1, 44, 11038),
+        (1, 44, 11038),
+        (1, 44, 11038),
+        (2, 180, 5),
+        (2, 181, 7),
+        (2, 182, 8),
+        (2, 183, 9),
+        (2, 184, 11),
+        (2, 185, 23),
+        (2, 186, 4),
+        (2, 187, 4),
+        (1, 47, 18),
+        (1, 47, 18),
+        (1, 47, 18),
+        (1, 47, 18),
+    ]
+    expected[176:192] = tuples
+
+    npt.assert_array_equal(result, expected)
+
+
+def test_generate_alignment_mem_size_dropping_siblings():
+    """Create alignment based on memory size histogram"""
+    initial_row_count_histogram = hist.empty_histogram(2)
+    filled_pixels = [0, 0, 14, 0, 5, 7, 8, 9, 11, 23, 4, 4, 17, 0, 1, 0]
+    initial_row_count_histogram[176:] = filled_pixels[:]
+
+    initial_mem_size_histogram = hist.empty_histogram(2)
+    filled_mem_sizes = [
+        0,
+        0,
+        4_000,
+        0,
+        5_000,
+        7_000,
+        8_000,
+        9_000,
+        1_000,
+        2_000,
+        4_000,
+        4_000,
+        7_000,
+        0,
+        1_000,
+        0,
+    ]
+    initial_mem_size_histogram[176:] = filled_mem_sizes[:]
+
+    # Generate alignment based on memory size thresholding.
+    result = hist.generate_alignment(
+        initial_row_count_histogram,
+        highest_order=2,
+        threshold=10_000,
+        drop_empty_siblings=True,
+        mem_size_histogram=initial_mem_size_histogram,
+    )
+
+    expected = np.full(hp.order2npix(2), None)
+    tuples = [
+        None,
+        None,
+        (2, 178, 14),
+        None,
+        (2, 180, 5),
+        (2, 181, 7),
+        (2, 182, 8),
+        (2, 183, 9),
+        (2, 184, 11),
+        (2, 185, 23),
+        (2, 186, 4),
+        (2, 187, 4),
+        (1, 47, 18),
+        (1, 47, 18),
+        (1, 47, 18),
+        (1, 47, 18),
+    ]
+    expected[176:192] = tuples
+
+    npt.assert_array_equal(result, expected)
+
+
+def test_generate_alignment_mem_size_exceeds_threshold():
+    """Create alignment based on memory size histogram"""
+    initial_row_count_histogram = hist.empty_histogram(2)
+    filled_pixels = [4, 11, 14, 13, 5, 7, 8, 9, 11, 23, 4, 4, 17, 0, 1, 0]
+    initial_row_count_histogram[176:] = filled_pixels[:]
+
+    initial_mem_size_histogram = hist.empty_histogram(2)
+    filled_mem_sizes = [
+        40_000,
+        1_000,
+        1_000,
+        1_000,
+        5_000,
+        7_000,
+        8_000,
+        9_000,
+        1_000,
+        2_000,
+        4_000,
+        4_000,
+        7_000,
+        0,
+        1_000,
+        0,
+    ]
+    initial_mem_size_histogram[176:] = filled_mem_sizes[:]
+
+    # We raise an error if any pixel exceeds the threshold.
+    with pytest.raises(ValueError, match="exceeds threshold"):
+        hist.generate_alignment(
+            initial_row_count_histogram,
+            highest_order=2,
+            threshold=5_000,
+            mem_size_histogram=initial_mem_size_histogram,
+        )