brighterFatterKernel.py

Go to the documentation of this file.

# This file is part of ip_isr.
#
# Developed for the LSST Data Management System.
# This product includes software developed by the LSST Project
# (https://www.lsst.org).
# See the COPYRIGHT file at the top-level directory of this distribution
# for details of code ownership.
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program.  If not, see <https://www.gnu.org/licenses/>.
#
"""Brighter Fatter Kernel calibration definition."""
 
 
__all__ = [
    'BrighterFatterKernel',
    'brighterFatterCorrection',
    'fluxConservingBrighterFatterCorrection']
 
 
import scipy
import numpy as np
from astropy.table import Table
 
import lsst.afw.math as afwMath
import lsst.afw.image as afwImage
 
from . import IsrCalib
from .isrFunctions import gainContext
 
 
class BrighterFatterKernel(IsrCalib):
    """Calibration of brighter-fatter kernels for an instrument.
 
    ampKernels are the kernels for each amplifier in a detector, as
    generated by having ``level == 'AMP'``.
 
    detectorKernel is the kernel generated for a detector as a
    whole, as generated by having ``level == 'DETECTOR'``.
 
    makeDetectorKernelFromAmpwiseKernels is a method to generate the
    kernel for a detector, constructed by averaging together the
    ampwise kernels in the detector.  The existing application code is
    only defined for kernels with ``level == 'DETECTOR'``, so this method
    is used if the supplied kernel was built with ``level == 'AMP'``.
 
    Parameters
    ----------
    camera : `lsst.afw.cameraGeom.Camera`
        Camera describing detector geometry.
    level : `str`
        Level the kernels will be generated for.
    log : `logging.Logger`, optional
        Log to write messages to.
    **kwargs :
        Parameters to pass to parent constructor.
 
    Notes
    -----
    Version 1.1 adds the `expIdMask` property, and substitutes
    `means` and `variances` for `rawMeans` and `rawVariances`
    from the PTC dataset.
 
    expIdMask : `dict`, [`str`,`numpy.ndarray`]
        Dictionary keyed by amp names containing the mask produced after
        outlier rejection.
    rawMeans : `dict`, [`str`, `numpy.ndarray`]
        Dictionary keyed by amp names containing the unmasked average of the
        means of the exposures in each flat pair.
    rawVariances : `dict`, [`str`, `numpy.ndarray`]
        Dictionary keyed by amp names containing the variance of the
        difference image of the exposures in each flat pair.
        Corresponds to rawVars of PTC.
    rawXcorrs : `dict`, [`str`, `numpy.ndarray`]
        Dictionary keyed by amp names containing an array of measured
        covariances per mean flux.
        Corresponds to covariances of PTC.
    badAmps : `list`
        List of bad amplifiers names.
    shape : `tuple`
        Tuple of the shape of the BFK kernels.
    gain : `dict`, [`str`,`float`]
        Dictionary keyed by amp names containing the fitted gains.
    noise : `dict`, [`str`,`float`]
        Dictionary keyed by amp names containing the fitted noise.
    meanXcorrs : `dict`, [`str`,`numpy.ndarray`]
        Dictionary keyed by amp names containing the averaged
        cross-correlations.
    valid : `dict`, [`str`,`bool`]
        Dictionary keyed by amp names containing validity of data.
    ampKernels : `dict`, [`str`, `numpy.ndarray`]
        Dictionary keyed by amp names containing the BF kernels.
    detKernels : `dict`
        Dictionary keyed by detector names containing the BF kernels.
    """
    _OBSTYPE = 'bfk'
    _SCHEMA = 'Brighter-fatter kernel'
    _VERSION = 1.1
 
    def __init__(self, camera=None, level=None, **kwargs):
        self.level = level
 
        # Things inherited from the PTC
        self.expIdMask = dict()
        self.rawMeans = dict()
        self.rawVariances = dict()
        self.rawXcorrs = dict()
        self.badAmps = list()
        self.shape = (17, 17)
        self.gain = dict()
        self.noise = dict()
 
        # Things calculated from the PTC
        self.meanXcorrs = dict()
        self.valid = dict()
 
        # Things that are used downstream
        self.ampKernels = dict()
        self.detKernels = dict()
 
        super().__init__(**kwargs)
 
        if camera:
            self.initFromCamera(camera, detectorId=kwargs.get('detectorId', None))
 
        self.requiredAttributes.update(['level', 'expIdMask', 'rawMeans', 'rawVariances', 'rawXcorrs',
                                        'badAmps', 'gain', 'noise', 'meanXcorrs', 'valid',
                                        'ampKernels', 'detKernels'])
 
    def updateMetadata(self, setDate=False, **kwargs):
        """Update calibration metadata.
 
        This calls the base class's method after ensuring the required
        calibration keywords will be saved.
 
        Parameters
        ----------
        setDate : `bool`, optional
            Update the CALIBDATE fields in the metadata to the current
            time. Defaults to False.
        kwargs :
            Other keyword parameters to set in the metadata.
        """
        kwargs['LEVEL'] = self.level
        kwargs['KERNEL_DX'] = self.shape[0]
        kwargs['KERNEL_DY'] = self.shape[1]
 
        super().updateMetadata(setDate=setDate, **kwargs)
 
    def initFromCamera(self, camera, detectorId=None):
        """Initialize kernel structure from camera.
 
        Parameters
        ----------
        camera : `lsst.afw.cameraGeom.Camera`
            Camera to use to define geometry.
        detectorId : `int`, optional
            Index of the detector to generate.
 
        Returns
        -------
        calib : `lsst.ip.isr.BrighterFatterKernel`
            The initialized calibration.
 
        Raises
        ------
        RuntimeError
            Raised if no detectorId is supplied for a calibration with
            ``level='AMP'``.
        """
        self._instrument = camera.getName()
 
        if detectorId is not None:
            detector = camera[detectorId]
            self._detectorId = detectorId
            self._detectorName = detector.getName()
            self._detectorSerial = detector.getSerial()
 
        if self.level == 'AMP':
            if detectorId is None:
                raise RuntimeError("A detectorId must be supplied if level='AMP'.")
 
            self.badAmps = []
 
            for amp in detector:
                ampName = amp.getName()
                self.expIdMask[ampName] = []
                self.rawMeans[ampName] = []
                self.rawVariances[ampName] = []
                self.rawXcorrs[ampName] = []
                self.gain[ampName] = amp.getGain()
                self.noise[ampName] = amp.getReadNoise()
                self.meanXcorrs[ampName] = []
                self.ampKernels[ampName] = []
                self.valid[ampName] = []
        elif self.level == 'DETECTOR':
            if detectorId is None:
                for det in camera:
                    detName = det.getName()
                    self.detKernels[detName] = []
            else:
                self.detKernels[self._detectorName] = []
 
        return self
 
    def getLengths(self):
        """Return the set of lengths needed for reshaping components.
 
        Returns
        -------
        kernelLength : `int`
            Product of the elements of self.shape.
        smallLength : `int`
            Size of an untiled covariance.
        nObs : `int`
            Number of observation pairs used in the kernel.
        """
        kernelLength = self.shape[0] * self.shape[1]
        smallLength = int((self.shape[0] - 1)*(self.shape[1] - 1)/4)
        if self.level == 'AMP':
            nObservations = set([len(self.rawMeans[amp]) for amp in self.rawMeans])
            if len(nObservations) != 1:
                raise RuntimeError("Inconsistent number of observations found.")
            nObs = nObservations.pop()
        else:
            nObs = 0
 
        return (kernelLength, smallLength, nObs)
 
    @classmethod
    def fromDict(cls, dictionary):
        """Construct a calibration from a dictionary of properties.
 
        Parameters
        ----------
        dictionary : `dict`
            Dictionary of properties.
 
        Returns
        -------
        calib : `lsst.ip.isr.BrighterFatterKernel`
            Constructed calibration.
 
        Raises
        ------
        RuntimeError
            Raised if the supplied dictionary is for a different
            calibration.
            Raised if the version of the supplied dictionary is 1.0.
        """
        calib = cls()
 
        if calib._OBSTYPE != (found := dictionary['metadata']['OBSTYPE']):
            raise RuntimeError(f"Incorrect brighter-fatter kernel supplied.  Expected {calib._OBSTYPE}, "
                               f"found {found}")
        calib.setMetadata(dictionary['metadata'])
        calib.calibInfoFromDict(dictionary)
 
        calib.level = dictionary['metadata'].get('LEVEL', 'AMP')
        calib.shape = (dictionary['metadata'].get('KERNEL_DX', 0),
                       dictionary['metadata'].get('KERNEL_DY', 0))
 
        calibVersion = dictionary['metadata']['bfk_VERSION']
        if calibVersion == 1.0:
            calib.log.debug("Old Version of brighter-fatter kernel found. Current version: "
                            f"{calib._VERSION}. The new attribute 'expIdMask' will be "
                            "populated with 'True' values, and the new attributes 'rawMeans' "
                            "and 'rawVariances' will be populated with the masked 'means' "
                            "and 'variances' values."
                            )
            # use 'means', because 'expIdMask' does not exist.
            calib.expIdMask = {amp: np.repeat(True, len(dictionary['means'][amp])) for amp in
                               dictionary['means']}
            calib.rawMeans = {amp: np.array(dictionary['means'][amp]) for amp in dictionary['means']}
            calib.rawVariances = {amp: np.array(dictionary['variances'][amp]) for amp in
                                  dictionary['variances']}
        elif calibVersion == 1.1:
            calib.expIdMask = {amp: np.array(dictionary['expIdMask'][amp]) for amp in dictionary['expIdMask']}
            calib.rawMeans = {amp: np.array(dictionary['rawMeans'][amp]) for amp in dictionary['rawMeans']}
            calib.rawVariances = {amp: np.array(dictionary['rawVariances'][amp]) for amp in
                                  dictionary['rawVariances']}
        else:
            raise RuntimeError(f"Unknown version for brighter-fatter kernel: {calibVersion}")
 
        # Lengths for reshape:
        _, smallLength, nObs = calib.getLengths()
        smallShapeSide = int(np.sqrt(smallLength))
 
        calib.rawXcorrs = {amp: np.array(dictionary['rawXcorrs'][amp]).reshape((nObs,
                                                                                smallShapeSide,
                                                                                smallShapeSide))
                           for amp in dictionary['rawXcorrs']}
 
        calib.gain = dictionary['gain']
        calib.noise = dictionary['noise']
 
        calib.meanXcorrs = {amp: np.array(dictionary['meanXcorrs'][amp]).reshape(calib.shape)
                            for amp in dictionary['rawXcorrs']}
        calib.ampKernels = {amp: np.array(dictionary['ampKernels'][amp]).reshape(calib.shape)
                            for amp in dictionary['ampKernels']}
        calib.valid = {amp: bool(value) for amp, value in dictionary['valid'].items()}
        calib.badAmps = [amp for amp, valid in dictionary['valid'].items() if valid is False]
 
        calib.detKernels = {det: np.array(dictionary['detKernels'][det]).reshape(calib.shape)
                            for det in dictionary['detKernels']}
 
        calib.updateMetadata()
        return calib
 
    def toDict(self):
        """Return a dictionary containing the calibration properties.
 
        The dictionary should be able to be round-tripped through
        `fromDict`.
 
        Returns
        -------
        dictionary : `dict`
            Dictionary of properties.
        """
        self.updateMetadata()
 
        outDict = {}
        metadata = self.getMetadata()
        outDict['metadata'] = metadata
 
        # Lengths for ravel:
        kernelLength, smallLength, nObs = self.getLengths()
 
        outDict['expIdMask'] = {amp: np.array(self.expIdMask[amp]).tolist() for amp in self.expIdMask}
        outDict['rawMeans'] = {amp: np.array(self.rawMeans[amp]).tolist() for amp in self.rawMeans}
        outDict['rawVariances'] = {amp: np.array(self.rawVariances[amp]).tolist() for amp in
                                   self.rawVariances}
 
        for amp in self.rawXcorrs.keys():
            # Check to see if we need to repack the data.
            correlationShape = np.array(self.rawXcorrs[amp]).shape
            if nObs != correlationShape[0]:
                if correlationShape[0] == np.sum(self.expIdMask[amp]):
                    # Repack data.
                    self.repackCorrelations(amp, correlationShape)
                else:
                    raise ValueError("Could not coerce rawXcorrs into appropriate shape "
                                     "(have %d correlations, but expect to see %d.",
                                     correlationShape[0], np.sum(self.expIdMask[amp]))
 
        outDict['rawXcorrs'] = {amp: np.array(self.rawXcorrs[amp]).reshape(nObs*smallLength).tolist()
                                for amp in self.rawXcorrs}
        outDict['badAmps'] = self.badAmps
        outDict['gain'] = self.gain
        outDict['noise'] = self.noise
 
        outDict['meanXcorrs'] = {amp: self.meanXcorrs[amp].reshape(kernelLength).tolist()
                                 for amp in self.meanXcorrs}
        outDict['ampKernels'] = {amp: self.ampKernels[amp].reshape(kernelLength).tolist()
                                 for amp in self.ampKernels}
        outDict['valid'] = self.valid
 
        outDict['detKernels'] = {det: self.detKernels[det].reshape(kernelLength).tolist()
                                 for det in self.detKernels}
        return outDict
 
    @classmethod
    def fromTable(cls, tableList):
        """Construct calibration from a list of tables.
 
        This method uses the `fromDict` method to create the
        calibration, after constructing an appropriate dictionary from
        the input tables.
 
        Parameters
        ----------
        tableList : `list` [`astropy.table.Table`]
            List of tables to use to construct the brighter-fatter
            calibration.
 
        Returns
        -------
        calib : `lsst.ip.isr.BrighterFatterKernel`
            The calibration defined in the tables.
        """
        ampTable = tableList[0]
 
        metadata = ampTable.meta
        inDict = dict()
        inDict['metadata'] = metadata
 
        amps = ampTable['AMPLIFIER']
 
        # Determine version for expected values.  The ``fromDict``
        # method can unpack either, but the appropriate fields need to
        # be supplied.
        calibVersion = metadata['bfk_VERSION']
 
        if calibVersion == 1.0:
            # We expect to find ``means`` and ``variances`` for this
            # case, and will construct an ``expIdMask`` from these
            # parameters in the ``fromDict`` method.
            rawMeanList = ampTable['MEANS']
            rawVarianceList = ampTable['VARIANCES']
 
            inDict['means'] = {amp: mean for amp, mean in zip(amps, rawMeanList)}
            inDict['variances'] = {amp: var for amp, var in zip(amps, rawVarianceList)}
        elif calibVersion == 1.1:
            # This will have ``rawMeans`` and ``rawVariances``, which
            # are filtered via the ``expIdMask`` fields.
            expIdMaskList = ampTable['EXP_ID_MASK']
            rawMeanList = ampTable['RAW_MEANS']
            rawVarianceList = ampTable['RAW_VARIANCES']
 
            inDict['expIdMask'] = {amp: mask for amp, mask in zip(amps, expIdMaskList)}
            inDict['rawMeans'] = {amp: mean for amp, mean in zip(amps, rawMeanList)}
            inDict['rawVariances'] = {amp: var for amp, var in zip(amps, rawVarianceList)}
        else:
            raise RuntimeError(f"Unknown version for brighter-fatter kernel: {calibVersion}")
 
        rawXcorrs = ampTable['RAW_XCORRS']
        gainList = ampTable['GAIN']
        noiseList = ampTable['NOISE']
 
        meanXcorrs = ampTable['MEAN_XCORRS']
        ampKernels = ampTable['KERNEL']
        validList = ampTable['VALID']
 
        inDict['rawXcorrs'] = {amp: kernel for amp, kernel in zip(amps, rawXcorrs)}
        inDict['gain'] = {amp: gain for amp, gain in zip(amps, gainList)}
        inDict['noise'] = {amp: noise for amp, noise in zip(amps, noiseList)}
        inDict['meanXcorrs'] = {amp: kernel for amp, kernel in zip(amps, meanXcorrs)}
        inDict['ampKernels'] = {amp: kernel for amp, kernel in zip(amps, ampKernels)}
        inDict['valid'] = {amp: bool(valid) for amp, valid in zip(amps, validList)}
 
        inDict['badAmps'] = [amp for amp, valid in inDict['valid'].items() if valid is False]
 
        if len(tableList) > 1:
            detTable = tableList[1]
            inDict['detKernels'] = {det: kernel for det, kernel
                                    in zip(detTable['DETECTOR'], detTable['KERNEL'])}
        else:
            inDict['detKernels'] = {}
 
        return cls.fromDict(inDict)
 
    def toTable(self):
        """Construct a list of tables containing the information in this
        calibration.
 
        The list of tables should create an identical calibration
        after being passed to this class's fromTable method.
 
        Returns
        -------
        tableList : `list` [`lsst.afw.table.Table`]
            List of tables containing the crosstalk calibration
            information.
 
        """
        tableList = []
        self.updateMetadata()
 
        # Lengths
        kernelLength, smallLength, nObs = self.getLengths()
 
        ampList = []
        expIdMaskList = []
        rawMeanList = []
        rawVarianceList = []
        rawXcorrs = []
        gainList = []
        noiseList = []
 
        meanXcorrsList = []
        kernelList = []
        validList = []
 
        if self.level == 'AMP':
            for amp in self.rawMeans.keys():
                ampList.append(amp)
                expIdMaskList.append(self.expIdMask[amp])
                rawMeanList.append(self.rawMeans[amp])
                rawVarianceList.append(self.rawVariances[amp])
 
                correlationShape = np.array(self.rawXcorrs[amp]).shape
                if nObs != correlationShape[0]:
                    if correlationShape[0] == np.sum(self.expIdMask[amp]):
                        # Repack data.
                        self.repackCorrelations(amp, correlationShape)
                    else:
                        raise ValueError("Could not coerce rawXcorrs into appropriate shape "
                                         "(have %d correlations, but expect to see %d.",
                                         correlationShape[0], np.sum(self.expIdMask[amp]))
 
                rawXcorrs.append(np.array(self.rawXcorrs[amp]).reshape(nObs*smallLength).tolist())
                gainList.append(self.gain[amp])
                noiseList.append(self.noise[amp])
 
                meanXcorrsList.append(self.meanXcorrs[amp].reshape(kernelLength).tolist())
                kernelList.append(self.ampKernels[amp].reshape(kernelLength).tolist())
                validList.append(int(self.valid[amp] and not (amp in self.badAmps)))
 
        ampTable = Table({'AMPLIFIER': ampList,
                          'EXP_ID_MASK': expIdMaskList,
                          'RAW_MEANS': rawMeanList,
                          'RAW_VARIANCES': rawVarianceList,
                          'RAW_XCORRS': rawXcorrs,
                          'GAIN': gainList,
                          'NOISE': noiseList,
                          'MEAN_XCORRS': meanXcorrsList,
                          'KERNEL': kernelList,
                          'VALID': validList,
                          })
 
        ampTable.meta = self.getMetadata().toDict()
        tableList.append(ampTable)
 
        if len(self.detKernels):
            detList = []
            kernelList = []
            for det in self.detKernels.keys():
                detList.append(det)
                kernelList.append(self.detKernels[det].reshape(kernelLength).tolist())
 
            detTable = Table({'DETECTOR': detList,
                              'KERNEL': kernelList})
            detTable.meta = self.getMetadata().toDict()
            tableList.append(detTable)
 
        return tableList
 
    def repackCorrelations(self, amp, correlationShape):
        """If the correlations were masked, they need to be repacked into the
        correct shape.
 
        Parameters
        ----------
        amp : `str`
            Amplifier needing repacked.
        correlationShape : `tuple` [`int`], (3, )
            Shape the correlations are expected to take.
        """
        repackedCorrelations = []
        idx = 0
        for maskValue in self.expIdMask[amp]:
            if maskValue:
                repackedCorrelations.append(self.rawXcorrs[amp][idx])
                idx += 1
            else:
                repackedCorrelations.append(np.full((correlationShape[1], correlationShape[2]), np.nan))
        self.rawXcorrs[amp] = repackedCorrelations
 
    # Implementation methods
    def makeDetectorKernelFromAmpwiseKernels(self, detectorName, ampsToExclude=[]):
        """Average the amplifier level kernels to create a detector level
        kernel.  There is no change in index ordering/orientation from
        this averaging.
 
        Parameters
        ----------
        detectorName : `str`
            Detector for which the averaged kernel will be used.
        ampsToExclude : `list` [`str`], optional
            Amps that should not be included in the average.
        """
        inKernels = np.array([self.ampKernels[amp] for amp in
                              self.ampKernels if (self.valid[amp] and amp not in ampsToExclude)])
        avgKernel = np.zeros_like(inKernels[0])
        sctrl = afwMath.StatisticsControl()
        sctrl.setNumSigmaClip(5.0)
        for i in range(np.shape(avgKernel)[0]):
            for j in range(np.shape(avgKernel)[1]):
                avgKernel[i, j] = afwMath.makeStatistics(inKernels[:, i, j],
                                                         afwMath.MEANCLIP, sctrl).getValue()
 
        self.detKernels[detectorName] = avgKernel
 
    def replaceDetectorKernelWithAmpKernel(self, ampName, detectorName):
        self.detKernel[detectorName] = self.ampKernel[ampName]
 
 
