<reponame>Xudewang/scarlet
import numpy as np
from astropy.visualization.lupton_rgb import LinearMapping, AsinhMapping
import matplotlib.pyplot as plt
from matplotlib.patches import Rectangle, Polygon
from matplotlib.ticker import MaxNLocator
from .bbox import Box
from .component import Component


def channels_to_rgb(channels):
    """Get the linear mapping of multiple channels to RGB channels
    The mapping created here assumes the the channels are ordered in wavelength
    direction, starting with the shortest wavelength. The mapping seeks to produce
    a relatively even weights for across all channels. It does not consider e.g.
    signal-to-noise variations across channels or human perception.
    Parameters
    ----------
    channels: int in range(0,7)
        Number of channels
    Returns
    -------
    array (3, channels) to map onto RGB
    """
    assert channels in range(
        0, 8
    ), "No mapping has been implemented for more than {} channels".format(channels)

    channel_map = np.zeros((3, channels))
    if channels == 1:
        channel_map[0, 0] = channel_map[1, 0] = channel_map[2, 0] = 1
    elif channels == 2:
        channel_map[0, 1] = 0.667
        channel_map[1, 1] = 0.333
        channel_map[1, 0] = 0.333
        channel_map[2, 0] = 0.667
        channel_map /= 0.667
    elif channels == 3:
        channel_map[0, 2] = 1
        channel_map[1, 1] = 1
        channel_map[2, 0] = 1
    elif channels == 4:
        channel_map[0, 3] = 1
        channel_map[0, 2] = 0.333
        channel_map[1, 2] = 0.667
        channel_map[1, 1] = 0.667
        channel_map[2, 1] = 0.333
        channel_map[2, 0] = 1
        channel_map /= 1.333
    elif channels == 5:
        channel_map[0, 4] = 1
        channel_map[0, 3] = 0.667
        channel_map[1, 3] = 0.333
        channel_map[1, 2] = 1
        channel_map[1, 1] = 0.333
        channel_map[2, 1] = 0.667
        channel_map[2, 0] = 1
        channel_map /= 1.667
    elif channels == 6:
        channel_map[0, 5] = 1
        channel_map[0, 4] = 0.667
        channel_map[0, 3] = 0.333
        channel_map[1, 4] = 0.333
        channel_map[1, 3] = 0.667
        channel_map[1, 2] = 0.667
        channel_map[1, 1] = 0.333
        channel_map[2, 2] = 0.333
        channel_map[2, 1] = 0.667
        channel_map[2, 0] = 1
        channel_map /= 2
    elif channels == 7:
        channel_map[:, 6] = 2/3.
        channel_map[0, 5] = 1
        channel_map[0, 4] = 0.667
        channel_map[0, 3] = 0.333
        channel_map[1, 4] = 0.333
        channel_map[1, 3] = 0.667
        channel_map[1, 2] = 0.667
        channel_map[1, 1] = 0.333
        channel_map[2, 2] = 0.333
        channel_map[2, 1] = 0.667
        channel_map[2, 0] = 1
        channel_map /= 2
    return channel_map


class LinearPercentileNorm(LinearMapping):
    def __init__(self, img, percentiles=[1, 99]):
        """Create norm that is linear between lower and upper percentile of img
        Parameters
        ----------
        img: array_like
            Image to normalize
        percentile: array_like, default=[1,99]
            Lower and upper percentile to consider. Pixel values below will be
            set to zero, above to saturated.
        """
        assert len(percentiles) == 2
        vmin, vmax = np.percentile(img, percentiles)
        super().__init__(minimum=vmin, maximum=vmax)


class AsinhPercentileNorm(AsinhMapping):
    def __init__(self, img, percentiles=[1, 99]):
        """Create norm that is linear between lower and upper percentile of img
        Parameters
        ----------
        img: array_like
            Image to normalize
        percentile: array_like, default=[1,99]
            Lower and upper percentile to consider. Pixel values below will be
            set to zero, above to saturated.
        """
        assert len(percentiles) == 2
        vmin, vmax = np.percentile(img, percentiles)
        # solution for beta assumes flat spectrum at vmax
        stretch = vmax - vmin
        beta = stretch / np.sinh(1)
        super().__init__(minimum=vmin, stretch=stretch, Q=beta)


def img_to_3channel(img, channel_map=None, fill_value=0):
    """Convert multi-band image cube into 3 RGB channels
    Parameters
    ----------
    img: array_like
        This should be an array with dimensions (channels, height, width).
    channel_map: array_like
        Linear mapping with dimensions (3, channels)
    fill_value: float, default=`0`
        Value to use for any masked pixels.
    Returns
    -------
    RGB: numpy array with dtype float
    """
    # expand single img into cube
    assert len(img.shape) in [2, 3]
    if len(img.shape) == 2:
        ny, nx = img.shape
        img_ = img.reshape(1, ny, nx)
    elif len(img.shape) == 3:
        img_ = img
    C = len(img_)

    # filterWeights: channel x band
    if channel_map is None:
        channel_map = channels_to_rgb(C)
    else:
        assert channel_map.shape == (3, len(img))

    # map channels onto RGB channels
    _, ny, nx = img_.shape
    rgb = np.dot(channel_map, img_.reshape(C, -1)).reshape(3, ny, nx)

    if hasattr(rgb, "mask"):
        rgb = rgb.filled(fill_value)

    return rgb


