Compressed Sensing, minimal example

cose.py

#!/usr/bin/env python
# -*-coding:utf-8-*-


"""
This module contains a minimal example to illustrate one of the main
implications of the compressed sensing paradigm: That a perfect recovery of a
sparse signal is possible with a number of measurements that depends on the
signal sparsity, and not on the signal length.

More sources:

* Blog entry by Terry Tao: https://terrytao.wordpress.com/2007/04/13/compressed-sensing-and-single-pixel-cameras/
* Original CoSe paper: https://arxiv.org/abs/math/0409186
* Follow-up paper, CoSe under noise: https://arxiv.org/abs/math/0503066


MIT License below.

--------------------------------------------------------------------------------
Copyright 2022 aferro (OrcID: 0000-0003-3830-3595)

Permission is hereby granted, free of charge, to any person obtaining a copy of
this software and associated documentation files (the "Software"), to deal in
the Software without restriction, including without limitation the rights to
use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies
of the Software, and to permit persons to whom the Software is furnished to do
so, subject to the following conditions:

The above copyright notice and this permission notice shall be included in all
copies or substantial portions of the Software.

THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
SOFTWARE.
"""


import numpy as np
from scipy.optimize import linprog


# ##############################################################################
# # HELPER FUNCTIONS
# ##############################################################################
def get_dft_coefficients(ndims, *indexes, as_sincos=True):
    """
    :param int ndims: If we were to take the DFT of an n-dimensional signal,
      this would be ``n``.
    :param indexes: Collection of integer indexes between 0 and ``n-1``.
    :param bool as_sincos: If true, the returned matrix will have twice as
      many columns; first half of columns containing real components and
      second half imaginary ones. Otherwise, complex matrix returned.
    :returns: If ``as_sincos`` is true, a real-valued matrix of shape
      ``(ndims, 2*len(indexes))``, otherwise, a complex matrix of shape
      ``(ndims, len(indexes))``, where each column is the ``ith`` element
      from the DFT basis.
    """
    result = np.fft.fft(np.eye(ndims))[:, indexes]
    if as_sincos:
        result = np.hstack((result.real, result.imag))
    return result


def cose_recovery(original_ndim, measurements, sensing_matrix):
    """
    The compressed sensing recovery consists in solving a linear program:
    find ``x`` that minimizes ``x @ c``, for ``c`` being all ones, subject
    to ``sensing_matrix @ x = measurements``.

    In other words, we want to find the ``x`` with minimal L1 norm that
    satisfies the measurements provided by our sensing matrix. If the sparsity
    assumptions are met, and the sensing and sparse basis are uncorrelated,
    the solution to this program is very likely to be exact.

    :param int original_ndim: Dimensionality of the returned solution.
    :param measurements: Real-valued vector of size ``m`` containing the
      CoSe measurements
    :param sensing_matrix: Real-valued matrix of shape ``(original_ndim, m)``
      used to obtain measurements via ``mat.T @ original = measurements``.
    :returns: The estimated ``original`` signal of expected length of
      ``original_ndim``.
    """
    # technicality: since the LP solver requires a non-negative solution,
    # we make 2 "flipped" copies and then we ensemble them together
    coeffs = np.ones(2 * original_ndim)
    eq_matrix = np.vstack((sensing_matrix, -sensing_matrix))
    # solve linear program and assert it finds solution
    lp = linprog(coeffs, A_eq=eq_matrix.T, b_eq=measurements,
                 bounds=(0, None), method="highs")
    assert lp.success, "Ill-posed LP?"
    # reassemble positive and negative solutions and return
    solution = lp.x[:original_ndim] - lp.x[original_ndim:]
    return solution


# ##############################################################################
# # MAIN ROUTINE
# ##############################################################################
if __name__ == "__main__":

    # define globals, you can play around with these parameters
    dim1, x1, y1 = 51,        3,  372.32839
    dim2, x2, y2 = 10_000, 8173, -537.72937
    measurement_idxs_1 = (2, 7, 8, 10,
                          11, 18, 19, 20)
    measurement_idxs_2 = (69, 1233, 1782, 2022,
                          4205, 6794, 9902, 9997)

    # create two signals of different lengths but with same sparsity level
    original1 = np.zeros(dim1)
    original1[x1] = y1
    original2 = np.zeros(dim2)
    original2[x2] = y2

    # gather random FFT CoSe matrices for s1 and s2. These matrices are "tall":
    # the height corresponds to the length of the original signals, which
    # varies, but the width to the (much smaller) number of measurements
    cosemat1 = get_dft_coefficients(dim1, *measurement_idxs_1)
    cosemat2 = get_dft_coefficients(dim2, *measurement_idxs_2)

    # perform "measurements", which is equivalent to computing dot products
    # between original signal and each basis from the cose matrices.
    # note that both measurements have same length, despite the length
    # differences between original1 and original2
    measurements1 = cosemat1.T @ original1
    measurements2 = cosemat2.T @ original2

    # this is the interesting step. Given the (short) measurements, our sensing
    # matrices and the original length of the signal (all things we readily
    # know), we recover the original signal exactly! i.e. all zeros except
    # ori[x] = y, with the exact location and value of the spike.
    # This is true regardless of the length of the original signal!
    recovery1 = cose_recovery(dim1, measurements1, cosemat1)
    recovery2 = cose_recovery(dim2, measurements2, cosemat2)

    # Finally, check that recovered signals are equal to originals
    assert np.allclose(original1, recovery1), "This should not happen!"
    assert np.allclose(original2, recovery2), "This should not happen!"
    print("Recovered signals were identical!")