1#!/usr/bin/env python3
2"""
3Test script to isolate wavpack lossy compression issue with multi-channel data.
4Tests compression ratio for single channel vs all channels with bps=3.
5"""
7import numpy as np
8import requests
9import sys
10from io import BytesIO
12# URL for test data
13DATA_URL = "https://tempory.net/ephys-compression-tests/aind/aind_compression_np2_probeB_ch101-110.raw.npy"
15def download_data():
16 """Download test data from URL."""
17 print(f"Downloading data from {DATA_URL}...")
18 response = requests.get(DATA_URL)
19 response.raise_for_status()
20 arr = np.load(BytesIO(response.content))
21 print(f"Data shape: {arr.shape}, dtype: {arr.dtype}")
22 return arr
24def wavpack_encode(x: np.ndarray, bps: float = None) -> bytes:
25 """Encode array using WavPack."""
26 from wavpack_numcodecs import WavPack
27 if bps is not None:
28 codec = WavPack(bps=bps)
29 else:
30 codec = WavPack()
31 encoded = codec.encode(x)
32 assert isinstance(encoded, bytes)
33 return encoded
35def test_compression(data: np.ndarray, bps: int = 3):
36 """Test compression for single channel vs all channels."""
38 print(f"\n{'='*60}")
39 print(f"Testing WavPack with bps={bps}")
40 print(f"{'='*60}\n")
42 # Test single channel (first channel)
43 print("--- Single Channel Test ---")
44 single_channel = data[:, 0:1].copy() # Keep 2D shape (N, 1) and ensure contiguous
45 print(f"Single channel shape: {single_channel.shape}")
46 print(f"Single channel size: {single_channel.nbytes} bytes")
48 encoded_single = wavpack_encode(single_channel, bps=bps)
49 size_single = len(encoded_single)
50 ratio_single = single_channel.nbytes / size_single
52 print(f"Compressed size: {size_single} bytes")
53 print(f"Compression ratio: {ratio_single:.3f}x")
55 # Test all channels
56 print(f"\n--- All Channels Test ---")
57 all_channels = np.ascontiguousarray(data) # Ensure contiguous
58 print(f"All channels shape: {all_channels.shape}")
59 print(f"All channels size: {all_channels.nbytes} bytes")
61 encoded_all = wavpack_encode(all_channels, bps=bps)
62 size_all = len(encoded_all)
63 ratio_all = all_channels.nbytes / size_all
65 print(f"Compressed size: {size_all} bytes")
66 print(f"Compression ratio: {ratio_all:.3f}x")
68 # Compare
69 print(f"\n--- Comparison ---")
70 print(f"Single channel compression ratio: {ratio_single:.3f}x")
71 print(f"All channels compression ratio: {ratio_all:.3f}x")
73 # Expected: similar ratios if working correctly
74 # If all-channel ratio is much worse, there may be an issue
75 if ratio_all < ratio_single * 0.5:
76 print(f"\n⚠️ WARNING: All-channel compression ratio is significantly worse!")
77 print(f" This suggests a potential issue with multi-channel compression.")
78 elif ratio_all < ratio_single * 0.9:
79 print(f"\n⚠️ NOTICE: All-channel compression ratio is somewhat worse.")
80 else:
81 print(f"\n✓ Compression ratios are similar - working as expected.")
83 return {
84 'single_channel_ratio': ratio_single,
85 'all_channels_ratio': ratio_all,
86 'single_channel_size': size_single,
87 'all_channels_size': size_all
88 }
90def main():
91 """Main test function."""
92 try:
93 # Download data
94 data = download_data()
96 # Test with bps=3 (the reported issue)
97 results = test_compression(data, bps=3)
99 # Also test with lossless for comparison
100 print(f"\n\n{'='*60}")
101 print(f"For comparison, testing lossless compression:")
102 print(f"{'='*60}\n")
104 # Lossless single channel
105 single_channel = data[:, 0:1].copy()
106 encoded_single_lossless = wavpack_encode(single_channel)
107 ratio_single_lossless = single_channel.nbytes / len(encoded_single_lossless)
108 print(f"Single channel lossless ratio: {ratio_single_lossless:.3f}x")
110 # Lossless all channels
111 all_channels = np.ascontiguousarray(data)
112 encoded_all_lossless = wavpack_encode(all_channels)
113 ratio_all_lossless = all_channels.nbytes / len(encoded_all_lossless)
114 print(f"All channels lossless ratio: {ratio_all_lossless:.3f}x")
116 print(f"\n{'='*60}")
117 print("Test completed successfully!")
118 print(f"{'='*60}")
120 except Exception as e:
121 print(f"\n❌ Error during test: {e}", file=sys.stderr)
122 import traceback
123 traceback.print_exc()
124 sys.exit(1)
126if __name__ == "__main__":
127 main()