def brighterFatterCorrection(exposure, kernel, maxIter, threshold, applyGain, gains=None):
    """Apply brighter fatter correction in place for the image.
 
    Parameters
    ----------
    exposure : `lsst.afw.image.Exposure`
        Exposure to have brighter-fatter correction applied.  Modified
        by this method.
    kernel : `numpy.ndarray`
        Brighter-fatter kernel to apply.
    maxIter : scalar
        Number of correction iterations to run.
    threshold : scalar
        Convergence threshold in terms of the sum of absolute
        deviations between an iteration and the previous one.
    applyGain : `Bool`
        If True, then the exposure values are scaled by the gain prior
        to correction.
    gains : `dict` [`str`, `float`]
        A dictionary, keyed by amplifier name, of the gains to use.
        If gains is None, the nominal gains in the amplifier object are used.
 
    Returns
    -------
    diff : `float`
        Final difference between iterations achieved in correction.
    iteration : `int`
        Number of iterations used to calculate correction.
 
    Notes
    -----
    This correction takes a kernel that has been derived from flat
    field images to redistribute the charge.  The gradient of the
    kernel is the deflection field due to the accumulated charge.
 
    Given the original image I(x) and the kernel K(x) we can compute
    the corrected image Ic(x) using the following equation:
 
    Ic(x) = I(x) + 0.5*d/dx(I(x)*d/dx(int( dy*K(x-y)*I(y))))
 
    To evaluate the derivative term we expand it as follows:
 
    0.5 * ( d/dx(I(x))*d/dx(int(dy*K(x-y)*I(y)))
        + I(x)*d^2/dx^2(int(dy* K(x-y)*I(y))) )
 
    Because we use the measured counts instead of the incident counts
    we apply the correction iteratively to reconstruct the original
    counts and the correction.  We stop iterating when the summed
    difference between the current corrected image and the one from
    the previous iteration is below the threshold.  We do not require
    convergence because the number of iterations is too large a
    computational cost.  How we define the threshold still needs to be
    evaluated, the current default was shown to work reasonably well
    on a small set of images.  For more information on the method see
    DocuShare Document-19407.
 
    The edges as defined by the kernel are not corrected because they
    have spurious values due to the convolution.
    """
    image = exposure.getMaskedImage().getImage()
 
    # The image needs to be units of electrons/holes
    with gainContext(exposure, image, applyGain, gains):
 
        kLx = np.shape(kernel)[0]
        kLy = np.shape(kernel)[1]
        kernelImage = afwImage.ImageD(kLx, kLy)
        kernelImage.getArray()[:, :] = kernel
        tempImage = afwImage.ImageD(image, deep=True)
 
        nanIndex = np.isnan(tempImage.getArray())
        tempImage.getArray()[nanIndex] = 0.
 
        corr = np.zeros(image.array.shape, dtype=np.float64)
        prev_image = np.zeros(image.array.shape, dtype=np.float64)
 
        # Define boundary by convolution region.  The region that the
        # correction will be calculated for is one fewer in each dimension
        # because of the second derivative terms.
        # NOTE: these need to use integer math, as we're using start:end as
        # numpy index ranges.
        startX = kLx//2
        endX = -kLx//2
        startY = kLy//2
        endY = -kLy//2
 
        for iteration in range(maxIter):
 
            outArray = scipy.signal.convolve(
                tempImage.array,
                kernelImage.array,
                mode="same",
                method="fft",
            )
            tmpArray = tempImage.getArray()
 
            with np.errstate(invalid="ignore", over="ignore"):
                # First derivative term
                gradTmp = np.gradient(tmpArray[startY:endY, startX:endX])
                gradOut = np.gradient(outArray[startY:endY, startX:endX])
                first = (gradTmp[0]*gradOut[0] + gradTmp[1]*gradOut[1])[1:-1, 1:-1]
 
                # Second derivative term
                diffOut20 = np.diff(outArray, 2, 0)[startY:endY, startX + 1:endX - 1]
                diffOut21 = np.diff(outArray, 2, 1)[startY + 1:endY - 1, startX:endX]
                second = tmpArray[startY + 1:endY - 1, startX + 1:endX - 1]*(diffOut20 + diffOut21)
 
                corr[startY + 1:endY - 1, startX + 1:endX - 1] = 0.5*(first + second)
 
                tmpArray[:, :] = image.getArray()[:, :]
                tmpArray[nanIndex] = 0.
                tmpArray[startY:endY, startX:endX] += corr[startY:endY, startX:endX]
 
            if iteration > 0:
                diff = np.sum(np.abs(prev_image - tmpArray), dtype=np.float64)
 
                if diff < threshold:
                    break
                prev_image[:, :] = tmpArray[:, :]
 
        image.getArray()[startY + 1:endY - 1, startX + 1:endX - 1] += \
            corr[startY + 1:endY - 1, startX + 1:endX - 1]
 
    return diff, iteration
 
 
