Linear uncertainty propagation¶

import numpy as np
from jax import numpy as jnp
from jax import random, nn, value_and_grad, jit, grad, vmap
from jax import lax
from jax.example_libraries import optimizers
from functools import partial
from tqdm import tqdm
from matplotlib import pyplot as plt
from luprox import linear_uncertainty, mc_uncertainty

from jwave.geometry import Domain, Medium, TimeAxis, Sources, _circ_mask, _points_on_circle, Sensors
from jwave.signal_processing import gaussian_window, apply_ramp, smooth, analytic_signal
from jwave.acoustics import ongrid_wave_propagation

import cv2

# load image as gray scale and convert to numpy
def read_img(path):
    img = cv2.imread(path, cv2.IMREAD_GRAYSCALE)

    # Rescale image to 192x192
    img = cv2.resize(img, (192, 192))

    # Convert to numpy array
    img = np.asarray(img)
    img = img/255.
    img[img<0.4] = 0.

    # pad image to 256 symmetric
    img = np.pad(img, ((32, 32), (32, 32)), 'constant', constant_values=(0, 0))

    return img

img = read_img("./ct.jpg")

# Settings
N = (256, 256)
dx = (.5, .5)
cfl = 0.3
num_sources = 16
source_freq = .5
source_mag = 5
random_seed = random.PRNGKey(42)

# Define domain
domain = Domain(N, dx)

# Define medium
sound_speed = img + 1

medium = Medium(domain=domain, sound_speed=sound_speed, density=jnp.ones(N),
    attenuation=0.0, pml_size=20.)

# Time axis
time_axis = TimeAxis.from_medium(medium, cfl=cfl, t_end=120.)

# Sources
source_mag = source_mag/time_axis.dt
t = time_axis.to_array()
s1 = source_mag * jnp.sin(2 * jnp.pi * source_freq * t)
signal = gaussian_window(apply_ramp(s1, time_axis.dt, source_freq), t, 5, 2)
x, y = _points_on_circle(num_sources,110,(128,128))
source_positions = (jnp.array(x), jnp.array(y))

src_signal = jnp.stack([signal]*num_sources)
sources = Sources(source_positions, src_signal)

from jwave.acoustics import ongrid_simple_wave_propagation
# Run simulations
params, solver = ongrid_simple_wave_propagation(
    medium=medium,
    time_array=time_axis,
    sources=sources,
    output_t_axis = time_axis,
    backprop=True,
    checkpoint=True
)

def get_field(bone_sos):
    sound_speed = img*bone_sos + 1
    
    new_params = params.copy()
    new_params["acoustic_params"]["speed_of_sound"] = jnp.expand_dims(sound_speed,-1)
    
    p = solver(new_params)
    p_max = jnp.mean(jnp.abs(p)**2, axis=0)
    return jnp.sqrt(p_max + 1e-6)

p = get_field(.7)

get_field_lup = linear_uncertainty(get_field)
x = jnp.array([1.])
covariance = jnp.array([[(0.05)**2]])
mu_linear, cov_linear = get_field_lup(x, covariance)

from jax import random

get_field_lup = mc_uncertainty(get_field, 20)
mu_mc, cov_mc = get_field_lup(x, jnp.sqrt(covariance), random.PRNGKey(0))