py-pluto 1.1.4__py3-none-any.whl

pyPLUTO/imagefuncs/display.py

@@ -0,0 +1,306 @@
+"""Module to manage the display of 2D plots in the image."""
+
+from typing import Any
+
+import numpy as np
+from matplotlib.collections import QuadMesh
+from numpy.typing import ArrayLike
+
+from pyPLUTO.imagefuncs.colorbar import ColorbarManager
+from pyPLUTO.imagefuncs.imagetools import ImageToolsManager
+from pyPLUTO.imagefuncs.range import RangeManager
+from pyPLUTO.imagefuncs.set_axis import AxisManager
+from pyPLUTO.imagemixin import ImageMixin
+from pyPLUTO.imagestate import ImageState
+from pyPLUTO.utils.inspector import track_kwargs
+
+
+class DisplayManager(ImageMixin):
+    """Class to manage the display of 2D plots in the image.
+
+    This class provides methods to create and manage 2D plots using matplotlib's
+    pcolormesh function. It allows for customization of the plot's appearance,
+    colorbar, axes, and other properties.
+    """
+
+    exposed_methods = ("display",)
+
+    def __init__(self, state: ImageState) -> None:
+        """Initialize the DisplayManager with the given state."""
+        self.state = state
+        self.ColorbarManager = ColorbarManager(state)
+        self.ImageToolsManager = ImageToolsManager(state)
+        self.RangeManager = RangeManager(state)
+        self.AxisManager = AxisManager(state)
+
+    @track_kwargs
+    def display(
+        self,
+        var: ArrayLike,
+        check: bool = True,
+        **kwargs: Any,
+    ) -> QuadMesh:
+        """Plot for a 2D function using the matplotlib's pcolormesh function.
+
+        A simple figure and a single axis can also be
+        created.
+
+        Returns
+        -------
+        - The 2D plot
+
+        Parameters
+        ----------
+        - alpha: float, default 1.0
+            Sets the transparency of the plot.
+        - aspect: {'auto', 'equal', float}, default 'auto'
+            Sets the aspect ratio of the plot.
+            The 'auto' keyword is the default option (most likely the plot will
+            be squared). The 'equal' keyword will set the same scaling for
+            x and y. A float will fix the ratio between the y-scale and the
+            x-scale (1.0 is the same as 'equal').
+        - ax: ax | int | None, default None
+            The axis where to plot the lines. If None, a new axis is created.
+            If 'old', the last considered axis will be used.
+        - bottom: float, default 0.1
+            The space from the bottom border to the plot.
+        - clabel: str, default None
+            Sets the label of the colorbar.
+        - cmap: str, default 'plasma'
+            Selects the colormap. If not defined, the colormap 'plasma' will be
+            adopted. Some useful colormaps are: plasma, magma, seismic. Please
+            avoid using colorbars like jet or rainbow, which are not
+            perceptively uniform and not suited for people with vision
+            deficiencies.
+        - cpad: float, default 0.07
+            Fraction of original axes between colorbar and the axes (in case cax
+            is not defined).
+        - cpos: {'top','bottom','left','right'}, default None
+            Enables the colorbar (if defined), default position on the right.
+        - cscale: {'linear','log','symlog','twoslope'}, default 'linear'
+            Sets the colorbar scale. Default is the linear ('norm') scale.
+        - cticks: {[float], None}, default None
+            If enabled (and different from None), sets manually ticks on the
+            colorbar.
+        - ctickslabels: str, default None
+            If enabled, sets manually ticks labels on the colorbar.
+        - extend: {'neither','both','min','max'}, default 'neither'
+            Sets the extension of the triangular colorbar extension.
+        - extendrect: bool, default False
+            If True, the colorbar extension will be rectangular.
+        - figsize: [float, float], default [6*sqrt(ncol),5*sqrt(nrow)]
+            Sets the figure size. The default value is computed from the number
+            of rows and columns.
+        - fontsize: float, default 17.0
+            Sets the fontsize for all the axes.
+        - grid: Bool, default False
+            Enables the grid on the plot.
+        - labelsize: float, default fontsize
+            Sets the labels fontsize (which is the same for both labels).
+            The default value corresponds to the value of the keyword
+            'fontsize'.
+        - left: float, default 0.125
+            The space from the left border to the plot.
+        - minorticks: str, default None
+            If not None enables the minor ticks on the plot (for both grid
+            axes).
+        - proj: str, default None
+            Custom projection for the plot (e.g. 3D). Recommended only if
+            needed.
+            This keyword should be used only if the axis is created.
+            WARNING: pyPLUTO does not support 3D plotting for now, only 3D axes.
+            The 3D plot feature will be available in future releases.
+        - right: float, default 0.9
+            The space from the right border to the plot.
+        - shading: {'flat,'nearest','auto','gouraud'}, default 'auto'
+            The shading between the grid points. If not defined, the shading
+            will one between 'flat' and 'nearest' depending on the size of the
+            x,y and z arrays. The 'flat' shading works only if, given a NxM
+            z-array, the x- and y-arrays have sizes of, respectively, N+1 and
+            M+1. All the other shadings require a N x-array and a M y-array.
+        - ticksdir: {'in', 'out'}, default 'in'
+            Sets the ticks direction. The default option is 'in'.
+        - tickssize: float, default fontsize
+            Sets the ticks fontsize (which is the same for both grid axes).
+            The default value corresponds to the value of the keyword
+            'fontsize'.
+        - title: str, default None
+            Places the title of the plot on top of it.
+        - titlesize: float, default fontsize
+            Sets the title fontsize. The default value corresponds to the value
+            of the keyword 'fontsize'.
+        - top: float, default 0.9
+            The space from the top border to the plot.
+        - transpose: True/False, default False
+            Transposes the variable matrix. Use is not recommended if not really
+            necessary (e.g. in case of highly customized variables and plots)
+        - tresh: float, default max(abs(vmin),vmax)*0.01
+            Sets the threshold for the colormap. If not defined, the threshold
+            will be set to 1% of the maximum absolute value of the variable.
+            The default cases are the following:
+            - twoslope colorscale: sets the limit between the two linear
+            regimes.
+            - symlog: sets the limit between the logaitrhmic and the linear
+            regime.
+        - var (not optional): 2D array
+            The array to be plotted.
+        - vmax: float, default max(var)
+            The maximum value of the colormap. If not defined, the maximum value
+            of z will be taken.
+        - vmin: float, default min(var)
+            The minimum value of the colormap. If not defined, the minimum value
+            of z will be taken.
+        - x1: np.ndarray, default 'Default'
+            the 'x' array. If not defined, a default array will be generated
+            depending on the size of z.
+        - x2: np.ndarray, default 'Default'
+            the 'y' array. If not defined, a default array will be generated
+            depending on the size of z.
+        - xrange: [float, float], default 'Default'
+            Sets the range in the x-direction. If not defined or set to
+            'Default' the code will compute the range while plotting the data by
+            taking the minimum and the maximum values of the x1-array.
+        - xscale: {'linear','log'}, default 'linear'
+            If enabled (and different from 'Default'), sets automatically the
+            scale on the x-axis. Data in log scale should be used with the
+            keyword 'log', while data in linear scale should be used with the
+            keyword 'linear'.
+        - xticks: [float] | None | bool, default True
+            If enabled (and different from 'Default'), sets manually ticks on
+            x-axis. In order to completely remove the ticks the keyword should
+            be used with None.
+        - xtickslabels: [str] | None | bool, default True
+            If enabled (and different from 'Default'), sets manually the ticks
+            labels on the x-axis. In order to completely remove the ticks the
+            keyword should be used with None. Note that fixed tickslabels should
+            always correspond to fixed ticks.
+        - xtitle: str, default None
+            Sets and places the label of the x-axis.
+        - yrange: [float, float], default 'Default'
+            Sets the range in the y-direction. If not defined or set to
+            'Default' the code will compute the range while plotting the data by
+            taking the minimum and the maximum values of the x2-array.
+        - yscale: {'linear','log'}, default 'linear'
+            If enabled (and different from 'Default'), sets automatically the
+            scale on the y-axis. Data in log scale should be used with the
+            keyword 'log', while data in linear scale should be used with the
+            keyword 'linear'.
+        - yticks: [float] | None | bool, default True
+            If enabled (and different from 'Default'), sets manually ticks on
+            y-axis. In order to completely remove the ticks the keyword should
+            be used with None.
+        - ytickslabels: [float] | None | bool, default True
+            If enabled (and different from 'Default'), sets manually the ticks
+            labels on the y-axis. In order to completely remove the ticks the
+            keyword should be used with None. Note that fixed tickslabels should
+            always correspond to fixed ticks.
+        - ytitle: str, default None
+            Sets and places the label of the y-axis.
+
+        ----
+
+        Examples
+        --------
+        - Example #1: create a simple 2d plot with title and colorbar on the
+            right
+
+            >>> import pyPLUTO as pp
+            >>> I = pp.Image()
+            >>> I.display(var, title="title", cpos="right")
+
+        - Example #2: create a 2d plot with title on the axes, bottom colorbar
+            and custom shading
+
+            >>> import pyPLUTO as pp
+            >>> I = pp.Image()
+            >>> I.display(x1, x2, var, xtitle = 'x', ytitle = 'y',
+                    cpos = 'bottom', shading = 'gouraud', cpad = 0.3)
+
+        - Example #3: create a 2d plot con custom range on axes and logarithmic
+            scale colorbar
+
+            >>> import pyPLUTO as pp
+            >>> I = pp.Image()
+            >>> I.display(var, xrange = [2,3], yrange = [2,4], cbar = 'right',
+                        cscale = 'log')
+
+        - Example #4: create a 2d plot with a custom symmetric logarithmic
+            colorbar with custom ticks.
+
+            >>> import pyPLUTO as pp
+            >>> I = pp.Image()
+            >>> I.display(var, cpos = 'right', cmap = 'RdBu_r',
+                        cscale = 'symlog', tresh = 0.001, vmin = -1, vmax = 1)
+
+        """
+        kwargs.pop("check", check)
+
+        # Set or create figure and axes
+        ax, nax = self.ImageToolsManager.assign_ax(
+            kwargs.pop("ax", None), **kwargs
+        )
+
+        if self.fig is None:
+            raise ValueError(
+                "No figure is present. Please create a figure first."
+            )
+        # Keyword x1 and x2
+        var = np.asarray(var)
+        if kwargs.get("transpose", False) is True:
+            var = var.T
+        x = np.asarray(kwargs.get("x1", np.arange(len(var[:, 0]) + 1)))
+        y = np.asarray(kwargs.get("x2", np.arange(len(var[0, :]) + 1)))
+
+        # Keywords xrange and yrange
+        if not kwargs.get("xrange") and self.setax[nax] != 1:
+            kwargs["xrange"] = [x.min(), x.max()]
+        if not kwargs.get("yrange") and self.setay[nax] != 1:
+            kwargs["yrange"] = [y.min(), y.max()]
+        # Set ax parameters
+        self.AxisManager.set_axis(ax=ax, check=False, **kwargs)
+        self.ImageToolsManager.hide_text(nax, ax.texts)
+
+        # Keywords vmin and vmax
+        vmin = kwargs.get("vmin", np.nanmin(var))
+        vmax = kwargs.get("vmax", np.nanmax(var))
+
+        # Keyword for colorbar and colorscale
+        cpos = kwargs.get("cpos")
+        cscale = kwargs.get("cscale", "norm")
+        tresh = kwargs.get("tresh", max(np.abs(vmin), vmax) * 0.01)
+        lint = kwargs.get("lint")
+        self.vlims[nax] = [vmin, vmax, tresh]
+
+        # Set the colorbar scale (put in function)
+        norm = self.ImageToolsManager.set_cscale(
+            cscale, vmin, vmax, tresh, lint
+        )
+
+        # Select shading
+        shade = kwargs.get("shading", "auto")
+        alpha = kwargs.get("alpha", 1.0)
+
+        cmap = self.ImageToolsManager.find_cmap(kwargs.get("cmap", "plasma"))
+
+        # Display the image
+        pcm = ax.pcolormesh(
+            x,
+            y,
+            var.T,
+            shading=shade,
+            cmap=cmap,
+            norm=norm,
+            linewidth=0,
+            rasterized=True,
+            alpha=alpha,
+        )
+
+        # Place the colorbar (use colorbar function)
+        if cpos is not None:
+            self.ColorbarManager.colorbar(pcm, check=False, **kwargs)
+
+        # If tight_layout is enabled, is re-inforced
+        if self.tight:
+            self.fig.tight_layout()
+
+        return pcm

