from typing import TYPE_CHECKING, Any, List, Optional, Tuple if TYPE_CHECKING: from ..classes.generated import PackedBitVector def reshape(data: list, shape: Optional[Tuple[int, ...]] = None) -> List[Any]: if shape is None: return data if len(shape) == 1: m = shape[0] return [data[i : i + m] for i in range(0, len(data), m)] elif len(shape) == 2: m, n = shape return [[data[i + j : i + j + n] for j in range(0, m * n, n)] for i in range(0, len(data), m * n)] else: raise ValueError("Invalid shape") def unpack_ints( packed: "PackedBitVector", start: int = 0, count: Optional[int] = None, shape: Optional[Tuple[int, ...]] = None, ) -> List[Any]: assert packed.m_BitSize is not None m_BitSize = packed.m_BitSize m_Data = packed.m_Data bitPos = m_BitSize * start indexPos = bitPos // 8 bitPos %= 8 if count is None: count = packed.m_NumItems # if m_BitSize <= 8: # dtype = np.uint8 # elif m_BitSize <= 16: # dtype = np.uint16 # elif m_BitSize <= 32: # dtype = np.uint32 # elif m_BitSize <= 64: # dtype = np.uint64 # else: # raise ValueError("Invalid bit size") # data = np.zeros(packed.m_NumItems, dtype=dtype) data = [0] * count for i in range(count): bits = 0 value = 0 while bits < m_BitSize: value |= (m_Data[indexPos] >> bitPos) << bits num = min(m_BitSize - bits, 8 - bitPos) bitPos += num bits += num if bitPos == 8: indexPos += 1 bitPos = 0 data[i] = value & ((1 << m_BitSize) - 1) return reshape(data, shape) def unpack_floats( packed: "PackedBitVector", start: int = 0, count: Optional[int] = None, shape: Optional[Tuple[int, ...]] = None, ) -> List[Any]: assert packed.m_BitSize is not None and packed.m_Range is not None and packed.m_Start is not None # avoid zero division of scale if packed.m_BitSize == 0: quantized = [packed.m_Start] * (packed.m_NumItems if count is None else count) else: # read as int and cast up to double to prevent loss of precision quantized_f64 = unpack_ints(packed, start, count) scale = packed.m_Range / ((1 << packed.m_BitSize) - 1) quantized = [x * scale + packed.m_Start for x in quantized_f64] return reshape(quantized, shape) # def pack_ints( # data: npt.NDArray[np.uint], bitsize: Optional[int] = 0 # ) -> PackedBitVector: # # ensure that the data type is unsigned # assert "uint" in data.dtype.name # m_NumItems = data.size # maxi = data.max() # # Prevent overflow # if bitsize: # m_BitSize = bitsize # else: # m_BitSize = (32 if maxi == 0xFFFFFFFF else np.ceil(np.log2(maxi + 1))) % 256 # m_Data = np.zeros((m_NumItems * m_BitSize + 7) // 8, dtype=np.uint8) # indexPos = 0 # bitPos = 0 # for x in data: # bits = 0 # while bits < m_BitSize: # m_Data[indexPos] |= (x >> bits) << bitPos # num = min(m_BitSize - bits, 8 - bitPos) # bitPos += num # bits += num # if bitPos == 8: # indexPos += 1 # bitPos = 0 # return PackedBitVector(m_NumItems=m_NumItems, m_BitSize=m_BitSize, m_Data=m_Data) # def pack_floats( # data: npt.NDArray[np.floating[Any]], # bitsize: Optional[int] = None, # ) -> PackedBitVector: # min = data.min() # max = data.max() # range = max - min # data_f64 = data.astype(np.float64) # # rebase to 0 # data_f64 -= min # # scale to [0, 1] # data_f64 /= range # # quantize to [0, 2^bit_size - 1] # bitsize = bitsize or max(data.itemsize, 32) # assert bitsize is not None # data_f64 *= (1 << bitsize) - 1 # # pack the data # packed = pack_ints(data_f64.astype(np.uint32), bitsize) # packed.m_Start = min # packed.m_Range = range # return packed __all__ = ("unpack_ints", "unpack_floats")