From 590c372ba63e57597531dac5e4534bb30593e520 Mon Sep 17 00:00:00 2001 From: Alec Thomson Date: Tue, 28 Jan 2025 15:22:29 +0800 Subject: [PATCH] Tjgalvin beams (#76) * mock up of target beams * beams mock / added to cli * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * added a test on _get_target_beam * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * corrected inverted check for nchans * typo and assert fixes * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * added checks for nan in target beam props / tests * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Ruff * Ruff --------- Co-authored-by: tgalvin Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: Tim Galvin --- .pre-commit-config.yaml | 32 ++-- README.md | 2 +- pyproject.toml | 32 +++- racs_tools/au2.py | 8 +- racs_tools/beamcon_2D.py | 8 +- racs_tools/beamcon_3D.py | 306 ++++++++++++++++++++++++------------ racs_tools/convolve_uv.py | 19 ++- racs_tools/gaussft.py | 3 +- racs_tools/getnoise_list.py | 16 +- racs_tools/logging.py | 5 +- racs_tools/parallel.py | 7 +- tests/hellompi.py | 3 +- tests/test_2d.py | 24 ++- tests/test_3d.py | 61 +++++-- tests/test_noise.py | 1 - 15 files changed, 348 insertions(+), 179 deletions(-) diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index c0882aa..c274332 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -1,24 +1,22 @@ repos: -- repo: https://github.com/pre-commit/pre-commit-hooks - rev: v4.6.0 - hooks: - - id: trailing-whitespace - - id: end-of-file-fixer - - id: check-yaml - - id: check-added-large-files -- repo: https://github.com/psf/black - rev: 24.4.2 - hooks: - - id: black -- repo: https://github.com/PyCQA/isort - rev: 5.13.2 - hooks: - - id: isort - args: ["--profile=black"] + - repo: https://github.com/pre-commit/pre-commit-hooks + rev: v4.6.0 + hooks: + - id: trailing-whitespace + - id: end-of-file-fixer + - id: check-yaml + - id: check-added-large-files + + - repo: https://github.com/astral-sh/ruff-pre-commit + rev: "v0.9.2" + hooks: + - id: ruff + args: ["--fix", "--show-fixes"] + - id: ruff-format + ci: autofix_commit_msg: | [pre-commit.ci] auto fixes from pre-commit.com hooks - for more information, see https://pre-commit.ci autofix_prs: true autoupdate_branch: '' diff --git a/README.md b/README.md index 29e465b..76feeb7 100644 --- a/README.md +++ b/README.md @@ -1,4 +1,4 @@ -[![Docker Build and Push](https://github.com/AlecThomson/RACS-tools/actions/workflows/docker.yml/badge.svg)](https://github.com/AlecThomson/RACS-tools/actions/workflows/docker.yml) ![Tests](https://github.com/AlecThomson/RACS-tools/actions/workflows/pytest.yml/badge.svg) [![Python package](https://github.com/AlecThomson/RACS-tools/actions/workflows/python-package.yml/badge.svg)](https://github.com/AlecThomson/RACS-tools/actions/workflows/python-package.yml) [![PyPi](https://github.com/AlecThomson/RACS-tools/actions/workflows/pypi.yml/badge.svg)](https://github.com/AlecThomson/RACS-tools/actions/workflows/pypi.yml) [![Code style: black](https://img.shields.io/badge/code%20style-black-000000.svg)](https://github.com/psf/black)[![pre-commit.ci status](https://results.pre-commit.ci/badge/github/AlecThomson/RACS-tools/master.svg)](https://results.pre-commit.ci/latest/github/AlecThomson/RACS-tools/master) +[![Docker Build and Push](https://github.com/AlecThomson/RACS-tools/actions/workflows/docker.yml/badge.svg)](https://github.com/AlecThomson/RACS-tools/actions/workflows/docker.yml) ![Tests](https://github.com/AlecThomson/RACS-tools/actions/workflows/pytest.yml/badge.svg) [![Python package](https://github.com/AlecThomson/RACS-tools/actions/workflows/python-package.yml/badge.svg)](https://github.com/AlecThomson/RACS-tools/actions/workflows/python-package.yml) [![PyPi](https://github.com/AlecThomson/RACS-tools/actions/workflows/pypi.yml/badge.svg)](https://github.com/AlecThomson/RACS-tools/actions/workflows/pypi.yml)[![pre-commit.ci status](https://results.pre-commit.ci/badge/github/AlecThomson/RACS-tools/master.svg)](https://results.pre-commit.ci/latest/github/AlecThomson/RACS-tools/master) # RACS-tools Useful scripts for RACS diff --git a/pyproject.toml b/pyproject.toml index b648640..a4937f3 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -19,8 +19,7 @@ numba = "*" mpi4py = {version = ">=3", optional = true} [tool.poetry.dev-dependencies] -black = "*" -isort = "*" +ruff = {version = "*", optional = true} [tool.poetry.extras] mpi = ["mpi4py"] @@ -38,3 +37,32 @@ vcs = "git" [build-system] requires = ["poetry-core>=1.0.0", "poetry-dynamic-versioning>=1.0.0,<2.0.0"] build-backend = "poetry_dynamic_versioning.backend" + +[tool.ruff] +src = ["racs_tools"] + +[tool.ruff.lint] +extend-select = [ + # "B", # flake8-bugbear + "I", # isort + "ARG", # flake8-unused-arguments + # "C4", # flake8-comprehensions + # "EM", # flake8-errmsg + "ICN", # flake8-import-conventions + # "G", # flake8-logging-format + # "PGH", # pygrep-hooks + # "PIE", # flake8-pie + # "PL", # pylint + # "PT", # flake8-pytest-style + # "PTH", # flake8-use-pathlib + # "RET", # flake8-return + # "RUF", # Ruff-specific + # "SIM", # flake8-simplify + # "T20", # flake8-print + "UP", # pyupgrade + # "YTT", # flake8-2020 + # "EXE", # flake8-executable + # "NPY", # NumPy specific rules + # "PD", # pandas-vet, + # "RET", # flake8-return +] diff --git a/racs_tools/au2.py b/racs_tools/au2.py index 49a76b0..d910db7 100644 --- a/racs_tools/au2.py +++ b/racs_tools/au2.py @@ -1,9 +1,9 @@ #!/usr/bin/env python3 -""" For getting fluxes right in Jy/beam """ +"""For getting fluxes right in Jy/beam""" + __author__ = "Tessa Vernstrom" import math -from typing import List, Tuple import numpy as np @@ -75,8 +75,8 @@ def gaussianDeconvolve(smaj, smin, spa, bmaj, bmin, bpa): def gauss_factor( - beamConv: List[float], beamOrig: List[float], dx1: float = 1, dy1: float = 1 -) -> Tuple[float, float, float, float, float]: + beamConv: list[float], beamOrig: list[float], dx1: float = 1, dy1: float = 1 +) -> tuple[float, float, float, float, float]: """ Calculates the scaling factor to be applied after convolving a map in Jy/beam with a gaussian to get fluxes in Jy/beam again. diff --git a/racs_tools/beamcon_2D.py b/racs_tools/beamcon_2D.py index 04c655c..c1b836b 100644 --- a/racs_tools/beamcon_2D.py +++ b/racs_tools/beamcon_2D.py @@ -1,5 +1,6 @@ #!/usr/bin/env python -""" Convolve ASKAP images to common resolution """ +"""Convolve ASKAP images to common resolution""" + from __future__ import annotations __author__ = "Alec Thomson" @@ -7,9 +8,8 @@ import logging import sys -import traceback from pathlib import Path -from typing import Literal, NamedTuple, Optional +from typing import Literal, NamedTuple import numpy as np from astropy import units as u @@ -290,7 +290,7 @@ def beamcon_2d_on_fits( def get_common_beam( files: list[Path], conv_mode: Literal["robust", "scipy", "astropy", "astropy_fft"] = "robust", - target_beam: Optional[Beam] = None, + target_beam: Beam | None = None, cutoff: float | None = None, tolerance: float = 0.0001, nsamps: float = 200, diff --git a/racs_tools/beamcon_3D.py b/racs_tools/beamcon_3D.py index 30fb3a2..58fcc6f 100644 --- a/racs_tools/beamcon_3D.py +++ b/racs_tools/beamcon_3D.py @@ -1,5 +1,6 @@ #!/usr/bin/env python3 -""" Convolve ASKAP cubes to common resolution """ +"""Convolve ASKAP cubes to common resolution""" + from __future__ import annotations __author__ = "Alec Thomson" @@ -134,45 +135,44 @@ def get_beams(file: Path, header: fits.Header) -> CubeBeams: beams = Table.read(hdu) # Otherwise use beamlog file + elif beamlog.exists(): + logger.info("No CASA beamtable found in header - looking for beamlogs") + logger.info(f"Getting beams from {beamlog}") + + beams = Table.read(beamlog, format="ascii.commented_header") + # Header looks like: + # colnames=['Channel', 'BMAJarcsec', 'BMINarcsec', 'BPAdeg'] + # But needs some fix up - astropy can't read the header properly + for col in beams.colnames: + idx = col.find("[") + if idx == -1: + new_col = col + unit = u.Unit("") + else: + new_col = col[:idx] + unit = u.Unit(col[idx + 1 : -1]) + beams[col].unit = unit + beams[col].name = new_col else: - if beamlog.exists(): - logger.info("No CASA beamtable found in header - looking for beamlogs") - logger.info(f"Getting beams from {beamlog}") - - beams = Table.read(beamlog, format="ascii.commented_header") - # Header looks like: - # colnames=['Channel', 'BMAJarcsec', 'BMINarcsec', 'BPAdeg'] - # But needs some fix up - astropy can't read the header properly - for col in beams.colnames: - idx = col.find("[") - if idx == -1: - new_col = col - unit = u.Unit("") - else: - new_col = col[:idx] - unit = u.Unit(col[idx + 1 : -1]) - beams[col].unit = unit - beams[col].name = new_col - else: - logger.warning("No beamlog found") - logger.warning("Using header beam - assuming a constant beam") - try: - beam: Beam = Beam.from_fits_header(header) - wcs = WCS(header) - spec_axis = wcs.spectral - _nchan = spec_axis.array_shape[0] - beams = Table() - beams.add_column(np.arange(_nchan), name="Channel") - beams.add_column([beam.major.to(u.arcsec)] * _nchan, name="BMAJ") - beams.add_column( - [beam.minor.to(u.arcsec)] * _nchan, - name="BMIN", - ) - beams.add_column([beam.pa.to(u.deg)] * _nchan, name="BPA") + logger.warning("No beamlog found") + logger.warning("Using header beam - assuming a constant beam") + try: + beam: Beam = Beam.from_fits_header(header) + wcs = WCS(header) + spec_axis = wcs.spectral + _nchan = spec_axis.array_shape[0] + beams = Table() + beams.add_column(np.arange(_nchan), name="Channel") + beams.add_column([beam.major.to(u.arcsec)] * _nchan, name="BMAJ") + beams.add_column( + [beam.minor.to(u.arcsec)] * _nchan, + name="BMIN", + ) + beams.add_column([beam.pa.to(u.deg)] * _nchan, name="BPA") - except Exception as e: - logger.error("Couldn't get beam from header") - raise e + except Exception as e: + logger.error("Couldn't get beam from header") + raise e nchan = len(beams) return CubeBeams(beams, nchan, beamlog) @@ -252,7 +252,7 @@ def smooth_plane( slicer = tuple(slicer) plane = cube.unmasked_data[slicer].value.astype(np.float32) - logger.debug(f"Size of plane is {(plane.nbytes*u.byte).to(u.MB)}") + logger.debug(f"Size of plane is {(plane.nbytes * u.byte).to(u.MB)}") if mask: logger.info(f"Masking channel {idx}") @@ -316,41 +316,21 @@ def make_data(files: list[Path], outdir: list[Path]) -> list[CubeData]: return cube_data_list -def commonbeamer( +def _get_commonbeams( cube_data_list: list[CubeData], nchans: int, conv_mode: Literal["robust", "scipy", "astropy", "astropy_fft"] = "robust", mode: Literal["natural", "total"] = "natural", - suffix: str | None = None, target_beam: Beam | None = None, circularise: bool = False, tolerance: float = 0.0001, nsamps: int = 200, epsilon: float = 0.0005, -) -> list[CommonBeamData]: - """Find common beam for all channels. - Computed beams will be written to convolving beam logger. - - Args: - cube_data_list (list[CubeData]): list of cube data. - nchans (int): Number of channels. - conv_mode (Literal["robust", "scipy", "astropy", "astropy_fft"], optional): Convolution mode. Defaults to "robust". - mode (Literal["natural", "total"], optional): Frequency mode. Defaults to "natural". - target_beam (Beam, optional): Target PSF. Defaults to None. - circularise (bool, optional): Circularise PSF. Defaults to False. - tolerance (float, optional): Common beam tolerance. Defaults to 0.0001. - nsamps (int, optional): Common beam samples. Defaults to 200. - epsilon (float, optional): Common beam epsilon. Defaults to 0.0005. - - Raises: - Exception: If convolving beam will be undersampled on pixel grid. +) -> Beams: + assert isinstance(target_beam, Beam) or target_beam is None, ( + f"Expected target_beam to be type Beam or None, got {type(target_beam)}" + ) - Returns: - list[CommonBeamData]: Common beam data for each channel and cube. - """ - if suffix is None: - suffix = mode - ### Natural mode ### if mode == "natural": big_beams = [] for n in trange( @@ -453,8 +433,8 @@ def commonbeamer( f"Smallest common Nyquist sampled beam is: {nyq_beam!r}" ) - logger.warn("COMMON BEAM WILL BE UNDERSAMPLED!") - logger.warn("SETTING COMMON BEAM TO NYQUIST BEAM") + logger.warning("COMMON BEAM WILL BE UNDERSAMPLED!") + logger.warning("SETTING COMMON BEAM TO NYQUIST BEAM") commonbeam = nyq_beam bmaj_common.append(commonbeam.major.to(u.arcsec).value) @@ -500,12 +480,13 @@ def commonbeamer( tolerance=tolerance, nsamps=nsamps, epsilon=epsilon ) except BeamError: - logger.warn("Couldn't find common beam with defaults") - logger.warn("Trying again with smaller tolerance") + logger.warning("Couldn't find common beam with defaults") + logger.warning("Trying again with smaller tolerance") commonbeam = big_beams[~np.isnan(big_beams)].common_beam( tolerance=tolerance * 0.1, nsamps=nsamps, epsilon=epsilon ) + if target_beam is not None: commonbeam = target_beam else: @@ -543,11 +524,11 @@ def commonbeamer( if target_beam is not None: commonbeam = target_beam if target_beam < nyq_beam: - logger.warn("TARGET BEAM WILL BE UNDERSAMPLED!") + logger.warning("TARGET BEAM WILL BE UNDERSAMPLED!") raise Exception("CAN'T UNDERSAMPLE BEAM - EXITING") else: - logger.warn("COMMON BEAM WILL BE UNDERSAMPLED!") - logger.warn("SETTING COMMON BEAM TO NYQUIST BEAM") + logger.warning("COMMON BEAM WILL BE UNDERSAMPLED!") + logger.warning("SETTING COMMON BEAM TO NYQUIST BEAM") commonbeam = nyq_beam # Make Beams object @@ -565,6 +546,61 @@ def commonbeamer( pa=commonbeams.pa * 0, ) + return commonbeams + + +def commonbeamer( + cube_data_list: list[CubeData], + nchans: int, + conv_mode: Literal["robust", "scipy", "astropy", "astropy_fft"] = "robust", + mode: Literal["natural", "total"] = "natural", + suffix: str | None = None, + target_beam: Beam | Beams | None = None, + circularise: bool = False, + tolerance: float = 0.0001, + nsamps: int = 200, + epsilon: float = 0.0005, +) -> list[CommonBeamData]: + """Find common beam for all channels. + Computed beams will be written to convolving beam logger. + + Args: + cube_data_list (List[CubeData]): List of cube data. + nchans (int): Number of channels. + conv_mode (Literal["robust", "scipy", "astropy", "astropy_fft"], optional): Convolution mode. Defaults to "robust". + mode (Literal["natural", "total"], optional): Frequency mode. Defaults to "natural". + target_beam (Optional[Union[Beam,List[Beam]]], optional): Target PSF. If list of beams, each corresponds to a channel resolution. Defaults to None. + circularise (bool, optional): Circularise PSF. Defaults to False. + tolerance (float, optional): Common beam tolerance. Defaults to 0.0001. + nsamps (int, optional): Common beam samples. Defaults to 200. + epsilon (float, optional): Common beam epsilon. Defaults to 0.0005. + + Raises: + Exception: If convolving beam will be undersampled on pixel grid. + + Returns: + List[CommonBeamData]: Common beam data for each channel and cube. + """ + if suffix is None: + suffix = mode + ### Natural mode ### + + if isinstance(target_beam, Beams): + logger.info(f"Setting the target beam for {len(target_beam)} channels") + commonbeams = target_beam + else: + commonbeams = _get_commonbeams( + cube_data_list=cube_data_list, + nchans=nchans, + conv_mode=conv_mode, + mode=mode, + target_beam=target_beam, + circularise=circularise, + tolerance=tolerance, + nsamps=nsamps, + epsilon=epsilon, + ) + logger.info("Final beams are:") for i, commonbeam in enumerate(commonbeams): logger.info(f"Channel {i}: {commonbeam!r}") @@ -731,9 +767,9 @@ def initfiles( spec_axis = wcs.spectral crpix = int(spec_axis.wcs.crpix) nchans = spec_axis.array_shape[0] - assert nchans == len( - commonbeams - ), "Number of channels in header and commonbeams do not match" + assert nchans == len(commonbeams), ( + "Number of channels in header and commonbeams do not match" + ) chans = np.arange(nchans) if ref_chan is None: # Get reference channel, and attach PSF there @@ -911,6 +947,71 @@ def smooth_and_write_plane( logger.info(f"{outfile} - channel {chan} - Done") +def _get_target_beam( + bmaj: float | list[float] | None = None, + bmin: float | list[float] | None = None, + bpa: float | list[float] | None = None, +) -> None | Beam | Beams: + """Appropriately handle the input target beam specification + + Args: + bmaj (Optional[Union[float,List[float]]], optional): Target beam major axis in arcsec. If a list, this should be the same length as the number of channels to smooth. Defaults to None. + bmin (Optional[Union[float,List[float]]], optional): Target beam minor axis in arcsec. If a list, this should be the same length as the number of channels to smooth. Defaults to None. + bpa (Optional[Union[float,List[float]]], optional): Target beam position angle in deg. If a list, this should be the same length as the number of channels to smooth. Defaults to None. + + Raises: + ValueError: Occurs if only subset of beam parameters specified + + Returns: + Union[None, Beam, Beams]: The appropriate set of processed beams + """ + + nonetest = [param is None for param in (bmaj, bmin, bpa)] + if all(nonetest): + return None + + if any(nonetest): + raise ValueError("Please specify all target beam params!") + + if all([isinstance(b, float) for b in (bmaj, bmin, bpa)]): + if any(~np.isfinite((bmaj, bmin, bpa))): + logger.debug( + f"Non-finite value detected in beam ({bmaj, bmin, bpa}). Setting all to nan." + ) + bmaj = bmin = bpa = float("nan") + + target_beam = Beam(bmaj * u.arcsec, bmin * u.arcsec, bpa * u.deg) + logger.info(f"Target beam is {target_beam!r}") + return target_beam + else: + assert all([isinstance(b, list) for b in (bmaj, bmin, bpa)]), ( + "Something is not a list" + ) + assert len(bmaj) == len(bmin) == len(bpa), "Unequal target beam lengths" + + # Deal with moments where, potentially, a beam has a NaN but not all NaN. + # Maybe unnecessary. + for idx, beam_props in enumerate(zip(bmaj, bmin, bpa)): + if any(~np.isfinite(beam_props)): + logger.debug( + f"Non-finite beam property detected in {beam_props=} at {idx=}. Setting all to nan. " + ) + bmaj[idx] = float("nan") + bmin[idx] = float("nan") + bpa[idx] = float("nan") + + target_beams = Beams( + np.array(bmaj) * u.arcsec, + np.array(bmin) * u.arcsec, + np.array(bpa) * u.degree, + ) + + logger.info(f"{len(target_beams)} target beams have been constructed") + return target_beams + + raise ValueError("Beam parameters were not successfully processed") + + def smooth_fits_cube( infiles_list: list[Path], uselogs: bool = False, @@ -920,9 +1021,9 @@ def smooth_fits_cube( prefix: str | None = None, suffix: str | None = None, outdir: Path | None = None, - bmaj: float | None = None, - bmin: float | None = None, - bpa: float | None = None, + bmaj: float | list[float] | None = None, + bmin: float | list[float] | None = None, + bpa: float | list[float] | None = None, cutoff: float | None = None, circularise: bool = False, ref_chan: int | None = None, @@ -943,11 +1044,11 @@ def smooth_fits_cube( dryrun (bool, optional): Do not write any images. Defaults to False. prefix (str, optional): Output filename prefix. Defaults to None. suffix (str, optional): Output filename suffix. Defaults to None. - outdir (Path | None, optional): Output directory. Defaults to None. - bmaj (float | None, optional): Target beam major axis in arcsec. Defaults to None. - bmin (float | None, optional): Target beam minor axis in arcsec. Defaults to None. - bpa (float | None, optional): Target beam position angle in deg. Defaults to None. - cutoff (float | None, optional): Beam cutoff in arcsec. Defaults to None. + outdir (Optional[Path], optional): Output directory. Defaults to None. + bmaj (Optional[Union[float,List[float]]], optional): Target beam major axis in arcsec. If a list, this should be the same length as the number of channels to smooth. Defaults to None. + bmin (Optional[Union[float,List[float]]], optional): Target beam minor axis in arcsec. If a list, this should be the same length as the number of channels to smooth. Defaults to None. + bpa (Optional[Union[float,List[float]]], optional): Target beam position angle in deg. If a list, this should be the same length as the number of channels to smooth. Defaults to None. + cutoff (Optional[float], optional): Beam cutoff in arcsec. Defaults to None. circularise (bool, optional): Set minor axis to major axis. Defaults to False. ref_chan (int | None, optional): Reference channel for PSF in header. Defaults to None. tolerance (float, optional): Radio beam tolerance. Defaults to 0.0001. @@ -990,14 +1091,7 @@ def smooth_fits_cube( # Check target conv_mode = parse_conv_mode(conv_mode) - nonetest = [param is None for param in (bmaj, bmin, bpa)] - if all(nonetest): - target_beam = None - elif any(nonetest): - raise ValueError("Please specify all target beam params!") - else: - target_beam = Beam(bmaj * u.arcsec, bmin * u.arcsec, bpa * u.deg) - logger.info(f"Target beam is {target_beam!r}") + target_beam = _get_target_beam(bmaj=bmaj, bmin=bmin, bpa=bpa) files = sorted(infiles_list) if len(files) == 0: @@ -1011,11 +1105,14 @@ def smooth_fits_cube( # Sanity check channel counts nchans = np.array([cube_data.nchan for cube_data in cube_data_list]) - check = all(nchans == nchans[0]) - - if not check: + if not all(nchans == nchans[0]): raise ValueError(f"Unequal number of spectral channels! Got {nchans}") + if isinstance(target_beam, Beams) and any(nchans != len(target_beam)): + raise ValueError( + f"Unequal length of target beams ({len(target_beam)}) to channels ({nchans})" + ) + nchans = nchans[0] # Check suffix @@ -1215,7 +1312,8 @@ def cli(): dest="bmaj", type=float, default=None, - help="BMAJ to convolve to [max BMAJ from given image(s)].", + nargs="+", + help="BMAJ to convolve to [max BMAJ from given image(s)]. If multiple are give they will be matched to each channel image.", ) parser.add_argument( @@ -1223,11 +1321,17 @@ def cli(): dest="bmin", type=float, default=None, - help="BMIN to convolve to [max BMAJ from given image(s)].", + nargs="+", + help="BMIN to convolve to [max BMAJ from given image(s)] .If multiple are give they will be matched to each channel image.", ) parser.add_argument( - "--bpa", dest="bpa", type=float, default=None, help="BPA to convolve to [0]." + "--bpa", + dest="bpa", + type=float, + default=None, + nargs="+", + help="BPA to convolve to [0]. If multiple are give they will be matched to each channel image.", ) parser.add_argument( @@ -1309,6 +1413,10 @@ def cli(): verbosity=args.verbosity, ) + bmaj = args.bmaj if args.bmaj is None or len(args.bmaj) > 1 else args.bmaj[0] + bmin = args.bmin if args.bmin is None or len(args.bmin) > 1 else args.bmin[0] + bpa = args.bpa if args.bpa is None or len(args.bpa) > 1 else args.bpa[0] + _ = smooth_fits_cube( infiles_list=args.infile, uselogs=args.uselogs, @@ -1318,9 +1426,9 @@ def cli(): prefix=args.prefix, suffix=args.suffix, outdir=args.outdir, - bmaj=args.bmaj, - bmin=args.bmin, - bpa=args.bpa, + bmaj=bmaj, + bmin=bmin, + bpa=bpa, cutoff=args.cutoff, circularise=args.circularise, ref_chan=args.ref_chan, diff --git a/racs_tools/convolve_uv.py b/racs_tools/convolve_uv.py index 48c25e5..0c90001 100644 --- a/racs_tools/convolve_uv.py +++ b/racs_tools/convolve_uv.py @@ -1,14 +1,14 @@ #!/usr/bin/env python3 -""" Fast convolution in the UV domain """ +"""Fast convolution in the UV domain""" + __author__ = "Wasim Raja" import gc -from typing import Literal, NamedTuple, Optional, Tuple +from typing import Literal, NamedTuple, Optional -import astropy.units as units import numpy as np import scipy.signal -from astropy import convolution +from astropy import convolution, units from astropy import units as u from astropy.io import fits from astropy.wcs import WCS @@ -16,8 +16,7 @@ from radio_beam import Beam, Beams from radio_beam.utils import BeamError -import racs_tools.gaussft as gaussft -from racs_tools import au2 +from racs_tools import au2, gaussft from racs_tools.logging import logger ZERO_BEAM = Beam(0 * u.deg) @@ -184,7 +183,7 @@ def convolve( # Create a mask for the NaNs mask = np.isnan(image).astype(np.int16) logger.warning( - f"Image contains {mask.sum()} ({mask.sum()/ mask.size *100 :0.1f}%) NaNs" + f"Image contains {mask.sum()} ({mask.sum() / mask.size * 100:0.1f}%) NaNs" ) image = np.nan_to_num(image) @@ -236,7 +235,7 @@ def convolve( # Need this to get around numerical issues mask_conv = ~(mask_conv + 1 < 2) logger.warning( - f"Convolved image contains {mask_conv.sum()} ({mask_conv.sum()/ mask_conv.size *100 :0.1f}%) NaNs" + f"Convolved image contains {mask_conv.sum()} ({mask_conv.sum() / mask_conv.size * 100:0.1f}%) NaNs" ) im_conv[mask_conv > 0] = np.nan @@ -414,7 +413,7 @@ def parse_conv_mode( str: Convolution mode. """ logger.info(f"Convolution mode: {conv_mode}") - if conv_mode not in CONVOLUTION_FUNCTIONS.keys(): + if conv_mode not in CONVOLUTION_FUNCTIONS: raise ValueError( f"Please select valid convolution method! Expected one of {list(CONVOLUTION_FUNCTIONS.keys())}, got {conv_mode}" ) @@ -436,7 +435,7 @@ def get_convolving_beam( dx: u.Quantity, dy: u.Quantity, cutoff: Optional[float] = None, -) -> Tuple[Beam, float]: +) -> tuple[Beam, float]: """Get the beam to use for smoothing Args: diff --git a/racs_tools/gaussft.py b/racs_tools/gaussft.py index 8ce660b..43d930e 100644 --- a/racs_tools/gaussft.py +++ b/racs_tools/gaussft.py @@ -6,7 +6,6 @@ Python version of gaussft.f by Wasim Raja """ -from typing import Tuple import numba as nb import numpy as np @@ -22,7 +21,7 @@ def gaussft( bpa: float, u: np.ndarray, v: np.ndarray, -) -> Tuple[np.ndarray, float]: +) -> tuple[np.ndarray, float]: """ Compute the Fourier transform of a 2D Gaussian for convolution. diff --git a/racs_tools/getnoise_list.py b/racs_tools/getnoise_list.py index 626c798..3548b4d 100644 --- a/racs_tools/getnoise_list.py +++ b/racs_tools/getnoise_list.py @@ -1,9 +1,9 @@ #!/usr/bin/env python3 -""" Find bad channels by checking statistics of each channel image. """ +"""Find bad channels by checking statistics of each channel image.""" import argparse import warnings -from typing import Tuple, Union +from typing import Union import astropy.units as u import numpy as np @@ -39,7 +39,7 @@ def getbadchans( qcube: SpectralCube, ucube: SpectralCube, cliplev: float = 5, -) -> Tuple[np.ndarray, u.Quantity, u.Quantity]: +) -> tuple[np.ndarray, u.Quantity, u.Quantity]: """Find bad channels in Stokes Q and U cubes Args: @@ -96,7 +96,7 @@ def getbadchans( def blankchans( qcube: SpectralCube, ucube: SpectralCube, totalbad: np.ndarray, blank: bool = False -) -> Tuple[SpectralCube, SpectralCube]: +) -> tuple[SpectralCube, SpectralCube]: """Mask out bad channels Args: @@ -162,13 +162,13 @@ def main( qcube = getcube(qfile) ucube = getcube(ufile) - assert len(ucube.spectral_axis) == len( - qcube.spectral_axis - ), "Cubes have different number of channels" + assert len(ucube.spectral_axis) == len(qcube.spectral_axis), ( + "Cubes have different number of channels" + ) # Iterate for i in range(iterate): - logger.info(f"Flagging iteration {i+1} of {iterate}") + logger.info(f"Flagging iteration {i + 1} of {iterate}") totalbad, qnoisevals, unoisevals = getbadchans( qcube, ucube, diff --git a/racs_tools/logging.py b/racs_tools/logging.py index 2fe1693..f5e33ce 100644 --- a/racs_tools/logging.py +++ b/racs_tools/logging.py @@ -1,10 +1,9 @@ #!/usr/bin/env python3 -# -*- coding: utf-8 -*- import logging import multiprocessing as mp from logging.handlers import QueueHandler, QueueListener -from typing import Optional, Tuple +from typing import Optional logging.captureWarnings(True) @@ -14,7 +13,7 @@ def setup_logger( filename: Optional[str] = None, -) -> Tuple[logging.Logger, QueueListener, mp.Queue]: +) -> tuple[logging.Logger, QueueListener, mp.Queue]: """Setup a logger Args: diff --git a/racs_tools/parallel.py b/racs_tools/parallel.py index e7daef6..ff883af 100644 --- a/racs_tools/parallel.py +++ b/racs_tools/parallel.py @@ -1,7 +1,8 @@ #!/usr/bin/env python3 -""" Utilities for parallelism """ +"""Utilities for parallelism""" + from concurrent.futures import ProcessPoolExecutor, ThreadPoolExecutor -from typing import Dict, Literal, TypeVar +from typing import Literal, TypeVar from racs_tools.logging import logger @@ -15,7 +16,7 @@ Executor = TypeVar("Executor", ThreadPoolExecutor, ProcessPoolExecutor, MPIPoolExecutor) -EXECUTORS: Dict[Literal["thread", "process", "mpi"], Executor] = { +EXECUTORS: dict[Literal["thread", "process", "mpi"], Executor] = { "thread": ThreadPoolExecutor, "process": ProcessPoolExecutor, "mpi": MPIPoolExecutor, diff --git a/tests/hellompi.py b/tests/hellompi.py index bc872c9..924bf83 100644 --- a/tests/hellompi.py +++ b/tests/hellompi.py @@ -12,4 +12,5 @@ rank = MPI.COMM_WORLD.Get_rank() name = MPI.Get_processor_name() -sys.stdout.write("Hello, World! I am process %d of %d on %s.\n" % (rank, size, name)) +msg = f"Hello, World! I am process {rank:d} of {size:d} on {name:s}.\n" +sys.stdout.write(msg) diff --git a/tests/test_2d.py b/tests/test_2d.py index 0696052..ce79cf6 100644 --- a/tests/test_2d.py +++ b/tests/test_2d.py @@ -1,5 +1,4 @@ #!/usr/bin/env python3 -# -*- coding: utf-8 -*- import logging import shutil @@ -11,11 +10,10 @@ import numpy as np import pytest from astropy.io import fits -from radio_beam import Beam, Beams - from racs_tools import au2, beamcon_2D from racs_tools.convolve_uv import smooth from racs_tools.logging import logger +from radio_beam import Beam, Beams logger.setLevel(logging.DEBUG) @@ -27,7 +25,7 @@ class TestImage(NamedTuple): pix_scale: u.Quantity -@pytest.fixture +@pytest.fixture() def make_2d_image(tmpdir) -> TestImage: """Make a fake 2D image from with a Gaussian beam. @@ -76,7 +74,7 @@ def make_2d_image(tmpdir) -> TestImage: outf.unlink() -@pytest.fixture +@pytest.fixture() def make_2d_image_smaller(tmpdir) -> TestImage: """Make a fake 2D image from with a Gaussian beam. @@ -125,7 +123,7 @@ def make_2d_image_smaller(tmpdir) -> TestImage: outf.unlink() -@pytest.fixture +@pytest.fixture() def mirsmooth(make_2d_image: TestImage) -> TestImage: """Smooth a FITS image to a target beam using MIRIAD. @@ -139,15 +137,15 @@ def mirsmooth(make_2d_image: TestImage) -> TestImage: target_beam = Beam(40 * u.arcsec, 40 * u.arcsec, 0 * u.deg) outim = make_2d_image.path.with_suffix(".im") cmd = f"fits op=xyin in={make_2d_image.path.as_posix()} out={outim.as_posix()}" - sp.run(cmd.split()) + sp.run(cmd.split(), check=False) smoothim = make_2d_image.path.with_suffix(".smooth.im") cmd = f"convol map={outim} fwhm={target_beam.major.to(u.arcsec).value},{target_beam.minor.to(u.arcsec).value} pa={target_beam.pa.to(u.deg).value} options=final out={smoothim}" - sp.run(cmd.split()) + sp.run(cmd.split(), check=False) smoothfits = outim.with_suffix(".mirsmooth.fits") cmd = f"fits op=xyout in={smoothim} out={smoothfits}" - sp.run(cmd.split()) + sp.run(cmd.split(), check=False) yield TestImage( path=smoothfits, @@ -265,9 +263,9 @@ def test_smooth(make_2d_image: TestImage, mirsmooth: TestImage): dy=make_2d_image.pix_scale, conv_mode=conv_mode, ) - assert np.allclose( - smooth_data, mirsmooth.data, atol=1e-5 - ), f"Smooth with {conv_mode} does not match miriad" + assert np.allclose(smooth_data, mirsmooth.data, atol=1e-5), ( + f"Smooth with {conv_mode} does not match miriad" + ) def test_get_common_beam(make_2d_image, make_2d_image_smaller): @@ -286,7 +284,7 @@ def test_flags(): flags = np.array([False for beam in beams]) cutoff = 1 - with pytest.raises(TypeError) as e_info: + with pytest.raises(TypeError): flags = beams.major.to(u.arcsec) > cutoff * u.arcsec | flags flags = beams.major.to(u.arcsec).value > cutoff | flags diff --git a/tests/test_3d.py b/tests/test_3d.py index 835db8c..5372ecd 100644 --- a/tests/test_3d.py +++ b/tests/test_3d.py @@ -1,5 +1,4 @@ #!/usr/bin/env python3 -# -*- coding: utf-8 -*- import shutil import subprocess as sp @@ -10,14 +9,13 @@ import pytest from astropy.io import fits from astropy.table import Table -from radio_beam import Beam, Beams - from racs_tools import beamcon_3D +from radio_beam import Beam, Beams from .test_2d import TestImage, check_images -@pytest.fixture +@pytest.fixture() def make_3d_image( tmpdir: str, beams: Beams = Beams( @@ -125,7 +123,7 @@ def make_3d_image( outf.unlink() -@pytest.fixture +@pytest.fixture() def mirsmooth_3d(make_3d_image: TestImage) -> TestImage: """Smooth a FITS image to a target beam using MIRIAD. @@ -139,15 +137,15 @@ def mirsmooth_3d(make_3d_image: TestImage) -> TestImage: target_beam = Beam(60 * u.arcsec, 60 * u.arcsec, 0 * u.deg) outim = make_3d_image.path.with_suffix(".im") cmd = f"fits op=xyin in={make_3d_image.path.as_posix()} out={outim.as_posix()}" - sp.run(cmd.split()) + sp.run(cmd.split(), check=False) smoothim = make_3d_image.path.with_suffix(".smooth.im") cmd = f"convol map={outim} fwhm={target_beam.major.to(u.arcsec).value},{target_beam.minor.to(u.arcsec).value} pa={target_beam.pa.to(u.deg).value} options=cube out={smoothim}" - sp.run(cmd.split()) + sp.run(cmd.split(), check=False) smoothfits = outim.with_suffix(".mirsmooth_3d.fits") cmd = f"fits op=xyout in={smoothim} out={smoothfits}" - sp.run(cmd.split()) + sp.run(cmd.split(), check=False) yield TestImage( path=smoothfits, @@ -176,9 +174,50 @@ def test_robust_3d(make_3d_image, mirsmooth_3d): ) fname_beamcon = make_3d_image.path.with_suffix(".robust.fits") - assert check_images( - mirsmooth_3d.path, fname_beamcon - ), "Beamcon does not match Miriad" + assert check_images(mirsmooth_3d.path, fname_beamcon), ( + "Beamcon does not match Miriad" + ) + + +def test_get_target_beam(): + """Ensure that the processing of the target beam parameters behaves correctly""" + + none_test = beamcon_3D._get_target_beam() + assert none_test is None + + with pytest.raises(ValueError): + _ = beamcon_3D._get_target_beam(bmaj=10.0) + + single_beam = beamcon_3D._get_target_beam(bmaj=10.0, bmin=10.0, bpa=10.0) + assert isinstance(single_beam, Beam) + assert not isinstance(single_beam, Beams) + + single_nan_beam = beamcon_3D._get_target_beam(bmaj=np.nan, bmin=10.0, bpa=10.0) + assert isinstance(single_nan_beam, Beam) + assert not isinstance(single_nan_beam, Beams) + assert all( + np.isnan(i) + for i in (single_nan_beam.major, single_nan_beam.minor, single_nan_beam.pa) + ) + + many_beam = beamcon_3D._get_target_beam( + bmaj=[9, 8, 10.0], bmin=[9, 8, 10.0], bpa=[9, 8, 10.0] + ) + assert not isinstance(many_beam, Beam) + assert isinstance(many_beam, Beams) + assert len(many_beam) == 3 + + many_nan_beam = beamcon_3D._get_target_beam( + bmaj=[np.nan, 8, 10.0], bmin=[9, 8, 10.0], bpa=[9, 8, 10.0] + ) + assert not isinstance(many_beam, Beam) + assert isinstance(many_nan_beam, Beams) + assert len(many_nan_beam) == 3 + single_nan_beam = many_nan_beam[0] + assert all( + np.isnan(i) + for i in (single_nan_beam.major, single_nan_beam.minor, single_nan_beam.pa) + ) # def test_astropy(self): diff --git a/tests/test_noise.py b/tests/test_noise.py index 9a1a865..2e3bc2b 100644 --- a/tests/test_noise.py +++ b/tests/test_noise.py @@ -1,5 +1,4 @@ # #!/usr/bin/env python3 -# # -*- coding: utf-8 -*- # import subprocess as sp # import unittest