Rays and autodiff

2D ray bending

Trace 2D rays through a smooth speed field and plot them over c(x). Requires beamax[viz-mpl].

r"""Trace 2D rays through a smooth field with $G(\mathbf{x},\mathbf{p})=c(\mathbf{x})|\mathbf{p}|$.

Example extras: viz-mpl
Example smoke: false
"""

import jax
import jax.numpy as jnp
import matplotlib.pyplot as plt
import numpy as np
from pathlib import Path

from beamax.gb import gb_solvers
from beamax.plotter import use_beamax_style


jax.config.update("jax_enable_x64", True)


def speed_field(x: jnp.ndarray) -> jnp.ndarray:
    lens_center = jnp.array([0.46, 0.52])
    lens = jnp.exp(-35.0 * jnp.sum((x - lens_center) ** 2, axis=-1))
    vertical_gradient = 0.18 * (x[..., 1] - 0.5)
    return 1.05 + vertical_gradient - 0.28 * lens


def solve_rays() -> tuple[jnp.ndarray, jnp.ndarray, jnp.ndarray, jnp.ndarray]:
    n_rays = 60
    ts = jnp.linspace(0.0, 0.75, 48)
    y0 = jnp.linspace(0.18, 0.82, n_rays)
    x0 = jnp.stack([jnp.full((n_rays,), 0.08), y0], axis=-1)
    p0 = jnp.tile(jnp.array([1.0, 0.0]), (n_rays, 1))
    m0 = 1j * jnp.eye(2)[None, :, :].repeat(n_rays, axis=0)
    a0 = jnp.ones((n_rays,))
    mode = jnp.ones((n_rays,))

    solver_config = gb_solvers.SolverConfig(
        rtol=1e-4,
        atol=1e-6,
        max_steps=1024,
    )
    xt, pt, _, _ = gb_solvers.solve_ODE_base(
        x0,
        p0,
        m0,
        a0,
        mode,
        ts,
        speed_field,
        0.0,
        solver_config,
    )
    return xt, pt, x0, ts


def main() -> None:
    plot_dir = Path("plots/rays")
    plot_dir.mkdir(parents=True, exist_ok=True)
    use_beamax_style()

    xt, pt, x0, ts = solve_rays()
    final = xt[:, -1, :]
    direction_angles = jnp.arctan2(pt[:, -1, 1], pt[:, -1, 0])
    initial_angles = jnp.arctan2(pt[:, 0, 1], pt[:, 0, 0])

    grid_n = 128
    x = jnp.linspace(0.0, 1.0, grid_n)
    y = jnp.linspace(0.0, 1.0, grid_n)
    xy = jnp.stack(jnp.meshgrid(x, y, indexing="ij"), axis=-1)
    c_values = speed_field(xy)

    fig, ax = plt.subplots(figsize=(6.5, 5.2))
    im = ax.imshow(
        np.asarray(c_values.T),
        extent=[0.0, 1.0, 0.0, 1.0],
        origin="lower",
        cmap="viridis",
        aspect="equal",
    )
    fig.colorbar(im, ax=ax, label="c(x)")

    ray_color = "#d94801"
    for ray in np.asarray(xt):
        ax.plot(ray[:, 0], ray[:, 1], color=ray_color, lw=1.6, alpha=0.88)
    ax.scatter(
        np.asarray(x0[:, 0]),
        np.asarray(x0[:, 1]),
        s=26,
        color="white",
        edgecolor="black",
        zorder=3,
    )
    ax.scatter(
        np.asarray(final[:, 0]),
        np.asarray(final[:, 1]),
        s=24,
        color=ray_color,
        edgecolor="black",
        zorder=3,
    )
    ax.set(xlim=(0.0, 1.0), ylim=(0.0, 1.0), xlabel="x", ylabel="y")
    ax.set_title("2D ray bending in a smooth speed field")
    fig.tight_layout()

    out_path = plot_dir / "2d_ray_bending.png"
    fig.savefig(out_path, dpi=180, bbox_inches="tight")
    plt.close(fig)

    mean_lateral_shift = float(jnp.mean(jnp.abs(final[:, 1] - x0[:, 1])))
    max_angle_change = float(jnp.max(jnp.abs(direction_angles - initial_angles)))
    print(f"Rays solved: {xt.shape[0]}, time samples: {ts.shape[0]}")
    print(
        f"Speed range on plot grid: [{float(c_values.min()):.3f}, {float(c_values.max()):.3f}]"
    )
    print(f"Mean lateral displacement: {mean_lateral_shift:.3f}")
    print(f"Max direction change: {max_angle_change:.3f} rad")
    print(f"Saved ray-bending plot to {out_path}")


