xradio 0.0.39__py3-none-any.whl → 0.0.41__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.
Files changed (52) hide show
  1. xradio/__init__.py +1 -1
  2. xradio/_utils/schema.py +14 -3
  3. xradio/{vis → correlated_data}/__init__.py +3 -2
  4. xradio/{vis/_vis_utils → correlated_data/_utils}/_ms/_tables/load_main_table.py +2 -2
  5. xradio/{vis/_vis_utils → correlated_data/_utils}/_ms/_tables/read.py +15 -1
  6. xradio/{vis/_vis_utils → correlated_data/_utils}/_ms/_tables/read_main_table.py +1 -1
  7. xradio/{vis/_vis_utils → correlated_data/_utils}/_ms/_tables/read_subtables.py +1 -1
  8. xradio/{vis/_vis_utils → correlated_data/_utils}/_ms/_tables/write.py +1 -1
  9. xradio/{vis/_vis_utils → correlated_data/_utils}/_ms/conversion.py +117 -58
  10. xradio/{vis/_vis_utils → correlated_data/_utils}/_ms/create_antenna_xds.py +196 -168
  11. xradio/{vis/_vis_utils → correlated_data/_utils}/_ms/create_field_and_source_xds.py +234 -200
  12. xradio/{vis/_vis_utils → correlated_data/_utils}/_ms/msv2_to_msv4_meta.py +1 -1
  13. xradio/correlated_data/_utils/_ms/msv4_info_dicts.py +203 -0
  14. xradio/correlated_data/_utils/_ms/msv4_sub_xdss.py +516 -0
  15. xradio/{vis/_vis_utils → correlated_data/_utils}/_ms/partition_queries.py +1 -1
  16. xradio/{vis/_vis_utils → correlated_data/_utils}/_ms/partitions.py +1 -1
  17. xradio/{vis/_vis_utils → correlated_data/_utils}/_ms/subtables.py +2 -2
  18. xradio/{vis/_vis_utils → correlated_data/_utils}/_utils/xds_helper.py +1 -1
  19. xradio/{vis/_vis_utils → correlated_data/_utils}/_zarr/read.py +1 -1
  20. xradio/{vis/_vis_utils → correlated_data/_utils}/_zarr/write.py +1 -1
  21. xradio/{vis/_vis_utils → correlated_data/_utils}/ms.py +1 -1
  22. xradio/{vis/_vis_utils → correlated_data/_utils}/zarr.py +4 -4
  23. xradio/{vis → correlated_data}/convert_msv2_to_processing_set.py +10 -3
  24. xradio/correlated_data/correlated_xds.py +13 -0
  25. xradio/{vis → correlated_data}/load_processing_set.py +13 -17
  26. xradio/{vis/read_processing_set.py → correlated_data/open_processing_set.py} +21 -23
  27. xradio/{vis/_processing_set.py → correlated_data/processing_set.py} +11 -12
  28. xradio/{vis → correlated_data}/schema.py +572 -186
  29. xradio/correlated_data/test__processing_set.py +74 -0
  30. xradio/image/_util/_casacore/xds_from_casacore.py +1 -1
  31. xradio/image/_util/_zarr/xds_from_zarr.py +1 -1
  32. xradio/image/_util/_zarr/zarr_low_level.py +1 -1
  33. {xradio-0.0.39.dist-info → xradio-0.0.41.dist-info}/METADATA +10 -10
  34. xradio-0.0.41.dist-info/RECORD +75 -0
  35. {xradio-0.0.39.dist-info → xradio-0.0.41.dist-info}/WHEEL +1 -1
  36. xradio/vis/_vis_utils/_ms/msv4_infos.py +0 -0
  37. xradio/vis/_vis_utils/_ms/msv4_sub_xdss.py +0 -351
  38. xradio-0.0.39.dist-info/RECORD +0 -73
  39. /xradio/{vis/_vis_utils → correlated_data/_utils}/__init__.py +0 -0
  40. /xradio/{vis/_vis_utils → correlated_data/_utils}/_ms/_tables/load.py +0 -0
  41. /xradio/{vis/_vis_utils → correlated_data/_utils}/_ms/_tables/table_query.py +0 -0
  42. /xradio/{vis/_vis_utils → correlated_data/_utils}/_ms/_tables/write_exp_api.py +0 -0
  43. /xradio/{vis/_vis_utils → correlated_data/_utils}/_ms/chunks.py +0 -0
  44. /xradio/{vis/_vis_utils → correlated_data/_utils}/_ms/descr.py +0 -0
  45. /xradio/{vis/_vis_utils → correlated_data/_utils}/_ms/msv2_msv3.py +0 -0
  46. /xradio/{vis/_vis_utils → correlated_data/_utils}/_ms/optimised_functions.py +0 -0
  47. /xradio/{vis/_vis_utils → correlated_data/_utils}/_utils/cds.py +0 -0
  48. /xradio/{vis/_vis_utils → correlated_data/_utils}/_utils/partition_attrs.py +0 -0
  49. /xradio/{vis/_vis_utils → correlated_data/_utils}/_utils/stokes_types.py +0 -0
  50. /xradio/{vis/_vis_utils → correlated_data/_utils}/_zarr/encoding.py +0 -0
  51. {xradio-0.0.39.dist-info → xradio-0.0.41.dist-info}/LICENSE.txt +0 -0
  52. {xradio-0.0.39.dist-info → xradio-0.0.41.dist-info}/top_level.txt +0 -0
