Add math helpers to align to nearest 2**N

odc/geo/math.py

@@ -7,7 +7,7 @@
 
 Minimal dependencies in this module.
 """
-from math import ceil, floor, fmod, isfinite
+from math import ceil, floor, fmod, isfinite, log2
 from typing import (
     Any,
     Iterator,
@@ -122,6 +122,31 @@ def align_up(x: int, align: int) -> int:
     return align_down(x + (align - 1), align)
 
 
+def align_up_pow2(x: int) -> int:
+    """
+    Align up to the nearest power of 2.
+
+    :return:
+       Smallest ``y`` such that ``y >= x`` and ``y == 2**n`` for some integer ``n``.
+    """
+    if x <= 0:
+        return 1
+    return 2 ** int(ceil(log2(x)))
+
+
+def align_down_pow2(x: int) -> int:
+    """
+    Align down to the nearest power of 2.
+
+    :return:
+       Largest ``y`` such that ``y <= x`` and ``y == 2**n`` for some integer ``n``.
+    """
+    y = align_up_pow2(x)
+    if y > x:
+        y = y // 2
+    return y
+
+
 def clamp(x, lo, up):
     """Clamp ``x`` to be ``lo <= x <= up``."""
     assert lo <= up

tests/test_math.py

@@ -1,3 +1,4 @@
+import math
 from typing import Tuple
 
 import numpy as np
@@ -10,7 +11,9 @@
     Poly2d,
     affine_from_axis,
     align_down,
+    align_down_pow2,
     align_up,
+    align_up_pow2,
     apply_affine,
     data_resolution_and_offset,
     is_almost_int,
@@ -358,3 +361,33 @@ def test_poly2d_not_enough_points():
 
     with pytest.raises(ValueError):
         _ = Poly2d.fit(pts[:2], pts[:2])
+
+
+@pytest.mark.parametrize(
+    "x",
+    [
+        *[2**n for n in range(20)],
+        *[2**n - 1 for n in range(20)],
+        *[2**n + 1 for n in range(20)],
+    ],
+)
+def test_align_up_pow2(x: int):
+    y = align_up_pow2(x)
+    assert isinstance(y, int)
+    assert y >= x
+    assert 2 ** int(math.log2(y)) == y
+
+
+@pytest.mark.parametrize(
+    "x",
+    [
+        *[2**n for n in range(20)],
+        *[2**n - 1 for n in range(1, 20)],
+        *[2**n + 1 for n in range(20)],
+    ],
+)
+def test_align_down_pow2(x: int):
+    y = align_down_pow2(x)
+    assert isinstance(y, int)
+    assert y <= x
+    assert 2 ** int(math.log2(y)) == y