if __name__ == "__main__":
    main()

2D rays autodiff

Optimize a neural c(x) field through the ray ODE and plot the rays, loss, and \(\Delta c\). Requires beamax[viz-mpl,autodiff].

"""Optimise a neural sound-speed field through the Gaussian beam ray ODE.

Example extras: viz-mpl,autodiff
Example smoke: false
"""

import jax
import jax.numpy as jnp
import matplotlib.pyplot as plt
import numpy as np
import optax
from pathlib import Path

from beamax.gb import gb_solvers
from beamax.plotter import use_beamax_style


jax.config.update("jax_enable_x64", True)

XMIN, XMAX = -10.0, 10.0
YMIN, YMAX = -10.0, 10.0
EXTENT = [XMIN, XMAX, YMIN, YMAX]


class NeuralSoundSpeed:
    def __init__(self, hidden_dim: int = 32, base_c: float = 1.0):
        self.base_c = base_c
        key = jax.random.PRNGKey(42)
        k1, k2, k3 = jax.random.split(key, 3)
        self.params = {
            "w1": 0.1 * jax.random.normal(k1, (2, hidden_dim)),
            "b1": jnp.zeros(hidden_dim),
            "w2": 0.1 * jax.random.normal(k2, (hidden_dim, hidden_dim)),
            "b2": jnp.zeros(hidden_dim),
            "w3": 0.1 * jax.random.normal(k3, (hidden_dim, 1)),
            "b3": jnp.zeros(1),
        }

    def __call__(self, x: jnp.ndarray, params: dict[str, jnp.ndarray]) -> jnp.ndarray:
        x_norm = jnp.stack(
            [
                2.0 * (x[..., 0] - XMIN) / (XMAX - XMIN) - 1.0,
                2.0 * (x[..., 1] - YMIN) / (YMAX - YMIN) - 1.0,
            ],
            axis=-1,
        )
        h = jnp.tanh(x_norm @ params["w1"] + params["b1"])
        h = jnp.tanh(h @ params["w2"] + params["b2"])
        delta_c = 0.3 * jnp.tanh(h @ params["w3"] + params["b3"])[..., 0]
        return self.base_c + delta_c


def ray_setup():
    n_rays = 20
    source_x = jnp.linspace(-5.0, 5.0, n_rays)
    source_y = -7.0
    x0 = jnp.stack([source_x, jnp.full(n_rays, source_y)], axis=-1)
    p0 = jnp.stack([jnp.zeros(n_rays), jnp.ones(n_rays)], axis=-1)
    focus = jnp.array([0.0, 5.0])

    d = 2
    alpha0 = jnp.ones((n_rays, d))
    m0 = 1j * jnp.einsum("bd,dj->bdj", alpha0, jnp.eye(d))
    a0 = jnp.ones((n_rays, 1))
    mode = jnp.ones((n_rays, 1))
    ts = jnp.linspace(0.0, 15.0, 300)
    return x0, p0, m0, a0, mode, ts, focus


