From 1d85a45dfe12b83c67cbc29c3b070e242fef1e45 Mon Sep 17 00:00:00 2001
From: Jeremy Kubica <104161096+jeremykubica@users.noreply.github.com>
Date: Thu, 11 Jan 2024 08:19:23 -0500
Subject: [PATCH 1/3] Refactor masking

---
 src/kbmod/masking.py                         | 267 +++++--------------
 src/kbmod/run_search.py                      |  57 +---
 src/kbmod/search/image_stack.cpp             |  50 +---
 src/kbmod/search/image_stack.h               |  13 +-
 src/kbmod/search/layered_image.cpp           | 104 ++++++--
 src/kbmod/search/layered_image.h             |  27 +-
 src/kbmod/search/pydocs/image_stack_docs.h   |  60 +----
 src/kbmod/search/pydocs/layered_image_docs.h |  89 ++++++-
 src/kbmod/search/pydocs/raw_image_docs.h     |  13 +-
 src/kbmod/search/raw_image.cpp               |  26 --
 src/kbmod/search/raw_image.h                 |   3 -
 tests/test_image_stack.py                    |  81 +-----
 tests/test_layered_image.py                  |  91 +++++--
 tests/test_masking.py                        | 155 +----------
 tests/test_raw_image.py                      |  25 --
 tests/test_search.py                         |   8 +-
 tests/test_stamp_parity.py                   |   2 +-
 17 files changed, 371 insertions(+), 700 deletions(-)

diff --git a/src/kbmod/masking.py b/src/kbmod/masking.py
index 7a370037d..a25f10092 100644
--- a/src/kbmod/masking.py
+++ b/src/kbmod/masking.py
@@ -1,95 +1,11 @@
-"""Classes for performing masking on images from FITS files.
-
-ImageMasker provides an abstract base class that can be overridden to define masking
-algorithms for specific studies, instruments, or FITS headers. Specific masking classes
-are provided to support common masking operations including: masking based on a bit vector,
-masking based on a dictionary, masking based on a threshold, and growing a current mask.
+"""Functions for performing masking operations as specified in the configuration.
 """
 
-import abc
-
 import kbmod.search as kb
 
-
-def apply_mask_operations(stack, mask_list):
-    """Apply a series of masking operations defined by a list of
-    ImageMasker objects.
-
-    Parameters
-    ----------
-    stack : `kbmod.ImageStack`
-        The stack before the masks have been applied.
-    mask_list : `list`
-        A list of mask_list objects.
-
-    Returns
-    -------
-    stack : `kbmod.ImageStack`
-        The same stack object to allow chaining.
-    """
-    for mask in mask_list:
-        stack = mask.apply_mask(stack)
-    return stack
-
-
-class ImageMasker(abc.ABC):
-    """The base class for masking operations."""
-
-    def __init__(self, *args, **kwargs):
-        pass
-
-    @abc.abstractmethod
-    def apply_mask(self, stack):
-        """Apply the mask to an image stack.
-
-        Parameters
-        ----------
-        stack : `kbmod.ImageStack`
-            The stack before the masks have been applied.
-
-        Returns
-        -------
-        stack : `kbmod.ImageStack`
-            The same stack object to allow chaining.
-        """
-        pass
-
-
-class BitVectorMasker(ImageMasker):
-    """Apply a mask given a bit vector of masking flags to use
-    and vector of bit vectors to ignore.
-
-    Attributes
-    ----------
-    flags : `int`
-        A bit vector of masking flags to apply.
-    """
-
-    def __init__(self, flags, *args, **kwargs):
-        super().__init__(*args, **kwargs)
-        self.flags = flags
-
-    def apply_mask(self, stack):
-        """Apply the mask to an image stack.
-
-        Parameters
-        ----------
-        stack : `kbmod.ImageStack`
-            The stack before the masks have been applied.
-
-        Returns
-        -------
-        stack : `kbmod.ImageStack`
-            The same stack object to allow chaining.
-        """
-        if self.flags != 0:
-            stack.apply_mask_flags(self.flags)
-        return stack
-
-
-class DictionaryMasker(BitVectorMasker):
-    """Apply a mask given a dictionary of masking condition to key
-    and a list of masking conditions to use.
+def mask_flags_from_dict(mask_bits_dict, flag_keys):
+    """Generate a bitmask integer of flag keys from a dictionary
+    of masking reasons and a list of reasons to use.
 
     Attributes
     ----------
@@ -98,126 +14,71 @@ class DictionaryMasker(BitVectorMasker):
         number in the masking bit vector.
     flag_keys : `list`
         A list of masking keys to use.
-    """
-
-    def __init__(self, mask_bits_dict, flag_keys, *args, **kwargs):
-        self.mask_bits_dict = mask_bits_dict
-        self.flag_keys = flag_keys
-
-        # Convert the dictionary into a bit vector.
-        bitvector = 0
-        for bit in self.flag_keys:
-            bitvector += 2 ** self.mask_bits_dict[bit]
-
-        # Initialize the BitVectorMasker parameters.
-        super().__init__(bitvector, *args, **kwargs)
-
 
-class GlobalDictionaryMasker(ImageMasker):
-    """Apply a mask given a dictionary of masking condition to key
-    and a list of masking conditions to use. Masks pixels in every image
-    if they are masked in *multiple* images in the stack.
-
-    Attributes
-    ----------
-    mask_bits_dict : `dict`
-        A dictionary mapping a masking key (string) to the bit
-        number in the masking bit vector.
-    global_flag_keys : `list`
-        A list of masking keys to use.
-    mask_num_images : `int`
-        The number of images that need to be masked in the stack
-        to apply the mask to all images.
+    Returns
+    -------
+    bitmask : `int`
+        The bitmask to use for masking operations.
     """
+    bitmask = 0
+    for bit in flag_keys:
+        bitmask += 2 ** mask_bits_dict[bit]
+    return bitmask
 
-    def __init__(self, mask_bits_dict, global_flag_keys, mask_num_images, *args, **kwargs):
-        super().__init__(*args, **kwargs)
+            
+def apply_mask_operations(config, stack):
+    """Perform all the masking operations based on the search's configuration parameters.
 
-        self.mask_bits_dict = mask_bits_dict
-        self.global_flag_keys = global_flag_keys
-        self.mask_num_images = mask_num_images
-
-        # Convert the dictionary into a bit vector.
-        self.global_flags = 0
-        for bit in self.global_flag_keys:
-            self.global_flags += 2 ** self.mask_bits_dict[bit]
-
-    def apply_mask(self, stack):
-        """Apply the mask to an image stack.
-
-        Parameters
-        ----------
-        stack : `kbmod.ImageStack`
-            The stack before the masks have been applied.
-
-        Returns
-        -------
-        stack : `kbmod.ImageStack`
-            The same stack object to allow chaining.
-        """
-        if self.global_flags != 0:
-            stack.apply_global_mask(self.global_flags, self.mask_num_images)
-        return stack
-
-
-class ThresholdMask(ImageMasker):
-    """Mask pixels over a given value.
-
-    Attributes
+    Parameters
     ----------
-    mask_threshold : `float`
-        The flux threshold for a pixel.
-    """
-
-    def __init__(self, mask_threshold, *args, **kwargs):
-        super().__init__(*args, **kwargs)
-        self.mask_threshold = mask_threshold
-
-    def apply_mask(self, stack):
-        """Apply the mask to an image stack.
-
-        Parameters
-        ----------
-        stack : `kbmod.ImageStack`
-            The stack before the masks have been applied.
+    config : `SearchConfiguration`
+        The configuration parameters
+    stack : `ImageStack`
+        The stack before the masks have been applied. Modified in-place.
 
