add mypy

python/lsst/meas/algorithms/computeExPsf.py

@@ -21,8 +21,10 @@
 
 from lsst.pex.config import Config, Field
 from lsst.pipe.base import Task
+import lsst.pipe.base as pipeBase
 import treecorr
 import copy
+import numpy.typing as npt
 
 
 __all__ = (
@@ -53,35 +55,44 @@ class ComputeExPsfTask(Task):
 
     Parameters
     ----------
-    de1: `numpy.array`
+    de1: `np.ndarray`
         PSF ellipticity residuals component 1.
-    de2: `numpy.array`
+    de2: `np.ndarray`
         PSF ellipticity residuals component 2.
-    ra: `numpy.array`
+    ra: `np.ndarray`
         Right ascension coordinate.
-    dec: `numpy.array`
+    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
     -------
-    kk_E1: `float`
-        <de1 de1> scalar correlation function, compute
-        in an angular bin define in ComputeExPsfConfig.
-    kk_E2: `float`
-        <de2 de2> scalar correlation function, compute
-        in an angular bin define in ComputeExPsfConfig.
-    kk_Ex: `float`
-        <de1 de2> scalar cross-correlation function, compute
-        in an angular bin define in ComputeExPsfConfig.
+    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 ComputeExPsfConfig.
+        ``E2``: `float`
+            <de2 de2> scalar correlation function, compute
+            in an angular bin define in ComputeExPsfConfig.
+        ``Ex``: `float`
+            <de1 de2> scalar cross-correlation function, compute
+            in an angular bin define in ComputeExPsfConfig.
     """
 
     ConfigClass = ComputeExPsfConfig
     _DefaultName = "computeExPsf"
 
-    def run(self, de1, de2, ra, dec, units="arcmin"):
+    def run(
+        self,
+        de1: npt.NDArray,
+        de2: npt.NDArray,
+        ra: npt.NDArray,
+        dec: npt.NDArray,
+        units: str = "arcmin",
+    ) -> pipeBase.Struct:
 
         kwargs_cat = {
             "ra": ra,
@@ -110,4 +121,4 @@ def run(self, de1, de2, ra, dec, units="arcmin"):
         kk.process(cat1, cat2)
         kk_Ex = copy.deepcopy(kk.xi[0])
 
-        return kk_E1, kk_E2, kk_Ex
+        return pipeBase.Struct(E1=kk_E1, E2=kk_E2, Ex=kk_Ex)

tests/test_computeExPsf.py

@@ -21,15 +21,20 @@
 
 
 import unittest
+from typing import Optional
+import numpy.typing as npt
 
 import numpy as np
 
 import lsst.utils.tests
 from lsst.meas.algorithms import GaussianProcessTreegp
 from lsst.meas.algorithms import ComputeExPsfTask
+import lsst.pipe.base as pipeBase
 
 
-def rbf_kernel(x1, x2, sigma, correlation_length):
+def rbf_kernel(
+    x1: npt.NDArray, x2: npt.NDArray, sigma: float, correlation_length: float
+) -> npt.NDArray:
     """
     Computes the radial basis function (RBF) kernel matrix.
 
@@ -46,7 +51,7 @@ def rbf_kernel(x1, x2, sigma, correlation_length):
 
     Returns:
     --------
-    kernel : `np.array`
+    kernel : `np.ndarray`
         RBF kernel matrix with shape (n_samples, n_samples).
     """
     distance_squared = np.sum((x1[:, None, :] - x2[None, :, :]) ** 2, axis=-1)
@@ -55,18 +60,18 @@ def rbf_kernel(x1, x2, sigma, correlation_length):
 
 
 def generate_gaussian_random_field(
-    xmin=0,
-    xmax=2000,
-    ymin=0,
-    ymax=2000,
-    npoints=10000,
-    nmax=1000,
-    std=1.0,
-    correlation_length=10.0,
-    white_noise=1.0,
-    seed=42,
-    input_coord=None,
-):
+    xmin: int = 0,
+    xmax: int = 2000,
+    ymin: int = 0,
+    ymax: int = 2000,
+    npoints: int = 10000,
+    nmax: int = 1000,
+    std: float = 1.0,
+    correlation_length: float = 10.0,
+    white_noise: float = 1.0,
+    seed: int = 42,
+    input_coord: Optional[npt.NDArray] = None,
+) -> pipeBase.Struct:
     """Generate a Gaussian Random Field.
 
     Function to generate a Gaussian Random Field.
@@ -108,16 +113,18 @@ def generate_gaussian_random_field(
     seed: `int`
         Seed of the random generator.
         Default: ``42``
-    input_coord: `np.array`
+    input_coord: `np.ndarray`
         Take a input coord to generate the Gaussian Random field
         Default: ``None``
 
     Returns
     -------
-    coord: `np.array`
-        2D coordinate of the gaussian random field
-    z: `np.array`
-        Scalar value of the gaussian random field
+    struct : `lsst.pipe.base.Struct`
+            The struct contains the following data:
+        ``coord``: `np.ndarray`
+            2D coordinate of the gaussian random field
+        ``z``: `np.ndarray`
+            Scalar value of the gaussian random field
     """
     np.random.seed(seed)
 
@@ -129,7 +136,7 @@ def generate_gaussian_random_field(
     else:
         ngenerated = npoints
 
-    if input_coord is not None or npoints > nmax:
+    if input_coord is None or npoints > nmax:
         x1 = np.random.uniform(xmin, xmax, ngenerated)
         x2 = np.random.uniform(ymin, ymax, ngenerated)
         coord1 = np.array([x1, x2]).T
@@ -164,46 +171,49 @@ def generate_gaussian_random_field(
         coord = coord1
         z = z1
 
-    return coord, z
+    return pipeBase.Struct(coord=coord, z=z)
 
 
 class ComputeExPsfTestCase(lsst.utils.tests.TestCase):
     """Test ComputeExPsfTask."""
 
-    def setUp(self):
+    def setUp(self) -> None:
         super().setUp()
 
-        kwargs_grf = {
-            "xmin": 0,
-            "xmax": 2000,
-            "ymin": 0,
-            "ymax": 2000,
-            "npoints": 10000,
-            "nmax": 2000,
-            "std": 1.0,
-            "white_noise": 0.01,
-        }
-
-        coord1, de1 = generate_gaussian_random_field(
+        output1 = generate_gaussian_random_field(
+            xmin=0,
+            xmax=2000,
+            ymin=0,
+            ymax=2000,
+            npoints=10000,
+            nmax=2000,
+            std=1.0,
             correlation_length=200.0,
+            white_noise=0.01,
             seed=42,
             input_coord=None,
-            **kwargs_grf,
         )
 
-        coord2, de2 = generate_gaussian_random_field(
+        output2 = generate_gaussian_random_field(
+            xmin=0,
+            xmax=2000,
+            ymin=0,
+            ymax=2000,
+            npoints=10000,
+            nmax=2000,
+            std=1.0,
             correlation_length=100.0,
+            white_noise=0.01,
             seed=44,
-            input_coord=coord1,
-            **kwargs_grf,
+            input_coord=output1.coord,
         )
 
-        self.coord1 = coord1
-        self.coord2 = coord2
-        self.de1 = de1
-        self.de2 = de2
+        self.coord1 = output1.coord
+        self.coord2 = output2.coord
+        self.de1 = output1.z
+        self.de2 = output2.z
 
-    def test_comp_ex_psf(self):
+    def test_comp_ex_psf(self) -> None:
         """Test that ex metric are compute and make sense."""
 
         np.testing.assert_equal(self.coord1, self.coord2)
@@ -216,15 +226,15 @@ def test_comp_ex_psf(self):
         config.thetaMax = 5.0
 
         task = ComputeExPsfTask(config)
-        kk_E1, kk_E2, kk_Ex = task.run(self.de1, self.de2, ra, dec, units="arcmin")
+        output1 = task.run(self.de1, self.de2, ra, dec, units="arcmin")
 
         # At small scale, expect the scalar two-point correlation function
         # to be close to the input variance for de1 and de2. Cross correlation
         # between de1 and de2 should be zeros are they are 2 indendant field.
 
-        np.testing.assert_allclose(kk_E1, 1.0, atol=2e-1)
-        np.testing.assert_allclose(kk_E2, 1.0, atol=2e-1)
-        np.testing.assert_allclose(kk_Ex, 0.0, atol=1e-1)
+        np.testing.assert_allclose(output1.E1, 1.0, atol=2e-1)
+        np.testing.assert_allclose(output1.E2, 1.0, atol=2e-1)
+        np.testing.assert_allclose(output1.Ex, 0.0, atol=1e-1)
 
         config = ComputeExPsfTask.ConfigClass()
         config.thetaMin = 600.0
@@ -234,11 +244,11 @@ def test_comp_ex_psf(self):
         # be all close to 0.
 
         task = ComputeExPsfTask(config)
-        kk_E1, kk_E2, kk_Ex = task.run(self.de1, self.de2, ra, dec, units="arcmin")
+        output2 = task.run(self.de1, self.de2, ra, dec, units="arcmin")
 
-        np.testing.assert_allclose(kk_E1, 0.0, atol=1e-1)
-        np.testing.assert_allclose(kk_E2, 0.0, atol=1e-1)
-        np.testing.assert_allclose(kk_Ex, 0.0, atol=1e-1)
+        np.testing.assert_allclose(output2.E1, 0.0, atol=1e-1)
+        np.testing.assert_allclose(output2.E2, 0.0, atol=1e-1)
+        np.testing.assert_allclose(output2.Ex, 0.0, atol=1e-1)
 
 
 def setup_module(module):