@@ -0,0 +1,516 @@
1
+ import toolviper.utils.logger as logger
2
+ import os
3
+ import time
4
+ from typing import Tuple, Union
5
+
6
+ import numpy as np
7
+ import xarray as xr
8
+
9
+ from xradio._utils.schema import (
10
+ column_description_casacore_to_msv4_measure,
11
+ convert_generic_xds_to_xradio_schema,
12
+ )
13
+ from .subtables import subt_rename_ids
14
+ from ._tables.read import (
15
+ load_generic_table,
16
+ make_taql_where_between_min_max,
17
+ table_exists,
18
+ table_has_column,
19
+ )
20
+
21
+
22
+ def interpolate_to_time(
23
+ xds: xr.Dataset,
24
+ interp_time: Union[xr.DataArray, None],
25
+ message_prefix: str,
26
+ time_name: str = "time",
27
+ ) -> xr.Dataset:
28
+ """
29
+ Interpolate the time coordinate of the input xarray dataset to the
30
+ a data array. This can be used for example to interpolate a pointing_xds
31
+ to the time coord of the (main) MSv4, or similarly the ephemeris
32
+ data variables of a field_and_source_xds.
33
+
34
+ Uses interpolation method "linear", unless the source number of points is
35
+ 1 in which case "nearest" is used, to avoid divide-by-zero issues.
36
+
37
+ Parameters:
38
+ ----------
39
+ xds : xr.Dataset
40
+ Xarray dataset to interpolate (presumably a pointing_xds or an xds of
41
+ ephemeris variables)
42
+ interp_time : Union[xr.DataArray, None]
43
+ Time axis to interpolate the dataset to (usually main MSv4 time)
44
+ message_prefix: str
45
+ A prefix for info/debug/etc. messages
46
+
47
+ Returns:
48
+ -------
49
+ interpolated_xds : xr.Dataset
50
+ xarray dataset with time axis interpolated to interp_time.
51
+ """
52
+ if interp_time is not None:
53
+ points_before = xds[time_name].size
54
+ if points_before > 1:
55
+ method = "linear"
56
+ else:
57
+ method = "nearest"
58
+ xds = xds.interp({time_name: interp_time}, method=method, assume_sorted=True)
59
+ # scan_number sneaks in as a coordinate of the main time axis, drop it
60
+ if "scan_number" in xds.coords:
61
+ xds = xds.drop_vars("scan_number")
62
+ points_after = xds[time_name].size
63
+ logger.debug(
64
+ f"{message_prefix}: interpolating the time coordinate "
65
+ f"from {points_before} to {points_after} points"
66
+ )
67
+
68
+ return xds
69
+
70
+
71
+ def make_taql_where_weather(
72
+ in_file: str, ant_xds_station_name_ids: xr.DataArray
73
+ ) -> str:
74
+ """
75
+ The use of taql_where with WEATHER is complicated because of ALMA and its
76
+ (ab)use of NS_WX_STATION_ID vs. ANTENNA_ID (all -1) (see read.py). We
77
+ cannot simply use a 'WHERE ANTENNA_ID=...'. This function produces a TaQL
78
+ where string that uses ANTENNA_ID/NS_WX_STATION_ID depending on waht is
79
+ found in the WEATHER subtable.
80
+
81
+ Parameters
82
+ ----------
83
+ in_file : str
84
+ Input MS name.
85
+ ant_xds_station_name_ids : xr.DataArray
86
+ station name data array from antenna_xds, with name/id information
87
+
88
+ Returns
89
+ -------
90
+ taql_where: str
91
+ WHERE substring safe to use in taql_where when loading WEATHER subtable
92
+ """
93
+ weather_path = os.path.join(in_file, "WEATHER")
94
+ if table_exists(weather_path) and table_has_column(
95
+ weather_path, "NS_WX_STATION_ID"
96
+ ):
97
+ taql_where = f"WHERE (NS_WX_STATION_ID IN [{','.join(map(str, ant_xds_station_name_ids.antenna_id.values))}])"
98
+ else:
99
+ taql_where = f"WHERE (ANTENNA_ID IN [{','.join(map(str, ant_xds_station_name_ids.antenna_id.values))}])"
100
+
101
+ return taql_where
102
+
103
+
104
+ def create_weather_xds(in_file: str, ant_xds_station_name_ids: xr.DataArray):
105
+ """
106
+ Creates a Weather Xarray Dataset from a MS v2 WEATHER table.
107
+
108
+ Parameters
109
+ ----------
110
+ in_file : str
111
+ Input MS name.
112
+ ant_xds_station_name_ids : xr.DataArray
113
+ station name data array from antenna_xds, with name/id information
114
+
115
+ Returns
116
+ -------
117
+ xr.Dataset
118
+ Weather Xarray Dataset.
119
+ """
120
+
121
+ try:
122
+ taql_where = make_taql_where_weather(in_file, ant_xds_station_name_ids)
123
+ generic_weather_xds = load_generic_table(
124
+ in_file,
125
+ "WEATHER",
126
+ rename_ids=subt_rename_ids["WEATHER"],
127
+ taql_where=taql_where,
128
+ )
129
+ except ValueError as _exc:
130
+ return None
131
+
132
+ if not generic_weather_xds.data_vars:
133
+ # for example when the weather subtable only has info for antennas/stations
134
+ # not present in the MSv4 (no overlap between antennas loaded in ant_xds and weather)
135
+ return None
136
+
137
+ weather_xds = xr.Dataset(attrs={"type": "weather"})
138
+ stations_present = ant_xds_station_name_ids.sel(
139
+ antenna_id=generic_weather_xds["ANTENNA_ID"]
140
+ )
141
+ coords = {
142
+ "station_name": stations_present.data,
143
+ "antenna_name": stations_present.coords["antenna_name"].data,
144
+ }
145
+ weather_xds = weather_xds.assign_coords(coords)
146
+
147
+ dims_station_time = ["station_name", "time"]
148
+ to_new_data_variables = {
149
+ "H20": ["H2O", dims_station_time],
150
+ "IONOS_ELECTRON": ["IONOS_ELECTRON", dims_station_time],
151
+ "PRESSURE": ["PRESSURE", dims_station_time],
152
+ "REL_HUMIDITY": ["REL_HUMIDITY", dims_station_time],
153
+ "TEMPERATURE": ["TEMPERATURE", dims_station_time],
154
+ "DEW_POINT": ["DEW_POINT", dims_station_time],
155
+ "WIND_DIRECTION": ["WIND_DIRECTION", dims_station_time],
156
+ "WIND_SPEED": ["WIND_SPEED", dims_station_time],
157
+ }
158
+
159
+ to_new_coords = {
160
+ "TIME": ["time", ["time"]],
161
+ }
162
+
163
+ weather_xds = convert_generic_xds_to_xradio_schema(
164
+ generic_weather_xds, weather_xds, to_new_data_variables, to_new_coords
165
+ )
166
+
167
+ return weather_xds
168
+
169
+
170
+ def correct_generic_pointing_xds(
171
+ generic_pointing_xds: xr.Dataset, to_new_data_variables: dict[str, list]
172
+ ) -> xr.Dataset:
173
+ """
174
+ Takes a (generic) pointing_xds as read from a POINTING subtable of an MSv2
175
+ and tries to correct several deviations from the MSv2 specs seen in
176
+ common test data.
177
+ The problems fixed here include wrong dimensions:
178
+ - for example transposed dimensions with respect to the MSv2 specs (output
179
+ from CASA simulator),
180
+ - missing/additional unexpected dimensions when some of the columns are
181
+ empty (in the sense of "empty casacore cells").
182
+
183
+ This function modifies the data arrays of the data vars affected by such
184
+ issues.
185
+
186
+ Parameters
187
+ ----------
188
+ generic_pointing_xds: xr.Dataset
189
+ The generic pointing dataset (loaded from MSv2) to be fixed
190
+ to_new_data_variables: dict
191
+ The dict used for convert_generic_xds_to_xradio_schema, which gives all
192
+ the data variables relevant for the final MSv4 dataset.
193
+
194
+ Returns:
195
+ --------
196
+ xr.Dataset
197
+ Corrected dataset with dimensions conforming to MSv2 specs.
198
+ """
199
+
200
+ correct_pointing_xds = generic_pointing_xds.copy()
201
+
202
+ for key in generic_pointing_xds:
203
+ if key in to_new_data_variables:
204
+ data_var_name = to_new_data_variables[key]
205
+ # Corrects dim sizes of "empty cell" variables, such as empty DIRECTION, TARGET, etc.
206
+ if (
207
+ "dim_2" in generic_pointing_xds.sizes
208
+ and generic_pointing_xds.sizes["dim_2"] == 0
209
+ ):
210
+ # When all direction variables are "empty"
211
+ data_var_data = xr.DataArray(
212
+ [[[[np.nan, np.nan]]]],
213
+ dims=generic_pointing_xds.dims,
214
+ ).isel(n_polynomial=0, drop=True)
215
+ correct_pointing_xds[data_var_name].data = data_var_data
216
+
217
+ elif (
218
+ "dir" in generic_pointing_xds.sizes
219
+ and generic_pointing_xds.sizes["dir"] == 0
220
+ ):
221
+ # When some direction variables are "empty" but some are populated properly
222
+ if "dim_2" in generic_pointing_xds[key].sizes:
223
+ data_var_data = xr.DataArray(
224
+ generic_pointing_xds[key].values,
225
+ dims=generic_pointing_xds[key].dims,
226
+ )
227
+ else:
228
+ shape = tuple(
229
+ generic_pointing_xds.sizes[dim]
230
+ for dim in ["TIME", "ANTENNA_ID"]
231
+ ) + (2,)
232
+ data_var_data = xr.DataArray(
233
+ np.full(shape, np.nan),
234
+ dims=generic_pointing_xds[key].dims,
235
+ )
236
+ correct_pointing_xds[data_var_name].data = data_var_data
237
+
238
+ return correct_pointing_xds
239
+
240
+
241
+ def create_pointing_xds(
242
+ in_file: str,
243
+ ant_xds_name_ids: xr.DataArray,
244
+ time_min_max: Union[Tuple[np.float64, np.float64], None],
245
+ interp_time: Union[xr.DataArray, None] = None,
246
+ ) -> xr.Dataset:
247
+ """
248
+ Creates a Pointing Xarray Dataset from an MS v2 POINTING (sub)table.
249
+
250
+ WIP: details of a few direction variables (and possibly moving some to attributes) to be
251
+ settled (see MSv4 spreadsheet).
252
+
253
+ Parameters
254
+ ----------
255
+ in_file : str
256
+ Input MS name.
257
+ ant_xds_name_ids : xr.DataArray
258
+ antenna_name data array from antenna_xds, with name/id information
259
+ time_min_max : tuple
260
+ min / max times values to constrain loading (from the TIME column)
261
+ interp_time : Union[xr.DataArray, None] (Default value = None)
262
+ interpolate time to this (presumably main dataset time)
263
+
264
+ Returns
265
+ -------
266
+ xr.Dataset
267
+ Pointing Xarray dataset
268
+ """
269
+ start = time.time()
270
+
271
+ taql_time_range = make_taql_where_between_min_max(
272
+ time_min_max, in_file, "POINTING", "TIME"
273
+ )
274
+
275
+ if taql_time_range is None:
276
+ taql_where = f"WHERE (ANTENNA_ID IN [{','.join(map(str, ant_xds_name_ids.antenna_id.values))}])"
277
+ else:
278
+ taql_where = (
279
+ taql_time_range
280
+ + f" AND (ANTENNA_ID IN [{','.join(map(str, ant_xds_name_ids.antenna_id.values))}])"
281
+ )
282
+ # Read POINTING table into a Xarray Dataset.
283
+ generic_pointing_xds = load_generic_table(
284
+ in_file,
285
+ "POINTING",
286
+ rename_ids=subt_rename_ids["POINTING"],
287
+ taql_where=taql_where,
288
+ )
289
+
290
+ if not generic_pointing_xds.data_vars:
291
+ # apparently empty MS/POINTING table => produce empty xds
292
+ return xr.Dataset()
293
+
294
+ # Checking a simple way of using only the one single coefficient of the polynomials
295
+ if "n_polynomial" in generic_pointing_xds.sizes:
296
+ size = generic_pointing_xds.sizes["n_polynomial"]
297
+ if size == 1:
298
+ generic_pointing_xds = generic_pointing_xds.sel({"n_polynomial": 0})
299
+
300
+ time_ant_dims = ["time", "antenna_name"]
301
+ time_ant_dir_dims = time_ant_dims + ["local_sky_dir_label"]
302
+ to_new_data_variables = {
303
+ "DIRECTION": ["POINTING_BEAM", time_ant_dir_dims],
304
+ "ENCODER": ["POINTING_DISH_MEASURED", time_ant_dir_dims],
305
+ "OVER_THE_TOP": ["POINTING_OVER_THE_TOP", time_ant_dims],
306
+ }
307
+
308
+ to_new_coords = {
309
+ "TIME": ["time", ["time"]],
310
+ "dim_2": ["local_sky_dir_label", ["local_sky_dir_label"]],
311
+ }
312
+
313
+ generic_pointing_xds = correct_generic_pointing_xds(
314
+ generic_pointing_xds, to_new_data_variables
315
+ )
316
+ pointing_xds = xr.Dataset(attrs={"type": "pointing"})
317
+ coords = {
318
+ "antenna_name": ant_xds_name_ids.sel(
319
+ antenna_id=generic_pointing_xds["ANTENNA_ID"]
320
+ ).data,
321
+ "local_sky_dir_label": ["az", "alt"],
322
+ }
323
+ pointing_xds = pointing_xds.assign_coords(coords)
324
+ pointing_xds = convert_generic_xds_to_xradio_schema(
325
+ generic_pointing_xds, pointing_xds, to_new_data_variables, to_new_coords
326
+ )
327
+
328
+ pointing_xds = interpolate_to_time(pointing_xds, interp_time, "pointing_xds")
329
+
330
+ logger.debug(f"create_pointing_xds() execution time {time.time() - start:0.2f} s")
331
+
332
+ return pointing_xds
333
+
334
+
335
+ def prepare_generic_sys_cal_xds(generic_sys_cal_xds: xr.Dataset) -> xr.Dataset:
336
+ """
337
+ A generic_sys_cal_xds loaded with load_generic_table() cannot be easily
338
+ used in convert_generic_xds_to_xradio_schema() to produce an MSv4
339
+ sys_cal_xds dataset, as their structure differs in dimensions and order
340
+ of dimensions.
341
+ This function performs various prepareation steps, such as:
342
+ - filter out dimensions not neeed for an individual MSv4 (SPW, FEED),
343
+ - drop variables loaded from columns with all items set to empty array,
344
+ - transpose the dimensions frequency,receptor,
345
+ - fix dimension names when needed.
346
+
347
+ Parameters
348
+ ----------
349
+ generic_sys_cal_xds : xr.Dataset
350
+ generic dataset read from an MSv2 SYSCAL subtable
351
+
352
+ Returns
353
+ -------
354
+ generic_sys_cal_xds: xr.Dataset
355
+ System calibration Xarray Dataset prepared for generic conversion
356
+ to MSv4.
357
+ """
358
+
359
+ # drop SPW and feed dims
360
+ generic_sys_cal_xds = generic_sys_cal_xds.isel(SPECTRAL_WINDOW_ID=0, drop=True)
361
+ generic_sys_cal_xds = generic_sys_cal_xds.isel(FEED_ID=0, drop=True)
362
+
363
+ # Often some of the T*_SPECTRUM are present but all the cells are populated
364
+ # with empty arrays
365
+ empty_arrays_vars = []
366
+ for data_var in generic_sys_cal_xds.data_vars:
367
+ if generic_sys_cal_xds[data_var].size == 0:
368
+ empty_arrays_vars.append(data_var)
369
+ if empty_arrays_vars:
370
+ generic_sys_cal_xds = generic_sys_cal_xds.drop_vars(empty_arrays_vars)
371
+
372
+ # Re-arrange receptor and frequency dims depending on input structure
373
+ if (
374
+ "receptor" in generic_sys_cal_xds.sizes
375
+ and "frequency" in generic_sys_cal_xds.sizes
376
+ ):
377
+ # From MSv2 tables we get (...,frequency, receptor)
378
+ # -> transpose to (...,receptor,frequency) ready for MSv4 sys_cal_xds
379
+ generic_sys_cal_xds = generic_sys_cal_xds.transpose(
380
+ "ANTENNA_ID", "TIME", "receptor", "frequency"
381
+ )
382
+ else:
383
+ # because order is (...,frequency,receptor), when frequency is missing
384
+ # receptor can get wrongly labeled as frequency
385
+ generic_sys_cal_xds = generic_sys_cal_xds.rename_dims({"frequency": "receptor"})
386
+
387
+ return generic_sys_cal_xds
388
+
389
+
390
+ def create_system_calibration_xds(
391
+ in_file: str,
392
+ main_xds_frequency: xr.DataArray,
393
+ ant_xds_name_ids: xr.DataArray,
394
+ sys_cal_interp_time: Union[xr.DataArray, None] = None,
395
+ ):
396
+ """
397
+ Creates a system calibration Xarray Dataset from a MSv2 SYSCAL table.
398
+
399
+ Parameters
400
+ ----------
401
+ in_file: str
402
+ Input MS name.
403
+ main_xds_frequency: xr.DataArray
404
+ frequency array of the main xds (MSv4), containing among other things
405
+ spectral_window_id and measures metadata
406
+ ant_xds_name_ids : xr.Dataset
407
+ antenna_name data array from antenna_xds, with name/id information
408
+ sys_cal_interp_time: Union[xr.DataArray, None] = None,
409
+ Time axis to interpolate the data vars to (usually main MSv4 time)
410
+
411
+ Returns
412
+ -------
413
+ sys_cal_xds: xr.Dataset
414
+ System calibration Xarray Dataset.
415
+ """
416
+
417
+ spectral_window_id = main_xds_frequency.attrs["spectral_window_id"]
418
+ try:
419
+ generic_sys_cal_xds = load_generic_table(
420
+ in_file,
421
+ "SYSCAL",
422
+ rename_ids=subt_rename_ids["SYSCAL"],
423
+ taql_where=(
424
+ f" where (SPECTRAL_WINDOW_ID = {spectral_window_id})"
425
+ f" AND (ANTENNA_ID IN [{','.join(map(str, ant_xds_name_ids.antenna_id.values))}])"
426
+ ),
427
+ )
428
+ except ValueError as _exc:
429
+ return None
430
+
431
+ if not generic_sys_cal_xds.data_vars:
432
+ # even though SYSCAL is an optional subtable, some write it empty
433
+ return None
434
+
435
+ generic_sys_cal_xds = prepare_generic_sys_cal_xds(generic_sys_cal_xds)
436
+
437
+ mandatory_dimensions = ["antenna_name", "time_cal", "receptor_label"]
438
+ if "frequency" not in generic_sys_cal_xds.sizes:
439
+ dims_all = mandatory_dimensions
440
+ else:
441
+ dims_all = mandatory_dimensions + ["frequency_cal"]
442
+
443
+ to_new_data_variables = {
444
+ "PHASE_DIFF": ["PHASE_DIFFERENCE", ["antenna_name", "time_cal"]],
445
+ "TCAL": ["TCAL", dims_all],
446
+ "TCAL_SPECTRUM": ["TCAL", dims_all],
447
+ "TRX": ["TRX", dims_all],
448
+ "TRX_SPECTRUM": ["TRX", dims_all],
449
+ "TSKY": ["TSKY", dims_all],
450
+ "TSKY_SPECTRUM": ["TSKY", dims_all],
451
+ "TSYS": ["TSYS", dims_all],
452
+ "TSYS_SPECTRUM": ["TSYS", dims_all],
453
+ "TANT": ["TANT", dims_all],
454
+ "TANT_SPECTRUM": ["TANT", dims_all],
455
+ "TAN_TSYS": ["TAN_TSYS", dims_all],
456
+ "TANT_SYS_SPECTRUM": ["TANT_TSYS", dims_all],
457
+ }
458
+
459
+ to_new_coords = {
460
+ "TIME": ["time_cal", ["time_cal"]],
461
+ "receptor": ["receptor_label", ["receptor_label"]],
462
+ "frequency": ["frequency_cal", ["frequency_cal"]],
463
+ }
464
+
465
+ sys_cal_xds = xr.Dataset(attrs={"type": "sys_cal"})
466
+ coords = {
467
+ "antenna_name": ant_xds_name_ids.sel(
468
+ antenna_id=generic_sys_cal_xds["ANTENNA_ID"]
469
+ ).data,
470
+ "receptor_label": generic_sys_cal_xds.coords["receptor"].data,
471
+ }
472
+ sys_cal_xds = sys_cal_xds.assign_coords(coords)
473
+ sys_cal_xds = convert_generic_xds_to_xradio_schema(
474
+ generic_sys_cal_xds, sys_cal_xds, to_new_data_variables, to_new_coords
475
+ )
476
+
477
+ # Add frequency coord and its measures data, if present
478
+ if "frequency_cal" in dims_all:
479
+ frequency_coord = {
480
+ "frequency_cal": generic_sys_cal_xds.coords["frequency"].data
481
+ }
482
+ sys_cal_xds = sys_cal_xds.assign_coords(frequency_coord)
483
+ frequency_measure = {
484
+ "type": main_xds_frequency.attrs["type"],
485
+ "units": main_xds_frequency.attrs["units"],
486
+ "frame": main_xds_frequency.attrs["frame"],
487
+ "reference_value": main_xds_frequency.attrs["reference_frequency"],
488
+ }
489
+ sys_cal_xds.coords["frequency_cal"].attrs.update(frequency_measure)
490
+
491
+ if sys_cal_interp_time is not None:
492
+ sys_cal_xds = interpolate_to_time(
493
+ sys_cal_xds,
494
+ sys_cal_interp_time,
495
+ "system_calibration_xds",
496
+ time_name="time_cal",
497
+ )
498
+
499
+ time_coord_attrs = {
500
+ "type": "time",
501
+ "units": ["s"],
502
+ "scale": "UTC",
503
+ "format": "UNIX",
504
+ }
505
+ # If interpolating time, rename time_cal => time
506
+ time_coord = {"time": ("time_cal", sys_cal_interp_time.data)}
507
+ sys_cal_xds = sys_cal_xds.assign_coords(time_coord)
508
+ sys_cal_xds.coords["time"].attrs.update(time_coord_attrs)
509
+ sys_cal_xds = sys_cal_xds.swap_dims({"time_cal": "time"}).drop_vars("time_cal")
510
+
511
+ # correct expected types
512
+ for data_var in sys_cal_xds:
513
+ if sys_cal_xds.data_vars[data_var].dtype != np.float64:
514
+ sys_cal_xds[data_var] = sys_cal_xds[data_var].astype(np.float64)
515
+
516
+ return sys_cal_xds
@@ -1,5 +1,5 @@
1
1
  import itertools