def img_to_rgb(img, channel_map=None, fill_value=0, norm=None, mask=None):
    """Convert images to normalized RGB.
    If normalized values are outside of the range [0..255], they will be
    truncated such as to preserve the corresponding color.
    Parameters
    ----------
    img: array_like
        This should be an array with dimensions (channels, height, width).
    channel_map: array_like
        Linear mapping with dimensions (3, channels)
    fill_value: float, default=`0`
        Value to use for any masked pixels.
    norm: `scarlet.display.Norm`, default `None`
        Norm to use for mapping in the allowed range [0..255]. If `norm=None`,
        `scarlet.display.LinearPercentileNorm` will be used.
    mask: array_like
        A [0,1] binary mask to apply over the top of the image,
        where pixels with mask==1 are masked out.
    Returns
    -------
    rgb: numpy array with dimensions (3, height, width) and dtype uint8
    """
    RGB = img_to_3channel(img, channel_map=channel_map)
    if norm is None:
        norm = LinearMapping(image=RGB)
    rgb = norm.make_rgb_image(*RGB)
    if mask is not None:
        rgb = np.dstack([rgb, ~mask * 255])
    return rgb


panel_size = 4.0


def show_likelihood(blend, figsize=None, **kwargs):
    fig, ax = plt.subplots(1, 1, figsize=figsize)
    ax.plot(blend.log_likelihood, **kwargs)
    ax.set_xlabel("Iteration")
    ax.xaxis.set_major_locator(MaxNLocator(integer=True))
    ax.set_ylabel("log-Likelihood")
    return fig


def show_observation(
    observation,
    norm=None,
    channel_map=None,
    sky_coords=None,
    show_psf=False,
    add_labels=True,
    figsize=None,
):
    """Plot observation in standardized form.
    """
    panels = 1 if show_psf is False else 2
    if figsize is None:
        figsize = (panel_size * panels, panel_size)
    fig, ax = plt.subplots(1, panels, figsize=figsize)
    if not hasattr(ax, "__iter__"):
        ax = (ax,)

    # Mask any pixels with zero weight in all bands
    mask = np.sum(observation.weights, axis=0) == 0
    # if there are no masked pixels, do not use a mask
    if np.all(mask == 0):
        mask = None

    panel = 0
    extent = get_extent(observation.bbox)
    ax[panel].imshow(
        img_to_rgb(observation.data, norm=norm, channel_map=channel_map, mask=mask),
        extent=extent,
        origin="lower",
    )
    ax[panel].set_title("Observation")

    if add_labels:
        assert sky_coords is not None, "Provide sky_coords for labeled objects"

        for k, center in enumerate(sky_coords):
            center_ = observation.get_pixel(center)
            color = "w" if observation.C > 1 else "r"
            ax[panel].text(*center_[::-1], k, color=color, ha="center", va="center")

    panel += 1
    if show_psf:
        psf_image = np.zeros(observation.data.shape)

        if observation.psf is not None:
            psf_model = observation.psf.get_model()
            # make PSF as bright as the brightest pixel of the observation
            psf_model *= (
                observation.data.mean(axis=0).max() / psf_model.mean(axis=0).max()
            )
            # insert into middle of "blank" observation
            full_box = Box(psf_image.shape)
            shift = tuple(
                psf_image.shape[c] // 2 - psf_model.shape[c] // 2
                for c in range(full_box.D)
            )
            model_box = Box(psf_model.shape) + shift
            model_box.insert_into(psf_image, psf_model)
            # slices = scarlet.box.overlapped_slices
        ax[panel].imshow(img_to_rgb(psf_image, norm=norm), origin="lower")
        ax[panel].set_title("PSF")

    fig.tight_layout()
    return fig


