Refine variance correction method and add comprehensive unit test

python/lsst/meas/algorithms/__init__.py

@@ -1,10 +1,10 @@
-# This file is part of meas_algorithms.
 #
-# 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.
+# LSST Data Management System
+#
+# Copyright 2008-2017  AURA/LSST.
+#
+# This product includes software developed by the
+# LSST Project (http://www.lsst.org/).
 #
 # 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
@@ -16,8 +16,10 @@
 # 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/>.
+# You should have received a copy of the LSST License Statement and
+# the GNU General Public License along with this program.  If not,
+# see <https://www.lsstcorp.org/LegalNotices/>.
+#
 """lsst.meas.algorithms
 """
 
@@ -60,12 +62,9 @@
 from .scaleVariance import *
 from .noise_covariance import *
 from .reinterpolate_pixels import *
-<<<<<<< HEAD
 from .setPrimaryFlags import *
-=======
 from .variance_plane import *
->>>>>>> 3553a488 (Add function to remove Poisson contribution from source from variance plane.)
 
 from .version import *
 
-import lsst.utils
+import lsst.utils

python/lsst/meas/algorithms/variance_plane.py

@@ -25,55 +25,99 @@
 import numpy as np
 
 
-__all__ = ['remove_signal_from_variance']
+__all__ = ["remove_signal_from_variance"]
 
 
-def remove_signal_from_variance(exposure, gain=None, average_across_amps=False, in_place=False):
-    """Remove the Poisson contribution by actual sources from the variance
-    plane of an Exposure.
+def remove_signal_from_variance(exposure, gain=None, gains=None, average_across_amps=False, in_place=False):
+    """
+    Removes the Poisson contribution from actual sources in the variance plane
+    of an Exposure.
 
-    If no gain value is provided, it will be estimated as a linear fit of
-    variance versus image plane. This estimation can be carried out on the
-    whole image at once, or separately for each amplifier.
+    If neither gain nor gains are provided, the function estimates the gain(s).
+    If 'average_across_amps' is True, a single gain value for the entire image
+    is estimated. If False, individual gain values for each amplifier are
+    estimated. The estimation involves a linear fit of variance versus image
+    plane.
 
     Parameters
-    -----------
-    exposure : `afw.image.Exposure`
-        Exposure that contains a variance plane that should be corrected for
+    ----------
+    exposure : `~lsst.afw.image.Exposure`
+        The exposure containing a variance plane to be corrected for
         source contributions.
     gain : `float`, optional
-        The gain value for the whole image. If not provided (the default),
-        will be estimated from the image and variance planes.
+        The gain value for the entire image. This parameter is used if 'gains'
+        is not provided. If both 'gain' and 'gains' are None, and
+        'average_across_amps' is True, 'gain' is estimated from the image and
+        variance planes.
+    gains : `list` of `float`, optional
+        A list of gain values for each amplifier. This parameter is used if
+        'gain' is not provided. If both 'gain' and 'gains' are None, and
+        'average_across_amps' is False, 'gains' are estimated from the image
+        and variance planes.
     average_across_amps : `bool`, optional
-        Whether the gain should be estimated on the whole image at once. If
-        False (the default), a different gain value is estimated for each
-        amplifier. Ignored if gain is not None.
+        Determines the gain estimation strategy. If True, the gain for the
+        entire image is estimated at once. If False, individual gains for each
+        amplifier are estimated. This parameter is ignored if either 'gain' or
+        'gains' is specified.
     in_place : `bool`, optional
-        If True, the variance plane is changed in place. Defaults to False.
+        If True, the variance plane of the input Exposure is modified in place.
+        If False (default), a modified copy of the variance plane is returned.
+
+    Returns
+    -------
+    variance_plane : `~lsst.afw.image.MaskedImage`, optional
+        The corrected variance plane, with the signal contribution removed.
+        This is only returned if 'in_place' is False.
+
+    Raises
+    ------
+    AttributeError
+        If amplifiers cannot be retrieved from the exposure without providing
+        'average_across_amps=True'.
+    ValueError
+        If both 'gain' and 'gains' are provided, or if the number of provided
+        'gains' does not match the number of amplifiers.
     """
     variance_plane = exposure.variance if in_place else exposure.variance.clone()
-    if average_across_amps:
-        amp_bboxes = [exposure.getBBox()]
-    else:
-        try:
-            amps = exposure.getDetector().getAmplifiers()
-            amp_bboxes = [amp.getBBox() for amp in amps]
-        except AttributeError:
-            raise AttributeError("Could not retrieve amplifiers from exposure. To compute a simple gain "
-                                 "value across the entire image, use average_across_amps=True.")
-    if gain is None:
-        # Fit a straight line to variance vs (sky-subtracted) signal.
-        # The evaluate that line at zero signal to get an estimate of the
+    if gain is None and gains is None:
+        if average_across_amps:
+            amp_bboxes = [exposure.getBBox()]
+        else:
+            try:
+                amps = exposure.getDetector().getAmplifiers()
+                amp_bboxes = [amp.getBBox() for amp in amps]
+            except AttributeError:
+                raise AttributeError(
+                    "Could not retrieve amplifiers from exposure. To compute a simple gain value across the "
+                    "entire image, use average_across_amps=True."
+                )
+        # Fit a straight line to variance vs (sky-subtracted) signal. Then
+        # evaluate that line at zero signal to get an estimate of the
         # signal-free variance.
         for amp_bbox in amp_bboxes:
             amp_im_arr = exposure[amp_bbox].image.array
             amp_var_arr = variance_plane[amp_bbox].array
             good = (amp_var_arr != 0) & np.isfinite(amp_var_arr) & np.isfinite(amp_im_arr)
             fit = np.polyfit(amp_im_arr[good], amp_var_arr[good], deg=1)
-            # fit is [1/gain, sky_var]
-            gain = 1./fit[0]
-            variance_plane[amp_bbox].array[good] -= amp_im_arr[good]/gain
-    else:
+            # Fit is [1/gain, sky_var].
+            amp_gain = 1.0 / fit[0]
+            variance_plane[amp_bbox].array[good] -= amp_im_arr[good] / amp_gain
+    elif gain is None and gains is not None:
+        amps = exposure.getDetector().getAmplifiers()
+        amp_bboxes = [amp.getBBox() for amp in amps]
+        namps = len(amps)
+        if len(gains) != namps:
+            raise ValueError(
+                f"Incorrect number of gains provided: {len(gains)} values for {namps} amplifiers."
+            )
+        for amp_bbox, amp_gain in zip(amp_bboxes, gains):
+            im_arr = exposure[amp_bbox].image.array
+            variance_plane[amp_bbox].array -= im_arr / amp_gain
+    elif gain is not None and gains is None:
         im_arr = exposure.image.array
-        variance_plane.array -= im_arr/gain
-    return variance_plane
+        variance_plane.array -= im_arr / gain
+    elif gain is not None and gains is not None:
+        raise ValueError("Both gain and gains are provided. Please provide only one of them or none at all "
+                         "in case of automatic gain estimation from the image and variance planes.")
+    if not in_place:
+        return variance_plane