1// @ts-nocheck
2import pako from "pako";
3import { inverseTransform } from "./fft";
5type QfcCompressionOpts = {
6 compression_method: "zlib" | "zstd";
7 dtype: "float32" | "int16";
8 id: "qfc";
9 quant_scale_factor: number;
10 segment_length: number;
11 zlib_level: number;
12 zstd_level: number;
13};
15const isQfcCompressionOpts = (x: any): x is QfcCompressionOpts => {
16 if (!x) return false;
17 if (typeof x !== "object") return false;
18 if (x.compression_method !== "zlib" && x.compression_method !== "zstd")
19 return false;
20 if (x.dtype !== "float32" && x.dtype !== "int16") return false;
21 if (x.id !== "qfc") return false;
22 if (typeof x.quant_scale_factor !== "number") return false;
23 if (typeof x.segment_length !== "number") return false;
24 if (typeof x.zlib_level !== "number") return false;
25 if (typeof x.zstd_level !== "number") return false;
26 return true;
27};
29export const qfcDecompress = async (
30 buf: ArrayBuffer,
31 shape: number[],
32 compressor: QfcCompressionOpts,
33): Promise<any> => {
34 if (!isQfcCompressionOpts(compressor)) {
35 console.warn(compressor);
36 throw Error("Invalid qfc compressor");
37 }
39 const header = new Int32Array(buf, 0, 5);
40 if (header[0] !== 7364182) {
41 throw Error(`Invalid header[0]: ${header[0]}`);
42 }
43 if (header[1] !== 1) {
44 throw Error(`Invalid header[1]: ${header[1]}`);
45 }
46 const num_samples = header[2];
47 const num_channels = header[3];
48 if (num_samples !== shape[0]) {
49 throw Error(
50 `Unexpected num samples in header. Expected ${shape[0]}, got ${num_samples}`,
51 );
52 }
53 if (num_channels !== shape[1]) {
54 throw Error(
55 `Unexpected num channels in header. Expected ${shape[1]}, got ${num_channels}`,
56 );
57 }
58 if (header[4] !== compressor.segment_length) {
59 throw Error(
60 `Unexpected segment length in header. Expected ${compressor.segment_length}, got ${header[4]}`,
61 );
62 }
64 const decompressed_buf = await qfc_multi_segment_decompress({
65 buf: buf.slice(4 * 5),
66 dtype: compressor.dtype,
67 num_channels,
68 num_samples,
69 segment_length: compressor.segment_length,
70 quant_scale_factor: compressor.quant_scale_factor,
71 compression_method: compressor.compression_method,
72 });
74 if (
75 decompressed_buf.byteLength !==
76 num_samples * num_channels * (compressor.dtype === "float32" ? 4 : 2)
77 ) {
78 console.warn("compressor", compressor);
79 throw Error(
80 `Unexpected decompressed buffer length. Expected ${num_samples * num_channels * (compressor.dtype === "float32" ? 4 : 2)}, got ${decompressed_buf.byteLength}`,
81 );
82 }
84 return decompressed_buf;
85};
87const qfc_multi_segment_decompress = async (o: {
88 buf: ArrayBuffer;
89 dtype: "float32" | "int16";
90 num_channels: number;
91 num_samples: number;
92 segment_length: number;
93 quant_scale_factor: number;
94 compression_method: "zlib" | "zstd";
95}): Promise<ArrayBuffer> => {
96 const {
97 buf,
98 dtype,
99 num_channels,
100 num_samples,
101 segment_length,
102 quant_scale_factor,
103 compression_method,
104 } = o;
106 let decompressedArray: Int16Array;
107 if (compression_method === "zlib") {
108 decompressedArray = new Int16Array(pako.inflate(buf).buffer);
109 } else if (compression_method === "zstd") {
110 throw Error("zstd decompression not implemented");
111 } else {
112 throw Error(`Unexpected compression method: ${compression_method}`);
113 }
115 return await qfc_multi_segment_inv_pre_compress({
116 array: decompressedArray,
117 quant_scale_factor,
118 segment_length,
119 dtype,
120 num_channels,
121 num_samples,
122 });
123};
125const qfc_multi_segment_inv_pre_compress = async (o: {
126 array: Int16Array;
127 quant_scale_factor: number;
128 segment_length: number;
129 dtype: "int16" | "float32";
130 num_samples: number;
131 num_channels: number;
132}): Promise<ArrayBuffer> => {
133 const {
134 array,
135 quant_scale_factor,
136 segment_length,
137 dtype,
138 num_samples,
139 num_channels,
140 } = o;
141 if (segment_length > 0 && segment_length < num_samples) {
142 const segment_ranges = _get_segment_ranges(num_samples, segment_length);
143 const prepared_segments = await Promise.all(
144 segment_ranges.map(
145 async (segment_range) =>
146 await qfc_inv_pre_compress({
147 array: array.slice(
148 segment_range[0] * num_channels,
149 segment_range[1] * num_channels,
150 ),
151 quant_scale_factor,
152 dtype,
153 num_samples: segment_range[1] - segment_range[0],
154 num_channels,
155 }),
156 ),
157 );
158 if (dtype === "int16") {
159 return concatenateInt16Arrays(prepared_segments as Int16Array[]);
160 } else if (dtype === "float32") {
161 return concatenateFloat32Arrays(prepared_segments as Float32Array[]);
162 } else {
163 throw Error(`Unexpected dtype: ${dtype}`);
164 }
165 } else {
166 return await qfc_inv_pre_compress({
167 array,
168 quant_scale_factor,
169 dtype,
170 num_samples,
171 num_channels,
172 });
173 }
174};
176const _get_segment_ranges = (
177 total_length: number,
178 segment_length: number,
179): [number, number][] => {
180 const segment_ranges: [number, number][] = [];
181 for (
182 let start_index = 0;
183 start_index < total_length;
184 start_index += segment_length
185 ) {
186 segment_ranges.push([
187 start_index,
188 Math.min(start_index + segment_length, total_length),
189 ]);
190 }
191 const size_of_final_segment =
192 segment_ranges[segment_ranges.length - 1][1] -
193 segment_ranges[segment_ranges.length - 1][0];
194 const half_segment_length = Math.floor(segment_length / 2);
195 if (
196 size_of_final_segment < half_segment_length &&
197 segment_ranges.length > 1
198 ) {
199 const adjustment = half_segment_length - size_of_final_segment;
200 segment_ranges[segment_ranges.length - 2] = [
201 segment_ranges[segment_ranges.length - 2][0],
202 segment_ranges[segment_ranges.length - 2][1] - adjustment,
203 ];
204 segment_ranges[segment_ranges.length - 1] = [
205 segment_ranges[segment_ranges.length - 1][0] - adjustment,
206 segment_ranges[segment_ranges.length - 1][1],
207 ];
208 }
209 return segment_ranges;
210};
212const qfc_inv_pre_compress = async (o: {
213 array: Int16Array;
214 quant_scale_factor: number;
215 dtype: "int16" | "float32";
216 num_samples: number;
217 num_channels: number;
218}): Promise<Int16Array | Float32Array> => {
219 const { array, quant_scale_factor, dtype, num_samples, num_channels } = o;
221 const m = Math.floor(num_samples / 2);
222 const isEvenNumberOfSamples = num_samples % 2 === 0;
223 const qs = quant_scale_factor;
224 const x_re = new Float32Array((m + 1) * num_channels);
225 for (let i = 0; i < m + 1; i++) {
226 for (let j = 0; j < num_channels; j++) {
227 x_re[i * num_channels + j] = array[i * num_channels + j] / qs;
228 }
229 }
230 // ns - (ns // 2 + 1) + 2 = ns - ns // 2 - 1 + 2 = ns - ns // 2 + 1 = ns // 2 + 1
231 const x_im = new Float32Array((m + 1) * num_channels);
232 x_im.fill(0); // probably not necessary
233 const mm = isEvenNumberOfSamples ? m : m + 1;
234 for (let i = 1; i < mm; i++) {
235 for (let j = 0; j < num_channels; j++) {
236 x_im[i * num_channels + j] =
237 array[(m + 1 + (i - 1)) * num_channels + j] / qs;
238 }
239 }
241 const x_fft = await irfftMultiChannel(x_re, x_im, num_samples, num_channels);
242 for (let i = 0; i < x_fft.length; i++) {
243 x_fft[i] = x_fft[i] * Math.sqrt(num_samples);
244 }
245 if (dtype === "int16") {
246 const ret = new Int16Array(x_fft.byteLength);
247 for (let i = 0; i < x_fft.length; i++) {
248 ret[i] = Math.round(x_fft[i]);
249 }
250 return ret;
251 } else if (dtype === "float32") {
252 return x_fft;
253 } else {
254 throw Error(`Unexpected dtype: ${dtype}`);
255 }
256};
258const irfftMultiChannel = async (
259 x_re: Float32Array,
260 x_im: Float32Array,
261 num_samples: number,
262 num_channels: number,
263): Promise<Float32Array> => {
264 const ns = num_samples;
265 const nc = num_channels;
266 const ret = new Float32Array(ns * nc);
267 for (let j = 0; j < nc; j++) {
268 const a_re = new Float32Array(x_re.length / nc);
269 const a_im = new Float32Array(x_im.length / nc);
270 for (let i = 0; i < x_re.length / num_channels; i++) {
271 a_re[i] = x_re[i * nc + j];
272 a_im[i] = x_im[i * nc + j];
273 }
274 const b = await irfft(a_re, a_im, num_samples);
275 for (let i = 0; i < ns; i++) {
276 ret[i * nc + j] = b[i];
277 }
278 }
279 return ret;
280};
282const irfft = async (
283 x_re: Float32Array,
284 x_im: Float32Array,
285 num_samples: number,
286): Promise<Float32Array> => {
287 if (x_re.length != Math.floor(num_samples / 2) + 1) {
288 throw Error(
289 `Unexpected x_re length. Expected ${Math.floor(num_samples / 2) + 1}, got ${x_re.length}`,
290 );
291 }
292 if (x_im.length != Math.floor(num_samples / 2) + 1) {
293 throw Error(
294 `Unexpected x_im length. Expected ${Math.floor(num_samples / 2) + 1}, got ${x_im.length}`,
295 );
296 }
297 const x_re_copy = new Float32Array(num_samples);
298 const x_im_copy = new Float32Array(num_samples);
299 for (let i = 0; i < x_re.length; i++) {
300 x_re_copy[i] = x_re[i];
301 x_im_copy[i] = x_im[i];
302 if (i > 0) {
303 // the last case is i = x_re.length - 1
304 // in this case we are filling in (num_samples - x_re.length + 1)
305 // x_re.length = num_samples // 2 + 1
306 // so i = num_samples // 2
307 // and we're filling in (num_samples - num_samples // 2)
308 // in the case where num_samples is even, this is num_samples // 2, and imag part is zero there, so it's correct
309 // in the case where num_samples is odd, this is num_samples // 2 + 1, which is correct
310 x_re_copy[num_samples - i] = x_re[i];
311 x_im_copy[num_samples - i] = -x_im[i];
312 }
313 }
314 inverseTransform(x_re_copy, x_im_copy); // in place
315 // let's verify that imaginary part is close to zero
316 for (let i = 0; i < num_samples; i++) {
317 if (Math.abs(x_im_copy[i]) > 1e-5) {
318 throw Error("Unexpected non-zero imaginary part after inverse transform");
319 }
320 }
321 return x_re_copy;
322};
324const concatenateInt16Arrays = (arrays: Int16Array[]): ArrayBuffer => {
325 const total_length = arrays.reduce((acc, x) => acc + x.length, 0);
326 const concatenated = new Int16Array(total_length);
327 let offset = 0;
328 for (const a of arrays) {
329 concatenated.set(a, offset);
330 offset += a.length;
331 }
332 return concatenated.buffer;
333};
335const concatenateFloat32Arrays = (arrays: Float32Array[]): ArrayBuffer => {
336 const total_length = arrays.reduce((acc, x) => acc + x.length, 0);
337 const concatenated = new Float32Array(total_length);
338 let offset = 0;
339 for (const a of arrays) {
340 concatenated.set(a, offset);
341 offset += a.length;
342 }
343 return concatenated.buffer;
344};