Практическое руководство по Zarr: чанки, сжатие, индексация и визуализация больших массивов

!pip install zarr numcodecs -q
import zarr
import numpy as np
import matplotlib.pyplot as plt
from numcodecs import Blosc, Delta, FixedScaleOffset
import tempfile
import shutil
import os
from pathlib import Path


print(f"Zarr version: {zarr.__version__}")
print(f"NumPy version: {np.__version__}")


print("=== BASIC ZARR OPERATIONS ===")

tutorial_dir = Path(tempfile.mkdtemp(prefix="zarr_tutorial_"))
print(f"Working directory: {tutorial_dir}")


z1 = zarr.zeros((1000, 1000), chunks=(100, 100), dtype='f4',
               store=str(tutorial_dir / 'basic_array.zarr'), zarr_format=2)
z2 = zarr.ones((500, 500, 10), chunks=(100, 100, 5), dtype='i4',
              store=str(tutorial_dir / 'multi_dim.zarr'), zarr_format=2)


print(f"2D Array shape: {z1.shape}, chunks: {z1.chunks}, dtype: {z1.dtype}")
print(f"3D Array shape: {z2.shape}, chunks: {z2.chunks}, dtype: {z2.dtype}")


z1[100:200, 100:200] = np.random.random((100, 100)).astype('f4')
z2[:, :, 0] = np.arange(500*500).reshape(500, 500)


print(f"Memory usage estimate: {z1.nbytes_stored() / 1024**2:.2f} MB")

print("\n=== ADVANCED CHUNKING ===")


time_steps, height, width = 365, 1000, 2000
time_series = zarr.zeros(
   (time_steps, height, width),
   chunks=(30, 250, 500),
   dtype='f4',
   store=str(tutorial_dir / 'time_series.zarr'),
   zarr_format=2
)


for t in range(0, time_steps, 30):
   end_t = min(t + 30, time_steps)
   seasonal = np.sin(2 * np.pi * np.arange(t, end_t) / 365)[:, None, None]
   spatial = np.random.normal(20, 5, (end_t - t, height, width))
   time_series[t:end_t] = (spatial + 10 * seasonal).astype('f4')


print(f"Time series created: {time_series.shape}")
print(f"Approximate chunks created")


import time
start = time.time()
temporal_slice = time_series[:, 500, 1000]
temporal_time = time.time() - start


start = time.time()
spatial_slice = time_series[100, :200, :200]
spatial_time = time.time() - start


print(f"Temporal access time: {temporal_time:.4f}s")
print(f"Spatial access time: {spatial_time:.4f}s")

print("\n=== COMPRESSION AND CODECS ===")


data = np.random.randint(0, 1000, (1000, 1000), dtype='i4')


from zarr.codecs import BloscCodec, BytesCodec


z_none = zarr.array(data, chunks=(100, 100),
                  codecs=[BytesCodec()],
                  store=str(tutorial_dir / 'no_compress.zarr'))


z_lz4 = zarr.array(data, chunks=(100, 100),
                  codecs=[BytesCodec(), BloscCodec(cname='lz4', clevel=5)],
                  store=str(tutorial_dir / 'lz4_compress.zarr'))


z_zstd = zarr.array(data, chunks=(100, 100),
                   codecs=[BytesCodec(), BloscCodec(cname='zstd', clevel=9)],
                   store=str(tutorial_dir / 'zstd_compress.zarr'))


sequential_data = np.cumsum(np.random.randint(-5, 6, (1000, 1000)), axis=1)
z_delta = zarr.array(sequential_data, chunks=(100, 100),
                    codecs=[BytesCodec(), BloscCodec(cname='zstd', clevel=5)],
                    store=str(tutorial_dir / 'sequential_compress.zarr'))


sizes = {
   'No compression': z_none.nbytes_stored(),
   'LZ4': z_lz4.nbytes_stored(),
   'ZSTD': z_zstd.nbytes_stored(),
   'Sequential+ZSTD': z_delta.nbytes_stored()
}


print("Compression comparison:")
original_size = data.nbytes
for name, size in sizes.items():
   ratio = size / original_size
   print(f"{name}: {size/1024**2:.2f} MB (ratio: {ratio:.3f})")

print("\n=== HIERARCHICAL DATA ORGANIZATION ===")


root = zarr.open_group(str(tutorial_dir / 'experiment.zarr'), mode='w')


raw_data = root.create_group('raw_data')
processed = root.create_group('processed')
metadata = root.create_group('metadata')


raw_data.create_dataset('images', shape=(100, 512, 512), chunks=(10, 128, 128), dtype='u2')
raw_data.create_dataset('timestamps', shape=(100,), dtype='datetime64[ns]')


processed.create_dataset('normalized', shape=(100, 512, 512), chunks=(10, 128, 128), dtype='f4')
processed.create_dataset('features', shape=(100, 50), chunks=(20, 50), dtype='f4')


