Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
9 changes: 5 additions & 4 deletions bioptim/gui/plot.py
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand All @@ -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]:
"""
Expand Down
29 changes: 29 additions & 0 deletions tests/shard1/test_plot.py
Original file line number Diff line number Diff line change
@@ -1,4 +1,7 @@
from types import SimpleNamespace

import numpy as np
import matplotlib.pyplot as plt
from casadi import DM

from bioptim import (
Expand All @@ -11,6 +14,8 @@
InitialGuessList,
PlotType,
CustomPlot,
Bounds,
InterpolationType,
)
from bioptim.gui.plot import DEFAULT_COLORS, PlotOcp

Expand Down Expand Up @@ -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
Expand Down
Loading