-        Returns
-        -------
-        stack : `kbmod.ImageStack`
-            The same stack object to allow chaining.
-        """
-        stack.apply_mask_threshold(self.mask_threshold)
-        return stack
-
-
-class GrowMask(ImageMasker):
-    """Apply a mask that grows the current max out a given number of pixels.
-
-    Attributes
-    ----------
-    num_pixels : `int`
-        The number of pixels to extend the mask.
+    Returns
+    -------
+    stack : `ImageStack`
+        The stack after the masks have been applied.
     """
+    # Generate the global mask before we start modifying the individual masks.
+    if config["repeated_flag_keys"] and len(config["repeated_flag_keys"]) > 0:
+        global_flags = mask_flags_from_dict(config["mask_bits_dict"], config["repeated_flag_keys"])
+        global_binary_mask = stack.make_global_mask(global_flags, config["mask_num_images"])
+    else:
+        global_binary_mask = None
+
+    # Start by creating a binary mask out of the primary flag values. Prioritize
+    # the config's mask_bit_vector over the dictionary based version.
+    if config["mask_bit_vector"]:
+        mask_flags = config["mask_bit_vector"]
+    elif config["flag_keys"] and len(config["flag_keys"]) > 0:
+        mask_flags = mask_flags_from_dict(config["mask_bits_dict"], config["flag_keys"])
+    else:
+        mask_flags = 0
+
+    # Apply the primary mask.
+    for i in range(stack.img_count()):
+        stack.get_single_image(i).binarize_mask(mask_flags)
+
+    # If the threshold is set, mask those pixels.
+    if config["mask_threshold"]:
+        for i in range(stack.img_count()):
+            stack.get_single_image(i).union_threshold_masking(config["mask_threshold"])
+
+    # Union in the global masking if there was one.
+    if global_binary_mask is not None:
+        for i in range(stack.img_count()):
+            stack.get_single_image(i).union_masks(global_binary_mask)
+
+    # Grow the masks.
+    if config["mask_grow"] and config["mask_grow"] > 0:
+        for i in range(stack.img_count()):
+            stack.get_single_image(i).grow_mask(config["mask_grow"])
+
+    # Apply the masks to the images.
+    for i in range(stack.img_count()):
+        stack.get_single_image(i).apply_mask(0xFFFFFF)
+            
+    return stack
 
-    def __init__(self, num_pixels, *args, **kwargs):
-        super().__init__(*args, **kwargs)
-
-        if num_pixels <= 0:
-            raise ValueError(f"Invalid num_pixels={num_pixels} for GrowMask")
-        self.num_pixels = num_pixels
-
-    def apply_mask(self, stack):
-        """Apply the mask to an image stack.
-
-        Parameters
-        ----------
-        stack : `kbmod.ImageStack`
-            The stack before the masks have been applied.
-
-        Returns
-        -------
-        stack : `kbmod.ImageStack`
-            The same stack object to allow chaining.
-        """
-        stack.grow_mask(self.num_pixels)
-        return stack
diff --git a/src/kbmod/run_search.py b/src/kbmod/run_search.py
index 42d6cff27..961f2f99c 100644
--- a/src/kbmod/run_search.py
+++ b/src/kbmod/run_search.py
@@ -15,14 +15,7 @@
 from .analysis_utils import PostProcess
 from .data_interface import load_input_from_config
 from .configuration import SearchConfiguration
-from .masking import (
-    BitVectorMasker,
-    DictionaryMasker,
-    GlobalDictionaryMasker,
-    GrowMask,
-    ThresholdMask,
-    apply_mask_operations,
-)
+from .masking import apply_mask_operations
 from .result_list import *
 from .filters.sigma_g_filter import SigmaGClipping
 from .work_unit import WorkUnit
@@ -34,52 +27,6 @@ class SearchRunner:
     def __init__(self):
         pass
 
-    def do_masking(self, config, stack):
-        """Perform the masking based on the search's configuration parameters.
-
-        Parameters
-        ----------
-        config : `SearchConfiguration`
-            The configuration parameters
-        stack : `ImageStack`
-            The stack before the masks have been applied. Modified in-place.
-
-        Returns
-        -------
-        stack : `ImageStack`
-            The stack after the masks have been applied.
-        """
-        mask_steps = []
-
-        # Prioritize the mask_bit_vector over the dictionary based version.
-        if config["mask_bit_vector"]:
-            mask_steps.append(BitVectorMasker(config["mask_bit_vector"]))
-        elif config["flag_keys"] and len(config["flag_keys"]) > 0:
-            mask_steps.append(DictionaryMasker(config["mask_bits_dict"], config["flag_keys"]))
-
-        # Add the threshold mask if it is set.
-        if config["mask_threshold"]:
-            mask_steps.append(ThresholdMask(config["mask_threshold"]))
-
-        # Add the global masking if it is set.
-        if config["repeated_flag_keys"] and len(config["repeated_flag_keys"]) > 0:
-            mask_steps.append(
-                GlobalDictionaryMasker(
-                    config["mask_bits_dict"],
-                    config["repeated_flag_keys"],
-                    config["mask_num_images"],
-                )
-            )
-
-        # Grow the mask.
-        if config["mask_grow"] and config["mask_grow"] > 0:
-            mask_steps.append(GrowMask(config["mask_grow"]))
-
-        # Apply the masks.
-        stack = apply_mask_operations(stack, mask_steps)
-
-        return stack
-
     def get_angle_limits(self, config):
         """Compute the angle limits based on the configuration information.
 
@@ -199,7 +146,7 @@ def run_search(self, config, stack):
 
         # Apply the mask to the images.
         if config["do_mask"]:
-            stack = self.do_masking(config, stack)
+            stack = apply_mask_operations(config, stack)
 
         # Perform the actual search.
         search = kb.StackSearch(stack)