def show_scene(
    sources,
    observation=None,
    norm=None,
    channel_map=None,
    show_model=True,
    show_observed=False,
    show_rendered=False,
    show_residual=False,
    add_labels=True,
    add_boxes=False,
    figsize=None,
    linear=True,
):
    """Plot all sources to recreate the scence.
    The functions provides a fast way of evaluating the quality of the entire model,
    i.e. the combination of all scences that seek to fit the observation.
    Parameters
    ----------
    sources: list of source models
    observation: `~scarlet.Observation`
    norm: norm to compress image intensity to the range [0,255]
    channel_map: array_like
        Linear mapping with dimensions (3, channels)
    show_model: bool
        Whether the model is shown in the model frame
    show_observed: bool
        Whether the observation is shown
    show_rendered: bool
        Whether the model, rendered to match the observation, is shown
    show_residual: bool
        Whether the residuals between rendered model and observation is shown
    add_label: bool
        Whether each source is labeled with its numerical index in the source list
    add_boxes: bool
        Whether each source box is shown
    figsize: matplotlib figsize argument
    linear: bool
        Whether or not to display the scene in a single line (`True`) or
        on multiple lines (`False`).
    Returns
    -------
    matplotlib figure
    """
    if show_observed or show_rendered or show_residual:
        assert (
            observation is not None
        ), "Provide matched observation to show observed frame"

    panels = sum((show_model, show_observed, show_rendered, show_residual))
    if linear:
        if figsize is None:
            figsize = (panel_size * panels, panel_size)
        fig, ax = plt.subplots(1, panels, figsize=figsize)
    else:
        columns = int(np.ceil(panels / 2))
        if figsize is None:
            figsize = (panel_size * columns, panel_size * 2)
        fig = plt.figure(figsize=figsize)
        ax = [fig.add_subplot(2, columns, n + 1) for n in range(panels)]
    if not hasattr(ax, "__iter__"):
        ax = (ax,)

    # Mask any pixels with zero weight in all bands
    if observation is not None:
        mask = np.sum(observation.weights, axis=0) == 0
        # if there are no masked pixels, do not use a mask
        if np.all(mask == 0):
            mask = None

    model_frame = sources[0].frame
    model = np.zeros(model_frame.shape)
    for src in sources:
        model += src.get_model(frame=model_frame)

    panel = 0
    if show_model:
        extent = get_extent(model_frame.bbox)
        ax[panel].imshow(
            img_to_rgb(model, norm=norm, channel_map=channel_map),
            extent=extent,
            origin="lower",
        )
        ax[panel].set_title("Model")
        panel += 1

    if show_rendered or show_residual:
        model = observation.render(model)
        extent = get_extent(observation.bbox)

    if show_rendered:
        ax[panel].imshow(
            img_to_rgb(model, norm=norm, channel_map=channel_map, mask=mask),
            extent=extent,
            origin="lower",
        )
        ax[panel].set_title("Model Rendered")
        panel += 1

    if show_observed:
        ax[panel].imshow(
            img_to_rgb(observation.data, norm=norm, channel_map=channel_map, mask=mask),
            extent=extent,
            origin="lower",
        )
        ax[panel].set_title("Observation")
        panel += 1

    if show_residual:
        residual = observation.data - model
        norm_ = LinearPercentileNorm(residual)
        ax[panel].imshow(
            img_to_rgb(residual, norm=norm_, channel_map=channel_map, mask=mask),
            extent=extent,
            origin="lower",
        )
        ax[panel].set_title("Residual")
        panel += 1

    for k, src in enumerate(sources):
        if add_boxes:
            panel = 0
            box_kwargs = {"facecolor": "none", "edgecolor": "w", "lw": 0.5}
            if show_model:
                extent = get_extent(src.bbox)
                rect = Rectangle(
                    (extent[0], extent[2]),
                    extent[1] - extent[0],
                    extent[3] - extent[2],
                    **box_kwargs
                )
                ax[panel].add_artist(rect)
                panel = 1
            if observation is not None:
                start, stop = src.bbox.start[-2:][::-1], src.bbox.stop[-2:][::-1]
                points = (start, (start[0], stop[1]), stop, (stop[0], start[1]))
                coords = [
                    observation.get_pixel(model_frame.get_sky_coord(p)) for p in points
                ]
                for panel in range(panel, panels):
                    poly = Polygon(coords, closed=True, **box_kwargs)
                    ax[panel].add_artist(poly)

        if add_labels and hasattr(src, "center") and src.center is not None:
            center = src.center
            panel = 0
            if show_model:
                ax[panel].text(*center[::-1], k, color="w", ha="center", va="center")
                panel = 1
            if observation is not None:
                center_ = observation.get_pixel(model_frame.get_sky_coord(center))
                for panel in range(panel, panels):
                    ax[panel].text(
                        *center_[::-1], k, color="w", ha="center", va="center"
                    )

    fig.tight_layout()
    return fig


def get_extent(bbox):
    return [bbox.start[-1], bbox.stop[-1], bbox.start[-2], bbox.stop[-2]]


