job-shop-lib 0.5.1__py3-none-any.whl → 1.0.0__py3-none-any.whl

job_shop_lib/visualization/gantt/_gantt_chart_video_and_gif_creation.py

@@ -0,0 +1,422 @@
+"""Module for creating a GIF or a video of the schedule being built."""
+
+import os
+import pathlib
+import shutil
+from typing import Sequence, Protocol, Optional, List, Any
+
+import imageio
+import matplotlib.pyplot as plt
+from matplotlib.figure import Figure
+import numpy as np
+
+from job_shop_lib import (
+    JobShopInstance,
+    Schedule,
+    Operation,
+    ScheduledOperation,
+)
+from job_shop_lib.exceptions import ValidationError
+from job_shop_lib.dispatching import (
+    Dispatcher,
+    HistoryObserver,
+)
+from job_shop_lib.dispatching.rules import DispatchingRuleSolver
+from job_shop_lib.visualization.gantt._plot_gantt_chart import plot_gantt_chart
+
+
+# This class serves as a more meaningful type hint than simply:
+# PlotFunction = Callable[
+#     [Schedule, int | None, list[Operation] | None, int | None], Figure
+# ]
+# That's why it doesn't have more methods or attributes. It is a protocol
+# for functions, not for classes.
+# pylint: disable=too-few-public-methods
+class PartialGanttChartPlotter(Protocol):
+    """A protocol for a function that plots an uncompleted Gantt chart
+    for a schedule.
+
+    This kind of functions are created using the
+    :func:`plot_gantt_chart_wrapper` function.
+    """
+
+    def __call__(
+        self,
+        schedule: Schedule,
+        makespan: Optional[int] = None,
+        available_operations: Optional[List[Operation]] = None,
+        current_time: Optional[int] = None,
+    ) -> Figure:
+        """Plots a Gantt chart for an unfinished schedule.
+
+        Args:
+            schedule:
+                The schedule to plot.
+            makespan:
+                The makespan of the schedule if known. Can be used to fix
+                the x-axis limits.
+            available_operations:
+                A list of available operations. If ``None``,
+                the available operations are not shown.
+            current_time:
+                The current time in the schedule. If provided, a red
+                vertical line is plotted at this time.
+        """
+
+
+def get_partial_gantt_chart_plotter(
+    title: Optional[str] = None,
+    cmap: str = "viridis",
+    show_available_operations: bool = False,
+    **kwargs: Any,
+) -> PartialGanttChartPlotter:
+    """Returns a function that plots a Gantt chart for an unfinished schedule.
+
+    Args:
+        title: The title of the Gantt chart.
+        cmap: The name of the colormap to use.
+        show_available_operations:
+            Whether to show the available operations in the Gantt chart.
+        **kwargs: Additional keyword arguments to pass to the
+            :func:`plot_gantt_chart` function.
+
+    Returns:
+        A function that plots a Gantt chart for a schedule. The function takes
+        the following arguments:
+
+        - schedule: The schedule to plot.
+        - makespan: The makespan of the schedule.
+        - available_operations: A list of available operations. If None,
+          the available operations are not shown.
+        - current_time: The current time in the schedule. If provided, a
+          red vertical line is plotted at this time.
+
+    """
+
+    def plot_function(
+        schedule: Schedule,
+        makespan: Optional[int] = None,
+        available_operations: Optional[List[Operation]] = None,
+        current_time: Optional[int] = None,
+    ) -> Figure:
+        fig, ax = plot_gantt_chart(
+            schedule, title=title, cmap_name=cmap, xlim=makespan, **kwargs
+        )
+
+        if show_available_operations and available_operations is not None:
+
+            operations_text = "\n".join(
+                str(operation) for operation in available_operations
+            )
+            text = f"Available operations:\n{operations_text}"
+            # Print the available operations at the bottom right corner
+            # of the Gantt chart
+            fig.text(
+                1.25,
+                0.05,
+                text,
+                ha="right",
+                va="bottom",
+                transform=ax.transAxes,
+                bbox={
+                    "facecolor": "white",
+                    "alpha": 0.5,
+                    "boxstyle": "round,pad=0.5",
+                },
+            )
+        if current_time is not None:
+            ax.axvline(current_time, color="red", linestyle="--")
+        return fig
+
+    return plot_function
+
+
+# Most of the arguments are optional with default values. There is no way to
+# reduce the number of arguments without losing functionality.
+# pylint: disable=too-many-arguments
+def create_gantt_chart_gif(
+    instance: JobShopInstance,
+    gif_path: Optional[str] = None,
+    solver: Optional[DispatchingRuleSolver] = None,
+    plot_function: Optional[PartialGanttChartPlotter] = None,
+    fps: int = 1,
+    remove_frames: bool = True,
+    frames_dir: Optional[str] = None,
+    plot_current_time: bool = True,
+    schedule_history: Optional[Sequence[ScheduledOperation]] = None,
+) -> None:
+    """Creates a GIF of the schedule being built.
+
+    Args:
+        instance:
+            The instance of the job shop problem to be scheduled.
+        gif_path:
+            The path to save the GIF file. It should end with ".gif". If not
+            provided, the name of the instance is used.
+        solver:
+            The dispatching rule solver to use. If not provided, the history
+            of scheduled operations should be provided.
+        plot_function:
+            A :class:`PlotFunction` that plots a Gantt chart for a schedule. It
+            should take a `Schedule` object and the makespan of the schedule as
+            input and return a `Figure` object. If not provided, a default
+            function is used.
+        fps:
+            The number of frames per second in the GIF.
+        remove_frames:
+            Whether to remove the frames after creating the GIF.
+        frames_dir:
+            The directory to save the frames in. If not provided,
+            ``gif_path.replace(".gif", "") + "_frames"`` is used.
+        plot_current_time:
+            Whether to plot a vertical line at the current time.
+        schedule_history:
+            A sequence of scheduled operations. If not provided, the solver
+            will be used to generate the history.
+    """
+    if gif_path is None:
+        gif_path = f"{instance.name}_gantt_chart.gif"
+
+    if plot_function is None:
+        plot_function = get_partial_gantt_chart_plotter()
+
+    if frames_dir is None:
+        # Use the name of the GIF file as the directory name
+        frames_dir = gif_path.replace(".gif", "") + "_frames"
+    path = pathlib.Path(frames_dir)
+    path.mkdir(exist_ok=True)
+    frames_dir = str(path)
+    create_gantt_chart_frames(
+        frames_dir,
+        instance,
+        solver,
+        plot_function,
+        plot_current_time,
+        schedule_history,
+    )
+    create_gif_from_frames(frames_dir, gif_path, fps)
+
+    if remove_frames:
+        shutil.rmtree(frames_dir)
+
+
+# Most of the arguments are optional with default values. There is no way to
+# reduce the number of arguments without losing functionality.
+# pylint: disable=too-many-arguments
+def create_gantt_chart_video(
+    instance: JobShopInstance,
+    video_path: Optional[str] = None,
+    solver: Optional[DispatchingRuleSolver] = None,
+    plot_function: Optional[PartialGanttChartPlotter] = None,
+    fps: int = 1,
+    remove_frames: bool = True,
+    frames_dir: Optional[str] = None,
+    plot_current_time: bool = True,
+    schedule_history: Optional[Sequence[ScheduledOperation]] = None,
+) -> None:
+    """Creates a video of the schedule being built.
+
+    Args:
+        instance:
+            The instance of the job shop problem to be scheduled.
+        video_path:
+            The path to save the video file.
+        solver:
+            The dispatching rule solver to use. If not provided, the history
+            of scheduled operations should be provided.
+        plot_function:
+            A function that plots a Gantt chart for a schedule. It
+            should take a :class:`Schedule` object and the makespan of the
+            schedule as input and return a ``Figure`` object. If not provided,
+            a default function is used.
+        fps:
+            The number of frames per second in the video.
+        remove_frames:
+            Whether to remove the frames after creating the video.
+        frames_dir:
+            The directory to save the frames in. If not provided,
+            ``name_without_the_extension + "_frames"`` is used.
+        plot_current_time:
+            Whether to plot a vertical line at the current time.
+        schedule_history:
+            A sequence of scheduled operations. If not provided, the solver
+            will be used to generate the history.
+    """
+    if video_path is None:
+        video_path = f"{instance.name}_gantt_chart.mp4"
+
+    if plot_function is None:
+        plot_function = get_partial_gantt_chart_plotter()
+
+    if frames_dir is None:
+        extension = video_path.split(".")[-1]
+        frames_dir = video_path.replace(f".{extension}", "") + "_frames"
+    path = pathlib.Path(frames_dir)
+    path.mkdir(exist_ok=True)
+    frames_dir = str(path)
+    create_gantt_chart_frames(
+        frames_dir,
+        instance,
+        solver,
+        plot_function,
+        plot_current_time,
+        schedule_history,
+    )
+    create_video_from_frames(frames_dir, video_path, fps)
+
+    if remove_frames:
+        shutil.rmtree(frames_dir)
+
+
+def create_gantt_chart_frames(
+    frames_dir: str,
+    instance: JobShopInstance,
+    solver: Optional[DispatchingRuleSolver],
+    plot_function: PartialGanttChartPlotter,
+    plot_current_time: bool = True,
+    schedule_history: Optional[Sequence[ScheduledOperation]] = None,
+) -> None:
+    """Creates frames of the Gantt chart for the schedule being built.
+
+    Args:
+        frames_dir:
+            The directory to save the frames in.
+        instance:
+            The instance of the job shop problem to be scheduled.
+        solver:
+            The dispatching rule solver to use. If not provided, the history
+            of scheduled operations should be provided.
+        plot_function:
+            A function that plots a Gantt chart for a schedule. It
+            should take a `Schedule` object and the makespan of the schedule as
+            input and return a `Figure` object.
+        plot_current_time:
+            Whether to plot a vertical line at the current time.
+        scheduled_history:
+            A sequence of scheduled operations. If not provided, the solver
+    """
+    if solver is not None and schedule_history is None:
+        dispatcher = Dispatcher(
+            instance, ready_operations_filter=solver.ready_operations_filter
+        )
+        history_tracker = HistoryObserver(dispatcher)
+        makespan = solver.solve(instance, dispatcher).makespan()
+        dispatcher.unsubscribe(history_tracker)
+        dispatcher.reset()
+        schedule_history = history_tracker.history
+    elif schedule_history is not None and solver is None:
+        dispatcher = Dispatcher(instance)
+        makespan = max(
+            scheduled_operation.end_time
+            for scheduled_operation in schedule_history
+        )
+    elif schedule_history is not None and solver is not None:
+        raise ValidationError(
+            "Only one of 'solver' and 'history' should be provided."
+        )
+    else:
+        raise ValidationError(
+            "Either 'solver' or 'history' should be provided."
+        )
+
+    for i, scheduled_operation in enumerate(schedule_history, start=1):
+        dispatcher.dispatch(
+            scheduled_operation.operation, scheduled_operation.machine_id
+        )
+        current_time = (
+            None if not plot_current_time else dispatcher.current_time()
+        )
+        fig = plot_function(
+            dispatcher.schedule,
+            makespan,
+            dispatcher.available_operations(),
+            current_time,
+        )
+        _save_frame(fig, frames_dir, i)
+
+
+def _save_frame(figure: Figure, frames_dir: str, number: int) -> None:
+    figure.savefig(f"{frames_dir}/frame_{number:02d}.png", bbox_inches="tight")
+    plt.close(figure)
+
+
+def create_gif_from_frames(
+    frames_dir: str, gif_path: str, fps: int, loop: int = 0
+) -> None:
+    """Creates a GIF or video from the frames in the given directory.
+
+    Args:
+        frames_dir:
+            The directory containing the frames to be used in the GIF.
+        gif_path:
+            The path to save the GIF file. It should end with ".gif".
+        fps:
+            The number of frames per second in the GIF.
+        loop:
+            The number of times the GIF should loop. Default is 0, which means
+            the GIF will loop indefinitely. If set to 1, the GIF will loop
+            once. Added in version 0.6.0.
+    """
+    images = _load_images(frames_dir)
+    imageio.mimsave(gif_path, images, fps=fps, loop=loop)
+
+
+def create_video_from_frames(
+    frames_dir: str, gif_path: str, fps: int, macro_block_size: int = 16
+) -> None:
+    """Creates a GIF or video from the frames in the given directory.
+
+    Args:
+        frames_dir:
+            The directory containing the frames to be used in the video.
+        gif_path:
+            The path to save the video.
+        fps:
+            The number of frames per second.
+    """
+    images = _load_images(frames_dir)
+    resized_images = [
+        resize_image_to_macro_block(image, macro_block_size=macro_block_size)
+        for image in images
+    ]
+    imageio.mimsave(
+        gif_path, resized_images, fps=fps  # type: ignore[arg-type]
+    )
+
+
+def resize_image_to_macro_block(
+    image: np.ndarray, macro_block_size: int = 16
+) -> np.ndarray:
+    """Resizes the image to ensure its dimensions are divisible by the macro
+    block size.
+
+    Args:
+        image (numpy.ndarray): The input image.
+        macro_block_size (int): The macro block size, typically 16.
+
+    Returns:
+        numpy.ndarray: The resized image.
+    """
+    height, width, channels = image.shape
+    new_height = (
+        (height + macro_block_size - 1) // macro_block_size * macro_block_size
+    )
+    new_width = (
+        (width + macro_block_size - 1) // macro_block_size * macro_block_size
+    )
+
+    if (new_height, new_width) != (height, width):
+        resized_image = np.zeros(
+            (new_height, new_width, channels), dtype=image.dtype
+        )
+        resized_image[:height, :width] = image
+        return resized_image
+    return image
+
+
+def _load_images(frames_dir: str) -> List:
+    frames = [
+        os.path.join(frames_dir, frame)
+        for frame in sorted(os.listdir(frames_dir))
+    ]
+    return [imageio.imread(frame) for frame in frames]

