Single-Level Discrete Wavelet Transform

def dwt(data, wavelet, mode: str = 'symmetric', axis: int = -1):
    """
    Single-level 1D discrete wavelet transform.
    
    Parameters:
    - data: Input 1D array or array with axis specified
    - wavelet: Wavelet specification (string name, Wavelet object, or filter bank)
    - mode: Signal extension mode for boundary handling
    - axis: Axis along which to perform DWT (default: -1, last axis)
    
    Returns:
    (cA, cD) - approximation and detail coefficients as tuple
    """

def idwt(cA, cD, wavelet, mode: str = 'symmetric', axis: int = -1):
    """
    Single-level 1D inverse discrete wavelet transform.
    
    Parameters:
    - cA: Approximation coefficients (can be None)
    - cD: Detail coefficients (can be None)  
    - wavelet: Wavelet specification matching forward transform
    - mode: Signal extension mode matching forward transform
    - axis: Axis along which to perform IDWT
    
    Returns:
    Reconstructed 1D signal
    """

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

# Create test signal with noise
t = np.linspace(0, 1, 1000)
signal = np.sin(2 * np.pi * 5 * t) + 0.5 * np.sin(2 * np.pi * 20 * t)
noisy_signal = signal + 0.3 * np.random.randn(len(signal))

# Single-level DWT
cA, cD = pywt.dwt(noisy_signal, 'db4')
print(f"Original signal length: {len(noisy_signal)}")
print(f"Approximation coefficients: {len(cA)}")
print(f"Detail coefficients: {len(cD)}")

# Perfect reconstruction
reconstructed = pywt.idwt(cA, cD, 'db4')
print(f"Reconstruction error: {np.max(np.abs(noisy_signal - reconstructed))}")

# Denoising by thresholding detail coefficients
threshold = 0.1 * np.max(np.abs(cD))
cD_thresh = np.where(np.abs(cD) > threshold, cD, 0)
denoised = pywt.idwt(cA, cD_thresh, 'db4')

# Plotting
plt.figure(figsize=(12, 8))
plt.subplot(2, 2, 1)
plt.plot(t, signal, 'b-', label='Original')
plt.plot(t, noisy_signal, 'r-', alpha=0.7, label='Noisy')
plt.title('Original vs Noisy Signal')
plt.legend()

plt.subplot(2, 2, 2)
plt.plot(cA, 'g-', label='Approximation')
plt.title('Approximation Coefficients')
plt.legend()

plt.subplot(2, 2, 3)
plt.plot(cD, 'r-', label='Detail')
plt.plot(cD_thresh, 'k-', label='Thresholded')
plt.title('Detail Coefficients')
plt.legend()

plt.subplot(2, 2, 4)
plt.plot(t, signal, 'b-', label='Original')
plt.plot(t, denoised, 'g-', label='Denoised')
plt.title('Original vs Denoised')
plt.legend()
plt.tight_layout()
plt.show()

# Multi-dimensional data along specific axis
data_2d = np.random.randn(50, 1000)
cA_2d, cD_2d = pywt.dwt(data_2d, 'haar', axis=1)  # Transform along columns
print(f"2D data shape: {data_2d.shape}")
print(f"2D coefficients shape: {cA_2d.shape}, {cD_2d.shape}")

def dwt2(data, wavelet, mode: str = 'symmetric', axes=(-2, -1)):
    """
    Single-level 2D discrete wavelet transform.
    
    Parameters:
    - data: Input 2D array or multi-dimensional array
    - wavelet: Wavelet specification
    - mode: Signal extension mode for boundary handling
    - axes: Pair of axes along which to perform 2D DWT (default: last two axes)
    
    Returns:
    (cA, (cH, cV, cD)) - approximation and horizontal/vertical/diagonal detail coefficients
    """

def idwt2(coeffs, wavelet, mode: str = 'symmetric', axes=(-2, -1)):
    """
    Single-level 2D inverse discrete wavelet transform.
    
    Parameters:
    - coeffs: Coefficient tuple (cA, (cH, cV, cD)) from dwt2
    - wavelet: Wavelet specification matching forward transform
    - mode: Signal extension mode matching forward transform
    - axes: Pair of axes along which to perform 2D IDWT
    
    Returns:
    Reconstructed 2D array
    """

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

