Merge pull request #389 from lsst/tickets/DM-45899

DM-45899: Write a Task to compute Ex for TEx

python/lsst/meas/algorithms/computeExPsf.py

@@ -0,0 +1,129 @@
+# 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.
+#
+# 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/>.
+
+
+from lsst.meas.algorithms.treecorrUtils import TreecorrConfig
+from lsst.pex.config import Config, ConfigField
+from lsst.pipe.base import Task
+import lsst.pipe.base as pipeBase
+import treecorr
+import copy
+import numpy.typing as npt
+
+
+__all__ = ("ComputeExPsfTask", "ComputeExPsfConfig")
+
+
+class ComputeExPsfConfig(Config):
+    """Config class of ComputeExPsfTask."""
+
+    treecorr = ConfigField(
+        dtype=TreecorrConfig,
+        doc="treecorr config.",
+    )
+
+    def setDefaults(self):
+        super().setDefaults()
+
+        self.treecorr.min_sep = 1.0 / 60.0
+        self.treecorr.max_sep = 5.0 / 60.0
+        self.treecorr.nbins = 1
+        self.treecorr.bin_type = "Linear"
+        self.treecorr.sep_units = "degree"
+
+
+class ComputeExPsfTask(Task):
+    """Compute Ex for PSF.
+
+    Compute scalar correlation function from
+    PSF ellipticity residuals to compute TEx
+    metrics.
+
+    Parameters
+    ----------
+    de1: `np.ndarray`
+        PSF ellipticity residuals component 1.
+    de2: `np.ndarray`
+        PSF ellipticity residuals component 2.
+    ra: `np.ndarray`
+        Right ascension coordinate.
+    dec: `np.ndarray`
+        Declination coordinate.
+    units: `str`
+        In which units are ra and dec. units supported
+        are the same as the one in treecorr.
+
+    Returns
+    -------
+    struct : `lsst.pipe.base.Struct`
+            The struct contains the following data:
+        ``E1``: `float`
+            <de1 de1> scalar correlation function, compute
+            in an angular bin define in TreecorrConfig.
+        ``E2``: `float`
+            <de2 de2> scalar correlation function, compute
+            in an angular bin define in TreecorrConfig.
+        ``Ex``: `float`
+            <de1 de2> scalar cross-correlation function, compute
+            in an angular bin define in TreecorrConfig.
+    """
+
+    ConfigClass = ComputeExPsfConfig
+    _DefaultName = "computeExPsf"
+
+    def run(
+        self,
+        de1: npt.NDArray,
+        de2: npt.NDArray,
+        ra: npt.NDArray,
+        dec: npt.NDArray,
+        units: str = "arcmin",
+    ) -> pipeBase.Struct:
+
+        if units != self.config.treecorr.sep_units:
+            raise ValueError(
+                "units from ComputeExPsfConfig and"
+                "ComputeExPsfTask are not the same (%s != %s)"
+                % ((units, self.config.treecorr.sep_units))
+            )
+
+        kwargs_cat = {
+            "ra": ra,
+            "dec": dec,
+            "ra_units": units,
+            "dec_units": units,
+        }
+
+        cat1 = treecorr.Catalog(k=de1, **kwargs_cat)
+        cat2 = treecorr.Catalog(k=de2, **kwargs_cat)
+
+        config_kk = self.config.treecorr.toDict()
+
+        kk = treecorr.KKCorrelation(config_kk)
+
+        kk.process(cat1)
+        kk_E1 = copy.deepcopy(kk.xi[0])
+        kk.process(cat2)
+        kk_E2 = copy.deepcopy(kk.xi[0])
+        kk.process(cat1, cat2)
+        kk_Ex = copy.deepcopy(kk.xi[0])
+
+        return pipeBase.Struct(metric_E1=kk_E1, metric_E2=kk_E2, metric_Ex=kk_Ex)

python/lsst/meas/algorithms/treecorrUtils.py

