ivp_1D_simulation.py

"""
Simulations In One Dimension Example

Ported from: k-Wave/examples/example_ivp_1D_simulation.m

Simulates the pressure field generated by an initial pressure distribution
(a smoothly shaped sinusoidal pulse) within a one-dimensional heterogeneous
propagation medium. It builds on the Homogeneous and Heterogeneous
Propagation Medium examples.

The domain has a region of higher sound speed on the left (2000 m/s vs
1500 m/s) and higher density on the right (1500 kg/m^3 vs 1000 kg/m^3).
You should observe the initial pulse splitting into left- and right-
travelling waves, with partial reflections and transmission at the
impedance boundaries.
"""
# %%
import numpy as np

from kwave.data import Vector
from kwave.kgrid import kWaveGrid
from kwave.kmedium import kWaveMedium
from kwave.ksensor import kSensor
from kwave.ksource import kSource
from kwave.kspaceFirstOrder import kspaceFirstOrder


# %%
def setup():
    """Set up the simulation physics (grid, medium, source).

    Returns:
        tuple: (kgrid, medium, source)
    """

    # create the computational grid
    Nx = 512  # number of grid points in the x direction
    dx = 0.05e-3  # grid point spacing in the x direction [m]
    kgrid = kWaveGrid(Vector([Nx]), Vector([dx]))

    # define the properties of the propagation medium
    c = 1500 * np.ones(Nx)  # [m/s]
    c[: round(Nx / 3)] = 2000  # MATLAB: 1:round(Nx/3)
    rho = 1000 * np.ones(Nx)  # [kg/m^3]
    rho[round(4 * Nx / 5) - 1 :] = 1500  # MATLAB: round(4*Nx/5):end (1-based -> 0-based)
    medium = kWaveMedium(sound_speed=c, density=rho)

    # create initial pressure distribution using a smoothly shaped sinusoid
    x_pos = 280  # starting grid point for the pulse [grid points]
    width = 100  # pulse width [grid points]
    height = 1.0  # pulse amplitude [au]
    in_arr = np.linspace(0, 2 * np.pi, width + 1)  # exactly 101 points
    p0 = np.concatenate(
        [
            np.zeros(x_pos),
            (height / 2) * np.sin(in_arr - np.pi / 2) + (height / 2),
            np.zeros(Nx - x_pos - (width + 1)),
        ]
    )
    source = kSource()
    source.p0 = p0

    # set the simulation time to capture reflections
    t_end = 2.5 * (Nx * dx) / np.max(c)  # [s]
    kgrid.makeTime(c, t_end=t_end)

    return kgrid, medium, source


# %%
def run(backend="python", device="cpu", quiet=True):
    """Run with the original Cartesian two-point sensor.

    The MATLAB original uses sensor.mask = [-10e-3, 10e-3], which places
    two Cartesian sensor points at x = -10 mm and x = +10 mm.

    Returns:
        dict: Simulation results with key 'p' (2 x time_steps).
    """
    kgrid, medium, source = setup()

    # create a Cartesian sensor mask (two points along x-axis)
    sensor = kSensor(mask=np.array([[-10e-3, 10e-3]]))  # [m]

    return kspaceFirstOrder(
        kgrid,
        medium,
        source,
        sensor,
        backend=backend,
        device=device,
        quiet=quiet,
        pml_inside=True,
    )


# %%
if __name__ == "__main__":
    import matplotlib.pyplot as plt

    result = run(quiet=False)
    p = np.asarray(result["p"])

    fig, ax = plt.subplots(figsize=(10, 5))

    ax.plot(p[0, :], "b-", label="Sensor Position 1 (x = -10 mm)")
    ax.plot(p[1, :], "r-", label="Sensor Position 2 (x = +10 mm)")
    ax.set_ylim(-0.1, 0.7)
    ax.set_xlabel("Time Step")
    ax.set_ylabel("Pressure")
    ax.legend()
    ax.set_title("Simulations In One Dimension")

    fig.tight_layout()
    plt.show()