Multi-Level Discrete Wavelet Transform

def wavedec(data, wavelet, mode: str = 'symmetric', level: int = None, axis: int = -1):
    """
    Multi-level 1D discrete wavelet decomposition.
    
    Parameters:
    - data: Input 1D array or multi-dimensional array
    - wavelet: Wavelet specification (string name, Wavelet object, or filter bank)
    - mode: Signal extension mode for boundary handling
    - level: Decomposition level (default: maximum possible level)
    - axis: Axis along which to perform decomposition (default: -1, last axis)
    
    Returns:
    List [cAn, cDn, cDn-1, ..., cD1] where:
    - cAn: Approximation coefficients at level n
    - cDi: Detail coefficients at level i (from finest to coarsest)
    """

def waverec(coeffs, wavelet, mode: str = 'symmetric', axis: int = -1):
    """
    Multi-level 1D discrete wavelet reconstruction.
    
    Parameters:
    - coeffs: Coefficient list from wavedec [cAn, cDn, ..., cD1]
    - wavelet: Wavelet specification matching decomposition
    - mode: Signal extension mode matching decomposition
    - axis: Axis along which to perform reconstruction
    
    Returns:
    Reconstructed 1D signal
    """

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

# Create test signal with multiple frequency components
t = np.linspace(0, 2, 2048)
signal = (np.sin(2 * np.pi * 1 * t) +           # Low frequency
          0.5 * np.sin(2 * np.pi * 10 * t) +    # Medium frequency  
          0.3 * np.sin(2 * np.pi * 100 * t))    # High frequency
noise = 0.1 * np.random.randn(len(signal))
noisy_signal = signal + noise

# Multi-level decomposition
coeffs = pywt.wavedec(noisy_signal, 'db8', level=6)
print(f"Number of coefficient arrays: {len(coeffs)}")
print(f"Coefficient array lengths: {[len(c) for c in coeffs]}")

# Access approximation and detail coefficients
cA6 = coeffs[0]  # Approximation at level 6 (coarsest)
cD6 = coeffs[1]  # Detail at level 6 (coarsest detail)
cD5 = coeffs[2]  # Detail at level 5
cD4 = coeffs[3]  # Detail at level 4
cD3 = coeffs[4]  # Detail at level 3
cD2 = coeffs[5]  # Detail at level 2
cD1 = coeffs[6]  # Detail at level 1 (finest detail)

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

# Selective reconstruction - remove high-frequency noise (cD1, cD2)
denoised_coeffs = coeffs.copy()
denoised_coeffs[-1][:] = 0  # Zero out cD1 (finest details)
denoised_coeffs[-2][:] = 0  # Zero out cD2 (second finest details)
denoised_signal = pywt.waverec(denoised_coeffs, 'db8')

# Visualization
fig, axes = plt.subplots(4, 2, figsize=(15, 12))

# Original signals
axes[0, 0].plot(t, signal, 'b-', label='Clean')
axes[0, 0].plot(t, noisy_signal, 'r-', alpha=0.7, label='Noisy')
axes[0, 0].set_title('Original Signals')
axes[0, 0].legend()

axes[0, 1].plot(t, denoised_signal, 'g-', label='Denoised')
axes[0, 1].plot(t, signal, 'b--', alpha=0.7, label='True Signal')
axes[0, 1].set_title('Denoising Result')
axes[0, 1].legend()

