py-pluto 1.1.4__py3-none-any.whl

pyPLUTO/imagefuncs/scatter.py

@@ -0,0 +1,270 @@
+"""ScatterManager class."""
+
+from typing import Any
+
+import numpy as np
+from matplotlib.collections import PathCollection
+from numpy.typing import NDArray
+
+from pyPLUTO.imagefuncs.colorbar import ColorbarManager
+from pyPLUTO.imagefuncs.imagetools import ImageToolsManager
+from pyPLUTO.imagefuncs.legend import LegendManager
+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 ScatterManager(ImageMixin):
+    """Manager for the scatter plot of a 2D function.
+
+    A simple figure and a single axis can also be created.
+    """
+
+    exposed_methods = ("scatter",)
+
+    def __init__(self, state: ImageState) -> None:
+        """Initialize the ScatterManager with the given state."""
+        self.state = state
+
+        self.AxisManager = AxisManager(state)
+        self.ColorbarManager = ColorbarManager(state)
+        self.ImageToolsManager = ImageToolsManager(state)
+        self.LegendManager = LegendManager(state)
+        self.RangeManager = RangeManager(state)
+
+    @track_kwargs
+    def scatter(
+        self,
+        x: NDArray[np.generic] | list[float],
+        y: NDArray[np.generic] | list[float],
+        check: bool = True,
+        **kwargs: Any,
+    ) -> PathCollection:
+        """Scatter plot for a 2D function (or a 2D slice).
+
+        A simple figure and a single axis can also be created.
+
+        Returns
+        -------
+        - The scatter 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: axis object
+        The axis where to plot the scatter. If not given, the last considered
+        axis will be used.
+        - c: str, default self.color
+            Determines the scatter plot color. If not defined, the program will
+            loop over an array of 6 color which are different for the most
+            common vision deficiencies.
+        - clabel: str, default None
+            Sets the colorbar label.
+        - 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.
+        - edgecolors: str, default None
+            Enables a contouring color for the markers.
+        - 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.
+        - label: str, default None
+            Associates a label to be used for the creation of the legend.
+        - 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'.
+        - legpos: int/str, default None
+            If enabled, creates a legend. This keyword selects the legend
+            location.
+        - marker: {'o', 'v', '^', '<', '>', 'X', ' ', etc.}, default ' '
+            Sets an optional symbol for every point. The default value is no
+            marker (' ').
+        - minorticks: str, default None
+            If not None enables the minor ticks on the plot (for both grid
+            axes).
+        - ms: float, default 3
+            Sets the marker size.
+        - 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'.
+        - 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.
+        - vmax: float
+            The maximum value of the colormap.
+        - vmin: float
+            The minimum value of the colormap.
+        - x (not optional): 1D array
+            The x-axis variable.
+        - 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 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.
+        - y (not optional): 1D array
+            The y-axis variable.
+        - 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 scatter plot of a variable
+
+            >>> I.scatter(x, y)
+
+        - Example #2: Plot a scatter plot of a variable with a colorbar
+
+            >>> I.scatter(x, y, cmap="hot", c=x**2 + y**2, cpos="right")
+
+        """
+        # Convert x and y to numpy arrays (if necessary)
+        x = np.asarray(x)
+        y = np.asarray(y)
+
+        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."
+            )
+        # 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
+        c = kwargs.get("c")
+        # If c is a list convert to array
+        vmin = (
+            kwargs.get("vmin", 0.0)
+            if c is None or isinstance(c, str)
+            else kwargs.get("vmin", np.nanmin(np.asarray(c)))
+        )
+        vmax = (
+            kwargs.get("vmax", 0.0)
+            if c is None or isinstance(c, str)
+            else kwargs.get("vmax", np.nanmax(np.asarray(c)))
+        )
+
+        # 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)
+        self.vlims[nax] = [vmin, vmax, tresh]
+
+        # Set the colorbar scale
+        if not isinstance(c, str) and c is not None:
+            norm = self.ImageToolsManager.set_cscale(cscale, vmin, vmax, tresh)
+            cmap = self.ImageToolsManager.find_cmap(kwargs.get("cmap"))
+        else:
+            norm = None
+            cmap = None
+
+        # Start scatter plot procedure
+        pcm = ax.scatter(
+            x,
+            y,
+            cmap=cmap,
+            norm=norm,
+            c=c,
+            s=kwargs.get("ms", 3),
+            edgecolors=kwargs.get("edgecolors", "none"),
+            alpha=kwargs.get("alpha", 1.0),
+            marker=kwargs.get("marker", "o"),
+        )
+
+        # Creation of the legend
+        self.legpos[nax] = kwargs.get("legpos", self.legpos[nax])
+        if self.legpos[nax] is not None:
+            copy_label = kwargs.get("label")
+            kwargs["label"] = None
+            self.LegendManager.legend(ax, check=False, fromplot=True, **kwargs)
+            kwargs["label"] = copy_label
+
+        # 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