Better wasp pipeline (#1051)

mirar/data/cache.py

@@ -10,7 +10,7 @@
 
 logger = logging.getLogger(__name__)
 
-USE_CACHE: bool = os.getenv("USE_MIRAR_CACHE", "true") in ["true", "True", True]
+USE_CACHE: bool = os.getenv("USE_MIRAR_CACHE", "true") in ["true", "True", True, 1, "1"]
 
 if not USE_CACHE:
     if os.getenv("USE_WINTER_CACHE") is not None:
@@ -19,7 +19,7 @@
             "Please use 'USE_MIRAR_CACHE' instead. "
             "This will be removed in a future version."
         )
-        USE_CACHE = os.getenv("USE_WINTER_CACHE") in ["true", "True", True]
+        USE_CACHE = os.getenv("USE_WINTER_CACHE") in ["true", "True", True, 1, "1"]
 
 
 class CacheError(Exception):

mirar/pipelines/wasp/blocks.py

@@ -14,9 +14,12 @@
     sextractor_astrometry_config,
     sextractor_photometry_config,
     swarp_config_path,
+    wasp_cal_requirements,
 )
 from mirar.pipelines.wasp.config.constants import WASP_PIXEL_SCALE
 from mirar.pipelines.wasp.generator import (
+    annotate_target_coordinates,
+    label_stack_id,
     wasp_astrometric_catalog_generator,
     wasp_photometric_catalog_generator,
     wasp_reference_image_generator,
@@ -39,22 +42,31 @@
 from mirar.processors.sources import (
     CSVExporter,
     ForcedPhotometryDetector,
+    ImageUpdater,
     ParquetWriter,
 )
+from mirar.processors.sources.utils import RegionsWriter
 from mirar.processors.utils import (
+    CustomImageBatchModifier,
     ImageBatcher,
     ImageDebatcher,
     ImageLoader,
+    ImageRebatcher,
     ImageSaver,
     ImageSelector,
 )
+from mirar.processors.utils.cal_hunter import CalHunter
 from mirar.processors.zogy.zogy import ZOGY, ZOGYPrepare
 
 load_raw = [
     ImageLoader(input_sub_dir="raw", load_image=load_raw_wasp_image),
-    ImageBatcher(BASE_NAME_KEY),
+    CustomImageBatchModifier(label_stack_id),
+    ImageRebatcher(split_key="stackid"),
+    CustomImageBatchModifier(annotate_target_coordinates),
+    ImageRebatcher(BASE_NAME_KEY),
 ]
 
+
 build_log = [  # pylint: disable=duplicate-code
     CSVLog(
         export_keys=[
@@ -63,6 +75,7 @@
             "OBJDEC",
             "DATE-OBS",
             "FILTER",
+            "STACKID",
             OBSCLASS_KEY,
             BASE_NAME_KEY,
         ]
@@ -71,8 +84,9 @@
 ]  # pylint: disable=duplicate-code
 
 calibrate = [
-    ImageSelector((OBSCLASS_KEY, ["bias", "flat", "science"])),
     ImageDebatcher(),
+    CalHunter(load_image=load_raw_wasp_image, requirements=wasp_cal_requirements),
+    ImageSelector((OBSCLASS_KEY, ["bias", "flat", "science"])),
     BiasCalibrator(),
     ImageSelector((OBSCLASS_KEY, ["flat", "science"])),
     ImageBatcher(split_key="filter"),
@@ -89,8 +103,7 @@
         scamp_config_path=scamp_path,
         cache=False,
     ),
-    ImageDebatcher(),
-    ImageBatcher(split_key=["target", "filter"]),
+    ImageRebatcher(split_key=["stackid"]),
     Swarp(
         swarp_config_path=swarp_config_path,
         include_scamp=True,
@@ -141,6 +154,8 @@
     ZOGY(output_sub_dir="zogy"),
     ImageSaver(output_dir_name="diff"),
     ForcedPhotometryDetector(ra_header_key="OBJRA", dec_header_key="OBJDEC"),
+    RegionsWriter(output_dir_name="diff"),
+    RegionsWriter(output_dir_name="processed"),
     AperturePhotometry(
         aper_diameters=[
             2 / WASP_PIXEL_SCALE,
@@ -170,4 +185,5 @@
     PSFPhotometry(),
     ParquetWriter(output_dir_name="sources"),
     CSVExporter(output_dir_name="sources"),
+    ImageUpdater(modify_dir_name="diff"),
 ]

mirar/pipelines/wasp/config/__init__.py

@@ -7,6 +7,7 @@
 import os
 
 from mirar.pipelines.wasp.config.constants import PIPELINE_NAME
+from mirar.processors.utils.cal_hunter import CalRequirement
 
 wasp_dir = os.path.dirname(__file__)
 
@@ -31,3 +32,9 @@
 swarp_config_path = os.path.join(astromatic_config_dir, "config.swarp")
 
 psfex_sci_config_path = os.path.join(astromatic_config_dir, "photom_sci.psfex")
+
+wasp_cal_requirements = [
+    CalRequirement(
+        target_name="bias", required_field="EXPTIME", required_values=["0.0"]
+    ),
+]

mirar/pipelines/wasp/generator.py → mirar/pipelines/wasp/generator/__init__.py

@@ -16,6 +16,8 @@
     sextractor_photometry_config,
     swarp_config_path,
 )
+from mirar.pipelines.wasp.generator.stacks import label_stack_id
+from mirar.pipelines.wasp.generator.target import annotate_target_coordinates
 from mirar.processors.astromatic import PSFex, Sextractor, Swarp
 from mirar.processors.astromatic.sextractor.sextractor import SEXTRACTOR_HEADER_KEY
 from mirar.references import BaseReferenceGenerator, PS1Ref, SDSSRef
@@ -58,10 +60,11 @@ def wasp_photometric_catalog_generator(image: Image) -> BaseCatalog:
     :return: catalog at image position
     """
     filter_name = image["FILTER"].replace("'", "")
-    dec = image["OBJDEC"]
+
+    ra, dec = image["CRVAL1"], image["CRVAL2"]
 
     if filter_name in ["u", "U"]:
-        if in_sdss(image["OBJRA"], image["OBJDEC"]):
+        if in_sdss(ra, dec):
             return SDSS(
                 min_mag=10,
                 max_mag=WASP_PHOTOMETRIC_MAX_MAG,

mirar/pipelines/wasp/generator/stacks.py

@@ -0,0 +1,65 @@
+"""
+Module to group images based on the target coordinates into planned stack groups
+"""
+
+from astropy import coordinates as coords
+from astropy import units as u
+from astropy.coordinates import Angle
+
+from mirar.data import ImageBatch
+from mirar.paths import OBSCLASS_KEY
+from mirar.processors.utils.image_selector import split_images_into_batches
+
+MAX_RADIUS_DEG = 9.0 / 60.0  # WASP is 18 arc minutes each side
+
+
+def label_stack_id(batch: ImageBatch) -> ImageBatch:
+    """
+    Label the stack id of the images in the batch
+    :param batch: Original batch of images
+    :return: Labeled batch of images
+    """
+
+    ras = []
+    decs = []
+
+    for image in batch:
+
+        if image[OBSCLASS_KEY] != "science":
+            image["targnum"] = -1
+            continue
+
+        target_ra = Angle(image["CRVAL1"], unit="hourangle").degree
+        target_dec = Angle(image["CRVAL2"], unit="degree").degree
+
+        position = coords.SkyCoord(target_ra, target_dec, unit="deg")
+
+        match = None
+
+        if len(ras) > 0:
+            idx, d2d, _ = position.match_to_catalog_sky(
+                coords.SkyCoord(ra=ras, dec=decs, unit="deg")
+            )
+
+            mask = d2d < (MAX_RADIUS_DEG * u.deg)
+            if mask:
+                match = idx
+
+        if match is None:
+            ras.append(target_ra)
+            decs.append(target_dec)
+            image["targnum"] = int(len(ras) - 1)
+        else:
+            image["targnum"] = int(match)
+
+    new_batches = split_images_into_batches(batch, ["targnum", "filter", "exptime"])
+
+    combined_batch = ImageBatch()
+
+    for i, split_batch in enumerate(new_batches):
+        label = f"stack{i}"
+        for image in split_batch:
+            image["stackid"] = label
+            combined_batch.append(image)
+
+    return combined_batch

mirar/pipelines/wasp/generator/target.py

@@ -0,0 +1,31 @@
+"""
+Module to annotate target coordinates on images.
+"""
+
+from mirar.data import ImageBatch
+from mirar.paths import TARGET_KEY, TIME_KEY
+
+
+def annotate_target_coordinates(image_batch: ImageBatch) -> ImageBatch:
+    """
+    Function to annotate target coordinates on images.
+    For WASP, this should be the value of RA/DEC in the header of the first image.
+
+    :param image_batch: ImageBatch object
+    :return: ImageBatch object
+    """
+
+    times = [image[TIME_KEY] for image in image_batch]
+    min_time = min(times)
+    first_image = image_batch[times.index(min_time)]
+
+    # In case one of the dithers is mis-named, we'll use the most common name
+    names = [x[TARGET_KEY] for x in image_batch]
+    most_common_name = max(set(names), key=names.count)
+
+    for image in image_batch:
+        image["OBJRA"] = first_image["OBJRA"]
+        image["OBJDEC"] = first_image["OBJDEC"]
+        image[TARGET_KEY] = most_common_name
+
+    return image_batch

mirar/pipelines/wasp/wasp_pipeline.py

@@ -8,7 +8,7 @@
 from mirar.data import Image
 from mirar.pipelines.base_pipeline import Pipeline
 from mirar.pipelines.wasp.blocks import build_log, load_raw, reduce, subtract
-from mirar.pipelines.wasp.config import PIPELINE_NAME
+from mirar.pipelines.wasp.config import PIPELINE_NAME, wasp_cal_requirements
 from mirar.pipelines.wasp.load_wasp_image import load_raw_wasp_image
 
 logger = logging.getLogger(__name__)
@@ -28,6 +28,8 @@ class WASPPipeline(Pipeline):
         "reduce": load_raw + reduce,
     }
 
+    defalut_cal_requirements = wasp_cal_requirements
+
     @staticmethod
     def download_raw_images_for_night(night: str | int):
         """

mirar/processors/sources/utils/regions_writer.py

@@ -7,9 +7,14 @@
 from pathlib import Path
 from typing import Optional
 
+from astropy import units as u
+from astropy.coordinates import Angle
+
 from mirar.data import SourceBatch
 from mirar.paths import (
     BASE_NAME_KEY,
+    CAND_DEC_KEY,
+    CAND_RA_KEY,
     XPOS_KEY,
     YPOS_KEY,
     base_output_dir,
@@ -32,11 +37,13 @@ def __init__(
         output_dir_name: Optional[str] = None,
         region_pix_radius: float = 8,
         output_dir: str | Path = base_output_dir,
+        use_ra_dec: bool = True,
     ):
         super().__init__()
         self.output_dir_name = output_dir_name
         self.region_pix_radius = region_pix_radius
         self.output_dir = Path(output_dir)
+        self.use_ra_dec = use_ra_dec
 
     def description(self) -> str:
         return (
@@ -63,12 +70,29 @@ def _apply_to_sources(
 
             regions_path.parent.mkdir(parents=True, exist_ok=True)
 
-            with open(f"{regions_path}", "w", encoding="utf8") as regions_f:
-                regions_f.write("image\n")
-                for _, row in candidate_table.iterrows():
-                    regions_f.write(
-                        f"CIRCLE({row[XPOS_KEY]},{row[YPOS_KEY]},"
-                        f"{self.region_pix_radius})\n"
-                    )
+            if self.use_ra_dec:
+                # Write regions file in ra/dec coordinates
+                with open(f"{regions_path}", "w", encoding="utf8") as regions_f:
+                    for _, row in candidate_table.iterrows():
+                        ra = Angle(row[CAND_RA_KEY] * u.deg).to_string(
+                            unit=u.hourangle, sep=":"
+                        )
+                        dec = Angle(row[CAND_DEC_KEY] * u.deg).to_string(
+                            unit=u.deg, sep=":"
+                        )
+
+                        regions_f.write(
+                            f"CIRCLE({ra},{dec}," f"{self.region_pix_radius})\n"
+                        )
+
+            else:
+                # Write regions file in pixel coordinates
+                with open(f"{regions_path}", "w", encoding="utf8") as regions_f:
+                    regions_f.write("image\n")
+                    for _, row in candidate_table.iterrows():
+                        regions_f.write(
+                            f"CIRCLE({row[XPOS_KEY]},{row[YPOS_KEY]},"
+                            f"{self.region_pix_radius})\n"
+                        )
 
         return batch

mirar/processors/utils/cal_hunter.py

@@ -195,6 +195,8 @@ class CalHunter(ImageLoader):
     by searching previous nights of data
     """
 
+    max_n_cpu = 1
+
     base_key = "calhunt"
 
     def __init__(