{
  "cells": [
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "%matplotlib inline"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "\n\n# Sparse Binary Sensing Matrices\n    :depth: 2\n    :local:\n\n\nA (random) sparse binary sensing matrix has a very simple design.\nAssume that the signal space is $\\RR^N$ and the measurement\nspace is $\\RR^M$.\nEvery column of a sparse binary sensing matrix has a 1 in\nexactly $d$ positions and 0s elsewhere. The indices\nat which ones are present are randomly selected for each column.\n\nFollowing is an example sparse binary matrix with 3 ones in\neach column\n\n\\begin{align}\\begin{bmatrix}\n    1 & 0 & 0 & 0 & 0 & 0 & 1 & 1 & 0 & 0 & 1 & 0 & 0 & 0 & 0 & 1\\\\\n    0 & 0 & 0 & 0 & 0 & 1 & 0 & 0 & 0 & 1 & 1 & 1 & 0 & 1 & 0 & 0\\\\\n    0 & 0 & 0 & 1 & 1 & 1 & 0 & 0 & 1 & 0 & 0 & 0 & 0 & 1 & 0 & 0\\\\\n    0 & 1 & 0 & 0 & 0 & 0 & 0 & 0 & 0 & 1 & 0 & 0 & 0 & 0 & 0 & 0\\\\\n    0 & 0 & 0 & 1 & 0 & 0 & 0 & 1 & 0 & 0 & 0 & 0 & 1 & 0 & 0 & 0\\\\\n    0 & 1 & 0 & 0 & 0 & 0 & 0 & 0 & 1 & 0 & 1 & 0 & 1 & 0 & 1 & 1\\\\\n    1 & 0 & 0 & 1 & 1 & 0 & 0 & 0 & 0 & 0 & 0 & 1 & 0 & 0 & 1 & 0\\\\\n    1 & 0 & 1 & 0 & 0 & 1 & 1 & 1 & 0 & 0 & 0 & 1 & 0 & 0 & 1 & 0\\\\\n    0 & 1 & 1 & 0 & 1 & 0 & 0 & 0 & 1 & 1 & 0 & 0 & 0 & 1 & 0 & 1\\\\\n    0 & 0 & 1 & 0 & 0 & 0 & 1 & 0 & 0 & 0 & 0 & 0 & 1 & 0 & 0 & 0\\\\\n    \\end{bmatrix}\\end{align}\n\nFrom the perspective of algorithm design, we often require that\nthe sensing matrix have unit form columns. This can be easily\nattained for sparse binary matrices by scaling them with\n$\\frac{1}{\\sqrt{d}}$.\n\n\nJAX provides an efficient way of storing sparse matrices in BCOO\nformat. By default we employ this storage format for the (random) sparse\nbinary matrices.\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Necessary imports\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "import math\nimport cr.nimble as crn\nimport cr.sparse as crs\nimport cr.sparse.dict as crdict\nimport cr.sparse.data as crdata\nimport cr.sparse.lop as crlop\nimport cr.sparse.plots as crplots\nimport numpy as np\nimport jax\nimport jax.numpy as jnp\nfrom jax import random"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Some random number generation keys\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "key = random.PRNGKey(3)\nkeys = random.split(key, 5)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Creating Sparse Binary Sensing Matrices\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "# Matrix parameters\nm = 10\nn = 16\nd = 4\n\n# construct a sparse binary sensing matrix\nA = crdict.sparse_binary_mtx(keys[0], m, n, d, normalize_atoms=False)\n# It is stored in a compressed BCOO format\nprint(A)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "If we wish to see its contents\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "Ad = A.todense()\nprint(Ad)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "We can quickly check that all columns have d ones\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "print(jnp.sum(Ad, 0))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "By convention, we generate normalized sensing matrices by default.\nHowever, in the case of sparse binary matrices, it is more efficient\nto work with the unnormalized sensing matrix\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "A = crdict.sparse_binary_mtx(keys[0], m, n, d)\nprint(A.todense())"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Sparse Binary Sensing Linear Operators\nIt is often advantageous to work with matrices wrapped\nin our linear operator design.\nLet us construct the sensing matrix as a linear operator\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "T = crlop.sparse_binary_dict(keys[0], m, n, d, normalize_atoms=False)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "We can extract the contents by multiplying with an identity matrix\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "print(T.times(jnp.eye(n, dtype=int)))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "We can keep the normalization of the sensing matrix has a separate\nscaling operator\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "# Let us construct a scaling operator\nd_scale = 1/ math.sqrt(d)\nT_scale = crlop.scalar_mult(d_scale, m)\n# Let us combine the scaling operator with the unnormalized sensing operator\nT_normed = crlop.compose(T_scale, T)\n# Verify the normalized operator\nprint(T_normed.times(jnp.eye(n, dtype=int)))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Compressive Sensing\nWe shall use a larger problem to demonstrate the sensing\ncapabilities of the sparse binary sensing operator.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "# Signal space dimension\nn = 1024\n# Measurement space dimension\nm = 256\n# Number of ones in each column\nd = 16\n# Sparsity level of the signal to be sensed\nk = 40"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Let us construct the unnormalized as well as normalized sensing operators.\nWe shall use the unnormalized one during sensing but the normalized\none during reconstruction.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "d_scale = 1/ math.sqrt(d)\nT = crlop.sparse_binary_dict(keys[1], m, n, d, normalize_atoms=False)\nT_scale = crlop.scalar_mult(d_scale, m)\nT_normed = crlop.compose(T_scale, T)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "We can quickly visualize the sparsity pattern of this sensing matrix\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "A = T.times(jnp.eye(n, dtype=int))\nax = crplots.one_plot()\nax.spy(A);"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Let us construct a sparse signal\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "x, omega = crdata.sparse_normal_representations(keys[2], n, k)\nax = crplots.one_plot()\ncrplots.plot_sparse_signal(ax, x)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "y = T.times(x)\nax = crplots.one_plot()\ncrplots.plot_sparse_signal(ax, y)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Sparse Recovery\nWe shall use various algorithms for reconstructing the original signal\n\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "### CoSaMP\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "# Import the algorithm\nfrom cr.sparse.pursuit import cosamp\n# Solve the problem\nsol =  cosamp.operator_solve_jit(T_normed, y, k)\nprint(sol)\n# We need to scale the solution since the measurements were unscaled\nx_hat = sol.x * d_scale\n# Compute the SNR and PRD\nprint(f'SNR: {crn.signal_noise_ratio(x, x_hat):.2f} dB, PRD: {crn.percent_rms_diff(x, x_hat):.0f} %')\n# Plot the original and the reconstructed signal\nax = crplots.h_plots(2)\ncrplots.plot_sparse_signals(ax, x, x_hat)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "### Subspace Pursuit\nImport the algorithm\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "from cr.sparse.pursuit import sp\n# Solve the problem\nsol =  sp.operator_solve_jit(T_normed, y, k)\nprint(sol)\n# We need to scale the solution since the measurements were unscaled\nx_hat = sol.x * d_scale\n# Compute the SNR and PRD\nprint(f'SNR: {crn.signal_noise_ratio(x, x_hat):.2f} dB, PRD: {crn.percent_rms_diff(x, x_hat):.0f} %')\n# Plot the original and the reconstructed signal\nax = crplots.h_plots(2)\ncrplots.plot_sparse_signals(ax, x, x_hat)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "### Hard Thresholding Pursuit\nImport the algorithm\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "from cr.sparse.pursuit import htp\n# Solve the problem\nsol =  htp.operator_solve_jit(T_normed, y, k)\nprint(sol)\n# We need to scale the solution since the measurements were unscaled\nx_hat = sol.x * d_scale\n# Compute the SNR and PRD\nprint(f'SNR: {crn.signal_noise_ratio(x, x_hat):.2f} dB, PRD: {crn.percent_rms_diff(x, x_hat):.0f} %')\n# Plot the original and the reconstructed signal\nax = crplots.h_plots(2)\ncrplots.plot_sparse_signals(ax, x, x_hat)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "### Truncated Newton Interior Points Method\nImport the algorithm\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "from cr.sparse.cvx import l1ls\n# Solve the problem\n# Note that this algorithm doesn't require sparsity level k as input\nsol =  l1ls.solve_jit(T_normed, y, 1e-2)\nprint(sol)\n# We need to scale the solution since the measurements were unscaled\nx_hat = sol.x * d_scale\n# Compute the SNR and PRD\nprint(f'SNR: {crn.signal_noise_ratio(x, x_hat):.2f} dB, PRD: {crn.percent_rms_diff(x, x_hat):.0f} %')\n# Plot the original and the reconstructed signal\nax = crplots.h_plots(2)\ncrplots.plot_sparse_signals(ax, x, x_hat)"
      ]
    }
  ],
  "metadata": {
    "kernelspec": {
      "display_name": "Python 3",
      "language": "python",
      "name": "python3"
    },
    "language_info": {
      "codemirror_mode": {
        "name": "ipython",
        "version": 3
      },
      "file_extension": ".py",
      "mimetype": "text/x-python",
      "name": "python",
      "nbconvert_exporter": "python",
      "pygments_lexer": "ipython3",
      "version": "3.8.13"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}