Fix linear bounds rendering in plots

Author: mickaelbegon

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]:
         """

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