diff --git a/arrakis/cutout.py b/arrakis/cutout.py index 1d1acf14..4f37e090 100644 --- a/arrakis/cutout.py +++ b/arrakis/cutout.py @@ -16,12 +16,13 @@ import numpy as np import pandas as pd import pymongo -from astropy.coordinates import Latitude, Longitude, SkyCoord +from astropy.coordinates import SkyCoord from astropy.io import fits from astropy.utils import iers from astropy.utils.exceptions import AstropyWarning from astropy.wcs.utils import skycoord_to_pixel from prefect import flow, task +from prefect.task_runners import ConcurrentTaskRunner from spectral_cube import SpectralCube from spectral_cube.utils import SpectralCubeWarning from tqdm.auto import tqdm @@ -55,9 +56,9 @@ class CutoutArgs(Struct): """Arguments for cutout function""" """Name of the source""" - ra_high: float + ra_left: float """Upper RA bound in degrees""" - ra_low: float + ra_right: float """Lower RA bound in degrees""" dec_high: float """Upper DEC bound in degrees""" @@ -148,19 +149,38 @@ def cutout_image( padder: float = cube.header["BMAJ"] * u.deg * pad - xlo = Longitude(cutout_args.ra_low * u.deg) - Longitude(padder) - xhi = Longitude(cutout_args.ra_high * u.deg) + Longitude(padder) - ylo = Latitude(cutout_args.dec_low * u.deg) - Latitude(padder) - yhi = Latitude(cutout_args.dec_high * u.deg) + Latitude(padder) + top_right = SkyCoord(cutout_args.ra_right * u.deg, cutout_args.dec_high * u.deg) + bottom_left = SkyCoord(cutout_args.ra_left * u.deg, cutout_args.dec_low * u.deg) + top_left = SkyCoord(cutout_args.ra_left * u.deg, cutout_args.dec_high * u.deg) + bottom_right = SkyCoord(cutout_args.ra_right * u.deg, cutout_args.dec_low * u.deg) - xp_lo, yp_lo = skycoord_to_pixel(SkyCoord(xlo, ylo), cube.wcs) - xp_hi, yp_hi = skycoord_to_pixel(SkyCoord(xhi, yhi), cube.wcs) + top_right_off = top_right.directional_offset_by( + position_angle=-45 * u.deg, + separation=padder * np.sqrt(2), + ) + bottom_left_off = bottom_left.directional_offset_by( + position_angle=135 * u.deg, + separation=padder * np.sqrt(2), + ) + top_left_off = top_left.directional_offset_by( + position_angle=45 * u.deg, + separation=padder * np.sqrt(2), + ) + bottom_right_off = bottom_right.directional_offset_by( + position_angle=-135 * u.deg, + separation=padder * np.sqrt(2), + ) - # Round for cutout - yp_lo_idx = int(np.floor(yp_lo)) - yp_hi_idx = int(np.ceil(yp_hi)) - xp_lo_idx = int(np.floor(xp_hi)) - xp_hi_idx = int(np.ceil(xp_lo)) + x_left, y_bottom = skycoord_to_pixel(bottom_left_off, cube.wcs.celestial) + x_right, y_top = skycoord_to_pixel(top_right_off, cube.wcs.celestial) + _x_left, _y_top = skycoord_to_pixel(top_left_off, cube.wcs.celestial) + _x_right, _y_bottom = skycoord_to_pixel(bottom_right_off, cube.wcs.celestial) + + # Compare all points in case of insanity at the poles + yp_lo_idx = int(np.floor(min(y_bottom, _y_bottom, y_top, _y_top))) + yp_hi_idx = int(np.ceil(max(y_bottom, _y_bottom, y_top, _y_top))) + xp_lo_idx = int(np.floor(min(x_left, x_right, _x_left, _x_right))) + xp_hi_idx = int(np.ceil(max(x_left, x_right, _x_left, _x_right))) # Use subcube for header transformation cutout_cube = cube[:, yp_lo_idx:yp_hi_idx, xp_lo_idx:xp_hi_idx] @@ -230,39 +250,78 @@ def get_args( ras: u.Quantity = comps.RA.values * u.deg decs: u.Quantity = comps.Dec.values * u.deg majs: List[float] = comps.Maj.values * u.arcsec - - coords = SkyCoord(ras, decs) + coords = SkyCoord(ras, decs, frame="icrs") + padder = np.max(majs) try: - ra_max = np.max(coords.ra) - ra_i_max = np.argmax(coords.ra) - ra_off = Longitude(majs[ra_i_max]) - ra_high = ra_max + ra_off - - ra_min = np.min(coords.ra) - ra_i_min = np.argmin(coords.ra) - ra_off = Longitude(majs[ra_i_min]) - ra_low = ra_min - ra_off - - dec_max = np.max(coords.dec) - dec_i_max = np.argmax(coords.dec) - dec_off = Longitude(majs[dec_i_max]) - dec_high = dec_max + dec_off - - dec_min = np.min(coords.dec) - dec_i_min = np.argmin(coords.dec) - dec_off = Longitude(majs[dec_i_min]) - dec_low = dec_min - dec_off + # If in South - Northern boundary will have greatest RA range + # And vice versa + northern_boundary = SkyCoord(ras, np.ones_like(ras.value) * decs.max()) + southern_boundary = SkyCoord(ras, np.ones_like(ras.value) * decs.min()) + + if np.abs(decs.max()) > np.abs(decs.min()): + # North - use Southern + longest_boundary = southern_boundary + else: + # South - use Northern + longest_boundary = northern_boundary + + # Compute separations and position angles between all pairs of coordinates + separations = longest_boundary[:, np.newaxis].separation(longest_boundary) + position_angles = longest_boundary[:, np.newaxis].position_angle( + longest_boundary + ) + + # Find the index of the pair with the largest separation + largest_i, largest_j = np.unravel_index( + np.argmax(separations), separations.shape + ) + + # Determine left and right points based on position angle + if position_angles[largest_i, largest_j] > 180 * u.deg: + left_point = longest_boundary[largest_i] + right_point = longest_boundary[largest_j] + else: + left_point = longest_boundary[largest_j] + right_point = longest_boundary[largest_i] + + # Compute coordinates for top right, top left, bottom right, and bottom left points + top_right = SkyCoord(right_point.ra, decs.max()) + top_left = SkyCoord(left_point.ra, decs.max()) + bottom_right = SkyCoord(right_point.ra, decs.min()) + bottom_left = SkyCoord(left_point.ra, decs.min()) + + # Compute offsets for top right, top left, bottom right, and bottom left points + offset = padder * np.sqrt(2) + top_right_off = top_right.directional_offset_by( + position_angle=-45 * u.deg, separation=offset + ) + bottom_left_off = bottom_left.directional_offset_by( + position_angle=135 * u.deg, separation=offset + ) + top_left_off = top_left.directional_offset_by( + position_angle=45 * u.deg, separation=offset + ) + bottom_right_off = bottom_right.directional_offset_by( + position_angle=-135 * u.deg, separation=offset + ) + + # Compute dec_low, dec_high, ra_right, and ra_left + dec_low = min(bottom_left_off.dec.value, bottom_right_off.dec.value) * u.deg + dec_high = max(top_left_off.dec.value, top_right_off.dec.value) * u.deg + ra_right = top_right_off.ra + ra_left = top_left_off.ra + except Exception as e: logger.debug(f"coords are {coords=}") logger.debug(f"comps are {comps=}") raise e return CutoutArgs( - ra_high=ra_high.deg, - ra_low=ra_low.deg, - dec_high=dec_high.deg, - dec_low=dec_low.deg, + ra_left=ra_left.to(u.deg).value, + ra_right=ra_right.to(u.deg).value, + dec_high=dec_high.to(u.deg).value, + dec_low=dec_low.to(u.deg).value, outdir=outdir, ) @@ -523,7 +582,7 @@ def main(args: argparse.Namespace) -> None: args (argparse.Namespace): Command-line args verbose (bool, optional): Verbose output. Defaults to True. """ - cutout_islands( + cutout_islands.with_options(task_runner=ConcurrentTaskRunner)( field=args.field, directory=args.datadir, host=args.host, diff --git a/arrakis/rmsynth_oncuts.py b/arrakis/rmsynth_oncuts.py index 6a669b35..2162e6ab 100644 --- a/arrakis/rmsynth_oncuts.py +++ b/arrakis/rmsynth_oncuts.py @@ -26,6 +26,7 @@ from astropy.wcs import WCS from astropy.wcs.utils import proj_plane_pixel_scales from prefect import flow, task +from prefect.task_runners import ConcurrentTaskRunner from radio_beam import Beam from RMtools_1D import do_RMsynth_1D from RMtools_3D import do_RMsynth_3D @@ -296,7 +297,11 @@ def extract_single_spectrum( rms[np.isnan(rms)] = np.nanmedian(rms) wcs = WCS(header) x, y = np.array(wcs.celestial.world_to_pixel(coord)).round().astype(int) - spectrum_arr = np.array(data[:, y, x]) + try: + spectrum_arr = np.array(data[:, y, x]) + except Exception as e: + logger.error(f"Error extracting spectrum from {filename}") + raise e spectrum_arr[spectrum_arr == 0] = np.nan rms[~np.isfinite(spectrum_arr)] = np.nan bkg[~np.isfinite(spectrum_arr)] = np.nan @@ -1205,7 +1210,7 @@ def cli(): password=args.password, ) - main( + main.with_options(task_runner=ConcurrentTaskRunner)( field=args.field, outdir=Path(args.datadir), host=args.host,