job_shop_lib/visualization/{gantt_chart.py → gantt/_plot_gantt_chart.py}

@@ -1,6 +1,6 @@
 """Module for plotting static Gantt charts for job shop schedules."""
 
-from typing import Optional
+from typing import Optional, List, Tuple, Dict
 
 from matplotlib.figure import Figure
 import matplotlib.pyplot as plt
@@ -9,27 +9,49 @@ from matplotlib.patches import Patch
 
 from job_shop_lib import Schedule, ScheduledOperation
 
-
 _BASE_Y_POSITION = 1
 _Y_POSITION_INCREMENT = 10
 
 
 def plot_gantt_chart(
     schedule: Schedule,
-    title: str | None = None,
+    title: Optional[str] = None,
     cmap_name: str = "viridis",
-    xlim: int | None = None,
+    xlim: Optional[int] = None,
     number_of_x_ticks: int = 15,
-) -> tuple[Figure, plt.Axes]:
+    job_labels: Optional[List[str]] = None,
+    machine_labels: Optional[List[str]] = None,
+    legend_title: str = "",
+    x_label: str = "Time units",
+    y_label: str = "Machines",
+) -> Tuple[Figure, plt.Axes]:
     """Plots a Gantt chart for the schedule.
 
+    This function generates a Gantt chart that visualizes the schedule of jobs
+    across multiple machines. Each job is represented with a unique color,
+    and operations are plotted as bars on the corresponding machines over time.
+
+    The Gantt chart helps to understand the flow of jobs on machines and
+    visualize the makespan of the schedule, i.e., the time it takes to
+    complete all jobs.
+
+    The Gantt chart includes:
+
+    - X-axis: Time units, representing the progression of the schedule.
+    - Y-axis: Machines, which are assigned jobs at various time slots.
+    - Legend: A list of jobs, labeled and color-coded for clarity.
+
+    .. note::
+        The last tick on the x-axis always represents the makespan for easy
+        identification of the completion time.
+
     Args:
         schedule:
             The schedule to plot.
         title:
             The title of the plot. If not provided, the title:
-            `f"Gantt Chart for {schedule.instance.name} instance"`
-            is used.
+            ``f"Gantt Chart for {schedule.instance.name} instance"``
+            is used. To remove the title, provide an empty string.
         cmap_name:
             The name of the colormap to use. Default is "viridis".
         xlim:
@@ -37,21 +59,45 @@ def plot_gantt_chart(
             the schedule is used.
         number_of_x_ticks:
             The number of ticks to use in the x-axis.
+        job_labels:
+            A list of labels for each job. If ``None``, the labels are
+            automatically generated as "Job 0", "Job 1", etc.
+        machine_labels:
+            A list of labels for each machine. If ``None``, the labels are
+            automatically generated as "0", "1", etc.
+        legend_title:
+            The title of the legend. If not provided, the legend will not have
+            a title.
+        x_label:
+            The label for the x-axis. Default is "Time units". To remove the
+            label, provide an empty string.
+        y_label:
+            The label for the y-axis. Default is "Machines". To remove the
+            label, provide an empty string.
+
+    Returns:
+        - A ``matplotlib.figure.Figure`` object.
+        - A ``matplotlib.axes.Axes`` object where the Gantt chart is plotted.
     """
