Support providing masks for star-finding

regularizepsf/fitter.py

@@ -5,7 +5,7 @@
 from collections import namedtuple
 from collections.abc import Callable
 from numbers import Real
-from typing import Any, Dict, Generator, List, Tuple
+from typing import Any, Dict, Generator, List, Optional, Tuple
 
 import deepdish as dd
 import numpy as np
@@ -271,7 +271,9 @@ def find_stars_and_average(cls,
                                interpolation_scale: int = 1,
                                average_mode: str = "median",
                                percentile: float = 10,
-                               star_threshold: int = 3, hdu_choice: int=0) -> CoordinatePatchCollection:
+                               star_threshold: int = 3,
+                               star_mask: Optional[list[str] | np.ndarray | Generator] = None,
+                               hdu_choice: int=0) -> CoordinatePatchCollection:
         """Loads a series of images, finds stars in each, 
             and builds a CoordinatePatchCollection with averaged stars
 
@@ -297,6 +299,14 @@ def find_stars_and_average(cls,
         star_threshold : int
             SEP's threshold for finding stars. See `threshold`
                 in https://sep.readthedocs.io/en/v1.1.x/api/sep.extract.html#sep-extract
+        star_mask : List[str] or np.ndarray or Generator
+            Masks to apply during star-finding. Can be a list of FITS filenames, a
+            numpy array of shape (n_images, ny, nx), or a Generator that yields
+            each mask array in turn. Where the mask pixel is `True`, the
+            corresponding data array pixel will not be selected as a star. See
+            `mask` in
+            https://sep.readthedocs.io/en/v1.1.x/api/sep.extract.html#sep-extract
+            for more details.
         hdu_choice : int
             Which HDU from each image will be used, 
                 default of 0 is most common but could be 1 for compressed images
@@ -312,24 +322,48 @@ def find_stars_and_average(cls,
             for large images can dramatically slow down the execution.
         """
         if isinstance(images, Generator):
-            iterator = images
+            data_iterator = images
         elif isinstance(images, np.ndarray):
             if len(images.shape) == 3:
                 def generator():
                     for image in images:
                         yield image
-                iterator = generator()
+                data_iterator = generator()
             else:
                 raise ValueError("Image data array must be 3D")
         elif isinstance(images, List) and isinstance(images[0], str):
             def generator():
                 for image_path in images:
                     with fits.open(image_path) as hdul:
                         yield hdul[hdu_choice].data.astype(float)
-            iterator = generator()
+            data_iterator = generator()
         else:
             raise ValueError("Unsupported type for `images`")
 
+        if star_mask is None:
+            def generator():
+                while True:
+                    yield None
+            star_mask_iterator = generator()
+        elif isinstance(star_mask, Generator):
+            star_mask_iterator = star_mask
+        elif isinstance(star_mask, np.ndarray):
+            if len(star_mask.shape) == 3:
+                def generator():
+                    for mask in star_mask:
+                        yield mask
+                star_mask_iterator = generator()
+            else:
+                raise ValueError("Star mask array must be 3D")
+        elif isinstance(star_mask, List) and isinstance(star_mask[0], str):
+            def generator():
+                for mask_path in star_mask:
+                    with fits.open(mask_path) as hdul:
+                        yield hdul[hdu_choice].data.astype(bool)
+            star_mask_iterator = generator()
+        else:
+            raise ValueError("Unsupported type for `star_mask`")
+
         # the output collection to return
         this_collection = cls({})
 
@@ -338,7 +372,7 @@ def generator():
         image_shape = None
 
         # for each image do the magic
-        for i, image in enumerate(iterator):
+        for i, (image, star_mask) in enumerate(zip(data_iterator, star_mask_iterator)):
             if image_shape is None:
                 image_shape = image.shape
             elif image.shape != image_shape:
@@ -363,7 +397,8 @@ def generator():
             image_background_removed = image - background
             image_star_coords = sep.extract(image_background_removed, 
                                             star_threshold, 
-                                            err=background.globalrms)
+                                            err=background.globalrms,
+                                            mask=star_mask)
 
             coordinates = [CoordinateIdentifier(i,
                                                 int(round(y - psf_size * interpolation_scale / 2)),

tests/test_fitter.py

@@ -160,17 +160,26 @@ def test_find_stars_and_average_powers_of_2_mean():
 
 
 def test_find_stars_and_average_image_formats():
-    # Run find_stars_and_average with the three possible input-data formats
+    """Runs find_stars_and_average with the three possible input-data formats"""
     img_paths = [str(TEST_DIR / "data/DASH.fits")]
-    example_list = CoordinatePatchCollection.find_stars_and_average(img_paths, 32, 100)
-
-    def generator():
-        yield fits.getdata(img_paths[0]).astype(float)
-    example_generator = CoordinatePatchCollection.find_stars_and_average(generator(), 32, 100)
 
     imgs_array = fits.getdata(img_paths[0]).astype(float)
     imgs_array = imgs_array.reshape((1, *imgs_array.shape))
-    example_ndarray = CoordinatePatchCollection.find_stars_and_average(imgs_array, 32, 100)
+
+    # Use a mask to only process part of the image, to speed up this test
+    mask = np.ones_like(imgs_array, dtype=bool)
+    mask[:, :800, :800] = 0
+
+    example_ndarray = CoordinatePatchCollection.find_stars_and_average(
+            imgs_array, 32, 100, star_mask=mask)
+
+    example_list = CoordinatePatchCollection.find_stars_and_average(
+            img_paths, 32, 100, star_mask=mask)
+
+    def generator():
+        yield imgs_array[0]
+    example_generator = CoordinatePatchCollection.find_stars_and_average(
+            generator(), 32, 100, star_mask=mask)
 
     # Check that we got the correct output type for each one
     for example in (example_list, example_generator, example_ndarray):
@@ -185,3 +194,49 @@ def generator():
         assert np.all(example_list.patches[loc] == example_generator.patches[loc])
         assert np.all(example_list.patches[loc] == example_ndarray.patches[loc])
 
+
+def test_find_stars_and_average_mask_formats(tmp_path):
+    """ Test star-finding masks, in all accepted formats"""
+    img_paths = [str(TEST_DIR / "data/DASH.fits")]
+    imgs_array = fits.getdata(img_paths[0])
+    # Cut down the data size so this test runs quicker
+    imgs_array = imgs_array[:800, :800].astype(float)
+    imgs_array = imgs_array.reshape((1, *imgs_array.shape))
+
+    # Mask out everything except one corner
+    mask_array = np.ones_like(imgs_array, dtype=bool)
+    mask_array[:, :402, :402] = 0
+
+    # Try all the formats in which masks are accepted
+
+    example_ndarray = CoordinatePatchCollection.find_stars_and_average(
+            imgs_array, 32, 200, star_mask=mask_array)
+
+    def generator():
+        yield mask_array[0]
+    example_generator = CoordinatePatchCollection.find_stars_and_average(
+            imgs_array, 32, 200, star_mask=generator())
+
+    mask_fname = str(tmp_path / "mask.fits")
+    fits.writeto(mask_fname, mask_array[0].astype(int))
+    example_file = CoordinatePatchCollection.find_stars_and_average(
+            imgs_array, 32, 200, star_mask=[mask_fname])
+
+    example_no_mask = CoordinatePatchCollection.find_stars_and_average(
+            imgs_array, 32, 200, star_mask=None)
+
+    for loc in example_file.patches.keys():
+        # Check that the patches are the same for all three mask formats
+        assert np.all(example_file.patches[loc] == example_generator.patches[loc])
+        assert np.all(example_file.patches[loc] == example_ndarray.patches[loc])
+
+        if loc.x >= 400 or loc.y >= 400:
+            # This patch should be fully masked
+            assert np.all(example_file[loc] == 0)
+        elif loc.x <= 200 and loc.y <= 200:
+            # This patch should be fully unmasked
+            assert np.all(example_file[loc] == example_no_mask[loc])
+        else:
+            # This patche covers a mix of masked and non-masked areas
+            assert np.any(example_file[loc] != 0)
+