def show_sources(
    sources,
    observation=None,
    norm=None,
    channel_map=None,
    show_model=True,
    show_observed=False,
    show_rendered=False,
    show_spectrum=True,
    figsize=None,
    model_mask=None,
    add_markers=True,
    add_boxes=False,
):
    """Plot each source individually.
    The functions provides an more detailed inspection of every source in the list.
    Parameters
    ----------
    sources: list of source models
    observation: `~scarlet.Observation`
    norm: norm to compress image intensity to the range [0,255]
    channel_map: array_like
        Linear mapping with dimensions (3, channels)
    show_model: bool
        Whether the model is shown in the model frame
    show_observed: bool
        Whether the observation is shown
    show_rendered: bool
        Whether the model, rendered to match the observation, is shown
    show_spectrum: bool
        Whether source specturm is shown.
        For multi-component sources, spectra are shown separately.
    figsize: matplotlib figsize argument
    model_mask: array
        Mask used to hide pixels in the model only.
    add_markers: bool
        Whether or not to mark the centers of the sources
        with their source number.
    add_boxes: bool
        Whether source boxes are shown
    Returns
    -------
    matplotlib figure
    """
    if show_observed or show_rendered:
        assert (
            observation is not None
        ), "Provide matched observation to show observed frame"

    panels = sum((show_model, show_observed, show_rendered, show_spectrum))
    if figsize is None:
        figsize = (panel_size * panels, panel_size * len(list(sources)))
    fig, ax = plt.subplots(len(list(sources)), panels, figsize=figsize, squeeze=False)

    marker_kwargs = {"mew": 1, "ms": 10}
    box_kwargs = {"facecolor": "none", "edgecolor": "w", "lw": 0.5}

    for k, src in enumerate(sources):

        model_frame = src.frame

        if hasattr(src, "center") and src.center is not None:
            center = np.array(src.center)[::-1]
        else:
            center = None

        if add_boxes:
            start, stop = src.bbox.start[-2:][::-1], src.bbox.stop[-2:][::-1]
            points = (start, (start[0], stop[1]), stop, (stop[0], start[1]))
            box_coords = [
                observation.get_pixel(model_frame.get_sky_coord(p)) for p in points
            ]

        # model in its bbox
        panel = 0
        model = src.get_model()

        if show_model:
            # Show the unrendered model in it's bbox
            extent = get_extent(src.bbox)
            ax[k][panel].imshow(
                img_to_rgb(model, norm=norm, channel_map=channel_map, mask=model_mask),
                extent=extent,
                origin="lower",
            )
            ax[k][panel].set_title("Model Source {}".format(k))
            if center is not None and add_markers:
                ax[k][panel].plot(*center, "wx", **marker_kwargs)
            panel += 1

        # model in observation frame
        if show_rendered:
            # Center and show the rendered model
            model_ = src.get_model(frame=model_frame)
            model_ = observation.render(model_)
            extent = get_extent(observation.bbox)
            ax[k][panel].imshow(
                img_to_rgb(model_, norm=norm, channel_map=channel_map),
                extent=extent,
                origin="lower",
            )
            ax[k][panel].set_title("Model Source {} Rendered".format(k))

            if center is not None and add_markers:
                center_ = observation.get_pixel(model_frame.get_sky_coord(center))
                ax[k][panel].plot(*center_, "wx", **marker_kwargs)
            if add_boxes:
                poly = Polygon(box_coords, closed=True, **box_kwargs)
                ax[k][panel].add_artist(poly)
            panel += 1

        if show_observed:
            # Center the observation on the source and display it
            _images = observation.data
            ax[k][panel].imshow(
                img_to_rgb(_images, norm=norm, channel_map=channel_map),
                extent=extent,
                origin="lower",
            )
            ax[k][panel].set_title("Observation".format(k))
            if center is not None and add_markers:
                center_ = observation.get_pixel(model_frame.get_sky_coord(center))
                ax[k][panel].plot(*center_, "wx", **marker_kwargs)
            if add_boxes:
                poly = Polygon(box_coords, closed=True, **box_kwargs)
                ax[k][panel].add_artist(poly)
            panel += 1

        if show_spectrum:
            # needs to be evaluated in the source box to prevent truncation
            if hasattr(src, "__iter__") and isinstance(src[0], Component):
                spectra = []
                for component in src:
                    model_ = component.get_model()
                    spectra.append(model_.sum(axis=(1, 2)))
            else:
                spectra = [model.sum(axis=(1, 2))]

            for spectrum in spectra:
                ax[k][panel].plot(spectrum)
            ax[k][panel].set_xticks(range(len(spectrum)))
            if hasattr(src.frame, "channels") and src.frame.channels is not None:
                ax[k][panel].set_xticklabels(src.frame.channels)
            ax[k][panel].set_title("Spectrum")
            ax[k][panel].set_xlabel("Channel")
            ax[k][panel].set_ylabel("Intensity")

    fig.tight_layout()
    return fig

import numpy as np
from astropy.visualization.lupton_rgb import LinearMapping, AsinhMapping
import matplotlib.pyplot as plt
from matplotlib.patches import Rectangle, Polygon
from matplotlib.ticker import MaxNLocator
from .bbox import Box
from .component import Component


def channels_to_rgb(channels):
    """Get the linear mapping of multiple channels to RGB channels
    The mapping created here assumes the the channels are ordered in wavelength
    direction, starting with the shortest wavelength. The mapping seeks to produce
    a relatively even weights for across all channels. It does not consider e.g.
    signal-to-noise variations across channels or human perception.
    Parameters
    ----------
    channels: int in range(0,7)
        Number of channels
    Returns
    -------
    array (3, channels) to map onto RGB
    """
    assert channels in range(
        0, 8
    ), "No mapping has been implemented for more than {} channels".format(channels)

    channel_map = np.zeros((3, channels))
    if channels == 1:
        channel_map[0, 0] = channel_map[1, 0] = channel_map[2, 0] = 1
    elif channels == 2:
        channel_map[0, 1] = 0.667
        channel_map[1, 1] = 0.333
        channel_map[1, 0] = 0.333
        channel_map[2, 0] = 0.667
        channel_map /= 0.667
    elif channels == 3:
        channel_map[0, 2] = 1
        channel_map[1, 1] = 1
        channel_map[2, 0] = 1
    elif channels == 4:
        channel_map[0, 3] = 1
        channel_map[0, 2] = 0.333
        channel_map[1, 2] = 0.667
        channel_map[1, 1] = 0.667
        channel_map[2, 1] = 0.333
        channel_map[2, 0] = 1
        channel_map /= 1.333
    elif channels == 5:
        channel_map[0, 4] = 1
        channel_map[0, 3] = 0.667
        channel_map[1, 3] = 0.333
        channel_map[1, 2] = 1
        channel_map[1, 1] = 0.333
        channel_map[2, 1] = 0.667
        channel_map[2, 0] = 1
        channel_map /= 1.667
    elif channels == 6:
        channel_map[0, 5] = 1
        channel_map[0, 4] = 0.667
        channel_map[0, 3] = 0.333
        channel_map[1, 4] = 0.333
        channel_map[1, 3] = 0.667
        channel_map[1, 2] = 0.667
        channel_map[1, 1] = 0.333
        channel_map[2, 2] = 0.333
        channel_map[2, 1] = 0.667
        channel_map[2, 0] = 1
        channel_map /= 2
    elif channels == 7:
        channel_map[:, 6] = 2/3.
        channel_map[0, 5] = 1
        channel_map[0, 4] = 0.667
        channel_map[0, 3] = 0.333
        channel_map[1, 4] = 0.333
        channel_map[1, 3] = 0.667
        channel_map[1, 2] = 0.667
        channel_map[1, 1] = 0.333
        channel_map[2, 2] = 0.333
        channel_map[2, 1] = 0.667
        channel_map[2, 0] = 1
        channel_map /= 2
    return channel_map


