xradio 0.0.52__tar.gz → 0.0.53__tar.gz

{xradio-0.0.52/src/xradio.egg-info → xradio-0.0.53}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: xradio
-Version: 0.0.52
+Version: 0.0.53
 Summary:  Xarray Radio Astronomy Data IO
 Author-email: Jan-Willem Steeb <jsteeb@nrao.edu>, Federico Montesino Pouzols <pouzols@eso.edu>, Dave Mehringer <dmehring@nrao.edu>, Peter Wortmann <peter.wortmann@skao.int>
 License: BSD 3-Clause License

{xradio-0.0.52 → xradio-0.0.53}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "xradio"
-version = "0.0.52"
+version = "0.0.53"
 description = " Xarray Radio Astronomy Data IO"
 authors = [
     {name = "Jan-Willem Steeb", email="jsteeb@nrao.edu"},

{xradio-0.0.52 → xradio-0.0.53}/src/xradio/image/_util/_casacore/xds_to_casacore.py

@@ -198,7 +198,9 @@ def _imageinfo_dict_from_xds(xds: xr.Dataset) -> dict:
     ii["image_type"] = (
         xds[ap_sky].attrs["image_type"] if "image_type" in xds[ap_sky].attrs else ""
     )
-    ii["objectname"] = xds[ap_sky].attrs[_object_name]
+    ii["objectname"] = (
+        xds[ap_sky].attrs[_object_name] if _object_name in xds[ap_sky].attrs else ""
+    )
     if "BEAM" in xds.data_vars:
         # multi beam
         pp = {}

{xradio-0.0.52 → xradio-0.0.53}/src/xradio/image/_util/_fits/xds_from_fits.py

@@ -82,8 +82,6 @@ def _add_freq_attrs(xds: xr.Dataset, helpers: dict) -> xr.Dataset:
     if helpers["has_freq"]:
         meta["rest_frequency"] = make_quantity(helpers["restfreq"], "Hz")
         meta["rest_frequencies"] = [meta["rest_frequency"]]
-        # meta["frame"] = helpers["specsys"]
-        # meta["units"] = "Hz"
         meta["type"] = "frequency"
         meta["wave_unit"] = "mm"
         freq_axis = helpers["freq_axis"]
@@ -123,16 +121,6 @@ def _add_l_m_attrs(xds: xr.Dataset, helpers: dict) -> xr.Dataset:
 
 
 def _add_lin_attrs(xds: xr.Dataset, helpers: dict) -> xr.Dataset:
-    """
-    if helpers["sphr_dims"]:
-        for i, name in zip(helpers["dir_axes"], helpers["sphr_axis_names"]):
-            meta = {
-                "units": "rad",
-                "crval": helpers["crval"][i],
-                "cdelt": helpers["cdelt"][i],
-            }
-            xds.coords[name].attrs = meta
-    """
     if not helpers["sphr_dims"]:
         for i, j in zip(helpers["dir_axes"], ("u", "v")):
             meta = {
@@ -162,7 +150,7 @@ def _xds_direction_attrs_from_header(helpers: dict, header) -> dict:
     direction["projection"] = p0
     helpers["projection"] = p0
     ref_sys = header["RADESYS"]
-    ref_eqx = header["EQUINOX"]
+    ref_eqx = None if ref_sys.upper() == "ICRS" else header["EQUINOX"]
     if ref_sys == "FK5" and ref_eqx == 2000:
         ref_eqx = "J2000.0"
     helpers["ref_sys"] = ref_sys
@@ -180,9 +168,10 @@ def _xds_direction_attrs_from_header(helpers: dict, header) -> dict:
         ddata.append(x.value)
         # direction["reference"]["value"][i] = x.value
         x = helpers["cdelt"][i] * u.Unit(_get_unit(helpers["cunit"][i]))
-        dunits.append(x.to("rad"))
+        dunits.append("rad")
     direction["reference"] = make_skycoord_dict(ddata, units=dunits, frame=ref_sys)
-    direction["reference"]["attrs"]["equinox"] = ref_eqx.lower()
+    if ref_eqx is not None:
+        direction["reference"]["attrs"]["equinox"] = ref_eqx.lower()
     direction["latpole"] = make_quantity(
         header["LATPOLE"] * _deg_to_rad, "rad", dims=["l", "m"]
     )
@@ -285,6 +274,9 @@ def _user_attrs_from_header(header) -> dict:
         "ALTRPIX",
         "ALTRVAL",
         "BITPIX",
+        "BMAJ",
+        "BMIN",
+        "BPA",
         "BSCALE",
         "BTYPE",
         "BUNIT",
@@ -310,21 +302,20 @@ def _user_attrs_from_header(header) -> dict:
     ]
     regex = r"|".join(
         [
-            "^NAXIS\\d?$",
-            "^CRVAL\\d$",
-            "^CRPIX\\d$",
-            "^CTYPE\\d$",
-            "^CDELT\\d$",
-            "^CUNIT\\d$",
-            "^OBSGEO-(X|Y|Z)$",
-            "^P(C|V)\\d_\\d$",
+            r"^NAXIS\d?$",
+            r"^CRVAL\d$",
+            r"^CRPIX\d$",
+            r"^CTYPE\d$",
+            r"^CDELT\d$",
+            r"^CUNIT\d$",
+            r"^OBSGEO-(X|Y|Z)$",
+            r"^P(C|V)0?\d_0?\d",
         ]
     )
     user = {}
     for k, v in header.items():
-        if re.search(regex, k) or k in exclude:
-            continue
-        user[k.lower()] = v
+        if not (re.search(regex, k) or k in exclude):
+            user[k.lower()] = v
     return user
 
 

{xradio-0.0.52 → xradio-0.0.53}/src/xradio/image/_util/_zarr/xds_to_zarr.py

@@ -1,3 +1,4 @@
+import dask.array as da
 import logging
 import numpy as np
 import xarray as xr
@@ -6,6 +7,22 @@ from .common import _np_types, _top_level_sub_xds
 
 
 def _write_zarr(xds: xr.Dataset, zarr_store: str):
+    max_chunk_size = 0.95 * 2**30
+    for dv in xds.data_vars:
+        obj = xds[dv]
+        if isinstance(obj, xr.core.dataarray.DataArray) and isinstance(
+            obj.data, da.Array
+        ):
+            # get chunk size to make sure it is small enough to be compressed
+            ary = obj.data
+            chunk_size_bytes = np.prod(ary.chunksize) * np.dtype(ary.dtype).itemsize
+            if chunk_size_bytes > max_chunk_size:
+                raise ValueError(
+                    f"Chunk size of {chunk_size_bytes/1e9} GB for data variable {dv} "
+                    "bytes is too large for compression. To fix this, "
+                    "reduce the chunk size of the dask array in the data variable "
+                    f"by at least a factor of {chunk_size_bytes/max_chunk_size}."
+                )
     xds_copy = xds.copy(deep=True)
     sub_xds_dict = _encode(xds_copy, zarr_store)
     z_obj = xds_copy.to_zarr(store=zarr_store, compute=True)

{xradio-0.0.52 → xradio-0.0.53}/src/xradio/image/image.py

@@ -7,6 +7,9 @@ import warnings
 from typing import List, Union
 import copy
 import numpy as np
+import os
+import shutil
+import toolviper.utils.logger as logger
 import xarray as xr
 
 # from .._utils.zarr.common import _load_no_dask_zarr
@@ -182,27 +185,36 @@ def load_image(infile: str, block_des: dict = {}, do_sky_coords=True) -> xr.Data
 
 
 def write_image(
-    xds: xr.Dataset, imagename: str, out_format: str = "casa", overwrite=False
+    xds: xr.Dataset, imagename: str, out_format: str = "casa", overwrite: bool = False
 ) -> None:
     """
     Convert an xds image to CASA or zarr image.
     xds : xarray.Dataset
         XDS to convert
     imagename : str
-        Path to output CASA image
+        Path to output image
     out_format : str
         Format of output image, currently "casa" and "zarr" are supported
+    overwrite : bool
+        If True, overwrite existing image. Default is False.
     Returns
     -------
     None
     """
+    if os.path.exists(imagename):
+        if overwrite:
+            logger.warning(
+                f"Because overwrite=True, removing existing path {imagename}"
+            )
+            if os.path.isdir(imagename):
+                shutil.rmtree(imagename)
+            else:
+                os.remove(imagename)
+        else:
+            raise FileExistsError(
+                f"Path {imagename} already exists. Set overwrite=True to remove it."
+            )
     my_format = out_format.lower()
-
-    if overwrite:
-        import os
-
-        os.system("rm -rf " + imagename)
-
     if my_format == "casa":
         _xds_to_casa_image(xds, imagename)
     elif my_format == "zarr":

{xradio-0.0.52 → xradio-0.0.53}/src/xradio/measurement_set/_utils/_msv2/conversion.py

@@ -16,6 +16,7 @@ from xradio.measurement_set._utils._msv2.msv4_sub_xdss import (
     create_pointing_xds,
     create_system_calibration_xds,
     create_weather_xds,
+    create_phased_array_xds,
 )
 from .msv4_info_dicts import create_info_dicts
 from xradio.measurement_set.schema import MSV4_SCHEMA_VERSION
@@ -1243,6 +1244,14 @@ def convert_and_write_partition(
                 + str(time.time() - start)
             )
 
+        # Create phased array xds
+        phased_array_xds = create_phased_array_xds(
+            in_file,
+            ant_xds.antenna_name,
+            ant_xds.receptor_label,
+            ant_xds.polarization_type,
+        )
+
         start = time.time()
 
         # Time and frequency should always be increasing
@@ -1358,6 +1367,9 @@ def convert_and_write_partition(
         if weather_xds:
             ms_xdt["/weather_xds"] = weather_xds
 
+        if phased_array_xds:
+            ms_xdt["/phased_array_xds"] = phased_array_xds
+
         if storage_backend == "zarr":
             ms_xdt.to_zarr(store=os.path.join(out_file, ms_v4_name))
         elif storage_backend == "netcdf":

{xradio-0.0.52 → xradio-0.0.53}/src/xradio/measurement_set/_utils/_msv2/msv4_sub_xdss.py

@@ -1,10 +1,11 @@
 import toolviper.utils.logger as logger
 import os
 import time
-from typing import Tuple, Union
+from typing import Optional, Tuple, Union
 
 import numpy as np
 import xarray as xr
+from numpy.typing import ArrayLike
 
 from xradio._utils.coord_math import convert_to_si_units
 from xradio._utils.schema import (
@@ -741,3 +742,117 @@ def create_system_calibration_xds(
             sys_cal_xds[data_var] = sys_cal_xds[data_var].astype(np.float64)
 
     return sys_cal_xds
+
+
+def create_phased_array_xds(
+    in_file: str,
+    antenna_names: list[str],
+    receptor_label: list[str],
+    polarization_type: ArrayLike,
+) -> Optional[xr.Dataset]:
+    """
+    Create an Xarray Dataset containing phased array information.
+
+    Parameters
+    ----------
+    in_file : str
+        Path to the input MSv2.
+    antenna_names: DataArray or Sequence[str]
+        Content of the antenna_name coordinate of the antenna_xds.
+    receptor_label: DataArray or Sequence[str]
+        Content of the receptor_label coordinate of the antenna_xds. Used to
+        label the corresponding axis of ELEMENT_FLAG.
+    polarization_type: DataArray or ArrayLike
+        Contents of the polarization_type coordinate of the antenna_xds.
+        Array-like of shape (num_antennas, 2) containing the polarization
+        hands for each antenna.
+
+    Returns
+    ----------
+        xr.Dataset or None: If the input MS contains a PHASED_ARRAY table,
+           returns the Xarray Dataset containing the phased array information.
+           Otherwise, return None.
+    """
+
+    def extract_data(dataarray_or_sequence):
+        if hasattr(dataarray_or_sequence, "data"):
+            return dataarray_or_sequence.data.tolist()
+        return dataarray_or_sequence
+
+    antenna_names = extract_data(antenna_names)
+    receptor_label = extract_data(receptor_label)
+    polarization_type = extract_data(polarization_type)
+
+    # NOTE: We cannot use the dimension renaming option of `load_generic_table`
+    # here, because it leads to a dimension name collision. This is caused by
+    # the presence of two dimensions of size 3 in multiple arrays.
+    # Instead, we do the renaming manually below.
+    try:
+        raw_xds = load_generic_table(
+            in_file,
+            "PHASED_ARRAY",
+            # Some MSes carry COORDINATE_SYSTEM as a copy of COORDINATE_AXES
+            # due to a past ambiguity on the PHASED_ARRAY schema
+            ignore=["COORDINATE_SYSTEM", "ANTENNA_ID"],
+        )
+    except ValueError:
+        return None
+
+    # Defend against empty PHASED_ARRAY table.
+    # The test MS "AA2-Mid-sim_00000.ms" has that problem.
+    required_keys = {"COORDINATE_AXES", "ELEMENT_OFFSET", "ELEMENT_FLAG"}
+    if not all(k in raw_xds for k in required_keys):
+        return None
+
+    def msv4_measure(raw_name: str) -> dict:
+        coldesc = raw_xds.attrs["other"]["msv2"]["ctds_attrs"]["column_descriptions"]
+        return column_description_casacore_to_msv4_measure(coldesc[raw_name])
+
+    def make_data_variable(raw_name: str, dim_names: list[str]) -> xr.DataArray:
+        da = raw_xds[raw_name]
+        da = xr.DataArray(da.data, dims=tuple(dim_names))
+        return da.assign_attrs(msv4_measure(raw_name))
+
+    raw_datavar_names_and_dims = [
+        (
+            "COORDINATE_AXES",
+            ("antenna_name", "cartesian_pos_label_local", "cartesian_pos_label"),
+        ),
+        ("ELEMENT_OFFSET", ("antenna_name", "cartesian_pos_label_local", "element_id")),
+        ("ELEMENT_FLAG", ("antenna_name", "receptor_label", "element_id")),
+    ]
+
+    data_vars = {
+        name: make_data_variable(name, dims)
+        for name, dims in raw_datavar_names_and_dims
+    }
+    data_vars["COORDINATE_AXES"].attrs = {
+        "type": "rotation_matrix",
+        "units": ["undimensioned", "undimensioned", "undimensioned"],
+    }
+    # Remove the "frame" attribute if it exists, because ELEMENT_OFFSET is
+    # defined in a station-local frame for which no standard name exists
+    data_vars["ELEMENT_OFFSET"].attrs.pop("frame", None)
+    data_vars["ELEMENT_OFFSET"].attrs.update(
+        {
+            "coordinate_system": "topocentric",
+            "origin": "ANTENNA_POSITION",
+        }
+    )
+
+    num_elements = data_vars["ELEMENT_OFFSET"].sizes["element_id"]
+
+    data_vars = {"PHASED_ARRAY_" + key: val for key, val in data_vars.items()}
+    coords = {
+        "antenna_name": antenna_names,
+        "element_id": np.arange(num_elements),
+        "receptor_label": receptor_label,
+        "polarization_type": (
+            ("antenna_name", "receptor_label"),
+            polarization_type,
+        ),
+        "cartesian_pos_label": ["x", "y", "z"],
+        "cartesian_pos_label_local": ["p", "q", "r"],
+    }
+    attrs = {"type": "phased_array"}
+    return xr.Dataset(data_vars, coords, attrs)

{xradio-0.0.52 → xradio-0.0.53}/src/xradio/measurement_set/schema.py

@@ -26,6 +26,8 @@ TimeWeather = Literal["time_weather"]
 """ time dimension of weather dataset (when not interpolated to main time) """
 AntennaName = Literal["antenna_name"]
 """ Antenna name dimension """
+ElementId = Literal["element_id"]
+""" Element Id dimension of phased_array_xds"""
 StationName = Literal["station_name"]
 """ Station name dimension """
 ReceptorLabel = Literal["receptor_label"]
@@ -56,6 +58,8 @@ EllipsoidPosLabel = Literal["ellipsoid_pos_label"]
 """ Coordinate labels of geodetic earth location data (typically shape 3 and 'lon', 'lat', 'height')"""
 CartesianPosLabel = Literal["cartesian_pos_label"]
 """ Coordinate labels of geocentric earth location data (typically shape 3 and 'x', 'y', 'z')"""
+CartesianPosLabelLocal = Literal["cartesian_pos_label_local"]
+""" Coordinate labels for phased array elements positions relative to their parent station position; defined in a station-local frame (typically shape 3 and 'p', 'q', 'r')"""
 nPolynomial = Literal["n_polynomial"]
 """ For data that is represented as variable in time using Taylor expansion """
 PolyTerm = Literal["poly_term"]
@@ -76,9 +80,12 @@ SkyCoord = Literal["sky_coord"]
 SpectralCoord = Literal["spectral_coord"]
 Location = Literal["location"]
 Doppler = Literal["doppler"]
-
+RotationMatrix = Literal["rotation_matrix"]
 
 # Units of quantities and measures
+UnitsUndimensioned = list[
+    Literal["undimensioned"]
+]  # name consistent with casacore measures
 UnitsSeconds = list[Literal["s"]]
 UnitsHertz = list[Literal["Hz"]]
 UnitsMeters = list[Literal["m"]]
@@ -562,6 +569,7 @@ AllowedLocationCoordinateSystems = Literal[
     "geodetic",
     "planetodetic",
     "orbital",
+    "topocentric",
 ]
 
 
@@ -2217,3 +2225,98 @@ class SpectrumXds:
     """The channel bandwidth that includes the effects of missing data."""
     FREQUENCY_CENTROID: Optional[Dataof[FrequencyCentroidArray]] = None
     """Includes the effects of missing data unlike ``frequency``."""
+
+
+@xarray_dataarray_schema
+class PhasedArrayElementOffsetArray:
+    """
+    Schema for PHASED_ARRAY_ELEMENT_OFFSET.
+    """
+
+    data: Data[
+        tuple[AntennaName, CartesianPosLabelLocal, ElementId],
+        float,
+    ]
+
+    units: Attr[list[Literal["m"]]]
+
+    type: Attr[Location]
+    """ Measure type. Should be ``"location"``."""
+
+    coordinate_system: Attr[Literal["topocentric"]]
+    """
+    Coordinate system in which the element offsets are expressed.
+    Should be ``"topocentric"``.
+    """
+
+    origin: Attr[Literal["ANTENNA_POSITION"]]
+    """ Origin of the coordinate system. Should be ``"ANTENNA_POSITION"``."""
+
+
+@xarray_dataarray_schema
+class PhasedArrayCoordinateAxesArray:
+    """
+    Schema for PHASED_ARRAY_COORDINATE_AXES
+    """
+
+    data: Data[tuple[AntennaName, CartesianPosLabelLocal, CartesianPosLabel], float]
+
+    units: Attr[UnitsUndimensioned]
+
+    type: Attr[RotationMatrix]
+    """ Measure type. Should be ``"rotation_matrix"``."""
+
+
+@xarray_dataset_schema
+class PhasedArrayXds:
+    """
+    Phased array dataset: define stations made of multiple receiver elements.
+    """
+
+    # Coordinates
+    antenna_name: Coordof[AntennaNameArray]
+    """ Antenna name """
+
+    element_id: Coord[ElementId, Union[numpy.int64, numpy.int32]]
+    """ Element Id within a station/antenna """
+
+    receptor_label: Coord[ReceptorLabel, str]
+    """ Names of receptors, i.e. polarization hands. """
+
+    polarization_type: Coord[tuple[AntennaName, ReceptorLabel], str]
+    """ Polarization type to which each receptor responds (e.g. ”R”,”L”,”X” or ”Y”).
+    This is the receptor polarization type as recorded in the final correlated data (e.g. ”RR”); i.e.
+    as measured after all polarization combiners. ['X','Y'], ['R','L'] """
+
+    cartesian_pos_label: Coord[CartesianPosLabel, str]
+    """ (x,y,z) - either cartesian or ellipsoid """
+
+    cartesian_pos_label_local: Coord[CartesianPosLabelLocal, str]
+    """ (p,q,r) - cartesian station-local frame of reference """
+
+    # Data variables
+    PHASED_ARRAY_COORDINATE_AXES: Dataof[PhasedArrayCoordinateAxesArray]
+    """
+    3x3 Rotation M such that X_geo = M X_local.
+    Used to convert PHASED_ARRAY_ELEMENT_OFFSET coordinates from a station-local
+    frame to a geocentric frame.
+    """
+
+    PHASED_ARRAY_ELEMENT_OFFSET: Dataof[PhasedArrayElementOffsetArray]
+    """
+    Offsets of each array element from its parent station position, expressed
+    in a station-local frame. Station positions are stored in
+    antenna_xds.ANTENNA_POSITION.
+    """
+
+    PHASED_ARRAY_ELEMENT_FLAG: Data[tuple[AntennaName, ReceptorLabel, ElementId], bool]
+    """
+    Boolean flag set to True if the data from a given polarisation receptor of a station element
+    should be ignored.
+    """
+
+    # Attributes
+    type: Attr[Literal["phased_array"]]
+    """
+    Type of dataset. Expected to be ``phased_array``
+    """

{xradio-0.0.52 → xradio-0.0.53}/src/xradio/schema/check.py

@@ -628,14 +628,15 @@ def check_datatree(
             continue
 
         # Look up schema
-        schema = _DATASET_TYPES.get(node.attrs.get("type"))
+        typ = node.attrs.get("type")
+        schema = _DATASET_TYPES.get(typ)
         if schema is None:
             issues.add(
                 SchemaIssue(
                     [("", xds_name)],
                     message="Unknown dataset type!",
                     found=typ,
-                    expected=list(schemas.keys()),
+                    expected=None,
                 )
             )
             continue

{xradio-0.0.52 → xradio-0.0.53/src/xradio.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: xradio
-Version: 0.0.52
+Version: 0.0.53
 Summary:  Xarray Radio Astronomy Data IO
 Author-email: Jan-Willem Steeb <jsteeb@nrao.edu>, Federico Montesino Pouzols <pouzols@eso.edu>, Dave Mehringer <dmehring@nrao.edu>, Peter Wortmann <peter.wortmann@skao.int>
 License: BSD 3-Clause License