array.py

"""
Standalone function for plotting a 2D array (image) directly with matplotlib.
"""

import os
from typing import List, Optional, Tuple

import numpy as np
from autoarray.plot.utils import (
    _FAST_PLOTS,
    subplots,
    apply_extent,
    apply_labels,
    conf_figsize,
    save_figure,
    zoom_array,
    auto_mask_edge,
    numpy_grid,
    numpy_lines,
    numpy_positions,
    norm_from,
    _apply_contours,
    _conf_imshow_origin,
    _overlay_yx_for_origin,
    _vector_yx_for_origin,
    plot_regions,
)

_zoom_array_2d = zoom_array
_mask_edge_coords = auto_mask_edge


def plot_array(
    array,
    ax=None,
    # --- spatial metadata -------------------------------------------------------
    extent: Optional[Tuple[float, float, float, float]] = None,
    # --- overlays ---------------------------------------------------------------
    mask: Optional[np.ndarray] = None,
    border=None,
    origin=None,
    grid=None,
    mesh_grid=None,
    positions=None,
    lines=None,
    vector_yx: Optional[np.ndarray] = None,
    array_overlay=None,
    patches: Optional[List] = None,
    fill_region: Optional[List] = None,
    contours: Optional[int] = None,
    regions: Optional[List] = None,
    region_colors: Optional[List] = None,
    region_alpha: float = 0.25,
    region_labels: Optional[List[str]] = None,
    # --- cosmetics --------------------------------------------------------------
    title: str = "",
    xlabel: str = "",
    ylabel: str = "",
    colormap: Optional[str] = None,
    vmin: Optional[float] = None,
    vmax: Optional[float] = None,
    use_log10: bool = False,
    cb_unit: Optional[str] = None,
    line_colors: Optional[List] = None,
    origin_imshow: Optional[str] = None,
    arcsec_after_decimal: Optional[bool] = None,
    # --- figure control (used only when ax is None) -----------------------------
    figsize: Optional[Tuple[int, int]] = None,
    output_path: Optional[str] = None,
    output_filename: str = "array",
    output_format: str = None,
    stretch: Optional[float] = None,
) -> None:
    """
    Plot a 2D array (image) using ``plt.imshow``.

    Parameters
    ----------
    array
        2D numpy array of pixel values.
    ax
        Existing matplotlib ``Axes`` to draw onto.  If ``None`` a new figure
        is created and saved / shown according to *output_path*.
    extent
        ``[xmin, xmax, ymin, ymax]`` spatial extent in data coordinates.
    mask
        Array of shape ``(N, 2)`` with ``(y, x)`` coordinates of masked
        pixels to overlay as black dots (auto-derived from array.mask by caller).
    border
        Array of shape ``(N, 2)`` with ``(y, x)`` border pixel coordinates.
    origin
        ``(y, x)`` origin coordinate(s) to scatter as a marker.
    grid
        Array of shape ``(N, 2)`` with ``(y, x)`` coordinates to scatter.
    mesh_grid
        Array of shape ``(N, 2)`` mesh grid coordinates to scatter.
    positions
        List of ``(N, 2)`` arrays; each is scattered as a distinct group.
    lines
        List of ``(N, 2)`` arrays with ``(y, x)`` columns to plot as lines.
    vector_yx
        Array of shape ``(N, 4)`` — ``(y, x, vy, vx)`` — plotted as quiver.
    array_overlay
        A second 2D array rendered on top of *array* with partial alpha.
    patches
        List of matplotlib ``Patch`` objects to draw over the image.
    fill_region
        List of two arrays ``[y1_arr, y2_arr]`` passed to ``ax.fill_between``.
    regions
        List of regions to overlay as filled polygons, each region being a list of ``(N, 2)``
        arrays of ``(y, x)`` coordinates (e.g. the ``contours`` of an ``ImageRegion``).  Every
        region is drawn in its own colour, so that the same colour identifies the same source
        structure across figures.
    region_colors
        The colour of each region, cycled if there are more regions than colours.  ``None`` uses
        ``["r", "g", "b", "m", "c", "y"]``.
    region_alpha
        The alpha of each region's fill; the outline is always drawn opaque.
    region_labels
        A text label drawn at the centre of each region, e.g. ``["1", "2"]``.
    title
        Figure title string.
    xlabel, ylabel
        Axis label strings.
    colormap
        Matplotlib colormap name.
    vmin, vmax
        Explicit color scale limits.
    use_log10
        When ``True`` a ``LogNorm`` is applied.
    stretch
        Optional positive asinh linear-width fraction of the displayed range,
        e.g. 0.01 to reveal faint structure beside a bright peak. Only changes
        the colour mapping; incompatible with ``use_log10=True``.
    origin_imshow
        Passed directly to ``imshow`` (``"upper"`` or ``"lower"``).
    arcsec_after_decimal
        Per-call override of the ``ticks.symbol_over_decimal`` config flag.
        When ``True`` the arcsecond tick labels place the ``″`` symbol over the
        decimal point (``3.″8``) instead of suffixing it (``3.8"``); when ``False``
        the suffix form is forced.  ``None`` (the default) reads the config.
    figsize
        Figure size in inches.
    output_path
        Directory to save the figure.  When empty / ``None`` ``plt.show()``
        is called instead.
    output_filename
        Base file name (without extension).
    output_format
        File format, e.g. ``"png"``.
    """
    if origin_imshow is None:
        origin_imshow = _conf_imshow_origin()

    # --- autoarray extraction --------------------------------------------------
    array = zoom_array(array)
    try:
        if extent is None:
            extent = array.geometry.extent
        # The mask-edge overlay derives the edge grid of the mask on EVERY call
        # (~0.4s on a 2000x100 CTI frame, times ~100 on-the-fly figures per
        # bypassed fit). PYAUTO_FAST_PLOTS already drops the figure before it
        # is rasterised or saved, so under it the overlay is never seen:
        # skip deriving it (PyAutoBrain /ci_speedup, 2026-09-08).
        if mask is None and not _FAST_PLOTS:
            mask = auto_mask_edge(array)
        array = array.native.array
    except AttributeError:
        array = np.asarray(array)

    if array is None or array.size == 0:
        return

    is_rgb = array.ndim == 3 and array.shape[2] in (3, 4)

    if colormap is None:
        from autoarray.plot.utils import _default_colormap

        colormap = _default_colormap()

    # convert overlay params (safe for None and already-numpy inputs)
    border = numpy_grid(border)
    origin = numpy_grid(origin)
    grid = numpy_grid(grid)
    mesh_grid = numpy_grid(mesh_grid)
    positions = numpy_positions(positions)
    lines = numpy_lines(lines)
    if array_overlay is not None:
        try:
            array_overlay = array_overlay.native.array
        except AttributeError:
            array_overlay = np.asarray(array_overlay)

    owns_figure = ax is None
    if owns_figure:
        figsize = figsize or conf_figsize("figures")
        fig, ax = subplots(1, 1, figsize=figsize)
    else:
        fig = ax.get_figure()

    # --- colour normalisation --------------------------------------------------
    norm = norm_from(array=array, use_log10=use_log10, vmin=vmin, vmax=vmax, stretch=stretch)

    # Compute the axes-box aspect ratio from the data extent so that the
    # physical cell is correctly shaped and tight_layout has no whitespace
    # to absorb.  This reproduces the old "square" subplot behaviour where
    # ratio = x_range / y_range was passed to plt.subplot(aspect=ratio).
    if extent is not None:
        x_range = abs(extent[1] - extent[0])
        y_range = abs(extent[3] - extent[2])
        _box_aspect = (x_range / y_range) if y_range > 0 else 1.0
    else:
        h, w = array.shape[:2]
        _box_aspect = (w / h) if h > 0 else 1.0

    if is_rgb:
        im = ax.imshow(
            array,
            extent=extent,
            aspect="auto",
            origin=origin_imshow,
        )
    else:
        im = ax.imshow(
            array,
            cmap=colormap,
            norm=norm,
            extent=extent,
            aspect="auto",  # image fills the axes box; box shape set below
            origin=origin_imshow,
        )

    # Shape the axes box to match the data so there is no surrounding
    # whitespace when the panel is embedded in a subplot grid.
    ax.set_aspect(_box_aspect, adjustable="box")

    if not is_rgb:
        from autoarray.plot.utils import _apply_colorbar

        _apply_colorbar(im, ax, cb_unit=cb_unit, is_subplot=not owns_figure)

    # --- overlays --------------------------------------------------------------
    # Every (y, x) overlay is reflected for display when ``origin_imshow`` is
    # ``"lower"`` so it stays registered to the raster (PyAutoArray#565). When
    # ``extent`` is ``None`` the raster is drawn in pixel-index coordinates,
    # where matplotlib already places row ``i`` at ``y = i`` under either
    # origin, so no reflection is applied.
    def _for_origin(yx):
        return _overlay_yx_for_origin(yx, extent, origin_imshow)

    mask = _for_origin(mask)
    border = _for_origin(border)
    origin = _for_origin(origin)
    grid = _for_origin(grid)
    mesh_grid = _for_origin(mesh_grid)
    if positions is not None:
        positions = [_for_origin(np.asarray(pos).reshape(-1, 2)) for pos in positions]
    if lines is not None:
        lines = [
            (
                _for_origin(np.asarray(line).reshape(-1, 2))
                if line is not None and len(line) > 0
                else line
            )
            for line in lines
        ]
    if regions is not None:
        regions = [
            (
                None
                if region is None
                else (
                    _for_origin(region)
                    if isinstance(region, np.ndarray) and region.ndim == 2
                    else [_for_origin(polygon) for polygon in region]
                )
            )
            for region in regions
        ]
    vector_yx = _vector_yx_for_origin(vector_yx, extent, origin_imshow)

    if array_overlay is not None:
        ax.imshow(
            array_overlay,
            cmap="Greys",
            alpha=0.5,
            extent=extent,
            aspect="auto",
            origin=origin_imshow,
        )

    if mask is not None:
        ax.scatter(mask[:, 1], mask[:, 0], s=1, c="k")

    if border is not None:
        ax.scatter(border[:, 1], border[:, 0], s=1, c="b")

    if origin is not None:
        origin_arr = np.asarray(origin)
        if origin_arr.ndim == 1:
            origin_arr = origin_arr[np.newaxis, :]
        ax.scatter(
            origin_arr[:, 1], origin_arr[:, 0], s=20, c="r", marker="x", zorder=6
        )

    if grid is not None:
        ax.scatter(grid[:, 1], grid[:, 0], s=1, c="k")

    if mesh_grid is not None:
        ax.scatter(mesh_grid[:, 1], mesh_grid[:, 0], s=1, c="w", alpha=0.5)

    if positions is not None:
        colors = ["k", "g", "b", "m", "c", "y"]
        for i, pos in enumerate(positions):
            pos = np.asarray(pos).reshape(-1, 2)
            ax.scatter(pos[:, 1], pos[:, 0], s=20, c=colors[i % len(colors)], zorder=5)

    if lines is not None:
        for i, line in enumerate(lines):
            if line is not None and len(line) > 0:
                line = np.asarray(line).reshape(-1, 2)
                color = (
                    line_colors[i]
                    if (line_colors is not None and i < len(line_colors))
                    else None
                )
                kw = {"linewidth": 1}
                if color is not None:
                    kw["color"] = color
                ax.plot(line[:, 1], line[:, 0], **kw)

    plot_regions(
        ax,
        regions=regions,
        region_colors=region_colors,
        region_alpha=region_alpha,
        region_labels=region_labels,
    )

    if vector_yx is not None:
        ax.quiver(
            vector_yx[:, 1],
            vector_yx[:, 0],
            vector_yx[:, 3],
            vector_yx[:, 2],
        )

    # Patches are drawn as given: they are not origin-aware.
    if patches is not None:
        for patch in patches:
            import copy

            ax.add_patch(copy.copy(patch))

    if fill_region is not None:
        y1, y2 = fill_region[0], fill_region[1]
        x_fill = np.arange(len(y1))
        ax.fill_between(x_fill, y1, y2, alpha=0.3)

    # Contours: auto-enabled for log10 plots; explicit int enables linear contours.
    if use_log10 or (contours is not None and contours > 0):
        _apply_contours(
            ax,
            array,
            extent,
            use_log10=use_log10,
            n=contours if (contours is not None and contours > 0) else None,
        )

    # --- labels / ticks --------------------------------------------------------
    apply_labels(
        ax, title=title, xlabel=xlabel, ylabel=ylabel, is_subplot=not owns_figure
    )

    if extent is not None:
        apply_extent(ax, extent, symbol_over_decimal=arcsec_after_decimal)

    # --- output ----------------------------------------------------------------
    if owns_figure:
        save_figure(
            fig,
            path=output_path or "",
            filename=output_filename,
            format=output_format,
        )