Skip to content
1 change: 1 addition & 0 deletions CHANGELOG.md
Original file line number Diff line number Diff line change
Expand Up @@ -73,6 +73,7 @@ Attention: The newest changes should be on top -->

### Fixed

- BUG: Add wraparound logic for wind direction in environment plots [#939](https://github.com/RocketPy-Team/RocketPy/pull/939)
- BUG: Restore `Rocket.power_off_drag` and `Rocket.power_on_drag` as `Function` objects while preserving raw inputs in `power_off_drag_input` and `power_on_drag_input` [#941](https://github.com/RocketPy-Team/RocketPy/pull/941)
- BUG: Add explicit timeouts to ThrustCurve API requests [#935](https://github.com/RocketPy-Team/RocketPy/pull/935)
- BUG: Fix hard-coded radius value for parachute added mass calculation [#889](https://github.com/RocketPy-Team/RocketPy/pull/889)
Expand Down
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