class LinearPercentileNorm(LinearMapping):
    def __init__(self, img, percentiles=[1, 99]):
        """Create norm that is linear between lower and upper percentile of img
        Parameters
        ----------
        img: array_like
            Image to normalize
        percentile: array_like, default=[1,99]
            Lower and upper percentile to consider. Pixel values below will be
            set to zero, above to saturated.
        """
        assert len(percentiles) == 2
        vmin, vmax = np.percentile(img, percentiles)
        super().__init__(minimum=vmin, maximum=vmax)


class AsinhPercentileNorm(AsinhMapping):
    def __init__(self, img, percentiles=[1, 99]):
        """Create norm that is linear between lower and upper percentile of img
        Parameters
        ----------
        img: array_like
            Image to normalize
        percentile: array_like, default=[1,99]
            Lower and upper percentile to consider. Pixel values below will be
            set to zero, above to saturated.
        """
        assert len(percentiles) == 2
        vmin, vmax = np.percentile(img, percentiles)
        # solution for beta assumes flat spectrum at vmax
        stretch = vmax - vmin
        beta = stretch / np.sinh(1)
        super().__init__(minimum=vmin, stretch=stretch, Q=beta)


def img_to_3channel(img, channel_map=None, fill_value=0):
    """Convert multi-band image cube into 3 RGB channels
    Parameters
    ----------
    img: array_like
        This should be an array with dimensions (channels, height, width).
    channel_map: array_like
        Linear mapping with dimensions (3, channels)
    fill_value: float, default=`0`
        Value to use for any masked pixels.
    Returns
    -------
    RGB: numpy array with dtype float
    """
    # expand single img into cube
    assert len(img.shape) in [2, 3]
    if len(img.shape) == 2:
        ny, nx = img.shape
        img_ = img.reshape(1, ny, nx)
    elif len(img.shape) == 3:
        img_ = img
    C = len(img_)

    # filterWeights: channel x band
    if channel_map is None:
        channel_map = channels_to_rgb(C)
    else:
        assert channel_map.shape == (3, len(img))

    # map channels onto RGB channels
    _, ny, nx = img_.shape
    rgb = np.dot(channel_map, img_.reshape(C, -1)).reshape(3, ny, nx)

    if hasattr(rgb, "mask"):
        rgb = rgb.filled(fill_value)

    return rgb


def img_to_rgb(img, channel_map=None, fill_value=0, norm=None, mask=None):
    """Convert images to normalized RGB.
    If normalized values are outside of the range [0..255], they will be
    truncated such as to preserve the corresponding color.
    Parameters
    ----------
    img: array_like
        This should be an array with dimensions (channels, height, width).
    channel_map: array_like
        Linear mapping with dimensions (3, channels)
    fill_value: float, default=`0`
        Value to use for any masked pixels.
    norm: `scarlet.display.Norm`, default `None`
        Norm to use for mapping in the allowed range [0..255]. If `norm=None`,
        `scarlet.display.LinearPercentileNorm` will be used.
    mask: array_like
        A [0,1] binary mask to apply over the top of the image,
        where pixels with mask==1 are masked out.
    Returns
    -------
    rgb: numpy array with dimensions (3, height, width) and dtype uint8
    """
    RGB = img_to_3channel(img, channel_map=channel_map)
    if norm is None:
        norm = LinearMapping(image=RGB)
    rgb = norm.make_rgb_image(*RGB)
    if mask is not None:
        rgb = np.dstack([rgb, ~mask * 255])
    return rgb


panel_size = 4.0


def show_likelihood(blend, figsize=None, **kwargs):
    fig, ax = plt.subplots(1, 1, figsize=figsize)
    ax.plot(blend.log_likelihood, **kwargs)
    ax.set_xlabel("Iteration")
    ax.xaxis.set_major_locator(MaxNLocator(integer=True))
    ax.set_ylabel("log-Likelihood")
    return fig