2
- import graphviper.utils.logger as logger
2
+ import toolviper.utils.logger as logger
3
3
  from pathlib import Path
4
4
  from typing import Dict, List, Tuple, Union
5
5
 
@@ -1,4 +1,4 @@
1
- import graphviper.utils.logger as logger
1
+ import toolviper.utils.logger as logger
2
2
  from typing import Any, Dict, List, Tuple, Union
3
3
 
4
4
  import numpy as np
@@ -1,6 +1,6 @@
1
1
  import os
2
2
 
3
- import graphviper.utils.logger as logger
3
+ import toolviper.utils.logger as logger
4
4
 
5
5
  from pathlib import Path
6
6
  from typing import Dict, List
@@ -23,7 +23,7 @@ subt_rename_ids = {
23
23
  "SPECTRAL_WINDOW": {"row": "spectral_window_id", "dim_1": "chan"},
24
24
  "SOURCE": {"dim_1": "ra/dec", "dim_2": "line"},
25
25
  "STATE": {"row": "state_id"},
26
- "SYSCAL": {"dim_1": "channel", "dim_2": "receiver"},
26
+ "SYSCAL": {"dim_1": "frequency", "dim_2": "receptor"},
27
27
  # Would make sense for non-std "WS_NX_STATION_POSITION"
28
28
  "WEATHER": {"dim_1": "xyz"},
29
29
  }
