Plot Functions¶

def iso_plot(
    data: pd.DataFrame,
    x: str = "ISOPleasant",
    y: str = "ISOEventful",
    title: str | None = "Soundscapy Plot",
    plot_layers: Literal["scatter", "density", "simple_density"]
    | Sequence[Literal["scatter", "density", "simple_density"]] = (
        "scatter",
        "density",
    ),
    **kwargs,
) -> Axes:
    """
    Plot a soundscape visualization based on the specified metrics using different
    combinations of layers such as scatter, density, or simple density plots.

    The function generates a 2D plot (via Matplotlib Axes object) based on the `x` and
    `y` metrics provided. Users can choose between individual layers or specify a
    combination of the supported plot layers. It supports automatic handling for
    specific layer combinations, such as "scatter + density". The core plotting
    functionality is delegated to other helper functions (`scatter` and `density`).

    Parameters
    ----------
    data
        The dataset containing the metrics to be plotted. Must include the columns
        specified for `x` and `y`.
    x
        The column name within `data` to be used for the x-axis. Defaults to
        "ISOPleasant".
    y
        The column name within `data` to be used for the y-axis. Defaults to
        "ISOEventful".
    title
        The title of the plot. If not specified, defaults to "Soundscapy Plot".
    plot_layers
        Specifies the type or combination of plot layers to generate. Valid options
        include:

         - "scatter": A scatter plot
         - "density": A density plot (without scatter points, unless combined)
         - "simple_density": A simplified density plot (without scatter points,
           unless combined).

        Can be passed as a string (single layer) or as a sequence of strings
        (combination of layers). Defaults to ("scatter", "density").
    **kwargs
        Additional keyword arguments to be passed to the underlying plotting
        functions (`scatter` or `density`). These allow for customization such as
        marker styles, colors, etc.

    Returns
    -------
    :
        A Matplotlib Axes object corresponding to the generated plot.

    Notes
    -----
    This function supports only specific combinations of layers. If an unsupported
    combination is specified, an exception will be raised. Layer compatibility
    rules:

      - Single layers: "scatter", "density", or "simple_density".
      - Dual layers:
        - "scatter" + "density"
        - "scatter" + "simple_density"

    Raises
    ------
    TypeError
        If the `plot_layers` argument is not a string or a sequence of strings.
    ValueError
        If the `plot_layers` argument specifies an unsupported plot type or
        combination of plot layers.

    Examples
    --------
    Basic density and scatter plot with default settings:

    >>> import soundscapy as sspy
    >>> import matplotlib.pyplot as plt
    >>> data = sspy.isd.load()
    >>> data = sspy.add_iso_coords(data)
    >>> ax = sspy.iso_plot(data)
    >>> plt.show()  # doctest: +SKIP

    Basic scatter plot:

    >>> ax = sspy.iso_plot(data, plot_layers="scatter")
    >>> plt.show()  # doctest: +SKIP

    Simple density plot with fewer contour levels:

    >>> ax = sspy.iso_plot(data, plot_layers="simple_density")
    >>> plt.show() # doctest: +SKIP

    Simple density with scatter points

    >>> ax = sspy.iso_plot(data, plot_layers=["scatter", "simple_density"])
    >>> plt.show() # doctest: +SKIP

    Density plot with custom styling:

    >>> sub_data = sspy.isd.select_location_ids(
    ...    data, ['CamdenTown', 'PancrasLock', 'RegentsParkJapan', 'RegentsParkFields'])
    >>> ax = sspy.iso_plot(
    ...     sub_data,
    ...     hue="SessionID",
    ...     plot_layers = ["scatter", "simple_density"],
    ...     legend_loc="upper right",
    ...     fill = False,
    ... )
    >>> plt.show() # doctest: +SKIP

    Add density to existing plots:

    >>> fig, axes = plt.subplots(1, 2, figsize=(12, 6))
    >>> axes[0] = sspy.iso_plot(
    ...     sspy.isd.select_location_ids(data, ['CamdenTown', 'PancrasLock']),
    ...     ax=axes.flatten()[0], title="CamdenTown and PancrasLock", hue="LocationID",
    ... )
    >>> axes[1] = sspy.iso_plot(
    ...     sspy.isd.select_location_ids(data, ['RegentsParkJapan']),
    ...     ax=axes.flatten()[1], title="RegentsParkJapan"
    ... )
    >>> plt.tight_layout()
    >>> plt.show() # doctest: +SKIP
    >>> plt.close('all')

    """  # noqa: D205
    if isinstance(plot_layers, str):
        plot_layers = [plot_layers]

    if not isinstance(plot_layers, Sequence):
        raise TypeError(PLOT_LAYER_TYPE_ERROR.format(type=type(plot_layers)))

    # Handle single layer case
    if len(plot_layers) == 1:
        layer_type = plot_layers[0]
        if layer_type == "scatter":
            return scatter(data, x=x, y=y, title=title, **kwargs)
        if layer_type == "simple_density":
            return density(
                data,
                x=x,
                y=y,
                title=title,
                density_type="simple",
                incl_scatter=False,
                **kwargs,
            )
        if layer_type == "density":
            return density(data, x=x, y=y, title=title, incl_scatter=False, **kwargs)

        raise ValueError(PLOT_LAYER_VALUE_ERROR.format(layers=plot_layers))

    # Handle two layer case
    if len(plot_layers) == 2:  # noqa: PLR2004
        layers_set = set(plot_layers)

        if "scatter" in layers_set and "density" in layers_set:
            return density(data, x=x, y=y, title=title, incl_scatter=True, **kwargs)

        if "scatter" in layers_set and "simple_density" in layers_set:
            return density(
                data,
                x=x,
                y=y,
                title=title,
                density_type="simple",
                incl_scatter=True,
                **kwargs,
            )

        # Default case for unrecognized but valid length combinations
        return density(data, x=x, y=y, title=title, incl_scatter=True, **kwargs)

    # More than 2 layers is not supported
    raise ValueError(PLOT_LAYER_VALUE_ERROR.format(layers=plot_layers))