diff --git a/src/kbmod/search/image_stack.cpp b/src/kbmod/search/image_stack.cpp
index afb313ce0..c8cbfe0a8 100644
--- a/src/kbmod/search/image_stack.cpp
+++ b/src/kbmod/search/image_stack.cpp
@@ -5,17 +5,11 @@ ImageStack::ImageStack(const std::vector<std::string>& filenames, const std::vec
     verbose = true;
     images = std::vector<LayeredImage>();
     load_images(filenames, psfs);
-
-    global_mask = RawImage(get_width(), get_height());
-    global_mask.set_all(0.0);
 }
 
 ImageStack::ImageStack(const std::vector<LayeredImage>& imgs) {
     verbose = true;
     images = imgs;
-
-    global_mask = RawImage(get_width(), get_height());
-    global_mask.set_all(0.0);
 }
 
 void ImageStack::load_images(const std::vector<std::string>& filenames, const std::vector<PSF>& psfs) {
@@ -65,43 +59,23 @@ void ImageStack::convolve_psf() {
     for (auto& i : images) i.convolve_psf();
 }
 
-void ImageStack::save_global_mask(const std::string& path) { global_mask.to_fits(path); }
-
 void ImageStack::save_images(const std::string& path) {
     for (auto& i : images) i.save_layers(path);
 }
 
-const RawImage& ImageStack::get_global_mask() const { return global_mask; }
-
-void ImageStack::apply_mask_flags(int flags) {
-    for (auto& i : images) {
-        i.apply_mask_flags(flags);
-    }
-}
-
-void ImageStack::apply_global_mask(int flags, int threshold) {
-    create_global_mask(flags, threshold);
-    for (auto& i : images) {
-        i.apply_global_mask(global_mask);
-    }
-}
-
-void ImageStack::apply_mask_threshold(float thresh) {
-    for (auto& i : images) i.apply_mask_threshold(thresh);
-}
-
-void ImageStack::grow_mask(int steps) {
-    for (auto& i : images) i.grow_mask(steps);
-}
-
-void ImageStack::create_global_mask(int flags, int threshold) {
+RawImage ImageStack::make_global_mask(int flags, int threshold) {
     int npixels = get_npixels();
 
+    // Start with an empty global mask.
+    RawImage global_mask = RawImage(get_width(), get_height());
+    global_mask.set_all(0.0);
+
     // For each pixel count the number of images where it is masked.
     std::vector<int> counts(npixels, 0);
     for (unsigned int img = 0; img < images.size(); ++img) {
         auto imgMask = images[img].get_mask().get_image().reshaped();
-        // Count the number of times a pixel has any of the flags
+
+        // Count the number of times a pixel has any of the given flags
         for (unsigned int pixel = 0; pixel < npixels; ++pixel) {
             if ((flags & static_cast<int>(imgMask[pixel])) != 0) counts[pixel]++;
         }
@@ -112,6 +86,8 @@ void ImageStack::create_global_mask(int flags, int threshold) {
     for (unsigned int p = 0; p < npixels; ++p) {
         global_m[p] = counts[p] < threshold ? 0.0 : 1.0;
     }
+
+    return global_mask;
 }
 
 #ifdef Py_PYTHON_H
@@ -130,14 +106,8 @@ static void image_stack_bindings(py::module& m) {
             .def("get_zeroed_time", &is::get_zeroed_time, pydocs::DOC_ImageStack_get_zeroed_time)
             .def("build_zeroed_times", &is::build_zeroed_times, pydocs::DOC_ImageStack_build_zeroed_times)
             .def("img_count", &is::img_count, pydocs::DOC_ImageStack_img_count)
-            .def("apply_mask_flags", &is::apply_mask_flags, pydocs::DOC_ImageStack_apply_mask_flags)
-            .def("apply_mask_threshold", &is::apply_mask_threshold,
-                 pydocs::DOC_ImageStack_apply_mask_threshold)
-            .def("apply_global_mask", &is::apply_global_mask, pydocs::DOC_ImageStack_apply_global_mask)
-            .def("grow_mask", &is::grow_mask, pydocs::DOC_ImageStack_grow_mask)
-            .def("save_global_mask", &is::save_global_mask, pydocs::DOC_ImageStack_save_global_mask)
+            .def("make_global_mask", &is::make_global_mask, pydocs::DOC_ImageStack_make_global_mask)
             .def("save_images", &is::save_images, pydocs::DOC_ImageStack_save_images)
-            .def("get_global_mask", &is::get_global_mask, pydocs::DOC_ImageStack_get_global_mask)
             .def("convolve_psf", &is::convolve_psf, pydocs::DOC_ImageStack_convolve_psf)
             .def("get_width", &is::get_width, pydocs::DOC_ImageStack_get_width)
             .def("get_height", &is::get_height, pydocs::DOC_ImageStack_get_height)
diff --git a/src/kbmod/search/image_stack.h b/src/kbmod/search/image_stack.h
index 4cb0b2fd8..a89d54359 100644
--- a/src/kbmod/search/image_stack.h
+++ b/src/kbmod/search/image_stack.h
@@ -29,26 +29,19 @@ class ImageStack {
     float get_zeroed_time(int index) const;
     std::vector<float> build_zeroed_times() const;  // Linear cost.
 
-    // Apply makes to all the images.
-    void apply_global_mask(int flags, int threshold);
-    void apply_mask_flags(int flags);
-    void apply_mask_threshold(float thresh);
-    void grow_mask(int steps);
-    const RawImage& get_global_mask() const;
-
     void convolve_psf();
 
+    // Make and return a global mask.
+    RawImage make_global_mask(int flags, int threshold);
+
     // Save data to files.
-    void save_global_mask(const std::string& path);
     void save_images(const std::string& path);
 
     virtual ~ImageStack(){};
 
 private:
     void load_images(const std::vector<std::string>& filenames, const std::vector<PSF>& psfs);
-    void create_global_mask(int flags, int threshold);
     std::vector<LayeredImage> images;
-    RawImage global_mask;
     bool verbose;
 };
 
diff --git a/src/kbmod/search/layered_image.cpp b/src/kbmod/search/layered_image.cpp
index 54066b65c..efe335fdd 100644
--- a/src/kbmod/search/layered_image.cpp
+++ b/src/kbmod/search/layered_image.cpp
@@ -67,11 +67,6 @@ LayeredImage::LayeredImage(std::string name, unsigned w, unsigned h, float noise
 
 void LayeredImage::set_psf(const PSF& new_psf) { psf = new_psf; }
 
-void LayeredImage::grow_mask(int steps) {
-    science.grow_mask(steps);
-    variance.grow_mask(steps);
-}
-
 void LayeredImage::convolve_given_psf(const PSF& given_psf) {
     science.convolve(given_psf);
 
@@ -83,25 +78,81 @@ void LayeredImage::convolve_given_psf(const PSF& given_psf) {
 
 void LayeredImage::convolve_psf() { convolve_given_psf(psf); }
 
-void LayeredImage::apply_mask_flags(int flags) {
+void LayeredImage::binarize_mask(int flags_to_use) {
+    const int num_pixels = get_npixels();
+    float* mask_pixels = mask.data();
+
+    for (int i = 0; i < num_pixels; ++i) {
+        int current_flags = static_cast<int>(mask_pixels[i]);
+
+        // Use a bitwise AND to only keep flags that are set in the current pixel
+        // and in the flags_to_use bitmask.
+        mask_pixels[i] = (flags_to_use & current_flags) > 0 ? 1 : 0;
+    }
+}
+
+void LayeredImage::apply_mask(int flags) {
     science.apply_mask(flags, mask);
     variance.apply_mask(flags, mask);
 }
 
-/* Mask all pixels that are not 0 in global mask */
-void LayeredImage::apply_global_mask(const RawImage& global_mask) {
-    science.apply_mask(0xFFFFFF, global_mask);
-    variance.apply_mask(0xFFFFFF, global_mask);
+void LayeredImage::union_masks(RawImage& new_mask) {
+    const int num_pixels = get_npixels();
+    if (num_pixels != new_mask.get_npixels()) {
+        throw std::runtime_error("Mismatched number of pixels between image and global mask.");
+    }
+
+    float* mask_pixels = mask.data();
+    float* new_pixels = new_mask.data();
+    for (int i = 0; i < num_pixels; ++i) {
+        int current_flags = static_cast<int>(mask_pixels[i]);
+        int new_flags = static_cast<int>(new_pixels[i]);
+
+        // Use a bitwise OR to keep flags set in the current pixel or the new mask.
+        mask_pixels[i] = current_flags | new_flags;
+    }
 }
 
-void LayeredImage::apply_mask_threshold(float thresh) {
+void LayeredImage::union_threshold_masking(float thresh) {
     const int num_pixels = get_npixels();
     float* sci_pixels = science.data();
-    float* var_pix = variance.data();
+    float* mask_pixels = mask.data();
+
     for (int i = 0; i < num_pixels; ++i) {
         if (sci_pixels[i] > thresh) {
-            sci_pixels[i] = NO_DATA;
-            var_pix[i] = NO_DATA;
+            // Use a logical OR to preserve all other flags.
+            mask_pixels[i] = static_cast<int>(mask_pixels[i]) | 1;
+        }
+    }
+}
+
+/* This implementation of grow_mask is optimized for steps > 1
+   (which is how the code is generally used. If you are only
+   growing the mask by 1, the extra copy will be a little slower.
+*/
+void LayeredImage::grow_mask(int steps) {
+    ImageI bitmask = ImageI::Constant(height, width, -1);
+    bitmask = (mask.get_image().array() > 0).select(0, bitmask);
+
+    for (int itr = 1; itr <= steps; ++itr) {
+        for (int j = 0; j < height; ++j) {
+            for (int i = 0; i < width; ++i) {
+                if (bitmask(j, i) == -1) {
+                    if (((j - 1 >= 0) && (bitmask(j - 1, i) == itr - 1)) ||
+                        ((i - 1 >= 0) && (bitmask(j, i - 1) == itr - 1)) ||
+                        ((j + 1 < height) && (bitmask(j + 1, i) == itr - 1)) ||
+                        ((i + 1 < width) && (bitmask(j, i + 1) == itr - 1))) {
+                        bitmask(j, i) = itr;
+                    }
+                }
+            }  // for i
+        }      // for j
+    }          // for step
+
+    // Overwrite the mask with the expanded one.
+    for (int j = 0; j < height; ++j) {
+        for (int i = 0; i < width; ++i) {
+            mask.set_pixel({j, i}, (bitmask(j, i) == -1) ? 0 : 1);
         }
     }
 }
@@ -226,12 +277,31 @@ static void layered_image_bindings(py::module& m) {
             .def(py::init<const ri&, const ri&, const ri&, pf&>())
             .def(py::init<std::string, int, int, double, float, float, pf&>())
             .def(py::init<std::string, int, int, double, float, float, pf&, int>())
+            .def("contains", &li::contains, pydocs::DOC_LayeredImage_cointains)
+            .def("get_science_pixel", &li::get_science_pixel, pydocs::DOC_LayeredImage_get_science_pixel)
+            .def("get_variance_pixel", &li::get_variance_pixel, pydocs::DOC_LayeredImage_get_variance_pixel)
+            .def(
+                    "contains",
+                    [](li& cls, int i, int j) {
+                        return cls.contains({i, j});
+                    })
+            .def(
+                    "get_science_pixel",
+                    [](li& cls, int i, int j) {
+                        return cls.get_science_pixel({i, j});
+                    })
+            .def(
+                    "get_variance_pixel",
+                    [](li& cls, int i, int j) {
+                        return cls.get_variance_pixel({i, j});
+                    })
             .def("set_psf", &li::set_psf, pydocs::DOC_LayeredImage_set_psf)
             .def("get_psf", &li::get_psf, py::return_value_policy::reference_internal,
                  pydocs::DOC_LayeredImage_get_psf)
-            .def("apply_mask_flags", &li::apply_mask_flags, pydocs::DOC_LayeredImage_apply_mask_flags)
-            .def("apply_mask_threshold", &li::apply_mask_threshold,
-                 pydocs::DOC_LayeredImage_apply_mask_threshold)
+            .def("binarize_mask", &li::binarize_mask, pydocs::DOC_LayeredImage_binarize_mask)
+            .def("apply_mask", &li::apply_mask, pydocs::DOC_LayeredImage_apply_mask)
+            .def("union_masks", &li::union_masks, pydocs::DOC_LayeredImage_union_masks)
+            .def("union_threshold_masking", &li::union_threshold_masking, pydocs::DOC_LayeredImage_union_threshold_masking)
             .def("sub_template", &li::subtract_template, pydocs::DOC_LayeredImage_sub_template)
             .def("save_layers", &li::save_layers, pydocs::DOC_LayeredImage_save_layers)
             .def("get_science", &li::get_science, py::return_value_policy::reference_internal,
diff --git a/src/kbmod/search/layered_image.h b/src/kbmod/search/layered_image.h
index 771f2fbe8..d3527fe31 100644
--- a/src/kbmod/search/layered_image.h
+++ b/src/kbmod/search/layered_image.h
@@ -39,11 +39,28 @@ class LayeredImage {
     RawImage& get_mask() { return mask; }
     RawImage& get_variance() { return variance; }
 
-    // Applies the mask functions to each of the science and variance layers.
-    void apply_mask_flags(int flag);
-    void apply_global_mask(const RawImage& global_mask);
-    void apply_mask_threshold(float thresh);
+    // Getter functions for the pixels of the science and variance layers that check
+    // the mask layer for any set bits.
+    inline float get_science_pixel(const Index& idx) const {
+        return contains(idx) ? (mask.get_pixel(idx) == 0 ? science.get_pixel(idx) : NO_DATA) : NO_DATA;
+    }
+
+    inline float get_variance_pixel(const Index& idx) const {
+        return contains(idx) ? 
+            (mask.get_pixel(idx) == 0 ? variance.get_pixel(idx) : NO_DATA) :
+             NO_DATA;
+    }
+
+    inline bool contains(const Index& idx) const {
+        return idx.i >= 0 && idx.i < height && idx.j >= 0 && idx.j < width;
+    }
+
+    // Masking functions.
+    void binarize_mask(int flags_to_keep);
+    void union_masks(RawImage& new_mask);
+    void union_threshold_masking(float thresh);
     void grow_mask(int steps);
+    void apply_mask(int flags);
 
     // Subtracts a template image from the science layer.
     void subtract_template(RawImage& sub_template);
@@ -76,7 +93,7 @@ class LayeredImage {
     PSF psf;
     RawImage science;
     RawImage mask;
-    RawImage variance;
+    RawImage variance;   
 };
 
 } /* namespace search */
diff --git a/src/kbmod/search/pydocs/image_stack_docs.h b/src/kbmod/search/pydocs/image_stack_docs.h
index 4f0c9f048..26c38f49a 100644
--- a/src/kbmod/search/pydocs/image_stack_docs.h
+++ b/src/kbmod/search/pydocs/image_stack_docs.h
@@ -32,57 +32,25 @@ static const auto DOC_ImageStack_build_zeroed_times = R"doc(
   in the stack and the first image.
   ")doc";
 
-static const auto DOC_ImageStack_apply_mask_flags = R"doc(
-  Applies a mask to each image by comparing the given bit vector with the
-  values in the mask layer and marking pixels NO_DATA. Modifies the image in-place.
+static const auto DOC_ImageStack_make_global_mask = R"doc(
+  Create a new global mask that from a set of flags and a threshold.
+  The global mask marks a pixel as masked if and only if it is masked
+  by one of the given flags in at least ``threshold`` individual images.
+  The returned mask is binary.
 
   Parameters
   ----------
-  flag : `int`
-      The bit mask of mask flags to use.
-  )doc";
-
-static const auto DOC_ImageStack_apply_mask_threshold = R"doc(
-  Applies a threshold mask to each image by setting pixel values over
-  a given threshold to NO_DATA. Modifies the images in-place.
-
-  Parameters
-  ----------
-  thresh : `float`
-      The threshold value to use.
-  )doc";
-
-static const auto DOC_ImageStack_apply_global_mask = R"doc(
-  Createas a global mask an applies it to each image. A global mask
-  masks a pixel if and only if that pixel is masked in at least ``threshold``
-  individual images.  Modifies the images in-place and creates the global mask.
-
-  Parameters
-  ----------
-  flag : `int`
-      The bit mask of mask flags to use.
+  flags : `int`
+      A bit mask of mask flags to use when counting.
   threshold : `int`
       The minimum number of images in which a pixel must be masked to be
       part of the global mask.
-  )doc";
 
-static const auto DOC_ImageStack_grow_mask = R"doc(
-  Expands the NO_DATA tags to nearby pixels for all images.
-  Modifies the images in-place.
-
-  Parameters
-  ----------
-  steps : `int`
-     The number of pixels by which to grow the masked regions.
-  )doc";
-
-static const auto DOC_ImageStack_save_global_mask = R"doc(
-  Saves the global mask created by apply_global_mask to a FITS file.
-
-  Parameters
-  ----------
-  path : `str`
-      The directory in which to store the global mask file.
+  Returns
+  -------
+  global_mask : `RawImage`
+      A RawImage containing the global mask with 1 for masked pixels
+      and 0 for unmasked pixels.
   )doc";
 
 static const auto DOC_ImageStack_save_images = R"doc(
@@ -97,10 +65,6 @@ static const auto DOC_ImageStack_save_images = R"doc(
       The file path to use. 
   )doc";
 
-static const auto DOC_ImageStack_get_global_mask = R"doc(
-  Returns a reference to the global mask created by apply_global_mask.
-  )doc";
-
 static const auto DOC_ImageStack_convolve_psf = R"doc(
   Convolves each image (science and variance layers) with the PSF
   stored in the LayeredImage object.
diff --git a/src/kbmod/search/pydocs/layered_image_docs.h b/src/kbmod/search/pydocs/layered_image_docs.h
index 0ffeb5b95..cd54decd4 100644
--- a/src/kbmod/search/pydocs/layered_image_docs.h
+++ b/src/kbmod/search/pydocs/layered_image_docs.h
@@ -55,20 +55,41 @@ static const auto DOC_LayeredImage_get_psf = R"doc(
   Returns the PSF object.
   )doc";
 
-static const auto DOC_LayeredImage_apply_mask_flags = R"doc(
-  Applies a mask to each image by comparing the given bit vector with the
-  values in the mask layer and marking pixels NO_DATA. 
-  Modifies the science and variance layers in-place.
+static const auto DOC_LayeredImage_binarize_mask = R"doc(
+  Convert the bitmask of flags into a single binary value of 1
+  for pixels that match one of the flags to use and 0 otherwise.
+  Modifies the mask layer in-place. Used to select which masking
+  flags are applied.
+
+  Note: This is a no-op for masks that are already binary and it is
+  safe to call this function multiple times.
 
   Parameters
   ----------
-  flag : `int`
-      The bit mask of mask flags to use.
+  flags_to_use : `int`
+      The bit mask of mask flags to keep.
+  )doc";
+
+static const auto DOC_LayeredImage_apply_mask = R"doc(
+  Applies the mask layer to each of the science and variance layers
+  by checking whether the pixel in the mask layer is 0 (no masking)
+  or non-zero (masked). Applies all flags. To use a subset of flags
+  call binarize_mask() first.
   )doc";
 
-static const auto DOC_LayeredImage_apply_mask_threshold = R"doc(
-  Applies a threshold mask by setting pixel values over a given threshold
-  to NO_DATA. Modifies the science and variance layers in-place.
+static const auto DOC_LayeredImage_union_masks = R"doc(
+  Unions the masked pixel flags from the a given second mask layer onto
+  this image's mask layer. Modifies the mask layer in place.
+
+  Parameters
+  ----------
+  global_mask : `RawImage`
+      The `RawImage` of global mask values (binary) for each pixel.
+  )doc";
+
+static const auto DOC_LayeredImage_union_threshold_masking = R"doc(
+  Masks any pixel whose corresponding value in the science layer is
+  above the given threshold using mask flag = 1.
 
   Parameters
   ----------
@@ -166,6 +187,56 @@ static const auto DOC_LayeredImage_set_obstime = R"doc(
   Set the image's observation time.
   )doc";
 
+static const auto DOC_LayeredImage_cointains = R"doc(
+  Returns a Boolean indicating whether the image contains the given coordinates.
+
+  Parameters
+  ----------
+  i : `int`
+      Row index.
+  j : `int`
+      Col index.
+
+  Returns
+  -------
+  result : `bool`
+      A Boolean indicating whether the image contains the given coordinates.
+  )doc";
+
+static const auto DOC_LayeredImage_get_science_pixel = R"doc(
+  Get the science pixel value at given index, checking the mask layer.
+  Returns NO_DATA if any of the mask bits are set.
+
+  Parameters
+  ----------
+  i : `int`
+      Row index.
+  j : `int`
+      Col index.
+
+  Returns
+  -------
+  value : `float`
+      Pixel value.
+  )doc";
+
+static const auto DOC_LayeredImage_get_variance_pixel = R"doc(
+  Get the variance pixel value at given index, checking the mask layer.
+  Returns NO_DATA if any of the mask bits are set.
+
+  Parameters
+  ----------
+  i : `int`
+      Row index.
+  j : `int`
+      Col index.
+
+  Returns
+  -------
+  value : `float`
+      Pixel value.
+  )doc";
+
 static const auto DOC_LayeredImage_generate_psi_image = R"doc(
   todo
   )doc";
