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 coefficient_panels(
panels: Sequence[dict],
shape: tuple[int, int],
*,
figsize: tuple[float, float] = (12, 9),
sharex: bool = True,
sharey: bool = False,
suptitle: str | None = None,
suptitle_fontsize: float = 11,
default_x=None,
default_xlim: tuple[float, float] | None = None,
default_ylim: tuple[float, float] | None = None,
default_xlabel: str | None = None,
legend_fontsize: float | str = 8,
grid_alpha: float = 0.3,
):
r"""Plot coefficient curves in a configurable panel grid.
Parameters
----------
panels : sequence of dict
One panel specification per subplot in row-major order. Each
panel dictionary may contain:
``x`` : array_like, optional
Shared abscissa for all curves in the panel. If omitted,
*default_x* is used.
``curves`` : list of dict, optional
Curve specifications. Each curve dict must contain ``y`` and
may contain ``label`` plus ``plot_kwargs``. To style complex-valued
coefficient regions, provide ``complex_from`` as the complex
coefficient array used for detection. Samples satisfying
``abs(imag(C)) > complex_tol * max(1, abs(C))`` are plotted with
``complex_plot_kwargs`` (default ``{'ls': '--'}``). To style
regions where the outgoing branch itself is evanescent, provide
``evanescent_mask``; those samples are plotted with
``evanescent_plot_kwargs`` (default ``{'ls': '-.'}``).
``markers`` : list of dict, optional
Vertical-line marker specifications. Each marker dict must
contain ``angle`` and may contain ``label`` plus
``line_kwargs``.
``title``, ``xlabel``, ``ylabel`` : str, optional
Per-panel labels.
``xlim``, ``ylim`` : tuple, optional
Per-panel axis limits.
``legend`` : bool, optional
Whether to draw a legend for this panel.
``legend_loc`` : str, optional
Legend location. Default ``'upper right'``.
``legend_fontsize`` : float or str, optional
Legend font size override for this panel.
``grid`` : bool, optional
Whether to draw the grid. Default *True*.
``grid_alpha`` : float, optional
Grid alpha override for this panel.
shape : tuple of int
Panel layout ``(nrows, ncols)``.
figsize : tuple of float, optional
Figure size passed to :func:`matplotlib.pyplot.subplots`.
sharex, sharey : bool, optional
Forwarded to :func:`matplotlib.pyplot.subplots`.
suptitle : str, optional
Figure title.
suptitle_fontsize : float, optional
Font size for *suptitle*.
default_x : array_like, optional
Shared x-axis array used when a panel does not define ``x``.
default_xlim, default_ylim : tuple, optional
Default axis limits for panels that do not override them.
default_xlabel : str, optional
Default x-axis label for panels that do not define ``xlabel``.
legend_fontsize : float or str, optional
Default legend font size.
grid_alpha : float, optional
Default grid alpha.
Returns
-------
tuple
``(fig, axes)`` from matplotlib.
"""
import matplotlib.pyplot as plt
nrows, ncols = shape
if len(panels) != nrows * ncols:
raise ValueError("Number of panels must match shape[0] * shape[1].")
def _complex_mask(values, tol):
values = np.asarray(values)
scale = np.maximum(1.0, np.abs(values))
return np.isfinite(values) & (np.abs(np.imag(values)) > tol * scale)
def _connect_segment_starts(mask):
mask = np.asarray(mask, dtype=bool).copy()
starts = np.flatnonzero(mask & ~np.r_[False, mask[:-1]])
starts = starts[starts > 0]
mask[starts - 1] = True
return mask
fig, axes = plt.subplots(
nrows,
ncols,
figsize=figsize,
sharex=sharex,
sharey=sharey,
squeeze=False,
)
for ax, panel in zip(axes.flat, panels):
x = panel.get("x", default_x)
curves = panel.get("curves", [])
markers = panel.get("markers", [])
has_labeled_artist = False
if curves and x is None:
raise ValueError("Each panel with curves must define x or use default_x.")
for curve in curves:
plot_kwargs = dict(curve.get("plot_kwargs", {}))
label = curve.get("label")
if label is not None:
plot_kwargs.setdefault("label", label)
has_labeled_artist = True
complex_from = curve.get("complex_from")
evanescent_mask = curve.get("evanescent_mask")
if complex_from is None and evanescent_mask is None:
ax.plot(x, curve["y"], **plot_kwargs)
continue
x_arr = np.asarray(x)
y_arr = np.asarray(curve["y"])
if x_arr.shape != y_arr.shape:
raise ValueError(
"x, y, complex_from, and evanescent_mask must have the same shape."
)
if complex_from is None:
complex_state = np.zeros_like(y_arr, dtype=bool)
else:
complex_arr = np.asarray(complex_from)
if complex_arr.shape != y_arr.shape:
raise ValueError(
"x, y, complex_from, and evanescent_mask must have the same shape."
)
tol = curve.get("complex_tol", 1e-10)
complex_state = _complex_mask(complex_arr, tol)
if evanescent_mask is None:
evanescent_state = np.zeros_like(y_arr, dtype=bool)
else:
evanescent_state = np.asarray(evanescent_mask, dtype=bool)
if evanescent_state.shape != y_arr.shape:
raise ValueError(
"x, y, complex_from, and evanescent_mask must have the same shape."
)
states = [
(~complex_state) & (~evanescent_state),
complex_state & (~evanescent_state),
evanescent_state,
]
styles = [
{},
{"ls": "--", **curve.get("complex_plot_kwargs", {})},
{"ls": "-.", **curve.get("evanescent_plot_kwargs", {})},
]
base_plot_kwargs = dict(plot_kwargs)
base_plot_kwargs.pop("label", None)
label_pending = label is not None
for state, style in zip(states, styles):
if not np.any(state):
continue
segment_mask = _connect_segment_starts(state)
segment_kwargs = dict(base_plot_kwargs)
segment_kwargs.update(style)
if label_pending:
segment_kwargs["label"] = label
label_pending = False
ax.plot(
x_arr,
np.where(segment_mask, y_arr, np.nan),
**segment_kwargs,
)
for marker in markers:
angle = marker.get("angle")
if angle is None or not np.isfinite(angle):
continue
line_kwargs = dict(marker.get("line_kwargs", {}))
label = marker.get("label")
if label is not None:
line_kwargs.setdefault("label", label)
has_labeled_artist = True
ax.axvline(angle, **line_kwargs)
xlim = panel.get("xlim", default_xlim)
if xlim is not None:
ax.set_xlim(xlim)
ylim = panel.get("ylim", default_ylim)
if ylim is not None:
ax.set_ylim(*ylim)
xlabel = panel.get("xlabel", default_xlabel)
if xlabel is not None:
ax.set_xlabel(xlabel)
ylabel = panel.get("ylabel")
if ylabel is not None:
ax.set_ylabel(ylabel)
title = panel.get("title")
if title is not None:
ax.set_title(title)
if panel.get("grid", True):
ax.grid(True, alpha=panel.get("grid_alpha", grid_alpha))
show_legend = panel.get("legend", has_labeled_artist)
if show_legend and has_labeled_artist:
ax.legend(
fontsize=panel.get("legend_fontsize", legend_fontsize),
loc=panel.get("legend_loc", "upper right"),
)
if suptitle is not None:
fig.suptitle(suptitle, fontsize=suptitle_fontsize)
fig.tight_layout(rect=(0.0, 0.0, 1.0, 0.96))
else:
fig.tight_layout()
return fig, axes
[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,
ray_linewidth: float = 0.8,
xlim: tuple | None = None,
ylim: tuple | None = None,
unit: str = "m",
plot_model: bool = True,
add_colorbar: bool = False,
discrete_colorbar: bool = False,
layer_colors: Sequence[str] | None = None,
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
ray_linewidth : float
Ray line width.
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.
layer_colors : sequence of str, optional
Custom colors for the velocity model layers. If omitted,
``viridis`` is used.
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, ListedColormap
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)
# 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]
layer_cmap = None
if layer_colors is not None:
layer_colors = list(layer_colors)
if len(layer_colors) < len(unique_vels):
raise ValueError(
"layer_colors must provide at least one color per unique layer value"
)
layer_cmap = ListedColormap(layer_colors[:len(unique_vels)])
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 = layer_cmap or 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 = layer_cmap or cm.get_cmap("viridis", 1)
norm = BoundaryNorm(bounds, cmap.N)
else:
cmap = layer_cmap or 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=ray_linewidth, **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