Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
42 changes: 38 additions & 4 deletions rocketpy/plots/environment_plots.py
Original file line number Diff line number Diff line change
Expand Up @@ -33,6 +33,30 @@ def __init__(self, environment):
self.grid = np.linspace(environment.elevation, environment.max_expected_height)
self.environment = environment

def _break_direction_wraparound(self, directions, altitudes):
"""Inserts NaN into direction and altitude arrays at 0°/360° wraparound
points so matplotlib does not draw a horizontal line across the plot.

Parameters
----------
directions : numpy.ndarray
Wind direction values in degrees, dtype float.
altitudes : numpy.ndarray
Altitude values corresponding to each direction, dtype float.

Returns
-------
directions : numpy.ndarray
Direction array with NaN inserted at wraparound points.
altitudes : numpy.ndarray
Altitude array with NaN inserted at wraparound points.
"""
WRAP_THRESHOLD = 180 # degrees; half the full circle
wrap_indices = np.where(np.abs(np.diff(directions)) > WRAP_THRESHOLD)[0] + 1
directions = np.insert(directions, wrap_indices, np.nan)
altitudes = np.insert(altitudes, wrap_indices, np.nan)
return directions, altitudes

def __wind(self, ax):
"""Adds wind speed and wind direction graphs to the same axis.

Expand All @@ -55,9 +79,14 @@ def __wind(self, ax):
ax.set_xlabel("Wind Speed (m/s)", color="#ff7f0e")
ax.tick_params("x", colors="#ff7f0e")
axup = ax.twiny()
directions = np.array(
[self.environment.wind_direction(i) for i in self.grid], dtype=float
)
altitudes = np.array(self.grid, dtype=float)
directions, altitudes = self._break_direction_wraparound(directions, altitudes)
axup.plot(
[self.environment.wind_direction(i) for i in self.grid],
self.grid,
directions,
altitudes,
color="#1f77b4",
label="Wind Direction",
)
Expand Down Expand Up @@ -311,9 +340,14 @@ def ensemble_member_comparison(self, *, filename=None):
ax8 = plt.subplot(324)
for i in range(self.environment.num_ensemble_members):
self.environment.select_ensemble_member(i)
dirs = np.array(
[self.environment.wind_direction(j) for j in self.grid], dtype=float
)
alts = np.array(self.grid, dtype=float)
dirs, alts = self._break_direction_wraparound(dirs, alts)
ax8.plot(
[self.environment.wind_direction(i) for i in self.grid],
self.grid,
dirs,
alts,
label=i,
)
ax8.set_ylabel("Height Above Sea Level (m)")
Expand Down
53 changes: 53 additions & 0 deletions tests/integration/environment/test_environment.py
Original file line number Diff line number Diff line change
Expand Up @@ -92,6 +92,59 @@
assert example_plain_env.prints.print_earth_details() is None


@patch("matplotlib.pyplot.show")
def test_wind_plots_wrapping_direction(mock_show, example_plain_env): # pylint: disable=unused-argument
"""Tests that wind direction plots handle 360°→0° wraparound without
drawing a horizontal line across the graph.

Parameters
----------
mock_show : mock
Mock object to replace matplotlib.pyplot.show() method.
example_plain_env : rocketpy.Environment
Example environment object to be tested.
"""
# Set a custom atmosphere where wind direction wraps from ~350° to ~10°
# across the altitude range by choosing wind_u and wind_v to create a
# direction near 350° at low altitude and ~10° at higher altitude.
# wind_direction = (180 + atan2(wind_u, wind_v)) % 360
# For direction ~350°: need atan2(wind_u, wind_v) ≈ 170° → wind_u>0, wind_v<0
# For direction ~10°: need atan2(wind_u, wind_v) ≈ -170° → wind_u<0, wind_v<0
example_plain_env.set_atmospheric_model(
type="custom_atmosphere",
pressure=None,
temperature=300,
wind_u=[(0, 1), (5000, -1)], # changes sign across altitude
wind_v=[(0, -6), (5000, -6)], # stays negative → heading near 350°/10°
)
# Verify that the wind direction actually wraps through 0°/360° in this
# atmosphere so the test exercises the wraparound code path.
low_dir = example_plain_env.wind_direction(0)
high_dir = example_plain_env.wind_direction(5000)
assert abs(low_dir - high_dir) > 180, (
"Test setup error: wind direction should cross 0°/360° boundary"
)
# Verify that the helper inserts NaN breaks into the direction and altitude
# arrays at the wraparound point, which is the core of the fix.
directions = np.array(

Check failure on line 129 in tests/integration/environment/test_environment.py

View workflow job for this annotation

GitHub Actions / lint (3.10)

ruff (F821)

tests/integration/environment/test_environment.py:129:18: F821 Undefined name `np`
[example_plain_env.wind_direction(i) for i in example_plain_env.plots.grid],
dtype=float,
)
altitudes = np.array(example_plain_env.plots.grid, dtype=float)

Check failure on line 133 in tests/integration/environment/test_environment.py

View workflow job for this annotation

GitHub Actions / lint (3.10)

ruff (F821)

tests/integration/environment/test_environment.py:133:17: F821 Undefined name `np`
directions_broken, altitudes_broken = (
example_plain_env.plots._break_direction_wraparound(directions, altitudes)
)
assert np.any(np.isnan(directions_broken)), (

Check failure on line 137 in tests/integration/environment/test_environment.py

View workflow job for this annotation

GitHub Actions / lint (3.10)

ruff (F821)

tests/integration/environment/test_environment.py:137:19: F821 Undefined name `np`

Check failure on line 137 in tests/integration/environment/test_environment.py

View workflow job for this annotation

GitHub Actions / lint (3.10)

ruff (F821)

tests/integration/environment/test_environment.py:137:12: F821 Undefined name `np`
"Expected NaN breaks in direction array at 0°/360° wraparound"
)
assert np.any(np.isnan(altitudes_broken)), (

Check failure on line 140 in tests/integration/environment/test_environment.py

View workflow job for this annotation

GitHub Actions / lint (3.10)

ruff (F821)

tests/integration/environment/test_environment.py:140:19: F821 Undefined name `np`

Check failure on line 140 in tests/integration/environment/test_environment.py

View workflow job for this annotation

GitHub Actions / lint (3.10)

ruff (F821)

tests/integration/environment/test_environment.py:140:12: F821 Undefined name `np`
"Expected NaN breaks in altitude array at 0°/360° wraparound"
)
# Verify info() and atmospheric_model() plots complete without error
assert example_plain_env.info() is None
assert example_plain_env.plots.atmospheric_model() is None


@pytest.mark.parametrize(
"model_name",
[
Expand Down
Loading