ortools.sat.colab.visualization

  1#!/usr/bin/env python3
  2# Copyright 2010-2025 Google LLC
  3# Licensed under the Apache License, Version 2.0 (the "License");
  4# you may not use this file except in compliance with the License.
  5# You may obtain a copy of the License at
  6#
  7#     http://www.apache.org/licenses/LICENSE-2.0
  8#
  9# Unless required by applicable law or agreed to in writing, software
 10# distributed under the License is distributed on an "AS IS" BASIS,
 11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 12# See the License for the specific language governing permissions and
 13# limitations under the License.
 14
 15"""Collection of helpers to visualize cp_model solutions in colab."""
 16
 17# pylint: disable=g-import-not-at-top
 18import random
 19
 20try:
 21    from IPython.display import display
 22    from IPython.display import SVG
 23    import plotly.figure_factory as ff
 24    import svgwrite
 25
 26    correct_imports = True
 27except ImportError:
 28    correct_imports = False
 29
 30
 31def RunFromIPython():
 32    if not correct_imports:
 33        return False
 34    try:
 35        return __IPYTHON__ is not None
 36    except NameError:
 37        return False
 38
 39
 40def ToDate(v):
 41    return "2016-01-01 6:%02i:%02i" % (v / 60, v % 60)
 42
 43
 44class ColorManager:
 45    """Utility to create colors to use in visualization."""
 46
 47    def ScaledColor(self, sr, sg, sb, er, eg, eb, num_steps, step):
 48        """Creates an interpolated rgb color between two rgb colors."""
 49        num_intervals = num_steps - 1
 50        dr = (er - sr) / num_intervals
 51        dg = (eg - sg) / num_intervals
 52        db = (eb - sb) / num_intervals
 53        r = sr + dr * step
 54        g = sg + dg * step
 55        b = sb + db * step
 56        return "rgb(%i, %i, %i)" % (r, g, b)
 57
 58    def SeedRandomColor(self, seed=0):
 59        random.seed(seed)
 60
 61    def RandomColor(self):
 62        return "rgb(%i,%i,%i)" % (
 63            random.randint(0, 255),
 64            random.randint(0, 255),
 65            random.randint(0, 255),
 66        )
 67
 68
 69def DisplayJobshop(starts, durations, machines, name):
 70    """Simple function to display a jobshop solution using plotly."""
 71
 72    jobs_count = len(starts)
 73    machines_count = len(starts[0])
 74    all_machines = range(0, machines_count)
 75    all_jobs = range(0, jobs_count)
 76    df = []
 77    for i in all_jobs:
 78        for j in all_machines:
 79            df.append(
 80                dict(
 81                    Task="Resource%i" % machines[i][j],
 82                    Start=ToDate(starts[i][j]),
 83                    Finish=ToDate(starts[i][j] + durations[i][j]),
 84                    Resource="Job%i" % i,
 85                )
 86            )
 87
 88    sorted_df = sorted(df, key=lambda k: k["Task"])
 89
 90    colors = {}
 91    cm = ColorManager()
 92    cm.SeedRandomColor(0)
 93    for i in all_jobs:
 94        colors["Job%i" % i] = cm.RandomColor()
 95
 96    fig = ff.create_gantt(
 97        sorted_df,
 98        colors=colors,
 99        index_col="Resource",
100        title=name,
101        show_colorbar=False,
102        showgrid_x=True,
103        showgrid_y=True,
104        group_tasks=True,
105    )
106    fig.show()
107
108
109class SvgWrapper:
110    """Simple SVG wrapper to use in colab."""
111
112    def __init__(self, sizex, sizey, scaling=20.0):
113        self.__sizex = sizex
114        self.__sizey = sizey
115        self.__scaling = scaling
116        self.__offset = scaling
117        self.__dwg = svgwrite.Drawing(
118            size=(
119                self.__sizex * self.__scaling + self.__offset,
120                self.__sizey * self.__scaling + self.__offset * 2,
121            )
122        )
123
124    def Display(self):
125        display(SVG(self.__dwg.tostring()))
126
127    def AddRectangle(self, x, y, dx, dy, fill, stroke="black", label=None):
128        """Draw a rectangle, dx and dy must be >= 0."""
129        s = self.__scaling
130        o = self.__offset
131        corner = (x * s + o, (self.__sizey - y - dy) * s + o)
132        size = (dx * s - 1, dy * s - 1)
133        self.__dwg.add(
134            self.__dwg.rect(insert=corner, size=size, fill=fill, stroke=stroke)
135        )
136        self.AddText(x + 0.5 * dx, y + 0.5 * dy, label)
137
138    def AddText(self, x, y, label):
139        text = self.__dwg.text(
140            label,
141            insert=(
142                x * self.__scaling + self.__offset,
143                (self.__sizey - y) * self.__scaling + self.__offset,
144            ),
145            text_anchor="middle",
146            font_family="sans-serif",
147            font_size="%dpx" % (self.__scaling / 2),
148        )
149        self.__dwg.add(text)
150
151    def AddXScale(self, step=1):
152        """Add an scale on the x axis."""
153        o = self.__offset
154        s = self.__scaling
155        y = self.__sizey * s + o / 2.0 + o
156        dy = self.__offset / 4.0
157        self.__dwg.add(
158            self.__dwg.line((o, y), (self.__sizex * s + o, y), stroke="black")
159        )
160        for i in range(0, int(self.__sizex) + 1, step):
161            self.__dwg.add(
162                self.__dwg.line(
163                    (o + i * s, y - dy), (o + i * s, y + dy), stroke="black"
164                )
165            )
166
167    def AddYScale(self, step=1):
168        """Add an scale on the y axis."""
169        o = self.__offset
170        s = self.__scaling
171        x = o / 2.0
172        dx = self.__offset / 4.0
173        self.__dwg.add(
174            self.__dwg.line((x, o), (x, self.__sizey * s + o), stroke="black")
175        )
176        for i in range(0, int(self.__sizey) + 1, step):
177            self.__dwg.add(
178                self.__dwg.line(
179                    (x - dx, i * s + o), (x + dx, i * s + o), stroke="black"
180                )
181            )
182
183    def AddTitle(self, title):
184        """Add a title to the drawing."""
185        text = self.__dwg.text(
186            title,
187            insert=(
188                self.__offset + self.__sizex * self.__scaling / 2.0,
189                self.__offset / 2,
190            ),
191            text_anchor="middle",
192            font_family="sans-serif",
193            font_size="%dpx" % (self.__scaling / 2),
194        )
195        self.__dwg.add(text)