def transferFlux(cFunc, fStep, correctionMode=True):
    """Take the input convolved deflection potential and the flux array
    to compute and apply the flux transfer into the correction array.
 
    Parameters
    ----------
    cFunc: `np.array`
        Deflection potential, being the convolution of the flux F with the
        kernel K.
    fStep: `np.array`
        The array of flux values which act as the source of the flux transfer.
    correctionMode: `bool`
        Defines if applying correction (True) or generating sims (False).
 
    Returns
    -------
    corr:
        BFE correction array
    """
 
    if cFunc.shape != fStep.shape:
        raise RuntimeError(f'transferFlux: array shapes do not match: {cFunc.shape}, {fStep.shape}')
 
    # set the sign of the correction and set its value for the
    # time averaged solution
    if correctionMode:
        # negative sign if applying BFE correction
        factor = -0.5
    else:
        # positive sign if generating BFE simulations
        factor = 0.5
 
    # initialise the BFE correction image to zero
    corr = np.zeros(cFunc.shape, dtype=np.float64)
 
    # Generate a 2D mesh of x,y coordinates
    yDim, xDim = cFunc.shape
    y = np.arange(yDim, dtype=int)
    x = np.arange(xDim, dtype=int)
    xc, yc = np.meshgrid(x, y)
 
    # process each axis in turn
    for ax in [0, 1]:
 
        # gradient of phi on right/upper edge of pixel
        diff = np.diff(cFunc, axis=ax)
 
        # expand array back to full size with zero gradient at the end
        gx = np.zeros(cFunc.shape, dtype=np.float64)
        yDiff, xDiff = diff.shape
        gx[:yDiff, :xDiff] += diff
 
        # select pixels with either positive gradients on the right edge,
        # flux flowing to the right/up
        # or negative gradients, flux flowing to the left/down
        for i, sel in enumerate([gx > 0, gx < 0]):
            xSelPixels = xc[sel]
            ySelPixels = yc[sel]
            # and add the flux into the pixel to the right or top
            # depending on which axis we are handling
            if ax == 0:
                xPix = xSelPixels
                yPix = ySelPixels+1
            else:
                xPix = xSelPixels+1
                yPix = ySelPixels
            # define flux as the either current pixel value or pixel
            # above/right
            # depending on whether positive or negative gradient
            if i == 0:
                # positive gradients, flux flowing to higher coordinate values
                flux = factor * fStep[sel]*gx[sel]
            else:
                # negative gradients, flux flowing to lower coordinate values
                flux = factor * fStep[yPix, xPix]*gx[sel]
            # change the fluxes of the donor and receiving pixels
            # such that flux is conserved
            corr[sel] -= flux
            corr[yPix, xPix] += flux
 
    # return correction array
    return corr
 
 
