py-pluto 1.1.4__py3-none-any.whl

pyPLUTO/imagefuncs/streamplot.py

@@ -0,0 +1,297 @@
+"""Module to manage the streamplot function from matplotlib.pyplot."""
+
+import warnings
+from typing import Any
+
+import numpy as np
+from matplotlib.collections import LineCollection
+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 StreamplotManager(ImageMixin):
+    """Manages the streamplot function from matplotlib.pyplot."""
+
+    exposed_methods = ("streamplot",)
+
+    def __init__(self, state: ImageState) -> None:
+        """Initialize the StreamplotManager 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 streamplot(
+        self,
+        var1: NDArray[np.generic],
+        var2: NDArray[np.generic],
+        check: bool = True,
+        **kwargs: Any,
+    ) -> LineCollection:
+        """Plot a streamplot of a vector field.
+
+        The function uses the streamplot function from matplotlib.pyplot.
+
+        Returns
+        -------
+        - strm: LineCollection
+            The streamplot of the given vector field.
+
+        Parameters
+        ----------
+        - alpha: float, default 1.0
+            Sets the opacity of the plot, where 1.0 means total opaque and 0.0
+            means total transparent.
+        - arrowsize: float, default 1.0
+            Sets the size of the arrows of the streamline.
+        - arrowstyle: str, default '-|>'
+            Sets the style of the arrows of the streamline.
+        - 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: axis object
+            The axis where to plot the streamlines. If not given, the last
+            considered axis will be used.
+        - brokenlines: bool, default True
+            Splits the streamlines in shorter segments.
+        - c: str, default self.color
+            Determines the streamplot color. 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.
+            All the colormap available are listed in the following link:
+            https://matplotlib.org/stable/tutorials/colors/colormaps.html
+        - 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.
+        - density: float, default 1.0
+            Sets the density and closeness of the streamlines. The domain is
+            divided in a 30x30 grid. When set as default, each cell contains at
+            most a number of crossing streamplot line equal to this keyword.
+        - 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.
+        - integration_direction: {'forward', 'backward', 'both'}, default:'both'
+            Sets the streamlines integration in the forward direction, backward
+            direction, or both.
+        - 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'.
+        - lw: float, default 1.0
+            Sets the width of the streamlines.
+        - maxlength: float, default 5.0
+            Sets the maximum length of a streamline in coordinates units.
+        - minlength: float, default 0.1
+            Sets the minimum length of a streamline in coordinates units.
+        - minorticks: str, default None
+            If not None enables the minor ticks on the plot (for both grid
+            axes).
+        - start_points: np.ndarray, default None
+            Sets the starting points of the streamlines, if a more controlled
+            plot is wanted.
+        - 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.
+        - var1 (not optional): np.ndarray
+            The x1-component of the vector field.
+        - var2 (not optional): np.ndarray
+            The x2-component of the vector field.
+        - vmax: float
+            The maximum value of the vector field norm.
+        - vmin: float
+            The minimum value of the vector field norm.
+        - x1: np.ndarray
+            The x-axis variable.
+        - x2: np.ndarray
+            The y-axis variable.
+        - 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 streamplot of a vector field
+
+            >>> I.streamplot(D.Bx1, D.Bx2)
+
+        """
+        kwargs.pop("check", check)
+
+        if np.shape(var1) != np.shape(var2):
+            raise ValueError("The shapes of the variables are different.")
+
+        # 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."
+            )
+
+        x = np.asarray(kwargs.get("x1", np.arange(len(var1[:, 0]))))
+        y = np.asarray(kwargs.get("x2", np.arange(len(var1[0, :]))))
+
+        # Keyword x1 and x2
+        varx, vary = (
+            np.array(var2.T, dtype=float, copy=True),
+            np.array(var1.T, dtype=float, copy=True),
+        )
+        if kwargs.get("transpose", False) is True:
+            varx, vary = varx.T, vary.T
+
+        fieldmod = np.sqrt(varx**2 + vary**2)
+        vmax = kwargs.get("vmax", np.nanmax(fieldmod))
+        vmin = kwargs.get("vmin", np.nanmin(fieldmod))
+
+        # Apply the masks to set the corresponding elements
+        # in varx and vary to NaN
+        mask = np.logical_or(fieldmod > vmax, fieldmod < vmin)
+        varx[mask] = vary[mask] = np.nan
+
+        # Set ax parameters
+        self.AxisManager.set_axis(ax=ax, check=False, **kwargs)
+        self.ImageToolsManager.hide_text(nax, ax.texts)
+
+        # Keyword for colorbar and colorscale
+        color = 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")
+
+        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 lines properties
+        linewidth = kwargs.get("lw", 1)
+        density = kwargs.get("density", 1)
+        arrowstyle = kwargs.get("arrowstyle", "-|>")
+        arrowsize = kwargs.get("arrowsize", 1)
+        minlength = kwargs.get("minlength", 0.1)
+        integration_direction = kwargs.get("integration_direction", "both")
+        start_points = kwargs.get("start_points")
+        maxlength = kwargs.get("maxlength", 5)
+        broken_streamlines = kwargs.get("brokenlines", True)
+
+        # Set the colorbar scale (put in function)
+        norm = self.ImageToolsManager.set_cscale(
+            cscale, vmin, vmax, tresh, lint
+        )
+
+        # Plot the streamplot
+        strm = ax.streamplot(
+            x,
+            y,
+            vary,
+            varx,
+            norm=norm,
+            cmap=cmap,
+            color=color,
+            linewidth=linewidth,
+            density=density,
+            arrowsize=arrowsize,
+            minlength=minlength,
+            maxlength=maxlength,
+            start_points=start_points,
+            arrowstyle=arrowstyle,
+            integration_direction=integration_direction,
+            broken_streamlines=broken_streamlines,
+        )
+
+        if cpos is not None:
+            self.ColorbarManager.colorbar(strm.lines, check=False, **kwargs)
+
+        # If tight_layout is enabled, is re-inforced
+        if self.tight:
+            self.fig.tight_layout()
+
+        del varx, vary
+
+        return strm.lines