Compressed Sensing and Signal Recovery (OMP)

In this example, we are considering a classical signal processing problem: there is a long-length signal. n = 256 but it is known that it is sparse — only k = 20 its counts are non-zero, the rest are zero. Instead of measuring all 256 values, we only do m = 128 linear measurements of the form y = A·x, where A — a random Gaussian matrix of size m × n. The task is based on "compressed" measurements y restore the original signal x.

What is the essence of the technology. The classical theory (Kotelnikov's theorem) requires a number of measurements on the order of the signal length. However, if the signal is sparse, the theory of Compressed Sensing guarantees accurate recovery already at m ≈ k·log(n) that is , when m ≪ n. The key idea is to use a nonlinear sparse solver instead of a conventional MNC.

The OMP algorithm (Orthogonal Matching Pursuit) is one of the simplest greedy recovery algorithms. At each iteration, he:

finds the column of the matrix A, most correlated with the current balance r = y − A·x;
adds its index to the supported set;
Solves the least squares problem only on selected columns;
Updates the balance.

Through k iterations of the algorithm returns a sparse solution with k non-zero components. In our example, it restores the signal with a relative order error. 10⁻¹⁵ — that is, it is almost perfect, which is clearly demonstrated by the graphs of the original and restored signals.

Realization

The example uses three libraries:

LinearAlgebra — the standard library of linear algebra: matrix operations, norms, solving linear systems (used for norm, A', Asub \ y).
SparseArrays is a standard library for working with sparse vectors and matrices (used for sparsevec, findnz, efficient signal storage with a small number of non-zero elements).
Random — the standard library for generating random numbers (used for randn, randperm and Random.seed! for reproducibility).

Pkg.add(["LinearAlgebra", "SparseArrays", "Random"])
using LinearAlgebra, SparseArrays, Random

omp(A, y, k) — implementation of the Orthogonal Matching Pursuit (OMP) algorithm for greedy sparse vector reconstruction x from the equation y = A·x, where k — the expected number of non-zero elements.

Algorithm:

Initializes: r = y, support = [], x = zeros(n).
Repeats k once:
- calculates correlations A' * r and takes their module;
- resets already selected indexes;
- selects the index j with maximum correlation and adds it to support;
builds a submatrix Asub = A[:, support];
solves the MNC problem Asub \ y, gets the coefficients coef;
updates x: resets and writes coef in the position support;
recalculates the balance r = y - Asub * coef;
if norm(r) < 1e-12, interrupts the cycle.
Returns x.

function omp(A, y, k)
    m, n = size(A)
    r = copy(y)
    support = Int[]
    x = zeros(n)
    for _ in 1:k
        c = abs.(A' * r)
        c[support] .= 0
        j = argmax(c)
        push!(support, j)
        Asub = A[:, support]
        coef = Asub \ y
        x .= 0
        x[support] = coef
        r = y - Asub * coef
        if norm(r) < 1e-12
            break
        end
    end
    return x
end

omp (generic function with 1 method)

Preparing data for compressed sensing:

Setting the parameters: n = 256, m = 128, k = 20, where m < n provides compression.
Fixing the generator: Random.seed!(42) for reproducibility.
Creating a sparse signal: x_true = zeros(n), choose k random indexes via randperm(n)[1:k] and we write it down there randn(k).
We are building a measurement matrix A = randn(m, n) and we normalize its columns: A = A ./ sqrt.(sum(A.^2, dims=1)) — so that OMP does not give an advantage to columns with more energy.
Calculate compressed measurements: y = A * x_true.

n = 256                # signal length
m = 128                # number of measurements (m < n)
k = 20                 # sparsity (number of non-zero elements)

Random.seed!(42)
x_true = zeros(n)
support = randperm(n)[1:k]
x_true[support] = randn(k)
A = randn(m, n)
A = A ./ sqrt.(sum(A.^2, dims=1))
y = A * x_true;

Calculating the reconstructed signal: x_rec = omp(A, y, k), parameters:

A — measurement matrix;
y — measurement vector;
k — sparsity.

The function returns x_rec lengths n with no more than k non-zero components. Next, we compare x_rec with x_true.

x_rec = omp(A, y, k);

After that, we calculate the relative recovery error, output it, the number of non-zero elements in the original and restored signals, and check if their supports match.

err = norm(x_true - x_rec) / norm(x_true)
println("Relative recovery error: ", err)
println("The number of non-zero elements in the original signal: ", k)
println("The number of non-zero elements in the reconstructed: ",
        count(!iszero, x_rec))
println("Support has been restored correctly: ",
        Set(findall(!iszero, x_true)) == Set(findall(!iszero, x_rec)))

Относительная ошибка восстановления: 3.772586544601344e-16
Число ненулевых элементов в исходном сигнале: 20
Число ненулевых элементов в восстановленном: 20
Поддержка восстановлена верно: true

Next, we plot the original signal x_true and on top of it, the graph of the restored signal is dotted in red. x_rec to compare them.

plot(x_true, label="The original", ylabel="The amplitude",
          title="Signal comparison", linewidth=2)
plot!(x_rec, label="Restored",
      linewidth=2, linestyle=:dash, color=:red)

Conclusion

In this example, the problem of compressed sensing has been successfully solved: by incomplete measurements y = A·x (where m = 128 < n = 256) a sparse signal has been restored x with k = 20 non-zero components. The OMP algorithm showed high accuracy:

the relative recovery error was 3.77e-16 which is close to machine precision;
the number of non-zero elements in the original and reconstructed signals is the same (20);
support (positions of non-zero elements) has been restored correctly.

Thus, the example clearly confirms the main provisions of the theory of compressed sensing: with sufficient sparsity of the signal, it can be accurately reconstructed from a significantly smaller number of measurements, and the greedy OMP algorithm is an effective tool for solving such a problem.