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

Improve src net rucio script #612

Merged
merged 14 commits into from
Sep 16, 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
2 changes: 1 addition & 1 deletion karabo/data/external_data.py
Original file line number Diff line number Diff line change
Expand Up @@ -83,7 +83,7 @@ def download(
os.makedirs(download_dir, exist_ok=True)

desc = f"Downloading {url} to {local_file_path}"
response.raw.read = functools.partial( # type: ignore
response.raw.read = functools.partial( # type: ignore[method-assign]
response.raw.read, decode_content=True
)
with tqdm.wrapattr(
Expand Down
23 changes: 10 additions & 13 deletions karabo/data/obscore.py
Original file line number Diff line number Diff line change
Expand Up @@ -500,7 +500,7 @@
freq_inc_hz = obs.frequency_increment_hz
min_freq_hz = obs.start_frequency_hz - freq_inc_hz / 2
end_freq_hz = min_freq_hz + freq_inc_hz * obs.number_of_channels
b = float(tel.longest_baseline())
b = float(tel.max_baseline())
ocm.s_resolution = tel.ang_res(freq=end_freq_hz, b=b)
return ocm

Expand All @@ -509,7 +509,7 @@
cls,
img: Image,
*,
fits_axes: FitsHeaderAxes = FitsHeaderAxes(), # immutable default
fits_axes: Optional[FitsHeaderAxes] = None,
) -> Self:
"""Update fields from `Image`.

Expand All @@ -534,12 +534,12 @@
"""
file = img.path
assert_valid_ending(path=file, ending=".fits")
if fits_axes is None:
fits_axes = FitsHeaderAxes()

Check warning on line 538 in karabo/data/obscore.py

View check run for this annotation

Codecov / codecov/patch

karabo/data/obscore.py#L538

Added line #L538 was not covered by tests
header = img.header
if not header["SIMPLE"]:
wmsg = f"{file} doesn't follow .fits standard! Info extraction might fail!"
warn(message=wmsg, category=UserWarning, stacklevel=1)
ra_deg = fits_axes.x.crval(header=header).to(u.deg).value # center
dec_deg = fits_axes.y.crval(header=header).to(u.deg).value # center
warn(message=wmsg, category=UserWarning, stacklevel=2)

Check warning on line 542 in karabo/data/obscore.py

View check run for this annotation

Codecov / codecov/patch

karabo/data/obscore.py#L542

Added line #L542 was not covered by tests
freq_center_hz = fits_axes.freq.crval(header=header).to(u.Hz).value # center
x_inc_deg = fits_axes.x.cdelt(header=header).to(u.deg).value # inc at center
y_inc_deg = fits_axes.y.cdelt(header=header).to(u.deg).value # inc at center
Expand All @@ -553,7 +553,7 @@
f"Pixel-size is not square for `s_pixel_scale`: {x_inc_deg=}, "
+ f"{y_inc_deg}. `s_pixel_scale` set to {s_pixel_scale=}."
)
warn(message=wmsg, category=UserWarning, stacklevel=1)
warn(message=wmsg, category=UserWarning, stacklevel=2)

Check warning on line 556 in karabo/data/obscore.py

View check run for this annotation

Codecov / codecov/patch

karabo/data/obscore.py#L556

Added line #L556 was not covered by tests
min_freq_hz = freq_center_hz - freq_inc_hz / 2
c = const.c.value
min_wavelength_m = c / min_freq_hz
Expand All @@ -562,14 +562,11 @@
fov_deg = np.sqrt(
(s_pixel_scale * x_pixel) ** 2 + (y_inc_deg * y_pixel) ** 2
) # circular fov of flattened image
half_width_deg = s_pixel_scale * x_pixel / 2
half_height_deg = abs(y_inc_deg) * y_pixel / 2
bottom_left = (ra_deg - half_width_deg, dec_deg - half_height_deg)
top_left = (ra_deg - half_width_deg, dec_deg + half_height_deg)
top_right = (ra_deg + half_width_deg, dec_deg + half_height_deg)
bottom_right = (ra_deg + half_width_deg, dec_deg - half_height_deg)
world_coords = Image.get_corners_in_world(header=header)
# assuming `get_corners_in_world` has circling corner order
corners = [(world_coord[0], world_coord[1]) for world_coord in world_coords]
spoly = cls.spoly(
poly=(bottom_left, top_left, top_right, bottom_right),
poly=corners,
ndigits=3,
suffix="d",
)
Expand Down
203 changes: 122 additions & 81 deletions karabo/examples/SRCNet_rucio_meta.py
Original file line number Diff line number Diff line change
@@ -1,119 +1,160 @@
"""Example script to attach Rucio ObsCore metadata data-products for ingestion.

This script probably needs adaption for your use-case. The parameters e.g. are not
customizable through an API. It also assumes that there's already a visibility
and image file available. Otherwise, you have to create them first. The manually
added values in this script are arbitrary to some extent and should be set (or not)
by yourself.

An end-to-end workflow would add the needed parts at the end of its simulation.
However, we just operate on existing files to avoid example-duplication.
Here, we create a simulated visibilities and a cleaned image. This is just an example
script which should be highly adapted, e.g. you custom-sky, simulation params (for
larger data-products, simulation params as user-input, multiple simulations, etc.).

Be aware that API-changes can take place in future Karabo-versions.

If not specified further (e.g. using `XDG_CACHE_HOME` or `FileHandler.root_stm`),
Karabo is using /tmp. Thus if you have a script which is producing large and/or many
data products, we suggest to adapt the cache-root to a volume with more space.
"""

from __future__ import annotations

import os
from argparse import ArgumentParser
from datetime import datetime

import numpy as np
from astropy import constants as const
from astropy import units as u

from karabo.data.obscore import FitsHeaderAxes, FitsHeaderAxis, ObsCoreMeta
from karabo.data.src import RucioMeta
from karabo.imaging.image import Image
from karabo.imaging.imager_wsclean import WscleanImageCleaner, WscleanImageCleanerConfig
from karabo.simulation.interferometer import InterferometerSimulation
from karabo.simulation.observation import Observation
from karabo.simulation.sky_model import SkyModel
from karabo.simulation.telescope import Telescope
from karabo.simulation.visibility import Visibility
from karabo.simulator_backend import SimulatorBackend
from karabo.util.file_handler import FileHandler
from karabo.util.helpers import get_rnd_str


def main() -> None:
parser = ArgumentParser()
parser.add_argument(
"--dp-path",
required=True,
type=str,
help="Path to data product inode (most likely file).",
# sky-to-visibilities simulation
# the params for this example are highly adapted to not create large
# data products, because this is not the focus of this script.
sky = SkyModel.get_GLEAM_Sky(min_freq=72e6, max_freq=231e6) # in Hz
phase_center = [250, -80] # RA,DEC in deg
filter_radius_deg = 0.8
sky = sky.filter_by_radius(0, filter_radius_deg, phase_center[0], phase_center[1])
tel = Telescope.constructor("ASKAP", backend=SimulatorBackend.OSKAR)
start_freq_hz = 76e6
num_chan = 16
freq_inc_hz = 1e8

obs = Observation(
start_frequency_hz=start_freq_hz,
start_date_and_time=datetime(2024, 3, 15, 10, 46, 0),
phase_centre_ra_deg=phase_center[0],
phase_centre_dec_deg=phase_center[1],
number_of_channels=num_chan,
frequency_increment_hz=freq_inc_hz,
number_of_time_steps=24,
)
parser.add_argument(
"--dp-type",
required=True,
type=str,
choices=["image", "visibility"],
help="Data product type. See `ObsCoreMeta` which file-formats are supported.",
# define any unique (required for ingestion) output-file-path
vis_path = os.path.join(FileHandler.stm(), "my-unique-vis-fname.vis")
interferometer_sim = InterferometerSimulation(
vis_path=vis_path, channel_bandwidth_hz=freq_inc_hz
)
vis = interferometer_sim.run_simulation(
telescope=tel,
sky=sky,
observation=obs,
backend=SimulatorBackend.OSKAR,
)
args = parser.parse_args()
dp_path: str = args.dp_path
dp_type: str = args.dp_type

if not os.path.exists(dp_path):
err_msg = f"Inode {dp_path=} doesn't exist!"
raise RuntimeError(err_msg)
dp_path_meta = RucioMeta.get_meta_fname(fname=dp_path)
if os.path.exists(dp_path_meta):
err_msg = f"{dp_path_meta=} already exists!"
raise FileExistsError(err_msg)
if dp_type == "image":
image = Image(path=dp_path)
# `FitsHeaderAxes` may need adaption based on the structure of your .fits image
axes = FitsHeaderAxes(freq=FitsHeaderAxis(axis=4, unit=u.Hz))
ocm = ObsCoreMeta.from_image(img=image, fits_axes=axes)
elif dp_type == "visibility":
vis = Visibility(vis_path=dp_path) # .vis supported, .ms not atm [07/2024]
# To extract additional information, `Telescope` & `Observation` should be
# provided with the same settings as `vis` was created. As mentioned in the
# module docstring, this is only necessary because we don't show the whole
# workflow here.
telescope = Telescope.constructor("ASKAP")
observation = Observation( # settings from notebook, of `minimal_visibility`
start_frequency_hz=100e6,
start_date_and_time=datetime(2024, 3, 15, 10, 46, 0),
phase_centre_ra_deg=250.0,
phase_centre_dec_deg=-80.0,
number_of_channels=16,
frequency_increment_hz=1e6,
number_of_time_steps=24,
)
ocm = ObsCoreMeta.from_visibility(
vis=vis,
calibrated=False,
tel=telescope,
obs=observation,
)
else:
err_msg = f"Unexpected {dp_type=}, allowed are only `dp-type` choices."
raise RuntimeError(err_msg)

# adapt each field according to your needs

# be sure that name & namespace together are unique, e.g. by having different fnames
name = os.path.split(dp_path)[-1]
rm = RucioMeta(
# create metadata of visibility (currently [08-24], .vis supported, casa .ms not)
vis_ocm = ObsCoreMeta.from_visibility(
vis=vis,
calibrated=False,
tel=tel,
obs=obs,
)
vis_rm = RucioMeta(
namespace="testing", # needs to be specified by Rucio service
name=name,
name=os.path.split(vis.vis_path)[-1], # remove path-infos for `name`
lifetime=86400, # 1 day
dataset_name=None,
meta=ocm,
meta=vis_ocm,
)

# ObsCore mandatory fields
ocm.obs_collection = "MRO/ASKAP"
obs_sim_id = 0 # unique observation-simulation ID of `USER`
vis_ocm.obs_collection = "MRO/ASKAP"
sfiruch marked this conversation as resolved.
Show resolved Hide resolved
obs_sim_id = 0 # inc/change for new simulation
user_rnd_str = get_rnd_str(k=7, seed=os.environ.get("USER"))
ocm.obs_id = f"karabo-{user_rnd_str}-{obs_sim_id}"
obs_id = f"karabo-{user_rnd_str}-{obs_sim_id}" # unique ID per user & simulation
vis_ocm.obs_id = obs_id
obs_publisher_did = RucioMeta.get_ivoid( # rest args are defaults
namespace=vis_rm.namespace,
name=vis_rm.name,
)
vis_ocm.obs_publisher_did = obs_publisher_did

# fill/correct other fields of `ObsCoreMeta` here!
# #####START#####
# HERE
# #####END#######

vis_path_meta = RucioMeta.get_meta_fname(fname=vis.vis_path)
_ = vis_rm.to_dict(fpath=vis_path_meta)
print(f"Created {vis_path_meta=}")

# -----Imaging-----

# calc imaging params
mean_freq = start_freq_hz + freq_inc_hz * (num_chan - 1) / 2
wavelength = const.c.value / mean_freq # in m
synthesized_beam = wavelength / tel.max_baseline() # in rad
imaging_cellsize = synthesized_beam / 3 # consider nyquist sampling theorem
fov_deg = 2 * filter_radius_deg # angular fov
imaging_npixel_estimate = fov_deg / np.rad2deg(imaging_cellsize) # not even|rounded
imaging_npixel = int(np.floor((imaging_npixel_estimate + 1) / 2.0) * 2.0)

print(f"Imaging: {imaging_npixel=}, {imaging_cellsize=} ...", flush=True)
restored_path = os.path.join(FileHandler.stm(), "my-unique-image-fname.fits")
restored = WscleanImageCleaner(
WscleanImageCleanerConfig(
imaging_npixel=imaging_npixel,
imaging_cellsize=imaging_cellsize,
niter=5000, # 10 times less than default
)
).create_cleaned_image( # currently, wsclean needs casa .ms, which is also created
ms_file_path=vis.ms_file_path,
Lukas113 marked this conversation as resolved.
Show resolved Hide resolved
output_fits_path=restored_path,
)

# create metadata for restored .fits image
# `FitsHeaderAxes` may need adaption based on the structure of your .fits image
axes = FitsHeaderAxes(freq=FitsHeaderAxis(axis=3, unit=u.Hz))
restored_ocm = ObsCoreMeta.from_image(img=restored, fits_axes=axes)
restored_rm = RucioMeta(
namespace="testing", # needs to be specified by Rucio service
name=os.path.split(restored.path)[-1], # remove path-infos for `name`
lifetime=86400, # 1 day
dataset_name=None,
meta=restored_ocm,
)
# ObsCore mandatory fields
# some of the metadata is taken from above, since both data-products originate
# from the same observation
restored_ocm.obs_collection = vis_ocm.obs_collection
restored_ocm.obs_id = vis_ocm.obs_id
obs_publisher_did = RucioMeta.get_ivoid( # rest args are defaults
namespace=rm.namespace,
name=rm.name,
namespace=restored_rm.namespace,
name=restored_rm.name,
)
ocm.obs_publisher_did = obs_publisher_did
restored_ocm.obs_publisher_did = obs_publisher_did

# fill other fields of `ObsCoreMeta` here!
# fill/correct other fields of `ObsCoreMeta` here!
# #####START#####
# HERE
# #####END#######

_ = rm.to_dict(fpath=dp_path_meta)
print(f"Created {dp_path_meta}")
restored_path_meta = RucioMeta.get_meta_fname(fname=restored.path)
_ = restored_rm.to_dict(fpath=restored_path_meta)
print(f"Created {restored_path_meta=}")


if __name__ == "__main__":
Expand Down
43 changes: 43 additions & 0 deletions karabo/imaging/image.py
Original file line number Diff line number Diff line change
Expand Up @@ -20,6 +20,7 @@
import matplotlib
import matplotlib.pyplot as plt
import numpy as np
from astropy.coordinates.sky_coordinate import SkyCoord
from astropy.io import fits
from astropy.io.fits.header import Header
from astropy.nddata import Cutout2D, NDData
Expand Down Expand Up @@ -747,6 +748,48 @@
wcs_2d = wcs.sub(ra_dec_axis)
return wcs_2d

@classmethod
def get_corners_in_world(cls, header: Header) -> NDArray[np.float64]:
"""Get the image corners RA,DEC of bl, br, tr & tl from `header` in deg.

Assumes that `header` has at least two axis RA & DEC.

Args:
header: Header to extract image-infos.

Returns:
Corners in world coordinates [deg] with shape 4x2.
"""
wcs = WCS(header)
if wcs.naxis < 2:
err_msg = (

Check warning on line 765 in karabo/imaging/image.py

View check run for this annotation

Codecov / codecov/patch

karabo/imaging/image.py#L765

Added line #L765 was not covered by tests
f"Header must have at least to axis (RA,DEC), but has {wcs.naxis=}"
)
raise RuntimeError(err_msg)

Check warning on line 768 in karabo/imaging/image.py

View check run for this annotation

Codecov / codecov/patch

karabo/imaging/image.py#L768

Added line #L768 was not covered by tests
naxis1 = header["NAXIS1"]
naxis2 = header["NAXIS2"]
corners = np.zeros(
shape=(4, wcs.naxis),
dtype=np.int64,
)
corners[1, 0] = naxis1 # bottom-right
corners[2, 0] = naxis1 # top-right
corners[2, 1] = naxis2 # # top-right
corners[3, 1] = naxis2 # top-left

world = wcs.pixel_to_world(*[corners[:, i] for i in range(corners.shape[1])])
if not isinstance(world, list): # typeguard would be safer, but was lazy
err_msg = (

Check warning on line 782 in karabo/imaging/image.py

View check run for this annotation

Codecov / codecov/patch

karabo/imaging/image.py#L782

Added line #L782 was not covered by tests
f"Expected list[{SkyCoord.__name__}], "
+ f"but got {type(world)=} of {world=}"
)
raise TypeError(err_msg)

Check warning on line 786 in karabo/imaging/image.py

View check run for this annotation

Codecov / codecov/patch

karabo/imaging/image.py#L786

Added line #L786 was not covered by tests
sky_coords: SkyCoord = world[0]
world_coords: NDArray[np.float64] = (
sky_coords.transform_to("icrs").to_table().to_pandas().to_numpy()
)
return world_coords


class ImageMosaicker:
"""
Expand Down
2 changes: 1 addition & 1 deletion karabo/simulation/telescope.py
Original file line number Diff line number Diff line change
Expand Up @@ -804,7 +804,7 @@ def get_baselines_dists(self) -> NDArray[np.float64]:
)
return dists

def longest_baseline(self) -> np.float64:
def max_baseline(self) -> np.float64:
"""Gets the longest baseline in meters.

Returns:
Expand Down
Loading