diff --git a/src/kbmod/search/pydocs/raw_image_docs.h b/src/kbmod/search/pydocs/raw_image_docs.h
index 39bcd6d8d..7fd2e6c88 100644
--- a/src/kbmod/search/pydocs/raw_image_docs.h
+++ b/src/kbmod/search/pydocs/raw_image_docs.h
@@ -87,7 +87,7 @@ static const auto DOC_RawImage_get_pixel = R"doc(
   Returns
   -------
   value : `float`
-      Pixel value/
+      Pixel value.
   )doc";
 
 static const auto DOC_RawImage_pixel_has_data = R"doc(
@@ -272,20 +272,11 @@ static const auto DOC_RawImage_apply_mask = R"doc(
   Parameters
   ----------
   flag : `int`
-      The bit mask of mask flags to use.
+      The bit mask of mask flags to use. Use 0xFFFFFF to apply all flags.
   mask : `RawImage`
       The image of pixel mask values.
   )doc";
 
-static const auto DOC_RawImage_grow_mask = R"doc(
-  Expands the NO_DATA tags to nearby pixels. Modifies the image in-place.
-
-  Parameters
-  ----------
-  steps : `int`
-     The number of pixels by which to grow the masked regions.
-  )doc";
-
 static const auto DOC_RawImage_convolve_gpu = R"doc(
   Convolve the image with a PSF on the GPU.
 