def show_observation(
    observation,
    norm=None,
    channel_map=None,
    sky_coords=None,
    show_psf=False,
    add_labels=True,
    figsize=None,
):
    """Plot observation in standardized form.
    """
    panels = 1 if show_psf is False else 2
    if figsize is None:
        figsize = (panel_size * panels, panel_size)
    fig, ax = plt.subplots(1, panels, figsize=figsize)
    if not hasattr(ax, "__iter__"):
        ax = (ax,)

    # Mask any pixels with zero weight in all bands
    mask = np.sum(observation.weights, axis=0) == 0
    # if there are no masked pixels, do not use a mask
    if np.all(mask == 0):
        mask = None

    panel = 0
    extent = get_extent(observation.bbox)
    ax[panel].imshow(
        img_to_rgb(observation.data, norm=norm, channel_map=channel_map, mask=mask),
        extent=extent,
        origin="lower",
    )
    ax[panel].set_title("Observation")

    if add_labels:
        assert sky_coords is not None, "Provide sky_coords for labeled objects"

        for k, center in enumerate(sky_coords):
            center_ = observation.get_pixel(center)
            color = "w" if observation.C > 1 else "r"
            ax[panel].text(*center_[::-1], k, color=color, ha="center", va="center")

    panel += 1
    if show_psf:
        psf_image = np.zeros(observation.data.shape)

        if observation.psf is not None:
            psf_model = observation.psf.get_model()
            # make PSF as bright as the brightest pixel of the observation
            psf_model *= (
                observation.data.mean(axis=0).max() / psf_model.mean(axis=0).max()
            )
            # insert into middle of "blank" observation
            full_box = Box(psf_image.shape)
            shift = tuple(
                psf_image.shape[c] // 2 - psf_model.shape[c] // 2
                for c in range(full_box.D)
            )
            model_box = Box(psf_model.shape) + shift
            model_box.insert_into(psf_image, psf_model)
            # slices = scarlet.box.overlapped_slices
        ax[panel].imshow(img_to_rgb(psf_image, norm=norm), origin="lower")
        ax[panel].set_title("PSF")

    fig.tight_layout()
    return fig


def show_scene(
    sources,
    observation=None,
    norm=None,
    channel_map=None,
    show_model=True,
    show_observed=False,
    show_rendered=False,
    show_residual=False,
    add_labels=True,
    add_boxes=False,
    figsize=None,
    linear=True,
):
    """Plot all sources to recreate the scence.
    The functions provides a fast way of evaluating the quality of the entire model,
    i.e. the combination of all scences that seek to fit the observation.
    Parameters
    ----------
    sources: list of source models
    observation: `~scarlet.Observation`
    norm: norm to compress image intensity to the range [0,255]
    channel_map: array_like
        Linear mapping with dimensions (3, channels)
    show_model: bool
        Whether the model is shown in the model frame
    show_observed: bool
        Whether the observation is shown
    show_rendered: bool
        Whether the model, rendered to match the observation, is shown
    show_residual: bool
        Whether the residuals between rendered model and observation is shown
    add_label: bool
        Whether each source is labeled with its numerical index in the source list
    add_boxes: bool
        Whether each source box is shown
    figsize: matplotlib figsize argument
    linear: bool
        Whether or not to display the scene in a single line (`True`) or
        on multiple lines (`False`).
    Returns
    -------
    matplotlib figure
    """
    if show_observed or show_rendered or show_residual:
        assert (
            observation is not None
        ), "Provide matched observation to show observed frame"

    panels = sum((show_model, show_observed, show_rendered, show_residual))
    if linear:
        if figsize is None:
            figsize = (panel_size * panels, panel_size)
        fig, ax = plt.subplots(1, panels, figsize=figsize)
    else:
        columns = int(np.ceil(panels / 2))
        if figsize is None:
            figsize = (panel_size * columns, panel_size * 2)
        fig = plt.figure(figsize=figsize)
        ax = [fig.add_subplot(2, columns, n + 1) for n in range(panels)]
    if not hasattr(ax, "__iter__"):
        ax = (ax,)

    # Mask any pixels with zero weight in all bands
    if observation is not None:
        mask = np.sum(observation.weights, axis=0) == 0
        # if there are no masked pixels, do not use a mask
        if np.all(mask == 0):
            mask = None

    model_frame = sources[0].frame
    model = np.zeros(model_frame.shape)
    for src in sources:
        model += src.get_model(frame=model_frame)

    panel = 0
    if show_model:
        extent = get_extent(model_frame.bbox)
        ax[panel].imshow(
            img_to_rgb(model, norm=norm, channel_map=channel_map),
            extent=extent,
            origin="lower",
        )
        ax[panel].set_title("Model")
        panel += 1

    if show_rendered or show_residual:
        model = observation.render(model)
        extent = get_extent(observation.bbox)

    if show_rendered:
        ax[panel].imshow(
            img_to_rgb(model, norm=norm, channel_map=channel_map, mask=mask),
            extent=extent,
            origin="lower",
        )
        ax[panel].set_title("Model Rendered")
        panel += 1

    if show_observed:
        ax[panel].imshow(
            img_to_rgb(observation.data, norm=norm, channel_map=channel_map, mask=mask),
            extent=extent,
            origin="lower",
        )
        ax[panel].set_title("Observation")
        panel += 1

    if show_residual:
        residual = observation.data - model
        norm_ = LinearPercentileNorm(residual)
        ax[panel].imshow(
            img_to_rgb(residual, norm=norm_, channel_map=channel_map, mask=mask),
            extent=extent,
            origin="lower",
        )
        ax[panel].set_title("Residual")
        panel += 1

    for k, src in enumerate(sources):
        if add_boxes:
            panel = 0
            box_kwargs = {"facecolor": "none", "edgecolor": "w", "lw": 0.5}
            if show_model:
                extent = get_extent(src.bbox)
                rect = Rectangle(
                    (extent[0], extent[2]),
                    extent[1] - extent[0],
                    extent[3] - extent[2],
                    **box_kwargs
                )
                ax[panel].add_artist(rect)
                panel = 1
            if observation is not None:
                start, stop = src.bbox.start[-2:][::-1], src.bbox.stop[-2:][::-1]
                points = (start, (start[0], stop[1]), stop, (stop[0], start[1]))
                coords = [
                    observation.get_pixel(model_frame.get_sky_coord(p)) for p in points
                ]
                for panel in range(panel, panels):
                    poly = Polygon(coords, closed=True, **box_kwargs)
                    ax[panel].add_artist(poly)

        if add_labels and hasattr(src, "center") and src.center is not None:
            center = src.center
            panel = 0
            if show_model:
                ax[panel].text(*center[::-1], k, color="w", ha="center", va="center")
                panel = 1
            if observation is not None:
                center_ = observation.get_pixel(model_frame.get_sky_coord(center))
                for panel in range(panel, panels):
                    ax[panel].text(
                        *center_[::-1], k, color="w", ha="center", va="center"
                    )

    fig.tight_layout()
    return fig