def create_iso_subplots(
    data: pd.DataFrame | list[pd.DataFrame],
    x: str = "ISOPleasant",
    y: str = "ISOEventful",
    subplot_by: str | None = None,
    title: str | None = "Soundscapy Plot",
    plot_layers: Literal["scatter", "density", "simple_density"]
    | Sequence[Literal["scatter", "simple_density", "density"]] = (
        "scatter",
        "density",
    ),
    *,
    subplot_size: tuple[int, int] = (4, 4),
    subplot_titles: Literal["by_group", "numbered"] | list[str] | None = "by_group",
    subplot_title_prefix: str = "Plot",  # Only used if subplot_titles = 'numbered'
    nrows: int | None = None,
    ncols: int | None = None,
    **kwargs,
) -> tuple[Figure, np.ndarray]:
    """
    Create a set of subplots displaying data visualizations for soundscape analysis.

    This function generates a collection of subplots, where each subplot corresponds
    to a subset of the input data. The subplots can display scatter plots, density
    plots, or simplified density plots, and can be organized by specific grouping
    criteria. Users can specify titles, overall size, row and column layout, and
    layering of plot types.

    Parameters
    ----------
    data
        Input data to be visualized. Can be a single data frame or a list of data
        frames for use in multiple subplots.
    x
        The name of the column in the data to be used for the x-axis. Default is
        "ISOPleasant".
    y
        The name of the column in the data to be used for the y-axis. Default is
        "ISOEventful".
    subplot_by
        The column name by which to group data into subplots. If None, data is not
        grouped and plotted in a single set of axes. Default is None.
    title
        The overarching title of the figure. If None, no overall title is added.
        Default is "Soundscapy Plot".
    plot_layers
        Type(s) of plot layers to include in each subplot. Can be a single type
        or a sequence of types. Default is ("scatter", "density").
    subplot_size
        Size of each subplot in inches as (width, height). Default is (4, 4).
    subplot_titles
        Determines how subplot titles are assigned. Options are "by_group" (titles
        derived from group names), "numbered" (titles as indices), or a list of
        custom titles. If None, no titles are added. Default is "by_group".
    subplot_title_prefix
        Prefix for subplot titles if "numbered" is selected as `subplot_titles`.
        Default is "Plot".
    nrows
        Number of rows for the subplot grid. If None, automatically calculated
        based on the number of subplots. Default is None.
    ncols
        Number of columns for the subplot grid. If None, automatically calculated
        based on the number of subplots. Default is None.
    **kwargs
        Additional keyword arguments to pass to matplotlib's `plt.subplots` or for
        customizing the figure and subplots.

    Returns
    -------
    :
        A tuple containing:

        - fig : matplotlib.figure.Figure
            The created matplotlib figure object containing the subplots.

        - np.ndarray
            An array of matplotlib.axes.Axes objects corresponding to the subplots.

    Examples
    --------
    Basic subplots with default settings:
    >>> import soundscapy as sspy
    >>> import matplotlib.pyplot as plt
    >>> import pandas as pd
    >>> data = sspy.isd.load()
    >>> data = sspy.add_iso_coords(data)
    >>> four_locs = sspy.isd.select_location_ids(data,
    ...     ['CamdenTown', 'PancrasLock', 'RegentsParkJapan', 'RegentsParkFields']
    ... )
    >>> fig, axes = sspy.create_iso_subplots(four_locs, subplot_by="LocationID")
    >>> plt.show() # doctest: +SKIP

    Create subplots by specifying a list of data
    >>> data1 = pd.DataFrame({'ISOPleasant': np.random.uniform(-1, 1, 50),
    ...                       'ISOEventful': np.random.uniform(-1, 1, 50)})
    >>> data2 = pd.DataFrame({'ISOPleasant': np.random.uniform(-1, 1, 50),
    ...                       'ISOEventful': np.random.uniform(-1, 1, 50)})
    >>> fig, axes = create_iso_subplots(
    ...     [data1, data2], plot_layers="scatter", nrows=1, ncols=2
    ... )
    >>> plt.show() # doctest: +SKIP
    >>> assert len(axes) == 2
    >>> plt.close('all')

    """
    # Process input data and prepare for subplot creation
    data_list, subplot_titles_list, n_subplots = _prepare_subplot_data(
        data=data, x=y, y=y, subplot_by=subplot_by, subplot_titles=subplot_titles
    )

    # Calculate subplot layout
    nrows, ncols, n_subplots = allocate_subplot_axes(nrows, ncols, n_subplots)

    # Set up figure and subplots
    vert_adjust = 1.2 if title else 1.0
    figsize = kwargs.pop(
        "figsize", (ncols * subplot_size[0], nrows * (vert_adjust * subplot_size[1]))
    )

    subplots_params = SubplotsParams()
    subplots_params.update(
        nrows=nrows,
        ncols=ncols,
        figsize=figsize,
        subplot_by=subplot_by,
        extra="ignore",
        **kwargs,
    )

    fig, axes = plt.subplots(**subplots_params.as_plt_subplots_args())

    # Create each subplot
    _create_subplots(
        data_list,
        axes,
        n_subplots,
        subplot_titles_list,
        x,
        y,
        plot_layers,
        subplot_title_prefix,
        **kwargs,
    )

    # Add overall title and adjust layout
    if title:
        fig.suptitle(title, fontsize=DEFAULT_STYLE_PARAMS["title_fontsize"])

    fig.tight_layout()

    return fig, axes

