2import numpy as np
3from helpers import bandpass_filter, estimate_noise_level, compute_entropy_per_sample
4from load_real import load_real_000876, load_real_000409, load_real_001290
5from compress_ints_lossless import compress_ints_lossless
7def linear_fit(x, y):
8 """Perform linear fit with constant term.
9 Returns coefficients and prediction function."""
10 from numpy.linalg import lstsq
11 X = np.column_stack([x, np.ones(len(x))])
12 coeffs = lstsq(X, y, rcond=None)[0]
14 def predict(x_new):
15 X_new = np.column_stack([x_new, np.ones(len(x_new))])
16 return np.dot(X_new, coeffs)
18 return coeffs, predict
20# %%
21N = 500_000
22# X = np.round(np.random.randn(N) * 500)
24# X = load_real_001290(num_samples=N, num_channels=1, start_channel=0)
25X = load_real_000409(num_samples=N, num_channels=1, start_channel=101)
26# X = load_real_000876(num_samples=N, num_channels=1, start_channel=45)
27# X = np.random.randn(len(X)) * 100
29X = X.astype(np.int16)
30X = X.flatten()
31# %%
32e1 = compute_entropy_per_sample(X)
33print(f'(raw) Bits per sample: {e1:.2f}')
34print(f'Ideal compression ratio: {X.itemsize * 8 / e1:.2f}')
36# %%
37e1 = compute_entropy_per_sample(np.diff(X))
38print(f'(raw diff) Bits per sample: {e1:.2f}')
39print(f'Ideal compression ratio: {X.itemsize * 8 / e1:.2f}')
41# %%
42# Actual compression ratio
43buf_zstd = compress_ints_lossless(np.diff(X), method='zstd')
44buf_zlib = compress_ints_lossless(np.diff(X), method='zlib')
45buf_lzma = compress_ints_lossless(np.diff(X), method='lzma')
46buf_ans = compress_ints_lossless(np.diff(X), method='simple_ans')
47print(f'Zstd compression ratio: {len(X) * X.itemsize / len(buf_zstd):.2f}')
48print(f'Zlib compression ratio: {len(X) * X.itemsize / len(buf_zlib):.2f}')
49print(f'Lzma compression ratio: {len(X) * X.itemsize / len(buf_lzma):.2f}')
50print(f'Simple ANS compression ratio: {len(X) * X.itemsize / len(buf_ans):.2f}')
52# %%
53M = 20
54# N - M + 1 x M
55sequences = np.array([X[i:i+M] for i in range(len(X) - 2 * M + 1)])
56predictors = sequences[:, :M - 1]
57target = sequences[:, M - 1]
59# Can choose either linear or quadratic fit
60# coeffs, predict = quadratic_fit(predictors, y)
61coeffs, predict = linear_fit(predictors, target)
62predictions = predict(predictors)
63predictions = np.round(predictions)
64residuals = target - predictions
65residuals = residuals.astype(np.int16)
66e3 = compute_entropy_per_sample(residuals)
67print(f'(raw adjusted) Bits per sample: {e3:.2f}')
68print(f'Ideal compression ratio: {X.itemsize * 8 / e3:.2f}')
70# %%
71v = 5
72lowcut = 300
73highcut = 6000
74X2 = bandpass_filter(X - np.median(X), sampling_frequency=30000, lowcut=lowcut, highcut=highcut)
75noise_level = estimate_noise_level(X2, sampling_frequency=30000)
76X2b = X2 / noise_level * v
77X2 = np.round(X2b).astype(np.int16)
78e2 = compute_entropy_per_sample(X2)
79print(f'(filtered) Bits per sample: {e2:.2f}')
80print(f'Ideal compression ratio: {X.itemsize * 8 / e2:.2f}')
82# %%
83e2 = compute_entropy_per_sample(np.diff(X2))
84print(f'(filtered diff) Bits per sample: {e2:.2f}')
85print(f'Ideal compression ratio: {X.itemsize * 8 / e2:.2f}')
87# %%
88M = 20
89# N - M + 1 x M
90sequences = np.array([X2[i:i+M] for i in range(len(X) - 2 * M + 1)])
91predictors = sequences[:, :M - 1]
92target = sequences[:, M - 1]
94coeffs, predict = linear_fit(predictors, target)
95predictions = predict(predictors)
96predictions = np.round(predictions)
97residuals = target - predictions
98residuals = residuals.astype(np.int16)
99e3 = compute_entropy_per_sample(residuals)
100print(f'(filtered adjusted) Bits per sample: {e3:.2f}')
101print(f'Ideal compression ratio: {X.itemsize * 8 / e3:.2f}')
102# %%
103# Get the actual compression ratio
104buf_zstd = compress_ints_lossless(residuals, method='zstd')
105buf_zlib = compress_ints_lossless(residuals, method='zlib')
106buf_lzma = compress_ints_lossless(residuals, method='lzma')
107buf_ans = compress_ints_lossless(residuals, method='simple_ans')
108print(f'Zstd compression ratio: {len(residuals) * residuals.itemsize / len(buf_zstd):.2f}')
109print(f'Zlib compression ratio: {len(residuals) * residuals.itemsize / len(buf_zlib):.2f}')
110print(f'Lzma compression ratio: {len(residuals) * residuals.itemsize / len(buf_lzma):.2f}')
111print(f'Simple ANS compression ratio: {len(residuals) * residuals.itemsize / len(buf_ans):.2f}')
113# %%
114import matplotlib.pyplot as plt
115plt.figure(figsize=(10, 6))
116plt.plot(X2[:600])
117# %%
118plt.plot(coeffs)
119# %%