# Create test image
image = np.zeros((128, 128))
image[30:98, 30:98] = 1  # Square in center
image[40:88, 40:88] = 0.5  # Smaller square inside
# Add some texture
image += 0.1 * np.random.randn(128, 128)

# 2D DWT
coeffs = pywt.dwt2(image, 'db2')
cA, (cH, cV, cD) = coeffs

print(f"Original image shape: {image.shape}")
print(f"Approximation shape: {cA.shape}")
print(f"Detail shapes: {cH.shape}, {cV.shape}, {cD.shape}")

# Perfect reconstruction
reconstructed = pywt.idwt2(coeffs, 'db2')
print(f"Reconstruction error: {np.max(np.abs(image - reconstructed))}")

# Visualize decomposition
fig, axes = plt.subplots(2, 3, figsize=(12, 8))

axes[0, 0].imshow(image, cmap='gray')
axes[0, 0].set_title('Original Image')

axes[0, 1].imshow(cA, cmap='gray')
axes[0, 1].set_title('Approximation (cA)')

axes[0, 2].imshow(reconstructed, cmap='gray')
axes[0, 2].set_title('Reconstructed')

axes[1, 0].imshow(cH, cmap='gray')
axes[1, 0].set_title('Horizontal Details (cH)')

axes[1, 1].imshow(cV, cmap='gray')
axes[1, 1].set_title('Vertical Details (cV)')

axes[1, 2].imshow(cD, cmap='gray')
axes[1, 2].set_title('Diagonal Details (cD)')

for ax in axes.flat:
    ax.axis('off')

plt.tight_layout()
plt.show()

# Image compression by coefficient thresholding
threshold = 0.1
cA_thresh = cA  # Keep all approximation coefficients
cH_thresh = np.where(np.abs(cH) > threshold, cH, 0)
cV_thresh = np.where(np.abs(cV) > threshold, cV, 0)
cD_thresh = np.where(np.abs(cD) > threshold, cD, 0)

compressed_coeffs = (cA_thresh, (cH_thresh, cV_thresh, cD_thresh))
compressed_image = pywt.idwt2(compressed_coeffs, 'db2')

# Calculate compression ratio
original_nonzero = np.count_nonzero(image)
compressed_nonzero = (np.count_nonzero(cA_thresh) + 
                     np.count_nonzero(cH_thresh) + 
                     np.count_nonzero(cV_thresh) + 
                     np.count_nonzero(cD_thresh))
compression_ratio = original_nonzero / compressed_nonzero

print(f"Compression ratio: {compression_ratio:.2f}")

def dwtn(data, wavelet, mode: str = 'symmetric', axes=None):
    """
    Single-level nD discrete wavelet transform.
    
    Parameters:
    - data: Input nD array
    - wavelet: Wavelet specification  
    - mode: Signal extension mode for boundary handling
    - axes: Axes along which to perform DWT (default: all axes)
    
    Returns:
    Dictionary of coefficients with string keys indicating coefficient type.
    For 2D: {'aa': cA, 'ad': cH, 'da': cV, 'dd': cD}
    For 3D: {'aaa': cAAA, 'aad': cAAD, 'ada': cADA, ...} (8 coefficients)
    """

def idwtn(coeffs, wavelet, mode: str = 'symmetric', axes=None):
    """
    Single-level nD inverse discrete wavelet transform.
    
    Parameters:
    - coeffs: Dictionary of coefficients from dwtn
    - wavelet: Wavelet specification matching forward transform
    - mode: Signal extension mode matching forward transform  
    - axes: Axes along which to perform IDWT (should match forward transform)
    
    Returns:
    Reconstructed nD array
    """

import pywt
import numpy as np

# 3D volume data example
volume = np.random.randn(32, 32, 32)
print(f"Original volume shape: {volume.shape}")

# 3D DWT
coeffs_3d = pywt.dwtn(volume, 'db2')
print(f"Number of coefficient arrays: {len(coeffs_3d)}")
print(f"Coefficient keys: {list(coeffs_3d.keys())}")

# Each coefficient array is 1/8 the size in 3D
for key, coeff in coeffs_3d.items():
    print(f"Coefficient '{key}' shape: {coeff.shape}")