def allocate_subplot_axes(
    nrows: int | None, ncols: int | None, n_subplots: int
) -> tuple[int, int, int]:
    """
    Allocate the subplot axes based on the number of data subsets.

    Parameters
    ----------
    nrows
        Number of rows for subplots. Can be None to auto-determine.
    ncols
        Number of columns for subplots. Can be None to auto-determine.
    n_subplots
        Total number of subplots needed.

    Returns
    -------
    :
        The number of rows and columns for the subplot grid.

    Notes
    -----
    Logic for determining subplot layout:

    - If both nrows and ncols are specified, use those values
    - If both are None, calculate a grid as close to square as possible
    - If only one is specified, calculate the other to fit all subplots

    """
    # If both dimensions are specified, return them as is
    if nrows is not None and ncols is not None:
        return nrows, ncols, n_subplots

    # If both dimensions are None, create a grid as close to square as possible
    if nrows is None and ncols is None:
        ncols = int(np.ceil(np.sqrt(n_subplots)))
        nrows = int(np.ceil(n_subplots / ncols))
    # If only one dimension is specified, calculate the other
    elif nrows is not None and ncols is None:
        ncols = int(np.ceil(n_subplots / nrows))
    elif nrows is None and ncols is not None:
        nrows = int(np.ceil(n_subplots / ncols))
    else:
        msg = (
            "We should never reach this point - both nrows and ncols are None, "
            "but were missed in earlier check."
        )
        raise ValueError(msg)

    n_subplots = nrows * ncols

    return nrows, ncols, n_subplots

