diff --git a/flint/naming.py b/flint/naming.py index 4730b783..34b831da 100644 --- a/flint/naming.py +++ b/flint/naming.py @@ -105,18 +105,25 @@ def create_name_from_common_fields( # TODO: Need to assess the mode argument, and define literals that are accepted def create_image_cube_name( - image_prefix: Path, mode: str, suffix: str = "cube.fits" + image_prefix: Path, + mode: Optional[Union[str, List[str]]] = None, + suffix: Optional[Union[str, List[str]]] = None, ) -> Path: """Create a consistent naming scheme when combining images into cube images. Intended to be used when combining many subband images together into a single cube. The name returned will be: - >> {image_prefix}.{mode}.{suffix} + >>> {image_prefix}.{mode}.{suffix}.cube.fits + + Should ``mode`` or ``suffix`` be a list, they will be joined with '.' separators. Hence, no + '.' should be added. + + This function will always output 'cube.fits' at the end of the returned file name. Args: image_prefix (Path): The unique path of the name. Generally this is the common part among the input planes - mode (str): Additional mode to add to the file name - suffix (str, optional): The final suffix to appended. Defaults to ".cube.fits". + mode (Optional[Union[str, List[str]]], optional): Additional mode/s to add to the file name. Defaults to None. + suffix (Optional[Union[str, List[str]]], optional): Additional suffix/s to add before the final 'cube.fits'. Defaults to None. Returns: Path: The final path and file name @@ -125,6 +132,23 @@ def create_image_cube_name( # it here for future proofing output_cube_name = f"{str(Path(image_prefix))}.{mode}.{suffix}" + output_components = [str(Path(image_prefix))] + if mode: + ( + output_components.append(mode) + if isinstance(mode, str) + else output_components.extend(mode) + ) + if suffix: + ( + output_components.append(suffix) + if isinstance(suffix, str) + else output_components.extend(suffix) + ) + + output_components.append("cube.fits") + + output_cube_name = ".".join(output_components) return Path(output_cube_name) diff --git a/flint/prefect/flows/subtract_cube_pipeline.py b/flint/prefect/flows/subtract_cube_pipeline.py index f354b910..90a04f46 100644 --- a/flint/prefect/flows/subtract_cube_pipeline.py +++ b/flint/prefect/flows/subtract_cube_pipeline.py @@ -187,13 +187,13 @@ def task_combine_all_linmos_images( images_to_combine = [ linmos_command.weight_fits for linmos_command in linmos_commands ] - output_suffix = "cube.weight.cube.fits" + output_suffix = "weight" else: logger.info("Combining image fits files") images_to_combine = [ linmos_command.image_fits for linmos_command in linmos_commands ] - output_suffix = "linmos.cube.fits" + output_suffix = "linmos" logger.info(f"Combining {len(images_to_combine)} FITS files together") diff --git a/tests/test_naming.py b/tests/test_naming.py index 5f1a66d7..856524cd 100644 --- a/tests/test_naming.py +++ b/tests/test_naming.py @@ -105,6 +105,33 @@ def test_create_image_cube_name(): "./57222/SB57222.RACS_1141-55.beam10.round3.i.residual.cube.fits" ) + name = create_image_cube_name( + image_prefix=Path("./57222/SB57222.RACS_1141-55.beam10.round3.i"), + mode=["residual", "pirate", "imaging"], + ) + assert isinstance(name, Path) + assert name == Path( + "./57222/SB57222.RACS_1141-55.beam10.round3.i.residual.pirate.imaging.cube.fits" + ) + name = create_image_cube_name( + image_prefix=Path("./57222/SB57222.RACS_1141-55.beam10.round3.i"), + mode=["residual", "pirate", "imaging"], + suffix="jackie", + ) + assert isinstance(name, Path) + assert name == Path( + "./57222/SB57222.RACS_1141-55.beam10.round3.i.residual.pirate.imaging.jackie.cube.fits" + ) + name = create_image_cube_name( + image_prefix=Path("./57222/SB57222.RACS_1141-55.beam10.round3.i"), + mode=["residual", "pirate", "imaging"], + suffix=["jackie", "boi"], + ) + assert isinstance(name, Path) + assert name == Path( + "./57222/SB57222.RACS_1141-55.beam10.round3.i.residual.pirate.imaging.jackie.boi.cube.fits" + ) + def test_get_beam_resolution_str(): """Map the known / support modes of beam resolution in file names"""