# Perfect reconstruction
reconstructed_3d = pywt.idwtn(coeffs_3d, 'db2')
print(f"3D reconstruction error: {np.max(np.abs(volume - reconstructed_3d))}")

# 2D example using dwtn (equivalent to dwt2)
image_2d = np.random.randn(64, 64)
coeffs_2d = pywt.dwtn(image_2d, 'haar')
print(f"2D coefficient keys: {list(coeffs_2d.keys())}")
print(f"Approximation 'aa' shape: {coeffs_2d['aa'].shape}")

# Compare with dwt2
coeffs_dwt2 = pywt.dwt2(image_2d, 'haar')
cA, (cH, cV, cD) = coeffs_dwt2

# Verify equivalence
print(f"dwtn 'aa' equals dwt2 cA: {np.allclose(coeffs_2d['aa'], cA)}")
print(f"dwtn 'ad' equals dwt2 cH: {np.allclose(coeffs_2d['ad'], cH)}")
print(f"dwtn 'da' equals dwt2 cV: {np.allclose(coeffs_2d['da'], cV)}")
print(f"dwtn 'dd' equals dwt2 cD: {np.allclose(coeffs_2d['dd'], cD)}")

# Partial transforms along specific axes
data_4d = np.random.randn(16, 16, 16, 100)  # 3D spatial + 1D time
coeffs_partial = pywt.dwtn(data_4d, 'db1', axes=(0, 1, 2))  # Transform only spatial dimensions
print(f"Partial transform result shape for 'aaa': {coeffs_partial['aaa'].shape}")

def dwt_max_level(data_len: int, filter_len: int) -> int:
    """
    Compute maximum useful decomposition level for 1D DWT.
    
    Parameters:
    - data_len: Length of input data
    - filter_len: Length of wavelet filter
    
    Returns:
    Maximum decomposition level
    """

def dwt_coeff_len(data_len: int, filter_len: int, mode: str) -> int:
    """
    Compute length of DWT coefficient arrays.
    
    Parameters:
    - data_len: Length of input data
    - filter_len: Length of wavelet filter
    - mode: Signal extension mode
    
    Returns:
    Length of coefficient arrays
    """

def dwtn_max_level(shape: tuple, wavelet, axes=None) -> int:
    """
    Compute maximum decomposition level for nD data.
    
    Parameters:
    - shape: Shape of input data
    - wavelet: Wavelet specification
    - axes: Axes for transform (default: all axes)
    
    Returns:
    Maximum decomposition level
    """

def pad(x, pad_widths, mode: str):
    """
    Extend signal using PyWavelets extension modes.
    
    Parameters:
    - x: Input array
    - pad_widths: Padding widths for each axis  
    - mode: PyWavelets extension mode
    
    Returns:
    Padded array
    """

def downcoef(part: str, data, wavelet, mode: str = 'symmetric', level: int = 1):
    """
    Partial discrete wavelet transform decomposition.
    
    Parameters:
    - part: Coefficients type ('a' for approximation, 'd' for details)
    - data: Input signal array
    - wavelet: Wavelet specification
    - mode: Signal extension mode
    - level: Decomposition level (default: 1)
    
    Returns:
    1D array of requested coefficients
    """

def upcoef(part: str, coeffs, wavelet, level: int = 1, take: int = 0):
    """
    Direct reconstruction from wavelet coefficients.
    
    Parameters:
    - part: Coefficients type ('a' for approximation, 'd' for details)
    - coeffs: Coefficient array to reconstruct from
    - wavelet: Wavelet specification
    - level: Multilevel reconstruction level (default: 1)
    - take: Take central part of specified length (0 for full length)
    
    Returns:
    1D array with reconstructed signal
    """

import pywt
import numpy as np

# Calculate maximum decomposition levels
data_1d = np.random.randn(1000)
wavelet = 'db4'
filter_len = pywt.Wavelet(wavelet).dec_len

max_level_1d = pywt.dwt_max_level(len(data_1d), filter_len)
print(f"Max level for 1D data (length {len(data_1d)}): {max_level_1d}")

