jan/klamath
1
0
Fork 0
Code Issues Pull Requests Releases Wiki Activity

Allow NDArray inputs to pack_* and avoid unnecesary copies

Browse Source
This commit is contained in:
Jan Petykiewicz 2024-07-28 23:04:12 -07:00
parent dc58159cdf
commit e7e42a2ef8
1 changed files with 6 additions and 6 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

12
klamath/basic.py
View File

@ -92,15 +92,15 @@ def pack_bitarray(data: int) -> bytes:
return struct.pack('>H', data) return struct.pack('>H', data)
def pack_int2(data: Sequence[int]) -> bytes: def pack_int2(data: NDArray[numpy.integer] | Sequence[int] | int) -> bytes:
arr = numpy.array(data) arr = numpy.array(data, copy=False)
if (arr > 32767).any() or (arr < -32768).any(): if (arr > 32767).any() or (arr < -32768).any():
raise KlamathError(f'int2 data out of range: {arr}') raise KlamathError(f'int2 data out of range: {arr}')
return arr.astype('>i2').tobytes() return arr.astype('>i2').tobytes()
def pack_int4(data: Sequence[int]) -> bytes: def pack_int4(data: NDArray[numpy.integer] | Sequence[int] | int) -> bytes:
arr = numpy.array(data) arr = numpy.array(data, copy=False)
if (arr > 2147483647).any() or (arr < -2147483648).any(): if (arr > 2147483647).any() or (arr < -2147483648).any():
raise KlamathError(f'int4 data out of range: {arr}') raise KlamathError(f'int4 data out of range: {arr}')
return arr.astype('>i4').tobytes() return arr.astype('>i4').tobytes()
@ -164,8 +164,8 @@ def encode_real8(fnums: NDArray[numpy.float64]) -> NDArray[numpy.uint64]:
return real8.astype(numpy.uint64, copy=False) return real8.astype(numpy.uint64, copy=False)
def pack_real8(data: Sequence[float]) -> bytes: def pack_real8(data: NDArray[numpy.floating] | Sequence[float] | float) -> bytes:
return encode_real8(numpy.array(data)).astype('>u8').tobytes() return encode_real8(numpy.array(data, copy=False)).astype('>u8').tobytes()
def pack_ascii(data: bytes) -> bytes: def pack_ascii(data: bytes) -> bytes: