diff --git a/CHANGELOG.md b/CHANGELOG.md index c2823d0a..d3a12595 100644 --- a/CHANGELOG.md +++ b/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 diff --git a/flint/coadd/linmos.py b/flint/coadd/linmos.py index 71f68082..fc9091dd 100644 --- a/flint/coadd/linmos.py +++ b/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)}.") diff --git a/flint/imager/wsclean.py b/flint/imager/wsclean.py index fed182db..d8599dca 100644 --- a/flint/imager/wsclean.py +++ b/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) diff --git a/flint/leakage.py b/flint/leakage.py index 9fb930f3..9a414f41 100644 --- a/flint/leakage.py +++ b/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: diff --git a/flint/naming.py b/flint/naming.py index da0c05d3..a59e96da 100644 --- a/flint/naming.py +++ b/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[0-9]+)\.(?P.+)\.beam(?P[0-9]+)((\.spw(?P[0-9]+))?)((\.round(?P[0-9]+))?)((\.pol(?P[a-zA-z]+))?)*" + r"^SB(?P[0-9]+)\.(?P.+)\.beam(?P[0-9]+)((\.spw(?P[0-9]+))?)((\.round(?P[0-9]+))?)((\.(?P(i|q|u|v|xx|yy|xy|yx)+))?)*" ) results = regex.match(in_name) diff --git a/flint/prefect/common/imaging.py b/flint/prefect/common/imaging.py index e949601a..65691526 100644 --- a/flint/prefect/common/imaging.py +++ b/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, ) ) diff --git a/flint/prefect/flows/continuum_pipeline.py b/flint/prefect/flows/continuum_pipeline.py index 22b3f7d4..673421ed 100644 --- a/flint/prefect/flows/continuum_pipeline.py +++ b/flint/prefect/flows/continuum_pipeline.py @@ -288,7 +288,7 @@ def process_science_fields( field_options=field_options, field_summary=field_summary, current_round=None, - additional_linmos_suffix_str="poli", + additional_linmos_suffix_str="i", ) archive_wait_for.extend(parsets) parset = parsets[-1] @@ -395,7 +395,7 @@ def process_science_fields( field_options=field_options, field_summary=field_summary, current_round=current_round, - additional_linmos_suffix_str="poli", + additional_linmos_suffix_str="i", ) archive_wait_for.extend(parsets_self) @@ -438,7 +438,7 @@ def process_science_fields( current_round=( field_options.rounds if field_options.rounds else None ), - additional_linmos_suffix_str="polv", + additional_linmos_suffix_str="v", ) archive_wait_for.extend(parsets) diff --git a/tests/test_leakage.py b/tests/test_leakage.py index a0fbc8d1..06602756 100644 --- a/tests/test_leakage.py +++ b/tests/test_leakage.py @@ -11,6 +11,7 @@ FITSImage, LeakageFilters, PixelCoords, + _get_output_catalogue_path, _load_fits_image, _load_component_table, filter_components, @@ -19,6 +20,24 @@ from flint.utils import get_packaged_resource_path +def test_get_catalogue_output_path(): + """Make sure the name of the catalogue comes out as expected""" + catalogue = Path("SB1234.JACK_1234+56.i.round2_comp.fits") + pol = "v" + + output_path = _get_output_catalogue_path(input_path=catalogue, pol=pol) + assert output_path == Path("SB1234.JACK_1234+56.i.round2_comp.v_leakage.fits") + + spec_path = Path("JackSparrow_leakage.fits") + output_path = _get_output_catalogue_path( + input_path=catalogue, pol=pol, output_path=spec_path + ) + assert output_path == spec_path + + with pytest.raises(AssertionError): + _get_output_catalogue_path(input_path="notapathjack", pol=pol) + + def _get_aegean_catalogue(): table_path = get_packaged_resource_path( package="flint.data.tests", diff --git a/tests/test_naming.py b/tests/test_naming.py index ca8c2d9f..fa767c95 100644 --- a/tests/test_naming.py +++ b/tests/test_naming.py @@ -31,22 +31,22 @@ def test_create_image_cube_name(): """Put together a consistent file cube name""" name = create_image_cube_name( image_prefix=Path( - "/jack/sparrow/worst/pirate/flint_fitscube/57222/SB57222.RACS_1141-55.beam10.round3.poli" + "/jack/sparrow/worst/pirate/flint_fitscube/57222/SB57222.RACS_1141-55.beam10.round3.i" ), mode="image", ) assert isinstance(name, Path) assert name == Path( - "/jack/sparrow/worst/pirate/flint_fitscube/57222/SB57222.RACS_1141-55.beam10.round3.poli.image.cube.fits" + "/jack/sparrow/worst/pirate/flint_fitscube/57222/SB57222.RACS_1141-55.beam10.round3.i.image.cube.fits" ) name = create_image_cube_name( - image_prefix=Path("./57222/SB57222.RACS_1141-55.beam10.round3.poli"), + image_prefix=Path("./57222/SB57222.RACS_1141-55.beam10.round3.i"), mode="residual", ) assert isinstance(name, Path) assert name == Path( - "./57222/SB57222.RACS_1141-55.beam10.round3.poli.residual.cube.fits" + "./57222/SB57222.RACS_1141-55.beam10.round3.i.residual.cube.fits" ) @@ -389,7 +389,7 @@ def test_formatted_name_components_withpol(): assert components.round is None assert components.pol is None - ex = "SB39400.RACS_0635-31.beam33.poli-MFS-image.conv.fits" + ex = "SB39400.RACS_0635-31.beam33.i-MFS-image.conv.fits" components = processed_ms_format(in_name=ex) assert isinstance(components, ProcessedNameComponents) @@ -400,7 +400,7 @@ def test_formatted_name_components_withpol(): assert components.round is None assert components.pol == "i" - ex = "SB39400.RACS_0635-31.beam33.round2.poli-MFS-image.conv.fits" + ex = "SB39400.RACS_0635-31.beam33.round2.i-MFS-image.conv.fits" components = processed_ms_format(in_name=ex) assert isinstance(components, ProcessedNameComponents) @@ -411,7 +411,7 @@ def test_formatted_name_components_withpol(): assert components.round == "2" assert components.pol == "i" - ex = "SB39400.RACS_0635-31.beam33.round2.poliq-MFS-image.conv.fits" + ex = "SB39400.RACS_0635-31.beam33.round2.iq-MFS-image.conv.fits" components = processed_ms_format(in_name=ex) assert isinstance(components, ProcessedNameComponents) diff --git a/tests/test_wsclean.py b/tests/test_wsclean.py index 8ef86b9c..9f091248 100644 --- a/tests/test_wsclean.py +++ b/tests/test_wsclean.py @@ -3,6 +3,7 @@ import os import shutil from pathlib import Path +from typing import Any, Dict import pytest @@ -11,12 +12,15 @@ ImageSet, WSCleanCommand, WSCleanOptions, + _rename_wsclean_title, + _rename_wsclean_file, _resolve_wsclean_key_value_to_cli_str, _wsclean_output_callback, combine_subbands_to_cube, create_wsclean_cmd, create_wsclean_name_argument, get_wsclean_output_names, + rename_wsclean_prefix_in_imageset, ) from flint.ms import MS from flint.naming import create_imaging_name_prefix @@ -46,6 +50,119 @@ def set_env(): os.environ["LOCALDIR"] = "Pirates/be/here" +def test_rename_wsclean_path_move(tmpdir: Any): + """Rename the wsclean supplied part of a filename while moving a file""" + test_path = Path(tmpdir) / "move_file/" + test_path.mkdir(parents=True, exist_ok=True) + + ex = test_path / Path("SB39400.RACS_0635-31.beam33.poli-MFS-image.fits") + out_ex = test_path / Path("SB39400.RACS_0635-31.beam33.poli.MFS.image.fits") + + with open(ex, "w") as out_file: + out_file.write("example") + + assert ex.exists() + assert not out_ex.exists() + assert _rename_wsclean_file(input_path=ex, rename_file=True) == out_ex + assert not ex.exists() + assert out_ex.exists() + + +def _write_test_image(items: Any): + for item in items: + with Path(item).open("w") as out_file: + out_file.write(str(item)) + + +def test_rename_wsclean_imageset(tmpdir: Any): + """Ensure that items described in an image set are able to be properly renamed""" + + test_dir = Path(tmpdir) / "imagesetrename" + test_dir.mkdir(parents=True, exist_ok=True) + + # create some test files and ensure they all exist + keys: Dict[Any, Any] = {} + prefix = f"{str(test_dir)}/SB39400.RACS_0635-31.beam33.i" + keys["prefix"] = prefix + for mode in ("image", "residual"): + items = [ + Path(f"{prefix}-{subband:04d}-{mode}.fits") for subband in range(4) + ] + [Path(f"{prefix}-MFS-{mode}.fits")] + _write_test_image(items=items) + keys[mode] = items + assert all([Path(f).exists() for f in items]) + + # form the image set that will have the wsclean appended properties string renamed + image_set = ImageSet(**keys) + assert isinstance(image_set, ImageSet) + new_image_set = rename_wsclean_prefix_in_imageset(input_imageset=image_set) + + # test to see thhat files exists + assert new_image_set.prefix == prefix + assert new_image_set.image is not None + assert all([file.exists() for file in new_image_set.image]) + assert new_image_set.residual is not None + assert all([file.exists() for file in new_image_set.residual]) + + # and ensure the originals no longer exist + assert all([not Path(file).exists() for file in keys["image"]]) + assert all([not (file).exists() for file in keys["residual"]]) + + +def test_rename_wsclean_path(): + """Rename the wsclean supplied part of a filename""" + + ex = Path("SB39400.RACS_0635-31.beam33.poli-MFS-image.fits") + out_ex = Path("SB39400.RACS_0635-31.beam33.poli.MFS.image.fits") + assert _rename_wsclean_file(input_path=ex) == out_ex + + ex = Path("SB39400.RACS_0635-31.beam33.poli-MFS-image") + out_ex = Path("SB39400.RACS_0635-31.beam33.poli.MFS.image") + assert _rename_wsclean_file(input_path=ex) == out_ex + + ex = Path("/a/path/that/is/a/parent/SB39400.RACS_0635-31.beam33.poli-MFS-image") + out_ex = Path("/a/path/that/is/a/parent/SB39400.RACS_0635-31.beam33.poli.MFS.image") + assert _rename_wsclean_file(input_path=ex) == out_ex + + +def test_regex_rename_wsclean_title(): + """Rename the wsclean supplied using regex""" + + ex = "SB39400.RACS_0635-31.beam33.poli-MFS-image.fits" + out_ex = "SB39400.RACS_0635-31.beam33.poli.MFS.image.fits" + assert _rename_wsclean_title(name_str=ex) == out_ex + + ex = "SB39400.RACS_0635-31.beam33.poli-MFS-image" + out_ex = "SB39400.RACS_0635-31.beam33.poli.MFS.image" + assert _rename_wsclean_title(name_str=ex) == out_ex + + ex = "SB39400.RACS_0635-31.beam33.poli-MFS-image" + out_ex = "SB39400.RACS_0635.31.beam33.poli.MFS.image" + assert not _rename_wsclean_title(name_str=ex) == out_ex + + ex = "SB39400.RACS_0635-31.beam33.poli.MFS.image.fits" + out_ex = "SB39400.RACS_0635-31.beam33.poli.MFS.image.fits" + assert _rename_wsclean_title(name_str=ex) == out_ex + assert _rename_wsclean_title(name_str=ex) is ex + + ex = "SB39400.RACS_0635-31.beam33.poli-i-MFS-image" + out_ex = "SB39400.RACS_0635-31.beam33.poli.i.MFS.image" + assert _rename_wsclean_title(name_str=ex) == out_ex + + +def test_regex_stokes_wsclean_title(): + """Test whether all stokes values are picked up properly""" + + prefix = "SB39400.RACS_0635-31.beam33.poli." + end = "-MFS-image.fits" + transformed = end.replace("-", ".") + + for stokes in ("i", "q", "u", "v", "xx", "xy", "yx", "yy"): + ex = f"{prefix}-{stokes}{end}" + out_ex = f"{prefix}.{stokes}{transformed}" + assert _rename_wsclean_title(name_str=ex) == out_ex + + def test_combine_subbands_to_cube(tmpdir): """Load in example fits images to combine into a cube""" files = [ @@ -120,7 +237,7 @@ def test_create_wsclean_name(ms_example): for pol in ("i", "I"): name = create_imaging_name_prefix(ms=ms_example, pol=pol) - assert name == "SB39400.RACS_0635-31.beam0.small.poli" + assert name == "SB39400.RACS_0635-31.beam0.small.i" def test_create_wsclean_name_argument(ms_example): @@ -134,16 +251,14 @@ def test_create_wsclean_name_argument(ms_example): parent = str(Path(ms_example).parent) assert isinstance(name_argument_path, Path) - assert f"{parent}/SB39400.RACS_0635-31.beam0.small.poli" == str(name_argument_path) + assert f"{parent}/SB39400.RACS_0635-31.beam0.small.i" == str(name_argument_path) wsclean_options_2 = WSCleanOptions(temp_dir="/jack/sparrow") name_argument_path = create_wsclean_name_argument( wsclean_options=wsclean_options_2, ms=ms ) - assert "/jack/sparrow/SB39400.RACS_0635-31.beam0.small.poli" == str( - name_argument_path - ) + assert "/jack/sparrow/SB39400.RACS_0635-31.beam0.small.i" == str(name_argument_path) def test_create_wsclean_command(ms_example):