xradio 0.0.48__py3-none-any.whl → 0.0.49__py3-none-any.whl

xradio/image/_util/image_factory.py

@@ -1,10 +1,16 @@
 from astropy.wcs import WCS
+from collections import OrderedDict
 import numpy as np
 import xarray as xr
 from typing import List, Union
 from .common import _c, _compute_world_sph_dims, _l_m_attr_notes
 from xradio._utils.coord_math import _deg_to_rad
-from xradio._utils.dict_helpers import make_quantity
+from xradio._utils.dict_helpers import (
+    make_frequency_reference_dict,
+    make_quantity,
+    make_skycoord_dict,
+    make_time_coord_attrs,
+)
 
 
 def _input_checks(
@@ -28,7 +34,7 @@ def _make_coords(
         chan_coords = [chan_coords]
     chan_coords = np.array(chan_coords, dtype=np.float64)
     restfreq = chan_coords[len(chan_coords) // 2]
-    vel_coords = (1 - chan_coords / restfreq) * _c
+    vel_coords = (1 - chan_coords / restfreq) * _c.to("m/s").value
     if not isinstance(time_coords, list) and not isinstance(time_coords, np.ndarray):
         time_coords = [time_coords]
     time_coords = np.array(time_coords, dtype=np.float64)
@@ -44,126 +50,124 @@ def _add_common_attrs(
     cell_size: Union[List[float], np.ndarray],
     projection: str,
 ) -> xr.Dataset:
-    xds.time.attrs = {"format": "MJD", "scale": "UTC", "units": "d"}
+    xds.time.attrs = make_time_coord_attrs(units=["d"], scale="utc", time_format="mjd")
     freq_vals = np.array(xds.frequency)
     xds.frequency.attrs = {
+        "observer": spectral_reference.lower(),
+        "reference_value": make_frequency_reference_dict(
+            value=freq_vals[len(freq_vals) // 2].item(),
+            units=["Hz"],
+            observer=spectral_reference.lower(),
+        ),
+        "rest_frequencies": make_quantity(restfreq, "Hz"),
         "rest_frequency": make_quantity(restfreq, "Hz"),
-        "frame": spectral_reference.upper(),
-        "units": "Hz",
-        "wave_unit": "mm",
-        # "crval": chan_coords[len(chan_coords) // 2],
-        "crval": freq_vals[len(freq_vals) // 2].item(),
-        "cdelt": (freq_vals[1] - freq_vals[0] if len(freq_vals) > 1 else 1000.0),
-        "pc": 1.0,
+        "type": "frequency",
+        "units": ["Hz"],
+        "wave_unit": ["mm"],
     }
-    xds.velocity.attrs = {"doppler_type": "RADIO", "units": "m/s"}
+    xds.velocity.attrs = {"doppler_type": "radio", "type": "doppler", "units": "m/s"}
+    reference = make_skycoord_dict(
+        data=phase_center, units=["rad", "rad"], frame=direction_reference
+    )
+    reference["attrs"].update({"equinox": "j2000.0"})
     xds.attrs = {
+        "data_groups": {"base": {}},
         "direction": {
-            "reference": {
-                "type": "sky_coord",
-                "frame": direction_reference,
-                "equinox": "J2000",
-                "value": list(phase_center),
-                "units": ["rad", "rad"],
-            },
-            "longpole": make_quantity(np.pi, "rad"),
-            "latpole": make_quantity(0.0, "rad"),
-            "pc": np.array([[1.0, 0.0], [0.0, 1.0]]),
+            "reference": reference,
+            "lonpole": make_quantity(np.pi, "rad", ["l", "m"]),
+            "latpole": make_quantity(0.0, "rad", ["l", "m"]),
+            "pc": [[1.0, 0.0], [0.0, 1.0]],
             "projection": projection,
             "projection_parameters": [0.0, 0.0],
         },
-        "active_mask": "",
-        "beam": None,
-        "object_name": "",
-        "obsdate": {
-            "type": "time",
-            "scale": "UTC",
-            "format": "MJD",
-            "value": np.array(xds.time)[0],
-            "units": "d",
-        },
-        "observer": "Karl Jansky",
-        "pointing_center": {"value": list(phase_center), "initial": True},
-        "description": "",
-        "telescope": {
-            "name": "ALMA",
-            "position": {
-                "type": "position",
-                "ellipsoid": "GRS80",
-                "units": ["rad", "rad", "m"],
-                "value": [-1.1825465955049892, -0.3994149869262738, 6379946.01326443],
-            },
-        },
-        "history": None,
     }
     return xds
 
 
-def _make_empty_sky_image(
-    phase_center: Union[list, np.ndarray],
-    image_size: Union[list, np.ndarray],
-    cell_size: Union[list, np.ndarray],
-    chan_coords: Union[list, np.ndarray],
+def _make_common_coords(
     pol_coords: Union[list, np.ndarray],
+    chan_coords: Union[list, np.ndarray],
     time_coords: Union[list, np.ndarray],
-    direction_reference: str,
-    projection: str,
-    spectral_reference: str,
-    do_sky_coords: bool,
-) -> xr.Dataset:
-    _input_checks(phase_center, image_size, cell_size)
+) -> dict:
     some_coords = _make_coords(chan_coords, time_coords)
-    if do_sky_coords:
-        long, lat = _compute_world_sph_dims(
-            projection=projection,
-            shape=image_size,
-            ctype=["RA", "Dec"],
-            crpix=[image_size[0] // 2, image_size[1] // 2],
-            crval=phase_center,
-            cdelt=[-abs(cell_size[0]), abs(cell_size[1])],
-            cunit=["rad", "rad"],
-        )["value"]
+    return {
+        "coords": {
+            "time": some_coords["time"],
+            "frequency": some_coords["chan"],
+            "velocity": ("frequency", some_coords["vel"]),
+            "polarization": pol_coords,
+        },
+        "restfreq": some_coords["restfreq"],
+    }
+
+
+def _make_lm_values(
+    image_size: Union[list, np.ndarray],
+    cell_size: Union[list, np.ndarray],
+) -> dict:
     # l follows RA as far as increasing/decreasing, see AIPS Meme 27, change in alpha
     # definition three lines below Figure 2 and the first of the pair of equations 10.
-    l_coords = [
+    l = [
         (i - image_size[0] // 2) * (-1) * abs(cell_size[0])
         for i in range(image_size[0])
     ]
-    m_coords = [
-        (i - image_size[1] // 2) * abs(cell_size[1]) for i in range(image_size[1])
-    ]
-    coords = {
-        "time": some_coords["time"],
-        "polarization": pol_coords,
-        "frequency": some_coords["chan"],
-        "velocity": ("frequency", some_coords["vel"]),
-        "l": l_coords,
-        "m": m_coords,
-    }
-    if do_sky_coords:
-        coords["right_ascension"] = (("l", "m"), long)
-        coords["declination"] = (("l", "m"), lat)
-    xds = xr.Dataset(coords=coords)
+    m = [(i - image_size[1] // 2) * abs(cell_size[1]) for i in range(image_size[1])]
+    return {"l": l, "m": m}
+
+
+def _make_sky_coords(
+    projection: str,
+    image_size: Union[list, np.ndarray],
+    cell_size: Union[list, np.ndarray],
+    phase_center: Union[list, np.ndarray],
+) -> dict:
+    long, lat = _compute_world_sph_dims(
+        projection=projection,
+        shape=image_size,
+        ctype=["RA", "Dec"],
+        crpix=[image_size[0] // 2, image_size[1] // 2],
+        crval=phase_center,
+        cdelt=[-abs(cell_size[0]), abs(cell_size[1])],
+        cunit=["rad", "rad"],
+    )["value"]
+    return {"right_ascension": (("l", "m"), long), "declination": (("l", "m"), lat)}
+
+
+def _add_lm_coord_attrs(xds: xr.Dataset) -> xr.Dataset:
     attr_note = _l_m_attr_notes()
     xds.l.attrs = {
-        "type": "quantity",
-        "crval": 0.0,
-        "cdelt": -abs(cell_size[0]),
-        "units": "rad",
-        "type": "quantity",
         "note": attr_note["l"],
     }
     xds.m.attrs = {
-        "type": "quantity",
-        "crval": 0.0,
-        "cdelt": abs(cell_size[1]),
-        "units": "rad",
-        "type": "quantity",
         "note": attr_note["m"],
     }
+
+
+def _make_empty_sky_image(
+    phase_center: Union[list, np.ndarray],
+    image_size: Union[list, np.ndarray],
+    cell_size: Union[list, np.ndarray],
+    chan_coords: Union[list, np.ndarray],
+    pol_coords: Union[list, np.ndarray],
+    time_coords: Union[list, np.ndarray],
+    direction_reference: str,
+    projection: str,
+    spectral_reference: str,
+    do_sky_coords: bool,
+) -> xr.Dataset:
+    _input_checks(phase_center, image_size, cell_size)
+    cc = _make_common_coords(pol_coords, chan_coords, time_coords)
+    coords = cc["coords"]
+    lm_values = _make_lm_values(image_size, cell_size)
+    coords.update(lm_values)
+    if do_sky_coords:
+        coords.update(_make_sky_coords(projection, image_size, cell_size, phase_center))
+    xds = xr.Dataset(coords=coords)
+    xds = _move_beam_param_dim_coord(xds)
+    _add_lm_coord_attrs(xds)
     _add_common_attrs(
         xds,
-        some_coords["restfreq"],
+        cc["restfreq"],
         spectral_reference,
         direction_reference,
         phase_center,
@@ -173,6 +177,34 @@ def _make_empty_sky_image(
     return xds
 
 
+def _make_uv_coords(
+    xds: xr.Dataset,
+    image_size: Union[list, np.ndarray],
+    sky_image_cell_size: Union[list, np.ndarray],
+) -> dict:
+    uv_values = _make_uv_values(image_size, sky_image_cell_size)
+    xds = xds.assign_coords(uv_values)
+    attr = make_quantity(0.0, "lambda")
+    xds.u.attrs = attr.copy()
+    xds.v.attrs = attr.copy()
+    return xds
+
+
+def _make_uv_values(
+    image_size: Union[list, np.ndarray],
+    sky_image_cell_size: Union[list, np.ndarray],
+) -> dict:
+    im_size_wave = 1 / np.array(sky_image_cell_size)
+    uv_cell_size = im_size_wave / np.array(image_size)
+    u_vals = [
+        (i - image_size[0] // 2) * abs(uv_cell_size[0]) for i in range(image_size[0])
+    ]
+    v_vals = [
+        (i - image_size[1] // 2) * abs(uv_cell_size[1]) for i in range(image_size[1])
+    ]
+    return {"u": u_vals, "v": v_vals}
+
+
 def _make_empty_aperture_image(
     phase_center: Union[list, np.ndarray],
     image_size: Union[list, np.ndarray],
@@ -185,44 +217,27 @@ def _make_empty_aperture_image(
     spectral_reference: str,
 ) -> xr.Dataset:
     _input_checks(phase_center, image_size, sky_image_cell_size)
-    some_coords = _make_coords(chan_coords, time_coords)
-    im_size_wave = 1 / np.array(sky_image_cell_size)
-    uv_cell_size = im_size_wave / np.array(image_size)
-    u_coords = [
-        (i - image_size[0] // 2) * abs(uv_cell_size[0]) for i in range(image_size[0])
-    ]
-    v_coords = [
-        (i - image_size[1] // 2) * abs(uv_cell_size[1]) for i in range(image_size[1])
-    ]
-    coords = {
-        "time": some_coords["time"],
-        "polarization": pol_coords,
-        "frequency": some_coords["chan"],
-        "velocity": ("frequency", some_coords["vel"]),
-        "u": u_coords,
-        "v": v_coords,
-    }
+    cc = _make_common_coords(pol_coords, chan_coords, time_coords)
+    coords = cc["coords"]
     xds = xr.Dataset(coords=coords)
-    xds.u.attrs = {
-        "crval": 0.0,
-        "cdelt": abs(uv_cell_size[0]),
-    }
-    xds.v.attrs = {
-        "crval": 0.0,
-        "cdelt": abs(uv_cell_size[1]),
-    }
+    xds = _make_uv_coords(xds, image_size, sky_image_cell_size)
     _add_common_attrs(
         xds,
-        some_coords["restfreq"],
+        cc["restfreq"],
         spectral_reference,
         direction_reference,
         phase_center,
         sky_image_cell_size,
         projection,
     )
+    xds = _move_beam_param_dim_coord(xds)
     return xds
 
 
+def _move_beam_param_dim_coord(xds: xr.Dataset) -> xr.Dataset:
+    return xds.assign_coords(beam_param=("beam_param", ["major", "minor", "pa"]))
+
+
 def _make_empty_lmuv_image(
     phase_center: Union[list, np.ndarray],
     image_size: Union[list, np.ndarray],
@@ -235,81 +250,18 @@ def _make_empty_lmuv_image(
     spectral_reference: str,
     do_sky_coords: bool,
 ) -> xr.Dataset:
-    _input_checks(phase_center, image_size, sky_image_cell_size)
-    some_coords = _make_coords(chan_coords, time_coords)
-    if do_sky_coords:
-        long, lat = _compute_world_sph_dims(
-            projection=projection,
-            shape=image_size,
-            ctype=["RA", "Dec"],
-            crpix=[image_size[0] // 2, image_size[1] // 2],
-            crval=phase_center,
-            cdelt=[-abs(sky_image_cell_size[0]), abs(sky_image_cell_size[1])],
-            cunit=["rad", "rad"],
-        )["value"]
-    # L follows RA as far as increasing/decreasing, see AIPS Meme 27, change in alpha
-    # definition three lines below Figure 2 and the first of the pair of equations 10.
-    l_coords = [
-        (i - image_size[0] // 2) * (-1) * abs(sky_image_cell_size[0])
-        for i in range(image_size[0])
-    ]
-    m_coords = [
-        (i - image_size[1] // 2) * abs(sky_image_cell_size[1])
-        for i in range(image_size[1])
-    ]
-    # im_size_wave = 1 / np.array(sky_image_cell_size)
-    # uv_cell_size = im_size_wave / np.array(image_size)
-    uv_cell_size = 1 / (np.array(image_size) * np.array(sky_image_cell_size))
-    u_coords = [
-        (i - image_size[0] // 2) * abs(uv_cell_size[0]) for i in range(image_size[0])
-    ]
-    v_coords = [
-        (i - image_size[1] // 2) * abs(uv_cell_size[1]) for i in range(image_size[1])
-    ]
-    coords = {
-        "time": some_coords["time"],
-        "polarization": pol_coords,
-        "frequency": some_coords["chan"],
-        "velocity": ("frequency", some_coords["vel"]),
-        "l": l_coords,
-        "m": m_coords,
-        "u": u_coords,
-        "v": v_coords,
-    }
-    if do_sky_coords:
-        coords["right_ascension"] = (("l", "m"), long)
-        coords["declination"] = (("l", "m"), lat)
-    xds = xr.Dataset(coords=coords)
-    attr_note = _l_m_attr_notes()
-    xds.l.attrs = {
-        "type": "quantity",
-        "crval": 0.0,
-        "cdelt": -abs(sky_image_cell_size[0]),
-        "units": "rad",
-        "note": attr_note["l"],
-    }
-    xds.m.attrs = {
-        "type": "quantity",
-        "crval": 0.0,
-        "cdelt": abs(sky_image_cell_size[1]),
-        "units": "rad",
-        "note": attr_note["m"],
-    }
-    xds.u.attrs = {
-        "crval": 0.0,
-        "cdelt": abs(uv_cell_size[0]),
-    }
-    xds.v.attrs = {
-        "crval": 0.0,
-        "cdelt": abs(uv_cell_size[1]),
-    }
-    _add_common_attrs(
-        xds,
-        some_coords["restfreq"],
-        spectral_reference,
-        direction_reference,
+    xds = _make_empty_sky_image(
         phase_center,
+        image_size,
         sky_image_cell_size,
+        chan_coords,
+        pol_coords,
+        time_coords,
+        direction_reference,
         projection,
+        spectral_reference,
+        do_sky_coords,
     )
+    xds = _make_uv_coords(xds, image_size, sky_image_cell_size)
+    xds = _move_beam_param_dim_coord(xds)
     return xds

xradio/image/image.py

@@ -82,7 +82,7 @@ def read_image(
         )
         do_casa = False
     if do_casa:
-        # next statement is for debug
+        # next statement is short circuit for debug, comment out when not debugging
         # return _read_casa_image(infile, chunks, verbose, do_sky_coords)
         try:
             return _read_casa_image(infile, chunks, verbose, do_sky_coords)
@@ -221,14 +221,14 @@ def make_empty_sky_image(
     chan_coords: Union[list, np.ndarray],
     pol_coords: Union[list, np.ndarray],
     time_coords: Union[list, np.ndarray],
-    direction_reference: str = "FK5",
+    direction_reference: str = "fK5",
     projection: str = "SIN",
     spectral_reference: str = "lsrk",
     do_sky_coords: bool = True,
 ) -> xr.Dataset:
     """
     Create an image xarray.Dataset with only coordinates (no datavariables).
-    The image dimensionality is time, pol, chan, l, m
+    The image dimensionality is time, frequency, polarization, l, m
 
     Parameters
     ----------
@@ -244,7 +244,7 @@ def make_empty_sky_image(
         The polarization code for each image polarization.
     time_coords : list or np.ndarray
         The time for each temporal plane in MJD.
-    direction_reference : str, default = 'FK5'
+    direction_reference : str, default = 'fk5'
     projection : str, default = 'SIN'
     spectral_reference : str, default = 'lsrk'
     do_sky_coords : bool
@@ -275,13 +275,13 @@ def make_empty_aperture_image(
     chan_coords: Union[List[float], np.ndarray],
     pol_coords: Union[List[str], np.ndarray],
     time_coords: Union[List[float], np.ndarray],
-    direction_reference: str = "FK5",
+    direction_reference: str = "fk5",
     projection: str = "SIN",
     spectral_reference: str = "lsrk",
 ) -> xr.Dataset:
     """
     Create an aperture (uv) mage xarray.Dataset with only coordinates (no datavariables).
-    The image dimensionality is time, pol, chan, u, v
+    The image dimensionality is time, frequency, polarization, u, v
 
     Parameters
     ----------
@@ -297,56 +297,7 @@ def make_empty_aperture_image(
         The polarization code for each image polarization.
     time_coords : list or np.ndarray
         The time for each temporal plane in MJD.
-    direction_reference : str, default = 'FK5'
-    projection : str, default = 'SIN'
-    spectral_reference : str, default = 'lsrk'
-    Returns
-    -------
-    xarray.Dataset
-    """
-    return _make_empty_aperture_image(
-        phase_center,
-        image_size,
-        sky_image_cell_size,
-        chan_coords,
-        pol_coords,
-        time_coords,
-        direction_reference,
-        projection,
-        spectral_reference,
-    )
-
-
-def make_empty_aperture_image(
-    phase_center: Union[List[float], np.ndarray],
-    image_size: Union[List[int], np.ndarray],
-    sky_image_cell_size: Union[List[float], np.ndarray],
-    chan_coords: Union[List[float], np.ndarray],
-    pol_coords: Union[List[str], np.ndarray],
-    time_coords: Union[List[float], np.ndarray],
-    direction_reference: str = "FK5",
-    projection: str = "SIN",
-    spectral_reference: str = "lsrk",
-) -> xr.Dataset:
-    """
-    Create an aperture (uv) mage xarray.Dataset with only coordinates (no datavariables).
-    The image dimensionality is time, pol, chan, u, v
-
-    Parameters
-    ----------
-    phase_center : array of float, length = 2, units = rad
-        Image phase center.
-    image_size : array of int, length = 2
-        Number of x and y axis pixels in image.
-    sky_image_cell_size : array of float, length = 2, units = rad
-        Cell size of x and y axis pixels in sky image, used to get cell size in uv image
-    chan_coords : list or np.ndarray
-        The center frequency in Hz of each image channel.
-    pol_coords : list or np.ndarray
-        The polarization code for each image polarization.
-    time_coords : list or np.ndarray
-        The time for each temporal plane in MJD.
-    direction_reference : str, default = 'FK5'
+    direction_reference : str, default = 'fk5'
     projection : str, default = 'SIN'
     spectral_reference : str, default = 'lsrk'
     Returns
@@ -373,14 +324,14 @@ def make_empty_lmuv_image(
     chan_coords: Union[List[float], np.ndarray],
     pol_coords: Union[List[float], np.ndarray],
     time_coords: Union[List[float], np.ndarray],
-    direction_reference: str = "FK5",
+    direction_reference: str = "fk5",
     projection: str = "SIN",
     spectral_reference: str = "lsrk",
     do_sky_coords: bool = True,
 ) -> xr.Dataset:
     """
     Create an image xarray.Dataset with only coordinates (no datavariables).
-    The image dimensionality is time, pol, chan, l, m, u, v
+    The image dimensionality is time, frequency, polarization, l, m, u, v
 
     Parameters
     ----------
@@ -397,7 +348,7 @@ def make_empty_lmuv_image(
         The polarization code for each image polarization.
     time_coords : list or np.ndarray
         The time for each temporal plane in MJD.
-    direction_reference : str, default = 'FK5'
+    direction_reference : str, default = 'fk5'
     projection : str, default = 'SIN'
     spectral_reference : str, default = 'lsrk'
     do_sky_coords : bool

xradio/measurement_set/__init__.py

@@ -1,17 +1,22 @@
-from .processing_set import ProcessingSet
+"""
+Processing Set and Measurement Set v4 API. Includes functions and classes to open, load,
+convert, and retrieve information from Processing Set and Measurement Sets nodes of the
+Processing Set DataTree
+"""
+
+from .processing_set_xdt import *
 from .open_processing_set import open_processing_set
-from .load_processing_set import load_processing_set, ProcessingSetIterator
+from .load_processing_set import load_processing_set  # , ProcessingSetIterator
 from .convert_msv2_to_processing_set import (
     convert_msv2_to_processing_set,
     estimate_conversion_memory_and_cores,
 )
-from .measurement_set_xds import MeasurementSetXds
-
+from .measurement_set_xdt import MeasurementSetXdt
 from .schema import SpectrumXds, VisibilityXds
 
 __all__ = [
-    "ProcessingSet",
-    "MeasurementSetXds",
+    "ProcessingSetXdt",
+    "MeasurementSetXdt",
     "open_processing_set",
     "load_processing_set",
     "ProcessingSetIterator",