@@ -0,0 +1,228 @@
+# 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.
+#
+# 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/>.
+
+
+from lsst.pex.config import ChoiceField, Config, Field, FieldValidationError
+
+
+class TreecorrConfig(Config):
+    """A Config class that holds some of the parameters supported by treecorr.
+
+    The fields in this class correspond to the parameters that can be passed to
+    any calls to `treecorr` methods, including catalog creation and two-point
+    correlation function calculations. The default values set for the fields
+    are identical to the default values set in v4.3 of `treecorr`.
+
+    A separate config class is used instead
+    of constructing a `~lsst.pex.config.DictField` so that mixed types can be
+    supported and the config can be validated at the beginning of the
+    execution.
+
+    Notes
+    -----
+    This is intended to be used with correlations of PSF residuals. It only supports
+    some of the fields that are relevant for rho-statistics calculations and the likes
+    of it.
+    """
+
+    nbins = Field[int](
+        doc=(
+            "How many bins to use. "
+            "(Exactly three of nbins, bin_size, min_sep, max_sep "
+            "are required. If nbins is not given, it will be "
+            "calculated from the values of the other three, "
+            "rounding up to the next highest integer. "
+            "In this case, bin_size will be readjusted to account "
+            "for this rounding up."
+        ),
+        optional=True,
+        check=lambda x: x > 0,
+    )
+
+    bin_size = Field[float](
+        doc=(
+            "The width of the bins in log(separation). "
+            "Exactly three of nbins, bin_size, min_sep, max_sep are required. "
+            "If bin_size is not given, it will be calculated from the values "
+            "of the other three."
+        ),
+        optional=True,
+    )
+
+    min_sep = Field[float](
+        doc=(
+            "The minimum separation in units of sep_units, if relevant. "
+            "Exactly three of nbins, bin_size, min_sep, max_sep are required. "
+            "If min_sep is not given, it will be calculated from the values "
+            "of the other three."
+        ),
+        optional=True,
+    )
+
+    max_sep = Field[float](
+        doc=(
+            "The maximum separation in units of sep_units, if relevant. "
+            "Exactly three of nbins, bin_size, min_sep, max_sep are required. "
+            "If max_sep is not given, it will be calculated from the values "
+            "of the other three."
+        ),
+        optional=True,
+    )
+
+    sep_units = ChoiceField[str](
+        doc=(
+            "The units to use for the separation values, given as a string. "
+            "This includes both min_sep and max_sep above, as well as the "
+            "units of the output distance values."
+        ),
+        default=None,
+        optional=True,
+        allowed={
+            units: units for units in ["arcsec", "arcmin", "degree", "hour", "radian"]
+        },
+    )
+
+    bin_slop = Field[float](
+        doc=(
+            "How much slop to allow in the placement of pairs in the bins. "
+            "If bin_slop = 1, then the bin into which a particular pair is "
+            "placed may be incorrect by at most 1.0 bin widths. "
+            r"If None, use a bin_slop that gives a maximum error of 10% on "
+            "any bin, which has been found to yield good results for most "
+            "applications."
+        ),
+        default=None,
+        optional=True,
+    )
+
+    precision = Field[int](
+        doc=(
+            "The precision to use for the output values. This specifies how many digits to write."
+        ),
+        default=4,
+        optional=True,
+        check=lambda x: x > 0,
+    )
+
+    metric = ChoiceField[str](
+        doc=(
+            "Which metric to use for distance measurements. For details, see "
+            "https://rmjarvis.github.io/TreeCorr/_build/html/metric.html"
+        ),
+        default="Euclidean",
+        optional=True,
+        allowed={
+            "Euclidean": "straight-line Euclidean distance between two points",
+            "FisherRperp": (
+                "the perpendicular component of the distance, "
+                "following the definitions in "
+                "Fisher et al, 1994 (MNRAS, 267, 927)"
+            ),
+            "OldRperp": (
+                "the perpendicular component of the distance using the "
+                "definition of Rperp from TreeCorr v3.x."
+            ),
+            "Rlens": (
+                "Distance from the first object (taken to be a lens) to "
+                "the line connecting Earth and the second object "
+                "(taken to be a lensed source)."
+            ),
+            "Arc": "the true great circle distance for spherical coordinates.",
+            "Periodic": "Like ``Euclidean``, but with periodic boundaries.",
+        },
+    )
+
+    bin_type = ChoiceField[str](
+        doc="What type of binning should be used?",
+        default="Log",
+        optional=True,
+        allowed={
+            "Log": (
+                "Logarithmic binning in the distance. The bin steps will "
+                "be uniform in log(r) from log(min_sep) .. log(max_sep)."
+            ),
+            "Linear": (
+                "Linear binning in the distance. The bin steps will be "
+                "uniform in r from min_sep .. max_sep."
+            ),
+            "TwoD": (
+                "2-dimensional binning from x = (-max_sep .. max_sep) "
+                "and y = (-max_sep .. max_sep). The bin steps will be "
+                "uniform in both x and y. (i.e. linear in x,y)"
+            ),
+        },
+    )
+
+    var_method = ChoiceField[str](
+        doc="Which method to use for estimating the variance",
+        default="shot",
+        optional=True,
+        allowed={
+            method: method
+            for method in [
+                "shot",
+                "jackknife",
+                "sample",
+                "bootstrap",
+                "marked_bootstrap",
+            ]
+        },
+    )
+
+    npatch = Field[int](
+        doc="How many patches to split the catalog into for the purpose of "
+        "jackknife variance or other options that involve running via "
+        "patches (boostrap, marked_boostrap etc.)",
+        default=1,
+        optional=True,
+    )
+
+    num_bootstrap = Field[int](
+        doc=(
+            "How many bootstrap samples to use for the 'bootstrap' and 'marked_bootstrap' var methods."
+        ),
+        default=500,
+        optional=True,
+    )
+
+    rng = Field[int](
+        doc="Value to seed the treecorr random number generator with. Used to generate patches.",
+        default=None,
+        optional=True,
+    )
+
+    def validate(self):
+        # Docs inherited from base class
+        super().validate()
+        req_params = (self.nbins, self.bin_size, self.min_sep, self.max_sep)
+        num_req_params = sum(param is not None for param in req_params)
+        if num_req_params != 3:
+            msg = (
+                "You must specify exactly three of ``nbins``, ``bin_size``, ``min_sep`` and ``max_sep``"
+                f" in treecorr_config. {num_req_params} parameters were set instead."
+            )
+            raise FieldValidationError(self.__class__.bin_size, self, msg)
+
+        if self.min_sep is not None and self.max_sep is not None:
+            if self.min_sep > self.max_sep:
+                raise FieldValidationError(
+                    self.__class__.min_sep, self, "min_sep must be <= max_sep"
+                )