From 2c08082b8d2e3c6c89fcabb6f0cc83f8cb59cd7f Mon Sep 17 00:00:00 2001
From: Robert Stein <rdstein@caltech.edu>
Date: Fri, 6 Dec 2024 15:21:56 -0500
Subject: [PATCH] Add gmos masking

---
 mirar/pipelines/gmos/blocks.py    |  8 ++++++-
 mirar/pipelines/gmos/gmos_mask.py | 36 +++++++++++++++++++++++++++++++
 2 files changed, 43 insertions(+), 1 deletion(-)
 create mode 100644 mirar/pipelines/gmos/gmos_mask.py

diff --git a/mirar/pipelines/gmos/blocks.py b/mirar/pipelines/gmos/blocks.py
index 2598ef428..2b902c99f 100644
--- a/mirar/pipelines/gmos/blocks.py
+++ b/mirar/pipelines/gmos/blocks.py
@@ -27,18 +27,21 @@
     gmos_reference_sextractor,
     gmos_zogy_catalogs_purifier,
 )
+from mirar.pipelines.gmos.gmos_mask import generate_gmos_mask
 from mirar.pipelines.gmos.load_gmos_image import load_detrended_gmos_image
 from mirar.processors.astromatic import PSFex, Scamp, Sextractor
 from mirar.processors.astromatic.swarp import Swarp
 from mirar.processors.astrometry.anet import AstrometryNet
 from mirar.processors.astrometry.autoastrometry import AutoAstrometry
 from mirar.processors.csvlog import CSVLog
+from mirar.processors.mask import MaskPixelsFromFunction
 from mirar.processors.photcal.photcalibrator import PhotCalibrator
 from mirar.processors.photometry import AperturePhotometry, PSFPhotometry
 from mirar.processors.reference import ProcessReference
 from mirar.processors.sources import (
     CSVExporter,
     ForcedPhotometryDetector,
+    ImageUpdater,
     ParquetWriter,
 )
 from mirar.processors.sources.utils import RegionsWriter
@@ -48,6 +51,8 @@
 load_raw = [
     ImageLoader(input_sub_dir="detrend", load_image=load_detrended_gmos_image),
     ImageBatcher(BASE_NAME_KEY),
+    MaskPixelsFromFunction(mask_function=generate_gmos_mask),
+    ImageSaver(output_dir_name="mask"),
 ]
 
 build_log = [  # pylint: disable=duplicate-code
@@ -128,7 +133,7 @@
     ZOGY(output_sub_dir="zogy"),
     ImageSaver(output_dir_name="diff"),
     ForcedPhotometryDetector(ra_header_key="OBJRA", dec_header_key="OBJDEC"),
-    RegionsWriter(output_dir_name="diffs"),
+    RegionsWriter(output_dir_name="diff"),
     AperturePhotometry(
         aper_diameters=[
             2 / GMOS_PIXEL_SCALE,
@@ -158,4 +163,5 @@
     PSFPhotometry(),
     ParquetWriter(output_dir_name="sources"),
     CSVExporter(output_dir_name="sources"),
+    ImageUpdater(modify_dir_name="diff"),
 ]
diff --git a/mirar/pipelines/gmos/gmos_mask.py b/mirar/pipelines/gmos/gmos_mask.py
new file mode 100644
index 000000000..1434c488f
--- /dev/null
+++ b/mirar/pipelines/gmos/gmos_mask.py
@@ -0,0 +1,36 @@
+"""
+Functions to create a mask for GMOS data.
+"""
+
+import logging
+
+import numpy as np
+
+from mirar.data import Image
+
+logger = logging.getLogger(__name__)
+
+
+def generate_gmos_mask(image: Image) -> np.array:
+    """
+    Function to create a mask for GMOS data
+
+    :param image: Image object
+    :return: Image object with mask
+    """
+
+    data = image.get_data()
+
+    median = np.nanmedian(data)
+    std = np.nanstd(data)
+
+    print("Median: ", median)
+    print("STD: ", std)
+
+    mask = data < 0.5 * median
+
+    print("Mask: ", mask)
+
+    logger.info(f"Masking {np.sum(mask)}/{len(mask)} pixels")
+
+    return mask