root.attrs['experiment_id'] = 'EXP_2024_001'
root.attrs['description'] = 'Advanced Zarr tutorial demonstration'
root.attrs['created'] = str(np.datetime64('2024-01-01'))


raw_data.attrs['instrument'] = 'Synthetic Camera'
raw_data.attrs['resolution'] = [512, 512]
processed.attrs['normalization'] = 'z-score'


timestamps = np.datetime64('2024-01-01') + np.arange(100) * np.timedelta64(1, 'h')
raw_data['timestamps'][:] = timestamps


for i in range(100):
   frame = np.random.poisson(100 + 50 * np.sin(2 * np.pi * i / 100), (512, 512)).astype('u2')
   raw_data['images'][i] = frame


print(f"Created hierarchical structure with {len(list(root.group_keys()))} groups")
print(f"Data arrays and groups created successfully")

print("\n=== ADVANCED INDEXING ===")


volume_data = zarr.zeros((50, 20, 256, 256), chunks=(5, 5, 64, 64), dtype='f4',
                       store=str(tutorial_dir / 'volume.zarr'), zarr_format=2)


for t in range(50):
   for z in range(20):
       y, x = np.ogrid[:256, :256]
       center_y, center_x = 128 + 20*np.sin(t*0.1), 128 + 20*np.cos(t*0.1)
       focus_quality = 1 - abs(z - 10) / 10
      
       signal = focus_quality * np.exp(-((y-center_y)**2 + (x-center_x)**2) / (50**2))
       noise = 0.1 * np.random.random((256, 256))
       volume_data[t, z] = (signal + noise).astype('f4')


print("Various slicing operations:")


max_projection = np.max(volume_data[:, 10], axis=0)
print(f"Max projection shape: {max_projection.shape}")


z_stack = volume_data[25, :, 100:156, 100:156]
print(f"Z-stack subset: {z_stack.shape}")


bright_pixels = volume_data[volume_data > 0.5]
print(f"Pixels above threshold: {len(bright_pixels)}")

print("\n=== PERFORMANCE OPTIMIZATION ===")


def process_chunk_serial(data, func):
   results = []
   for i in range(0, len(dt), 100):
       chunk = data[i:i+100]
       results.append(func(chunk))
   return np.concatenate(results)


def gaussian_filter_1d(x, sigma=1.0):
   kernel_size = int(4 * sigma)
   if kernel_size % 2 == 0:
       kernel_size += 1
   kernel = np.exp(-0.5 * ((np.arange(kernel_size) - kernel_size//2) / sigma)**2)
   kernel = kernel / kernel.sum()
   return np.convolve(x.astype(float), kernel, mode='same')


large_array = zarr.random.random((10000,), chunks=(1000,),
                              store=str(tutorial_dir / 'large.zarr'), zarr_format=2)


start_time = time.time()
chunk_size = 1000
filtered_data = []
for i in range(0, len(large_array), chunk_size):
   end_idx = min(i + chunk_size, len(large_array))
   chunk_data = large_array[i:end_idx]
   smoothed = np.convolve(chunk_data, np.ones(5)/5, mode='same')
   filtered_data.append(smoothed)


result = np.concatenate(filtered_data)
processing_time = time.time() - start_time


print(f"Chunk-aware processing time: {processing_time:.4f}s")
print(f"Processed {len(large_array):,} elements")

print("\n=== VISUALIZATION ===")


fig, axes = plt.subplots(2, 3, figsize=(15, 10))
fig.suptitle('Advanced Zarr Tutorial - Data Visualization', fontsize=16)


axes[0,0].plot(temporal_slice)
axes[0,0].set_title('Temporal Evolution (Single Pixel)')
axes[0,0].set_xlabel('Day of Year')
axes[0,0].set_ylabel('Temperature')


im1 = axes[0,1].imshow(spatial_slice, cmap='viridis')
axes[0,1].set_title('Spatial Pattern (Day 100)')
plt.colorbar(im1, ax=axes[0,1])


methods = list(sizes.keys())
ratios = [sizes[m]/original_size for m in methods]
axes[0,2].bar(range(len(methods)), ratios)
axes[0,2].set_xticks(range(len(methods)))
axes[0,2].set_xticklabels(methods, rotation=45)
axes[0,2].set_title('Compression Ratios')
axes[0,2].set_ylabel('Size Ratio')


axes[1,0].imshow(max_projection, cmap='hot')
axes[1,0].set_title('Max Intensity Projection')


z_profile = np.mean(volume_data[25, :, 120:136, 120:136], axis=(1,2))
axes[1,1].plot(z_profile, 'o-')
axes[1,1].set_title('Z-Profile (Center Region)')
axes[1,1].set_xlabel('Z-slice')
axes[1,1].set_ylabel('Mean Intensity')


axes[1,2].plot(result[:1000])
axes[1,2].set_title('Processed Signal (First 1000 points)')
axes[1,2].set_xlabel('Sample')
axes[1,2].set_ylabel('Amplitude')


plt.tight_layout()
plt.show()