py-pluto 1.1.4__py3-none-any.whl

pyPLUTO/imagefuncs/contour.py

@@ -0,0 +1,254 @@
+"""Module for managing contour plots in image displays."""
+
+import warnings
+from typing import Any
+
+import numpy as np
+from matplotlib.contour import QuadContourSet
+from numpy.typing import NDArray
+
+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 ContourManager(ImageMixin):
+    """Class to manage contour plots in the image.
+
+    This class provides methods to create contour plots of variables in the
+    image class. It allows for customization of the contour lines, colorbars,
+    and other properties.
+    """
+
+    exposed_methods = ("contour",)
+
+    def __init__(self, state: ImageState) -> None:
+        """Initialize the ContourManager with the given state."""
+        self.state = state
+        self.AxisManager = AxisManager(state)
+        self.ColorbarManager = ColorbarManager(state)
+        self.ImageToolsManager = ImageToolsManager(state)
+        self.RangeManager = RangeManager(state)
+
+    @track_kwargs
+    def contour(
+        self, var: NDArray[np.generic], check: bool = True, **kwargs: Any
+    ) -> QuadContourSet:
+        """Plot a contour plot of a given variable.
+
+        The function uses the matplotlib.pyplot.contour function. The function
+        returns None.
+
+        Returns
+        -------
+        - cnt: LineCollection
+            The set of contour lines of the given variable.
+
+        Parameters
+        ----------
+        - alpha: float, default 1.0
+            Sets the transparency of the contour lines.
+        - 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 object, 'old', 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.
+        - c: str, default self.color
+            Determines the contour lines plot. If not defined, the program will
+            loop over an array of 6 color which are different for the most
+            common vision deficiencies.
+        - cmap: str, default 'hot'
+            Selects the colormap. If not defined, the colormap 'hot' will be
+            adopted. Some useful colormaps are: plasma, magma, seismic. Please
+            avoid using colorbars like jjet or rainbow, which are not
+            perceptively uniform and not suited for people with vision
+            deficiencies.
+        - 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.
+        - 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.
+        - fontsize: float, default 17.0
+            Sets the fontsize for all the axis components (only for the current
+            axis).
+        - grid: bool, default False
+            Enables/disables 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'.
+        - levels: np.ndarray
+            The levels of the contour lines.
+        - minorticks: str, default None
+            If not None enables the minor ticks on the plot (for both grid
+            axes).
+        - sharex: Matplotlib axis | False, default False
+            Shares the x-axis with another axis.
+        - sharey: Matplotlib axis | False, default False
+            Shares the y-axis with another axis.
+        - ticksdir: {'in', 'out'}, default 'in'
+            Sets the ticks direction. The default option is 'in'.
+        - tickssize: float | bool, default True
+            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.
+        - titlepad: float, default 8.0
+            Sets the distance between the title and the top of the plot
+        - titlesize: float, default fontsize
+            Sets the title fontsize. The default value corresponds to the value
+            of the keyword 'fontsize'.
+        - 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): np.ndarray
+            The variable to be plotted.
+        - vmax: float
+            The maximum value of the colormap.
+        - vmin: float
+            The minimum value of the colormap.
+        - x1: 1D array, default 'Default'
+            The 'x' array.
+        - x2: 1D array, default 'Default'
+            The 'y' array.
+        - xrange: [float, float], default [0,1]
+            Sets the range in the x-direction. If not defined the code will
+            compute the range while plotting the data.
+        - xscale: {'linear','log'}, default 'linear'
+            If enabled (and different from True), 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, True}, default True
+            If enabled (and different from True), sets manually ticks on
+            x-axis. In order to completely remove the ticks the keyword should
+            be used with None.
+        - xtickslabels: {[str], None, True}, default True
+            If enabled (and different from True), 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 [0,1]
+            Sets the range in the y-direction. If not defined the code will
+            compute the range while plotting the data.
+        - yscale: {'linear','log'}, default 'linear'
+            If enabled (and different from True), 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, True}, default True
+            If enabled (and different from True), sets manually ticks on
+            y-axis. In order to completely remove the ticks the keyword should
+            be used with None.
+        - ytickslabels: {[str], None, True}, default True
+            If enabled (and different from True), 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: Plot a contour plot of a variable
+
+            >>> I.contour(D.rho, levels=10)
+
+        """
+        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
+        x = np.asarray(kwargs.get("x1", np.arange(len(var[:, 0]))))
+        y = np.asarray(kwargs.get("x2", np.arange(len(var[0, :]))))
+
+        # Transpose if needed
+        var = np.asarray(var.T)
+        if kwargs.get("transpose", False) is True:
+            var = var.T
+
+        # 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))
+
+        # Sets levels for the contour plot
+        levels = kwargs.get("levels", np.linspace(vmin, vmax, 10))
+
+        # Keyword for colorbar and colorscale
+        colors = kwargs.get("c")
+        cmap = self.ImageToolsManager.find_cmap(kwargs.get("cmap"))
+        cpos = kwargs.get("cpos")
+        cscale = kwargs.get("cscale", "norm")
+        tresh = kwargs.get("tresh", max(np.abs(vmin), vmax) * 0.01)
+        lint = kwargs.get("lint")
+        lw = kwargs.get("lw", 1.0)
+
+        if "colors" in kwargs and "cmap" in kwargs:
+            warn = "Both colors and cmap are defined. Using c."
+            warnings.warn(warn, UserWarning, stacklevel=2)
+
+        # Set the colorbar scale (put in function)
+        norm = self.ImageToolsManager.set_cscale(
+            cscale, vmin, vmax, tresh, lint
+        )
+
+        # Select shading
+        alpha = kwargs.get("alpha", 1.0)
+
+        # Plot the contour plot
+        cnt = ax.contour(
+            x,
+            y,
+            var,
+            levels=levels,
+            norm=norm,
+            cmap=cmap,
+            colors=colors,
+            alpha=alpha,
+            linewidths=lw,
+        )
+
+        if cpos is not None:
+            self.ColorbarManager.colorbar(cnt, check=False, **kwargs)
+
+        # If tight_layout is enabled, is re-inforced
+        if self.tight:
+            self.fig.tight_layout()
+
+        return cnt