__init__.py

"""
Shamrock plot utility functions.
"""
 
import glob
 
import shamrock.sys
 
__all__ = []
 
try:
    import matplotlib.pyplot as plt
    from matplotlib import animation
 
    _HAS_MATPLOTLIB = True
except ImportError:
    _HAS_MATPLOTLIB = False
    # print("Warning: matplotlib is not installed, some Shamrock functions will not be available")
 
try:
    from PIL import Image
 
    _HAS_PIL = True
except ImportError:
    _HAS_PIL = False
    # print("Warning: PIL is not installed, some Shamrock functions will not be available")
 
try:
    import graphviz
 
    _HAS_GRAPHVIZ = True
except ImportError:
    _HAS_GRAPHVIZ = False
    # print("Warning: graphviz is not installed, some Shamrock functions will not be available")
 
if _HAS_MATPLOTLIB:
    from .benchmark import add_end_labels, make_std_bench_plot
 
    __all__.extend(["add_end_labels", "make_std_bench_plot"])
 
if _HAS_MATPLOTLIB and _HAS_PIL:
    __all__.append("show_image_sequence")
 
    def show_image_sequence(
        glob_str, render_gif=True, dpi=200, interval=50, repeat_delay=10, fig=None
    ):
        """
        Create a matplotlib animation from a sequence of image files.
 
        Available only if matplotlib and PIL are installed.
 
        Parameters
        ----------
        glob_str : str
            Glob pattern matching image files.
        render_gif : bool, optional
            Whether to render the animation.
        dpi : int, optional
            Dots per inch for the figure.
        interval : int, optional
            Delay between frames in milliseconds.
        repeat_delay : int, optional
            Delay before repeating the animation.
 
        Raises
        ------
        FileNotFoundError : if no images are found for the glob pattern
 
        Returns
        -------
        matplotlib.animation.FuncAnimation or None
            Animation object on rank 0, otherwise None.
        """
 
        if not render_gif:
            return None
 
        if shamrock.sys.world_rank() != 0:
            return None
 
        files = sorted(glob.glob(glob_str))
 
        image_array = []
        for my_file in files:
            with Image.open(my_file) as image:
                image_array.append(image.copy())
 
        if not image_array:
            raise FileNotFoundError(f"No images found for glob pattern: {glob_str}")
 
        pixel_x, pixel_y = image_array[0].size
 
        if fig is None:
            fig = plt.figure(dpi=dpi)
        plt.gca().set_position((0, 0, 1, 1))
        plt.gcf().set_size_inches(pixel_x / dpi, pixel_y / dpi)
        plt.axis("off")
 
        im = plt.imshow(image_array[0], animated=True, aspect="auto")
 
        def update(i):
            im.set_array(image_array[i])
            return (im,)
 
        ani = animation.FuncAnimation(
            fig,
            update,
            frames=len(image_array),
            interval=interval,
            blit=True,
            repeat_delay=repeat_delay,
        )
 
        return ani
 
 
__all__.append("DotGraph")
 
 
class DotGraph:
    """
    Wrap a Graphviz DOT graph source so it renders as inline SVG.
 
    Meant to display the solver graphs produced by e.g.
    ``model.get_solver_dot_graph()`` or ``setup_node.get_dot()`` in the
    sphinx-gallery generated examples: sphinx-gallery captures the
    ``_repr_html_`` of an expression left bare as the last statement of a
    code block (the same mechanism Jupyter uses for rich display), so no
    custom scraper is needed. See https://stackoverflow.com/a/65117672 for
    the technique this is based on.
 
    ``_repr_html_`` is only defined if the graphviz python package (and the
    Graphviz ``dot`` executable) are installed; otherwise this falls back to
    ``__repr__``, i.e. the raw DOT source.
 
    Parameters
    ----------
    dot_source : str
        Source of the graph in the DOT language.
    """
 
    def __init__(self, dot_source):
        self.dot_source = dot_source
 
    def _repr_html_(self):
        source = self.dot_source
        if not source.lstrip().startswith("digraph"):
            source = "digraph G {\n" + source + "\n}"
        return graphviz.Source(source).pipe(format="svg").decode("utf-8")
 
    def __repr__(self):
        return self.dot_source
 
 
if not _HAS_GRAPHVIZ:
    del DotGraph._repr_html_