diff --git a/src/kbmod/search/raw_image.cpp b/src/kbmod/search/raw_image.cpp
index a70c34984..1dcb13214 100644
--- a/src/kbmod/search/raw_image.cpp
+++ b/src/kbmod/search/raw_image.cpp
@@ -227,31 +227,6 @@ void RawImage::apply_mask(int flags, const RawImage& mask) {
     }      // for j
 }
 
-/* This implementation of grow_mask is optimized for steps > 1
-   (which is how the code is generally used. If you are only
-   growing the mask by 1, the extra copy will be a little slower.
-*/
-void RawImage::grow_mask(const int steps) {
-    ImageI bitmask = ImageI::Constant(height, width, -1);
-    bitmask = (image.array() == NO_DATA).select(0, bitmask);
-
-    for (int itr = 1; itr <= steps; ++itr) {
-        for (int j = 0; j < height; ++j) {
-            for (int i = 0; i < width; ++i) {
-                if (bitmask(j, i) == -1) {
-                    if (((j - 1 >= 0) && (bitmask(j - 1, i) == itr - 1)) ||
-                        ((i - 1 >= 0) && (bitmask(j, i - 1) == itr - 1)) ||
-                        ((j + 1 < height) && (bitmask(j + 1, i) == itr - 1)) ||
-                        ((i + 1 < width) && (bitmask(j, i + 1) == itr - 1))) {
-                        bitmask(j, i) = itr;
-                    }
-                }
-            }  // for i
-        }      // for j
-    }          // for step
-    image = (bitmask.array() > -1).select(NO_DATA, image);
-}
-
 void RawImage::set_all(float value) { image.setConstant(value); }
 
 // The maximum value of the image and return the coordinates.
@@ -639,7 +614,6 @@ static void raw_image_bindings(py::module& m) {
                     },
                     pydocs::DOC_RawImage_get_interp_neighbors_and_weights)
             .def("apply_mask", &rie::apply_mask, pydocs::DOC_RawImage_apply_mask)
-            .def("grow_mask", &rie::grow_mask, pydocs::DOC_RawImage_grow_mask)
             .def("convolve_gpu", &rie::convolve, pydocs::DOC_RawImage_convolve_gpu)
             .def("convolve_cpu", &rie::convolve_cpu, pydocs::DOC_RawImage_convolve_cpu)
             .def("save_fits", &rie::to_fits, pydocs::DOC_RawImage_save_fits)