@@ -1,5 +1,5 @@
1
1
  from importlib.metadata import version
2
- import graphviper.utils.logger as logger, multiprocessing, psutil
2
+ import toolviper.utils.logger as logger, multiprocessing, psutil
3
3
  from typing import Any, Dict, List, Tuple, Union
4
4
 
5
5
  import numpy as np
@@ -1,4 +1,4 @@
1
- import graphviper.utils.logger as logger
1
+ import toolviper.utils.logger as logger
2
2
  import os
3
3
  from pathlib import Path
4
4
  from typing import Dict, List, Tuple, Union
@@ -1,4 +1,4 @@
1
- import graphviper.utils.logger as logger, numcodecs, os, time, warnings
1
+ import toolviper.utils.logger as logger, numcodecs, os, time, warnings
2
2
  from itertools import cycle
3
3
  from pathlib import Path
4
4
  from typing import Any, Dict, Union
@@ -1,5 +1,5 @@
1
1
  import os
2
- import graphviper.utils.logger as logger
2
+ import toolviper.utils.logger as logger
3
3
  from typing import List, Tuple, Union
4
4
 
5
5
  from ._utils.cds import CASAVisSet
@@ -3,14 +3,14 @@ from pathlib import Path
3
3
  from typing import Dict, Union
4
4
 
