fix ans headers
2 changed files+626−142
benchcompress/src/benchcompress/algorithms/ans/__init__.pymodified+141−142View file
@@ -7,6 +7,15 @@ from .lpc_predict import (
77 lpc_predict as lpc_predict_cpp,
88 )
99 from .get_run_lengths import get_run_lengths
10+from .header_utils import (
11+ create_header_with_schema,
12+ unpack_header_with_schema,
13+ create_ans_header,
14+ create_lpc_header,
15+ unpack_lpc_header,
16+ create_lpc_sparse_header,
17+ unpack_lpc_sparse_header,
18+)
1019
1120
1221 SOURCE_FILE = "ans/__init__.py"
@@ -23,9 +32,9 @@ LONG_DESCRIPTION = _load_long_description()
2332
2433
2534 def ans_encode(x: np.ndarray) -> bytes:
26- from simple_ans import ans_encode
35+ from simple_ans import ans_encode as ans_encode_0
2736
28- encoded = ans_encode(x)
37+ encoded = ans_encode_0(x)
2938 if x.dtype == np.uint8:
3039 dtype_code = 0
3140 elif x.dtype == np.uint16:
@@ -38,31 +47,41 @@ def ans_encode(x: np.ndarray) -> bytes:
3847 dtype_code = 4
3948 else:
4049 raise ValueError(f"Unsupported dtype: {x.dtype}")
41- header = (
42- [
43- dtype_code,
44- len(encoded.words),
45- encoded.signal_length,
46- encoded.state,
47- len(encoded.symbol_counts),
48- ]
49- + [c for c in encoded.symbol_counts]
50- + [v for v in encoded.symbol_values]
50+
51+ # Use the new header utilities
52+ header_bytes, schema = create_ans_header(
53+ dtype_code=dtype_code,
54+ num_words=len(encoded.words),
55+ signal_length=encoded.signal_length,
56+ state=encoded.state,
57+ symbol_counts=encoded.symbol_counts,
58+ symbol_values=encoded.symbol_values,
5159 )
52- header_bytes = np.array(header, dtype=np.int64).tobytes()
53- header_size = np.uint32(len(header_bytes))
54- return header_size.tobytes() + header_bytes + encoded.words.tobytes()
5560
61+ # Create self-describing header
62+ complete_header = create_header_with_schema(header_bytes, schema)
63+ header_size = np.uint32(len(complete_header))
5664
57-def ans0_decode(x: bytes, dtype: str, shape: tuple) -> np.ndarray:
58- from simple_ans import ans_decode, EncodedSignal
65+ return header_size.tobytes() + complete_header + encoded.words.tobytes()
66+
67+
68+def ans_decode(x: bytes, dtype: str, shape: tuple) -> np.ndarray:
69+ from simple_ans import ans_decode as ans_decode_0, EncodedSignal
5970
6071 header_size = np.frombuffer(x[:4], dtype=np.uint32)[0]
61- header = np.frombuffer(x[4 : 4 + header_size], dtype=np.int64)
62- dtype_code, num_words, signal_length, state, num_symbols = header[:5]
63- symbol_counts = header[5 : 5 + num_symbols]
64- symbol_values = header[5 + num_symbols :]
72+
73+ # Use the new header utilities
74+ header_dict, total_header_size = unpack_header_with_schema(x[4 : 4 + header_size])
75+
76+ dtype_code = header_dict["dtype_code"]
77+ num_words = header_dict["num_words"]
78+ signal_length = header_dict["signal_length"]
79+ state = header_dict["state"]
80+ symbol_counts = header_dict["symbol_counts"]
81+ symbol_values = header_dict["symbol_values"]
82+
6583 words_bytes = x[4 + header_size :]
84+
6685 if dtype_code == 0:
6786 assert dtype == "uint8"
6887 elif dtype_code == 1:
@@ -78,12 +97,12 @@ def ans0_decode(x: bytes, dtype: str, shape: tuple) -> np.ndarray:
7897
7998 encoded = EncodedSignal(
8099 signal_length=int(signal_length),
81- state=int(state),
100+ state=np.uint64(state),
82101 symbol_counts=symbol_counts.astype(np.uint32),
83102 symbol_values=symbol_values.astype(dtype),
84103 words=np.frombuffer(words_bytes, dtype=np.uint32, count=num_words),
85104 )
86- return ans_decode(encoded).reshape(shape)
105+ return ans_decode_0(encoded).reshape(shape)
87106
88107
89108 def ans_delta_encode(x: np.ndarray) -> bytes:
@@ -106,22 +125,24 @@ def ans_delta_encode(x: np.ndarray) -> bytes:
106125 dtype_code = 4
107126 else:
108127 raise ValueError(f"Unsupported dtype: {x.dtype}")
109- # Include x[0] in the header
110- header = (
111- [
112- dtype_code,
113- len(encoded.words),
114- encoded.signal_length,
115- encoded.state,
116- len(encoded.symbol_counts),
117- x[0], # Store first value in header
118- ]
119- + [c for c in encoded.symbol_counts]
120- + [v for v in encoded.symbol_values]
128+
129+ # Use the new header utilities with extra value for x[0]
130+ extra_values = [(x[0], x.dtype.name)]
131+ header_bytes, schema = create_ans_header(
132+ dtype_code=dtype_code,
133+ num_words=len(encoded.words),
134+ signal_length=encoded.signal_length,
135+ state=encoded.state,
136+ symbol_counts=encoded.symbol_counts,
137+ symbol_values=encoded.symbol_values,
138+ extra_values=extra_values,
121139 )
122- header_bytes = np.array(header, dtype=np.int64).tobytes()
123- header_size = np.uint32(len(header_bytes))
124- return header_size.tobytes() + header_bytes + encoded.words.tobytes()
140+
141+ # Create self-describing header
142+ complete_header = create_header_with_schema(header_bytes, schema)
143+ header_size = np.uint32(len(complete_header))
144+
145+ return header_size.tobytes() + complete_header + encoded.words.tobytes()
125146
126147
127148 def ans_delta_decode(x: bytes, dtype: str, shape: tuple) -> np.ndarray:
@@ -130,13 +151,20 @@ def ans_delta_decode(x: bytes, dtype: str, shape: tuple) -> np.ndarray:
130151 assert len(shape) == 1
131152
132153 header_size = np.frombuffer(x[:4], dtype=np.uint32)[0]
133- header = np.frombuffer(x[4 : 4 + header_size], dtype=np.int64)
134- dtype_code, num_words, signal_length, state, num_symbols, x0 = header[
135- :6
136- ] # Extract x0 from header
137- symbol_counts = header[6 : 6 + num_symbols]
138- symbol_values = header[6 + num_symbols :]
154+
155+ # Use the new header utilities
156+ header_dict, total_header_size = unpack_header_with_schema(x[4 : 4 + header_size])
157+
158+ dtype_code = header_dict["dtype_code"]
159+ num_words = header_dict["num_words"]
160+ signal_length = header_dict["signal_length"]
161+ state = header_dict["state"]
162+ symbol_counts = header_dict["symbol_counts"]
163+ symbol_values = header_dict["symbol_values"]
164+ x0 = header_dict["extra_values"][0] # First extra value is x[0]
165+
139166 words_bytes = x[4 + header_size :]
167+
140168 if dtype_code == 0:
141169 assert dtype == "uint8"
142170 elif dtype_code == 1:
@@ -152,7 +180,7 @@ def ans_delta_decode(x: bytes, dtype: str, shape: tuple) -> np.ndarray:
152180
153181 encoded = EncodedSignal(
154182 signal_length=int(signal_length),
155- state=int(state),
183+ state=state,
156184 symbol_counts=symbol_counts.astype(np.uint32),
157185 symbol_values=symbol_values.astype(dtype),
158186 words=np.frombuffer(words_bytes, dtype=np.uint32, count=num_words),
@@ -160,7 +188,9 @@ def ans_delta_decode(x: bytes, dtype: str, shape: tuple) -> np.ndarray:
160188 # Decode the differences
161189 diffs = ans_decode(encoded)
162190 # Insert x0 at the beginning and cumulatively sum the differences
163- return np.cumsum(np.insert(diffs, 0, x0))
191+ result = np.cumsum(np.insert(diffs, 0, x0))
192+ # Cast back to the original dtype
193+ return result.astype(dtype)
164194
165195
166196 def ans_lpc_encode(x: np.ndarray) -> bytes:
@@ -183,25 +213,24 @@ def ans_lpc_encode(x: np.ndarray) -> bytes:
183213 dtype_code = 4
184214 else:
185215 raise ValueError(f"Unsupported dtype: {x.dtype}")
186- # Include x[0] in the header
187- header = (
188- [
189- dtype_code,
190- len(encoded.words),
191- encoded.signal_length,
192- encoded.state,
193- len(encoded.symbol_counts),
194- len(coeffs),
195- len(initial),
196- ]
197- + [c for c in encoded.symbol_counts]
198- + [v for v in encoded.symbol_values]
199- + [c for c in coeffs]
200- + [v for v in initial]
216+
217+ # Use the new LPC header utilities
218+ header_bytes, schema = create_lpc_header(
219+ dtype_code=dtype_code,
220+ num_words=len(encoded.words),
221+ signal_length=encoded.signal_length,
222+ state=encoded.state,
223+ symbol_counts=encoded.symbol_counts,
224+ symbol_values=encoded.symbol_values,
225+ coeffs=coeffs,
226+ initial=initial,
201227 )
202- header_bytes = np.array(header, dtype=np.float64).tobytes()
203- header_size = np.uint32(len(header_bytes))
204- return header_size.tobytes() + header_bytes + encoded.words.tobytes()
228+
229+ # Create self-describing header
230+ complete_header = create_header_with_schema(header_bytes, schema)
231+ header_size = np.uint32(len(complete_header))
232+
233+ return header_size.tobytes() + complete_header + encoded.words.tobytes()
205234
206235
207236 def ans_lpc_decode(x: bytes, dtype: str, shape: tuple) -> np.ndarray:
@@ -210,33 +239,21 @@ def ans_lpc_decode(x: bytes, dtype: str, shape: tuple) -> np.ndarray:
210239 assert len(shape) == 1
211240
212241 header_size = np.frombuffer(x[:4], dtype=np.uint32)[0]
213- header = np.frombuffer(x[4 : 4 + header_size], dtype=np.float64)
214- (
215- dtype_code,
216- num_words,
217- signal_length,
218- state,
219- num_symbols,
220- num_coeffs,
221- num_initial,
222- ) = header[:7]
223- dtype_code = int(dtype_code)
224- signal_length = int(signal_length)
225- state = int(state)
226- num_symbols = int(num_symbols)
227- num_coeffs = int(num_coeffs)
228- num_initial = int(num_initial)
229-
230- pos = 7
231- symbol_counts = header[pos : pos + num_symbols]
232- pos += num_symbols
233- symbol_values = header[pos : pos + num_symbols]
234- pos += num_symbols
235- coeffs = header[pos : pos + num_coeffs].astype(np.float32)
236- pos += num_coeffs
237- initial = header[pos : pos + num_initial].astype(dtype)
238- pos += num_initial
242+
243+ # Use the new header utilities
244+ header_dict, total_header_size = unpack_header_with_schema(x[4 : 4 + header_size])
245+
246+ dtype_code = header_dict["dtype_code"]
247+ num_words = header_dict["num_words"]
248+ signal_length = header_dict["signal_length"]
249+ state = header_dict["state"]
250+ symbol_counts = header_dict["symbol_counts"]
251+ symbol_values = header_dict["symbol_values"]
252+ coeffs = header_dict["coeffs"]
253+ initial = header_dict["initial"]
254+
239255 bitstream = x[4 + header_size :]
256+
240257 if dtype_code == 0:
241258 assert dtype == "uint8"
242259 elif dtype_code == 1:
@@ -252,7 +269,7 @@ def ans_lpc_decode(x: bytes, dtype: str, shape: tuple) -> np.ndarray:
252269
253270 encoded = EncodedSignal(
254271 signal_length=int(signal_length),
255- state=int(state),
272+ state=state,
256273 symbol_counts=symbol_counts.astype(np.uint32),
257274 symbol_values=symbol_values.astype(dtype),
258275 words=np.frombuffer(bitstream, dtype=np.uint32, count=num_words),
@@ -313,29 +330,27 @@ def ans_lpc_sparse_encode(x: np.ndarray) -> bytes:
313330 else:
314331 raise ValueError(f"Unsupported run length dtype: {run_lengths.dtype}")
315332
316- header = (
317- [
318- dtype_code,
319- len(encoded.words),
320- encoded.signal_length,
321- encoded.state,
322- len(encoded.symbol_counts),
323- len(coeffs),
324- len(initial),
325- run_length_dtype_code,
326- len(run_lengths),
327- ]
328- + [c for c in encoded.symbol_counts]
329- + [v for v in encoded.symbol_values]
330- + [c for c in coeffs]
331- + [v for v in initial]
333+ # Use the new LPC sparse header utilities
334+ header_bytes, schema = create_lpc_sparse_header(
335+ dtype_code=dtype_code,
336+ num_words=len(encoded.words),
337+ signal_length=encoded.signal_length,
338+ state=encoded.state,
339+ symbol_counts=encoded.symbol_counts,
340+ symbol_values=encoded.symbol_values,
341+ coeffs=coeffs,
342+ initial=initial,
343+ run_length_dtype_code=run_length_dtype_code,
344+ num_run_lengths=len(run_lengths),
332345 )
333- header_bytes = np.array(header, dtype=np.float64).tobytes()
334- header_size = np.uint32(len(header_bytes))
346+
347+ # Create self-describing header
348+ complete_header = create_header_with_schema(header_bytes, schema)
349+ header_size = np.uint32(len(complete_header))
335350
336351 return (
337352 header_size.tobytes()
338- + header_bytes
353+ + complete_header
339354 + encoded.words.tobytes()
340355 + run_lengths.tobytes()
341356 )
@@ -347,36 +362,20 @@ def ans_lpc_sparse_decode(x: bytes, dtype: str, shape: tuple) -> np.ndarray:
347362 assert len(shape) == 1
348363
349364 header_size = np.frombuffer(x[:4], dtype=np.uint32)[0]
350- header = np.frombuffer(x[4 : 4 + header_size], dtype=np.float64)
351- (
352- dtype_code,
353- num_words,
354- signal_length,
355- state,
356- num_symbols,
357- num_coeffs,
358- num_initial,
359- run_length_dtype_code,
360- num_run_lengths,
361- ) = header[:10]
362- dtype_code = int(dtype_code)
363- signal_length = int(signal_length)
364- state = int(state)
365- num_symbols = int(num_symbols)
366- num_coeffs = int(num_coeffs)
367- num_initial = int(num_initial)
368- run_length_dtype_code = int(run_length_dtype_code)
369- num_run_lengths = int(num_run_lengths)
370-
371- pos = 10
372- symbol_counts = header[pos : pos + num_symbols]
373- pos += num_symbols
374- symbol_values = header[pos : pos + num_symbols]
375- pos += num_symbols
376- coeffs = header[pos : pos + num_coeffs].astype(np.float32)
377- pos += num_coeffs
378- initial = header[pos : pos + num_initial].astype(dtype)
379- pos += num_initial
365+
366+ # Use the new header utilities
367+ header_dict, total_header_size = unpack_header_with_schema(x[4 : 4 + header_size])
368+
369+ dtype_code = header_dict["dtype_code"]
370+ num_words = header_dict["num_words"]
371+ signal_length = header_dict["signal_length"]
372+ state = header_dict["state"]
373+ symbol_counts = header_dict["symbol_counts"]
374+ symbol_values = header_dict["symbol_values"]
375+ coeffs = header_dict["coeffs"]
376+ initial = header_dict["initial"]
377+ run_length_dtype_code = header_dict["run_length_dtype_code"]
378+ num_run_lengths = header_dict["num_run_lengths"]
380379
381380 words_end = 4 + header_size + num_words * 4
382381 words_bytes = x[4 + header_size : words_end]
@@ -411,7 +410,7 @@ def ans_lpc_sparse_decode(x: bytes, dtype: str, shape: tuple) -> np.ndarray:
411410
412411 encoded = EncodedSignal(
413412 signal_length=int(signal_length),
414- state=int(state),
413+ state=state,
415414 symbol_counts=symbol_counts.astype(np.uint32),
416415 symbol_values=symbol_values.astype(dtype),
417416 words=np.frombuffer(words_bytes, dtype=np.uint32, count=num_words),
@@ -450,7 +449,7 @@ algorithms = [
450449 "name": "ANS",
451450 "version": "3",
452451 "encode": lambda x: ans_encode(x),
453- "decode": lambda x, dtype, shape: ans0_decode(x, dtype, shape),
452+ "decode": lambda x, dtype, shape: ans_decode(x, dtype, shape),
454453 "description": "ANS compression via simple_ans for efficient data compression.",
455454 "tags": ["ANS", "integer"],
456455 "source_file": SOURCE_FILE,
benchcompress/src/benchcompress/algorithms/ans/header_utils.pyadded+485−0View file
@@ -0,0 +1,485 @@
1+"""
2+Utilities for serializing and deserializing mixed-type headers in ANS compression.
3+
4+This module provides functions to pack and unpack headers containing different
5+data types while preserving their original types and precision.
6+"""
7+
8+import numpy as np
9+from typing import List, Tuple, Any, Union
10+
11+
12+def pack_header_values(values: List[Tuple[Any, str]]) -> bytes:
13+ """
14+ Pack a list of (value, dtype) tuples into a binary format.
15+
16+ Args:
17+ values: List of (value, dtype_string) tuples where dtype_string is like 'int32', 'float32', etc.
18+
19+ Returns:
20+ bytes: Packed binary data
21+ """
22+ packed_data = []
23+
24+ for value, dtype_str in values:
25+ if isinstance(value, (list, np.ndarray)):
26+ # Handle arrays
27+ arr = np.array(value, dtype=dtype_str)
28+ packed_data.append(arr.tobytes())
29+ else:
30+ # Handle scalars
31+ scalar = np.array([value], dtype=dtype_str)
32+ packed_data.append(scalar.tobytes())
33+
34+ return b"".join(packed_data)
35+
36+
37+def unpack_header_values(data: bytes, schema: List[Tuple[str, int]]) -> List[Any]:
38+ """
39+ Unpack binary data according to a schema.
40+
41+ Args:
42+ data: Binary data to unpack
43+ schema: List of (dtype_string, count) tuples where count is number of elements
44+ Use count=1 for scalars, count>1 for arrays
45+
46+ Returns:
47+ List of unpacked values
48+ """
49+ values = []
50+ offset = 0
51+
52+ for dtype_str, count in schema:
53+ dtype = np.dtype(dtype_str)
54+ size = dtype.itemsize * count
55+
56+ if offset + size > len(data):
57+ raise ValueError(f"Not enough data to unpack {count} {dtype_str} values")
58+
59+ chunk = data[offset : offset + size]
60+ arr = np.frombuffer(chunk, dtype=dtype)
61+
62+ values.append(arr.flatten()) # Always return flattened array to ensure 1D
63+
64+ offset += size
65+
66+ return values
67+
68+
69+def create_ans_header(
70+ dtype_code: int,
71+ num_words: int,
72+ signal_length: int,
73+ state: np.uint64,
74+ symbol_counts: np.ndarray,
75+ symbol_values: np.ndarray,
76+ extra_values: Union[List[Tuple[Any, str]], None] = None,
77+) -> Tuple[bytes, List[Tuple[str, int]]]:
78+ """
79+ Create a header for ANS compression with proper type preservation.
80+
81+ Args:
82+ dtype_code: Integer code for the data type
83+ num_words: Number of words in the encoded data
84+ signal_length: Length of the original signal
85+ state: ANS state
86+ symbol_counts: Array of symbol counts (will be stored as uint32)
87+ symbol_values: Array of symbol values (will be stored with original dtype)
88+ extra_values: Optional list of (value, dtype_string) tuples for additional data
89+
90+ Returns:
91+ Tuple of (header_bytes, schema) where schema describes the unpacking format
92+ """
93+ # Build the values list with their types
94+ values: List[Tuple[Any, str]] = [
95+ (dtype_code, "int32"),
96+ (num_words, "int32"),
97+ (signal_length, "int32"),
98+ (state, "uint64"),
99+ (len(symbol_counts), "int32"),
100+ ]
101+
102+ # Add symbol counts as uint32
103+ values.append((symbol_counts, "uint32"))
104+
105+ # Add symbol values with their original dtype
106+ symbol_dtype = symbol_values.dtype.name
107+ values.append((symbol_values, symbol_dtype))
108+
109+ # Add any extra values
110+ if extra_values:
111+ values.extend(extra_values)
112+
113+ # Create schema for unpacking
114+ schema = [
115+ ("int32", 1), # dtype_code
116+ ("int32", 1), # num_words
117+ ("int32", 1), # signal_length
118+ ("uint64", 1), # state
119+ ("int32", 1), # num_symbols
120+ ("uint32", len(symbol_counts)), # symbol_counts
121+ (symbol_dtype, len(symbol_values)), # symbol_values
122+ ]
123+
124+ # Add schema entries for extra values
125+ if extra_values:
126+ for value, dtype_str in extra_values:
127+ if isinstance(value, (list, np.ndarray)):
128+ count = len(value)
129+ else:
130+ count = 1
131+ schema.append((dtype_str, count))
132+
133+ header_bytes = pack_header_values(values)
134+ return header_bytes, schema
135+
136+
137+def unpack_ans_header(data: bytes, schema: List[Tuple[str, int]]) -> dict:
138+ """
139+ Unpack an ANS header according to the provided schema.
140+
141+ Args:
142+ data: Binary header data
143+ schema: Schema describing the data layout
144+
145+ Returns:
146+ Dictionary with unpacked header values
147+ """
148+ values = unpack_header_values(data, schema)
149+
150+ result = {
151+ "dtype_code": values[0][0], # Extract scalar from 1-element array
152+ "num_words": values[1][0],
153+ "signal_length": values[2][0],
154+ "state": values[3][0],
155+ "num_symbols": values[4][0],
156+ "symbol_counts": values[5], # Keep as array
157+ "symbol_values": values[6], # Keep as array
158+ }
159+
160+ # Add any extra values
161+ if len(values) > 7:
162+ extra_values = []
163+ for val in values[7:]:
164+ if len(val) == 1:
165+ extra_values.append(val[0]) # Extract scalar
166+ else:
167+ extra_values.append(val) # Keep as array
168+ result["extra_values"] = extra_values
169+
170+ return result
171+
172+
173+def create_lpc_sparse_header(
174+ dtype_code: int,
175+ num_words: int,
176+ signal_length: int,
177+ state: np.uint64,
178+ symbol_counts: np.ndarray,
179+ symbol_values: np.ndarray,
180+ coeffs: np.ndarray,
181+ initial: np.ndarray,
182+ run_length_dtype_code: int,
183+ num_run_lengths: int,
184+) -> Tuple[bytes, List[Tuple[str, int]]]:
185+ """
186+ Create a header for ANS LPC sparse compression with proper type preservation.
187+
188+ Args:
189+ dtype_code: Integer code for the data type
190+ num_words: Number of words in the encoded data
191+ signal_length: Length of the original signal
192+ state: ANS state
193+ symbol_counts: Array of symbol counts (will be stored as uint32)
194+ symbol_values: Array of symbol values (will be stored with original dtype)
195+ coeffs: LPC coefficients (will be stored as float32)
196+ initial: Initial values (will be stored with original dtype)
197+ run_length_dtype_code: Integer code for run length data type
198+ num_run_lengths: Number of run length values
199+
200+ Returns:
201+ Tuple of (header_bytes, schema) where schema describes the unpacking format
202+ """
203+ # Build the values list with their types
204+ values: List[Tuple[Any, str]] = [
205+ (dtype_code, "int32"),
206+ (num_words, "int32"),
207+ (signal_length, "int32"),
208+ (state, "uint64"), # Use uint64 for state as it can be very large
209+ (len(symbol_counts), "int32"),
210+ (len(coeffs), "int32"),
211+ (len(initial), "int32"),
212+ (run_length_dtype_code, "int32"),
213+ (num_run_lengths, "int32"),
214+ ]
215+
216+ # Add symbol counts as uint32
217+ values.append((symbol_counts, "uint32"))
218+
219+ # Add symbol values with their original dtype
220+ symbol_dtype = symbol_values.dtype.name
221+ values.append((symbol_values, symbol_dtype))
222+
223+ # Add coefficients as float32
224+ values.append((coeffs, "float32"))
225+
226+ # Add initial values with their original dtype
227+ initial_dtype = initial.dtype.name
228+ values.append((initial, initial_dtype))
229+
230+ # Create schema for unpacking
231+ schema = [
232+ ("int32", 1), # dtype_code
233+ ("int32", 1), # num_words
234+ ("int32", 1), # signal_length
235+ ("uint64", 1), # state
236+ ("int32", 1), # num_symbols
237+ ("int32", 1), # num_coeffs
238+ ("int32", 1), # num_initial
239+ ("int32", 1), # run_length_dtype_code
240+ ("int32", 1), # num_run_lengths
241+ ("uint32", len(symbol_counts)), # symbol_counts
242+ (symbol_dtype, len(symbol_values)), # symbol_values
243+ ("float32", len(coeffs)), # coeffs
244+ (initial_dtype, len(initial)), # initial
245+ ]
246+
247+ header_bytes = pack_header_values(values)
248+ return header_bytes, schema
249+
250+
251+def unpack_lpc_sparse_header(data: bytes, schema: List[Tuple[str, int]]) -> dict:
252+ """
253+ Unpack an ANS LPC sparse header according to the provided schema.
254+
255+ Args:
256+ data: Binary header data
257+ schema: Schema describing the data layout
258+
259+ Returns:
260+ Dictionary with unpacked header values
261+ """
262+ values = unpack_header_values(data, schema)
263+
264+ result = {
265+ "dtype_code": values[0][0], # Extract scalar from 1-element array
266+ "num_words": values[1][0],
267+ "signal_length": values[2][0],
268+ "state": values[3][0],
269+ "num_symbols": values[4][0],
270+ "num_coeffs": values[5][0],
271+ "num_initial": values[6][0],
272+ "run_length_dtype_code": values[7][0],
273+ "num_run_lengths": values[8][0],
274+ "symbol_counts": values[9], # Keep as array
275+ "symbol_values": values[10], # Keep as array
276+ "coeffs": values[11], # Keep as array
277+ "initial": values[12], # Keep as array
278+ }
279+
280+ return result
281+
282+
283+def create_header_with_schema(
284+ header_bytes: bytes, schema: List[Tuple[str, int]]
285+) -> bytes:
286+ """
287+ Create a complete header that includes the schema information for self-describing headers.
288+
289+ Args:
290+ header_bytes: The actual header data
291+ schema: Schema describing the header layout
292+
293+ Returns:
294+ Complete header with embedded schema
295+ """
296+ # Serialize the schema
297+ schema_data = []
298+ for dtype_str, count in schema:
299+ # Store dtype string length, dtype string, and count
300+ dtype_bytes = dtype_str.encode("utf-8")
301+ schema_data.append(np.array([len(dtype_bytes)], dtype="uint8").tobytes())
302+ schema_data.append(dtype_bytes)
303+ schema_data.append(np.array([count], dtype="int32").tobytes())
304+
305+ schema_bytes = b"".join(schema_data)
306+ schema_length = np.array([len(schema_bytes)], dtype="uint32").tobytes()
307+
308+ # Combine: schema_length + schema + header_data
309+ return schema_length + schema_bytes + header_bytes
310+
311+
312+def unpack_header_with_schema(data: bytes) -> Tuple[dict, int]:
313+ """
314+ Unpack a self-describing header that includes schema information.
315+
316+ Args:
317+ data: Complete header data with embedded schema
318+
319+ Returns:
320+ Tuple of (unpacked_header_dict, total_header_size)
321+ """
322+ # Read schema length
323+ schema_length = np.frombuffer(data[:4], dtype="uint32")[0]
324+
325+ # Read and parse schema
326+ schema_data = data[4 : 4 + schema_length]
327+ schema = []
328+ offset = 0
329+
330+ while offset < len(schema_data):
331+ # Read dtype string length
332+ dtype_len = np.frombuffer(schema_data[offset : offset + 1], dtype="uint8")[0]
333+ offset += 1
334+
335+ # Read dtype string
336+ dtype_str = schema_data[offset : offset + dtype_len].decode("utf-8")
337+ offset += dtype_len
338+
339+ # Read count
340+ count = np.frombuffer(schema_data[offset : offset + 4], dtype="int32")[0]
341+ offset += 4
342+
343+ schema.append((dtype_str, count))
344+
345+ # Unpack the actual header
346+ header_data = data[4 + schema_length :]
347+
348+ # Determine header type based on schema structure
349+ if len(schema) >= 13 and schema[7][0] == "int32" and schema[8][0] == "int32":
350+ # This looks like an LPC sparse header (has run_length_dtype_code and num_run_lengths)
351+ header_dict = unpack_lpc_sparse_header(header_data, schema)
352+ elif len(schema) >= 11 and schema[5][0] == "int32" and schema[6][0] == "int32":
353+ # This looks like an LPC header
354+ header_dict = unpack_lpc_header(header_data, schema)
355+ else:
356+ # This is a basic ANS header
357+ header_dict = unpack_ans_header(header_data, schema)
358+
359+ total_size = 4 + schema_length + len(header_data)
360+ return header_dict, total_size
361+
362+
363+def create_lpc_header(
364+ dtype_code: int,
365+ num_words: int,
366+ signal_length: int,
367+ state: np.uint64,
368+ symbol_counts: np.ndarray,
369+ symbol_values: np.ndarray,
370+ coeffs: np.ndarray,
371+ initial: np.ndarray,
372+ extra_values: Union[List[Tuple[Any, str]], None] = None,
373+) -> Tuple[bytes, List[Tuple[str, int]]]:
374+ """
375+ Create a header for ANS LPC compression with proper type preservation.
376+
377+ Args:
378+ dtype_code: Integer code for the data type
379+ num_words: Number of words in the encoded data
380+ signal_length: Length of the original signal
381+ state: ANS state
382+ symbol_counts: Array of symbol counts (will be stored as uint32)
383+ symbol_values: Array of symbol values (will be stored with original dtype)
384+ coeffs: LPC coefficients (will be stored as float32)
385+ initial: Initial values (will be stored with original dtype)
386+ extra_values: Optional list of (value, dtype_string) tuples for additional data
387+
388+ Returns:
389+ Tuple of (header_bytes, schema) where schema describes the unpacking format
390+ """
391+ # Build the values list with their types
392+ values: List[Tuple[Any, str]] = [
393+ (dtype_code, "int32"),
394+ (num_words, "int32"),
395+ (signal_length, "int32"),
396+ (state, "uint64"), # Use uint64 for state as it can be very large
397+ (len(symbol_counts), "int32"),
398+ (len(coeffs), "int32"),
399+ (len(initial), "int32"),
400+ ]
401+
402+ # Add symbol counts as uint32
403+ values.append((symbol_counts, "uint32"))
404+
405+ # Add symbol values with their original dtype
406+ symbol_dtype = symbol_values.dtype.name
407+ values.append((symbol_values, symbol_dtype))
408+
409+ # Add coefficients as float32
410+ values.append((coeffs, "float32"))
411+
412+ # Add initial values with their original dtype
413+ initial_dtype = initial.dtype.name
414+ values.append((initial, initial_dtype))
415+
416+ # Add any extra values
417+ if extra_values:
418+ values.extend(extra_values)
419+
420+ # Create schema for unpacking
421+ schema = [
422+ ("int32", 1), # dtype_code
423+ ("int32", 1), # num_words
424+ ("int32", 1), # signal_length
425+ ("uint64", 1), # state
426+ ("int32", 1), # num_symbols
427+ ("int32", 1), # num_coeffs
428+ ("int32", 1), # num_initial
429+ ("uint32", len(symbol_counts)), # symbol_counts
430+ (symbol_dtype, len(symbol_values)), # symbol_values
431+ ("float32", len(coeffs)), # coeffs
432+ (initial_dtype, len(initial)), # initial
433+ ]
434+
435+ # Add schema entries for extra values
436+ if extra_values:
437+ for value, dtype_str in extra_values:
438+ if isinstance(value, (list, np.ndarray)):
439+ count = len(value)
440+ else:
441+ count = 1
442+ schema.append((dtype_str, count))
443+
444+ header_bytes = pack_header_values(values)
445+ return header_bytes, schema
446+
447+
448+def unpack_lpc_header(data: bytes, schema: List[Tuple[str, int]]) -> dict:
449+ """
450+ Unpack an ANS LPC header according to the provided schema.
451+
452+ Args:
453+ data: Binary header data
454+ schema: Schema describing the data layout
455+
456+ Returns:
457+ Dictionary with unpacked header values
458+ """
459+ values = unpack_header_values(data, schema)
460+
461+ result = {
462+ "dtype_code": values[0][0], # Extract scalar from 1-element array
463+ "num_words": values[1][0],
464+ "signal_length": values[2][0],
465+ "state": values[3][0],
466+ "num_symbols": values[4][0],
467+ "num_coeffs": values[5][0],
468+ "num_initial": values[6][0],
469+ "symbol_counts": values[7], # Keep as array
470+ "symbol_values": values[8], # Keep as array
471+ "coeffs": values[9], # Keep as array
472+ "initial": values[10], # Keep as array
473+ }
474+
475+ # Add any extra values
476+ if len(values) > 11:
477+ extra_values = []
478+ for val in values[11:]:
479+ if len(val) == 1:
480+ extra_values.append(val[0]) # Extract scalar
481+ else:
482+ extra_values.append(val) # Keep as array
483+ result["extra_values"] = extra_values
484+
485+ return result