lsurf 1.0.0__py3-none-any.whl

lsurf/visualization/atmospheric_plots.py

@@ -0,0 +1,754 @@
+# The Clear BSD License
+#
+# Copyright (c) 2026 Tobias Heibges
+# All rights reserved.
+#
+# Redistribution and use in source and binary forms, with or without
+# modification, are permitted (subject to the limitations in the disclaimer
+# below) provided that the following conditions are met:
+#
+#      * Redistributions of source code must retain the above copyright notice,
+#      this list of conditions and the following disclaimer.
+#
+#      * Redistributions in binary form must reproduce the above copyright
+#      notice, this list of conditions and the following disclaimer in the
+#      documentation and/or other materials provided with the distribution.
+#
+#      * Neither the name of the copyright holder nor the names of its
+#      contributors may be used to endorse or promote products derived from this
+#      software without specific prior written permission.
+#
+# NO EXPRESS OR IMPLIED LICENSES TO ANY PARTY'S PATENT RIGHTS ARE GRANTED BY
+# THIS LICENSE. THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND
+# CONTRIBUTORS "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
+# LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A
+# PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR
+# CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,
+# EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,
+# PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR
+# BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER
+# IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
+# ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
+# POSSIBILITY OF SUCH DAMAGE.
+
+"""
+Atmospheric Refraction Visualization
+
+Functions for plotting ray trajectories through inhomogeneous atmospheres
+and refractive index profiles.
+"""
+
+from pathlib import Path
+from typing import Any
+
+import matplotlib.pyplot as plt
+import numpy as np
+from numpy.typing import NDArray
+
+from ..materials.utils.constants import EARTH_RADIUS
+
+
+def plot_ray_trajectories(
+    ax: plt.Axes,
+    trajectories: dict[float, NDArray],
+    source_altitude: float = 0.0,
+    show_straight_lines: bool = True,
+    show_earth_surface: bool = True,
+    duct_center: float = 0.0,
+    duct_width: float = 0.0,
+    xlim: tuple[float, float] | None = None,
+    ylim: tuple[float, float] | None = None,
+    cmap: str = "turbo",
+    linewidth: float = 1.5,
+    impact_parameter_keys: bool = False,
+    use_colorbar: bool = False,
+) -> plt.cm.ScalarMappable | None:
+    """
+    Plot ray trajectories on a single axis.
+
+    Parameters
+    ----------
+    ax : matplotlib.axes.Axes
+        Axes to plot on.
+    trajectories : dict
+        Dictionary mapping initial angle (degrees) or impact parameter (meters)
+        to trajectory array. Each trajectory is shape (N, 2) with [x, z] in meters.
+    source_altitude : float
+        Source altitude in meters (for straight line reference).
+    show_straight_lines : bool
+        If True, show dashed straight-line references.
+    show_earth_surface : bool
+        If True, show Earth's curved surface.
+    xlim : tuple, optional
+        X-axis limits in km.
+    ylim : tuple, optional
+        Y-axis limits in km.
+    cmap : str
+        Colormap for ray colors.
+    linewidth : float
+        Line width for trajectories.
+    impact_parameter_keys : bool
+        If True, dictionary keys are impact parameters in meters (not angles).
+    use_colorbar : bool
+        If True, use colorbar instead of legend (better for many rays).
+
+    Returns
+    -------
+    sm : ScalarMappable or None
+        ScalarMappable for colorbar if use_colorbar=True, else None.
+
+    Examples
+    --------
+    >>> fig, ax = plt.subplots()
+    >>> plot_ray_trajectories(ax, trajectories, source_altitude=0)
+    >>> plt.show()
+    """
+    keys = list(trajectories.keys())
+    n_rays = len(keys)
+    colors = plt.get_cmap(cmap)(np.linspace(0.1, 0.9, n_rays))
+
+    # For colorbar, create a ScalarMappable
+    sm = None
+    if use_colorbar and impact_parameter_keys:
+        norm = plt.Normalize(vmin=min(keys) / 1000, vmax=max(keys) / 1000)
+        sm = plt.cm.ScalarMappable(cmap=cmap, norm=norm)
+        sm.set_array([])
+
+    for i, ((key, traj), color) in enumerate(zip(trajectories.items(), colors)):
+        x_km = traj[:, 0] / 1000
+        z_km = traj[:, 1] / 1000
+
+        # Determine label
+        if use_colorbar:
+            label = None
+        elif impact_parameter_keys:
+            label = f"{key / 1000:.1f} km"
+        else:
+            label = f"{key}°"
+
+        # Plot refracted ray
+        ax.plot(x_km, z_km, color=color, linewidth=linewidth, label=label)
+
+        # Plot straight-line reference (only for angle-based trajectories)
+        if show_straight_lines and not impact_parameter_keys:
+            angle_rad = np.radians(key)
+            if np.abs(np.cos(angle_rad)) > 0.1:
+                max_x = traj[-1, 0]
+                straight_z = (source_altitude + max_x * np.tan(angle_rad)) / 1000
+                ax.plot(
+                    [0, max_x / 1000],
+                    [source_altitude / 1000, straight_z],
+                    color=color,
+                    linewidth=linewidth * 0.6,
+                    linestyle="--",
+                    alpha=0.4,
+                )
+
+    # Plot Earth's surface
+    if show_earth_surface and xlim is not None:
+        x_surface = np.linspace(xlim[0] * 1000, xlim[1] * 1000, 500)
+        z_surface = -EARTH_RADIUS + np.sqrt(
+            np.maximum(EARTH_RADIUS**2 - x_surface**2, 0)
+        )
+        ax.fill_between(
+            x_surface / 1000, z_surface / 1000, -50, color="saddlebrown", alpha=0.3
+        )
+        ax.plot(x_surface / 1000, z_surface / 1000, "k-", linewidth=2)
+
+    # Plot duct region
+    if show_earth_surface and xlim is not None and duct_width > 0:
+        x_surface = np.linspace(xlim[0] * 1000, xlim[1] * 1000, 500)
+        z_surface_l = -EARTH_RADIUS + np.sqrt(
+            np.maximum(
+                (EARTH_RADIUS + duct_center - duct_width / 2) ** 2 - x_surface**2,
+                0,
+            )
+        )
+        z_surface_u = -EARTH_RADIUS + np.sqrt(
+            np.maximum(
+                (EARTH_RADIUS + duct_center + duct_width / 2) ** 2 - x_surface**2, 0
+            )
+        )
+        ax.fill_between(
+            x_surface / 1000,
+            z_surface_l / 1000,
+            z_surface_u / 1000,
+            color="blue",
+            alpha=0.1,
+        )
+        ax.plot(x_surface / 1000, z_surface_l / 1000, "k:", linewidth=1)
+        ax.plot(x_surface / 1000, z_surface_u / 1000, "k:", linewidth=1)
+
+    ax.set_xlabel("Horizontal distance (km)")
+    ax.set_ylabel("Altitude (km)")
+    ax.grid(True, alpha=0.3)
+
+    # Add legend only if not using colorbar and limited number of rays
+    if not use_colorbar:
+        title = "Impact param (km)" if impact_parameter_keys else "Initial angle"
+        ax.legend(title=title, loc="upper left", fontsize=8, ncol=2)
+
+    if xlim:
+        ax.set_xlim(xlim)
+    if ylim:
+        ax.set_ylim(ylim)
+
+    return sm
+
+
+def plot_refractive_index_profile(
+    ax: plt.Axes,
+    atmosphere: Any,
+    max_altitude_km: float = 100.0,
+    n_points: int = 500,
+    color: str = "b",
+    linewidth: float = 2.0,
+    annotate_sea_level: bool = True,
+    wavelength: float = 532e-9,
+) -> None:
+    """
+    Plot refractive index vs altitude profile.
+
+    Parameters
+    ----------
+    ax : matplotlib.axes.Axes
+        Axes to plot on.
+    atmosphere : MaterialField
+        Atmosphere material with n_at_altitude method.
+    max_altitude_km : float
+        Maximum altitude in km.
+    n_points : int
+        Number of points for profile.
+    color : str
+        Line color.
+    linewidth : float
+        Line width.
+    annotate_sea_level : bool
+        If True, annotate the sea level value.
+    wavelength : float
+        Wavelength in meters for spectral atmospheres (default: 532nm).
+
+    Examples
+    --------
+    >>> fig, ax = plt.subplots()
+    >>> plot_refractive_index_profile(ax, atmosphere)
+    >>> plt.show()
+    """
+    altitudes_km = np.linspace(0, max_altitude_km, n_points)
+
+    # Handle both simple (1-arg) and spectral (2-arg) atmospheres
+    def get_n(altitude_m: float) -> float:
+        try:
+            return atmosphere.n_at_altitude(altitude_m, wavelength)
+        except TypeError:
+            return atmosphere.n_at_altitude(altitude_m)
+
+    n_values = [get_n(h * 1000) for h in altitudes_km]
+
+    ax.plot(n_values, altitudes_km, color=color, linewidth=linewidth)
+    ax.set_xlabel("Refractive Index n")
+    ax.set_ylabel("Altitude (km)")
+    ax.set_title("Atmospheric Refractive Index Profile")
+    ax.grid(True, alpha=0.3)
+
+    # Annotate sea level
+    if annotate_sea_level:
+        n_sea = get_n(0)
+        ax.axhline(y=0, color="brown", linestyle="-", linewidth=1)
+        # Place text inside plot, above sea level line
+        ax.text(
+            n_sea,
+            max_altitude_km * 0.08,
+            f"n₀ = {n_sea:.6f}",
+            fontsize=9,
+            ha="center",
+            va="bottom",
+            bbox=dict(boxstyle="round,pad=0.3", facecolor="white", alpha=0.8),
+        )
+
+
+def create_atmospheric_refraction_figure(
+    trajectories: dict[float, NDArray],
+    atmosphere: Any,
+    title: str = "Atmospheric Refraction",
+    source_altitude: float = 0.0,
+    xlim: tuple[float, float] | None = None,
+    ylim: tuple[float, float] | None = None,
+    show_earth_surface: bool = True,
+    show_straight_lines: bool = True,
+    duct_center: float = 0.0,
+    duct_width: float = 0.0,
+    figsize: tuple[float, float] = (14, 5),
+) -> plt.Figure:
+    """
+    Create a two-panel figure with ray trajectories and refractive index profile.
+
+    Parameters
+    ----------
+    trajectories : dict
+        Dictionary mapping initial angle (degrees) to trajectory array.
+        Each trajectory is shape (N, 2) with columns [x, z] in meters.
+    atmosphere : MaterialField
+        Atmosphere material with n_at_altitude method.
+    title : str
+        Title for the trajectory plot.
+    source_altitude : float
+        Source altitude in meters.
+    xlim : tuple, optional
+        X-axis limits in km for trajectory plot.
+    ylim : tuple, optional
+        Y-axis limits in km for trajectory plot.
+    show_earth_surface : bool
+        If True, show Earth's curved surface in trajectory plot.
+    show_straight_lines : bool
+        If True, show dashed straight-line references.
+    figsize : tuple
+        Figure size (width, height) in inches.
+
+    Returns
+    -------
+    fig : matplotlib.figure.Figure
+        The created figure.
+
+    Examples
+    --------
+    >>> fig = create_atmospheric_refraction_figure(
+    ...     trajectories, atmosphere,
+    ...     title="Exponential Atmosphere",
+    ...     xlim=(0, 1000), ylim=(0, 100),
+    ... )
+    >>> plt.savefig("refraction.png")
+    """
+    fig, (ax1, ax2) = plt.subplots(1, 2, figsize=figsize)
+
+    # Left: Ray trajectories
+    plot_ray_trajectories(
+        ax1,
+        trajectories,
+        source_altitude=source_altitude,
+        show_straight_lines=show_straight_lines,
+        show_earth_surface=show_earth_surface,
+        duct_center=duct_center,
+        duct_width=duct_width,
+        xlim=xlim,
+        ylim=ylim,
+    )
+    subtitle = "(solid=refracted, dashed=straight)" if show_straight_lines else ""
+    ax1.set_title(f"{title}\n{subtitle}".strip())
+
+    # Right: Refractive index profile
+    max_alt = ylim[1] if ylim else 100.0
+    plot_refractive_index_profile(
+        ax=ax2, atmosphere=atmosphere, max_altitude_km=max_alt
+    )
+
+    plt.tight_layout()
+    return fig
+
+
+def plot_trajectory_comparison(
+    results: list[dict],
+    figsize: tuple[float, float] | None = None,
+    suptitle: str = "Atmospheric Model Comparison",
+) -> plt.Figure:
+    """
+    Create a comparison figure with multiple trajectory plots.
+
+    Parameters
+    ----------
+    results : list of dict
+        List of result dictionaries, each containing:
+        - 'trajectories': dict mapping angle to trajectory array
+        - 'title': plot title
+        - 'xlim': x-axis limits in km
+        - 'ylim': y-axis limits in km
+    figsize : tuple, optional
+        Figure size. Auto-determined if None.
+    suptitle : str
+        Super title for the figure.
+
+    Returns
+    -------
+    fig : matplotlib.figure.Figure
+        The created figure.
+
+    Examples
+    --------
+    >>> results = [
+    ...     {'trajectories': traj1, 'title': 'Model A', 'xlim': (0, 100), 'ylim': (0, 50)},
+    ...     {'trajectories': traj2, 'title': 'Model B', 'xlim': (0, 100), 'ylim': (0, 50)},
+    ... ]
+    >>> fig = plot_trajectory_comparison(results)
+    """
+    n_plots = len(results)
+    if n_plots == 0:
+        return plt.figure()
+
+    # Determine grid layout
+    if n_plots <= 2:
+        rows, cols = 1, n_plots
+    elif n_plots <= 4:
+        rows, cols = 2, 2
+    elif n_plots <= 6:
+        rows, cols = 2, 3
+    else:
+        rows, cols = 3, 3
+
+    if figsize is None:
+        figsize = (5 * cols, 4 * rows)
+
+    fig, axes = plt.subplots(rows, cols, figsize=figsize)
+    if n_plots == 1:
+        axes = [axes]
+    else:
+        axes = axes.flatten()
+
+    for ax, result in zip(axes, results):
+        traj = result["trajectories"]
+        colors = plt.get_cmap("turbo")(np.linspace(0.1, 0.9, len(traj)))
+
+        for (angle_deg, t), color in zip(traj.items(), colors):
+            ax.plot(t[:, 0] / 1000, t[:, 1] / 1000, color=color, linewidth=1.5)
+
+        ax.set_title(result.get("title", ""), fontsize=10)
+        if "xlim" in result:
+            ax.set_xlim(result["xlim"])
+        if "ylim" in result:
+            ax.set_ylim(result["ylim"])
+        ax.set_xlabel("Distance (km)", fontsize=9)
+        ax.set_ylabel("Altitude (km)", fontsize=9)
+        ax.grid(True, alpha=0.3)
+
+    # Hide unused subplots
+    for ax in axes[n_plots:]:
+        ax.axis("off")
+
+    plt.suptitle(suptitle, fontsize=14)
+    plt.tight_layout()
+    return fig
+
+
+def plot_trajectory_offset(
+    ax: plt.Axes,
+    trajectories: dict[float, NDArray],
+    source_altitude: float = 0.0,
+    cmap: str = "turbo",
+    linewidth: float = 1.5,
+    xlim: tuple[float, float] | None = None,
+) -> None:
+    """
+    Plot angular deviation from initial ray direction vs horizontal distance.
+
+    Shows how much the ray direction has changed from its initial launch angle
+    as it propagates through the atmosphere. Atmospheric refraction typically
+    bends rays downward (toward higher n), resulting in negative angular offsets.
+
+    Parameters
+    ----------
+    ax : matplotlib.axes.Axes
+        Axes to plot on.
+    trajectories : dict
+        Dictionary mapping initial angle (degrees) to trajectory array.
+        Each trajectory is shape (N, 2) with columns [x, z] in meters.
+    source_altitude : float
+        Source altitude in meters (unused, kept for API compatibility).
+    cmap : str
+        Colormap for ray colors.
+    linewidth : float
+        Line width for curves.
+    xlim : tuple, optional
+        X-axis limits in km.
+
+    Examples
+    --------
+    >>> fig, ax = plt.subplots()
+    >>> plot_trajectory_offset(ax, trajectories, source_altitude=0)
+    >>> plt.show()
+    """
+    colors = plt.get_cmap(cmap)(np.linspace(0.1, 0.9, len(trajectories)))
+
+    for (angle_deg, traj), color in zip(trajectories.items(), colors):
+        x_m = traj[:, 0]
+
+        # Compute direction angle at each point from finite differences
+        dx = np.diff(traj[:, 0])
+        dz = np.diff(traj[:, 1])
+
+        # Direction angle in degrees (angle above horizontal)
+        segment_angles = np.degrees(np.arctan2(dz, dx))
+
+        # Compute angular deviation from initial direction
+        # First point has zero deviation, subsequent points show the change
+        angular_deviation = np.zeros(len(traj))
+        angular_deviation[0] = 0.0  # No deviation at start
+        angular_deviation[1:] = segment_angles - angle_deg  # Deviation from initial
+
+        ax.plot(
+            x_m / 1000,
+            angular_deviation,
+            color=color,
+            linewidth=linewidth,
+            label=f"{angle_deg}°",
+        )
+
+    ax.axhline(y=0, color="gray", linestyle="--", linewidth=1, alpha=0.5)
+    ax.set_xlabel("Horizontal distance (km)")
+    ax.set_ylabel("Angular deviation (°)")
+    ax.set_title("Angular Deviation from Initial Direction")
+    ax.legend(title="Initial angle", loc="best", fontsize=9)
+    ax.grid(True, alpha=0.3)
+
+    if xlim:
+        ax.set_xlim(xlim)
+
+
+def create_atmospheric_refraction_figure_with_offset(
+    trajectories: dict[float, NDArray],
+    atmosphere: Any,
+    title: str = "Atmospheric Refraction",
+    source_altitude: float = 0.0,
+    xlim: tuple[float, float] | None = None,
+    ylim: tuple[float, float] | None = None,
+    show_earth_surface: bool = True,
+    show_straight_lines: bool = True,
+    duct_center: float = 0.0,
+    duct_width: float = 0.0,
+    figsize: tuple[float, float] = (16, 10),
+    horizontal_rays: bool = False,
+) -> plt.Figure:
+    """
+    Create a four-panel figure with trajectories, offset, and refractive index.
+
+    Panels:
+    - Top-left: Ray trajectories (refracted vs straight-line)
+    - Top-right: Refractive index profile vs altitude
+    - Bottom: Altitude offset between refracted and straight paths
+
+    Parameters
+    ----------
+    trajectories : dict
+        Dictionary mapping initial angle (degrees) to trajectory array.
+        Each trajectory is shape (N, 2) with columns [x, z] in meters.
+    atmosphere : MaterialField
+        Atmosphere material with n_at_altitude method.
+    title : str
+        Title for the figure.
+    source_altitude : float
+        Source altitude in meters.
+    xlim : tuple, optional
+        X-axis limits in km for trajectory plot.
+    ylim : tuple, optional
+        Y-axis limits in km for trajectory plot.
+    show_earth_surface : bool
+        If True, show Earth's curved surface in trajectory plot.
+    show_straight_lines : bool
+        If True, show dashed straight-line references.
+    figsize : tuple
+        Figure size (width, height) in inches.
+    horizontal_rays : bool
+        If True, use plot for horizontal rays with varying impact parameters.
+        If False (default), use plot for rays with varying initial angles.
+
+    Returns
+    -------
+    fig : matplotlib.figure.Figure
+        The created figure.
+
+    Examples
+    --------
+    >>> fig = create_atmospheric_refraction_figure_with_offset(
+    ...     trajectories, atmosphere,
+    ...     title="Exponential Atmosphere",
+    ...     xlim=(0, 1000), ylim=(0, 100),
+    ... )
+    >>> plt.savefig("refraction_with_offset.png")
+    """
+    fig = plt.figure(figsize=figsize)
+
+    # Create grid using GridSpec for better control
+    if horizontal_rays:
+        # Use gridspec to leave room for colorbar
+        gs = fig.add_gridspec(2, 3, width_ratios=[1, 1, 0.05], wspace=0.3, hspace=0.3)
+        ax1 = fig.add_subplot(gs[0, 0])  # Top-left: trajectories
+        ax2 = fig.add_subplot(gs[0, 1])  # Top-right: n profile
+        ax3 = fig.add_subplot(gs[1, 0:2])  # Bottom: deviation (spans 2 cols)
+        cax = fig.add_subplot(gs[:, 2])  # Colorbar (right side, full height)
+    else:
+        ax1 = fig.add_subplot(2, 2, 1)
+        ax2 = fig.add_subplot(2, 2, 2)
+        ax3 = fig.add_subplot(2, 1, 2)
+
+    # Top-left: Ray trajectories
+    sm = plot_ray_trajectories(
+        ax1,
+        trajectories,
+        source_altitude=source_altitude,
+        show_straight_lines=show_straight_lines and not horizontal_rays,
+        show_earth_surface=show_earth_surface,
+        duct_center=duct_center,
+        duct_width=duct_width,
+        xlim=xlim,
+        ylim=ylim,
+        impact_parameter_keys=horizontal_rays,
+        use_colorbar=horizontal_rays,
+    )
+    ax1.set_title("Ray Trajectories")
+
+    # Top-right: Refractive index profile
+    max_alt = ylim[1] if ylim else 100.0
+    plot_refractive_index_profile(
+        ax=ax2, atmosphere=atmosphere, max_altitude_km=max_alt
+    )
+
+    # Bottom: Angular deviation
+    if horizontal_rays:
+        plot_horizontal_ray_deviation(
+            ax3,
+            trajectories,
+            xlim=xlim,
+            use_colorbar=True,
+        )
+        # Add shared colorbar
+        if sm is not None:
+            cbar = fig.colorbar(sm, cax=cax)
+            cbar.set_label("Impact parameter (km)", fontsize=10)
+    else:
+        plot_trajectory_offset(
+            ax3,
+            trajectories,
+            source_altitude=source_altitude,
+            xlim=xlim,
+        )
+
+    fig.suptitle(title, fontsize=14, fontweight="bold")
+    plt.tight_layout()
+    return fig
+
+
+def plot_horizontal_ray_deviation(
+    ax: plt.Axes,
+    trajectories: dict[float, NDArray],
+    cmap: str = "turbo",
+    linewidth: float = 1.5,
+    xlim: tuple[float, float] | None = None,
+    use_colorbar: bool = False,
+) -> None:
+    """
+    Plot angular deviation from initial direction for rays with varying impact parameters.
+
+    For rays that start at different altitudes (impact parameters), this shows how
+    each ray bends relative to its initial direction as it propagates through the
+    atmosphere. Atmospheric refraction typically bends rays downward.
+
+    Parameters
+    ----------
+    ax : matplotlib.axes.Axes
+        Axes to plot on.
+    trajectories : dict
+        Dictionary mapping impact parameter (meters) to trajectory array.
+        Each trajectory is shape (N, 2) with columns [x, z] in meters.
+    cmap : str
+        Colormap for ray colors.
+    linewidth : float
+        Line width for curves.
+    xlim : tuple, optional
+        X-axis limits in km.
+    use_colorbar : bool
+        If True, skip legend (colorbar added externally).
+
+    Examples
+    --------
+    >>> fig, ax = plt.subplots()
+    >>> plot_horizontal_ray_deviation(ax, trajectories)
+    >>> plt.show()
+    """
+    colors = plt.get_cmap(cmap)(np.linspace(0.1, 0.9, len(trajectories)))
+
+    for (impact_param_m, traj), color in zip(trajectories.items(), colors):
+        x_m = traj[:, 0]
+
+        # Compute direction angle at each point from finite differences
+        dx = np.diff(traj[:, 0])
+        dz = np.diff(traj[:, 1])
+
+        # Direction angle in degrees (angle above horizontal)
+        segment_angles = np.degrees(np.arctan2(dz, dx))
+
+        # Get initial direction from first segment
+        initial_angle = segment_angles[0]
+
+        # Build full array: deviation = current angle - initial angle
+        angular_deviation = np.zeros(len(traj))
+        angular_deviation[0] = 0.0  # No deviation at start
+        angular_deviation[1:] = segment_angles - initial_angle
+
+        # Convert impact parameter to km for legend
+        impact_km = impact_param_m / 1000
+        label = None if use_colorbar else f"{impact_km:.1f} km"
+
+        ax.plot(
+            x_m / 1000,
+            angular_deviation,
+            color=color,
+            linewidth=linewidth,
+            label=label,
+        )
+
+    ax.axhline(y=0, color="gray", linestyle="--", linewidth=1, alpha=0.5)
+    ax.set_xlabel("Horizontal distance (km)")
+    ax.set_ylabel("Angular deviation from initial direction (°)")
+    ax.set_title("Angular Deviation from Initial Ray Direction")
+    if not use_colorbar:
+        ax.legend(title="Impact param", loc="best", fontsize=8, ncol=2)
+    ax.grid(True, alpha=0.3)
+
+    if xlim:
+        ax.set_xlim(xlim)
+
+
+def save_atmospheric_figure(
+    fig: plt.Figure,
+    filename: str | Path,
+    output_dir: Path | None = None,
+    dpi: int = 150,
+) -> Path:
+    """
+    Save an atmospheric refraction figure.
+
+    Parameters
+    ----------
+    fig : matplotlib.figure.Figure
+        Figure to save.
+    filename : str or Path
+        Output filename.
+    output_dir : Path, optional
+        Output directory. Created if doesn't exist.
+    dpi : int
+        Resolution in dots per inch.
+
+    Returns
+    -------
+    output_path : Path
+        Full path to saved file.
+    """
+    if output_dir is not None:
+        output_dir = Path(output_dir)
+        output_dir.mkdir(exist_ok=True)
+        output_path = output_dir / filename
+    else:
+        output_path = Path(filename)
+
+    fig.savefig(output_path, dpi=dpi)
+    return output_path
+
+
+__all__ = [
+    "plot_ray_trajectories",
+    "plot_refractive_index_profile",
+    "plot_trajectory_offset",
+    "plot_horizontal_ray_deviation",
+    "create_atmospheric_refraction_figure",
+    "create_atmospheric_refraction_figure_with_offset",
+    "plot_trajectory_comparison",
+    "save_atmospheric_figure",
+]