# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: Python 3
#     language: python
#     name: python3
# ---

# %load_ext autoreload
# %autoreload 2

import jax
import jax.numpy as jnp
import numpy as np
from jax import random, jit, grad
import scipy

import cr.sparse as crs
from cr.sparse import la
from cr.sparse import dict
from cr.sparse import pursuit
from cr.sparse import data

from cr.sparse.pursuit import omp

M = 256
N = 1024
K = 16
S = 32

key = random.PRNGKey(0)
Phi = dict.gaussian_mtx(key, M,N)

dict.coherence(Phi)

X, omega = data.sparse_normal_representations(key, N, K, S)
X.shape

omega

Y = Phi @ X
Y.shape

solution = omp.solve_multi(Phi, Y, K)

jnp.max(solution.r_norm_sqr)


def time_solve_multi():
    solution = omp.solve_multi(Phi, Y, K)
    solution.x_I.block_until_ready()
    solution.r.block_until_ready()
    solution.I.block_until_ready()
    solution.r_norm_sqr.block_until_ready()


# %timeit time_solve_multi()

solve_multi = jax.jit(omp.solve_multi, static_argnums=(2,))


def time_solve_multi_jit():
    solution = solve_multi(Phi, Y, K)
    solution.x_I.block_until_ready()
    solution.r.block_until_ready()
    solution.I.block_until_ready()
    solution.r_norm_sqr.block_until_ready()

# %timeit time_solve_multi_jit()



StarcoderdataJupyterScriptsDedupFiltered