-    fig, ax = _initialize_plot(schedule, title)
-    legend_handles = _plot_machine_schedules(schedule, ax, cmap_name)
-    _configure_legend(ax, legend_handles)
-    _configure_axes(schedule, ax, xlim, number_of_x_ticks)
+    fig, ax = _initialize_plot(schedule, title, x_label, y_label)
+    legend_handles = _plot_machine_schedules(
+        schedule, ax, cmap_name, job_labels
+    )
+    _configure_legend(ax, legend_handles, legend_title)
+    _configure_axes(schedule, ax, xlim, number_of_x_ticks, machine_labels)
     return fig, ax
 
 
 def _initialize_plot(
-    schedule: Schedule, title: str | None
-) -> tuple[Figure, plt.Axes]:
+    schedule: Schedule,
+    title: Optional[str],
+    x_label: str = "Time units",
+    y_label: str = "Machines",
+) -> Tuple[Figure, plt.Axes]:
     """Initializes the plot."""
     fig, ax = plt.subplots()
-    ax.set_xlabel("Time units")
-    ax.set_ylabel("Machines")
+    ax.set_xlabel(x_label)
+    ax.set_ylabel(y_label)
     ax.grid(True, which="both", axis="x", linestyle="--", linewidth=0.5)
     ax.yaxis.grid(False)
     if title is None:
