Enable adaptive bitround

Author: thodson-usgs

numcodecs/bitround.py

@@ -1,5 +1,6 @@
 import numpy as np
 
+from typing import Callable
 
 from .abc import Codec
 from .compat import ensure_ndarray_like, ndarray_copy
@@ -29,19 +30,30 @@ class BitRound(Codec):
     Parameters
     ----------
 
-    keepbits: int
+    keepbits: int or function
         The number of bits of the mantissa to keep. The range allowed
         depends on the dtype input data. If keepbits is
         equal to the maximum allowed for the data type, this is equivalent
-        to no transform.
+        to no transform. Alternatively, pass a function to determine the
+        number of bits to keep from the input data. The function should
+        take a single argument, the input data, and return an integer
+        specifying the number of bits to keep.
     """
 
     codec_id = 'bitround'
 
-    def __init__(self, keepbits: int):
-        if keepbits < 0:
+    def __init__(self, keepbits: [int, Callable]):
+        if isinstance(keepbits, int) and keepbits < 0:
             raise ValueError("keepbits must be zero or positive")
-        self.keepbits = keepbits
+
+        elif isinstance(keepbits, int):
+            self.keepbits = [lambda x: keepbits]
+
+        elif isinstance(keepbits, Callable):
+            self.keepbits = keepbits
+
+        else:
+            raise TypeError("keepbits must be an integer or function")
 
     def encode(self, buf):
         """Create int array by rounding floating-point data
@@ -56,12 +68,13 @@ def encode(self, buf):
         # cast float to int type of same width (preserve endianness)
         a_int_dtype = np.dtype(a.dtype.str.replace("f", "i"))
         all_set = np.array(-1, dtype=a_int_dtype)
-        if self.keepbits == bits:
+        buf_keepbits = self.keepbits(buf)
+        if buf_keepbits == bits:
             return a
-        if self.keepbits > bits:
+        if buf_keepbits > bits:
             raise ValueError("Keepbits too large for given dtype")
         b = a.view(a_int_dtype)
-        maskbits = bits - self.keepbits
+        maskbits = bits - buf_keepbits
         mask = (all_set >> maskbits) << maskbits
         half_quantum1 = (1 << (maskbits - 1)) - 1
         b += ((b >> maskbits) & 1) + half_quantum1