Source code for pynpoint.util.sdi

"""
Functions for spectral differential imaging.
"""

import numpy as np

from typeguard import typechecked

from pynpoint.util.image import scale_image, shift_image




@typechecked
def sdi_scaling(image_in: np.ndarray, scaling: np.ndarray) -> np.ndarray:
    """
    Function to rescale the images by their wavelength ratios.

    Parameters
    ----------
    image_in : np.ndarray
        Data to rescale
    scaling : np.ndarray
        Scaling factors.

    Returns
    -------
    np.ndarray
        Rescaled images with the same shape as ``image_in``.
    """

    if image_in.shape[0] != scaling.shape[0]:
        raise ValueError(
            "The number of wavelengths is not equal to the number of available "
            "scaling factors."
        )

    image_out = np.zeros(image_in.shape)

    for i in range(image_in.shape[0]):
        swaps = scale_image(image_in[i,], scaling[i], scaling[i])

        npix_del = swaps.shape[-1] - image_out.shape[-1]

        if npix_del == 0:
            image_out[i,] = swaps

        else:
            if npix_del % 2 == 0:
                npix_del_a = int(npix_del / 2)
                npix_del_b = int(npix_del / 2)

            else:
                npix_del_a = int((npix_del - 1) / 2)
                npix_del_b = int((npix_del + 1) / 2)

            image_out[i,] = swaps[npix_del_a:-npix_del_b, npix_del_a:-npix_del_b]

        if npix_del % 2 == 1:
            image_out[i,] = shift_image(
                image_out[i,], (-0.5, -0.5), interpolation="spline"
            )

    return image_out





@typechecked
def scaling_factors(wavelengths: np.ndarray) -> np.ndarray:
    """
    Function to calculate the scaling factors for SDI.

    Parameters
    ----------
    wavelengths : np.ndarray
        Array with the wavelength of each frame.

    Returns
    -------
    np.ndarray
        Scaling factors.
    """

    return max(wavelengths) / wavelengths