Display & Metrics¶

def display_psf(hdulist_or_filename, ext=0, vmin=1e-7, vmax=1e-1,
                scale='log', cmap=None, title=None, imagecrop=None,
                adjust_for_oversampling=False, normalize='None',
                crosshairs=False, markcentroid=False, colorbar=True,
                colorbar_orientation='vertical', pixelscale='PIXELSCL',
                ax=None, return_ax=False, interpolation=None, cube_slice=None,
                angular_coordinate_unit=u.arcsec) -> None:
    """Display a PSF image stored in a FITS HDUList or file.

    Renders a 2-D PSF on a linear or logarithmic colour scale with optional
    crosshairs, centroid marker, and colour bar.  Axes are labelled in angular
    units (default: arcseconds).

    Args:
        hdulist_or_filename (fits.HDUList | str): FITS object or path to a
            FITS file containing the PSF image.
        ext (int): HDU extension index to read. Defaults to ``0``.
        vmin (float): Lower colour-scale limit. Defaults to ``1e-7``.
        vmax (float): Upper colour-scale limit. Defaults to ``1e-1``.
        scale (str): Intensity scale — ``'log'`` or ``'linear'``.
            Defaults to ``'log'``.
        cmap (matplotlib.colors.Colormap | None): Colormap instance.
            ``None`` uses the library default. Defaults to ``None``.
        title (str | None): Axes title.  Defaults to ``None``.
        imagecrop (float | None): Half-width of the displayed region in
            arcseconds.  ``None`` shows the full image. Defaults to ``None``.
        adjust_for_oversampling (bool): Rescale pixel values to account for
            detector oversampling. Defaults to ``False``.
        normalize (str): Normalisation applied before display
            (``'None'``, ``'peak'``, or ``'total'``). Defaults to
            ``'None'``.
        crosshairs (bool): Draw crosshairs at the optical axis.
            Defaults to ``False``.
        markcentroid (bool): Compute and mark the PSF centroid.
            Defaults to ``False``.
        colorbar (bool): Draw a colour bar. Defaults to ``True``.
        colorbar_orientation (str): ``'vertical'`` or ``'horizontal'``.
            Defaults to ``'vertical'``.
        pixelscale (str | float): FITS header keyword or numeric value
            (arcsec/pixel) used to set axis scale. Defaults to
            ``'PIXELSCL'``.
        ax (matplotlib.axes.Axes | None): Target axes.  Creates a new
            figure when ``None``. Defaults to ``None``.
        return_ax (bool): Return the axes object after plotting.
            Defaults to ``False``.
        interpolation (str | None): Matplotlib interpolation method for
            ``imshow``. Defaults to ``None``.
        cube_slice (int | None): For PSF datacubes, the wavelength slice
            index to display. Defaults to ``None``.
        angular_coordinate_unit (astropy.units.Unit): Angular unit for the
            axis labels. Defaults to ``astropy.units.arcsec``.

    Returns:
        matplotlib.axes.Axes | None: The axes object when
            ``return_ax=True``, otherwise ``None``.

    Example:
        >>> psf = tel.calc_psf(fov_arcsec=2.0)
        >>> display_psf(psf, scale='log', title='My PSF')
    """
    if isinstance(hdulist_or_filename, str):
        hdulist = fits.open(hdulist_or_filename)
    elif isinstance(hdulist_or_filename, fits.HDUList):
        hdulist = hdulist_or_filename
    elif isinstance(hdulist_or_filename, (np.ndarray)):
        hdulist = None
    else : 
        raise ValueError("input must be a filename or FITS HDUList object")

    if hdulist is not None:
        # Get a handle on the input image
        if hdulist[ext].data.ndim == 2:
            im0 = hdulist[ext].data
            psf_array_shape = im0.shape
        elif hdulist[ext].data.ndim == 3:
            if cube_slice is None:
                raise ValueError("To display a PSF datacube, you must set cube_slice=<#>.")
            else:
                im0 = hdulist[ext].data[cube_slice]
                psf_array_shape = hdulist[ext].data.shape[1:]
        else:
            raise RuntimeError("Unsupported image dimensionality.")

    else : 
        im0 = hdulist_or_filename
        psf_array_shape = im0.shape

    # Normalization
    if adjust_for_oversampling:
        try:
            scalefactor = hdulist[ext].header['OVERSAMP'] ** 2
        except KeyError:
            _log.error("Could not determine oversampling scale factor; "
                       "therefore NOT rescaling fluxes.")
            scalefactor = 1
        im = im0 * scalefactor
    else:
        # don't change normalization of actual input array, work with a copy!
        im = im0.copy()

    if normalize.lower() == 'peak':
        _log.debug("Displaying image normalized to peak = 1")
        im /= im.max()
    elif normalize.lower() == 'total':
        _log.debug("Displaying image normalized to PSF total = 1")
        im /= im.sum()

    if scale == 'linear':
        norm = mpl.colors.Normalize(vmin=vmin, vmax=vmax)
    else:
        norm = mpl.colors.LogNorm(vmin=vmin, vmax=vmax)

    if hdulist is not None:
        if isinstance(pixelscale, str):
            pixelscale = hdulist[ext].header[pixelscale]
        else:
            pixelscale = float(pixelscale)
    else:
        pixelscale = 0.298 # Euclid NISP resolution by default (in arcsec/pixel)

    if angular_coordinate_unit != u.arcsec:
        coordinate_rescale = (1 * u.arcsec).to_value(angular_coordinate_unit)
        pixelscale *= coordinate_rescale
    else:
        coordinate_rescale = 1
    halffov_x = pixelscale * psf_array_shape[1] / 2.0
    halffov_y = pixelscale * psf_array_shape[0] / 2.0

    unit_label = str(angular_coordinate_unit)
    extent = [-halffov_x, halffov_x, -halffov_y, halffov_y]

    if cmap is None:
        cmap = getattr(mpl.cm, poppy.conf.cmap_sequential)
    # update and get (or create) image axes
    ax = imshow_with_mouseover(
        im,
        extent=extent,
        cmap=cmap,
        norm=norm,
        ax=ax,
        interpolation=interpolation,
        origin='lower'
    )
    ax.set_xlabel(unit_label)

    if markcentroid:
        _log.info("measuring centroid to mark on plot...")
        # ceny, cenx = measure_centroid(hdulist, ext=ext, units='arcsec', relativeto='center', boxsize=20, threshold=0.1)
        if hdulist is not None:
            pixelscale_measure,cent_of_mass = measure_centroid(hdulist, ext=ext, units='arcsec', relativeto='center', boxsize=20, threshold=0.1)
        else : 
            pixelscale_measure,cent_of_mass = measure_centroid(im, units='arcsec', relativeto='center', boxsize=20, threshold=0.1)

        if (pixelscale_measure == None) and (hdulist == None) :
            pass
        else : 
            cent_of_mass = tuple(np.array(cent_of_mass) * pixelscale)

        ceny, cenx = cent_of_mass

        ceny *= coordinate_rescale  # if display coordinate unit isn't arcseconds, rescale the centroid accordingly
        cenx *= coordinate_rescale
        ax.plot(cenx, ceny, 'k+', markersize=15, markeredgewidth=1)
        _log.info("centroid: (%f, %f) " % (cenx, ceny))

    if imagecrop is not None:
        halffov_x = min((imagecrop / 2.0, halffov_x))
        halffov_y = min((imagecrop / 2.0, halffov_y))
    ax.set_xbound(-halffov_x, halffov_x)
    ax.set_ybound(-halffov_y, halffov_y)
    if crosshairs:
        ax.axhline(0, ls=':', color='k')
        ax.axvline(0, ls=':', color='k')
    if title is None:
        if hdulist is not None:
            try:
                fspec = "%s, %s" % (hdulist[ext].header['INSTRUME'], hdulist[ext].header['FILTER'])
            except KeyError:
                fspec = str(hdulist_or_filename)
            title = "PSF sim for " + fspec
        else : 
            title = "PSF"
    ax.set_title(title)

    if colorbar:
        if ax.images[0].colorbar is not None:
            # Reuse existing colorbar axes (Issue #21)
            colorbar_axes = ax.images[0].colorbar.ax
            cb = plt.colorbar(
                ax.images[0],
                ax=ax,
                cax=colorbar_axes,
                orientation=colorbar_orientation,
                fraction=0.046,
                pad=0.04,
            )
        else:
            cb = plt.colorbar(
                ax.images[0],
                ax=ax,
                orientation=colorbar_orientation,
                fraction=0.046,
                pad=0.04,
            )
        if scale.lower() == 'log':
            ticks = np.logspace(np.log10(vmin), np.log10(vmax), int(np.round(np.log10(vmax / vmin) + 1)))
            if colorbar_orientation == 'horizontal' and vmax == 1e-1 and vmin == 1e-8:
                ticks = [1e-8, 1e-6, 1e-4, 1e-2, 1e-1]  # looks better
            cb.set_ticks(ticks)
            cb.set_ticklabels(ticks)
        if normalize.lower() == 'peak':
            cb.set_label('Intensity relative to peak pixel')
        else:
            cb.set_label('Fractional intensity per pixel')

    if return_ax:
        if colorbar:
            return ax, cb
        else:
            return ax

def display_ee(hdulist_or_filename=None, ext: int = 0, overplot: bool = False, ax=None, mark_levels: bool = True, **kwargs: object) -> None:
    """Plot the azimuthally averaged Encircled Energy (EE) curve for a PSF.

    The EE is computed via :func:`radial_profile` and plotted as a function
    of angular radius in arcseconds.  Reference lines are optionally drawn
    at EE = 50 %, 80 %, and 95 %.

    Args:
        hdulist_or_filename (fits.HDUList | str): FITS object or path to a
            FITS file containing the PSF image.
        ext (int): HDU extension index. Defaults to ``0``.
        overplot (bool): If ``True``, add the curve to the current axes
            without clearing them.  Defaults to ``False``.
        ax (matplotlib.axes.Axes | None): Target axes.  When ``None`` and
            ``overplot=False`` a new axes is created.  Defaults to ``None``.
        mark_levels (bool): Annotate the EE = 50 %, 80 %, and 95 % radii on
            the plot.  Defaults to ``True``.
        **kwargs: Extra keyword arguments forwarded to :func:`radial_profile`
            (e.g. ``center``, ``stddev``).

    Raises:
        ValueError: If ``hdulist_or_filename`` is neither a string nor an
            ``fits.HDUList``.

    Example:
        >>> psf = tel.calc_psf(fov_arcsec=3.0)
        >>> fig, ax = plt.subplots()
        >>> display_ee(psf, ax=ax, mark_levels=True)
    """
    if isinstance(hdulist_or_filename, str):
        hdu_list = fits.open(hdulist_or_filename)
    elif isinstance(hdulist_or_filename, fits.HDUList):
        hdu_list = hdulist_or_filename
    else:
        raise ValueError("input must be a filename or HDUlist")

    radius, profile, ee = radial_profile(hdu_list, ee=True, ext=ext, **kwargs)

    if not overplot:
        if ax is None:
            plt.clf()
            ax = plt.subplot(111)

    ax.plot(radius, ee)  # , nonposy='clip')
    if not overplot:
        ax.set_xlabel("Radius [arcsec]")
        ax.set_ylabel("Encircled Energy")

    if mark_levels:
        for level in [0.5, 0.8, 0.95]:
            try : # Here is the modified part, that allows to plot EE when the PSF is such that the EE90 isn't on the graph (thus np.where(ee > level) is empty)
                ee_lev = radius[np.where(ee > level)[0][0]]
                yoffset = 0 if level < 0.9 else -0.05
                ax.text(ee_lev + 0.1, level + yoffset, 'EE=%2d%% at r=%.3f"' % (level * 100, ee_lev))
            except :
                pass

