BUG: Add wraparound logic for wind direction in environment plots (#939)

* chore: added personal toolkit files

* update branch name in workflow

* chore: update toolkit files

* Fix: add wraparound logic for wind direction and related tests

* style: fix ruff formatting

* Remove unused import

Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com>

* refactor: move repetitive logic into helper method

* fix: update test logic in test_environment

* add changelog entry

---------

Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com>
Co-authored-by: Gui-FernandesBR <63590233+Gui-FernandesBR@users.noreply.github.com>

Author: 3 people

CHANGELOG.md

@@ -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)

rocketpy/plots/environment_plots.py

@@ -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.
 
@@ -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",
         )
@@ -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)")

tests/integration/environment/test_environment.py

@@ -92,6 +92,59 @@ def test_standard_atmosphere(mock_show, example_plain_env):  # pylint: disable=u
     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(
+        [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)
+    directions_broken, altitudes_broken = (
+        example_plain_env.plots._break_direction_wraparound(directions, altitudes)
+    )
+    assert np.any(np.isnan(directions_broken)), (
+        "Expected NaN breaks in direction array at 0°/360° wraparound"
+    )
+    assert np.any(np.isnan(altitudes_broken)), (
+        "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",
     [