diff --git a/bioptim/gui/plot.py b/bioptim/gui/plot.py index bfacd6fea..03a52bb5f 100644 --- a/bioptim/gui/plot.py +++ b/bioptim/gui/plot.py @@ -705,8 +705,6 @@ def _add_bounds_to_plot( else self.t ) - # TODO: introduce repeat for the COLLOCATIONS min/max_bounds only for states graphs. - # For now the plots in COLLOCATIONS with LINEAR are not giving the right values nlp.plot[variable].bounds.check_and_adjust_dimensions(n_elements=len(mapping_to_first_index), n_shooting=ns) idx = mapping_to_first_index.index(ctr) @@ -717,8 +715,11 @@ def _add_bounds_to_plot( bounds_min = np.concatenate((bounds_min, [bounds_min[-1]])) bounds_max = np.concatenate((bounds_max, [bounds_max[-1]])) - self.plots_bounds.append([ax.step(t[i], bounds_min, where="post", **self.plot_options["bounds"]), i]) - self.plots_bounds.append([ax.step(t[i], bounds_max, where="post", **self.plot_options["bounds"]), i]) + is_linear = nlp.plot[variable].bounds.type == InterpolationType.LINEAR + plot_function = ax.plot if is_linear else ax.step + plot_options = {} if is_linear else {"where": "post"} + self.plots_bounds.append([plot_function(t[i], bounds_min, **plot_options, **self.plot_options["bounds"]), i]) + self.plots_bounds.append([plot_function(t[i], bounds_max, **plot_options, **self.plot_options["bounds"]), i]) def _add_new_axis(self, variable: Str, nb: Int, n_rows: Int, n_cols: Int) -> np.ndarray[plt.Axes]: """ diff --git a/tests/shard1/test_plot.py b/tests/shard1/test_plot.py index 775436958..6b1c22f66 100644 --- a/tests/shard1/test_plot.py +++ b/tests/shard1/test_plot.py @@ -1,4 +1,7 @@ +from types import SimpleNamespace + import numpy as np +import matplotlib.pyplot as plt from casadi import DM from bioptim import ( @@ -11,6 +14,8 @@ InitialGuessList, PlotType, CustomPlot, + Bounds, + InterpolationType, ) from bioptim.gui.plot import DEFAULT_COLORS, PlotOcp @@ -91,6 +96,30 @@ def test_default_colors(): assert PlotType.POINT in DEFAULT_COLORS +def test_linear_bounds_are_plotted_with_linear_interpolation(): + """Linear bounds must not be rendered as steps on integrated state plots.""" + plot_ocp = PlotOcp.__new__(PlotOcp) + plot_ocp.t = [np.linspace(0, 1, 3)] + plot_ocp.plots_bounds = [] + plot_ocp.plot_options = {"bounds": {"color": "k"}} + + bounds = Bounds( + "q", + min_bound=np.array([[0.0, 1.0]]), + max_bound=np.array([[2.0, 4.0]]), + interpolation=InterpolationType.LINEAR, + ) + nlp = SimpleNamespace(ns=2, plot={"q_states": SimpleNamespace(bounds=bounds)}) + figure, axis = plt.subplots() + + plot_ocp._add_bounds_to_plot(0, nlp, "q_states", 0, axis, [0]) + + np.testing.assert_allclose(plot_ocp.plots_bounds[0][0][0].get_ydata(), [0.0, 0.5, 1.0]) + np.testing.assert_allclose(plot_ocp.plots_bounds[1][0][0].get_ydata(), [2.0, 3.0, 4.0]) + assert plot_ocp.plots_bounds[0][0][0].get_drawstyle() == "default" + plt.close(figure) + + def test_plot_options(): """Test the plot options of PlotOcp""" from tests.utils import TestUtils