def scatter(
    data: pd.DataFrame,
    title: str | None = "Soundscape Scatter Plot",
    ax: Axes | None = None,
    *,
    x: str | None = "ISOPleasant",
    y: str | None = "ISOEventful",
    hue: str | None = None,
    palette: SeabornPaletteType | None = "colorblind",
    legend: Literal["auto", "brief", "full", False] = "auto",
    prim_labels: bool | None = None,  # Alias for primary_labels, deprecated
    **kwargs,
) -> Axes:
    """
    Plot ISOcoordinates as scatter points on a soundscape circumplex grid.

    Creates a scatter plot of data on a standardized circumplex grid with the custom
    Soundscapy styling for soundscape circumplex visualisations.

    Parameters
    ----------
    data
        Input data structure containing coordinate data, typically with ISOPleasant
        and ISOEventful columns.
    x
        Column name for x variable, by default "ISOPleasant"
    y
        Column name for y variable, by default "ISOEventful"
    title
        Title to add to circumplex plot, by default "Soundscape Scatter Plot"
    ax
        Pre-existing matplotlib axes for the plot, by default None
        If `None` call `matplotlib.pyplot.subplots` with `figsize` internally.
    hue
        Grouping variable that will produce points with different colors.

        Can be either categorical or numeric,
        although color mapping will behave differently in latter case, by default None
    palette
        Method for choosing the colors to use when mapping the hue semantic.
        String values are passed to seaborn.color_palette().
        List or dict values imply categorical mapping, while a colormap object
        implies numeric mapping, by default "colorblind"
    legend
        How to draw the legend. If "brief", numeric hue and size variables will be
        represented with a sample of evenly spaced values. If "full", every group will
        get an entry in the legend. If "auto", choose between brief or full
        representation based on number of levels.

        If False, no legend data is added and no legend is drawn, by default "auto"
    prim_labels
        Deprecated. Use xlabel and ylabel parameters instead.

    **kwargs
        Additional style arguments. Common options include `color` (default
        `"#0173B2"` when hue mapping is unused), `figsize` (default `(5, 5)` when
        creating a new figure), `s` (default `20` for point size), axis label and
        limit controls such as `xlabel`, `ylabel`, `xlim`, `ylim`, legend placement
        via `legend_loc`, `diagonal_lines`, and font settings such as
        `prim_ax_fontdict`, `fontsize`, `fontweight`, `fontstyle`, `family`, `c`,
        `alpha`, and `parse_math`.

    Returns
    -------
    :
        Axes object containing the plot.

    Notes
    -----
    This function applies special styling appropriate for circumplex plots including
    gridlines, axis labels, and proportional axes.

    Examples
    --------
    Basic scatter plot with default settings:

    >>> import soundscapy as sspy
    >>> import matplotlib.pyplot as plt
    >>> data = sspy.isd.load()
    >>> data = sspy.add_iso_coords(data)
    >>> ax = sspy.scatter(data)
    >>> plt.show() # doctest: +SKIP

    Scatter plot with grouping by location:

    >>> ax = sspy.scatter(data, hue="LocationID", diagonal_lines=True, legend=False)
    >>> plt.show() # doctest: +SKIP
    >>> plt.close('all')

    """
    style_args, subplots_args, kwargs = _setup_style_and_subplots_args_from_kwargs(
        x=x, y=y, prim_labels=prim_labels, kwargs=kwargs
    )

    scatter_args = ScatterParams()
    scatter_args.update(
        data=data,
        x=x,
        y=y,
        palette=palette,
        hue=hue,
        legend=legend,
        extra="allow",
        ignore_null=False,
        **kwargs,
    )  # pass all the rest to scatter

    # Removes the palette if no hue is specified
    scatter_args.crosscheck_palette_hue()

    if ax is None:
        _, ax = plt.subplots(1, 1, figsize=subplots_args.figsize)

    p = sns.scatterplot(ax=ax, **scatter_args.as_dict())

    _set_style()
    _circumplex_grid(
        ax=ax,
        **style_args.get_multiple(
            ["xlim", "ylim", "xlabel", "ylabel", "diagonal_lines", "prim_ax_fontdict"]
        ),
    )
    if title is not None:
        _set_circum_title(
            ax=ax,
            title=title,
            xlabel=style_args.get("xlabel"),
            ylabel=style_args.get("ylabel"),
        )
    if legend is not None and hue is not None and style_args.legend_loc is not False:
        _move_legend(ax=ax, new_loc=style_args.get("legend_loc"))
    return p

