5.3. First Derivative Operator¶
import numpy as np
import jax.numpy as jnp
import pylops
from cr.sparse import lop
import cr.sparse as crs
n = 1000*1000
x_np = np.random.normal(0, 1, (n))
x_jax = jnp.array(x_np)
op_np = pylops.FirstDerivative(n, kind='forward')
y_np = op_np * x_np
op_jax = lop.first_derivative(n, kind='forward')
op_jax = lop.jit(op_jax)
y_jax = op_jax.times(x_jax)
np.allclose(y_np, y_jax, atol=1e-4)
True
np_time = %timeit -o op_np * x_np
2.8 ms ± 2.44 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)
jax_time = %timeit -o op_jax.times(x_jax).block_until_ready()
37.3 µs ± 713 ns per loop (mean ± std. dev. of 7 runs, 10000 loops each)
gain = np_time.average / jax_time.average
print(gain)
74.98343664823614
y1_np = op_np.H * x_np
y1_jax = op_jax.trans(x_jax).block_until_ready()
np.allclose(y1_np, y1_jax, atol=1e-4)
True
np_time = %timeit -o op_np.H * x_np
4.45 ms ± 9.09 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)
jax_time = %timeit -o op_jax.trans(x_jax).block_until_ready()
80.6 µs ± 554 ns per loop (mean ± std. dev. of 7 runs, 10000 loops each)
gain = np_time.average / jax_time.average
print(gain)
55.21886362958511