Add zero beam handling (#70)

Author: AlecThomson

racs_tools/beamcon_2D.py

@@ -21,6 +21,8 @@
 from tqdm import tqdm
 
 from racs_tools.convolve_uv import (
+    NAN_BEAM,
+    ZERO_BEAM,
     get_convolving_beam,
     get_nyquist_beam,
     my_ceil,
@@ -37,8 +39,6 @@
 )
 from racs_tools.parallel import get_executor
 
-# logger = setup_logger()
-
 
 #############################################
 #### ADAPTED FROM SCRIPT BY T. VERNSTROM ####
@@ -188,6 +188,9 @@ def savefile(
     """
     logger.info(f"Saving to {outfile.absolute()}")
     beam = new_beam
+    if np.isnan(beam):
+        logger.warning("Beam is NaN, setting to 0 in header")
+        beam = ZERO_BEAM
     header = beam.attach_to_header(header)
     fits.writeto(
         outfile.absolute(), newimage.astype(np.float32), header=header, overwrite=True
@@ -312,26 +315,23 @@ def get_common_beam(
         beam = Beam.from_fits_header(header)
         beams_list.append(beam)
 
-    beams = Beams(
-        [beam.major.to(u.deg).value for beam in beams_list] * u.deg,
-        [beam.minor.to(u.deg).value for beam in beams_list] * u.deg,
-        [beam.pa.to(u.deg).value for beam in beams_list] * u.deg,
-    )
+    beams = Beams(beams=beams_list)
+
+    # Init flags - False is good, True is bad
+    flags = np.array([False for beam in beams])
+
+    # Flag zero beams
+    flags = np.array([beam == ZERO_BEAM for beam in beams]) | flags
+
     if cutoff is not None:
-        flags = beams.major > cutoff * u.arcsec
+        flags = beams.major > cutoff * u.arcsec | flags
         if np.all(flags):
             logger.critical(
                 "All beams are larger than cutoff. All outputs will be blanked!"
             )
-            nan_beam = Beam(np.nan * u.deg, np.nan * u.deg, np.nan * u.deg)
-            nan_beams = Beams(
-                [np.nan for beam in beams_list] * u.deg,
-                [np.nan for beam in beams_list] * u.deg,
-                [np.nan for beam in beams_list] * u.deg,
-            )
-            return nan_beam, nan_beams
-    else:
-        flags = np.array([False for beam in beams])
+
+            nan_beams = Beams(beams=[NAN_BEAM for beam in beams_list])
+            return NAN_BEAM, nan_beams
 
     if target_beam is None:
         # Find the common beam

racs_tools/beamcon_3D.py

@@ -26,7 +26,7 @@
 
 from racs_tools import au2
 from racs_tools.beamcon_2D import my_ceil, round_up
-from racs_tools.convolve_uv import parse_conv_mode, smooth
+from racs_tools.convolve_uv import NAN_BEAM, ZERO_BEAM, parse_conv_mode, smooth
 from racs_tools.logging import (
     init_worker,
     log_listener,
@@ -194,7 +194,7 @@ def getfacs(
     """
     facs_list = []
     for conbm, old_beam in zip(convbeams, beams):
-        if conbm == Beam(major=0 * u.deg, minor=0 * u.deg, pa=0 * u.deg):
+        if conbm == ZERO_BEAM:
             fac = 1.0
         else:
             fac, _, _, _, _ = au2.gauss_factor(
@@ -393,9 +393,7 @@ def commonbeamer(
             total=nchans,
         ):
             if all(np.isnan(beams)):
-                commonbeam = Beam(
-                    major=np.nan * u.deg, minor=np.nan * u.deg, pa=np.nan * u.deg
-                )
+                commonbeam = NAN_BEAM
             else:
                 try:
                     commonbeam = beams[~np.isnan(beams)].common_beam(
@@ -585,9 +583,7 @@ def commonbeamer(
         masks = cube_data.mask
         for commonbeam, old_beam, mask in zip(commonbeams, old_beams, masks):
             if mask:
-                convbeam = Beam(
-                    major=np.nan * u.deg, minor=np.nan * u.deg, pa=np.nan * u.deg
-                )
+                convbeam = NAN_BEAM
             else:
                 old_beam_check = Beam(
                     major=my_ceil(old_beam.major.to(u.arcsec).value, precision=1)
@@ -597,12 +593,8 @@ def commonbeamer(
                     pa=round_up(old_beam.pa.to(u.deg), decimals=2),
                 )
                 if commonbeam == old_beam_check:
-                    convbeam = Beam(
-                        major=0 * u.deg,
-                        minor=0 * u.deg,
-                        pa=0 * u.deg,
-                    )
-                    logger.warn(
+                    convbeam = ZERO_BEAM
+                    logger.warning(
                         f"New beam {commonbeam!r} and old beam {old_beam_check!r} are the same. Won't attempt convolution."
                     )
 
@@ -682,7 +674,7 @@ def masking(
         list[CubeData]: list of masked cube data.
     """
     # Check for pipeline masking
-    nullbeam = Beam(major=0 * u.deg, minor=0 * u.deg, pa=0 * u.deg)
+    nullbeam = ZERO_BEAM
     tiny = np.finfo(np.float32).tiny  # Smallest positive number - used to mask
     smallbeam = Beam(major=tiny * u.deg, minor=tiny * u.deg, pa=tiny * u.deg)
     masked_cube_data_list: list[CubeData] = []
@@ -775,7 +767,7 @@ def initfiles(
         )
     ):
         logger.warning("Reference PSF is NaN - replacing with 0 in the header")
-        ref_psf = Beam(major=0 * u.deg, minor=0 * u.deg, pa=0 * u.deg)
+        ref_psf = ZERO_BEAM
         header["COMMENT"] = "Reference PSF is NaN"
         header["COMMENT"] = "- This is likely because the reference channel is masked."
         header["COMMENT"] = "- It has been replaced with 0 to keep FITS happy."

racs_tools/convolve_uv.py

@@ -20,6 +20,9 @@
 from racs_tools import au2
 from racs_tools.logging import logger
 
+ZERO_BEAM = Beam(0 * u.deg)
+NAN_BEAM = Beam(np.nan * u.deg)
+
 
 class ConvolutionResult(NamedTuple):
     """Result of convolution"""
@@ -170,7 +173,7 @@ def convolve(
         return ConvolutionResult(image * np.nan, sfactor)
     if np.isnan(image).all():
         return ConvolutionResult(image, sfactor)
-    if conbeam == Beam(major=0 * u.deg, minor=0 * u.deg, pa=0 * u.deg) and sfactor == 1:
+    if conbeam == ZERO_BEAM and sfactor == 1:
         return ConvolutionResult(image, sfactor)
     ###
 
@@ -278,7 +281,7 @@ def convolve_scipy(
         return ConvolutionResult(image * np.nan, sfactor)
     if np.isnan(image).all():
         return ConvolutionResult(image, sfactor)
-    if conbeam == Beam(major=0 * u.deg, minor=0 * u.deg, pa=0 * u.deg) and sfactor == 1:
+    if conbeam == ZERO_BEAM and sfactor == 1:
         return ConvolutionResult(image, sfactor)
 
     gauss_kern = conbeam.as_kernel(dy)
@@ -326,7 +329,7 @@ def convolve_astropy(
         return ConvolutionResult(image * np.nan, sfactor)
     if np.isnan(image).all():
         return ConvolutionResult(image, sfactor)
-    if conbeam == Beam(major=0 * u.deg, minor=0 * u.deg, pa=0 * u.deg) and sfactor == 1:
+    if conbeam == ZERO_BEAM and sfactor == 1:
         return ConvolutionResult(image, sfactor)
     gauss_kern = conbeam.as_kernel(dy)
     conbm1 = gauss_kern.array / gauss_kern.array.max()
@@ -377,7 +380,7 @@ def convolve_astropy_fft(
         return ConvolutionResult(image * np.nan, sfactor)
     if np.isnan(image).all():
         return ConvolutionResult(image, sfactor)
-    if conbeam == Beam(major=0 * u.deg, minor=0 * u.deg, pa=0 * u.deg) and sfactor == 1:
+    if conbeam == ZERO_BEAM and sfactor == 1:
         return ConvolutionResult(image, sfactor)
     gauss_kern = conbeam.as_kernel(dy)
     conbm1 = gauss_kern.array / gauss_kern.array.max()
@@ -453,20 +456,12 @@ def get_convolving_beam(
 
     if cutoff is not None and old_beam.major.to(u.arcsec) > cutoff * u.arcsec:
         return (
-            Beam(
-                major=np.nan * u.deg,
-                minor=np.nan * u.deg,
-                pa=np.nan * u.deg,
-            ),
+            NAN_BEAM,
             np.nan,
         )
 
     if new_beam == old_beam:
-        conbm = Beam(
-            major=0 * u.deg,
-            minor=0 * u.deg,
-            pa=0 * u.deg,
-        )
+        conbm = ZERO_BEAM
         fac = 1.0
         logger.warning(
             f"New beam {new_beam!r} and old beam {old_beam!r} are the same. Won't attempt convolution."