def fluxConservingBrighterFatterCorrection(exposure, kernel, maxIter, threshold, applyGain,
                                           gains=None, correctionMode=True):
    """Apply brighter fatter correction in place for the image.
 
    This version presents a modified version of the algorithm
    found in ``lsst.ip.isr.isrFunctions.brighterFatterCorrection``
    which conserves the image flux, resulting in improved
    correction of the cores of stars. The convolution has also been
    modified to mitigate edge effects.
 
    Parameters
    ----------
    exposure : `lsst.afw.image.Exposure`
        Exposure to have brighter-fatter correction applied.  Modified
        by this method.
    kernel : `np.ndarray`
        Brighter-fatter kernel to apply.
    maxIter : scalar
        Number of correction iterations to run.
    threshold : scalar
        Convergence threshold in terms of the sum of absolute
        deviations between an iteration and the previous one.
    applyGain : `Bool`
        If True, then the exposure values are scaled by the gain prior
        to correction.
    gains : `dict` [`str`, `float`]
        A dictionary, keyed by amplifier name, of the gains to use.
        If gains is None, the nominal gains in the amplifier object are used.
    correctionMode : `Bool`
        If True (default) the function applies correction for BFE.  If False,
        the code can instead be used to generate a simulation of BFE (sign
        change in the direction of the effect)
 
    Returns
    -------
    diff : `float`
        Final difference between iterations achieved in correction.
    iteration : `int`
        Number of iterations used to calculate correction.
 
    Notes
    -----
    Modified version of ``lsst.ip.isr.isrFunctions.brighterFatterCorrection``.
 
    This correction takes a kernel that has been derived from flat
    field images to redistribute the charge.  The gradient of the
    kernel is the deflection field due to the accumulated charge.
 
    Given the original image I(x) and the kernel K(x) we can compute
    the corrected image Ic(x) using the following equation:
 
    Ic(x) = I(x) + 0.5*d/dx(I(x)*d/dx(int( dy*K(x-y)*I(y))))
 
    Improved algorithm at this step applies the divergence theorem to
    obtain a pixelised correction.
 
    Because we use the measured counts instead of the incident counts
    we apply the correction iteratively to reconstruct the original
    counts and the correction.  We stop iterating when the summed
    difference between the current corrected image and the one from
    the previous iteration is below the threshold.  We do not require
    convergence because the number of iterations is too large a
    computational cost.  How we define the threshold still needs to be
    evaluated, the current default was shown to work reasonably well
    on a small set of images.
 
    Edges are handled in the convolution by padding.  This is still not
    a physical model for the edge, but avoids discontinuity in the correction.
 
    Author of modified version: Lance.Miller@physics.ox.ac.uk
    (see DM-38555).
    """
    image = exposure.getMaskedImage().getImage()
 
    # The image needs to be units of electrons/holes
    with gainContext(exposure, image, applyGain, gains):
 
        # get kernel and its shape
        kLy, kLx = kernel.shape
        kernelImage = afwImage.ImageD(kLx, kLy)
        kernelImage.getArray()[:, :] = kernel
        tempImage = afwImage.ImageD(image, deep=True)
 
        nanIndex = np.isnan(tempImage.getArray())
        tempImage.getArray()[nanIndex] = 0.
 
        outImage = afwImage.ImageD(image.getDimensions())
        corr = np.zeros(image.array.shape, dtype=np.float64)
        prevImage = np.zeros(image.array.shape, dtype=np.float64)
        convCntrl = afwMath.ConvolutionControl(False, False, 1)
        fixedKernel = afwMath.FixedKernel(kernelImage)
 
        # set the padding amount
        # ensure we pad by an even amount larger than the kernel
        kLy = 2 * ((1+kLy)//2)
        kLx = 2 * ((1+kLx)//2)
 
        # The deflection potential only depends on the gradient of
        # the convolution, so we can subtract the mean, which then
        # allows us to pad the image with zeros and avoid wrap-around effects
        # (although still not handling the image edges with a physical model)
        # This wouldn't be great if there were a strong image gradient.
        imYdimension, imXdimension = tempImage.array.shape
        imean = np.mean(tempImage.getArray()[~nanIndex], dtype=np.float64)
        # subtract mean from image
        tempImage -= imean
        tempImage.array[nanIndex] = 0.0
        padArray = np.pad(tempImage.getArray(), ((0, kLy), (0, kLx)))
        outImage = afwImage.ImageD(np.pad(outImage.getArray(), ((0, kLy), (0, kLx))))
        # Convert array to afw image so afwMath.convolve works
        padImage = afwImage.ImageD(padArray.shape[1], padArray.shape[0])
        padImage.array[:] = padArray
 
        for iteration in range(maxIter):
 
            # create deflection potential, convolution of flux with kernel
            # using padded counts array
            afwMath.convolve(outImage, padImage, fixedKernel, convCntrl)
            tmpArray = tempImage.getArray()
            outArray = outImage.getArray()
 
            # trim convolution output back to original shape
            outArray = outArray[:imYdimension, :imXdimension]
 
            # generate the correction array, with correctionMode set as input
            corr[...] = transferFlux(outArray, tmpArray, correctionMode=correctionMode)
 
            # update the arrays for the next iteration
            tmpArray[:, :] = image.getArray()[:, :]
            tmpArray += corr
            tmpArray[nanIndex] = 0.
            # update padded array
            # subtract mean
            tmpArray -= imean
            tempImage.array[nanIndex] = 0.
            padArray = np.pad(tempImage.getArray(), ((0, kLy), (0, kLx)))
 
            if iteration > 0:
                diff = np.sum(np.abs(prevImage - tmpArray), dtype=np.float64)
 
                if diff < threshold:
                    break
                prevImage[:, :] = tmpArray[:, :]
 
        image.getArray()[:] += corr[:]
 
    return diff, iteration