Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Polar #71

Merged
merged 6 commits into from
May 13, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
139 changes: 99 additions & 40 deletions arrakis/cutout.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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"""
Expand Down Expand Up @@ -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]
Expand Down Expand Up @@ -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,
)

Expand Down Expand Up @@ -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,
Expand Down
9 changes: 7 additions & 2 deletions arrakis/rmsynth_oncuts.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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,
Expand Down
Loading