def get_extent(bbox):
    return [bbox.start[-1], bbox.stop[-1], bbox.start[-2], bbox.stop[-2]]


def show_sources(
    sources,
    observation=None,
    norm=None,
    channel_map=None,
    show_model=True,
    show_observed=False,
    show_rendered=False,
    show_spectrum=True,
    figsize=None,
    model_mask=None,
    add_markers=True,
    add_boxes=False,
):
    """Plot each source individually.
    The functions provides an more detailed inspection of every source in the list.
    Parameters
    ----------
    sources: list of source models
    observation: `~scarlet.Observation`
    norm: norm to compress image intensity to the range [0,255]
    channel_map: array_like
        Linear mapping with dimensions (3, channels)
    show_model: bool
        Whether the model is shown in the model frame
    show_observed: bool
        Whether the observation is shown
    show_rendered: bool
        Whether the model, rendered to match the observation, is shown
    show_spectrum: bool
        Whether source specturm is shown.
        For multi-component sources, spectra are shown separately.
    figsize: matplotlib figsize argument
    model_mask: array
        Mask used to hide pixels in the model only.
    add_markers: bool
        Whether or not to mark the centers of the sources
        with their source number.
    add_boxes: bool
        Whether source boxes are shown
    Returns
    -------
    matplotlib figure
    """
    if show_observed or show_rendered:
        assert (
            observation is not None
        ), "Provide matched observation to show observed frame"

    panels = sum((show_model, show_observed, show_rendered, show_spectrum))
    if figsize is None:
        figsize = (panel_size * panels, panel_size * len(list(sources)))
    fig, ax = plt.subplots(len(list(sources)), panels, figsize=figsize, squeeze=False)

    marker_kwargs = {"mew": 1, "ms": 10}
    box_kwargs = {"facecolor": "none", "edgecolor": "w", "lw": 0.5}

    for k, src in enumerate(sources):

        model_frame = src.frame

        if hasattr(src, "center") and src.center is not None:
            center = np.array(src.center)[::-1]
        else:
            center = None

        if add_boxes:
            start, stop = src.bbox.start[-2:][::-1], src.bbox.stop[-2:][::-1]
            points = (start, (start[0], stop[1]), stop, (stop[0], start[1]))
            box_coords = [
                observation.get_pixel(model_frame.get_sky_coord(p)) for p in points
            ]

        # model in its bbox
        panel = 0
        model = src.get_model()

        if show_model:
            # Show the unrendered model in it's bbox
            extent = get_extent(src.bbox)
            ax[k][panel].imshow(
                img_to_rgb(model, norm=norm, channel_map=channel_map, mask=model_mask),
                extent=extent,
                origin="lower",
            )
            ax[k][panel].set_title("Model Source {}".format(k))
            if center is not None and add_markers:
                ax[k][panel].plot(*center, "wx", **marker_kwargs)
            panel += 1

        # model in observation frame
        if show_rendered:
            # Center and show the rendered model
            model_ = src.get_model(frame=model_frame)
            model_ = observation.render(model_)
            extent = get_extent(observation.bbox)
            ax[k][panel].imshow(
                img_to_rgb(model_, norm=norm, channel_map=channel_map),
                extent=extent,
                origin="lower",
            )
            ax[k][panel].set_title("Model Source {} Rendered".format(k))

            if center is not None and add_markers:
                center_ = observation.get_pixel(model_frame.get_sky_coord(center))
                ax[k][panel].plot(*center_, "wx", **marker_kwargs)
            if add_boxes:
                poly = Polygon(box_coords, closed=True, **box_kwargs)
                ax[k][panel].add_artist(poly)
            panel += 1

        if show_observed:
            # Center the observation on the source and display it
            _images = observation.data
            ax[k][panel].imshow(
                img_to_rgb(_images, norm=norm, channel_map=channel_map),
                extent=extent,
                origin="lower",
            )
            ax[k][panel].set_title("Observation".format(k))
            if center is not None and add_markers:
                center_ = observation.get_pixel(model_frame.get_sky_coord(center))
                ax[k][panel].plot(*center_, "wx", **marker_kwargs)
            if add_boxes:
                poly = Polygon(box_coords, closed=True, **box_kwargs)
                ax[k][panel].add_artist(poly)
            panel += 1

        if show_spectrum:
            # needs to be evaluated in the source box to prevent truncation
            if hasattr(src, "__iter__") and isinstance(src[0], Component):
                spectra = []
                for component in src:
                    model_ = component.get_model()
                    spectra.append(model_.sum(axis=(1, 2)))
            else:
                spectra = [model.sum(axis=(1, 2))]

            for spectrum in spectra:
                ax[k][panel].plot(spectrum)
            ax[k][panel].set_xticks(range(len(spectrum)))
            if hasattr(src.frame, "channels") and src.frame.channels is not None:
                ax[k][panel].set_xticklabels(src.frame.channels)
            ax[k][panel].set_title("Spectrum")
            ax[k][panel].set_xlabel("Channel")
            ax[k][panel].set_ylabel("Intensity")

    fig.tight_layout()
    return fig
