Coverage for moptipy / examples / jssp / plot_gantt_chart.py: 86%
124 statements
« prev ^ index » next coverage.py v7.12.0, created at 2025-11-24 08:49 +0000
1"""Plot a Gantt chart into one figure."""
2from typing import Callable, Final, Iterable
4from matplotlib.artist import Artist # type: ignore
5from matplotlib.axes import Axes # type: ignore
6from matplotlib.figure import Figure # type: ignore
7from matplotlib.lines import Line2D # type: ignore
8from matplotlib.patches import Rectangle # type: ignore
9from matplotlib.text import Text # type: ignore
10from matplotlib.ticker import MaxNLocator # type: ignore
11from pycommons.types import type_error
13import moptipy.utils.plot_defaults as pd
14import moptipy.utils.plot_utils as pu
15from moptipy.evaluation.axis_ranger import AxisRanger
16from moptipy.examples.jssp.gantt import Gantt
17from moptipy.examples.jssp.makespan import makespan
18from moptipy.utils.lang import Lang
21def marker_lb(x: Gantt) -> tuple[str, int | float]:
22 """
23 Compute the marker for the lower bound.
25 :param x: the Gantt chart
26 :return: the lower bound marker
27 """
28 return Lang.current()["lower_bound_short"], \
29 x.instance.makespan_lower_bound
32def marker_makespan(x: Gantt) -> tuple[str, int | float]:
33 """
34 Compute the marker for the makespan.
36 :param x: the Gantt chart
37 :return: the makespan marker
38 """
39 return Lang.current()["makespan"], makespan(x)
42#: the color for markers at the left end
43__LEFT_END_MARK: Final[tuple[float, float, float]] = (0.95, 0.02, 0.02)
44#: the color for markers at the right end
45__RIGHT_END_MARK: Final[tuple[float, float, float]] = (0.02, 0.02, 0.95)
46#: the color for markers in the middle
47__MIDDLE_MARK: Final[tuple[float, float, float]] = pd.COLOR_BLACK
50def plot_gantt_chart(
51 gantt: Gantt | str,
52 figure: Axes | Figure,
53 markers: Iterable[tuple[str, int | float] | Callable[
54 [Gantt], tuple[str, int | float]]] | None = (marker_lb,),
55 x_axis: AxisRanger | Callable[[Gantt], AxisRanger] =
56 lambda _: AxisRanger(chosen_min=0),
57 importance_to_line_width_func: Callable[[int], float] =
58 pd.importance_to_line_width,
59 importance_to_font_size_func: Callable[[int], float] =
60 pd.importance_to_font_size,
61 info: str | Callable[[Gantt], str] | None =
62 lambda gantt: Lang.current().format_str("gantt_info", gantt=gantt),
63 x_grid: bool = False,
64 y_grid: bool = False,
65 x_label: str | Callable[[Gantt], str] | None =
66 Lang.translate_call("time"),
67 x_label_inside: bool = True,
68 x_label_location: float = 1.0,
69 y_label: str | Callable[[Gantt], str] | None =
70 Lang.translate_call("machine"),
71 y_label_inside: bool = True,
72 y_label_location: float = 0.5) -> Axes:
73 """
74 Plot a Gantt chart.
76 :param gantt: the gantt chart or a path to a file to load it from
77 :param figure: the figure
78 :param markers: a set of markers
79 :param x_axis: the ranger for the x-axis
80 :param info: the optional info header
81 :param importance_to_font_size_func: the function converting
82 importance values to font sizes
83 :param importance_to_line_width_func: the function converting
84 importance values to line widths
85 :param x_grid: should we have a grid along the x-axis?
86 :param y_grid: should we have a grid along the y-axis?
87 :param x_label: a callable returning the label for
88 the x-axis, a label string, or `None` if no label should be put
89 :param x_label_inside: put the x-axis label inside the plot (so that
90 it does not consume additional vertical space)
91 :param x_label_location: the location of the x-label
92 :param y_label: a callable returning the label for
93 the y-axis, a label string, or `None` if no label should be put
94 :param y_label_inside: put the y-axis label inside the plot (so that
95 it does not consume additional horizontal space)
96 :param y_label_location: the location of the y-label
97 :returns: the axes object to allow you to add further plot elements
98 """
99 if isinstance(gantt, str):
100 gantt = Gantt.from_log(gantt)
101 if not isinstance(gantt, Gantt):
102 raise type_error(gantt, "gantt", (Gantt, str))
103 axes: Final[Axes] = pu.get_axes(figure)
105 # grab the data
106 jobs: Final[int] = gantt.instance.jobs
107 machines: Final[int] = gantt.instance.machines
109 # Set up the x-axis range.
110 if callable(x_axis):
111 x_axis = x_axis(gantt)
112 if not isinstance(x_axis, AxisRanger):
113 raise type_error(x_axis, "x_axis", AxisRanger)
115 # Compute all the marks
116 marks: dict[int | float, str] = {}
117 if markers is not None:
118 if not isinstance(markers, Iterable):
119 raise type_error(markers, "markers", Iterable)
120 for usemarker in markers:
121 marker = usemarker(gantt) if callable(usemarker) else usemarker
122 if not marker:
123 continue
124 if isinstance(marker, tuple):
125 name, val = marker
126 if (not name) or (not val):
127 continue
128 else:
129 raise type_error(marker, "marker", tuple, True)
130 if not isinstance(name, str):
131 raise type_error(name, "marker name", str)
132 if not isinstance(val, int | float):
133 raise type_error(val, "marker", (int, float))
134 if val in marks:
135 marks[val] = f"{marks[val]}/{name}"
136 else:
137 marks[val] = name
138 x_axis.register_value(val)
140 # Add x-axis data
141 x_axis.register_array(gantt[:, 0, 1].flatten()) # register start times
142 x_axis.register_array(gantt[:, -1, 2].flatten()) # register end times
143 x_axis.apply(axes, "x")
144 xmin, xmax = axes.get_xlim()
146 # Set up the y-axis range.
147 height: Final[float] = 0.7
148 bar_ofs: Final[float] = height / 2
149 y_min: Final[float] = -((7 * bar_ofs) / 6)
150 y_max: Final[float] = (machines - 1) + ((7 * bar_ofs) / 6)
152 axes.set_ylim(y_min, y_max)
153 axes.set_ybound(y_min, y_max)
154 axes.yaxis.set_major_locator(MaxNLocator(nbins="auto",
155 integer=True))
157 # get the color and font styles
158 colors: Final[tuple] = pd.distinct_colors(jobs)
159 font_size: Final[float] = importance_to_font_size_func(-1)
161 # get the transforms needed to obtain text dimensions
162 rend: Final = pu.get_renderer(axes)
163 inv: Final = axes.transData.inverted()
165 # draw the grid
166 if x_grid or y_grid:
167 grid_lwd = importance_to_line_width_func(-1)
168 if x_grid:
169 axes.grid(axis="x", color=pd.GRID_COLOR, linewidth=grid_lwd)
170 if y_grid:
171 axes.grid(axis="y", color=pd.GRID_COLOR, linewidth=grid_lwd)
173 z_order: int = 0
175 # print the marker lines
176 for val in marks:
177 axes.add_artist(Line2D(xdata=(val, val),
178 ydata=(y_min, y_max),
179 color=__LEFT_END_MARK if val <= xmin
180 else __RIGHT_END_MARK if val >= xmax
181 else __MIDDLE_MARK,
182 linewidth=2.0,
183 zorder=z_order))
184 z_order += 1
186 # plot the jobs
187 for machine in range(machines):
188 for jobi in range(jobs):
189 job, x_start, x_end = gantt[machine, jobi, :]
191 if x_end <= x_start: # skip operations that take 0 time
192 continue
194 background = colors[job]
195 foreground = pd.text_color_for_background(colors[job])
196 jobstr = str(job)
198 # first plot the colored rectangle
199 y_start = machine - bar_ofs
201 axes.add_artist(Rectangle(
202 xy=(x_start, y_start),
203 width=(x_end - x_start),
204 height=height,
205 color=background,
206 linewidth=0,
207 zorder=z_order))
208 z_order += 1
210 # Now we insert the job IDs, which is a bit tricky:
211 # First, the rectangle may be too small to hold the text.
212 # So we need to get the text bounding box size to compare it with
213 # the rectangle's size.
214 # If that fits, we can print the text.
215 # Second, the printed text tends to be slightly off vertically,
216 # even if we try to center it vertically. This is because fonts
217 # can extend below their baseline which seems to be considered
218 # during vertical alignment although job IDs are numbers and that
219 # does not apply here. Therefore, we try to re-adjust the boxes
220 # in a very, very crude way.
221 xp: float = 0.5 * (x_start + x_end)
222 yp: float | int = machine
224 # Get the size of the text using a temporary text
225 # that gets immediately deleted again.
226 tmp: Text = axes.text(x=xp, y=yp, s=jobstr,
227 fontsize=font_size,
228 color=foreground,
229 horizontalalignment="center",
230 verticalalignment="baseline")
231 bb_bl = inv.transform_bbox(tmp.get_window_extent(
232 renderer=rend))
233 Artist.set_visible(tmp, False)
234 Artist.remove(tmp)
235 del tmp
237 if (bb_bl.width < 0.97 * (x_end - x_start)) and \
238 (bb_bl.height < (0.97 * height)):
239 # OK, there is enough space. Let's re-compute the y
240 # offset to do proper alignment using another temporary
241 # text.
242 tmp = axes.text(x=xp, y=yp, s=jobstr,
243 fontsize=font_size,
244 color=foreground,
245 horizontalalignment="center",
246 verticalalignment="bottom")
247 bb_bt = inv.transform_bbox(tmp.get_window_extent(
248 renderer=rend))
249 Artist.set_visible(tmp, False)
250 Artist.remove(tmp)
251 del tmp
253 # Now we can really print the actual text with a more or less
254 # nice vertical alignment.
255 adj = bb_bl.y0 - bb_bt.y0
256 if adj < 0:
257 yp += adj / 3
259 axes.text(x=xp, y=yp, s=jobstr,
260 fontsize=font_size,
261 color=foreground,
262 horizontalalignment="center",
263 verticalalignment="center",
264 zorder=z_order)
265 z_order += 1
267 # print the marker labels
268 bbox = {"boxstyle": "round",
269 "color": "white",
270 "fill": True,
271 "linewidth": 0,
272 "alpha": 0.9}
273 y_mark: Final[float] = -0.1 # machines - 1 + (0.9 * bar_ofs) for top
274 for val, name in marks.items():
275 axes.annotate(text=f"{name}={val}",
276 xy=(val, y_mark),
277 xytext=(-4, -4),
278 verticalalignment="bottom",
279 horizontalalignment="right",
280 xycoords="data",
281 textcoords="offset points",
282 fontsize=font_size,
283 color=__LEFT_END_MARK if val <= xmin
284 else __RIGHT_END_MARK if val >= xmax
285 else __MIDDLE_MARK,
286 rotation=90,
287 bbox=bbox,
288 zorder=z_order)
289 z_order += 1
291 info_font_size: Final[float] = importance_to_font_size_func(0)
292 pu.label_axes(axes=axes,
293 x_label=x_label(gantt) if callable(x_label) else x_label,
294 x_label_inside=x_label_inside,
295 x_label_location=x_label_location,
296 y_label=y_label(gantt) if callable(y_label) else y_label,
297 y_label_inside=y_label_inside,
298 y_label_location=y_label_location,
299 font_size=info_font_size,
300 z_order=z_order)
301 z_order += 1
303 if callable(info):
304 info = info(gantt)
305 if info is not None:
306 if not isinstance(info, str):
307 raise type_error(info, "info", str)
308 pu.label_box(axes=axes,
309 text=info,
310 x=0.5,
311 y=1,
312 font_size=importance_to_font_size_func(1),
313 may_rotate_text=False,
314 z_order=z_order)
315 return axes