Source code for jwst.extract_1d.soss_extract.soss_boxextract

import logging
import warnings

import numpy as np

log = logging.getLogger(__name__)

__all__ = ["get_box_weights", "box_extract", "estim_error_nearest_data"]




def get_box_weights(centroid, n_pix, shape, cols):
    """
    Return the weights of a box aperture given the centroid and the width of the box in pixels.

    All pixels will have the same weights except at the ends of the box aperture.

    Parameters
    ----------
    centroid : ndarray
        Position of the centroid (in rows). Same shape as ``cols``.
    n_pix : float
        Width of the extraction box in pixels.
    shape : tuple
        Shape of the output image, i.e., ``(n_row, n_column)``.
    cols : array
        Column indices (int) of good columns. Used if the centroid is defined
        for specific columns or a sub-range of columns.

    Returns
    -------
    weights : ndarray
        An array of pixel weights to use with the box extraction.
    """
    nrows, _ = shape

    # Row centers of all pixels.
    rows = np.indices((nrows, len(cols)))[0]

    # Pixels that are entirely inside the box are set to one.
    cond = rows <= (centroid - 0.5 + n_pix / 2)
    cond &= (centroid + 0.5 - n_pix / 2) <= rows
    weights = cond.astype(float)

    # Fractional weights at the upper bound.
    cond = (centroid - 0.5 + n_pix / 2) < rows
    cond &= rows < (centroid + 0.5 + n_pix / 2)
    weights[cond] = (centroid + n_pix / 2 - (rows - 0.5))[cond]

    # Fractional weights at the lower bound.
    cond = rows < (centroid + 0.5 - n_pix / 2)
    cond &= (centroid - 0.5 - n_pix / 2) < rows
    weights[cond] = (rows + 0.5 - (centroid - n_pix / 2))[cond]

    # Return with the specified shape with zeros where the box is not defined.
    out = np.zeros(shape, dtype=float)
    out[:, cols] = weights

    return out





def box_extract(scidata, scierr, scimask, box_weights):
    """
    Perform a box extraction.

    Parameters
    ----------
    scidata : ndarray
        2D array of science data with shape ``(n_row, n_columns)``
    scierr : ndarray
        2D array of uncertainty map with same shape as ``scidata``
    scimask : ndarray
        2D boolean array of masked pixels with same shape as ``scidata``
    box_weights : ndarray
        2D array of pre-computed weights for box extraction,
        with same shape as ``scidata``

    Returns
    -------
    cols : ndarray
        Indices of extracted columns
    flux : ndarray
        The flux in each column
    flux_var : ndarray
        The variance of the flux in each column
    """
    cols = np.arange(scidata.shape[1])

    # Keep only required columns and make a copy.
    data = scidata[:, cols].copy()
    error = scierr[:, cols].copy()
    mask = scimask[:, cols].copy()
    box_weights = box_weights[:, cols].copy()

    # Check that all invalid values in the extraction region are masked.
    extract_region = box_weights > 0
    condition = extract_region & ~mask

    if not np.isfinite(data[condition]).all():
        message = "scidata contains un-masked invalid values."
        log.critical(message)
        raise ValueError(message)

    if not np.isfinite(error[condition]).all():
        message = "scierr contains un-masked invalid values."
        log.critical(message)
        raise ValueError(message)

    # Set all pixels values outside of extraction region to Nan
    # so it will be correctly handle by np.nansum.
    data = np.where(extract_region, data, np.nan)
    error = np.where(extract_region, error, np.nan)

    # Set the weights of masked pixels to zero.
    box_weights[mask] = 0.0

    with warnings.catch_warnings():
        warnings.filterwarnings("ignore", "invalid value encountered in multiply", RuntimeWarning)
        # Extract total flux (sum over columns).
        flux = np.nansum(box_weights * data, axis=0)
        npix = np.nansum(box_weights, axis=0)

        # Extract flux error (sum of variances).
        flux_var = np.nansum(box_weights * error**2, axis=0)
        flux_err = np.sqrt(flux_var)

        # Set empty columns to NaN.
        flux = np.where(npix > 0, flux, np.nan)
        flux_err = np.where(npix > 0, flux_err, np.nan)

    return cols, flux, flux_err, npix





def estim_error_nearest_data(err, data, pix_to_estim, valid_pix):
    """
    Estimate pixel error empirically using error on nearby pixel values in data.

    Intended to be used in a box extraction when the bad pixels are modeled.

    Parameters
    ----------
    err : ndarray
        Uncertainty map (2D) of the pixels.
    data : ndarray
        Pixel values (2D).
    pix_to_estim : ndarray
        Boolean map (2D) of the pixels where the uncertainty needs to be estimated.
    valid_pix : ndarray
        Boolean map (2D) of valid pixels to be used to find the error empirically.

    Returns
    -------
    err_filled : ndarray
        Same as ``err`` (2D), but the pixels to be estimated are filled with the estimated values.

    Notes
    -----
    For uncorrelated noise, the average error on the replaced pixels will be roughly
    half of the average error on the original good pixels.
    The reason is because this method chooses the lower error between the two nearest-flux
    data points, leading to a factor-of-2 decrease (assuming errors are normally distributed).
    Future work should follow up on whether this remains the desired behavior.
    """
    # Transform to 1d arrays
    data_to_estim = data[pix_to_estim]
    err_valid = err[valid_pix]
    data_valid = data[valid_pix]

    # Need to sort the arrays used to find similar values
    idx_sort = np.argsort(data_valid)
    err_valid = err_valid[idx_sort]
    data_valid = data_valid[idx_sort]

    # Searchsorted: gives the position of the nearest higher value,
    # not necessarily the closest value
    idx_higher = np.searchsorted(data_valid, data_to_estim)
    idx_higher = np.clip(idx_higher, 0, err_valid.size - 1)
    # The nearest lower value is given by the preceding index
    idx_lower = np.clip(idx_higher - 1, 0, err_valid.size - 1)

    # Find the best between index around the value (lower and higher index) ...
    idx_around = np.vstack([idx_lower, idx_higher])
    # ... using the one with the smallest error
    distance = np.abs(data_valid[idx_around] - data_to_estim[None, :])
    idx_best_of_2 = np.argmin(distance, axis=0)
    idx_closest = idx_around[idx_best_of_2, np.arange(idx_best_of_2.size)]

    # Get the corresponding error (that's what we want to find!)
    err_estimate = err_valid[idx_closest]

    # Replace estimated values in the output error 2D image
    err_out = err.copy()
    err_out[pix_to_estim] = err_estimate

    return err_out