Get the linear mapping of multiple channels to RGB channels     The mapping created here assumes the the channels are ordered in wavelength     direction, starting with the shortest wavelength. The mapping seeks to produce     a relatively even weights for across all channels. It does not consider e.g.     signal-to-noise variations across channels or human perception.     Parameters     ----------     channels: int in range(0,7)         Number of channels     Returns     -------     array (3, channels) to map onto RGB Create norm that is linear between lower and upper percentile of img         Parameters         ----------         img: array_like             Image to normalize         percentile: array_like, default=[1,99]             Lower and upper percentile to consider. Pixel values below will be             set to zero, above to saturated. Create norm that is linear between lower and upper percentile of img         Parameters         ----------         img: array_like             Image to normalize         percentile: array_like, default=[1,99]             Lower and upper percentile to consider. Pixel values below will be             set to zero, above to saturated. # solution for beta assumes flat spectrum at vmax Convert multi-band image cube into 3 RGB channels     Parameters     ----------     img: array_like         This should be an array with dimensions (channels, height, width).     channel_map: array_like         Linear mapping with dimensions (3, channels)     fill_value: float, default=`0`         Value to use for any masked pixels.     Returns     -------     RGB: numpy array with dtype float # expand single img into cube # filterWeights: channel x band # map channels onto RGB channels Convert images to normalized RGB.     If normalized values are outside of the range [0..255], they will be     truncated such as to preserve the corresponding color.     Parameters     ----------     img: array_like         This should be an array with dimensions (channels, height, width).     channel_map: array_like         Linear mapping with dimensions (3, channels)     fill_value: float, default=`0`         Value to use for any masked pixels.     norm: `scarlet.display.Norm`, default `None`         Norm to use for mapping in the allowed range [0..255]. If `norm=None`,         `scarlet.display.LinearPercentileNorm` will be used.     mask: array_like         A [0,1] binary mask to apply over the top of the image,         where pixels with mask==1 are masked out.     Returns     -------     rgb: numpy array with dimensions (3, height, width) and dtype uint8 Plot observation in standardized form. # Mask any pixels with zero weight in all bands # if there are no masked pixels, do not use a mask # make PSF as bright as the brightest pixel of the observation # insert into middle of "blank" observation # slices = scarlet.box.overlapped_slices Plot all sources to recreate the scence.     The functions provides a fast way of evaluating the quality of the entire model,     i.e. the combination of all scences that seek to fit the observation.     Parameters     ----------     sources: list of source models     observation: `~scarlet.Observation`     norm: norm to compress image intensity to the range [0,255]     channel_map: array_like         Linear mapping with dimensions (3, channels)     show_model: bool         Whether the model is shown in the model frame     show_observed: bool         Whether the observation is shown     show_rendered: bool         Whether the model, rendered to match the observation, is shown     show_residual: bool         Whether the residuals between rendered model and observation is shown     add_label: bool         Whether each source is labeled with its numerical index in the source list     add_boxes: bool         Whether each source box is shown     figsize: matplotlib figsize argument     linear: bool         Whether or not to display the scene in a single line (`True`) or         on multiple lines (`False`).     Returns     -------     matplotlib figure # Mask any pixels with zero weight in all bands # if there are no masked pixels, do not use a mask Plot each source individually.     The functions provides an more detailed inspection of every source in the list.     Parameters     ----------     sources: list of source models     observation: `~scarlet.Observation`     norm: norm to compress image intensity to the range [0,255]     channel_map: array_like         Linear mapping with dimensions (3, channels)     show_model: bool         Whether the model is shown in the model frame     show_observed: bool         Whether the observation is shown     show_rendered: bool         Whether the model, rendered to match the observation, is shown     show_spectrum: bool         Whether source specturm is shown.         For multi-component sources, spectra are shown separately.     figsize: matplotlib figsize argument     model_mask: array         Mask used to hide pixels in the model only.     add_markers: bool         Whether or not to mark the centers of the sources         with their source number.     add_boxes: bool         Whether source boxes are shown     Returns     -------     matplotlib figure # model in its bbox # Show the unrendered model in it's bbox # model in observation frame # Center and show the rendered model # Center the observation on the source and display it # needs to be evaluated in the source box to prevent truncation