# Calculate coefficient lengths for different modes
coeff_len_sym = pywt.dwt_coeff_len(len(data_1d), filter_len, 'symmetric')
coeff_len_per = pywt.dwt_coeff_len(len(data_1d), filter_len, 'periodization')
print(f"Coefficient length (symmetric): {coeff_len_sym}")
print(f"Coefficient length (periodization): {coeff_len_per}")

# nD maximum levels
image_2d = np.random.randn(256, 256)
max_level_2d = pywt.dwtn_max_level(image_2d.shape, wavelet)
print(f"Max level for 2D image {image_2d.shape}: {max_level_2d}")

volume_3d = np.random.randn(64, 64, 64)
max_level_3d = pywt.dwtn_max_level(volume_3d.shape, wavelet)
print(f"Max level for 3D volume {volume_3d.shape}: {max_level_3d}")

# Signal padding examples
signal = np.array([1, 2, 3, 4, 5])
pad_widths = [(2, 2)]  # Pad 2 elements on each side

padded_zero = pywt.pad(signal, pad_widths, 'zero')
padded_symmetric = pywt.pad(signal, pad_widths, 'symmetric')
padded_periodic = pywt.pad(signal, pad_widths, 'periodic')

print(f"Original: {signal}")
print(f"Zero padding: {padded_zero}")
print(f"Symmetric padding: {padded_symmetric}")
print(f"Periodic padding: {padded_periodic}")

# Partial decomposition with downcoef
test_signal = np.sin(2 * np.pi * np.linspace(0, 1, 256))

# Get only approximation coefficients at level 2
approx_only = pywt.downcoef('a', test_signal, 'db4', level=2)
print(f"Approximation coefficients shape: {approx_only.shape}")

# Get only detail coefficients at level 1  
detail_only = pywt.downcoef('d', test_signal, 'db4', level=1)
print(f"Detail coefficients shape: {detail_only.shape}")

# Compare with full DWT
coeffs_full = pywt.wavedec(test_signal, 'db4', level=2)
cA2, cD2, cD1 = coeffs_full
print(f"Full decomposition - cA2: {cA2.shape}, matches downcoef: {np.allclose(cA2, approx_only)}")

# Direct reconstruction with upcoef
# Reconstruct signal from approximation coefficients only
reconstructed_approx = pywt.upcoef('a', approx_only, 'db4', level=2)
print(f"Reconstructed from approx shape: {reconstructed_approx.shape}")

# Reconstruct signal from detail coefficients only
reconstructed_detail = pywt.upcoef('d', detail_only, 'db4', level=1)  
print(f"Reconstructed from detail shape: {reconstructed_detail.shape}")

# Visualization of partial reconstruction
plt.figure(figsize=(12, 8))
plt.subplot(2, 2, 1)
plt.plot(test_signal, 'b-', label='Original')
plt.title('Original Signal')
plt.legend()

plt.subplot(2, 2, 2) 
plt.plot(reconstructed_approx[:len(test_signal)], 'g-', label='From Approximation')
plt.title('Reconstructed from Approximation Only')
plt.legend()

plt.subplot(2, 2, 3)
plt.plot(reconstructed_detail[:len(test_signal)], 'r-', label='From Details')
plt.title('Reconstructed from Details Only')
plt.legend()

plt.subplot(2, 2, 4)
combined = reconstructed_approx[:len(test_signal)] + reconstructed_detail[:len(test_signal)]
plt.plot(test_signal, 'b-', alpha=0.7, label='Original')
plt.plot(combined, 'k--', label='Approx + Detail')
plt.title('Combined Reconstruction')
plt.legend()
plt.tight_layout()
plt.show()

# Single-level coefficient formats
Coeffs1D = Tuple[np.ndarray, np.ndarray]  # (cA, cD)
Coeffs2D = Tuple[np.ndarray, Tuple[np.ndarray, np.ndarray, np.ndarray]]  # (cA, (cH, cV, cD))
CoeffsND = Dict[str, np.ndarray]  # String keys like 'aa', 'ad', 'da', 'dd' for 2D

# Input data types
ArrayLike = Union[np.ndarray, list, tuple]

# Extension modes
Mode = Literal[
    'zero', 'constant', 'symmetric', 'periodic', 
    'smooth', 'periodization', 'reflect', 
    'antisymmetric', 'antireflect'
]