xradio 1.0.2__py3-none-any.whl → 1.1.1__py3-none-any.whl
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- xradio/_utils/_casacore/casacore_from_casatools.py +75 -9
- xradio/_utils/dict_helpers.py +38 -7
- xradio/_utils/list_and_array.py +26 -3
- xradio/_utils/schema.py +44 -0
- xradio/_utils/xarray_helpers.py +63 -0
- xradio/_utils/zarr/common.py +4 -2
- xradio/image/__init__.py +4 -2
- xradio/image/_util/_casacore/common.py +2 -1
- xradio/image/_util/_casacore/xds_from_casacore.py +144 -92
- xradio/image/_util/_casacore/xds_to_casacore.py +118 -53
- xradio/image/_util/_fits/xds_from_fits.py +125 -37
- xradio/image/_util/_zarr/common.py +0 -1
- xradio/image/_util/casacore.py +183 -25
- xradio/image/_util/common.py +10 -8
- xradio/image/_util/image_factory.py +469 -27
- xradio/image/image.py +72 -100
- xradio/image/image_xds.py +262 -0
- xradio/image/schema.py +85 -0
- xradio/measurement_set/__init__.py +5 -4
- xradio/measurement_set/_utils/_msv2/_tables/read.py +4 -3
- xradio/measurement_set/_utils/_msv2/conversion.py +6 -9
- xradio/measurement_set/_utils/_msv2/create_field_and_source_xds.py +1 -0
- xradio/measurement_set/_utils/_msv2/msv4_sub_xdss.py +1 -1
- xradio/measurement_set/_utils/_utils/interpolate.py +5 -0
- xradio/measurement_set/_utils/_utils/partition_attrs.py +0 -1
- xradio/measurement_set/convert_msv2_to_processing_set.py +9 -9
- xradio/measurement_set/load_processing_set.py +2 -2
- xradio/measurement_set/measurement_set_xdt.py +83 -93
- xradio/measurement_set/open_processing_set.py +1 -1
- xradio/measurement_set/processing_set_xdt.py +33 -26
- xradio/schema/check.py +70 -19
- xradio/schema/common.py +0 -1
- xradio/testing/__init__.py +0 -0
- xradio/testing/_utils/__template__.py +58 -0
- xradio/testing/measurement_set/__init__.py +58 -0
- xradio/testing/measurement_set/checker.py +131 -0
- xradio/testing/measurement_set/io.py +22 -0
- xradio/testing/measurement_set/msv2_io.py +1854 -0
- {xradio-1.0.2.dist-info → xradio-1.1.1.dist-info}/METADATA +65 -23
- xradio-1.1.1.dist-info/RECORD +75 -0
- {xradio-1.0.2.dist-info → xradio-1.1.1.dist-info}/WHEEL +1 -1
- xradio-1.0.2.dist-info/RECORD +0 -66
- {xradio-1.0.2.dist-info → xradio-1.1.1.dist-info}/licenses/LICENSE.txt +0 -0
- {xradio-1.0.2.dist-info → xradio-1.1.1.dist-info}/top_level.txt +0 -0
xradio/image/image.py
CHANGED
|
@@ -12,22 +12,25 @@ import shutil
|
|
|
12
12
|
import toolviper.utils.logger as logger
|
|
13
13
|
import xarray as xr
|
|
14
14
|
|
|
15
|
-
|
|
16
|
-
|
|
17
|
-
# from ._util.fits import _read_fits_image
|
|
18
|
-
from ._util.image_factory import (
|
|
15
|
+
from xradio.image._util.image_factory import (
|
|
19
16
|
_make_empty_aperture_image,
|
|
20
17
|
_make_empty_lmuv_image,
|
|
21
18
|
_make_empty_sky_image,
|
|
22
19
|
)
|
|
23
|
-
from ._util.zarr import
|
|
24
|
-
|
|
20
|
+
from xradio.image._util.zarr import (
|
|
21
|
+
_load_image_from_zarr_no_dask,
|
|
22
|
+
_xds_from_zarr,
|
|
23
|
+
_xds_to_zarr,
|
|
24
|
+
)
|
|
25
|
+
from xradio.image._util._fits.xds_from_fits import _fits_image_to_xds
|
|
25
26
|
|
|
26
|
-
|
|
27
|
+
from xradio.image._util.image_factory import create_image_xds_from_store
|
|
27
28
|
|
|
29
|
+
# warnings.filterwarnings("ignore", category=FutureWarning)
|
|
28
30
|
|
|
29
|
-
|
|
30
|
-
|
|
31
|
+
|
|
32
|
+
def open_image(
|
|
33
|
+
store: Union[str, dict],
|
|
31
34
|
chunks: dict = {},
|
|
32
35
|
verbose: bool = False,
|
|
33
36
|
do_sky_coords: bool = True,
|
|
@@ -65,9 +68,9 @@ def read_image(
|
|
|
65
68
|
|
|
66
69
|
Parameters
|
|
67
70
|
----------
|
|
68
|
-
|
|
69
|
-
Path to the input
|
|
70
|
-
|
|
71
|
+
store : str
|
|
72
|
+
Path to the input image
|
|
73
|
+
chunks : dict
|
|
71
74
|
The desired dask chunk size. Only applicable for casacore and fits images.
|
|
72
75
|
Supported optional keys are 'l', 'm', 'frequency', 'polarization', and 'time'.
|
|
73
76
|
The supported values are positive integers, indicating the length of a chunk
|
|
@@ -75,7 +78,7 @@ def read_image(
|
|
|
75
78
|
along that axis is equal to the number of pixels along that axis. For zarr
|
|
76
79
|
images, this parameter is ignored and the chunk size used to store the arrays
|
|
77
80
|
in the zarr image is used. 'l' represents the longitude like dimension, and 'm'
|
|
78
|
-
represents the latitude like dimension. For
|
|
81
|
+
represents the latitude like dimension. For aperture images, 'u' may be used in
|
|
79
82
|
place of 'l', and 'v' in place of 'm'.
|
|
80
83
|
verbose : bool
|
|
81
84
|
emit debugging messages? Default is False.
|
|
@@ -105,47 +108,36 @@ def read_image(
|
|
|
105
108
|
-------
|
|
106
109
|
xarray.Dataset
|
|
107
110
|
"""
|
|
108
|
-
#
|
|
109
|
-
#
|
|
110
|
-
|
|
111
|
-
|
|
112
|
-
|
|
113
|
-
|
|
114
|
-
|
|
115
|
-
|
|
116
|
-
|
|
117
|
-
|
|
118
|
-
|
|
119
|
-
|
|
120
|
-
|
|
121
|
-
|
|
122
|
-
|
|
123
|
-
|
|
124
|
-
|
|
125
|
-
|
|
126
|
-
|
|
127
|
-
|
|
128
|
-
|
|
129
|
-
|
|
130
|
-
|
|
131
|
-
|
|
132
|
-
except Exception as e:
|
|
133
|
-
emsgs.append(f"image format appears not to be fits {e.args}")
|
|
134
|
-
# when done debuggin comment out next line
|
|
135
|
-
# return _xds_from_zarr(infile, {"dv": "dask", "coords": "numpy"}, selection=selection)
|
|
136
|
-
try:
|
|
137
|
-
return _xds_from_zarr(
|
|
138
|
-
infile, {"dv": "dask", "coords": "numpy"}, selection=selection
|
|
139
|
-
)
|
|
140
|
-
except Exception as e:
|
|
141
|
-
emsgs.append(f"image format appears not to be zarr {e.args}")
|
|
142
|
-
emsgs.insert(
|
|
143
|
-
0, f"Unrecognized image format. Supported types are CASA, FITS, and zarr.\n"
|
|
111
|
+
# try:
|
|
112
|
+
# from ._util.casacore import _open_casa_image
|
|
113
|
+
# except ModuleNotFoundError as exc:
|
|
114
|
+
# logger.warning(
|
|
115
|
+
# "Could not import the function to convert from MSv2 to MSv4. "
|
|
116
|
+
# f"That functionality will not be available. Details: {exc}"
|
|
117
|
+
# )
|
|
118
|
+
# _open_casa_image = None
|
|
119
|
+
|
|
120
|
+
from ._util.casacore import _open_casa_image
|
|
121
|
+
|
|
122
|
+
img_xds = create_image_xds_from_store(
|
|
123
|
+
store,
|
|
124
|
+
_open_casa_image,
|
|
125
|
+
{"chunks": chunks, "verbose": verbose, "do_sky_coords": do_sky_coords},
|
|
126
|
+
_fits_image_to_xds,
|
|
127
|
+
{
|
|
128
|
+
"chunks": chunks,
|
|
129
|
+
"verbose": verbose,
|
|
130
|
+
"do_sky_coords": do_sky_coords,
|
|
131
|
+
"compute_mask": compute_mask,
|
|
132
|
+
},
|
|
133
|
+
_xds_from_zarr,
|
|
134
|
+
{"output": {"dv": "dask", "coords": "numpy"}, "selection": selection},
|
|
144
135
|
)
|
|
145
|
-
raise RuntimeError("\n".join(emsgs))
|
|
146
136
|
|
|
137
|
+
return img_xds
|
|
147
138
|
|
|
148
|
-
|
|
139
|
+
|
|
140
|
+
def load_image(store: str, block_des: dict = None, do_sky_coords=True) -> xr.Dataset:
|
|
149
141
|
"""
|
|
150
142
|
Load an image or portion of an image (subimage) into memory with data variables
|
|
151
143
|
being converted from dask to numpy arrays and coordinate arrays being converted
|
|
@@ -154,7 +146,7 @@ def load_image(infile: str, block_des: dict = None, do_sky_coords=True) -> xr.Da
|
|
|
154
146
|
|
|
155
147
|
Parameters
|
|
156
148
|
----------
|
|
157
|
-
|
|
149
|
+
store : str
|
|
158
150
|
Path to the input image, currently CASA and zarr images are supported
|
|
159
151
|
block_des : dict
|
|
160
152
|
The description of data to return, supported keys are time,
|
|
@@ -176,9 +168,6 @@ def load_image(infile: str, block_des: dict = None, do_sky_coords=True) -> xr.Da
|
|
|
176
168
|
-------
|
|
177
169
|
xarray.Dataset
|
|
178
170
|
"""
|
|
179
|
-
do_casa = True
|
|
180
|
-
emsgs = []
|
|
181
|
-
|
|
182
171
|
if block_des is None:
|
|
183
172
|
block_des = {}
|
|
184
173
|
|
|
@@ -187,52 +176,35 @@ def load_image(infile: str, block_des: dict = None, do_sky_coords=True) -> xr.Da
|
|
|
187
176
|
for k, v in selection.items():
|
|
188
177
|
if type(v) == int:
|
|
189
178
|
selection[k] = slice(v, v + 1)
|
|
190
|
-
|
|
191
|
-
|
|
192
|
-
|
|
193
|
-
|
|
194
|
-
|
|
195
|
-
|
|
196
|
-
|
|
197
|
-
|
|
198
|
-
|
|
199
|
-
|
|
200
|
-
|
|
201
|
-
try:
|
|
202
|
-
from ._util.casacore import _load_casa_image_block
|
|
203
|
-
|
|
204
|
-
return _load_casa_image_block(infile, selection, do_sky_coords)
|
|
205
|
-
except Exception as e:
|
|
206
|
-
emsgs.append(f"image format appears not to be casacore: {e.args}")
|
|
207
|
-
"""
|
|
208
|
-
try:
|
|
209
|
-
return __read_fits_image(infile, chunks, masks, history, verbose)
|
|
210
|
-
except Exception as e:
|
|
211
|
-
emsgs.append(f'image format appears not to be fits {e.args}')
|
|
212
|
-
"""
|
|
213
|
-
# when done debugging, comment out next line
|
|
214
|
-
# return _load_image_from_zarr_no_dask(infile, block_des)
|
|
215
|
-
# return _xds_from_zarr(infile, {"dv": "numpy"}, selection)
|
|
216
|
-
try:
|
|
217
|
-
return _load_image_from_zarr_no_dask(infile, block_des)
|
|
218
|
-
# return _xds_from_zarr(infile, {"dv": "numpy", "coords": "numpy"}, selection)
|
|
219
|
-
except Exception as e:
|
|
220
|
-
emsgs.append(f"image format appears not to be zarr {e.args}")
|
|
221
|
-
emsgs.insert(
|
|
222
|
-
0, f"Unrecognized image format. Supported formats are casacore and zarr.\n"
|
|
179
|
+
|
|
180
|
+
from ._util.casacore import _load_casa_image_block
|
|
181
|
+
|
|
182
|
+
img_xds = create_image_xds_from_store(
|
|
183
|
+
store,
|
|
184
|
+
_load_casa_image_block,
|
|
185
|
+
{"block_des": selection, "do_sky_coords": do_sky_coords},
|
|
186
|
+
None,
|
|
187
|
+
{},
|
|
188
|
+
_xds_from_zarr,
|
|
189
|
+
{"output": {"dv": "dask", "coords": "numpy"}, "selection": selection},
|
|
223
190
|
)
|
|
224
|
-
|
|
191
|
+
return img_xds
|
|
225
192
|
|
|
226
193
|
|
|
227
194
|
def write_image(
|
|
228
195
|
xds: xr.Dataset, imagename: str, out_format: str = "casa", overwrite: bool = False
|
|
229
196
|
) -> None:
|
|
230
197
|
"""
|
|
198
|
+
TODO: I think the user should be permitted to specify data groups to write.
|
|
231
199
|
Convert an xds image to CASA or zarr image.
|
|
232
200
|
xds : xarray.Dataset
|
|
233
201
|
XDS to convert
|
|
234
202
|
imagename : str
|
|
235
203
|
Path to output image
|
|
204
|
+
For writing to CASA, it is possible multiple images will be created, based
|
|
205
|
+
on what are in the data groups. If multiple images are created, the imagenames
|
|
206
|
+
will have identifying extensions added to the provided imagename. If only one
|
|
207
|
+
image is created, the provided imagename will be used as is.
|
|
236
208
|
out_format : str
|
|
237
209
|
Format of output image, currently "casa" and "zarr" are supported
|
|
238
210
|
overwrite : bool
|
|
@@ -256,9 +228,9 @@ def write_image(
|
|
|
256
228
|
)
|
|
257
229
|
my_format = out_format.lower()
|
|
258
230
|
if my_format == "casa":
|
|
259
|
-
from ._util.casacore import _xds_to_casa_image
|
|
231
|
+
from ._util.casacore import _xds_to_multiple_casa_images, _xds_to_casa_image
|
|
260
232
|
|
|
261
|
-
|
|
233
|
+
_xds_to_multiple_casa_images(xds, imagename)
|
|
262
234
|
elif my_format == "zarr":
|
|
263
235
|
_xds_to_zarr(xds, imagename)
|
|
264
236
|
else:
|
|
@@ -272,7 +244,7 @@ def make_empty_sky_image(
|
|
|
272
244
|
phase_center: Union[list, np.ndarray],
|
|
273
245
|
image_size: Union[list, np.ndarray],
|
|
274
246
|
cell_size: Union[list, np.ndarray],
|
|
275
|
-
|
|
247
|
+
frequency_coords: Union[list, np.ndarray],
|
|
276
248
|
pol_coords: Union[list, np.ndarray],
|
|
277
249
|
time_coords: Union[list, np.ndarray],
|
|
278
250
|
direction_reference: str = "fK5",
|
|
@@ -292,7 +264,7 @@ def make_empty_sky_image(
|
|
|
292
264
|
Number of x and y axis pixels in image.
|
|
293
265
|
cell_size : array of float, length = 2, units = rad
|
|
294
266
|
Cell size of x and y axis pixels in image.
|
|
295
|
-
|
|
267
|
+
frequency_coords : list or np.ndarray
|
|
296
268
|
The center frequency in Hz of each image channel.
|
|
297
269
|
pol_coords : list or np.ndarray
|
|
298
270
|
The polarization code for each image polarization.
|
|
@@ -312,7 +284,7 @@ def make_empty_sky_image(
|
|
|
312
284
|
phase_center,
|
|
313
285
|
image_size,
|
|
314
286
|
cell_size,
|
|
315
|
-
|
|
287
|
+
frequency_coords,
|
|
316
288
|
pol_coords,
|
|
317
289
|
time_coords,
|
|
318
290
|
direction_reference,
|
|
@@ -326,7 +298,7 @@ def make_empty_aperture_image(
|
|
|
326
298
|
phase_center: Union[List[float], np.ndarray],
|
|
327
299
|
image_size: Union[List[int], np.ndarray],
|
|
328
300
|
sky_image_cell_size: Union[List[float], np.ndarray],
|
|
329
|
-
|
|
301
|
+
frequency_coords: Union[List[float], np.ndarray],
|
|
330
302
|
pol_coords: Union[List[str], np.ndarray],
|
|
331
303
|
time_coords: Union[List[float], np.ndarray],
|
|
332
304
|
direction_reference: str = "fk5",
|
|
@@ -345,7 +317,7 @@ def make_empty_aperture_image(
|
|
|
345
317
|
Number of x and y axis pixels in image.
|
|
346
318
|
sky_image_cell_size : array of float, length = 2, units = rad
|
|
347
319
|
Cell size of x and y axis pixels in sky image, used to get cell size in uv image
|
|
348
|
-
|
|
320
|
+
frequency_coords : list or np.ndarray
|
|
349
321
|
The center frequency in Hz of each image channel.
|
|
350
322
|
pol_coords : list or np.ndarray
|
|
351
323
|
The polarization code for each image polarization.
|
|
@@ -362,7 +334,7 @@ def make_empty_aperture_image(
|
|
|
362
334
|
phase_center,
|
|
363
335
|
image_size,
|
|
364
336
|
sky_image_cell_size,
|
|
365
|
-
|
|
337
|
+
frequency_coords,
|
|
366
338
|
pol_coords,
|
|
367
339
|
time_coords,
|
|
368
340
|
direction_reference,
|
|
@@ -375,7 +347,7 @@ def make_empty_lmuv_image(
|
|
|
375
347
|
phase_center: Union[List[float], np.ndarray],
|
|
376
348
|
image_size: Union[List[int], np.ndarray],
|
|
377
349
|
sky_image_cell_size: Union[List[float], np.ndarray],
|
|
378
|
-
|
|
350
|
+
frequency_coords: Union[List[float], np.ndarray],
|
|
379
351
|
pol_coords: Union[List[float], np.ndarray],
|
|
380
352
|
time_coords: Union[List[float], np.ndarray],
|
|
381
353
|
direction_reference: str = "fk5",
|
|
@@ -396,7 +368,7 @@ def make_empty_lmuv_image(
|
|
|
396
368
|
sky_image_cell_size : array of float, length = 2, units = rad
|
|
397
369
|
Cell size of sky image. The cell size of the u,v image will be computed from
|
|
398
370
|
1/(image_size * sky_image_cell_size)
|
|
399
|
-
|
|
371
|
+
frequency_coords : list or np.ndarray
|
|
400
372
|
The center frequency in Hz of each image channel.
|
|
401
373
|
pol_coords : list or np.ndarray
|
|
402
374
|
The polarization code for each image polarization.
|
|
@@ -416,7 +388,7 @@ def make_empty_lmuv_image(
|
|
|
416
388
|
phase_center,
|
|
417
389
|
image_size,
|
|
418
390
|
sky_image_cell_size,
|
|
419
|
-
|
|
391
|
+
frequency_coords,
|
|
420
392
|
pol_coords,
|
|
421
393
|
time_coords,
|
|
422
394
|
direction_reference,
|
|
@@ -0,0 +1,262 @@
|
|
|
1
|
+
from collections.abc import Mapping, Iterable
|
|
2
|
+
import datetime
|
|
3
|
+
from typing import Any, Union
|
|
4
|
+
|
|
5
|
+
import numpy as np
|
|
6
|
+
import xarray as xr
|
|
7
|
+
|
|
8
|
+
from xradio._utils.list_and_array import to_list
|
|
9
|
+
|
|
10
|
+
IMAGE_DATASET_TYPES = {"image_dataset"}
|
|
11
|
+
|
|
12
|
+
from xradio._utils.xarray_helpers import get_data_group_name, create_new_data_group
|
|
13
|
+
|
|
14
|
+
|
|
15
|
+
class InvalidAccessorLocation(ValueError):
|
|
16
|
+
"""
|
|
17
|
+
Raised by ImageXds accessor functions called on a wrong Dataset (not image).
|
|
18
|
+
"""
|
|
19
|
+
|
|
20
|
+
pass
|
|
21
|
+
|
|
22
|
+
|
|
23
|
+
@xr.register_dataset_accessor("xr_img")
|
|
24
|
+
class ImageXds:
|
|
25
|
+
"""Accessor to the Image Dataset."""
|
|
26
|
+
|
|
27
|
+
_xds: xr.Dataset
|
|
28
|
+
|
|
29
|
+
def __init__(self, dataset: xr.Dataset):
|
|
30
|
+
"""
|
|
31
|
+
Initialize the ImageXds instance.
|
|
32
|
+
|
|
33
|
+
Parameters
|
|
34
|
+
----------
|
|
35
|
+
dataset: xarray.Dataset
|
|
36
|
+
The image Dataset node to construct an ImageXds accessor.
|
|
37
|
+
"""
|
|
38
|
+
|
|
39
|
+
self._xds = dataset
|
|
40
|
+
self.meta = {"summary": {}}
|
|
41
|
+
|
|
42
|
+
def test_func(self):
|
|
43
|
+
if self._xds.attrs.get("type") not in IMAGE_DATASET_TYPES:
|
|
44
|
+
raise InvalidAccessorLocation(f"{self._xds.path} is not a image node.")
|
|
45
|
+
|
|
46
|
+
return "Hallo"
|
|
47
|
+
|
|
48
|
+
def add_data_group(
|
|
49
|
+
self,
|
|
50
|
+
new_data_group_name: str,
|
|
51
|
+
new_data_group: dict = {},
|
|
52
|
+
data_group_dv_shared_with: str = None,
|
|
53
|
+
) -> xr.Dataset:
|
|
54
|
+
"""Adds a data group to the image Dataset, grouping the given data, weight, flag, etc. variables
|
|
55
|
+
and field_and_source_xds.
|
|
56
|
+
|
|
57
|
+
Parameters
|
|
58
|
+
----------
|
|
59
|
+
new_data_group_name : str
|
|
60
|
+
_description_
|
|
61
|
+
new_data_group : dict
|
|
62
|
+
_description_, by default Non
|
|
63
|
+
data_group_dv_shared_with : str, optional
|
|
64
|
+
_description_, by default "base"
|
|
65
|
+
|
|
66
|
+
Returns
|
|
67
|
+
-------
|
|
68
|
+
xr.Dataset
|
|
69
|
+
Image Dataset with the new group added
|
|
70
|
+
"""
|
|
71
|
+
|
|
72
|
+
if self._xds.attrs.get("type") not in IMAGE_DATASET_TYPES:
|
|
73
|
+
raise InvalidAccessorLocation(f"{self._xds.path} is not a image node.")
|
|
74
|
+
|
|
75
|
+
new_data_group_name, new_data_group = create_new_data_group(
|
|
76
|
+
self._xds,
|
|
77
|
+
"image",
|
|
78
|
+
new_data_group_name,
|
|
79
|
+
new_data_group,
|
|
80
|
+
data_group_dv_shared_with=data_group_dv_shared_with,
|
|
81
|
+
)
|
|
82
|
+
|
|
83
|
+
self._xds.attrs["data_groups"][new_data_group_name] = new_data_group
|
|
84
|
+
return self._xds
|
|
85
|
+
|
|
86
|
+
def get_lm_cell_size(self):
|
|
87
|
+
"""Get the lm cell size in radians from the image Dataset.
|
|
88
|
+
|
|
89
|
+
Returns
|
|
90
|
+
-------
|
|
91
|
+
float
|
|
92
|
+
The lm cell size in radians.
|
|
93
|
+
"""
|
|
94
|
+
|
|
95
|
+
if self._xds.attrs.get("type") not in IMAGE_DATASET_TYPES:
|
|
96
|
+
raise InvalidAccessorLocation(f"{self._xds.path} is not a image node.")
|
|
97
|
+
|
|
98
|
+
l_cell_size = self._xds.coords["l"][1].values - self._xds.coords["l"][0].values
|
|
99
|
+
m_cell_size = self._xds.coords["m"][1].values - self._xds.coords["m"][0].values
|
|
100
|
+
|
|
101
|
+
return np.array([l_cell_size, m_cell_size])
|
|
102
|
+
|
|
103
|
+
def add_uv_coordinates(self) -> xr.Dataset:
|
|
104
|
+
"""Adds the uv coordinates in wavelengths to the image Dataset.
|
|
105
|
+
|
|
106
|
+
Parameters
|
|
107
|
+
----------
|
|
108
|
+
|
|
109
|
+
Returns
|
|
110
|
+
-------
|
|
111
|
+
xr.Dataset
|
|
112
|
+
Image Dataset with the uv coordinates added.
|
|
113
|
+
"""
|
|
114
|
+
|
|
115
|
+
if self._xds.attrs.get("type") not in IMAGE_DATASET_TYPES:
|
|
116
|
+
raise InvalidAccessorLocation(f"{self._xds.path} is not a image node.")
|
|
117
|
+
|
|
118
|
+
from xradio.image._util.image_factory import _make_uv_coords
|
|
119
|
+
|
|
120
|
+
# self._xds = _make_uv_coords(self._xds,image_size=image_size, sky_image_cell_size=self.get_lm_cell_size())
|
|
121
|
+
|
|
122
|
+
# Calculate uv coordinates in meters based on l and m. _make_uv_coords assumes reference pixel at center (not necessary the case).
|
|
123
|
+
delta = self.get_lm_cell_size()
|
|
124
|
+
image_size = [self._xds.sizes["l"], self._xds.sizes["m"]]
|
|
125
|
+
|
|
126
|
+
u = self._xds.coords["l"].values / ((delta[0] ** 2) * image_size[0])
|
|
127
|
+
v = self._xds.coords["m"].values / ((delta[1] ** 2) * image_size[1])
|
|
128
|
+
|
|
129
|
+
self._xds = self._xds.assign_coords({"u": u, "v": v})
|
|
130
|
+
return self._xds
|
|
131
|
+
|
|
132
|
+
def get_uv_in_lambda(self, frequency: float):
|
|
133
|
+
"""Get the uv coordinates in wavelengths for a specific frequency from the image Dataset.
|
|
134
|
+
|
|
135
|
+
Parameters
|
|
136
|
+
----------
|
|
137
|
+
frequency : float
|
|
138
|
+
The frequency in Hz to calculate the uv coordinates in wavelengths.
|
|
139
|
+
|
|
140
|
+
Returns
|
|
141
|
+
-------
|
|
142
|
+
np.ndarray
|
|
143
|
+
The uv coordinates in wavelengths.
|
|
144
|
+
"""
|
|
145
|
+
|
|
146
|
+
if self._xds.attrs.get("type") not in IMAGE_DATASET_TYPES:
|
|
147
|
+
raise InvalidAccessorLocation(f"{self._xds.path} is not a image node.")
|
|
148
|
+
|
|
149
|
+
c = 299792458.0 # Speed of light in m/s
|
|
150
|
+
wavelength = c / frequency # Wavelength in meters
|
|
151
|
+
|
|
152
|
+
u_in_lambda = self._xds.coords["u"] / wavelength
|
|
153
|
+
v_in_lambda = self._xds.coords["v"] / wavelength
|
|
154
|
+
|
|
155
|
+
return u_in_lambda, v_in_lambda
|
|
156
|
+
|
|
157
|
+
def get_reference_pixel_indices(self):
|
|
158
|
+
"""Get the reference pixel indices from the image Dataset. The reference pixel is defined as the pixel where l=0 and m=0 or u=0 and v=0.
|
|
159
|
+
|
|
160
|
+
Returns
|
|
161
|
+
-------
|
|
162
|
+
dict
|
|
163
|
+
A dictionary with the reference pixel indices for each dimension.
|
|
164
|
+
"""
|
|
165
|
+
|
|
166
|
+
if self._xds.attrs.get("type") not in IMAGE_DATASET_TYPES:
|
|
167
|
+
raise InvalidAccessorLocation(f"{self._xds.path} is not a image node.")
|
|
168
|
+
|
|
169
|
+
image_center_index = None
|
|
170
|
+
|
|
171
|
+
if "l" in self._xds.coords:
|
|
172
|
+
l_index = np.where(self._xds.coords["l"].values == 0)[0][0]
|
|
173
|
+
m_index = np.where(self._xds.coords["m"].values == 0)[0][0]
|
|
174
|
+
|
|
175
|
+
lm_indexes = np.array([l_index, m_index])
|
|
176
|
+
image_center_index = lm_indexes
|
|
177
|
+
else:
|
|
178
|
+
lm_indexes = None
|
|
179
|
+
|
|
180
|
+
if "u" in self._xds.coords:
|
|
181
|
+
u_index = np.where(self._xds.coords["u"].values == 0)[0][0]
|
|
182
|
+
v_index = np.where(self._xds.coords["v"].values == 0)[0][0]
|
|
183
|
+
uv_indexes = np.array([u_index, v_index])
|
|
184
|
+
|
|
185
|
+
assert np.array_equal(
|
|
186
|
+
lm_indexes, uv_indexes
|
|
187
|
+
), "lm and uv reference pixel indices do not match."
|
|
188
|
+
image_center_index = uv_indexes
|
|
189
|
+
else:
|
|
190
|
+
uv_indexes = None
|
|
191
|
+
|
|
192
|
+
if image_center_index is None:
|
|
193
|
+
raise ValueError("No lm or uv coordinates found in the image Dataset.")
|
|
194
|
+
|
|
195
|
+
return image_center_index
|
|
196
|
+
|
|
197
|
+
def sel(
|
|
198
|
+
self,
|
|
199
|
+
indexers: Union[Mapping[Any, Any], None] = None,
|
|
200
|
+
method: Union[str, None] = None,
|
|
201
|
+
tolerance: Union[int, float, Iterable[Union[int, float]], None] = None,
|
|
202
|
+
drop: bool = False,
|
|
203
|
+
**indexers_kwargs: Any,
|
|
204
|
+
) -> xr.Dataset:
|
|
205
|
+
"""
|
|
206
|
+
Select data along dimension(s) by label. Alternative to `xarray.Dataset.sel <https://xarray.pydata.org/en/stable/generated/xarray.Dataset.sel.html>`__ so that a data group can be selected by name by using the `data_group_name` parameter.
|
|
207
|
+
For more information on data groups see `Data Groups <https://xradio.readthedocs.io/en/latest/measurement_set_overview.html#Data-Groups>`__ section. See `xarray.Dataset.sel <https://xarray.pydata.org/en/stable/generated/xarray.Dataset.sel.html>`__ for parameter descriptions.
|
|
208
|
+
|
|
209
|
+
Returns
|
|
210
|
+
-------
|
|
211
|
+
xarray.Dataset
|
|
212
|
+
xarray Dataset with ImageXds accessors
|
|
213
|
+
|
|
214
|
+
Examples
|
|
215
|
+
--------
|
|
216
|
+
>>> # Select data group 'robust0.5' and polarization 'XX'.
|
|
217
|
+
>>> selected_img_xds = img_xds.xr_img.sel(data_group_name='robust0.5', polarization='XX')
|
|
218
|
+
"""
|
|
219
|
+
|
|
220
|
+
if self._xds.attrs.get("type") not in IMAGE_DATASET_TYPES:
|
|
221
|
+
raise InvalidAccessorLocation(f"{self._xds.path} is not a image node.")
|
|
222
|
+
|
|
223
|
+
if "data_group_name" in indexers_kwargs:
|
|
224
|
+
data_group_name = indexers_kwargs["data_group_name"]
|
|
225
|
+
del indexers_kwargs["data_group_name"]
|
|
226
|
+
elif (indexers is not None) and ("data_group_name" in indexers):
|
|
227
|
+
data_group_name = indexers["data_group_name"]
|
|
228
|
+
del indexers["data_group_name"]
|
|
229
|
+
else:
|
|
230
|
+
data_group_name = None
|
|
231
|
+
|
|
232
|
+
if data_group_name is not None:
|
|
233
|
+
sel_data_group_set = set(
|
|
234
|
+
self._xds.attrs["data_groups"][data_group_name].values()
|
|
235
|
+
) - set(["date", "description"])
|
|
236
|
+
|
|
237
|
+
data_variables_to_drop = []
|
|
238
|
+
for dg_name, dg in self._xds.attrs["data_groups"].items():
|
|
239
|
+
# print(f"Data group: {dg_name}", dg)
|
|
240
|
+
dg_copy = dg.copy()
|
|
241
|
+
dg_copy.pop("date", None)
|
|
242
|
+
dg_copy.pop("description", None)
|
|
243
|
+
temp_set = set(dg_copy.values()) - sel_data_group_set
|
|
244
|
+
data_variables_to_drop.extend(list(temp_set))
|
|
245
|
+
|
|
246
|
+
data_variables_to_drop = list(set(data_variables_to_drop))
|
|
247
|
+
|
|
248
|
+
sel_img_xds = self._xds
|
|
249
|
+
|
|
250
|
+
sel_corr_xds = self._xds.sel(
|
|
251
|
+
indexers, method, tolerance, drop, **indexers_kwargs
|
|
252
|
+
).drop_vars(data_variables_to_drop)
|
|
253
|
+
|
|
254
|
+
sel_img_xds = sel_corr_xds
|
|
255
|
+
|
|
256
|
+
sel_img_xds.attrs["data_groups"] = {
|
|
257
|
+
data_group_name: self._xds.attrs["data_groups"][data_group_name]
|
|
258
|
+
}
|
|
259
|
+
|
|
260
|
+
return sel_img_xds
|
|
261
|
+
else:
|
|
262
|
+
return self._xds.sel(indexers, method, tolerance, drop, **indexers_kwargs)
|
xradio/image/schema.py
ADDED
|
@@ -0,0 +1,85 @@
|
|
|
1
|
+
from __future__ import annotations
|
|
2
|
+
|
|
3
|
+
from typing import Literal, Optional, Union
|
|
4
|
+
from xradio.schema.bases import (
|
|
5
|
+
xarray_dataset_schema,
|
|
6
|
+
xarray_dataarray_schema,
|
|
7
|
+
dict_schema,
|
|
8
|
+
)
|
|
9
|
+
|
|
10
|
+
|
|
11
|
+
@dict_schema
|
|
12
|
+
class DataGroupDict:
|
|
13
|
+
"""Defines a group of images."""
|
|
14
|
+
|
|
15
|
+
sky: Optional[str]
|
|
16
|
+
""" Name of the sky variable, for example 'SKY'. Derived from the gridded visibilities. On plane tangential to celestial sphere. """
|
|
17
|
+
flag: Optional[str]
|
|
18
|
+
""" Name of the sky pixels flags variable, for example 'FLAG_SKY'. For CASA images this is an internal mask. """
|
|
19
|
+
point_spread_function: Optional[str]
|
|
20
|
+
""" Name of the point spread function variable of the group, for example 'POINT_SPREAD_FUNCTION'. On plane tangential to celestial sphere. """
|
|
21
|
+
primary_beam: Optional[str]
|
|
22
|
+
""" Name of the primary beam variable of the group, for example 'PRIMARY_BEAM'. On plane tangential to celestial sphere. """
|
|
23
|
+
mask: Optional[str]
|
|
24
|
+
""" Name of the deconvolution mask variable of the group, for example 'MASK_DECONVOLVE'. On plane tangential to celestial sphere. """
|
|
25
|
+
beam_fit_params_sky: Optional[str]
|
|
26
|
+
""" Name of the beam fit parameters variable of the group, for example 'BEAM_FIT_PARAMETERS'. That applies to the sky, residual images and the point spread function if present. """
|
|
27
|
+
beam_fit_params_point_spread_function: Optional[str]
|
|
28
|
+
""" Name of the beam fit parameters variable of the group, for example 'BEAM_FIT_PARAMETERS'. That applies to the sky, residual images and the point spread function if present. """
|
|
29
|
+
visibility: Optional[str]
|
|
30
|
+
""" Name of the visibility variable of the group, for example 'VISIBILITY'. The gridded visibilities used to create the images using a Fourier transform. On aperture plane."""
|
|
31
|
+
visibility_normalization: Optional[str]
|
|
32
|
+
""" Normalization factor for the gridded visibility data. """
|
|
33
|
+
uv_sampling: Optional[str]
|
|
34
|
+
""" Name of the uv sampling variable of the group, for example 'UV_SAMPLING'. The gridded weights used to create the point spread function using a Fourier transform. On aperture plane."""
|
|
35
|
+
uv_sampling_normalization: Optional[str]
|
|
36
|
+
""" Normalization factor for the gridded weights. This is the sum of weights and the sensitivity can be calculated using 1/sqrt(uv_sampling_normalization)."""
|
|
37
|
+
aperture: Optional[str]
|
|
38
|
+
""" Name of the aperture variable of the group, for example 'APERTURE'. On aperture plane. The aperture is the Fourier transform of the primary beam."""
|
|
39
|
+
aperture_normalization: Optional[str]
|
|
40
|
+
""" Normalization factor for the aperture data. """
|
|
41
|
+
description: str
|
|
42
|
+
""" More details about the data group. """
|
|
43
|
+
date: str
|
|
44
|
+
""" Creation date-time, in ISO 8601 format: 'YYYY-MM-DDTHH:mm:ss.SSS'. """
|
|
45
|
+
|
|
46
|
+
# sky: Optional[str]
|
|
47
|
+
# """ Name of the sky variable, for example 'SKY'. Derived from the gridded visibilities. On plane tangential to celestial sphere. """
|
|
48
|
+
# flag_sky: Optional[str]
|
|
49
|
+
# """ Name of the sky pixels flags variable, for example 'FLAG_SKY'. For CASA images this is an internal mask. """
|
|
50
|
+
# model: Optional[str]
|
|
51
|
+
# """ Name of the model variable, for example 'MODEL'. On plane tangential to celestial sphere. """
|
|
52
|
+
# flag_model: Optional[str]
|
|
53
|
+
# """ Name of the model pixels flags variable, for example 'FLAG_MODEL'. For CASA images this is an internal mask. """
|
|
54
|
+
# residual: Optional[str]
|
|
55
|
+
# """ Name of the residual variable of the group, for example 'RESIDUAL'. residual = sky - model. On plane tangential to celestial sphere. """
|
|
56
|
+
# flag_residual: Optional[str]
|
|
57
|
+
# """ Name of the residual pixels flags variable, for example 'FLAG_RESIDUAL'. For CASA images this is an internal mask. """
|
|
58
|
+
# point_spread_function: Optional[str]
|
|
59
|
+
# """ Name of the point spread function variable of the group, for example 'POINT_SPREAD_FUNCTION'. On plane tangential to celestial sphere. """
|
|
60
|
+
# flag_point_spread_function: Optional[str]
|
|
61
|
+
# """ Name of the point spread function pixels flags variable, for example 'FLAG_POINT_SPREAD_FUNCTION'. For CASA images this is an internal mask. """
|
|
62
|
+
# primary_beam: Optional[str]
|
|
63
|
+
# """ Name of the primary beam variable of the group, for example 'PRIMARY_BEAM'. On plane tangential to celestial sphere. """
|
|
64
|
+
# flag_primary_beam: Optional[str]
|
|
65
|
+
# """ Name of the primary beam pixels flags variable, for example 'FLAG_PRIMARY_BEAM'. For CASA images this is an internal mask. """
|
|
66
|
+
# mask_deconvolve: Optional[str]
|
|
67
|
+
# """ Name of the deconvolution mask variable of the group, for example 'MASK_DECONVOLVE'. On plane tangential to celestial sphere. """
|
|
68
|
+
# beam_fit_params: Optional[str]
|
|
69
|
+
# """ Name of the beam fit parameters variable of the group, for example 'BEAM_FIT_PARAMETERS'. That applies to the sky, residual images and the point spread function if present. """
|
|
70
|
+
# visibility: Optional[str]
|
|
71
|
+
# """ Name of the visibility variable of the group, for example 'VISIBILITY'. The gridded visibilities used to create the images using a Fourier transform. On aperture plane."""
|
|
72
|
+
# visibility_normalization: Optional[str]
|
|
73
|
+
# """ Normalization factor for the gridded visibility data. """
|
|
74
|
+
# uv_sampling: Optional[str]
|
|
75
|
+
# """ Name of the uv sampling variable of the group, for example 'UV_SAMPLING'. The gridded weights used to create the point spread function using a Fourier transform. On aperture plane."""
|
|
76
|
+
# uv_sampling_normalization: Optional[str]
|
|
77
|
+
# """ Normalization factor for the gridded weights. This is the sum of weights and the sensitivity can be calculated using 1/sqrt(uv_sampling_normalization)."""
|
|
78
|
+
# aperture: Optional[str]
|
|
79
|
+
# """ Name of the aperture variable of the group, for example 'APERTURE'. On aperture plane. The aperture is the Fourier transform of the primary beam."""
|
|
80
|
+
# aperture_normalization: Optional[str]
|
|
81
|
+
# """ Normalization factor for the aperture data. """
|
|
82
|
+
# description: str
|
|
83
|
+
# """ More details about the data group. """
|
|
84
|
+
# date: str
|
|
85
|
+
# """ Creation date-time, in ISO 8601 format: 'YYYY-MM-DDTHH:mm:ss.SSS'. """
|