.. DO NOT EDIT. .. THIS FILE WAS AUTOMATICALLY GENERATED BY SPHINX-GALLERY. .. TO MAKE CHANGES, EDIT THE SOURCE PYTHON FILE: .. "gallery/0100_lop/wt_op.py" .. LINE NUMBERS ARE GIVEN BELOW. .. only:: html .. note:: :class: sphx-glr-download-link-note :ref:`Go to the end ` to download the full example code .. rst-class:: sphx-glr-example-title .. _sphx_glr_gallery_0100_lop_wt_op.py: .. _gallery:lop:wt_ops: Wavelet Transform Operators ============================== .. contents:: :depth: 2 :local: This example demonstrates following features: - ``cr.sparse.lop.convolve2D`` A 2D convolution linear operator - ``cr.sparse.sls.lsqr`` LSQR algorithm for solving a least square problem on 2D images .. GENERATED FROM PYTHON SOURCE LINES 20-21 Let's import necessary libraries .. GENERATED FROM PYTHON SOURCE LINES 21-36 .. code-block:: default import jax.numpy as jnp # For plotting diagrams import matplotlib.pyplot as plt ## CR-Sparse modules import cr.nimble as cnb # Linear operators from cr.sparse import lop # Sample images import skimage.data # Utilities from cr.nimble.dsp import time_values # Configure JAX for 64-bit computing from jax.config import config config.update("jax_enable_x64", True) .. GENERATED FROM PYTHON SOURCE LINES 37-39 1D Wavelet Transform Operator --------------------------------------------------- .. GENERATED FROM PYTHON SOURCE LINES 42-46 A signal consisting of multiple sinusoids '''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''' Individual sinusoids have different frequencies and amplitudes. Sampling frequency .. GENERATED FROM PYTHON SOURCE LINES 46-63 .. code-block:: default fs = 1000. # Time duration T = 2 # time values t = time_values(fs, T) # Number of samples n = t.size x = jnp.zeros(n) freqs = [25, 7, 9] amps = [1, -3, .8] for (f, amp) in zip(freqs, amps): sinusoid = amp * jnp.sin(2 * jnp.pi * f * t) x = x + sinusoid # Plot the signal plt.figure(figsize=(8,2)) plt.plot(t, x, 'k', label='Composite signal') .. image-sg:: /gallery/0100_lop/images/sphx_glr_wt_op_001.png :alt: wt op :srcset: /gallery/0100_lop/images/sphx_glr_wt_op_001.png :class: sphx-glr-single-img .. rst-class:: sphx-glr-script-out .. code-block:: none [] .. GENERATED FROM PYTHON SOURCE LINES 64-66 1D wavelet transform operator '''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''' .. GENERATED FROM PYTHON SOURCE LINES 66-68 .. code-block:: default DWT_op = lop.dwt(n, wavelet='dmey', level=5) .. GENERATED FROM PYTHON SOURCE LINES 69-71 Wavelet coefficients '''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''' .. GENERATED FROM PYTHON SOURCE LINES 71-75 .. code-block:: default alpha = DWT_op.times(x) plt.figure(figsize=(8,2)) plt.plot(alpha, label='Wavelet coefficients') .. image-sg:: /gallery/0100_lop/images/sphx_glr_wt_op_002.png :alt: wt op :srcset: /gallery/0100_lop/images/sphx_glr_wt_op_002.png :class: sphx-glr-single-img .. rst-class:: sphx-glr-script-out .. code-block:: none [] .. GENERATED FROM PYTHON SOURCE LINES 76-79 Compression '''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''' Let's keep only 10 percent of the coefficients .. GENERATED FROM PYTHON SOURCE LINES 79-84 .. code-block:: default cutoff = n // 10 alpha2 = alpha.at[cutoff:].set(0) plt.figure(figsize=(8,2)) plt.plot(alpha2, label='Wavelet coefficients after compression') .. image-sg:: /gallery/0100_lop/images/sphx_glr_wt_op_003.png :alt: wt op :srcset: /gallery/0100_lop/images/sphx_glr_wt_op_003.png :class: sphx-glr-single-img .. rst-class:: sphx-glr-script-out .. code-block:: none [] .. GENERATED FROM PYTHON SOURCE LINES 85-87 Reconstruction '''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''' .. GENERATED FROM PYTHON SOURCE LINES 87-101 .. code-block:: default x_rec = DWT_op.trans(alpha2) # RMSE rmse = cnb.root_mse(x, x_rec) print(rmse) # SNR snr = cnb.signal_noise_ratio(x, x_rec) print(snr) plt.figure(figsize=(8,2)) plt.plot(x, 'k', label='Original') plt.plot(x_rec, 'r', label='Reconstructed') plt.title('Reconstructed signal') plt.legend() .. image-sg:: /gallery/0100_lop/images/sphx_glr_wt_op_004.png :alt: Reconstructed signal :srcset: /gallery/0100_lop/images/sphx_glr_wt_op_004.png :class: sphx-glr-single-img .. rst-class:: sphx-glr-script-out .. code-block:: none 0.034259365997888466 36.56352988951201 .. GENERATED FROM PYTHON SOURCE LINES 102-104 2D Wavelet Transform Operator --------------------------------------------------- .. GENERATED FROM PYTHON SOURCE LINES 104-111 .. code-block:: default # Sample image image = skimage.data.grass() DWT2_op = lop.dwt2D(image.shape, wavelet='haar', level=5) DWT2_op = lop.jit(DWT2_op) .. GENERATED FROM PYTHON SOURCE LINES 112-114 Wavelet coefficients '''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''' .. GENERATED FROM PYTHON SOURCE LINES 114-116 .. code-block:: default coefs = DWT2_op.times(image) .. GENERATED FROM PYTHON SOURCE LINES 117-120 Compression '''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''' Let's keep only 1/16 of the coefficients .. GENERATED FROM PYTHON SOURCE LINES 120-124 .. code-block:: default h, w = coefs.shape coefs2 = jnp.zeros_like(coefs) coefs2 = coefs2.at[:h//4, :w//4].set(coefs[:h//4, :w//4]) .. GENERATED FROM PYTHON SOURCE LINES 125-127 Reconstruction '''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''' .. GENERATED FROM PYTHON SOURCE LINES 127-152 .. code-block:: default image_rec = DWT2_op.trans(coefs2) # RMSE rmse = cnb.root_mse(image, image_rec) print(rmse) # PSNR psnr = cnb.peak_signal_noise_ratio(image, image_rec) print(psnr) # Plot everything fig, axs = plt.subplots(1, 4, figsize=(16, 3)) axs[0].imshow(image, cmap='gray') axs[0].set_title('Image') axs[0].axis('tight') axs[1].imshow(coefs, cmap='gray_r', vmin=-1e2, vmax=1e2) axs[1].set_title('DWT2 coefficients') axs[1].axis('tight') axs[2].imshow(coefs2, cmap='gray_r', vmin=-1e2, vmax=1e2) axs[2].set_title('After compression') axs[2].axis('tight') axs[3].imshow(image_rec, cmap='gray') axs[3].set_title('Reconstructed image') axs[3].axis('tight') .. image-sg:: /gallery/0100_lop/images/sphx_glr_wt_op_005.png :alt: Image, DWT2 coefficients, After compression, Reconstructed image :srcset: /gallery/0100_lop/images/sphx_glr_wt_op_005.png :class: sphx-glr-single-img .. rst-class:: sphx-glr-script-out .. code-block:: none 27.383743287757877 19.380947312633058 (-0.5, 511.5, 511.5, -0.5) .. rst-class:: sphx-glr-timing **Total running time of the script:** (0 minutes 4.833 seconds) .. _sphx_glr_download_gallery_0100_lop_wt_op.py: .. only:: html .. container:: sphx-glr-footer sphx-glr-footer-example .. container:: sphx-glr-download sphx-glr-download-python :download:`Download Python source code: wt_op.py ` .. container:: sphx-glr-download sphx-glr-download-jupyter :download:`Download Jupyter notebook: wt_op.ipynb ` .. only:: html .. rst-class:: sphx-glr-signature `Gallery generated by Sphinx-Gallery `_