diff --git a/src/kbmod/search/raw_image.h b/src/kbmod/search/raw_image.h
index bacdfc005..43cec9d53 100644
--- a/src/kbmod/search/raw_image.h
+++ b/src/kbmod/search/raw_image.h
@@ -97,9 +97,6 @@ class RawImage {
     // to apply (use 0xFFFFFF to apply all flags).
     void apply_mask(int flags, const RawImage& mask);
 
-    // Grow the area of masked array.
-    void grow_mask(int steps);
-
     // The maximum value of the image and return the coordinates. The parameter
     // furthest_from_center indicates whether to break ties using the peak further
     // or closer to the center of the image.
diff --git a/tests/test_image_stack.py b/tests/test_image_stack.py
index b9fda14f9..314727c53 100644
--- a/tests/test_image_stack.py
+++ b/tests/test_image_stack.py
@@ -57,87 +57,22 @@ def test_times(self):
         for i in range(self.num_images):
             self.assertEqual(times[i], 2.0 * i)
 
-    def test_apply_mask(self):
-        # Nothing is initially masked.
+    def test_make_global_mask(self):
+        # Mask a single point in each image with flag=2.
         for i in range(self.num_images):
-            sci = self.im_stack.get_single_image(i).get_science()
-            for y in range(self.im_stack.get_height()):
-                for x in range(self.im_stack.get_width()):
-                    self.assertTrue(sci.pixel_has_data(y, x))
-
-        self.im_stack.apply_mask_flags(1)
-
-        # Check that one pixel is masked in each time.
-        for i in range(self.num_images):
-            sci = self.im_stack.get_single_image(i).get_science()
-            for y in range(self.im_stack.get_height()):
-                for x in range(self.im_stack.get_width()):
-                    if y == 10 and x == 10 + i:
-                        self.assertFalse(sci.pixel_has_data(y, x))
-                    else:
-                        self.assertTrue(sci.pixel_has_data(y, x))
-
-    def test_create_global_mask(self):
-        global_mask = self.im_stack.get_global_mask()
-        for y in range(self.im_stack.get_height()):
-            for x in range(self.im_stack.get_width()):
-                self.assertEqual(global_mask.get_pixel(y, x), 0.0)
-
-        # Apply the global mask for flag=1 and a threshold of the bit set
-        # in at least one mask.
-        self.im_stack.apply_global_mask(1, 1)
-
-        # Check that the correct pixels are masked in each time.
-        for i in range(self.num_images):
-            sci = self.im_stack.get_single_image(i).get_science()
-            for y in range(self.im_stack.get_height()):
-                for x in range(self.im_stack.get_width()):
-                    if y == 10 and x >= 10 and x <= 10 + (self.num_images - 1):
-                        self.assertFalse(sci.pixel_has_data(y, x))
-                    else:
-                        self.assertTrue(sci.pixel_has_data(y, x))
-
-        # Check that the global mask is now set.
-        global_mask = self.im_stack.get_global_mask()
-        for y in range(self.im_stack.get_height()):
-            for x in range(self.im_stack.get_width()):
-                if y == 10 and x >= 10 and x <= 10 + (self.num_images - 1):
-                    self.assertEqual(global_mask.get_pixel(y, x), 1.0)
-                else:
-                    self.assertEqual(global_mask.get_pixel(y, x), 0.0)
-
-    def test_create_global_mask_reset(self):
-        global_mask = self.im_stack.get_global_mask()
-        for y in range(self.im_stack.get_height()):
-            for x in range(self.im_stack.get_width()):
-                self.assertEqual(global_mask.get_pixel(y, x), 0.0)
+            self.images[i].get_mask().set_pixel(5, 5, 2)
 
         # Apply the global mask for flag=1 and a threshold of the bit set
-        # in at least one mask.
-        self.im_stack.apply_global_mask(1, 1)
+        # in at least one mask. Note this ignores flag=2 for the count.
+        mask = self.im_stack.make_global_mask(1, 1)
 
-        # Check that the global mask is set.
-        global_mask = self.im_stack.get_global_mask()
+        # Check that the correct pixels are masked
         for y in range(self.im_stack.get_height()):
             for x in range(self.im_stack.get_width()):
                 if y == 10 and x >= 10 and x <= 10 + (self.num_images - 1):
-                    self.assertEqual(global_mask.get_pixel(y, x), 1.0)
+                    self.assertEqual(mask.get_pixel(y, x), 1)
                 else:
-                    self.assertEqual(global_mask.get_pixel(y, x), 0.0)
-
-        # Unmask the pixels.
-        for i in range(self.num_images):
-            img = self.im_stack.get_single_image(i)
-            mask = img.get_mask()
-            mask.set_pixel(10, 10 + i, 0)
-
-        # Reapply the mask and check that nothing is masked.
-        # Note the science pixels will still be masked from the previous application.
-        self.im_stack.apply_global_mask(1, 1)
-        global_mask = self.im_stack.get_global_mask()
-        for y in range(self.im_stack.get_height()):
-            for x in range(self.im_stack.get_width()):
-                self.assertEqual(global_mask.get_pixel(y, x), 0.0)
+                    self.assertEqual(mask.get_pixel(y, x), 0)
 
     # WOW, this is the first test that caught the fact that interpolated_add
     # called add, and that add had flipped i and j by accident. The first one.
diff --git a/tests/test_layered_image.py b/tests/test_layered_image.py
index 906dbd3bb..391f81d0e 100644
--- a/tests/test_layered_image.py
+++ b/tests/test_layered_image.py
@@ -55,6 +55,7 @@ def test_create_from_layers(self):
 
         mask = RawImage(30, 40)
         mask.set_all(0.0)
+        mask.set_pixel(10, 12, 1)
 
         # Create the layered image.
         img2 = LayeredImage(sci, var, mask, PSF(2.0))
@@ -63,13 +64,31 @@ def test_create_from_layers(self):
         self.assertEqual(img2.get_npixels(), 30.0 * 40.0)
         self.assertEqual(img2.get_obstime(), -1.0)  # No time given
 
+        # Test the bounds checking.
+        self.assertTrue(img2.contains(0, 0))
+        self.assertTrue(img2.contains(39, 29))
+        self.assertFalse(img2.contains(39, 30))
+        self.assertFalse(img2.contains(40, 15))
+        self.assertFalse(img2.contains(15, -1))
+        self.assertFalse(img2.contains(-1, 0))
+
         # Check the layers.
         science = img2.get_science()
         variance = img2.get_variance()
         mask2 = img2.get_mask()
         for y in range(img2.get_height()):
             for x in range(img2.get_width()):
-                self.assertEqual(mask2.get_pixel(y, x), 0)
+                if x == 12 and y == 10:
+                    # The masked pixel should have no data.
+                    self.assertEqual(mask2.get_pixel(y, x), 1)
+                    self.assertAlmostEqual(img2.get_science_pixel(y, x), KB_NO_DATA)
+                    self.assertAlmostEqual(img2.get_variance_pixel(y, x), KB_NO_DATA)
+                else:
+                    self.assertEqual(mask2.get_pixel(y, x), 0)
+                    self.assertAlmostEqual(img2.get_science_pixel(y, x), x + 40.0 * y)
+                    self.assertAlmostEqual(img2.get_variance_pixel(y, x), 1.0)
+
+                # The individual layers do not have the masking until it is applied.
                 self.assertEqual(variance.get_pixel(y, x), 1.0)
                 self.assertAlmostEqual(science.get_pixel(y, x), x + 40.0 * y)
 
@@ -126,14 +145,56 @@ def test_overwrite_PSF(self):
         science_pixel_psf2 = self.image.get_science().get_pixel(50, 50)
         self.assertLess(science_pixel_psf1, science_pixel_psf2)
 
-    def test_mask_threshold(self):
+    def test_binarize_mask(self):
+        # Mask out a range of pixels.
+        mask = self.image.get_mask()
+        for x in range(9):
+            mask.set_pixel(10, x, x)
+            
+        # Only keep the mask at for pixels with flags at
+        # bit positions 0 and 2 (1 + 4 = 5).
+        self.image.binarize_mask(5)
+        self.assertEqual(mask.get_pixel(10, 0), 0)
+        self.assertEqual(mask.get_pixel(10, 1), 1)
+        self.assertEqual(mask.get_pixel(10, 2), 0)
+        self.assertEqual(mask.get_pixel(10, 3), 1)
+        self.assertEqual(mask.get_pixel(10, 4), 1)
+        self.assertEqual(mask.get_pixel(10, 5), 1)
+        self.assertEqual(mask.get_pixel(10, 6), 1)
+        self.assertEqual(mask.get_pixel(10, 7), 1)
+        self.assertEqual(mask.get_pixel(10, 8), 0)
+
+    def test_union_masks(self):
+        # Mask out a range of pixels.
+        mask = self.image.get_mask()
+        mask.set_pixel(15, 12, 1)
+        mask.set_pixel(15, 13, 2)
+        mask.set_pixel(15, 14, 3)
+
+        mask2 = RawImage(mask.width, mask.height)
+        mask2.set_all(0.0)
+        mask2.set_pixel(15, 11, 1)
+        mask2.set_pixel(15, 13, 1)
+        mask2.set_pixel(15, 14, 1)
+        mask2.set_pixel(15, 15, 8)
+
+        self.image.union_masks(mask2)
+        self.assertEqual(mask.get_pixel(15, 10), 0)
+        self.assertEqual(mask.get_pixel(15, 11), 1)  # bit 1 added
+        self.assertEqual(mask.get_pixel(15, 12), 1)
+        self.assertEqual(mask.get_pixel(15, 13), 3)  # bit 1 added
+        self.assertEqual(mask.get_pixel(15, 14), 3)
+        self.assertEqual(mask.get_pixel(15, 15), 8)
+        self.assertEqual(mask.get_pixel(15, 16), 0)
+
+    def test_add_threshold_mask_flags(self):
         masked_pixels = {}
         threshold = 20.0
 
         # Add an object brighter than the threshold.
         add_fake_object(self.image, 50, 50, 500.0, self.p)
 
-        # Find all the pixels that should be masked.
+        # Manually find all the pixels that should be masked.
         science = self.image.get_science()
         for y in range(self.image.get_height()):
             for x in range(self.image.get_width()):
@@ -141,21 +202,21 @@ def test_mask_threshold(self):
                 if value > threshold:
                     index = self.image.get_width() * y + x
                     masked_pixels[index] = True
-
-        # Do the masking and confirm we have masked
-        # at least 1 pixel.
-        self.image.apply_mask_threshold(threshold)
         self.assertGreater(len(masked_pixels), 0)
 
+        # Reset the mask an perform threshold masking.
+        mask = self.image.get_mask()
+        mask.set_all(0.0)
+        self.image.union_threshold_masking(threshold)
+
         # Check that we masked the correct pixels.
-        science = self.image.get_science()
         for y in range(self.image.get_height()):
             for x in range(self.image.get_width()):
                 index = self.image.get_width() * y + x
                 if index in masked_pixels:
-                    self.assertFalse(science.pixel_has_data(y, x))
+                    self.assertEqual(mask.get_pixel(y, x), 1)
                 else:
-                    self.assertTrue(science.pixel_has_data(y, x))
+                    self.assertEqual(mask.get_pixel(y, x), 0)
 
     def test_apply_mask(self):
         # Nothing is initially masked.
@@ -171,7 +232,7 @@ def test_apply_mask(self):
         mask.set_pixel(10, 13, 3)
 
         # Apply the mask flags to only (10, 11) and (10, 13)
-        self.image.apply_mask_flags(1)
+        self.image.apply_mask(1)
 
         science = self.image.get_science()
         for y in range(self.image.get_height()):
@@ -186,11 +247,9 @@ def test_grow_mask(self):
         mask.set_pixel(11, 10, 1)
         mask.set_pixel(12, 10, 1)
         mask.set_pixel(13, 10, 1)
-        self.image.apply_mask_flags(1)
         self.image.grow_mask(1)
 
         # Check that the mask has grown to all adjacent pixels.
-        science = self.image.get_science()
         for y in range(self.image.get_height()):
             for x in range(self.image.get_width()):
                 should_mask = (
@@ -198,24 +257,22 @@ def test_grow_mask(self):
                     or (x == 9 and y <= 13 and y >= 11)
                     or (x == 11 and y <= 13 and y >= 11)
                 )
-                self.assertEqual(science.pixel_has_data(y, x), not should_mask)
+                self.assertEqual(mask.get_pixel(y, x) == 0, not should_mask)
 
     def test_grow_mask_mult(self):
         mask = self.image.get_mask()
         mask.set_pixel(11, 10, 1)
         mask.set_pixel(12, 10, 1)
-        self.image.apply_mask_flags(1)
         self.image.grow_mask(3)
 
         # Check that the mask has grown to all applicable pixels.
-        science = self.image.get_science()
         for y in range(self.image.get_height()):
             for x in range(self.image.get_width()):
                 # Check whether the point is a manhattan distance of <= 3 from
                 # one of the original masked pixels.
                 dx = abs(x - 10)
                 dy = min(abs(y - 11), abs(y - 12))
-                self.assertEqual(science.pixel_has_data(y, x), dx + dy > 3)
+                self.assertEqual(mask.get_pixel(y, x) == 0, dx + dy > 3)
 
     def test_psi_and_phi_image(self):
         p = PSF(0.00000001)  # A point function.
diff --git a/tests/test_masking.py b/tests/test_masking.py
index 5bdbad7fb..6dda673cf 100644
--- a/tests/test_masking.py
+++ b/tests/test_masking.py
@@ -1,159 +1,9 @@
 import unittest
 
 from kbmod.configuration import SearchConfiguration
-from kbmod.masking import (
-    BitVectorMasker,
-    DictionaryMasker,
-    GlobalDictionaryMasker,
-    GrowMask,
-    ThresholdMask,
-    apply_mask_operations,
-)
-from kbmod.run_search import SearchRunner
+from kbmod.masking import apply_mask_operations
 from kbmod.search import *
 
-
-class test_masking_classes(unittest.TestCase):
-    def setUp(self):
-        # The configuration parameters.
-        self.mask_bits_dict = {
-            "BAD": 0,
-            "CLIPPED": 9,
-            "CR": 3,
-            "CROSSTALK": 10,
-            "DETECTED": 5,
-            "DETECTED_NEGATIVE": 6,
-            "EDGE": 4,
-            "INEXACT_PSF": 11,
-            "INTRP": 2,
-            "NOT_DEBLENDED": 12,
-            "NO_DATA": 8,
-            "REJECTED": 13,
-            "SAT": 1,
-            "SENSOR_EDGE": 14,
-            "SUSPECT": 7,
-            "UNMASKEDNAN": 15,
-        }
-        self.default_flag_keys = ["BAD", "EDGE", "NO_DATA", "SUSPECT", "UNMASKEDNAN"]
-
-        # Create the a fake layered image.
-        self.img_count = 5
-        self.dim_x = 20
-        self.dim_y = 20
-        self.noise_level = 0.1
-        self.variance = self.noise_level**2
-        self.p = PSF(1.0)
-        self.imlist = []
-        self.time_list = []
-        for i in range(self.img_count):
-            time = i / self.img_count
-            self.time_list.append(time)
-            im = LayeredImage(
-                str(i), self.dim_x, self.dim_y, self.noise_level, self.variance, time, self.p, i
-            )
-            self.imlist.append(im)
-        self.stack = ImageStack(self.imlist)
-
-    def test_threshold_masker(self):
-        # Set one science pixel per image above the threshold
-        for i in range(self.img_count):
-            img = self.stack.get_single_image(i)
-            sci = img.get_science()
-            sci.set_pixel(8, 2 + i, 501.0)
-            sci.set_pixel(9, 1 + i, 499.0)
-
-        # With a threshold of 500 one pixel per image should be masked.
-        mask = ThresholdMask(500)
-        self.stack = mask.apply_mask(self.stack)
-        for i in range(self.img_count):
-            sci = self.stack.get_single_image(i).get_science()
-            for x in range(self.dim_x):
-                for y in range(self.dim_y):
-                    if x == 2 + i and y == 8:
-                        self.assertFalse(sci.pixel_has_data(y, x))
-                    else:
-                        self.assertTrue(sci.pixel_has_data(y, x))
-
-    def test_per_image_dictionary_mask(self):
-        # Set each mask pixel in a row to one masking reason.
-        for i in range(self.img_count):
-            img = self.stack.get_single_image(i)
-            msk = img.get_mask()
-            for x in range(self.dim_x):
-                msk.set_pixel(3, x, 2**x)
-
-        # Mask with two keys.
-        mask = DictionaryMasker(self.mask_bits_dict, ["BAD", "EDGE"])
-        self.stack = mask.apply_mask(self.stack)
-        for i in range(self.img_count):
-            sci = self.stack.get_single_image(i).get_science()
-            for x in range(self.dim_x):
-                for y in range(self.dim_y):
-                    if y == 3 and (x == 0 or x == 4):
-                        self.assertFalse(sci.pixel_has_data(y, x))
-                    else:
-                        self.assertTrue(sci.pixel_has_data(y, x))
-
-        # Mask with all the default keys.
-        mask = DictionaryMasker(self.mask_bits_dict, self.default_flag_keys)
-        self.stack = mask.apply_mask(self.stack)
-        for i in range(self.img_count):
-            sci = self.stack.get_single_image(i).get_science()
-            for x in range(self.dim_x):
-                for y in range(self.dim_y):
-                    if y == 3 and (x == 0 or x == 4 or x == 7 or x == 8 or x == 15):
-                        self.assertFalse(sci.pixel_has_data(y, x))
-                    else:
-                        self.assertTrue(sci.pixel_has_data(y, x))
-
-    def test_mask_grow(self):
-        # Mask one pixel per image.
-        for i in range(self.img_count):
-            img = self.stack.get_single_image(i)
-            msk = img.get_mask()
-            for x in range(self.dim_x):
-                msk.set_pixel(8, 2 + i, 1)
-
-        # Apply the bit vector based mask and check that one pixel per image is masked.
-        self.stack = BitVectorMasker(1).apply_mask(self.stack)
-        for i in range(self.img_count):
-            sci = self.stack.get_single_image(i).get_science()
-            for x in range(self.dim_x):
-                for y in range(self.dim_y):
-                    self.assertEqual(sci.pixel_has_data(y, x), x != (2 + i) or y != 8)
-
-        # Grow the mask by two pixels and recheck.
-        self.stack = GrowMask(2).apply_mask(self.stack)
-        for i in range(self.img_count):
-            sci = self.stack.get_single_image(i).get_science()
-            for x in range(self.dim_x):
-                for y in range(self.dim_y):
-                    dist = abs(2 + i - x) + abs(y - 8)
-                    self.assertEqual(sci.pixel_has_data(y, x), dist > 2)
-
-    def test_global_mask(self):
-        # Set each mask pixel in a single row depending on the image number.
-        for i in range(self.img_count):
-            img = self.stack.get_single_image(i)
-            msk = img.get_mask()
-
-            # Set key "CR" on every other image.
-            if i % 2:
-                msk.set_pixel(1, 1, 8)
-
-            # Set key "INTRP" on only one image.
-            if i == 0:
-                msk.set_pixel(5, 5, 4)
-
-        mask = GlobalDictionaryMasker(self.mask_bits_dict, ["CR", "INTRP"], 2)
-        self.stack = mask.apply_mask(self.stack)
-        for i in range(self.img_count):
-            sci = self.stack.get_single_image(i).get_science()
-            for x in range(self.dim_x):
-                for y in range(self.dim_y):
-                    self.assertEqual(sci.pixel_has_data(y, x), x != 1 or y != 1)
-
-
 class test_run_search_masking(unittest.TestCase):
     def setUp(self):
         # Create the a fake layered image.
@@ -219,8 +69,7 @@ def test_apply_masks(self):
         config.set_multiple(overrides)
 
         # Do the actual masking.
-        rs = SearchRunner()
-        self.stack = rs.do_masking(config, self.stack)
+        self.stack = apply_mask_operations(config, self.stack)
 
         # Test the the correct pixels have been masked.
         for i in range(self.img_count):
diff --git a/tests/test_raw_image.py b/tests/test_raw_image.py
index 2e4173830..b6ba5d438 100644
--- a/tests/test_raw_image.py
+++ b/tests/test_raw_image.py
@@ -341,31 +341,6 @@ def test_convolve_psf_orientation_gpu(self):
         """Test convolution on GPU with a non-symmetric PSF"""
         self.convolve_psf_orientation_cpu("GPU")
 
-    def test_grow_mask(self):
-        """Test grow_mask based on manhattan distances."""
-        img = RawImage(self.array)
-        img.set_pixel(5, 7, KB_NO_DATA)
-        img.set_pixel(3, 7, KB_NO_DATA)
-
-        for y in range(img.height):
-            for x in range(img.width):
-                should_mask = (y == 3 and x == 7) or (y == 5 and x == 7)
-                self.assertEqual(img.pixel_has_data(y, x), not should_mask)
-
-        # Grow the mask by one pixel.
-        img.grow_mask(1)
-        for y in range(img.height):
-            for x in range(img.width):
-                dist = min([abs(7 - x) + abs(3 - y), abs(7 - x) + abs(5 - y)])
-                self.assertEqual(img.pixel_has_data(y, x), dist > 1)
-
-        # Grow the mask by an additional two pixels (for a total of 3).
-        img.grow_mask(2)
-        for y in range(img.height):
-            for x in range(img.width):
-                dist = min([abs(7 - x) + abs(3 - y), abs(7 - x) + abs(5 - y)])
-                self.assertEqual(img.pixel_has_data(y, x), dist > 3)
-
     # Stamp as is tested here and as it's used in StackSearch are heaven and earth
     # TODO: Add proper tests
     def test_make_stamp(self):
diff --git a/tests/test_search.py b/tests/test_search.py
index 746125f8c..02aef5a8c 100644
--- a/tests/test_search.py
+++ b/tests/test_search.py
@@ -69,7 +69,7 @@ def setUp(self):
             if i % 2 == 0:
                 mask = im.get_mask()
                 mask.set_pixel(self.masked_y, self.masked_x, 1)
-                im.apply_mask_flags(1)
+                im.apply_mask(1)
 
             self.imlist.append(im)
         self.stack = ImageStack(self.imlist)
@@ -104,7 +104,7 @@ def test_psiphi(self):
 
         mask = image2.get_mask()
         mask.set_pixel(9, 4, 1)
-        image2.apply_mask_flags(1)
+        image2.apply_mask(1)
 
         # Create a stack from the two objects.
         stack = ImageStack([image1, image2])
@@ -485,7 +485,7 @@ def test_coadd_cpu_simple(self):
                 mask.set_pixel(1, 1, 1)
             if i == 1:
                 mask.set_pixel(0, 1, 1)
-            im.apply_mask_flags(1)
+            im.apply_mask(1)
 
             imlist.append(im)
         stack = ImageStack(imlist)
@@ -545,7 +545,7 @@ def test_coadd_gpu_simple(self):
                 mask.set_pixel(1, 1, 1)
             if i == 1:
                 mask.set_pixel(0, 1, 1)
-            im.apply_mask_flags(1)
+            im.apply_mask(1)
 
             imlist.append(im)
         stack = ImageStack(imlist)
diff --git a/tests/test_stamp_parity.py b/tests/test_stamp_parity.py
index bedffcca2..c5e3eb318 100644
--- a/tests/test_stamp_parity.py
+++ b/tests/test_stamp_parity.py
@@ -72,7 +72,7 @@ def setUp(self):
             if i % 2 == 0:
                 mask = im.get_mask()
                 mask.set_pixel(self.masked_y, self.masked_x, 1)
-                im.apply_mask_flags(1)
+                im.apply_mask(1)
 
             self.imlist.append(im)
         self.stack = ImageStack(self.imlist)

From 75dfc65967049e0050a4c5f21c41b19f281ccb1c Mon Sep 17 00:00:00 2001
From: Jeremy Kubica <104161096+jeremykubica@users.noreply.github.com>
Date: Thu, 11 Jan 2024 08:38:54 -0500
Subject: [PATCH 2/3] Fix linting errors

---
 src/kbmod/masking.py        | 6 +++---
 tests/test_layered_image.py | 2 +-
 tests/test_masking.py       | 1 +
 3 files changed, 5 insertions(+), 4 deletions(-)

diff --git a/src/kbmod/masking.py b/src/kbmod/masking.py
index a25f10092..0623b0375 100644
--- a/src/kbmod/masking.py
+++ b/src/kbmod/masking.py
@@ -3,6 +3,7 @@
 
 import kbmod.search as kb
 
+
 def mask_flags_from_dict(mask_bits_dict, flag_keys):
     """Generate a bitmask integer of flag keys from a dictionary
     of masking reasons and a list of reasons to use.
@@ -25,7 +26,7 @@ def mask_flags_from_dict(mask_bits_dict, flag_keys):
         bitmask += 2 ** mask_bits_dict[bit]
     return bitmask
 
-            
+
 def apply_mask_operations(config, stack):
     """Perform all the masking operations based on the search's configuration parameters.
 
@@ -79,6 +80,5 @@ def apply_mask_operations(config, stack):
     # Apply the masks to the images.
     for i in range(stack.img_count()):
         stack.get_single_image(i).apply_mask(0xFFFFFF)
-            
-    return stack
 
+    return stack
diff --git a/tests/test_layered_image.py b/tests/test_layered_image.py
index 391f81d0e..cadcc8f16 100644
--- a/tests/test_layered_image.py
+++ b/tests/test_layered_image.py
@@ -150,7 +150,7 @@ def test_binarize_mask(self):
         mask = self.image.get_mask()
         for x in range(9):
             mask.set_pixel(10, x, x)
-            
+
         # Only keep the mask at for pixels with flags at
         # bit positions 0 and 2 (1 + 4 = 5).
         self.image.binarize_mask(5)
diff --git a/tests/test_masking.py b/tests/test_masking.py
index 6dda673cf..eb3b3353a 100644
--- a/tests/test_masking.py
+++ b/tests/test_masking.py
@@ -4,6 +4,7 @@
 from kbmod.masking import apply_mask_operations
 from kbmod.search import *
 
+
 class test_run_search_masking(unittest.TestCase):
     def setUp(self):
         # Create the a fake layered image.

From 53fb617d7e28cfaf5872934223e1f3300f57640a Mon Sep 17 00:00:00 2001
From: Jeremy Kubica <104161096+jeremykubica@users.noreply.github.com>
Date: Fri, 12 Jan 2024 08:56:23 -0500
Subject: [PATCH 3/3] Address PR comments

---
 src/kbmod/masking.py                       | 2 +-
 src/kbmod/search/layered_image.h           | 2 +-
 src/kbmod/search/pydocs/image_stack_docs.h | 2 +-
 tests/test_layered_image.py                | 4 ++--
 4 files changed, 5 insertions(+), 5 deletions(-)

diff --git a/src/kbmod/masking.py b/src/kbmod/masking.py
index 0623b0375..d359b0229 100644
--- a/src/kbmod/masking.py
+++ b/src/kbmod/masking.py
@@ -77,7 +77,7 @@ def apply_mask_operations(config, stack):
         for i in range(stack.img_count()):
             stack.get_single_image(i).grow_mask(config["mask_grow"])
 
-    # Apply the masks to the images.
+    # Apply the masks to the images. Use 0xFFFFFF to apply all active masking bits.
     for i in range(stack.img_count()):
         stack.get_single_image(i).apply_mask(0xFFFFFF)
 
diff --git a/src/kbmod/search/layered_image.h b/src/kbmod/search/layered_image.h
index d3527fe31..11fe159ee 100644
--- a/src/kbmod/search/layered_image.h
+++ b/src/kbmod/search/layered_image.h
@@ -93,7 +93,7 @@ class LayeredImage {
     PSF psf;
     RawImage science;
     RawImage mask;
-    RawImage variance;   
+    RawImage variance;
 };
 
 } /* namespace search */
