Merge pull request #161 from tjgalvin/polname

Changing the `pol` field name

CHANGELOG.md

@@ -26,6 +26,9 @@
     - list the reference catalogues that are expected
 - Added `flint.leakage` CLI program to attempt to characterise leakage over
   polarisations, e.g. V/I
+- removing the `pol` string in the polarisation field of the
+  `ProcessedNameComponents`
+  - `wclean` output `-` separation character chhanged to `.`
 
 # 0.2.5
 

flint/coadd/linmos.py

@@ -96,7 +96,7 @@ def trim_fits_image(
     """
     logger.info(f"Trimming {image_path.name}")
     with fits.open(image_path) as fits_image:
-        data = fits_image[0].data
+        data = fits_image[0].data  # type: ignore
         logger.info(f"Original data shape: {data.shape}")
 
         image_shape = (data.shape[-2], data.shape[-1])
@@ -118,7 +118,7 @@ def trim_fits_image(
             bounding_box.ymin : bounding_box.ymax,
         ]
 
-        header = fits_image[0].header
+        header = fits_image[0].header  # type: ignore
         header["CRPIX1"] -= bounding_box.ymin
         header["CRPIX2"] -= bounding_box.xmin
 
@@ -162,7 +162,7 @@ def get_image_weight(
     ), f"Invalid {mode=} specified. Available modes: {weight_modes}"
 
     with fits.open(image_path, memmap=True) as in_fits:
-        image_data = in_fits[image_slice].data
+        image_data = in_fits[image_slice].data  # type: ignore
 
         assert len(
             image_data.shape
@@ -185,7 +185,7 @@ def get_image_weight(
             )
 
     logger.info(f"Weight {weight:.3f} for {image_path}")
-    return weight
+    return float(weight)
 
 
 def generate_weights_list_and_files(
@@ -336,11 +336,10 @@ def generate_linmos_parameter_set(
         parset += (
             f"linmos.primarybeam      = ASKAP_PB\n"
             f"linmos.primarybeam.ASKAP_PB.image = {str(holofile.absolute())}\n"
-            # f"linmos.primarybeam.ASKAP_PB.alpha = {paf_alpha}\n"
             f"linmos.removeleakage    = true\n"
         )
-        if alpha:
-            parset += f"linmos.primarybeam.ASKAP_PB.alpha = {alpha} # in radians\n"
+        assert alpha is not None, f"{alpha=}, which should not happen"
+        parset += f"linmos.primarybeam.ASKAP_PB.alpha = {alpha} # in radians\n"
 
     # Now write the file, me hearty
     logger.info(f"Writing parset to {str(parset_output_path)}.")

flint/imager/wsclean.py

@@ -1,7 +1,22 @@
-"""Simple interface into wsclean"""
+"""Simple interface into wsclean
+
+Some notes around the file naming.
+
+A certain filenaming scheme is required for FITS files to perform leakage correction
+when co-adding them together in the yandasoft linmos application. The stokes field
+needs to be encoded as ``.i.``. For example:
+
+    >>> `SB1234.RACS_0123+43.beam01.i.image.fits`
+
+The wsclean formatted output string appends something denoted with ``-``. Within
+this code there is and attempt to rename the wsclean outputs to replace the ``-``
+with a ``.``.
+
+"""
 
 from __future__ import annotations
 
+import re
 from argparse import ArgumentParser
 from glob import glob
 from numbers import Number
@@ -164,6 +179,60 @@ def with_options(self, **kwargs) -> WSCleanCommand:
         return WSCleanCommand(**_dict)
 
 
+def _rename_wsclean_title(name_str: str) -> str:
+    """Construct an apply a regular expression that aims to identify
+    the wsclean appended properties string within a file and replace
+    the `-` separator with a `.`.
+
+    Args:
+        name_str (str): The name that will be extracted and modified
+
+    Returns:
+        str: The modified string if a wsclean string was matched, otherwise the input `name-str`
+    """
+    search_re = r"(-(i|q|u|v|xx|xy|yx|yy))?-((MFS|[0-9]{4}))(-t[0-9]{5})?-(image|dirty|model|residual|psf)"
+    match_re = re.compile(search_re)
+
+    logger.info(f"{name_str=} {type(name_str)=}")
+    result = match_re.search(str(name_str))
+
+    if result is None:
+        return name_str
+
+    name = name_str.replace(result[0], result[0].replace("-", "."))
+
+    return name
+
+
+def _rename_wsclean_file(
+    input_path: Path,
+    rename_file: bool = False,
+) -> Path:
+    """Rename a file with wsclean appended string information to convert its
+    `-` separation markers with `.`. This should handle skipping the field
+    name of the target field observed.
+
+    Args:
+        input_path (Path): The file path that would be examined and modified
+        clean_parts (Union[int, Tuple[int, ...]], optional): Which parts of a split string will be modified. Defaults to -2.
+        rename_file (bool, optional): Whether the file should be moved / renamed. Defaults to False.
+
+    Returns:
+        Path: Path to the renamed file
+    """
+    input_path = Path(input_path)
+    file_name = Path(input_path.name)
+    new_name = _rename_wsclean_title(name_str=str(file_name))
+
+    new_path = input_path.parent / new_name
+
+    if rename_file:
+        logger.info(f"Renaming {input_path} to {new_path}")
+        input_path.rename(new_path)
+
+    return new_path
+
+
 def _wsclean_output_callback(line: str) -> None:
     """Call back function used to detect clean divergence"""
 
@@ -212,7 +281,7 @@ def get_wsclean_output_names(
     Returns:
         ImageSet: The file paths that wsclean should create/has created.
     """
-    # TODO: NEED TESTS!
+    # TODO: Use a regular expression for this
     subband_strs = [f"{subband:04}" for subband in range(subbands)]
     if include_mfs:
         subband_strs.append("MFS")
@@ -560,6 +629,41 @@ def combine_subbands_to_cube(
     return ImageSet(**imageset_dict)
 
 
+def rename_wsclean_prefix_in_imageset(input_imageset: ImageSet) -> ImageSet:
+    """Given an input imageset, rename the files contained in it to
+    remove the `-` separator that wsclean uses and replace it with `.`.
+
+    Files will be renamed on disk appropriately.
+
+    Args:
+        input_imageset (ImageSet): The collection of output wsclean products
+
+    Returns:
+        ImageSet: The updated imageset after replacing the separator and renaming files
+    """
+
+    input_args = options_to_dict(input_options=input_imageset)
+
+    check_keys = ("prefix", "image", "residual", "model", "dirty")
+
+    output_args: Dict[str, Any] = {}
+
+    for key, value in input_args.items():
+        if key == "prefix":
+            output_args[key] = _rename_wsclean_title(name_str=value)
+        elif key in check_keys and value is not None:
+            output_args[key] = [
+                _rename_wsclean_file(input_path=fits_path, rename_file=True)
+                for fits_path in value
+            ]
+        else:
+            output_args[key] = value
+
+    output_imageset = ImageSet(**output_args)
+
+    return output_imageset
+
+
 def run_wsclean_imager(
     wsclean_cmd: WSCleanCommand,
     container: Path,
@@ -649,6 +753,8 @@ def run_wsclean_imager(
             imageset=imageset, remove_original_images=True
         )
 
+    imageset = rename_wsclean_prefix_in_imageset(input_imageset=imageset)
+
     logger.info(f"Constructed {imageset=}")
 
     return wsclean_cmd.with_options(imageset=imageset)

flint/leakage.py

@@ -29,7 +29,7 @@ class LeakageFilters(NamedTuple):
     """The upper limit on acceptable int/peak ratios"""
     lower_int_peak_ratio: float = 0.8
     """The lower limit on acceptable int/peak ratios"""
-    search_box_size: int = 4
+    search_box_size: int = 1
     """The size of a box to search for peak polarised signal in"""
     noise_box_size: int = 30
     """the size of a box to compute a local RMS noise measure from"""
@@ -75,8 +75,8 @@ def _load_fits_image(fits_path: Path) -> FITSImage:
     ), f"Unexpected file type for {fits_path=}, expected fits"
     logger.info(f"Opening {fits_path=}")
     with fits.open(fits_path) as in_fits:
-        image_data = in_fits[0].data
-        header = dict(in_fits[0].header.items())
+        image_data = in_fits[0].data  # type: ignore
+        header = dict(in_fits[0].header.items())  # type: ignore
         wcs = WCS(header)
 
     return FITSImage(data=image_data, header=header, wcs=wcs, path=fits_path)
@@ -298,11 +298,31 @@ def extract_pol_stats_in_box(
     return pol_peak, pol_noise
 
 
+def _get_output_catalogue_path(
+    input_path: Path, pol: str, output_path: Optional[Path] = None
+) -> Path:
+    """Create the output leakage catalogue name"""
+    # NOTE: This is a separate function to test against after a silly. Might move with the other named Pirates
+    assert isinstance(input_path, Path)
+    input_suffix = input_path.suffix
+
+    output_path = (
+        input_path.with_suffix(f".{pol}_leakage{input_suffix}")
+        if output_path is None
+        else output_path
+    )
+    assert (
+        output_path is not None
+    ), f"{output_path=} is empty, and no catalogue path provided"
+
+    return Path(output_path)
+
+
 def create_leakge_component_table(
     pol_image: Path,
     catalogue: Union[Table, Path],
     pol: str = "v",
-    output_base: Optional[Path] = None,
+    output_path: Optional[Path] = None,
 ) -> Path:
     """Create a component catalogue that includes enough information to describe the
     polarisation fraction of sources across a field. This is intended to be used
@@ -318,7 +338,7 @@ def create_leakge_component_table(
         pol_image (Path): The polarised image that will be used to extract peak polarised flux from
         catalogue (Union[Table, Path]): Component table describing positions to extract flux from
         pol (str, optional): The polarisation stokes being considered. Defaults to "v".
-        output_base (Optional[Path], optional): The base name of the new catalogue. If `None` it is derived from the input `catalogue` path. Defaults to None.
+        output_path (Optional[Path], optional): The path of the new catalogue. If `None` it is derived from the input `catalogue` path. Defaults to None.
 
     Returns:
         Path: Path to the new catalogue use for leakage
@@ -347,22 +367,16 @@ def create_leakge_component_table(
     components[f"{pol}_peak"] = pol_peak
     components[f"{pol}_noise"] = pol_noise
 
-    if isinstance(catalogue, Path):
-        catalogue_suffix = catalogue.suffix
-        output_base = (
-            catalogue.with_suffix(f".{pol}_leakage.{catalogue_suffix}")
-            if output_base is None
-            else output_base
-        )
-
-    assert (
-        output_base is not None
-    ), f"{output_base=} is empty, and no catalogue path provided"
+    output_path = _get_output_catalogue_path(
+        input_path=catalogue if isinstance(catalogue, Path) else pol_image,
+        pol=pol,
+        output_path=output_path,
+    )
 
-    logger.info(f"Writing {output_base}")
-    components.write(output_base, overwrite=True)
+    logger.info(f"Writing {output_path}")
+    components.write(output_path, overwrite=True)
 
-    return output_base
+    return output_path
 
 
 def get_parser() -> ArgumentParser:

flint/naming.py

@@ -51,7 +51,7 @@ def create_imaging_name_prefix(ms: Union[MS, Path], pol: Optional[str] = None) -
 
     name = ms_path.stem
     if pol:
-        name = f"{name}.pol{pol.lower()}"
+        name = f"{name}.{pol.lower()}"
 
     return name
 
@@ -283,7 +283,7 @@ def processed_ms_format(
     logger.debug(f"Matching {in_name}")
     # A raw string is used to avoid bad unicode escaping
     regex = re.compile(
-        r"^SB(?P<sbid>[0-9]+)\.(?P<field>.+)\.beam(?P<beam>[0-9]+)((\.spw(?P<spw>[0-9]+))?)((\.round(?P<round>[0-9]+))?)((\.pol(?P<pol>[a-zA-z]+))?)*"
+        r"^SB(?P<sbid>[0-9]+)\.(?P<field>.+)\.beam(?P<beam>[0-9]+)((\.spw(?P<spw>[0-9]+))?)((\.round(?P<round>[0-9]+))?)((\.(?P<pol>(i|q|u|v|xx|yy|xy|yx)+))?)*"
     )
     results = regex.match(in_name)
 

flint/prefect/common/imaging.py

@@ -172,7 +172,7 @@ def task_run_bane_and_aegean(
         logger.info(f"Have extracted image: {image_paths}")
 
         # For the moment, will only source find on an MFS image
-        image_paths = [image for image in image_paths if "-MFS-" in str(image)]
+        image_paths = [image for image in image_paths if ".MFS." in str(image)]
         assert (
             len(image_paths) == 1
         ), "More than one image found after filter for MFS only images. "
@@ -462,7 +462,7 @@ def task_convolve_image(
 def task_linmos_images(
     images: Collection[Collection[Path]],
     container: Path,
-    filter: str = "-MFS-",
+    filter: str = ".MFS.",
     field_name: Optional[str] = None,
     suffix_str: str = "noselfcal",
     holofile: Optional[Path] = None,
@@ -476,7 +476,7 @@ def task_linmos_images(
     Args:
         images (Collection[Collection[Path]]): Images that will be co-added together
         container (Path): Path to singularity container that contains yandasoft
-        filter (str, optional): Filter to extract the images that will be extracted from the set of input images. These will be co-added. Defaults to "-MFS-".
+        filter (str, optional): Filter to extract the images that will be extracted from the set of input images. These will be co-added. Defaults to ".MFS.".
         field_name (Optional[str], optional): Name of the field, which is included in the output images created. Defaults to None.
         suffix_str (str, optional): Additional string added to the prefix of the output linmos image products. Defaults to "noselfcal".
         holofile (Optional[Path], optional): The FITS cube with the beam corrections derived from ASKAP holography. Defaults to None.
@@ -550,7 +550,7 @@ def _convolve_linmos(
     cutoff: float = 0.05,
     field_summary: Optional[FieldSummary] = None,
     convol_mode: str = "image",
-    convol_filter: str = "-MFS-",
+    convol_filter: str = ".MFS.",
     convol_suffix_str: str = "conv",
 ) -> LinmosCommand:
     """An internal function that launches the convolution to a common resolution
@@ -564,7 +564,7 @@ def _convolve_linmos(
         cutoff (float, optional): The primary beam attenuation cutoff supplied to linmos when coadding. Defaults to 0.05.
         field_summary (Optional[FieldSummary], optional): The summary of the field, including (importantly) to orientation of the third-axis. Defaults to None.
         convol_mode (str, optional): The mode passed to the convol task to describe the images to extract. Support image or residual.  Defaults to image.
-        convol_filter (str, optional): A text file applied when assessing images to co-add. Defaults to '-MFS-'.
+        convol_filter (str, optional): A text file applied when assessing images to co-add. Defaults to '.MFS.'.
         convol_suffix_str (str, optional): The suffix added to the convolved images. Defaults to 'conv'.
 
     Returns:
@@ -636,7 +636,7 @@ def _create_convol_linmos_images(
         beam_shape = task_get_common_beam.submit(
             wsclean_cmds=wsclean_cmds,
             cutoff=field_options.beam_cutoff,
-            filter="-MFS-",
+            filter=".MFS.",
             fixed_beam_shape=round_beam_shape,
         )
         # NOTE: The order matters here. The last linmos file is used
@@ -652,7 +652,7 @@ def _create_convol_linmos_images(
                     cutoff=field_options.pb_cutoff,
                     field_summary=field_summary,
                     convol_mode="residual",
-                    convol_filter="-MFS-",
+                    convol_filter=".MFS.",
                     convol_suffix_str=convol_suffix_str,
                 )
             )
@@ -665,7 +665,7 @@ def _create_convol_linmos_images(
                 cutoff=field_options.pb_cutoff,
                 field_summary=field_summary,
                 convol_mode="image",
-                convol_filter="-MFS-",
+                convol_filter=".MFS.",
                 convol_suffix_str=convol_suffix_str,
             )
         )