def solve_rays(
    params: dict[str, jnp.ndarray],
    sound_speed: NeuralSoundSpeed,
    x0: jnp.ndarray,
    p0: jnp.ndarray,
    m0: jnp.ndarray,
    a0: jnp.ndarray,
    mode: jnp.ndarray,
    ts: jnp.ndarray,
) -> jnp.ndarray:
    def c_fn(x):
        return sound_speed(x, params)

    xt, _, _, _ = gb_solvers.solve_ODE_base(x0, p0, m0, a0, mode, ts, c_fn, 0.0, None)
    return xt


def focusing_loss(
    params: dict[str, jnp.ndarray],
    sound_speed: NeuralSoundSpeed,
    x0: jnp.ndarray,
    p0: jnp.ndarray,
    m0: jnp.ndarray,
    a0: jnp.ndarray,
    mode: jnp.ndarray,
    ts: jnp.ndarray,
    focus: jnp.ndarray,
):
    """Combine focus distance, ray spread, and field regularisation."""
    xt = solve_rays(params, sound_speed, x0, p0, m0, a0, mode, ts)
    dist_to_focus = jnp.linalg.norm(xt - focus[None, None, :], axis=-1)
    min_dist = jnp.min(dist_to_focus, axis=1)
    focus_loss = jnp.mean(min_dist)

    mean_dist = jnp.mean(dist_to_focus, axis=0)
    focus_time_idx = jnp.argmin(mean_dist)
    rays_at_focus = xt[:, focus_time_idx, :]
    spread_loss = jnp.mean(jnp.std(rays_at_focus, axis=0))
    smooth_loss = 0.01 * sum(jnp.mean(value**2) for value in params.values())

    total_loss = focus_loss + 0.5 * spread_loss + smooth_loss
    return total_loss, {
        "focus": focus_loss,
        "spread": spread_loss,
        "smooth": smooth_loss,
        "xt": xt,
    }


def speed_map(
    sound_speed: NeuralSoundSpeed,
    params: dict[str, jnp.ndarray],
    nx: int = 200,
    ny: int = 200,
) -> jnp.ndarray:
    xg = jnp.linspace(XMIN, XMAX, nx)
    yg = jnp.linspace(YMIN, YMAX, ny)
    xx, yy = jnp.meshgrid(xg, yg, indexing="xy")
    grid_points = jnp.stack([xx, yy], axis=-1)
    return sound_speed(grid_points, params)


def plot_rays_before_after(
    out_path,
    c_init_map,
    c_opt_map,
    xt_init,
    xt_final,
    x0,
    focus,
) -> None:
    vmin = float(jnp.minimum(jnp.min(c_init_map), jnp.min(c_opt_map)))
    vmax = float(jnp.maximum(jnp.max(c_init_map), jnp.max(c_opt_map)))

    fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(12, 5))
    panels = (
        (ax1, c_init_map, xt_init, r"$c_{\mathrm{init}}(\mathbf{x})$"),
        (ax2, c_opt_map, xt_final, r"$c_{\mathrm{opt}}(\mathbf{x})$"),
    )
    for ax, c_map, xt, title in panels:
        im = ax.imshow(
            np.asarray(c_map),
            extent=EXTENT,
            origin="lower",
            cmap="viridis",
            aspect="equal",
            vmin=vmin,
            vmax=vmax,
        )
        for ray in np.asarray(xt):
            ax.plot(ray[:, 0], ray[:, 1], "w-", lw=1.0, alpha=0.5)
        ax.scatter(
            np.asarray(x0[:, 0]),
            np.asarray(x0[:, 1]),
            s=20,
            c="red",
            marker="o",
            label="sources",
        )
        ax.scatter(
            float(focus[0]),
            float(focus[1]),
            s=100,
            c="red",
            marker="*",
            label="focus",
        )
        ax.set(title=title, xticks=[], yticks=[])
        ax.legend(frameon=True, fancybox=True, loc="lower right")

    fig.colorbar(im, ax=ax2)
    fig.tight_layout()
    fig.savefig(out_path, dpi=300, bbox_inches="tight")
    plt.close(fig)