diff --git a/src/kbmod/search/pydocs/image_stack_docs.h b/src/kbmod/search/pydocs/image_stack_docs.h
index 26c38f49a..a00e99332 100644
--- a/src/kbmod/search/pydocs/image_stack_docs.h
+++ b/src/kbmod/search/pydocs/image_stack_docs.h
@@ -33,7 +33,7 @@ static const auto DOC_ImageStack_build_zeroed_times = R"doc(
   ")doc";
 
 static const auto DOC_ImageStack_make_global_mask = R"doc(
-  Create a new global mask that from a set of flags and a threshold.
+  Create a new global mask from a set of flags and a threshold.
   The global mask marks a pixel as masked if and only if it is masked
   by one of the given flags in at least ``threshold`` individual images.
   The returned mask is binary.
diff --git a/tests/test_layered_image.py b/tests/test_layered_image.py
index cadcc8f16..bfc45c368 100644
--- a/tests/test_layered_image.py
+++ b/tests/test_layered_image.py
@@ -151,7 +151,7 @@ def test_binarize_mask(self):
         for x in range(9):
             mask.set_pixel(10, x, x)
 
-        # Only keep the mask at for pixels with flags at
+        # Only keep the mask for pixels with flags at
         # bit positions 0 and 2 (1 + 4 = 5).
         self.image.binarize_mask(5)
         self.assertEqual(mask.get_pixel(10, 0), 0)
@@ -204,7 +204,7 @@ def test_add_threshold_mask_flags(self):
                     masked_pixels[index] = True
         self.assertGreater(len(masked_pixels), 0)
 
-        # Reset the mask an perform threshold masking.
+        # Reset the mask and perform threshold masking.
         mask = self.image.get_mask()
         mask.set_all(0.0)
         self.image.union_threshold_masking(threshold)