r"""
Standalone visualisation helpers for LayTracer.
Provides matplotlib-based 2-D plots and a Plotly-based 3-D ray viewer.
"""
from __future__ import annotations
from typing import Sequence
import numpy as np
import pandas as pd
[docs]
def velocity_profile(
vel_df: pd.DataFrame,
param: str = "Vp",
ax=None,
color: str | None = None,
label: str | None = None,
xlim: tuple | None = None,
ylim: tuple | None = None,
unit: str = "m",
**kwargs,
):
r"""Plot a 1-D model parameter–depth step profile.
Parameters
----------
vel_df : pandas.DataFrame
Velocity model.
param : str, optional
``'Vp'`` (default), ``'Vs'``, ``'Rho'``, ``'Qp'``, or ``'Qs'``.
ax : matplotlib.axes.Axes, optional
Axes to plot on. Created if *None*.
color : str, optional
Line colour.
label : str, optional
Legend label.
xlim, ylim : tuple, optional
Axis limits ``(min, max)``.
unit : str, optional
``'m'`` (default) or ``'km'``. Scales the vertical depth axis.
Returns
-------
matplotlib.axes.Axes
"""
import matplotlib.pyplot as plt
if ax is None:
_, ax = plt.subplots(figsize=(4, 6))
depths = vel_df["Depth"].values
vals = vel_df[param].values
n = len(depths)
scale = 1000.0 if unit.lower() == "km" else 1.0
# Step profile
z_plot, v_plot = [], []
for i in range(n):
z_top = depths[i] / scale
if i + 1 < n:
z_bot = depths[i + 1] / scale
else:
# Last layer (half-space) extension
span = (depths[-1] - depths[0]) / scale
if span == 0:
span = 1000.0 / scale # Fallback for single-layer model
z_def = z_top + span * 0.3
if ylim:
# Extend to at least the plot limit if provided
z_bot = max(z_def, float(max(ylim)))
else:
z_bot = z_def
z_plot.extend([z_top, z_bot])
v_plot.extend([vals[i], vals[i]])
ax.plot(v_plot, z_plot, color=color, label=label, **kwargs)
if param in ("Vp", "Vs"):
xlabel_str = f"{param} (m/s)"
elif param == "Rho":
xlabel_str = r"$\rho$ (kg/m³)"
elif param in ("Qp", "Qs"):
xlabel_str = f"{param}"
else:
xlabel_str = param
ax.margins(y=0)
ax.set_xlabel(xlabel_str)
ax.set_ylabel(f"Depth ({unit})")
# Handle y-axis limits and inversion
if ylim:
ax.set_ylim(ylim)
# Only invert if top is less than bottom (matplotlib default puts 0 at bottom)
bottom, top = ax.get_ylim()
if bottom < top:
ax.invert_yaxis()
if xlim:
ax.set_xlim(xlim)
title_str = "Velocity profile" if param in ("Vp", "Vs") else f"{param} profile"
ax.set_title(title_str)
return ax
[docs]
def rays_2d(
vel_df: pd.DataFrame,
rays: Sequence[np.ndarray],
vel_type: str = "Vp",
sources: np.ndarray | None = None,
receivers: np.ndarray | None = None,
ax=None,
ray_color: str = "k",
ray_alpha: float = 0.6,
xlim: tuple | None = None,
ylim: tuple | None = None,
unit: str = "m",
plot_model: bool = True,
add_colorbar: bool = False,
discrete_colorbar: bool = False,
model_alpha: float = 1.0,
equal_scale: bool = True,
colorbar_orientation: str = "vertical",
**kwargs,
):
r"""Plot ray paths over a 2-D layered velocity cross-section.
Parameters
----------
vel_df : pandas.DataFrame
Velocity model.
rays : list of numpy.ndarray
Each element is shape ``(M, 2)`` or ``(M, 3)``. If 3-D, the
first two columns are treated as horizontal/depth.
vel_type : str
``'Vp'`` or ``'Vs'`` — used for layer colouring.
sources, receivers : numpy.ndarray, optional
Coordinate arrays for plotting markers.
ax : matplotlib.axes.Axes, optional
ray_color : str
ray_alpha : float
xlim, ylim : tuple, optional
plot_model : bool
If *True* (default), plot the velocity model background and
set axis labels/titles. If *False*, only plot the rays and
markers.
unit : str
``'m'`` (default) or ``'km'``. Scales coordinates and labels.
add_colorbar : bool
If *True* (default *False*), add a colorbar for the velocity
model. Only applies if *plot_model* is True.
discrete_colorbar : bool
If *True* (default *False*), quantize the colormap to the
unique velocity values in the model.
model_alpha : float
Opacity of the velocity model layers (0.0 to 1.0). Default 1.0.
equal_scale : bool
If *True* (default *True*), force equal scaling for x and y axes
using ``ax.set_aspect('equal')``.
colorbar_orientation : str
``'vertical'`` (default) or ``'horizontal'``.
Returns
-------
matplotlib.axes.Axes
"""
import matplotlib.pyplot as plt
from matplotlib.patches import Rectangle
from matplotlib.collections import PatchCollection
import matplotlib.cm as cm
from matplotlib.colors import Normalize, BoundaryNorm
if ax is None:
_, ax = plt.subplots(figsize=(10, 6))
depths = vel_df["Depth"].values
vels = vel_df[vel_type].values
n = len(depths)
scale = 1000.0 if unit.lower() == "km" else 1.0
# Determine x-range from rays (only if we need to plot model or set limits)
if plot_model:
if rays:
all_x = np.concatenate([r[:, 0] / scale for r in rays])
x_lo, x_hi = all_x.min(), all_x.max()
else:
x_lo, x_hi = 0, 1000 / scale # Default fallback
if xlim:
x_lo, x_hi = sorted(xlim)
pad = (x_hi - x_lo) * 0.05
x_lo -= pad
x_hi += pad
# Ensure model background covers the requested xlim if provided
if xlim:
x_lo = min(x_lo, min(xlim))
x_hi = max(x_hi, max(xlim))
# Layer rectangles
unique_vels = np.sort(np.unique(vels))
vmin, vmax = unique_vels[0], unique_vels[-1]
if discrete_colorbar and len(unique_vels) > 1:
# Create discrete boundaries
# Midpoints between values
mids = (unique_vels[:-1] + unique_vels[1:]) / 2.0
# Extend to cover first and last
# We can pick abitrary padding, e.g. estimated step
step = (vmax - vmin) / (len(unique_vels) - 1) if len(unique_vels) > 1 else 1.0
bounds = np.concatenate(([vmin - step/2], mids, [vmax + step/2]))
cmap = cm.get_cmap("viridis", len(unique_vels))
norm = BoundaryNorm(bounds, cmap.N)
elif discrete_colorbar and len(unique_vels) == 1:
# Single unique velocity — one discrete colour, tick at exact value
bounds = np.array([vmin - 0.5, vmax + 0.5])
cmap = cm.get_cmap("viridis", 1)
norm = BoundaryNorm(bounds, cmap.N)
else:
cmap = cm.get_cmap("viridis")
# Guard against vmin == vmax (e.g. single-velocity model)
if vmin == vmax:
norm = Normalize(vmin=vmin - 1.0, vmax=vmax + 1.0)
else:
norm = Normalize(vmin=vmin, vmax=vmax)
# Determine the deepest point across all rays, sources, and receivers
# so the half-space layer rectangle always covers the visible area.
z_max_data = 0.0
if rays:
all_z = np.concatenate(
[(r[:, -1] if r.shape[1] == 2 else r[:, 2]) / scale for r in rays]
)
z_max_data = max(z_max_data, float(all_z.max()))
if sources is not None:
z_max_data = max(z_max_data, float(np.atleast_2d(sources)[:, -1].max() / scale))
if receivers is not None:
z_max_data = max(z_max_data, float(np.atleast_2d(receivers)[:, -1].max() / scale))
patches = []
colors_list = []
for i in range(n):
z_top = depths[i] / scale
if i + 1 < n:
z_bot = depths[i + 1] / scale
else:
# Last layer (half-space) extension
span = (depths[-1] - depths[0]) / scale
if span == 0: span = 1000.0 / scale
z_def = z_top + span * 0.3
# Extend to cover the deepest data point (with 10 % padding)
z_def = max(z_def, z_max_data * 1.1)
if ylim:
z_bot = max(z_def, float(max(ylim)))
else:
z_bot = z_def
rect = Rectangle((x_lo, z_top), x_hi - x_lo, z_bot - z_top)
patches.append(rect)
# Use the norm to map velocity to color
colors_list.append(cmap(norm(vels[i])))
pc = PatchCollection(patches, facecolor=colors_list, alpha=model_alpha, edgecolor="grey", linewidth=0.5)
ax.add_collection(pc)
if add_colorbar:
from mpl_toolkits.axes_grid1 import make_axes_locatable
divider = make_axes_locatable(ax)
sm = cm.ScalarMappable(cmap=cmap, norm=norm)
sm.set_array([])
if colorbar_orientation == "horizontal":
cax = divider.append_axes("bottom", size="5%", pad=0.5)
else:
cax = divider.append_axes("right", size="5%", pad=0.1)
# Format ticks for discrete case
if discrete_colorbar and len(unique_vels) >= 1:
# Place ticks at the unique values
ticks = unique_vels
plt.colorbar(sm, cax=cax, orientation=colorbar_orientation, label=f"{vel_type} (m/s)", alpha=model_alpha, ticks=ticks)
else:
plt.colorbar(sm, cax=cax, orientation=colorbar_orientation, label=f"{vel_type} (m/s)", alpha=model_alpha)
# Rays
for ray in rays:
x = ray[:, 0] / scale
z = (ray[:, -1] if ray.shape[1] == 2 else ray[:, 2]) / scale
ax.plot(x, z, color=ray_color, alpha=ray_alpha, linewidth=0.8, **kwargs)
# Markers
if sources is not None:
src = np.atleast_2d(sources)
ax.scatter(src[:, 0] / scale, src[:, -1] / scale, marker="*", s=120, c="red", zorder=5, label="Source")
if receivers is not None:
rcv = np.atleast_2d(receivers)
ax.scatter(rcv[:, 0] / scale, rcv[:, -1] / scale, marker="v", s=60, c="blue", zorder=5, label="Receiver")
ax.margins(y=0)
if plot_model:
ax.invert_yaxis()
ax.set_xlabel(f"Horizontal distance ({unit})")
ax.set_ylabel(f"Depth ({unit})")
ax.set_title("Ray paths")
if xlim:
ax.set_xlim(xlim)
elif plot_model:
ax.set_xlim(x_lo, x_hi)
if ylim:
ax.set_ylim(ylim)
elif plot_model:
# Default depth range from model (0 to bottom)
# Find max depth of model or rays
z_max_model = depths[-1] / scale
if rays:
all_z = np.concatenate([(r[:, -1] if r.shape[1] == 2 else r[:, 2]) / scale for r in rays])
z_max_rays = all_z.max()
z_max_model = max(z_max_model, z_max_rays)
# Add slight padding at bottom
z_max_model *= 1.1
ax.set_ylim(z_max_model, 0)
if equal_scale:
ax.set_aspect("equal")
return ax
[docs]
def rays_3d(
vel_df: pd.DataFrame,
rays: Sequence[np.ndarray],
vel_type: str = "Vp",
sources: np.ndarray | None = None,
receivers: np.ndarray | None = None,
ray_color: str = "red",
opacity: float = 0.3,
**kwargs,
):
r"""Interactive 3-D ray visualisation using Plotly.
Parameters
----------
vel_df : pandas.DataFrame
Velocity model.
rays : list of numpy.ndarray
Each element is shape ``(M, 3)``.
vel_type : str
``'Vp'`` or ``'Vs'``.
sources, receivers : numpy.ndarray, optional
Coordinate arrays for plotting markers.
ray_color : str
Ray trace colour.
opacity : float
Layer surface opacity.
Returns
-------
plotly.graph_objects.Figure
"""
import plotly.graph_objects as go
fig = go.Figure()
# Rays
for ray in rays:
fig.add_trace(
go.Scatter3d(
x=ray[:, 0],
y=ray[:, 1],
z=ray[:, 2],
mode="lines",
line=dict(color=ray_color, width=2),
showlegend=False,
)
)
# Source / receiver markers
if sources is not None:
src = np.atleast_2d(sources)
fig.add_trace(
go.Scatter3d(
x=src[:, 0], y=src[:, 1], z=src[:, 2],
mode="markers",
marker=dict(size=6, color="red", symbol="diamond"),
name="Sources",
)
)
if receivers is not None:
rcv = np.atleast_2d(receivers)
fig.add_trace(
go.Scatter3d(
x=rcv[:, 0], y=rcv[:, 1], z=rcv[:, 2],
mode="markers",
marker=dict(size=4, color="blue", symbol="circle"),
name="Receivers",
)
)
fig.update_layout(
scene=dict(
xaxis_title="X (m)",
yaxis_title="Y (m)",
zaxis_title="Depth (m)",
zaxis=dict(autorange="reversed"),
aspectmode="data",
),
title="3-D Ray paths",
)
return fig