@@ -61,8 +107,11 @@ def _initialize_plot(
 
 
 def _plot_machine_schedules(
-    schedule: Schedule, ax: plt.Axes, cmap_name: str
-) -> dict[int, Patch]:
+    schedule: Schedule,
+    ax: plt.Axes,
+    cmap_name: str,
+    job_labels: Optional[List[str]],
+) -> Dict[int, Patch]:
     """Plots the schedules for each machine."""
     max_job_id = schedule.instance.num_jobs - 1
     cmap = plt.get_cmap(cmap_name, max_job_id + 1)
@@ -81,12 +130,20 @@ def _plot_machine_schedules(
             )
             if scheduled_op.job_id not in legend_handles:
                 legend_handles[scheduled_op.job_id] = Patch(
-                    facecolor=color, label=f"Job {scheduled_op.job_id}"
+                    facecolor=color,
+                    label=_get_job_label(job_labels, scheduled_op.job_id),
                 )
 
     return legend_handles
 
 
+def _get_job_label(job_labels: Optional[List[str]], job_id: int) -> str:
+    """Returns the label for the job."""
+    if job_labels is None:
+        return f"Job {job_id}"
+    return job_labels[job_id]
+
+
 def _plot_scheduled_operation(
     ax: plt.Axes,
     scheduled_op: ScheduledOperation,
@@ -103,7 +160,9 @@ def _plot_scheduled_operation(
     )
 
 
-def _configure_legend(ax: plt.Axes, legend_handles: dict[int, Patch]):
+def _configure_legend(
+    ax: plt.Axes, legend_handles: Dict[int, Patch], legend_title: str
+):
     """Configures the legend for the plot."""
     sorted_legend_handles = [
         legend_handles[job_id] for job_id in sorted(legend_handles)
@@ -111,7 +170,8 @@ def _configure_legend(ax: plt.Axes, legend_handles: dict[int, Patch]):
     ax.legend(
         handles=sorted_legend_handles,
         loc="upper left",
-        bbox_to_anchor=(1.01, 1),
+        bbox_to_anchor=(1, 1),
+        title=legend_title,
     )
 
 
@@ -120,6 +180,7 @@ def _configure_axes(
     ax: plt.Axes,
     xlim: Optional[int],
     number_of_x_ticks: int,
+    machine_labels: list[str] | None,
 ):
     """Sets the limits and labels for the axes."""
     num_machines = len(schedule.schedule)
@@ -132,7 +193,9 @@ def _configure_axes(
             for i in range(num_machines)
         ]
     )
-    ax.set_yticklabels([str(i) for i in range(num_machines)])
+    if machine_labels is None:
+        machine_labels = [str(i) for i in range(num_machines)]
+    ax.set_yticklabels(machine_labels)
     makespan = schedule.makespan()
     xlim = xlim if xlim is not None else makespan
     ax.set_xlim(0, xlim)

job_shop_lib/visualization/graphs/__init__.py

@@ -0,0 +1,29 @@
+"""Contains functions and classes for visualizing job shop scheduling problems.
+
+.. autosummary::
+
+    plot_disjunctive_graph
+    plot_resource_task_graph
+    three_columns_layout
+    duration_labeler
+    color_nodes_by_machine
+
+"""
+
+from ._plot_disjunctive_graph import (
+    plot_disjunctive_graph,
+    duration_labeler,
+)
+from ._plot_resource_task_graph import (
+    plot_resource_task_graph,
+    three_columns_layout,
+    color_nodes_by_machine,
+)
+
+__all__ = [
+    "plot_disjunctive_graph",
+    "plot_resource_task_graph",
+    "three_columns_layout",
+    "duration_labeler",
+    "color_nodes_by_machine",
+]