# Plot coefficients at different levels
for i, (coeff, title) in enumerate(zip(coeffs[1:], 
                                     ['cD6 (Coarsest Detail)', 'cD5', 'cD4', 
                                      'cD3', 'cD2', 'cD1 (Finest Detail)'])):
    row = (i // 2) + 1
    col = i % 2
    if row < 4:
        axes[row, col].plot(coeff)
        axes[row, col].set_title(title)

plt.tight_layout()
plt.show()

# Automatic vs manual level selection
max_level = pywt.dwt_max_level(len(noisy_signal), pywt.Wavelet('db8').dec_len)
print(f"Maximum possible decomposition level: {max_level}")

# Decompose with automatic level
coeffs_auto = pywt.wavedec(noisy_signal, 'db8')
print(f"Automatic decomposition levels: {len(coeffs_auto) - 1}")

def wavedec2(data, wavelet, mode: str = 'symmetric', level: int = None, axes=(-2, -1)):
    """
    Multi-level 2D discrete wavelet decomposition.
    
    Parameters:
    - data: Input 2D array or multi-dimensional array
    - wavelet: Wavelet specification
    - mode: Signal extension mode for boundary handling
    - level: Decomposition level (default: maximum possible level)
    - axes: Pair of axes along which to perform 2D DWT (default: last two axes)
    
    Returns:
    List [cAn, (cHn, cVn, cDn), (cHn-1, cVn-1, cDn-1), ..., (cH1, cV1, cD1)] where:
    - cAn: Approximation coefficients at level n
    - (cHi, cVi, cDi): Horizontal, vertical, diagonal detail coefficients at level i
    """

def waverec2(coeffs, wavelet, mode: str = 'symmetric', axes=(-2, -1)):
    """
    Multi-level 2D discrete wavelet reconstruction.
    
    Parameters:
    - coeffs: Coefficient list from wavedec2
    - wavelet: Wavelet specification matching decomposition
    - mode: Signal extension mode matching decomposition
    - 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 with different frequency content
x, y = np.mgrid[0:256, 0:256]
image = (np.sin(2 * np.pi * x / 64) * np.cos(2 * np.pi * y / 64) +    # Low frequency
         0.5 * np.sin(2 * np.pi * x / 16) * np.cos(2 * np.pi * y / 16) +  # Medium frequency
         0.3 * np.sin(2 * np.pi * x / 4) * np.cos(2 * np.pi * y / 4))     # High frequency

# Add noise
noisy_image = image + 0.2 * np.random.randn(*image.shape)

# Multi-level 2D decomposition
coeffs = pywt.wavedec2(noisy_image, 'db4', level=4)
print(f"Number of decomposition levels: {len(coeffs) - 1}")
print(f"Approximation shape: {coeffs[0].shape}")
for i, (cH, cV, cD) in enumerate(coeffs[1:], 1):
    print(f"Level {len(coeffs)-i} detail shapes: {cH.shape}, {cV.shape}, {cD.shape}")

# Perfect reconstruction
reconstructed = pywt.waverec2(coeffs, 'db4')
print(f"2D Reconstruction error: {np.max(np.abs(noisy_image - reconstructed))}")

# Image denoising by coefficient thresholding
def threshold_coeffs(coeffs, threshold):
    """Apply soft thresholding to detail coefficients."""
    coeffs_thresh = [coeffs[0]]  # Keep approximation unchanged
    for cH, cV, cD in coeffs[1:]:
        cH_thresh = pywt.threshold(cH, threshold, mode='soft')
        cV_thresh = pywt.threshold(cV, threshold, mode='soft')
        cD_thresh = pywt.threshold(cD, threshold, mode='soft')
        coeffs_thresh.append((cH_thresh, cV_thresh, cD_thresh))
    return coeffs_thresh

# Apply thresholding and reconstruct
threshold_value = 0.1
coeffs_thresh = threshold_coeffs(coeffs, threshold_value)
denoised_image = pywt.waverec2(coeffs_thresh, 'db4')

# Visualization
fig, axes = plt.subplots(2, 3, figsize=(15, 10))

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

axes[0, 1].imshow(noisy_image, cmap='gray')
axes[0, 1].set_title('Noisy Image')
axes[0, 1].axis('off')

axes[0, 2].imshow(denoised_image, cmap='gray')
axes[0, 2].set_title('Denoised Image')
axes[0, 2].axis('off')

# Show approximation at coarsest level
axes[1, 0].imshow(coeffs[0], cmap='gray')
axes[1, 0].set_title(f'Approximation Level {len(coeffs)-1}')
axes[1, 0].axis('off')

# Show details at finest level
cH1, cV1, cD1 = coeffs[-1]
axes[1, 1].imshow(np.abs(cH1), cmap='gray')
axes[1, 1].set_title('Horizontal Details (Level 1)')
axes[1, 1].axis('off')

axes[1, 2].imshow(np.abs(cD1), cmap='gray')
axes[1, 2].set_title('Diagonal Details (Level 1)')
axes[1, 2].axis('off')

plt.tight_layout()
plt.show()

# Progressive reconstruction - build up from coarse to fine
progressive_images = []
for level in range(len(coeffs)):
    # Reconstruct using only levels 0 to level
    partial_coeffs = coeffs[:level+1]
    if level < len(coeffs) - 1:
        # Add zero details for missing levels
        zero_shape = coeffs[level+1][0].shape
        for missing_level in range(level+1, len(coeffs)):
            zero_details = (np.zeros(zero_shape), np.zeros(zero_shape), np.zeros(zero_shape))
            partial_coeffs.append(zero_details)
            zero_shape = (zero_shape[0] * 2, zero_shape[1] * 2)
    
    progressive_image = pywt.waverec2(partial_coeffs, 'db4')
    progressive_images.append(progressive_image)

# Show progressive reconstruction
fig, axes = plt.subplots(1, len(progressive_images), figsize=(20, 4))
for i, (ax, img) in enumerate(zip(axes, progressive_images)):
    ax.imshow(img, cmap='gray')
    ax.set_title(f'Up to Level {i}')
    ax.axis('off')
plt.tight_layout()
plt.show()

def wavedecn(data, wavelet, mode: str = 'symmetric', level: int = None, axes=None):
    """
    Multi-level nD discrete wavelet decomposition.
    
    Parameters:
    - data: Input nD array
    - wavelet: Wavelet specification
    - mode: Signal extension mode for boundary handling
    - level: Decomposition level (default: maximum possible level)
    - axes: Axes along which to perform DWT (default: all axes)
    
    Returns:
    List [cAn, {details_level_n}, {details_level_n-1}, ..., {details_level_1}] where:
    - cAn: Approximation coefficients at level n
    - {details_level_i}: Dictionary of detail coefficients at level i
    """

def waverecn(coeffs, wavelet, mode: str = 'symmetric', axes=None):
    """
    Multi-level nD discrete wavelet reconstruction.
    
    Parameters:
    - coeffs: Coefficient list from wavedecn
    - wavelet: Wavelet specification matching decomposition
    - mode: Signal extension mode matching decomposition
    - axes: Axes along which to perform IDWT (should match decomposition)
    
    Returns:
    Reconstructed nD array
    """

import pywt
import numpy as np

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

# 3D multi-level decomposition
coeffs_3d = pywt.wavedecn(volume, 'db2', level=3)
print(f"Number of decomposition levels: {len(coeffs_3d) - 1}")
print(f"Approximation shape: {coeffs_3d[0].shape}")

# Each level has 2^n - 1 detail coefficient arrays (n = number of dimensions)
for level, details in enumerate(coeffs_3d[1:], 1):
    print(f"Level {len(coeffs_3d) - level} details:")
    for key, coeff in details.items():
        print(f"  '{key}': {coeff.shape}")

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

# 4D example - 3D spatial + 1D temporal
data_4d = np.random.randn(32, 32, 32, 100)
coeffs_4d = pywt.wavedecn(data_4d, 'haar', level=2)
print(f"4D data shape: {data_4d.shape}")
print(f"4D approximation shape: {coeffs_4d[0].shape}")
print(f"Number of detail types at each level: {len(coeffs_4d[1])}")  # 2^4 - 1 = 15

# Show some detail coefficient keys for 4D
print("4D detail coefficient keys (level 1):")
for key in sorted(coeffs_4d[-1].keys()):
    print(f"  '{key}': {coeffs_4d[-1][key].shape}")

def downcoef(part: str, data, wavelet, mode: str = 'symmetric', level: int = 1):
    """
    Partial DWT - compute only approximation or detail coefficients.
    
    Parameters:
    - part: 'a' for approximation, 'd' for detail coefficients
    - data: Input signal
    - wavelet: Wavelet specification
    - mode: Signal extension mode
    - level: Decomposition level
    
    Returns:
    Requested coefficient type at specified level
    """

def upcoef(part: str, coeffs, wavelet, level: int = 1, take: int = 0):
    """
    Direct reconstruction from single coefficient type.
    
    Parameters:
    - part: 'a' for approximation, 'd' for detail coefficients
    - coeffs: Input coefficients
    - wavelet: Wavelet specification  
    - level: Number of reconstruction levels
    - take: Number of coefficients to take (0 = all)
    
    Returns:
    Reconstructed signal from specified coefficient type only
    """

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

# Generate test signal
t = np.linspace(0, 1, 1024)
signal = np.sin(2 * np.pi * 5 * t) + 0.5 * np.sin(2 * np.pi * 50 * t)

# Extract only approximation at level 3
approx_level3 = pywt.downcoef('a', signal, 'db4', level=3)
print(f"Approximation at level 3 shape: {approx_level3.shape}")

# Extract only detail at level 2  
detail_level2 = pywt.downcoef('d', signal, 'db4', level=2)
print(f"Detail at level 2 shape: {detail_level2.shape}")

# Reconstruct signal from approximation only
reconstructed_approx = pywt.upcoef('a', approx_level3, 'db4', level=3)
print(f"Reconstructed from approximation shape: {reconstructed_approx.shape}")

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

# Visualization
plt.figure(figsize=(12, 8))

plt.subplot(2, 2, 1)
plt.plot(t, signal)
plt.title('Original Signal')

plt.subplot(2, 2, 2)
plt.plot(reconstructed_approx[:len(signal)])
plt.title('Reconstructed from Approximation (Level 3)')

plt.subplot(2, 2, 3)
plt.plot(reconstructed_detail[:len(signal)])
plt.title('Reconstructed from Detail (Level 2)')

plt.subplot(2, 2, 4)
plt.plot(t, signal, 'b-', label='Original')
plt.plot(reconstructed_approx[:len(signal)] + reconstructed_detail[:len(signal)], 
         'r--', label='Approx + Detail')
plt.title('Partial Reconstruction Sum')
plt.legend()

plt.tight_layout()
plt.show()

# Taking partial coefficients
partial_approx = pywt.upcoef('a', approx_level3[:len(approx_level3)//2], 'db4', level=3)
print(f"Reconstruction from half coefficients: {partial_approx.shape}")

def fswavedecn(data, wavelet, mode: str = 'symmetric', levels=None, axes=None):
    """
    Fully separable wavelet decomposition.
    
    Parameters:
    - data: Input nD array
    - wavelet: Wavelet specification or list of wavelets for each axis
    - mode: Signal extension mode or list of modes for each axis
    - levels: Decomposition levels as int or list of ints for each axis
    - axes: Axes along which to perform transform (default: all axes)
    
    Returns:
    FswavedecnResult object with coefficient access and reconstruction methods
    """

def fswaverecn(fswavedecn_result):
    """
    Inverse fully separable wavelet transform.
    
    Parameters:
    - fswavedecn_result: FswavedecnResult object from fswavedecn
    
    Returns:
    Reconstructed nD array
    """

class FswavedecnResult:
    """Container for fully separable decomposition results."""
    
    # Properties
    coeffs: np.ndarray  # Coefficient array
    coeff_slices: list  # Coefficient slice information
    axes: tuple  # Transform axes
    wavelets: list  # Wavelets used for each axis
    modes: list  # Extension modes for each axis
    levels: list  # Decomposition levels for each axis
    approx: np.ndarray  # Approximation coefficients
    
    def __getitem__(self, levels):
        """Access coefficients at specified levels."""
    
    def __setitem__(self, levels, x):
        """Set coefficients at specified levels."""
    
    def detail_keys(self):
        """List all detail coefficient keys."""

# Multi-level coefficient formats
MultiLevelCoeffs1D = List[np.ndarray]  # [cAn, cDn, cDn-1, ..., cD1]

MultiLevelCoeffs2D = List[
    Union[
        np.ndarray,  # Approximation coefficients (first element)
        Tuple[np.ndarray, np.ndarray, np.ndarray]  # (cH, cV, cD) detail tuples
    ]
]

MultiLevelCoeffsND = List[
    Union[
        np.ndarray,  # Approximation coefficients (first element)  
        Dict[str, np.ndarray]  # Detail coefficient dictionaries
    ]
]

# Coefficient part specifications
CoeffPart = Literal['a', 'd']  # 'a' for approximation, 'd' for detail