Rework chunking computation and expand test coverage

Author: BrianMichell

src/mdio/converters/segy.py

@@ -506,35 +506,125 @@ def _calculate_live_mask_chunksize(grid: Grid) -> Sequence[int]:
         grid: The grid to calculate the chunksize for.
 
     Returns:
-        A sequence of integers representing the optimal chunk size for each dimension of the grid.
+        A sequence of integers representing the optimal chunk size for each dimension
+        of the grid.
     """
-    # Use nbytes for the initial check, since we are limited by Blosc's maximum
-    # chunk size in bytes.
-    if grid.live_mask.nbytes < INT32_MAX:
-        # Base case where we don't need to chunk the live mask
-        return grid.live_mask.shape
-
-    # Calculate the optimal chunksize for the live mask
-    total_elements = np.prod(grid.shape[:-1])  # Exclude sample dimension
-    num_chunks = np.ceil(total_elements / INT32_MAX).astype(int)
-
-    # Calculate chunk size for each dimension
-    chunks = []
-    remaining_elements = total_elements
-
-    for dim_size in grid.shape[:-1]:  # Exclude sample dimension
-        # Calculate how many chunks we need in this dimension
-        # We want to distribute chunks evenly across dimensions
-        dim_chunks = max(
-            1,
-            int(
-                np.ceil(
-                    dim_size / np.ceil(np.power(num_chunks, 1 / len(grid.shape[:-1])))
-                )
-            ),
-        )
-        chunk_size = int(np.ceil(dim_size / dim_chunks))
-        chunks.append(chunk_size)
-        remaining_elements //= dim_chunks
+    return _calculate_optimal_chunksize(grid.live_mask, INT32_MAX)
+
+
+def _calculate_optimal_chunksize(  # noqa: C901
+    volume: np.ndarray | zarr.Array, n_bytes: int
+) -> Sequence[int]:
+    """Calculate a uniform chunk shape for an N-dimensional data volume such that...
+
+    0. The product of the chunk dimensions multiplied by the element size does not
+       exceed n_bytes.
+    1. The chunk shape is "regular" – each chunk dimension is a divisor of the
+       overall volume shape.
+    2. If an exact match is impossible, the chunk shape chosen maximizes the number of
+       elements (minimizing the unused bytes).
+    3. The computation is efficient.
+
+    The computation efficiency is broken down as follows:
+
+    - Divisor Computation: For each of the N dimensions (assume size ~ n), it checks
+      up to n numbers, so this part is roughly O(N * n).
+      For example, if you have a 3D array where each dimension is about 100,
+      it does around 3*100 = 300 steps.
+    - DFS Search: In the worst-case, the DFS explores about D choices per dimension
+      (D = average number of divisors) leading to O(D^N) combinations.
+      In practice, D is small (often < 10), so for a 2D array this is around 10^2
+      (about 100 combinations) and for a 3D array about 10^3 (roughly 1,000 combinations).
+      Since N is typically small (often <6), this exponential term behaves like a
+      constant factor.
+
+    Args:
+      volume : np.ndarray | zarr.Array
+          An N-dimensional array-like object (e.g. np.ndarray or zarr array).
+      n_bytes : int
+          Maximum allowed number of bytes per chunk (>= 1).
 