def plot_loss(out_path, loss_history: list[float]) -> None:
    fig, ax = plt.subplots(figsize=(8, 4))
    ax.plot(loss_history)
    ax.set(xlabel="iteration", ylabel=r"$\mathcal{L}$")
    ax.grid(True, alpha=0.3)
    fig.tight_layout()
    fig.savefig(out_path, dpi=300, bbox_inches="tight")
    plt.close(fig)


def plot_speed_delta(out_path, c_init_map, c_opt_map, focus) -> None:
    diff_map = c_opt_map - c_init_map
    delta_max = float(jnp.max(jnp.abs(diff_map)))

    fig, ax = plt.subplots(figsize=(6, 5))
    im = ax.imshow(
        np.asarray(diff_map),
        extent=EXTENT,
        origin="lower",
        cmap="RdBu_r",
        aspect="equal",
        vmin=-delta_max,
        vmax=delta_max,
    )
    ax.scatter(float(focus[0]), float(focus[1]), s=100, c="black", marker="*")
    ax.set(title=r"$\Delta c(\mathbf{x})$", xticks=[], yticks=[])
    fig.colorbar(im, ax=ax)
    fig.tight_layout()
    fig.savefig(out_path, dpi=300, bbox_inches="tight")
    plt.close(fig)


def main() -> None:
    plot_dir = Path("plots/rays")
    plot_dir.mkdir(parents=True, exist_ok=True)
    use_beamax_style()

    x0, p0, m0, a0, mode, ts, focus = ray_setup()
    sound_speed = NeuralSoundSpeed(hidden_dim=32)
    params = sound_speed.params

    loss_grad = jax.jit(
        jax.value_and_grad(
            lambda current_params: focusing_loss(
                current_params,
                sound_speed,
                x0,
                p0,
                m0,
                a0,
                mode,
                ts,
                focus,
            ),
            has_aux=True,
        )
    )

    xt_init = solve_rays(params, sound_speed, x0, p0, m0, a0, mode, ts)
    lr_schedule = optax.exponential_decay(
        init_value=0.02,
        transition_steps=50,
        decay_rate=0.9,
    )
    optimizer = optax.chain(optax.clip_by_global_norm(1.0), optax.adam(lr_schedule))
    opt_state = optimizer.init(params)

    loss_history: list[float] = []
    best_loss = float("inf")
    best_params = params
    best_xt = xt_init

    num_iters = 300
    print("Optimizing speed of sound field for ray focusing...")
    for step in range(num_iters):
        (loss_value, aux), grads = loss_grad(params)
        loss_float = float(loss_value)
        loss_history.append(loss_float)

        # Keep the parameters and trajectory from the same iterate.
        if loss_float < best_loss:
            best_loss = loss_float
            best_params = params
            best_xt = aux["xt"]

        updates, opt_state = optimizer.update(grads, opt_state, params)
        params = optax.apply_updates(params, updates)

        if step % 50 == 0:
            print(
                f"iter {step:03d} | loss {loss_float:.6f} | "
                f"focus {float(aux['focus']):.4f} | spread {float(aux['spread']):.4f}"
            )

    c_init_map = speed_map(sound_speed, sound_speed.params)
    c_opt_map = speed_map(sound_speed, best_params)

    rays_path = plot_dir / "focusing_rays_before_after.png"
    loss_path = plot_dir / "focusing_loss_convergence.png"
    delta_path = plot_dir / "focusing_sound_speed_delta.png"

    plot_rays_before_after(
        rays_path,
        c_init_map,
        c_opt_map,
        xt_init,
        best_xt,
        x0,
        focus,
    )
    plot_loss(loss_path, loss_history)
    plot_speed_delta(delta_path, c_init_map, c_opt_map, focus)

    print("Optimization complete.")
    print(f"Final focusing loss: {best_loss:.6f}")
    print(f"Saved figures to {plot_dir.resolve()}")


if __name__ == "__main__":
    main()