Heterogeneous medium¶
In [1]:
Copied!
import numpy as np
from jax import jit
from jax import numpy as jnp
from jaxdf import FourierSeries
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, show_positive_field
domain = Domain((128, 128), (0.1e-3, 0.1e-3))
import numpy as np
from jax import jit
from jax import numpy as jnp
from jaxdf import FourierSeries
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, show_positive_field
domain = Domain((128, 128), (0.1e-3, 0.1e-3))
In [2]:
Copied!
density = np.ones(domain.N) * 1000
density[50:90, 65:100] = 2300
density = FourierSeries(np.expand_dims(density, -1), domain)
show_positive_field(density)
_ = plt.title("Density")
density = np.ones(domain.N) * 1000
density[50:90, 65:100] = 2300
density = FourierSeries(np.expand_dims(density, -1), domain)
show_positive_field(density)
_ = plt.title("Density")
WARNING:absl:No GPU/TPU found, falling back to CPU. (Set TF_CPP_MIN_LOG_LEVEL=0 and rerun for more info.)
In [3]:
Copied!
medium = Medium(domain=domain, sound_speed=1500.0, density=density)
time_axis = TimeAxis.from_medium(medium, cfl=0.3)
medium = Medium(domain=domain, sound_speed=1500.0, density=density)
time_axis = TimeAxis.from_medium(medium, cfl=0.3)
In [4]:
Copied!
# Defining the initial pressure
N = domain.N
p0 = circ_mask(N, 3, (64, 30))
p0 = 1.0 * jnp.expand_dims(p0, -1)
p0 = FourierSeries(p0, domain)
# Defining the initial pressure
N = domain.N
p0 = circ_mask(N, 3, (64, 30))
p0 = 1.0 * jnp.expand_dims(p0, -1)
p0 = FourierSeries(p0, domain)
In [5]:
Copied!
show_field(p0)
plt.title("Initial pressure")
show_field(p0)
plt.title("Initial pressure")
Out[5]:
Text(0.5, 1.0, 'Initial pressure')
In [6]:
Copied!
@jit
def compiled_simulator(c, p0):
medium = Medium(domain=domain, sound_speed=1500.0, density=c)
return simulate_wave_propagation(medium, time_axis, p0=p0)
@jit
def compiled_simulator(c, p0):
medium = Medium(domain=domain, sound_speed=1500.0, density=c)
return simulate_wave_propagation(medium, time_axis, p0=p0)
In [7]:
Copied!
pressure = compiled_simulator(density, p0)
pressure = compiled_simulator(density, p0)
In [8]:
Copied!
t = 270
show_field(pressure[t])
plt.title(f"Pressure field at t={time_axis.to_array()[t]}")
t = 270
show_field(pressure[t])
plt.title(f"Pressure field at t={time_axis.to_array()[t]}")
Out[8]:
Text(0.5, 1.0, 'Pressure field at t=5.4000001910026185e-06')
In [9]:
Copied!
pressure = compiled_simulator(1000.0, p0) # Works with different input types
pressure = compiled_simulator(1000.0, p0) # Works with different input types
In [10]:
Copied!
t = 270
show_field(pressure[t])
plt.title(f"Pressure field at t={time_axis.to_array()[t]}")
t = 270
show_field(pressure[t])
plt.title(f"Pressure field at t={time_axis.to_array()[t]}")
Out[10]:
Text(0.5, 1.0, 'Pressure field at t=5.4000001910026185e-06')
In [11]:
Copied!
from jaxdf import FiniteDifferences
density = np.ones(domain.N) * 1000
density[50:90, 65:100] = 2300
density = FiniteDifferences(np.expand_dims(density, -1), domain, accuracy=8)
N = domain.N
p0 = circ_mask(N, 3, (64, 30))
p0 = 1.0 * jnp.expand_dims(p0, -1)
p0 = FiniteDifferences(p0, domain, accuracy=8)
from jaxdf import FiniteDifferences
density = np.ones(domain.N) * 1000
density[50:90, 65:100] = 2300
density = FiniteDifferences(np.expand_dims(density, -1), domain, accuracy=8)
N = domain.N
p0 = circ_mask(N, 3, (64, 30))
p0 = 1.0 * jnp.expand_dims(p0, -1)
p0 = FiniteDifferences(p0, domain, accuracy=8)
In [12]:
Copied!
pressure = compiled_simulator(density, p0)
pressure = compiled_simulator(density, p0)
In [13]:
Copied!
t = 270
show_field(pressure[t])
plt.title(f"Pressure field at t={time_axis.to_array()[t]}")
t = 270
show_field(pressure[t])
plt.title(f"Pressure field at t={time_axis.to_array()[t]}")
Out[13]:
Text(0.5, 1.0, 'Pressure field at t=5.4000001910026185e-06')
