# 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
[docs]
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
[docs]
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),
)
[docs]
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
[docs]
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)
[docs]
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)
__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",
]