-    return tuple(chunks)
+    Returns:
+      Sequence[int]
+          A tuple representing the optimal chunk shape (number of elements along each axis).
+
+    Raises:
+      ValueError if n_bytes is less than the number of bytes of one element.
+    """
+    # Get volume shape and element size.
+    shape = volume.shape
+
+    if volume.size == 0:
+        logging.warning("Chunking calculation received empty volume shape...")
+        return volume.shape
+
+    itemsize = volume.dtype.itemsize
+
+    # Maximum number of elements that can fit in a chunk
+    # (we ignore any extra bytes; must not exceed n_bytes).
+    max_elements_allowed = n_bytes // itemsize
+    if max_elements_allowed < 1:
+        raise ValueError("n_bytes is too small to hold even one element of the volume.")
+
+    n_dims = len(shape)
+
+    def get_divisors(n: int) -> list[int]:
+        """Return a sorted list of all positive divisors of n.
+
+        Args:
+            n: The number to compute the divisors of.
+
+        Returns:
+            A sorted list of all positive divisors of n.
+        """
+        divs = []
+        # It is efficient enough for typical dimension sizes.
+        for i in range(1, n + 1):
+            if n % i == 0:
+                divs.append(i)
+        return sorted(divs)
+
+    # For each dimension, compute the list of allowed chunk sizes (divisors).
+    divisors_list = [get_divisors(d) for d in shape]
+
+    # For pruning: precompute the maximum possible product achievable from axis i to N-1.
+    # This is the product of the maximum divisors for each remaining axis.
+    max_possible = [1] * (n_dims + 1)
+    for i in range(n_dims - 1, -1, -1):
+        max_possible[i] = max(divisors_list[i]) * max_possible[i + 1]
+
+    best_product = 0
+    best_combination = [None] * n_dims
+    current_chunk = [None] * n_dims
+
+    def dfs(dim: int, current_product: int) -> None:
+        """Depth-first search to find the optimal chunk shape.
+
+        Args:
+            dim: The current dimension to process.
+            current_product: The current product of the chunk dimensions.
+        """
+        nonlocal best_product
+        # If all dimensions have been processed, update best combination if needed.
+        if dim == n_dims:
+            if current_product > best_product:
+                best_product = current_product
+                best_combination[:] = current_chunk[:]
+            return
+
+        # Prune branches: even if we take the maximum allowed for all remaining dimensions,
+        # if we cannot exceed best_product, then skip.
+        if current_product * max_possible[dim] < best_product:
+            return
+
+        # Iterate over allowed divisors for the current axis,
+        # trying larger candidates first so that high products are found early.
+        for candidate in sorted(divisors_list[dim], reverse=True):
+            new_product = current_product * candidate
+            if new_product > max_elements_allowed:
+                continue  # This candidate would exceed the byte restriction.
+            current_chunk[dim] = candidate
+            dfs(dim + 1, new_product)
+
+    dfs(0, 1)
+    return tuple(best_combination)

tests/unit/test_live_mask_chunksize.py

@@ -1,9 +1,10 @@
 """Test live mask chunk size calculation."""
 
 import numpy as np
+import pytest
 
-from mdio.constants import INT32_MAX
 from mdio.converters.segy import _calculate_live_mask_chunksize
+from mdio.converters.segy import _calculate_optimal_chunksize
 from mdio.core import Dimension
 from mdio.core import Grid
 
@@ -20,7 +21,7 @@ def test_small_grid_no_chunking():
     grid.live_mask = np.ones((100, 100), dtype=bool)
 
     result = _calculate_live_mask_chunksize(grid)
-    assert result == -1
+    assert result == (100, 100)
 
 
 def test_large_2d_grid_chunking():
@@ -37,13 +38,8 @@ def test_large_2d_grid_chunking():
 
     result = _calculate_live_mask_chunksize(grid)
 
-    # Calculate expected values
-    total_elements = 50000 * 50000
-    num_chunks = np.ceil(total_elements / INT32_MAX).astype(int)
-    dim_chunks = int(np.ceil(50000 / np.ceil(np.power(num_chunks, 1 / 2))))
-    expected_chunk_size = int(np.ceil(50000 / dim_chunks))
-
-    assert result == (expected_chunk_size, expected_chunk_size)
+    # TODO(BrianMichell): Avoid magic numbers.
+    assert result == (50000, 25000)
 
 
 def test_large_3d_grid_chunking():
@@ -62,12 +58,13 @@ def test_large_3d_grid_chunking():
     result = _calculate_live_mask_chunksize(grid)
 
     # Calculate expected values
-    total_elements = 1500 * 1500 * 1500
-    num_chunks = np.ceil(total_elements / INT32_MAX).astype(int)
-    dim_chunks = int(np.ceil(1500 / np.ceil(np.power(num_chunks, 1 / 3))))
-    expected_chunk_size = int(np.ceil(1500 / dim_chunks))
+    # total_elements = 1500 * 1500 * 1500
+    # num_chunks = np.ceil(total_elements / INT32_MAX).astype(int)
+    # dim_chunks = int(np.ceil(1500 / np.ceil(np.power(num_chunks, 1 / 3))))
+    # expected_chunk_size = int(np.ceil(1500 / dim_chunks))
 
-    assert result == (expected_chunk_size, expected_chunk_size, expected_chunk_size)
+    # assert result == (expected_chunk_size, expected_chunk_size, expected_chunk_size)
+    assert result == (1500, 1500, 750)
 
 
 def test_uneven_dimensions_chunking():
@@ -84,14 +81,7 @@ def test_uneven_dimensions_chunking():
     grid.live_mask = np.ones((50000, 50000), dtype=bool)
 
     result = _calculate_live_mask_chunksize(grid)
-
-    # Calculate expected values
-    total_elements = 50000 * 50000
-    num_chunks = np.ceil(total_elements / INT32_MAX).astype(int)
-    dim_chunks = int(np.ceil(50000 / np.ceil(np.power(num_chunks, 1 / 2))))
-    expected_chunk_size = int(np.ceil(50000 / dim_chunks))
-
-    assert result == (expected_chunk_size, expected_chunk_size)
+    assert result == (50000, 25000)
 
 
 def test_prestack_land_survey_chunking():
@@ -114,21 +104,7 @@ def test_prestack_land_survey_chunking():
     grid.live_mask = np.ones((1000, 1000, 100, 36), dtype=bool)
 
     result = _calculate_live_mask_chunksize(grid)
-
-    # Calculate expected values
-    total_elements = 1000 * 1000 * 100 * 36
-    num_chunks = np.ceil(total_elements / INT32_MAX).astype(int)
-    dim_chunks = int(np.ceil(1000 / np.ceil(np.power(num_chunks, 1 / 4))))
-    expected_chunk_size = int(np.ceil(1000 / dim_chunks))
-
-    # For a 4D grid, we expect chunk sizes to be distributed across all dimensions
-    # The chunk size should be the same for all dimensions since they're all equally important
-    assert result == (
-        expected_chunk_size,
-        expected_chunk_size,
-        expected_chunk_size,
-        expected_chunk_size,
-    )
+    assert result == (1000, 1000, 100, 18)
 
 
 def test_edge_case_empty_grid():
@@ -142,4 +118,91 @@ def test_edge_case_empty_grid():
     grid.live_mask = np.zeros((0, 0), dtype=bool)
 
     result = _calculate_live_mask_chunksize(grid)
-    assert result == -1  # Empty grid shouldn't need chunking
+    assert result == (0, 0)
+
+
+# Additional tests for _calculate_optimal_chunksize function
+def test_empty_volume():
+    """Test that an empty volume returns its shape."""
+    empty_arr = np.zeros((0, 10), dtype=np.int8)
+    result = _calculate_optimal_chunksize(empty_arr, 100)
+    assert result == (0, 10)
+
+
+def test_nbytes_too_small():
+    """Test that a too-small n_bytes value raises a ValueError."""
+    arr = np.zeros((10,), dtype=np.int8)  # itemsize is 1
+    with pytest.raises(
+        ValueError, match=r"n_bytes is too small to hold even one element"
+    ):
+        _calculate_optimal_chunksize(arr, 0)
+
+
+def test_one_dim_full_chunk():
+    """Test one-dimensional volume where the whole dimension can be used as chunk."""
+    arr = np.zeros((100,), dtype=np.int8)
+    # With n_bytes = 100, max_elements_allowed = 100, thus optimal chunk should be (100,)
+    result = _calculate_optimal_chunksize(arr, 100)
+    assert result == (100,)
+
+
+def test_two_dim_optimal():
+    """Test two-dimensional volume with limited n_bytes.
+
+    For a shape of (8,6) with n_bytes=20, the optimal chunk is expected to be (8,2).
+    """
+    arr = np.zeros((8, 6), dtype=np.int8)
+    result = _calculate_optimal_chunksize(arr, 20)
+    assert result == (8, 2)
+
+
+def test_three_dim_optimal():
+    """Test three-dimensional volume optimal chunk calculation.
+
+    For a shape of (9,6,4) with n_bytes=100, the expected chunk is (9,2,4).
+    """
+    arr = np.zeros((9, 6, 4), dtype=np.int8)
+    result = _calculate_optimal_chunksize(arr, 100)
+    assert result == (9, 2, 4)
+
+
+def test_minimal_chunk_for_large_dtype():
+    """Test that n_bytes forcing minimal chunking returns all ones.
+
+    Using int32 (itemsize=4) with shape (4,5) and n_bytes=4 yields (1,1).
+    """
+    arr = np.zeros((4, 5), dtype=np.int32)
+    result = _calculate_optimal_chunksize(arr, 4)
+    assert result == (1, 1)
+
+
+def test_large_nbytes():
+    """Test that a very large n_bytes returns the full volume shape as the optimal chunk."""
+    arr = np.zeros((10, 10), dtype=np.int8)
+    result = _calculate_optimal_chunksize(arr, 1000)
+    assert result == (10, 10)
+
+
+def test_two_dim_non_int8():
+    """Test with a non-int8 dtype where n_bytes exactly covers the full volume in bytes."""
+    arr = np.zeros((6, 8), dtype=np.int16)  # int16 has itemsize 2
+    # Total bytes of full volume = 6*8*2 = 96, so optimal chunk should be (6,8)
+    result = _calculate_optimal_chunksize(arr, 96)
+    assert result == (6, 8)
+
+
+def test_irregular_dimensions():
+    """Test volume with prime dimensions where divisors are limited.
+
+    For shape (7,5) with n_bytes=35, optimal chunk should be (7,5) since 7*5 = 35.
+    """
+    arr = np.zeros((7, 5), dtype=np.int8)
+    result = _calculate_optimal_chunksize(arr, 35)
+    assert result == (7, 5)
+
+
+def test_primes():
+    """Test volume with prime dimensions where divisors are limited."""
+    arr = np.zeros((7, 5), dtype=np.int8)
+    result = _calculate_optimal_chunksize(arr, 23)
+    assert result == (7, 5)