A matrix backed linear operator

This example demonstrates how to construct a linear operator from an ordinary matrix.

Let’s import necessary libraries

import jax.numpy as jnp
from cr.sparse import lop

Setup

Create a small matrix

n = 4
A = jnp.ones((n, n))
print(A)

[[1. 1. 1. 1.]
 [1. 1. 1. 1.]
 [1. 1. 1. 1.]
 [1. 1. 1. 1.]]

Create a vector

x = jnp.ones(n)
print(x)

[1. 1. 1. 1.]

Operator construction

Convert A into an operator

T = lop.matrix(A)

Operator usage

Compute A x

print(T.times(x))

[4. 4. 4. 4.]

Compute A^H x

print(T.trans(x))

[4. 4. 4. 4.]

JIT Compilation

Wrap the times and trans functions with jit

T = lop.jit(T)

Compute A x

print(T.times(x))

[4. 4. 4. 4.]

Compute A^H x

print(T.trans(x))

[4. 4. 4. 4.]