Custom sensors¶
In [1]:
Copied!
from jax import jit
from jax import numpy as jnp
from matplotlib import pyplot as plt
from jwave import FourierSeries
from jwave.acoustics import simulate_wave_propagation
from jwave.geometry import *
from jwave.geometry import circ_mask
from jwave.utils import show_field
domain = Domain((128, 128), (0.1e-3, 0.1e-3))
medium = Medium(domain=domain, sound_speed=1500.0)
time_axis = TimeAxis.from_medium(medium, cfl=0.3)
from jax import jit
from jax import numpy as jnp
from matplotlib import pyplot as plt
from jwave import FourierSeries
from jwave.acoustics import simulate_wave_propagation
from jwave.geometry import *
from jwave.geometry import circ_mask
from jwave.utils import show_field
domain = Domain((128, 128), (0.1e-3, 0.1e-3))
medium = Medium(domain=domain, sound_speed=1500.0)
time_axis = TimeAxis.from_medium(medium, cfl=0.3)
WARNING:absl:No GPU/TPU found, falling back to CPU. (Set TF_CPP_MIN_LOG_LEVEL=0 and rerun for more info.)
In [2]:
Copied!
# Defining the initial pressure
N = domain.N
mask1 = circ_mask(N, 8, (50, 50))
mask2 = circ_mask(N, 5, (80, 60))
mask3 = circ_mask(N, 10, (64, 64))
mask4 = circ_mask(N, 30, (64, 64))
p0 = 5.0 * mask1 + 3.0 * mask2 + 4.0 * mask3 + 0.5 * mask4
p0 = 1.0 * jnp.expand_dims(p0, -1)
p0 = FourierSeries(p0, domain)
# Defining the initial pressure
N = domain.N
mask1 = circ_mask(N, 8, (50, 50))
mask2 = circ_mask(N, 5, (80, 60))
mask3 = circ_mask(N, 10, (64, 64))
mask4 = circ_mask(N, 30, (64, 64))
p0 = 5.0 * mask1 + 3.0 * mask2 + 4.0 * mask3 + 0.5 * mask4
p0 = 1.0 * jnp.expand_dims(p0, -1)
p0 = FourierSeries(p0, domain)
In [3]:
Copied!
show_field(p0)
plt.title("Initial pressure")
show_field(p0)
plt.title("Initial pressure")
Out[3]:
Text(0.5, 1.0, 'Initial pressure')
In [4]:
Copied!
def sensor_func(p, u, rho):
position = (64, 24)
# Returns the maximum velocity on a 16x16 neighborhood at
# each timestep
v = u.on_grid[position[0] : position[0] + 16, position[1] : position[1] + 16]
return jnp.amax(v, axis=(0, 1))
def sensor_func(p, u, rho):
position = (64, 24)
# Returns the maximum velocity on a 16x16 neighborhood at
# each timestep
v = u.on_grid[position[0] : position[0] + 16, position[1] : position[1] + 16]
return jnp.amax(v, axis=(0, 1))
In [5]:
Copied!
@jit
def compiled_simulator(medium, p0):
return simulate_wave_propagation(medium, time_axis, p0=p0, sensors=sensor_func)
@jit
def compiled_simulator(medium, p0):
return simulate_wave_propagation(medium, time_axis, p0=p0, sensors=sensor_func)
In [6]:
Copied!
sensors_data = compiled_simulator(medium, p0)
sensors_data = compiled_simulator(medium, p0)
In [7]:
Copied!
sensors_data.shape
sensors_data.shape
Out[7]:
(599, 2)
In [8]:
Copied!
plt.figure(figsize=(8, 3), dpi=100)
plt.plot(sensors_data[:, 0], label="Max $v_x$", color="k")
plt.plot(sensors_data[:, 1], label="Max $v_y$", color="k", linestyle="-.")
plt.legend()
plt.figure(figsize=(8, 3), dpi=100)
plt.plot(sensors_data[:, 0], label="Max $v_x$", color="k")
plt.plot(sensors_data[:, 1], label="Max $v_y$", color="k", linestyle="-.")
plt.legend()
Out[8]:
<matplotlib.legend.Legend at 0x7f9a082df880>
