Correlation Models for Wyner-Ziv Coding

This example demonstrates various correlation models used in distributed source coding and Wyner-Ziv compression using the new CorrelatedDataset.

We explore different correlation coefficients and visualize the relationship between source and side information signals.

import matplotlib.pyplot as plt
import numpy as np
import torch

from kaira.data import CorrelatedDataset

# Set random seed for reproducibility
torch.manual_seed(42)
np.random.seed(42)

Generate Correlated Data

Create datasets with different correlation coefficients

# Define correlation levels to test
correlations = [0.2, 0.5, 0.8, 0.95]
n_samples = 1000
signal_length = 128

# Generate correlated data for each correlation level
datasets = {}
for corr in correlations:
    datasets[corr] = CorrelatedDataset(length=n_samples, shape=(signal_length,), correlation=corr, noise_std=0.1, seed=42)

Visualize Signal Correlation

Plot source vs side information for different correlation levels

fig, axes = plt.subplots(2, 2, figsize=(12, 10))
axes = axes.ravel()

for i, corr in enumerate(correlations):
    # Get a sample from the dataset
    source, side_info = datasets[corr][0]

    # Convert to numpy for plotting
    source_np = source.numpy()
    side_info_np = side_info.numpy()

    # Scatter plot of first 100 samples
    axes[i].scatter(source_np[:100], side_info_np[:100], alpha=0.6, s=10)
    axes[i].set_title(f"Correlation = {corr}")
    axes[i].set_xlabel("Source Signal")
    axes[i].set_ylabel("Side Information")
    axes[i].grid(True, alpha=0.3)

    # Add correlation line
    x_range = np.linspace(source_np.min(), source_np.max(), 100)
    y_range = corr * x_range
    axes[i].plot(x_range, y_range, "r--", alpha=0.8, label=f"y = {corr}x")
    axes[i].legend()

plt.tight_layout()
plt.suptitle("Source-Side Information Correlation", y=1.02, fontsize=14)
plt.show()

Measure Empirical Correlation

Calculate actual correlation coefficients for validation

print("Empirical vs Theoretical Correlation:")
print("=====================================")

for corr in correlations:
    # Generate multiple samples and calculate correlation
    sources = []
    side_infos = []

    for i in range(100):  # Use 100 samples for statistics
        source, side_info = datasets[corr][i]
        sources.append(source.numpy().flatten())
        side_infos.append(side_info.numpy().flatten())

    # Combine all samples
    all_sources = np.concatenate(sources)
    all_side_infos = np.concatenate(side_infos)

    # Calculate empirical correlation
    empirical_corr = np.corrcoef(all_sources, all_side_infos)[0, 1]

    print(f"Theoretical: {corr:.2f}, Empirical: {empirical_corr:.3f}")

Empirical vs Theoretical Correlation:
=====================================
Theoretical: 0.20, Empirical: 0.189
Theoretical: 0.50, Empirical: 0.492
Theoretical: 0.80, Empirical: 0.797
Theoretical: 0.95, Empirical: 0.950

Time Series Visualization

Show how correlated signals evolve over time

plt.figure(figsize=(15, 8))

# Use high correlation for clearer visualization
high_corr_dataset = CorrelatedDataset(length=1, shape=(200,), correlation=0.85, noise_std=0.1, seed=42)

source, side_info = high_corr_dataset[0]
time_steps = np.arange(len(source))

plt.subplot(2, 1, 1)
plt.plot(time_steps, source.numpy(), "b-", label="Source Signal", linewidth=1.5)
plt.plot(time_steps, side_info.numpy(), "r--", label="Side Information", linewidth=1.5)
plt.title("Correlated Signals Over Time (ρ = 0.85)")
plt.xlabel("Time Step")
plt.ylabel("Amplitude")
plt.legend()
plt.grid(True, alpha=0.3)

# Show the difference signal
plt.subplot(2, 1, 2)
difference = source.numpy() - side_info.numpy()
plt.plot(time_steps, difference, "g-", linewidth=1.5)
plt.title("Difference Signal (Source - Side Information)")
plt.xlabel("Time Step")
plt.ylabel("Difference")
plt.grid(True, alpha=0.3)

plt.tight_layout()
plt.show()

Statistical Analysis

Analyze the statistical properties of the correlation model

print("\nStatistical Properties:")
print("======================")

for corr in [0.5, 0.8]:
    dataset = CorrelatedDataset(length=1000, shape=(64,), correlation=corr, noise_std=0.1, seed=42)

    # Collect statistics
    source_vars = []
    side_info_vars = []
    correlations_empirical = []

    for i in range(100):
        source, side_info = dataset[i]
        source_np = source.numpy()
        side_info_np = side_info.numpy()

        source_vars.append(np.var(source_np))
        side_info_vars.append(np.var(side_info_np))
        correlations_empirical.append(np.corrcoef(source_np, side_info_np)[0, 1])

    print(f"\nCorrelation {corr}:")
    print(f"  Source variance: {np.mean(source_vars):.3f} ± {np.std(source_vars):.3f}")
    print(f"  Side info variance: {np.mean(side_info_vars):.3f} ± {np.std(side_info_vars):.3f}")
    print(f"  Empirical correlation: {np.mean(correlations_empirical):.3f} ± {np.std(correlations_empirical):.3f}")

Statistical Properties:
======================

Correlation 0.5:
  Source variance: 0.988 ± 0.148
  Side info variance: 0.991 ± 0.176
  Empirical correlation: 0.504 ± 0.096

Correlation 0.8:
  Source variance: 0.988 ± 0.148
  Side info variance: 0.994 ± 0.161
  Empirical correlation: 0.801 ± 0.046