JIT in encoders

Author: dsehnal

ciftools/binary/decoder.py

@@ -13,7 +13,6 @@
 StringArrayEncoding,
 )
 
-
 defdecode_cif_data(encoded_data: EncodedCIFData) ->Union[np.ndarray, list[str]]:
 result=encoded_data["data"]
 forencodinginencoded_data["encoding"][::-1]:
@@ -48,7 +47,7 @@ def _decode_delta(data: np.ndarray, encoding: DeltaEncoding) -> np.ndarray:
 result[0] +=encoding["origin"]
 returnnp.cumsum(result, out=result)
 
-
+# TODO: JIT
 def_decode_integer_packing_signed(data: np.ndarray, encoding: IntegerPackingEncoding) ->np.ndarray:
 upper_limit=0x7Fifencoding["byteCount"] ==1else0x7FFF
 lower_limit=-upper_limit-1
@@ -70,6 +69,7 @@ def _decode_integer_packing_signed(data: np.ndarray, encoding: IntegerPackingEnc
 returnoutput
 
 
+# TODO: JIT
 def_decode_integer_packing_unsigned(data: np.ndarray, encoding: IntegerPackingEncoding) ->np.ndarray:
 upper_limit=0xFFifencoding["byteCount"] ==1else0xFFFF
 n=len(data)

ciftools/binary/encoder.py

@@ -2,6 +2,7 @@
 importsys
 fromtypingimportAny, List, Protocol, Union
 
+fromnumbaimportjit
 importnumpyasnp
 fromciftools.binary.data_typesimportDataType, DataTypeEnum
 fromciftools.binary.encoded_dataimportEncodedCIFData
@@ -105,8 +106,7 @@ def encode(self, data: np.ndarray, *args, **kwargs) -> EncodedCIFData:
 classIntegerPacking(BinaryCIFEncoder):
 defencode(self, data: np.ndarray) ->EncodedCIFData:
 
-# TODO: must be 32bit integer
-
+# TODO: must be 32bit integer?
 packing=_determine_packing(data)
 ifpacking.bytesPerElement==4:
 returnBYTE_ARRAY.encode(data)
@@ -130,8 +130,6 @@ def encode(self, data: np.ndarray) -> EncodedCIFData:
 
 lower_limit=-upper_limit-1
 
-# TODO: figure out if there is a way to implement this
-# better & faster with numpy methods.
 _pack_values(data, upper_limit, lower_limit, packed)
 
 byte_array_result=BYTE_ARRAY.encode(packed)
@@ -154,6 +152,7 @@ class _PackingInfo:
 bytesPerElement: int
 
 
+@jit(nopython=True)
 def_pack_values(data: np.ndarray, upper_limit: int, lower_limit: int, target: np.ndarray) ->None:
 offset=0
 forvalueindata:
@@ -200,6 +199,7 @@ def _determine_packing(data: np.ndarray) -> _PackingInfo:
 returnpacking
 
 
+@jit(nopython=True)
 def_packing_size_signed(data: np.ndarray, upper_limit: int) ->int:
 lower_limit=-upper_limit-1
 size=0
@@ -213,6 +213,7 @@ def _packing_size_signed(data: np.ndarray, upper_limit: int) -> int:
 returnsize+len(data)
 
 
+@jit(nopython=True)
 def_packing_size_unsigned(data: np.ndarray, upper_limit: int) ->int:
 size=0
 
@@ -299,34 +300,16 @@ def encode(self, data: np.ndarray) -> EncodedCIFData:
 
 classStringArray(BinaryCIFEncoder):
 defencode(self, data: Union[np.ndarray, list[str]]) ->EncodedCIFData:
-_map=dict()
-
 strings: list[str] = []
 offsets= [0]
 indices=np.empty(len(data), dtype="<i4")
 
-acc_len=0
-
-fori, sinenumerate(data):
-# handle null strings.
-ifnots:
-indices[i] =-1
-continue
-
-index=_map.get(s)
-ifindexisNone:
-# increment the length
-acc_len+=len(s)
-
-# store the string and index
-index=len(strings)
-strings.append(s)
-_map[s] =index
-
-# write the offset
-offsets.append(acc_len)
-
-indices[i] =index
+_pack_strings(
+data,
+indices,
+strings,
+offsets,
+ )
 
 encoded_offsets=_OFFSET_ENCODER.encode(np.array(offsets, dtype="<i4"))
 encoded_data=_DATA_ENCODER.encode(indices)
@@ -341,4 +324,31 @@ def encode(self, data: Union[np.ndarray, list[str]]) -> EncodedCIFData:
 
 returnEncodedCIFData(data=encoded_data["data"], encoding=[encoding])
 
+# TODO: benchmark if JIT helps here
+@jit(nopython=False, forceobj=True)
+def_pack_strings(data: List[str], indices: np.ndarray, strings: List[str], offsets: List[int]) ->None:
+acc_len=0
+str_map=dict()
+
+fori, sinenumerate(data):
+# handle null strings.
+ifnots:
+indices[i] =-1
+continue
+
+index=str_map.get(s)
+ifindexisNone:
+# increment the length
+acc_len+=len(s)
+
+# store the string and index
+index=len(strings)
+strings.append(s)
+str_map[s] =index
+
+# write the offset
+offsets.append(acc_len)
+
+indices[i] =index
+
 STRING_ARRAY=StringArray()