def measure_ee(hdulist_or_filename=None, ext=0, center=None, binsize=None, normalize='None'):
    """ measure encircled energy vs radius and return as an interpolator

    Returns a function object which when called returns the Encircled Energy inside a given radius,
    for any arbitrary desired radius smaller than the image size.



    Parameters
    ----------
    hdulist_or_filename : string
        Either a fits.HDUList object or a filename of a FITS file on disk
    ext : int
        Extension in that FITS file
    center : tuple of floats
        Coordinates (x,y) of PSF center. Default is image center.
    binsize:
        size of step for profile. Default is pixel size.
    normalize : string
        set to 'peak' to normalize peak intensity =1, or to 'total' to normalize total flux=1.
        Default is no normalization (i.e. retain whatever normalization was used in computing the PSF itself)

    Returns
    -------
    encircled_energy: function
        A function which will return the encircled energy interpolated to any desired radius.


    Examples
    --------
    >>> ee = measure_ee("someimage.fits")  # doctest: +SKIP
    >>> print("The EE at 0.5 arcsec is ", ee(0.5))  # doctest: +SKIP

    """

    rr, radialprofile2, ee = radial_profile(hdulist_or_filename, ext, ee=True, center=center, binsize=binsize,
                                            normalize=normalize)

    # append the zero at the center
    rr_ee = rr + (rr[1] - rr[0]) / 2.0  # add half a binsize to this, because the ee is measured inside the
    # outer edge of each annulus.
    rr0 = np.concatenate(([0], rr_ee))
    ee0 = np.concatenate(([0], ee))

    ee_fn = scipy.interpolate.interp1d(rr0, ee0, kind='cubic', bounds_error=False)

    return ee_fn

