A sparse model vector generator

Demonstrates how to create sparse model vectors with small number of non-zero entries sampled from Gaussian distribution

Let’s import necessary libraries

import matplotlib as mpl
import matplotlib.pyplot as plt
from jax import random
import jax.numpy as jnp
import cr.sparse as crs
import cr.sparse.data as crdata
from cr.nimble.dsp import (
    nonzero_indices,
    nonzero_values
)

Let’s define the size of model and number of sparse entries

# Model size
N = 1000
# Number of non-zero entries in the sparse model
K = 30

Let’s generate a random sparse model

key = random.PRNGKey(1)
x, omega = crdata.sparse_normal_representations(key, N, K, 1)
x = jnp.squeeze(x)

We can easily find the locations of non-zero entries

print(nonzero_indices(x))

[ 45  69 114 133 159 185 199 233 259 338 348 377 393 425 490 504 521 533
 537 627 722 730 782 790 817 846 891 947 949 968]

We can extract corresponding non-zero values in a compact vector

print(nonzero_values(x))

[-0.64617896 -1.2449833  -0.9160154   2.7065587   0.99159986  1.7746289
  0.8708357  -0.28911775  0.30582124 -1.0570332  -0.03854743 -2.4557247
 -0.04477294 -2.2054574  -1.835958    1.8202777  -0.3686923  -0.54682994
 -0.82268345  1.5236915   0.24205726 -0.01583949  1.150074   -0.00814029
 -0.6977474  -0.88545537  0.4539182   1.7888712  -0.24443631 -0.39677632]

Let’s plot the vector to see where the non-zero entries are

plt.figure(figsize=(8,6), dpi= 100, facecolor='w', edgecolor='k')
plt.stem(x, markerfmt='.');

<StemContainer object of 3 artists>