def density(
    data: pd.DataFrame,
    title: str | None = "Soundscape Density Plot",
    ax: Axes | None = None,
    *,
    x: str | None = "ISOPleasant",
    y: str | None = "ISOEventful",
    hue: str | None = None,
    incl_scatter: bool = True,
    density_type: str = "full",
    palette: SeabornPaletteType | None = "colorblind",
    scatter_kws: dict | None = None,
    legend: Literal["auto", "brief", "full", False] = "auto",
    prim_labels: bool | None = None,  # Alias for primary_labels, deprecated
    **kwargs,
) -> Axes:
    """
    Plot a density plot of ISOCoordinates.

    Creates a kernel density estimate visualization of data distribution on a
    circumplex grid with the custom Soundscapy styling for soundscape circumplex
    visualisations. Can optionally include a scatter plot of the underlying data points.

    Parameters
    ----------
    data
        Input data structure containing coordinate data, typically with ISOPleasant
        and ISOEventful columns.
    title
        Title to add to circumplex plot, by default "Soundscape Density Plot"
    ax
        Pre-existing axes object to use for the plot, by default None
        If `None` call `matplotlib.pyplot.subplots` with `figsize` internally.
    x
        Column name for x variable, by default "ISOPleasant"
    y
        Column name for y variable, by default "ISOEventful"
    hue
        Grouping variable that will produce density contours with different colors.
        Can be either categorical or numeric, although color mapping will behave
        differently in latter case, by default None
    incl_scatter
        Whether to include a scatter plot of the data points, by default True
    density_type
        Type of density plot to draw. "full" uses default parameters, "simple"
        uses a lower number of levels (2), higher threshold (0.5), and lower alpha (0.5)
        for a cleaner visualization, by default "full"
    palette
        Method for choosing the colors to use when mapping the hue semantic.
        String values are passed to seaborn.color_palette().
        List or dict values imply categorical mapping, while a colormap object
        implies numeric mapping, by default "colorblind"
    scatter_kws
        Keyword arguments to pass to `seaborn.scatterplot` if incl_scatter is True,
        by default {"s": 25, "linewidth": 0}
    legend
        How to draw the legend. If "brief", numeric hue variables will be
        represented with a sample of evenly spaced values. If "full", every group will
        get an entry in the legend. If "auto", choose between brief or full
        representation based on number of levels.
        If False, no legend data is added and no legend is drawn, by default "auto"
    prim_labels
        Deprecated. Use xlabel and ylabel parameters instead.

    **kwargs
        Additional styling parameters. Common options include `incl_outline`,
        density controls such as `alpha`, `fill`, `levels`, `thresh`, and
        `bw_adjust`, axis label and limit controls such as `xlabel`, `ylabel`,
        `xlim`, `ylim`, legend placement via `legend_loc`, `diagonal_lines`,
        `figsize` for newly created figures, font settings such as
        `prim_ax_fontdict`, `fontsize`, `fontweight`, `fontstyle`, `family`, `c`,
        `alpha`, and `parse_math`, plus any extra keyword arguments accepted by
        matplotlib's `contour` and `contourf`.

    Returns
    -------
    :
        Axes object containing the plot.

    Notes
    -----
    This function will raise a warning if the dataset has fewer than
    RECOMMENDED_MIN_SAMPLES (30) data points, as density plots are not reliable
    with small sample sizes.

    Examples
    --------
    Basic density plot with default settings:

    >>> import soundscapy as sspy
    >>> import matplotlib.pyplot as plt
    >>> data = sspy.isd.load()
    >>> data = sspy.add_iso_coords(data)
    >>> ax = sspy.density(data)
    >>> plt.show() # doctest: +SKIP

    Simple density plot with fewer contour levels:

    >>> ax = sspy.density(data, density_type="simple")
    >>> plt.show() # doctest: +SKIP

    Density plot with custom styling:

    >>> sub_data = sspy.isd.select_location_ids(
    ...    data, ['CamdenTown', 'PancrasLock', 'RegentsParkJapan', 'RegentsParkFields'])
    >>> ax = sspy.density(
    ...     sub_data,
    ...     hue="SessionID",
    ...     incl_scatter=True,
    ...     legend_loc="upper right",
    ...     fill = False,
    ...     density_type = "simple",
    ... )
    >>> plt.show() # doctest: +SKIP

    Add density to existing plots:

    >>> fig, axes = plt.subplots(1, 2, figsize=(12, 6))
    >>> axes[0] = sspy.density(
    ...     sspy.isd.select_location_ids(data, ['CamdenTown', 'PancrasLock']),
    ...     ax=axes[0], title="CamdenTown and PancrasLock", hue="LocationID",
    ...     density_type="simple"
    ... )
    >>> axes[1] = sspy.density(
    ...     sspy.isd.select_location_ids(data, ['RegentsParkJapan']),
    ...     ax=axes[1], title="RegentsParkJapan"
    ... )
    >>> plt.tight_layout()
    >>> plt.show() # doctest: +SKIP
    >>> plt.close('all')

    """
    style_args, subplots_args, kwargs = _setup_style_and_subplots_args_from_kwargs(
        x=x, y=y, prim_labels=prim_labels, kwargs=kwargs
    )

    # Set up density parameters
    density_args = _setup_density_params(
        data=data,
        x=x,
        y=y,
        hue=hue,
        density_type=density_type,
        palette=palette,
        legend=legend,
        **kwargs,
    )

    # Check if dataset is large enough for density plots
    _valid_density(data)

    if ax is None:
        _, ax = plt.subplots(1, 1, figsize=subplots_args.get("figsize"))

    # Removes the palette if no hue is specified
    if density_args.get("hue") is None:
        density_args.update(palette=None)

    # Set up scatter parameters if needed
    scatter_args = ScatterParams()
    scatter_args.update(
        data=data,
        x=x,
        y=y,
        palette=palette,
        hue=density_args.get("hue"),
        color=density_args.get("color"),
        **(scatter_kws or {}),
    )

    scatter_args.crosscheck_palette_hue()
    density_args.crosscheck_palette_hue()

    if incl_scatter:
        d = sns.scatterplot(ax=ax, **scatter_args.as_seaborn_kwargs())

    if density_type == "simple":
        d = sns.kdeplot(ax=ax, **density_args.as_seaborn_kwargs())
        d = sns.kdeplot(ax=ax, **density_args.to_outline().as_seaborn_kwargs())

    elif density_type == "full":
        d = sns.kdeplot(ax=ax, **density_args.as_seaborn_kwargs())
    else:
        raise ValueError

    _set_style()
    _circumplex_grid(
        ax=ax,
        xlim=style_args.get("xlim"),
        ylim=style_args.get("ylim"),
        xlabel=style_args.get("xlabel"),
        ylabel=style_args.get("ylabel"),
        diagonal_lines=style_args.get("diagonal_lines"),
        prim_ax_fontdict=style_args.get("prim_ax_fontdict"),
    )
    if title is not None:
        _set_circum_title(
            ax=ax,
            title=title,
            xlabel=style_args.get("xlabel"),
            ylabel=style_args.get("ylabel"),
        )
    if legend is not None and hue is not None:
        _move_legend(ax=ax, new_loc=style_args.get("legend_loc"))

    return d