def measure_centroid(HDUlist_or_filename=None, ext=0, slice=0, boxsize=20, verbose=False, units='pixels',
                     relativeto='origin', **kwargs):
    """ Measure the center of an image via center-of-mass

    The centroid method used is the floating-box center of mass algorithm by
    Jeff Valenti et al., which has been adopted for JWST target acquisition
    measurements on orbit.
    See JWST technical reports JWST-STScI-001117 and JWST-STScI-001134 for details.

    Parameters
    ----------
    HDUlist_or_filename : string
        Either a fits.HDUList object or a filename of a FITS file on disk
    ext : int
        Extension in that FITS file
    slice : int, optional
        If that extension is a 3D datacube, which slice (plane) of that datacube to use
    boxsize : int
        Half box size for centroid
    relativeto : string
        either 'origin' for relative to pixel (0,0) or 'center' for relative to image center. Default is 'origin'
    units : string
        either 'pixels' for position in pixels or 'arcsec' for arcseconds.
        Relative to the relativeto parameter point in either case.
    verbose : bool
        Be more verbose


    Returns
    -------
    CoM : array_like
        [Y, X] coordinates of center of mass.

    """
    from .fwcentroid import fwcentroid

    if isinstance(HDUlist_or_filename, str):
        HDUlist = fits.open(HDUlist_or_filename)
    elif isinstance(HDUlist_or_filename, fits.HDUList):
        HDUlist = HDUlist_or_filename
    else:
        raise ValueError("input must be a filename or HDUlist")

    image = HDUlist[ext].data

    if image.ndim >= 3:  # handle datacubes gracefully
        image = image[slice, :, :]

    cent_of_mass = fwcentroid(image, halfwidth=boxsize, **kwargs)

    if verbose:
        print("Center of mass: (%.4f, %.4f)" % (cent_of_mass[1], cent_of_mass[0]))

    if relativeto == 'center':
        imcen = np.array([(image.shape[0] - 1) / 2., (image.shape[1] - 1) / 2.])
        cent_of_mass = tuple(np.array(cent_of_mass) - imcen)

    if units == 'arcsec':
        pixelscale = HDUlist[ext].header['PIXELSCL']
        cent_of_mass = tuple(np.array(cent_of_mass) * pixelscale)

    return cent_of_mass