perf:Remove magic number functions

Author: mkitti

src/zarr/core/indexing.py

@@ -1467,141 +1467,6 @@ def decode_morton(z: int, chunk_shape: tuple[int, ...]) -> tuple[int, ...]:
     return tuple(out)
 
 
-# Magic numbers for 2D Morton decode: extract every 2nd bit and compact them.
-# Bits for dimension d are at positions d, d+2, d+4, ...
-_MORTON_2D_MASKS: tuple[int, ...] = (
-    0x5555555555555555,  # Extract mask: bits 0, 2, 4, ...
-    0x3333333333333333,  # Compact stage 1
-    0x0F0F0F0F0F0F0F0F,  # Compact stage 2
-    0x00FF00FF00FF00FF,  # Compact stage 3
-    0x0000FFFF0000FFFF,  # Compact stage 4
-)
-
-# Magic numbers for 3D Morton decode: extract every 3rd bit and compact them.
-# Bits for dimension d are at positions d, d+3, d+6, d+9, ...
-# Uses uint64 to handle large mask values correctly with numpy.
-_MORTON_3D_MASKS: tuple[np.uint64, ...] = (
-    np.uint64(0x1249249249249249),  # Extract mask: bits 0, 3, 6, 9, ...
-    np.uint64(0x30C30C30C30C30C3),  # Compact stage 1
-    np.uint64(0xF00F00F00F00F00F),  # Compact stage 2
-    np.uint64(0x00FF0000FF0000FF),  # Compact stage 3
-    np.uint64(0x00FF00000000FFFF),  # Compact stage 4
-    np.uint64(0x00000000001FFFFF),  # Compact stage 5
-)
-
-# Magic numbers for 4D Morton decode: extract every 4th bit and compact them.
-# Bits for dimension d are at positions d, d+4, d+8, d+12, ...
-_MORTON_4D_MASKS: tuple[np.uint64, ...] = (
-    np.uint64(0x1111111111111111),  # Extract mask: bits 0, 4, 8, 12, ...
-    np.uint64(0x0303030303030303),  # Compact stage 1
-    np.uint64(0x000F000F000F000F),  # Compact stage 2
-    np.uint64(0x000000FF000000FF),  # Compact stage 3
-    np.uint64(0x000000000000FFFF),  # Compact stage 4
-)
-
-# Magic numbers for 5D Morton decode: extract every 5th bit and compact them.
-# Bits for dimension d are at positions d, d+5, d+10, d+15, ...
-# Max 12 bits per dimension (64 // 5 = 12), supporting values up to 4095.
-_MORTON_5D_MASKS: tuple[np.uint64, ...] = (
-    np.uint64(0x1084210842108421),  # Extract mask: bits 0, 5, 10, 15, ...
-    np.uint64(0x0318C6318C6318C63),  # Compact stage 1
-    np.uint64(0xF03C0F03C0F03C0F),  # Compact stage 2
-    np.uint64(0x0000FF000FF000FF),  # Compact stage 3
-    np.uint64(0x0000000000000FFF),  # Compact stage 4 (12 bits max)
-)
-
-
-def _decode_morton_2d(z: npt.NDArray[np.intp]) -> npt.NDArray[np.intp]:
-    """Decode 2D Morton codes using magic number bit manipulation.
-
-    This extracts interleaved x,y coordinates from Morton codes using
-    parallel bit operations instead of bit-by-bit loops.
-    """
-    out = np.zeros((len(z), 2), dtype=np.intp)
-
-    # Extract x (bits 0, 2, 4, ...) and compact to (bits 0, 1, 2, ...)
-    x = z & _MORTON_2D_MASKS[0]
-    x = (x | (x >> 1)) & _MORTON_2D_MASKS[1]
-    x = (x | (x >> 2)) & _MORTON_2D_MASKS[2]
-    x = (x | (x >> 4)) & _MORTON_2D_MASKS[3]
-    x = (x | (x >> 8)) & _MORTON_2D_MASKS[4]
-    out[:, 0] = x
-
-    # Extract y (bits 1, 3, 5, ...) and compact
-    y = (z >> 1) & _MORTON_2D_MASKS[0]
-    y = (y | (y >> 1)) & _MORTON_2D_MASKS[1]
-    y = (y | (y >> 2)) & _MORTON_2D_MASKS[2]
-    y = (y | (y >> 4)) & _MORTON_2D_MASKS[3]
-    y = (y | (y >> 8)) & _MORTON_2D_MASKS[4]
-    out[:, 1] = y
-
-    return out
-
-
-def _decode_morton_3d(z: npt.NDArray[np.intp]) -> npt.NDArray[np.intp]:
-    """Decode 3D Morton codes using magic number bit manipulation.
-
-    This extracts interleaved x,y,z coordinates from Morton codes using
-    parallel bit operations instead of bit-by-bit loops.
-    """
-    # Convert to uint64 for bitwise operations with large masks
-    z_u64 = z.astype(np.uint64)
-    out = np.zeros((len(z), 3), dtype=np.intp)
-
-    for dim in range(3):
-        x = (z_u64 >> dim) & _MORTON_3D_MASKS[0]
-        x = (x ^ (x >> 2)) & _MORTON_3D_MASKS[1]
-        x = (x ^ (x >> 4)) & _MORTON_3D_MASKS[2]
-        x = (x ^ (x >> 8)) & _MORTON_3D_MASKS[3]
-        x = (x ^ (x >> 16)) & _MORTON_3D_MASKS[4]
-        x = (x ^ (x >> 32)) & _MORTON_3D_MASKS[5]
-        out[:, dim] = x
-
-    return out
-
-
-def _decode_morton_4d(z: npt.NDArray[np.intp]) -> npt.NDArray[np.intp]:
-    """Decode 4D Morton codes using magic number bit manipulation.
-
-    This extracts interleaved x,y,z,w coordinates from Morton codes using
-    parallel bit operations instead of bit-by-bit loops.
-    """
-    # Convert to uint64 for bitwise operations with large masks
-    z_u64 = z.astype(np.uint64)
-    out = np.zeros((len(z), 4), dtype=np.intp)
-
-    for dim in range(4):
-        x = (z_u64 >> dim) & _MORTON_4D_MASKS[0]
-        x = (x ^ (x >> 3)) & _MORTON_4D_MASKS[1]
-        x = (x ^ (x >> 6)) & _MORTON_4D_MASKS[2]
-        x = (x ^ (x >> 12)) & _MORTON_4D_MASKS[3]
-        x = (x ^ (x >> 24)) & _MORTON_4D_MASKS[4]
-        out[:, dim] = x
-
-    return out
-
-
-def _decode_morton_5d(z: npt.NDArray[np.intp]) -> npt.NDArray[np.intp]:
-    """Decode 5D Morton codes using magic number bit manipulation.
-
-    This extracts interleaved coordinates from Morton codes using
-    parallel bit operations instead of bit-by-bit loops.
-    """
-    # Convert to uint64 for bitwise operations with large masks
-    z_u64 = z.astype(np.uint64)
-    out = np.zeros((len(z), 5), dtype=np.intp)
-
-    for dim in range(5):
-        x = (z_u64 >> dim) & _MORTON_5D_MASKS[0]
-        x = (x ^ (x >> 4)) & _MORTON_5D_MASKS[1]
-        x = (x ^ (x >> 8)) & _MORTON_5D_MASKS[2]
-        x = (x ^ (x >> 16)) & _MORTON_5D_MASKS[3]
-        x = (x ^ (x >> 32)) & _MORTON_5D_MASKS[4]
-        out[:, dim] = x
-
-    return out
-
-
 def decode_morton_vectorized(
     z: npt.NDArray[np.intp], chunk_shape: tuple[int, ...]
 ) -> npt.NDArray[np.intp]:
@@ -1622,20 +1487,6 @@ def decode_morton_vectorized(
     n_dims = len(chunk_shape)
     bits = tuple((c - 1).bit_length() for c in chunk_shape)
 
-    # Use magic number optimization for 2D/3D/4D/5D with uniform bit widths.
-    # Bit limits are floor(64 / n_dims): 2D=32, 3D=21, 4D=16, 5D=12 bits per dimension.
-    if len(set(bits)) == 1:  # All dimensions have same bit width
-        max_bits = bits[0] if bits else 0
-        if n_dims == 2 and max_bits <= 32:
-            return _decode_morton_2d(z)
-        if n_dims == 3 and max_bits <= 21:
-            return _decode_morton_3d(z)
-        if n_dims == 4 and max_bits <= 16:
-            return _decode_morton_4d(z)
-        if n_dims == 5 and max_bits <= 12:
-            return _decode_morton_5d(z)
-
-    # Fall back to generic bit-by-bit decoding
     max_coords_bits = max(bits) if bits else 0
     out = np.zeros((len(z), n_dims), dtype=np.intp)