pyPLUTO/imagefuncs/figure.py

@@ -0,0 +1,395 @@
+"""Figure Manager Module."""
+
+import shutil
+import warnings
+from typing import Unpack
+
+import matplotlib as mpl
+import matplotlib.pyplot as plt
+from matplotlib.figure import Figure
+
+from pyPLUTO.imagemixin import ImageMixin
+from pyPLUTO.imagestate import ImageState
+from pyPLUTO.utils.annotator import AllKwargs
+from pyPLUTO.utils.inspector import track_kwargs
+
+
+class FigureManager(ImageMixin):
+    """Manages the figure and sets the style, size, and LaTeX settings."""
+
+    @track_kwargs
+    def __init__(
+        self,
+        state: ImageState,
+        **kwargs: Unpack[AllKwargs],
+    ) -> None:
+        """Initialize the FigureManager class.
+
+        It creates a new figure and sets the LaTeX conditions, as well as the
+        matplotlib style. Every Image is associated to a figure object and only
+        one in order to avoid confusion between images and figures. If you
+        want to create multiple figures, you have to create multiple
+        Image objects.
+
+        Returns
+        -------
+        - None
+
+        Parameters
+        ----------
+        - state: ImageState
+            The state of the image, which contains the figure and other
+            properties.
+        - kwargs: dict[str, Any]
+            Additional keyword arguments to customize the figure, such as
+            `figsize`, `fontsize`, `nwin`, `suptitle`, etc.
+
+        """
+        # needed because the __init__ is longer than simply self.state = state
+        self.state = state
+
+        # Extract specific kwargs for colorlines, with defaults if not provided
+        close = kwargs.pop("close", True)
+        fontweight = kwargs.pop("fontweight", "normal")
+        numcolors = kwargs.pop("numcolors", 10)
+        replace = kwargs.pop("replace", False)
+        suptitle = kwargs.pop("suptitle", None)
+        suptitlesize = kwargs.pop("suptitlesize", "large")
+        withblack = kwargs.pop("withblack", False)
+        withwhite = kwargs.pop("withwhite", False)
+
+        self.fig = kwargs.get("fig", self.fig)
+        self.figsize = kwargs.get("figsize", self.figsize)
+        self.fontsize = kwargs.get("fontsize", self.fontsize)
+        self.LaTeX = kwargs.get("LaTeX", self.LaTeX)
+        self.nwin = kwargs.get("nwin", self.nwin)
+        self.style = kwargs.get("style", self.style)
+        self.tight = kwargs.get("tight", self.tight)
+
+        self.check_previous_fig(close)
+        if "figsize" in kwargs:
+            self.set_size = True
+
+        self.setup_style()
+        self.color = self.choose_colorlines(numcolors, withblack, withwhite)
+        self.assign_LaTeX(fontweight)
+        self.create_figure(replace, suptitle, suptitlesize)
+
+    def setup_style(self) -> None:
+        """Set the matplotlib style."""
+        try:
+            plt.style.use(self.style)
+        except OSError:
+            warn = f"Warning: Style '{self.style}' not found. \
+                Switching to 'default'"
+            warnings.warn(warn, UserWarning, stacklevel=2)
+            self.style = "default"
+
+    def choose_colorlines(
+        self, numcolors: int, withblack: bool, withwhite: bool
+    ) -> list[str]:
+        """Choose the colors for the lines.
+
+        The colors are taken from a list of colors that are suitable for all
+        types of color vision deficiencies.
+
+        Returns
+        -------
+        - colors: list[str]
+            The list of colors for the lines.
+
+        Parameters
+        ----------
+        - numcolors: int, default 10
+            The number of colors.
+        - withblack: bool, default False
+            If True, the black color is used as first color.
+        - withwhite: bool default False
+            If True, the white color is used as first color.
+
+        ----
+
+        Examples
+        --------
+        - Example #1: withblack = True
+
+            >>> _choose_colorlines(6, True)
+
+        - Example #2: 12 colors, withwhite = True
+
+            >>> _choose_colorlines(12, False, True)
+
+        """
+        # New colors dictionary (black and white included)
+        self.dictcol = {
+            0: "#ffffff",
+            1: "#e8ecfb",
+            2: "#d9cce3",
+            3: "#d1bbd7",
+            4: "#caaccb",
+            5: "#ae76a3",
+            6: "#aa6f9e",
+            7: "#994f88",
+            8: "#882e72",
+            9: "#0104fe",
+            10: "#1e3888",
+            11: "#437dbf",
+            12: "#5289c7",
+            13: "#6195cf",
+            14: "#7bafde",
+            15: "#4eb265",
+            16: "#90c987",
+            17: "#cae0ab",
+            18: "#f7f056",
+            19: "#f7cb45",
+            20: "#f6c141",
+            21: "#f4a736",
+            22: "#f1932d",
+            23: "#ee8026",
+            24: "#e8601c",
+            25: "#e65518",
+            26: "#dc050c",
+            27: "#a5170e",
+            28: "#72190e",
+            29: "#42150a",
+            30: "#777777",
+            31: "#000000",
+            32: "#0104fe",
+        }
+
+        # Colors are ordered to avoid color vision deficiencies
+        lstc = [
+            9,
+            26,
+            15,
+            23,
+            14,
+            17,
+            6,
+            25,
+            28,
+            18,
+            11,
+            2,
+            8,
+            16,
+            10,
+            21,
+            7,
+            27,
+            4,
+            13,
+            19,
+            29,
+            1,
+            30,
+        ]
+
+        # Black and white addition
+        lstc = [0, *lstc] if withwhite else [31, *lstc] if withblack else lstc
+
+        # End of function, return the colors
+        return [self.dictcol[lstc[i]] for i in range(numcolors)]
+
+    def assign_LaTeX(self, fontweight: str) -> None:
+        """Set the LaTeX conditions.
+
+        The option 'pgf' requires XeLaTeX and should be used only to get
+        vectorial figures with minimal file size.
+
+        Returns
+        -------
+        - None
+
+        Parameters
+        ----------
+        - LaTeX (not optional): bool | str
+            The LaTeX option. Is True is selected, the default LaTeX font is
+            used. If 'pgf' is selected, the pgf backend is used to save pdf
+            figures with minimal file size. If XeLaTeX is not installed and the
+            'pgf' option is selected, the LaTeX option True is used as backup
+            strategy.
+
+        ----
+
+        Examples
+        --------
+        - Example #1: LaTeX option True
+
+            >>> _assign_LaTeX(True)
+
+        - Example #2: LaTeX option 'pgf'
+
+            >>> _assign_LaTeX("pgf")
+
+        """
+        # LaTeX option 'pgf' (requires XeLaTeX)
+        if self.LaTeX == "pgf" and not shutil.which("latex"):
+            warn = "LaTeX not installed, switching to LaTeX = True"
+            warnings.warn(warn, UserWarning, stacklevel=2)
+            self.LaTeX = True
+
+        if self.LaTeX == "pgf":
+            # Set the pgf backend
+            try:
+                plt.switch_backend("pgf")
+
+                # Preamble (LaTeX commands and packages)
+                pgf_preamble = r"""
+                \usepackage{amsmath}
+                \usepackage{amssymb}
+                \usepackage{mathptmx}
+                \usepackage{siunitx}
+                \usepackage[T1]{fontenc}
+                \newcommand{\DS}{\displaystyle}
+                """
+
+                # Update the rcParams
+                mpl.rcParams.update(
+                    {
+                        "pgf.preamble": pgf_preamble,
+                        "font.family": "serif",
+                        "font.weight": fontweight,
+                        "text.usetex": True,
+                    }
+                )
+
+            # If errors occur, the LaTeX option True is used and a warning
+            # message is displayed
+            except ImportError:
+                warn = "The pgf backend is not available, reverting to True\n"
+                warnings.warn(warn, UserWarning, stacklevel=2)
+                self.LaTeX = True
+
+        # LaTeX option True: default LaTeX font
+        if self.LaTeX is True:
+            try:
+                mpl.rcParams["mathtext.fontset"] = "stix"
+                mpl.rcParams["font.family"] = "STIXGeneral"
+            except ImportError:
+                warn = "The LaTeX = True option is not available."
+                warnings.warn(warn, UserWarning, stacklevel=2)
+
+        # End of the function
+
+    def check_previous_fig(self, close: bool) -> None:
+        """Check if there is an existing figure.
+
+        If it exists, the code will check if it is closed or not.
+
+        Returns
+        -------
+        - None
+
+        Parameters
+        ----------
+        - close: bool, default True
+            If True, the existing figure with the same window number is closed.
+
+        ----
+
+        Examples
+        --------
+        - Example #1: Check if there is an existing figure
+
+            >>> _check_previous_fig(True)
+
+        """
+        if isinstance(self.fig, Figure):
+            self.figsize = [
+                self.fig.get_figwidth(),
+                self.fig.get_figheight(),
+            ]
+            self.fontsize = plt.rcParams["font.size"]
+            try:
+                fignum = self.fig.number
+                if isinstance(fignum, int):
+                    self.state.nwin = fignum
+            except AttributeError:
+                warnings.warn(
+                    "The figure is not associated to a window number",
+                    UserWarning,
+                    stacklevel=2,
+                )
+                self.nwin = 1
+            self.tight = self.fig.get_tight_layout()
+
+        # Close the existing figure if it exists (and 'close' is enabled).
+        # `clf()` releases artists/axes payloads immediately, which prevents
+        # stale Image instances (still holding axes/text refs) from retaining
+        # heavy plot data in memory.
+        if plt.fignum_exists(self.nwin) and close is True:
+            existing_fig = plt.figure(self.nwin)
+            existing_fig.clf()
+            plt.close(existing_fig)
+
+    def create_figure(
+        self, replace: bool, suptitle: str | None, suptitlesize: int | str
+    ) -> None:
+        """Create the figure associated to an Image instance.
+
+        It is called by default when the Image class is instantiated.
+
+        Returns
+        -------
+        - None
+
+        Parameters
+        ----------
+        - close: bool, default True
+            If True, the existing figure with the same window number is closed.
+        - fig (not optional): Figure | None, default None
+            The figure instance. If not None, the figure is used (only if we
+            need to associate an Image to an existing figure).
+        - figsize: list[float], default [8,5]
+            The figure size.
+        - fontsize: int, default 17
+            The font size.
+        - nwin: int, default 1
+            The window number.
+        - suptitle: str | None, default None
+            The super title of the figure.
+        - suptitlesize: int | str, default 'large'
+            The figure title size.
+        - tight: bool, default True
+            If True, the tight layout is used.
+
+        ----
+
+        Examples
+        --------
+        - Example #1: Create a new figure
+
+            >>> _create_figure()
+
+        - Example #2: Associate an Image to an existing figure
+
+            >>> _create_figure(fig=fig)
+
+        - Example #3: Create a new figure with different size and a figure title
+
+            >>> _create_figure(suptitle="Super Title", figsize=[10, 5])
+
+        - Example #4: Create a new figure with a specific window number
+
+            >>> _create_figure(nwin=2)
+
+        """
+        # Create a new figure instance with the provided window number
+        if self.fig is None or replace is True:
+            self.fig = plt.figure(
+                self.nwin,
+                figsize=(self.figsize[0], self.figsize[1]),
+            )
+        plt.rcParams.update({"font.size": self.fontsize})
+
+        if self.fig is None:
+            raise ValueError("The figure could not be created.")
+
+        # Suptitle
+        if suptitle is not None:
+            self.fig.suptitle(suptitle, fontsize=suptitlesize)
+
+        # Tight layout
+        if self.tight is True:
+            self.fig.tight_layout()