diff --git a/pyproject.toml b/pyproject.toml index 456441f33..9133029e8 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -77,7 +77,16 @@ animation = [ "imageio-ffmpeg>=0.5" ] -all = ["rocketpy[env-analysis]", "rocketpy[monte-carlo]", "rocketpy[animation]"] +maps = [ + "folium>=0.14", +] + +all = [ + "rocketpy[env-analysis]", + "rocketpy[monte-carlo]", + "rocketpy[animation]", + "rocketpy[maps]", +] [tool.coverage.report] diff --git a/rocketpy/plots/flight_plots.py b/rocketpy/plots/flight_plots.py index 6dd5e8802..df5713565 100644 --- a/rocketpy/plots/flight_plots.py +++ b/rocketpy/plots/flight_plots.py @@ -144,6 +144,76 @@ def trajectory_3d(self, *, filename=None): # pylint: disable=too-many-statement ax1.set_box_aspect(None, zoom=0.95) # 95% for label adjustment show_or_save_plot(filename) + def trajectory_on_map(self, *, filename=None): + """Create an interactive Folium map of the flight trajectory. + + Draws the ground-track path from ``flight.latitude`` / + ``flight.longitude`` and marks the launch and landing sites. + Requires the optional ``folium`` dependency + (``pip install folium`` or ``pip install rocketpy[maps]``). + + Parameters + ---------- + filename : str | None, optional + Path to save the map as an HTML file. If None, the map is not + written to disk. Default is None. + + Returns + ------- + folium.Map + The interactive map object. In Jupyter, displaying the return + value renders the map. + """ + folium = import_optional_dependency("folium") + flight = self.flight + + latitudes = np.asarray(flight.latitude[:, 1], dtype=float) + longitudes = np.asarray(flight.longitude[:, 1], dtype=float) + path = list(zip(latitudes.tolist(), longitudes.tolist())) + if not path: + raise ValueError("Flight has no latitude/longitude samples to plot.") + + launch = path[0] + landing = path[-1] + center = [ + float(0.5 * (launch[0] + landing[0])), + float(0.5 * (launch[1] + landing[1])), + ] + + flight_map = folium.Map(location=center, zoom_start=13) + folium.PolyLine( + locations=path, + color="#1f77b4", + weight=3, + opacity=0.85, + tooltip="Flight trajectory", + ).add_to(flight_map) + folium.Marker( + location=launch, + popup="Launch", + tooltip="Launch", + icon=folium.Icon(color="green"), + ).add_to(flight_map) + folium.Marker( + location=landing, + popup="Landing", + tooltip="Landing", + icon=folium.Icon(color="red"), + ).add_to(flight_map) + + south = float(np.min(latitudes)) + north = float(np.max(latitudes)) + west = float(np.min(longitudes)) + east = float(np.max(longitudes)) + if abs(north - south) > 1e-12 or abs(east - west) > 1e-12: + flight_map.fit_bounds([[south, west], [north, east]]) + + if filename is not None: + flight_map.save(filename) + logger.info("File %s saved with success!", filename) + + return flight_map + def _resolve_animation_model_path(self, file_name): """Resolve model path, defaulting to the built-in STL when omitted.""" if file_name is not None: diff --git a/tests/unit/test_flight_trajectory_map.py b/tests/unit/test_flight_trajectory_map.py new file mode 100644 index 000000000..33eb687de --- /dev/null +++ b/tests/unit/test_flight_trajectory_map.py @@ -0,0 +1,84 @@ +"""Tests for optional Folium flight trajectory maps.""" + +from unittest.mock import MagicMock, patch + +import pytest + +from rocketpy.plots.flight_plots import _FlightPlots + + +def test_trajectory_on_map_requires_folium(flight_calisto_robust): + """Missing folium should raise a clear ImportError via optional import.""" + with patch( + "rocketpy.plots.flight_plots.import_optional_dependency", + side_effect=ImportError( + "folium is an optional dependency and is not installed.\n" + "\t\tUse 'pip install folium' to install it or " + "'pip install rocketpy[all]' to install all optional dependencies." + ), + ): + with pytest.raises(ImportError, match="folium"): + flight_calisto_robust.plots.trajectory_on_map() + + +def test_trajectory_on_map_builds_map_with_mocked_folium(flight_calisto_robust): + """Map construction should add a path and launch/landing markers.""" + mock_folium = MagicMock() + mock_map = MagicMock() + mock_folium.Map.return_value = mock_map + mock_polyline = MagicMock() + mock_folium.PolyLine.return_value = mock_polyline + mock_marker = MagicMock() + mock_folium.Marker.return_value = mock_marker + + with patch( + "rocketpy.plots.flight_plots.import_optional_dependency", + return_value=mock_folium, + ): + result = flight_calisto_robust.plots.trajectory_on_map() + + assert result is mock_map + mock_folium.Map.assert_called_once() + mock_folium.PolyLine.assert_called_once() + assert mock_folium.Marker.call_count == 2 + mock_polyline.add_to.assert_called_once_with(mock_map) + assert mock_marker.add_to.call_count == 2 + mock_map.save.assert_not_called() + + +def test_trajectory_on_map_saves_html_with_mocked_folium( + flight_calisto_robust, tmp_path +): + """filename= should call Map.save with the requested path.""" + mock_folium = MagicMock() + mock_map = MagicMock() + mock_folium.Map.return_value = mock_map + mock_folium.PolyLine.return_value = MagicMock() + mock_folium.Marker.return_value = MagicMock() + out = tmp_path / "trajectory.html" + + with patch( + "rocketpy.plots.flight_plots.import_optional_dependency", + return_value=mock_folium, + ): + result = flight_calisto_robust.plots.trajectory_on_map(filename=str(out)) + + assert result is mock_map + mock_map.save.assert_called_once_with(str(out)) + + +def test_trajectory_on_map_creates_html_file(flight_calisto_robust, tmp_path): + """With folium installed, save an HTML map and return a Map instance.""" + folium = pytest.importorskip("folium") + + out = tmp_path / "trajectory.html" + result = flight_calisto_robust.plots.trajectory_on_map(filename=str(out)) + + assert isinstance(result, folium.Map) + assert isinstance(flight_calisto_robust.plots, _FlightPlots) + assert out.is_file() + assert out.stat().st_size > 0 + html = out.read_text(encoding="utf-8") + assert "leaflet" in html.lower() + assert "Launch" in html + assert "Landing" in html