r""" .. _gallery:cs:bsbl:1: Block Sparse Bayesian Learning ================================= .. contents:: :depth: 2 :local: In this example, we demonstrate the BSBL (Block Sparse Bayesian Learning) algorithm :cite:`zhang2012recovery,zhang2013extension` for reconstruction of block sparse signals with intra block correlations from their compressive measurements. In particular, we show - Creation of block sparse signals with intra-block correlation - Compressive sampling of the signal with Gaussian and sparse binary sensing matrices. - Reconstruction using BSBL EM algorithm (Expectation Maximization). - Reconstruction using BSBL BO algorithm (Bound Optimization). - Reconstruction in the presence of high measurement noise Our implementation of BSBL is fully JIT compilable. To achieve this, we limit ourselves of equal sized blocks where the block size is user defined. This is not a problem in practice. As shown in :cite:`zhang2012compressed`, the reconstruction from compressive measurements of real life signals is not affected much by the block size. The basic compressive sensing model is given by .. math:: \by = \Phi \bx + \be where :math:`\by` is a known measurement vector, :math:`\Phi` is a known sensing matrix and :math:`\bx` is a sparse signal to be recovered from the measurements. We introduce the block/group structure on :math:`\bx` as .. math:: \bx = \begin{pmatrix} \bx_1 & \bx_2 & \dots & \bx_g \end{pmatrix} where each :math:`\bx_i` is a block of :math:`b` values. The signal :math:`\bx` consists of :math:`g` such blocks/groups. We only consider the case of equal sized blocks in our implementation. Under the block sparsity model, only a few :math:`k \ll g` blocks are nonzero (active) in the signal :math:`\bx` however, the locations of these blocks are unknown. We can rewrite the sensing equation as: .. math:: \by = \sum_{i=1}^g \Phi_i \bx_i + \be by splitting the sensing matrix into blocks of columns appropriately. Under the sparse Bayesian framework, each block is assumed to satisfy a parametrized multivariate Gaussian distribution: .. math:: \PP(\bx_i ; \gamma_i, \bB_i) = \NNN(\bzero, \gamma_i \bB_i), \Forall i=1,\dots,g. The covariance matrix :math:`\bB_i` captures the intra block correlations. We further assume that the blocks are mutually uncorrelated. The prior of :math:`\bx` can then be written as .. math:: \PP(\bx; \{ \gamma_i, \bB_i\}_i ) = \NNN(\bzero, \Sigma_0) where .. math:: \Sigma_0 = \text{diag}\{\gamma_1 \bB_1, \dots, \gamma_g \bB_g \}. We also model the correlation among the values within each active block as an AR-1 process. Under this assumption the matrices :math:`\bB_i` take the form of a Toeplitz matrix .. math:: \bB = \begin{bmatrix} 1 & r & \dots & r^{b-1}\\ r & 1 & \dots & r^{b-2}\\ \vdots & & \ddots & \vdots\\ r^{b-1} & r^{b-2} & \dots & 1 \end{bmatrix} where :math:`r` is the AR-1 model coefficient. This constraint significantly reduces the model parameters to be learned. Measurement noise is modeled as independent zero mean Gaussian noise :math:`\PP(\be; \lambda) \sim \NNN(\bzero, \lambda \bI)`. BSBL doesn't require you to provide the value of noise variance as input. It is able to estimate :math:`\lambda` within a algorithm. The estimate of :math:`\bx` under Bayesian learning framework is given by the posterior mean of :math:`\bx` given the measurements :math:`\by`. Please also refer to the `BSBL website `_ by the authors of the original algorithm for further information. Related Examples - :ref:`gallery:cs:sparse_binary_sensor` """ # %% # Setup # ------------------------------ # Configure JAX for 64-bit computing from jax.config import config config.update("jax_enable_x64", True) # %% # Let's import necessary libraries from matplotlib import pyplot as plt # jax imports import jax.numpy as jnp from jax import random, jit # cr-suite imports import cr.nimble as crn import cr.nimble.dsp as crdsp import cr.sparse as crs import cr.sparse.dict as crdict import cr.sparse.data as crdata import cr.sparse.plots as crplot import cr.sparse.block.bsbl as bsbl # %% # Problem Configuration # ------------------------------ # ambient dimension n = 300 # block length b = 4 # number of blocks nb = n // b # Block sparsity: number of nonzero blocks k = 6 # Number of measurements m = 100 # %% # Block Sparse Signal # ------------------------------ # Block sparse signal with intra block correlation x, blocks, indices = crdata.sparse_normal_blocks( crn.KEYS[2], n, k, b, cor=0.9, normalize_blocks=True) ax = crplot.one_plot() ax.stem(x); # %% # Gaussian Sensing # ------------------------------ # %% # Sensing matrix Phi = crdict.gaussian_mtx(crn.KEYS[0], m, n, normalize_atoms=True) ax = crplot.one_plot() ax.imshow(Phi); # %% # Measurements y = Phi @ x ax = crplot.one_plot() crplot.plot_signal(ax, y) # %% # Reconstruction using BSBL EM # ''''''''''''''''''''''''''''''''''' # We need to provide the sensing matrix, measurements # and the block size as parameters to the # reconstruction algorithm options = bsbl.bsbl_em_options(y, learn_lambda=0) sol = bsbl.bsbl_em_jit(Phi, y, b, options) print(sol) # %% # Reconstructed signal x_hat = sol.x print(f'BSBL-EM: PRD: {crn.prd(x, x_hat):.1f} %, SNR: {crn.signal_noise_ratio(x, x_hat):.2f} dB.' ) # %% # Plot the original and reconstructed signal ax = crplot.h_plots(2) ax[0].stem(x) ax[1].stem(x_hat) # %% # Reconstruction using BSBL BO # ''''''''''''''''''''''''''''''''''' # We need to provide the sensing matrix, measurements # and the block size as parameters to the # reconstruction algorithm options = bsbl.bsbl_bo_options(y, learn_lambda=0) sol = bsbl.bsbl_bo_jit(Phi, y, b, options) print(sol) # %% # Reconstructed signal x_hat = sol.x print(f'BSBL-BO: PRD: {crn.prd(x, x_hat):.1f} %, SNR: {crn.signal_noise_ratio(x, x_hat):.2f} dB.' ) # %% # Plot the original and reconstructed signal ax = crplot.h_plots(2) ax[0].stem(x) ax[1].stem(x_hat) # %% # Observations: # # * We specified ``learn_lambda=0`` since we knew that this is a noiseless problem. # * Note the nonzero blocks count. They have been identified correctly. # * Recovery is perfect for both algorithms. In other words, both the # nonzero coefficient values and locations have been # correctly estimated and identified respectively. # * BSBL-BO is faster compared to BSBL-EM. # See how it finished in far less number of iterations. # This is on expected lines as BSBL-BO accelerates the convergence # using bound optimization a.k.a. majorization-minimization. # %% # Sparse Binary Sensing # ------------------------------ # %% # We shall have just 12 ones in each column of the sensing matrix d = 12 # %% # Build the sensing matrix Phi = crdict.sparse_binary_mtx(crn.KEYS[0], m, n, d, normalize_atoms=True, dense=True) ax = crplot.one_plot() ax.spy(Phi); # %% # Measurements y = Phi @ x ax = crplot.one_plot() crplot.plot_signal(ax, y) # %% # Reconstruction using BSBL EM # ''''''''''''''''''''''''''''''''''' options = bsbl.bsbl_em_options(y, learn_lambda=0) sol = bsbl.bsbl_em_jit(Phi, y, b, options) print(sol) # %% # Reconstructed signal x_hat = sol.x print(f'BSBL-EM: PRD: {crn.prd(x, x_hat):.1f} %, SNR: {crn.signal_noise_ratio(x, x_hat):.2f} dB.' ) # %% # Plot the original and reconstructed signal ax = crplot.h_plots(2) ax[0].stem(x) ax[1].stem(x_hat) # %% # Reconstruction using BSBL BO # ''''''''''''''''''''''''''''''''''' options = bsbl.bsbl_bo_options(y, learn_lambda=0) sol = bsbl.bsbl_bo_jit(Phi, y, b, options) print(sol) # %% # Reconstructed signal x_hat = sol.x print(f'BSBL-BO: PRD: {crn.prd(x, x_hat):.1f} %, SNR: {crn.signal_noise_ratio(x, x_hat):.2f} dB.' ) # %% # Plot the original and reconstructed signal ax = crplot.h_plots(2) ax[0].stem(x) ax[1].stem(x_hat) # %% # Observations: # # * Recovery is perfect for both algorithms. # * BSBL-BO is much faster. # * Both algorithms are converging in same number of iterations # as Gaussian sensing matrices. # %% # Noisy Measurements # ------------------------------ # We now consider an example where compressive measurements # are corrupted with significant amount of noise. # rows (measurements) m = 80 # columns (signal space) n = 162 # block length b = 6 # number of nonzero blocks k = 5 # number of signals s = 1 # number of blocks nb = n // b # Signal to Noise Ratio in DB snr = 15 # %% # generate block sparse signal with high intra block correlation x, blocks, indices = crdata.sparse_normal_blocks( crn.KEYS[1], n, k, b, s, cor=0.95, normalize_blocks=True) ax = crplot.one_plot() ax.stem(x) # %% # Sensing matrix Phi = crdict.gaussian_mtx(crn.KEYS[2], m, n) # %% # Noiseless measurements y0 = Phi @ x # %% # Noise at an SNR of 15 dB noise = crdsp.awgn_at_snr_std(crn.KEYS[3], y0, snr) # %% # Addition of noise to measurements y = y0 + noise print(f'measurement SNR: {crn.signal_noise_ratio(y0, y):.2f} dB') # %% # Plot the noiseless and noisy measurements ax = crplot.h_plots(2) ax[0].plot(y0) ax[1].plot(y) # %% # Reconstruction Benchmark # '''''''''''''''''''''''''''''''''' # An oracle reconstruction is possible if one knows the # nonzero indices of x. One can then compute a least # square solution over these indices. # The reconstruction SNR of this solution gives us # a good benchmark against which we can evaluate # the quality of reconstruction by any other algorithm. # Find the least square solution x_ls_coeffs = jnp.linalg.pinv(Phi[:, indices]) @ y x_ls = crdsp.build_signal_from_indices_and_values(n, indices, x_ls_coeffs) print(f'Benchmark rec: PRD: {crn.prd(x, x_ls):.1f} %, SNR: {crn.signal_noise_ratio(x, x_ls):.2f} dB') # %% # Plot the oracle reconstruction ax = crplot.h_plots(2) ax[0].stem(x) ax[1].stem(x_ls) # %% # Reconstruction using BSBL EM # ''''''''''''''''''''''''''''''''''' options = bsbl.bsbl_em_options(y) sol = bsbl.bsbl_em_jit(Phi, y, b, options) print(sol) # %% # Reconstructed signal x_hat = sol.x print(f'BSBL-EM: PRD: {crn.prd(x, x_hat):.1f} %, SNR: {crn.signal_noise_ratio(x, x_hat):.2f} dB.' ) # %% # Plot the original and reconstructed signal ax = crplot.h_plots(2) ax[0].stem(x) ax[1].stem(x_hat) # %% # Reconstruction using BSBL BO # ''''''''''''''''''''''''''''''''''' options = bsbl.bsbl_bo_options(y) sol = bsbl.bsbl_bo_jit(Phi, y, b, options) print(sol) # %% # Reconstructed signal x_hat = sol.x print(f'BSBL-BO: PRD: {crn.prd(x, x_hat):.1f} %, SNR: {crn.signal_noise_ratio(x, x_hat):.2f} dB.' ) # %% # Plot the original and reconstructed signal ax = crplot.h_plots(2) ax[0].stem(x) ax[1].stem(x_hat) # %% # Observations: # # * Benchmark SNR is slightly better than measurement SNR. # This is due to the fact that we are using our knowledge of # support for x. # * Both BSBL-EM and BSBL-BO are able to detect the correct support of x. # * SNR for both BSBL-EM and BSBL-BO is better than the benchmark SNR. # This is due to the fact that BSBL is exploiting the intra-block correlation # modeled as an AR-1 process. # * The ordinary least squares solution is not attempting to exploit the # intra block correlation structure at all. # * In this example BSBL-BO is somewhat faster than BSBL-EM but not by much.