5
5
  import zarr
6
- import graphviper.utils.logger as logger
6
+ import toolviper.utils.logger as logger
7
7
 
8
8
  from ._utils.cds import CASAVisSet
9
9
  from ._zarr.read import read_part_keys, read_partitions, read_subtables
10
10
  from ._zarr.write import write_metainfo, write_part_keys, write_partitions
11
11
 
12
12
 
13
- def is_zarr_vis(inpath: str) -> bool:
13
+ def is_zarr_cor(inpath: str) -> bool:
14
14
  """
15
15
  Check if a given path has a visibilities dataset in Zarr format
16
16
 
@@ -32,7 +32,7 @@ def is_zarr_vis(inpath: str) -> bool:
32
32
  return False
33
33
 
34
34
 
35
- def read_vis(
35
+ def read_cor(
36
36
  inpath: str,
37
37
  subtables: bool = True,
38
38
  asdm_subtables: bool = False,
@@ -83,7 +83,7 @@ def read_vis(
83
83
  return cds
84
84
 
85
85
 
86
- def write_vis(
86
+ def write_cor(
87
87
  cds: CASAVisSet,
88
88
  outpath: str,
89
89
  chunks_on_disk: Union[Dict, None] = None,
@@ -1,11 +1,11 @@
1
- import graphviper.utils.logger as logger
1
+ import toolviper.utils.logger as logger
2
2
  import numcodecs
3
3
  from typing import Dict, Union
4
4
 
5
5
  import dask
6
6
 
7
- from xradio.vis._vis_utils._ms.partition_queries import create_partitions
8
- from xradio.vis._vis_utils._ms.conversion import convert_and_write_partition
7
+ from xradio.correlated_data._utils._ms.partition_queries import create_partitions
8
+ from xradio.correlated_data._utils._ms.conversion import convert_and_write_partition
9
9
 
10
10
 
11
11
  def convert_msv2_to_processing_set(
@@ -18,6 +18,7 @@ def convert_msv2_to_processing_set(
18
18
  pointing_interpolate: bool = False,
19
19
  ephemeris_interpolate: bool = False,
20
20
  phase_cal_interpolate: bool = False,
21
+ sys_cal_interpolate: bool = False,
21
22
  use_table_iter: bool = False,
22
23
  compressor: numcodecs.abc.Codec = numcodecs.Zstd(level=2),
23
24
  storage_backend: str = "zarr",
@@ -47,6 +48,10 @@ def convert_msv2_to_processing_set(
47
48
  Whether to interpolate the time axis of the pointing sub-dataset to the time axis of the main dataset
48
49
  ephemeris_interpolate : bool, optional
49
50
  Whether to interpolate the time axis of the ephemeris data variables (of the field_and_source sub-dataset) to the time axis of the main dataset
51
+ phase_cal_interpolate : bool, optional
52
+ Whether to interpolate the time axis of the phase calibration data variables to the time axis of the main dataset
53
+ sys_cal_interpolate : bool, optional
54
+ Whether to interpolate the time axis of the system calibration data variables (sys_cal_xds) to the time axis of the main dataset
50
55
  use_table_iter : bool, optional
51
56
  Whether to use the table iterator to read the main table of the MS v2. This should be set to True when reading datasets with large number of rows and few partitions, by default False.
52
57
  compressor : numcodecs.abc.Codec, optional
@@ -96,6 +101,7 @@ def convert_msv2_to_processing_set(
96
101
  pointing_interpolate=pointing_interpolate,
97
102
  ephemeris_interpolate=ephemeris_interpolate,
98
103
  phase_cal_interpolate=phase_cal_interpolate,
104
+ sys_cal_interpolate=sys_cal_interpolate,
99
105
  compressor=compressor,
100
106
  overwrite=overwrite,
101
107
  )
@@ -114,6 +120,7 @@ def convert_msv2_to_processing_set(
114
120
  pointing_interpolate=pointing_interpolate,
115
121
  ephemeris_interpolate=ephemeris_interpolate,
116
122
  phase_cal_interpolate=phase_cal_interpolate,
123
+ sys_cal_interpolate=sys_cal_interpolate,
117
124
  compressor=compressor,
118
125
  overwrite=overwrite,
119
126
  )
@@ -0,0 +1,13 @@
1
+ import pandas as pd
2
+ from xradio._utils.list_and_array import to_list
3
+ import xarray as xr
4
+ import numbers
5
+
6
+
7
+ class CorrelatedXds(xr.Dataset):
8
+
9
+ def __init__(self, *args, **kwargs):
10
+ super().__init__(*args, **kwargs)
11
+
12
+ def to_zarr():
13
+ pass