xarray-ms 0.2.0__py3-none-any.whl

xarray_ms/__init__.py

@@ -0,0 +1,3 @@
+__all__ = ["xds_from_zarr", "xdt_from_zarr", "xds_to_zarr", "xdt_to_zarr"]
+
+from xarray_ms.core import xds_from_zarr, xds_to_zarr, xdt_from_zarr, xdt_to_zarr

xarray_ms/backend/msv2/antenna_dataset_factory.py

@@ -0,0 +1,36 @@
+from typing import Mapping
+
+from xarray import Dataset, Variable
+
+from xarray_ms.backend.msv2.structure import MSv2StructureFactory
+
+
+class AntennaDatasetFactory:
+  _structure_factory: MSv2StructureFactory
+
+  def __init__(self, structure_factory: MSv2StructureFactory):
+    self._structure_factory = structure_factory
+
+  def get_dataset(self) -> Mapping[str, Variable]:
+    ants = self._structure_factory()._ant
+
+    import pyarrow.compute as pac
+
+    ant_pos = pac.list_flatten(ants["POSITION"]).to_numpy().reshape(-1, 3)
+
+    return Dataset(
+      data_vars={
+        "ANTENNA_POSITION": Variable(
+          ("antenna_name", "cartesian_pos_label/ellipsoid_pos_label"), ant_pos
+        )
+      },
+      coords={
+        "antenna_name": Variable("antenna_name", ants["NAME"].to_numpy()),
+        "station": Variable(
+          "antenna_name", ants["STATION"].to_numpy(), {"coordinates": "station"}
+        ),
+        "mount": Variable(
+          "antenna_name", ants["MOUNT"].to_numpy(), {"coordinates": "mount"}
+        ),
+      },
+    )

xarray_ms/backend/msv2/array.py

@@ -0,0 +1,75 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any, Callable, Tuple
+
+import numpy as np
+from xarray.backends import BackendArray
+from xarray.core.indexing import IndexingSupport, explicit_indexing_adapter
+
+if TYPE_CHECKING:
+  import numpy.typing as npt
+
+  from xarray_ms.backend.msv2.structure import MSv2StructureFactory, PartitionKeyT
+  from xarray_ms.backend.msv2.table_factory import TableFactory
+
+
+def slice_length(s, max_len):
+  start, stop, step = s.indices(max_len)
+  if step != 1:
+    raise NotImplementedError(f"Slicing with steps {s} other than 1 not supported")
+  return stop - start
+
+
+class MSv2Array(BackendArray):
+  """Backend array containing functionality for reading an MSv2 column"""
+
+  _table_factory: TableFactory
+  _structure_factory: MSv2StructureFactory
+  _partition: PartitionKeyT
+  _column: str
+  _shape: Tuple[int, ...]
+  _dtype: npt.DTypeLike
+  _default: Any | None
+  _transform: Callable[[npt.NDArray], npt.NDArray] | None
+
+  def __init__(
+    self,
+    table_factory: TableFactory,
+    structure_factory: MSv2StructureFactory,
+    partition: PartitionKeyT,
+    column: str,
+    shape: Tuple[int, ...],
+    dtype: npt.DTypeLike,
+    default: Any | None = None,
+    transform: Callable[[npt.NDArray], npt.NDArray] | None = None,
+  ):
+    self._table_factory = table_factory
+    self._structure_factory = structure_factory
+    self._partition = partition
+    self._column = column
+    self._default = default
+    self._transform = transform
+    self.shape = shape
+    self.dtype = np.dtype(dtype)
+
+    assert len(shape) >= 2, "(time, baseline) required"
+
+  def __getitem__(self, key) -> npt.NDArray:
+    return explicit_indexing_adapter(
+      key, self.shape, IndexingSupport.OUTER, self._getitem
+    )
+
+  def _getitem(self, key) -> npt.NDArray:
+    assert len(key) == len(self.shape)
+    expected_shape = tuple(slice_length(k, s) for k, s in zip(key, self.shape))
+    # Map the (time, baseline) coordinates onto row indices
+    rows = self._structure_factory()[self._partition].row_map[key[:2]]
+    xkey = (rows.ravel(),) + key[2:]
+    row_shape = (rows.size,) + expected_shape[2:]
+    result = np.full(row_shape, self._default, dtype=self.dtype)
+    self._table_factory().getcol(self._column, xkey, result)
+    result = result.reshape(rows.shape + expected_shape[2:])
+    return self._transform(result) if self._transform else result
+
+  def set_transform(self, transform: Callable[[npt.NDArray], npt.NDArray]):
+    self._transform = transform

xarray_ms/backend/msv2/encoders.py

@@ -0,0 +1,199 @@
+"""
+TODO(sjperkins): This implementation is incomplete and
+will need refactoring.
+
+In particular, the reference columns and codes for more complex Measures
+are not yet handled.
+
+This logic can be found here:
+
+casacore/measures/TableMeasRefDesc/TableMeasRefDesc.cc
+TableMeasRefDesc::TableMeasRefDesc
+"""
+
+from __future__ import annotations
+
+import sys
+from datetime import datetime
+from typing import TYPE_CHECKING, Any, Dict, List
+
+if TYPE_CHECKING:
+  if sys.version_info >= (3, 11):
+    from typing import Self
+  else:
+    from typing_extensions import Self
+
+
+import numpy as np
+import numpy.testing as npt
+from xarray import Variable
+from xarray.coding.variables import (
+  VariableCoder,
+  unpack_for_decoding,
+  unpack_for_encoding,
+)
+
+from xarray_ms.backend.msv2.array import MSv2Array
+from xarray_ms.errors import MissingMeasuresInfo, MissingQuantumUnits
+
+if TYPE_CHECKING:
+  from xarray.coding.variables import T_Name
+
+  from xarray_ms.casa_types import ColumnDesc
+
+
+class CasaCoder(VariableCoder):
+  """Base class for CASA Measures Coders"""
+
+  _column: str
+  _column_descs: Dict[str, ColumnDesc]
+
+  def __init__(self, column: str, column_descs: Dict[str, ColumnDesc]):
+    assert column in column_descs
+    self._column = column
+    self._column_descs = column_descs
+
+  @property
+  def column(self) -> str:
+    """Returns the column"""
+    return self._column
+
+  @property
+  def column_descs(self) -> Dict[str, ColumnDesc]:
+    """Returns the column descriptors"""
+    return self._column_descs
+
+  @property
+  def column_desc(self) -> ColumnDesc:
+    """Returns the column descriptor"""
+    try:
+      return self._column_descs[self._column]
+    except KeyError:
+      raise KeyError(f"No Column Descriptor exist for {self.column}")
+
+  @property
+  def measinfo(self) -> Dict[str, Any]:
+    """Returns the MEASINFO keyword in the column descriptor"""
+    kw = self.column_desc.keywords
+    try:
+      return kw["MEASINFO"]
+    except KeyError:
+      raise MissingMeasuresInfo(f"No MEASINFO found for {self.column}")
+
+  @property
+  def quantum_units(self) -> List[str]:
+    """REturns the QuantumUnits keyword in the column descriptor"""
+    kw = self.column_desc.keywords
+    try:
+      return kw["QuantumUnits"]
+    except KeyError:
+      raise MissingQuantumUnits(f"No QuantumUnits found for {self.column}")
+
+
+class QuantityCoder(CasaCoder):
+  def encode(self, variable: Variable, name: T_Name = None) -> Variable:
+    dims, data, attrs, encoding = unpack_for_encoding(variable)
+    attrs.pop("type", None)
+    attrs.pop("units", None)
+    return Variable(dims, data, attrs, encoding, fastpath=True)
+
+  def decode(self, variable: Variable, name: T_Name = None) -> Variable:
+    dims, data, attrs, encoding = unpack_for_decoding(variable)
+    attrs["type"] = "quantity"
+    attrs["units"] = self.quantum_units
+    return Variable(dims, data, attrs, encoding, fastpath=True)
+
+
+class TimeCoder(CasaCoder):
+  """Dispatches encoding functionality to sub-classes"""
+
+  @classmethod
+  def from_time_coder(cls, time_coder: TimeCoder) -> Self:
+    return cls(time_coder._column, time_coder._column_descs)
+
+  def dispatched_coder(self) -> TimeCoder:
+    measures = self.measinfo
+    assert measures["type"] == "epoch"
+    ref = measures["Ref"].upper()
+    if ref == "UTC":
+      cls = UTCCoder
+    else:
+      raise NotImplementedError(measures["Ref"])
+
+    return cls.from_time_coder(self)
+
+  def encode(self, variable: Variable, name: T_Name = None) -> Variable:
+    return self.dispatched_coder().encode(variable, name)
+
+  def decode(self, variable: Variable, name: T_Name = None) -> Variable:
+    return self.dispatched_coder().decode(variable, name)
+
+
+class UTCCoder(TimeCoder):
+  """Encode MJD UTC"""
+
+  MJD_EPOCH: datetime = datetime(1858, 11, 17)
+  UTC_EPOCH: datetime = datetime(1970, 1, 1)
+  MJD_OFFSET_SECONDS: float = (UTC_EPOCH - MJD_EPOCH).total_seconds()
+
+  def encode(self, variable: Variable, name: T_Name = None) -> Variable:
+    """Convert UTC in seconds to Modified Julian Date"""
+    dims, data, attrs, encoding = unpack_for_encoding(variable)
+    attrs.pop("type", None)
+    attrs.pop("units", None)
+    attrs.pop("scale", None)
+    attrs.pop("format", None)
+
+    if isinstance(data, MSv2Array):
+      data.set_transform(UTCCoder.encode_array)
+    elif isinstance(data, np.ndarray):
+      data = UTCCoder.encode_array(data)
+    else:
+      raise TypeError(f"Unknown data type {type(data)}")
+
+    return Variable(dims, data, attrs, encoding, fastpath=True)
+
+  def decode(self, variable: Variable, name: T_Name = None) -> Variable:
+    """Convert Modified Julian Date in seconds to UTC in seconds"""
+    dims, data, attrs, encoding = unpack_for_decoding(variable)
+    attrs["type"] = "time"
+    attrs["units"] = self.quantum_units
+    attrs["scale"] = "utc"
+    attrs["format"] = "unix"
+
+    if isinstance(data, MSv2Array):
+      data.set_transform(UTCCoder.decode_array)
+    elif isinstance(data, np.ndarray):
+      data = UTCCoder.decode_array(data)
+    else:
+      raise TypeError(f"Unknown data type {type(data)}")
+
+    return Variable(dims, data, attrs, encoding, fastpath=True)
+
+  @staticmethod
+  def encode_array(data: npt.NDArray) -> npt.NDArray:
+    return data + UTCCoder.MJD_OFFSET_SECONDS
+
+  @staticmethod
+  def decode_array(data: npt.NDArray) -> npt.NDArray:
+    return data - UTCCoder.MJD_OFFSET_SECONDS
+
+
+class SpectralCoordCoder(CasaCoder):
+  """Encode Measures Spectral Coordinates"""
+
+  def encode(self, variable: Variable, name: T_Name = None) -> Variable:
+    dims, data, attrs, encoding = unpack_for_encoding(variable)
+    attrs.pop("type", None)
+    return Variable(dims, data, attrs, encoding, fastpath=True)
+
+  def decode(self, variable: Variable, name: T_Name = None) -> Variable:
+    dims, data, attrs, encoding = unpack_for_decoding(variable)
+    measures = self.measinfo
+    assert measures["type"] == "frequency"
+    attrs["type"] = "spectral_coord"
+    # TODO(sjperkins): topo is hard-coded here and will almost
+    # certainly need extra work to support other frames
+    attrs["frame"] = "topo"
+    attrs["units"] = self.quantum_units
+    return Variable(dims, data, attrs, encoding, fastpath=True)

xarray_ms/backend/msv2/entrypoint.py

@@ -0,0 +1,354 @@
+from __future__ import annotations
+
+import os
+import warnings
+from datetime import datetime, timezone
+from typing import TYPE_CHECKING, Any, Dict, Iterable, List, Tuple
+from uuid import uuid4
+
+import xarray
+from arcae.lib.arrow_tables import Table
+from xarray.backends import BackendEntrypoint
+from xarray.backends.common import AbstractWritableDataStore, _normalize_path
+from xarray.backends.store import StoreBackendEntrypoint
+from xarray.core.dataset import Dataset
+from xarray.core.datatree import DataTree
+from xarray.core.utils import try_read_magic_number_from_file_or_path
+
+from xarray_ms.backend.msv2.antenna_dataset_factory import AntennaDatasetFactory
+from xarray_ms.backend.msv2.main_dataset_factory import MainDatasetFactory
+from xarray_ms.backend.msv2.structure import (
+  DEFAULT_PARTITION_COLUMNS,
+  MSv2Structure,
+  MSv2StructureFactory,
+)
+from xarray_ms.backend.msv2.table_factory import TableFactory
+from xarray_ms.errors import InvalidPartitionKey
+from xarray_ms.utils import format_docstring
+
+if TYPE_CHECKING:
+  from io import BufferedIOBase
+
+  from xarray.backends.common import AbstractDataStore
+
+  from xarray_ms.backend.msv2.structure import DEFAULT_PARTITION_COLUMNS, PartitionKeyT
+
+
+def promote_chunks(
+  structure: MSv2Structure, chunks: Dict | None
+) -> Dict[PartitionKeyT, Dict[str, int]] | None:
+  """Promotes a chunks dictionary into a
+  :code:`{partition_key: chunks}` dictionary.
+  """
+  if chunks is None:
+    return None
+
+  # Base case, no chunking
+  return_chunks: Dict[PartitionKeyT, Dict[str, int]] = {k: {} for k in structure.keys()}
+
+  if all(isinstance(k, str) for k in chunks.keys()):
+    # All keys are strings, try promote them to partition keys
+    # keys, may resolve to multiple partition keys
+    try:
+      crkeys = list(map(structure.resolve_key, chunks.keys()))
+    except InvalidPartitionKey:
+      # Apply a chunk dictionary to all partitions in the structure
+      return {k: chunks for k in structure.keys()}
+    else:
+      return_chunks.update((k, v) for rk, v in zip(crkeys, chunks.values()) for k in rk)
+  else:
+    for k, v in chunks.items():
+      rkeys = structure.resolve_key(k)
+      return_chunks.update((k, v) for k in rkeys)
+
+  return return_chunks
+
+
+def initialise_default_args(
+  ms: str,
+  ninstances: int,
+  auto_corrs: bool,
+  epoch: str | None,
+  table_factory: TableFactory | None,
+  partition_columns: List[str] | None,
+  structure_factory: MSv2StructureFactory | None,
+) -> Tuple[str, TableFactory, List[str], MSv2StructureFactory]:
+  """
+  Ensures consistency when initialising default arguments from multiple locations
+  """
+  if not os.path.exists(ms):
+    raise ValueError(f"MS {ms} does not exist")
+
+  table_factory = table_factory or TableFactory(
+    Table.from_filename,
+    ms,
+    ninstances=ninstances,
+    readonly=True,
+    lockoptions="nolock",
+  )
+  epoch = epoch or uuid4().hex[:8]
+  partition_columns = partition_columns or DEFAULT_PARTITION_COLUMNS
+  structure_factory = structure_factory or MSv2StructureFactory(
+    table_factory, partition_columns, auto_corrs=auto_corrs
+  )
+  return epoch, table_factory, partition_columns, structure_factory
+
+
+class MSv2Store(AbstractWritableDataStore):
+  """Store for reading and writing MSv2 data"""
+
+  __slots__ = (
+    "_table_factory",
+    "_structure_factory",
+    "_partition_columns",
+    "_partition_key",
+    "_auto_corrs",
+    "_ninstances",
+    "_epoch",
+  )
+
+  _table_factory: TableFactory
+  _structure_factory: MSv2StructureFactory
+  _partition_columns: List[str]
+  _partition: PartitionKeyT
+  _autocorrs: bool
+  _ninstances: int
+  _epoch: str
+
+  def __init__(
+    self,
+    table_factory: TableFactory,
+    structure_factory: MSv2StructureFactory,
+    partition_columns: List[str],
+    partition_key: PartitionKeyT,
+    auto_corrs: bool,
+    ninstances: int,
+    epoch: str,
+  ):
+    self._table_factory = table_factory
+    self._structure_factory = structure_factory
+    self._partition_columns = partition_columns
+    self._partition_key = partition_key
+    self._auto_corrs = auto_corrs
+    self._ninstances = ninstances
+    self._epoch = epoch
+
+  @classmethod
+  def open(
+    cls,
+    ms: str,
+    drop_variables=None,
+    partition_columns: List[str] | None = None,
+    partition_key: PartitionKeyT | None = None,
+    auto_corrs: bool = True,
+    ninstances: int = 1,
+    epoch: str | None = None,
+    structure_factory: MSv2StructureFactory | None = None,
+  ):
+    if not isinstance(ms, str):
+      raise ValueError("Measurement Sets paths must be strings")
+
+    epoch, table_factory, partition_columns, structure_factory = (
+      initialise_default_args(
+        ms,
+        ninstances,
+        auto_corrs,
+        epoch,
+        None,
+        partition_columns,
+        structure_factory,
+      )
+    )
+
+    structure = structure_factory()
+
+    if partition_key is None:
+      partition_key = next(iter(structure.keys()))
+      warnings.warn(
+        f"No partition_key was supplied. Selected first partition {partition_key}"
+      )
+    elif partition_key not in structure:
+      raise ValueError(f"{partition_key} not in {list(structure.keys())}")
+
+    return cls(
+      table_factory,
+      structure_factory,
+      partition_columns=partition_columns,
+      partition_key=partition_key,
+      auto_corrs=auto_corrs,
+      ninstances=ninstances,
+      epoch=epoch,
+    )
+
+  def close(self, **kwargs):
+    pass
+
+  def get_variables(self):
+    return MainDatasetFactory(
+      self._partition_key, self._table_factory, self._structure_factory
+    ).get_variables()
+
+  def get_attrs(self):
+    try:
+      ddid = next(iter(v for k, v in self._partition_key if k == "DATA_DESC_ID"))
+    except StopIteration:
+      raise KeyError("DATA_DESC_ID not found in partition")
+
+    antenna_factory = AntennaDatasetFactory(self._structure_factory)
+    ds = antenna_factory.get_dataset()
+
+    return {
+      "antenna_xds": ds,
+      "version": "0.0.1",
+      "creation_date": datetime.now(timezone.utc).isoformat(),
+      "data_description_id": ddid,
+    }
+
+  def get_dimensions(self):
+    return None
+
+  def get_encoding(self):
+    return {}
+
+
+class MSv2PartitionEntryPoint(BackendEntrypoint):
+  open_dataset_parameters = [
+    "filename_or_obj",
+    "partition_columns" "partition_key",
+    "auto_corrs",
+    "ninstances",
+    "epoch",
+    "structure_factory",
+  ]
+  description = "Opens v2 CASA Measurement Sets in Xarray"
+  url = "https://link_to/your_backend/documentation"
+
+  def guess_can_open(
+    self, filename_or_obj: str | os.PathLike[Any] | BufferedIOBase | AbstractDataStore
+  ) -> bool:
+    """Return true if this is a CASA table"""
+    if not isinstance(filename_or_obj, (str, os.PathLike)):
+      return False
+
+    # CASA Tables are directories containing a table.dat file
+    table_path = os.path.join(_normalize_path(filename_or_obj), "table.dat")
+    if not os.path.exists(table_path):
+      return False
+
+    # Check the magic number
+    if magic := try_read_magic_number_from_file_or_path(table_path, count=4):
+      return magic == b"\xbe\xbe\xbe\xbe"
+
+    return False
+
+  @format_docstring(DEFAULT_PARTITION_COLUMNS=DEFAULT_PARTITION_COLUMNS)
+  def open_dataset(
+    self,
+    filename_or_obj: str | os.PathLike[Any] | BufferedIOBase | AbstractDataStore,
+    *,
+    drop_variables: str | Iterable[str] | None = None,
+    partition_columns: List[str] | None = None,
+    partition_key: PartitionKeyT | None = None,
+    auto_corrs: bool = True,
+    ninstances: int = 8,
+    epoch: str | None = None,
+    structure_factory: MSv2StructureFactory | None = None,
+  ) -> Dataset:
+    """Create a :class:`~xarray.Dataset` presenting an MSv4 view
+    over a partition of a MSv2 CASA Measurement Set
+
+    Args:
+      filename_or_obj: The path to the MSv2 CASA Measurement Set file.
+      drop_variables: Variables to drop from the dataset.
+      partition_columns: The columns to use for partitioning the Measurement set.
+        Defaults to :code:`{DEFAULT_PARTITION_COLUMNS}`.
+      partition_key: A key corresponding to an individual partition.
+        For example :code:`(('DATA_DESC_ID', 0), ('FIELD_ID', 0))`.
+        If :code:`None`, the first partition will be opened.
+      auto_corrs: Include/Exclude auto-correlations.
+      ninstances: The number of Measurement Set instances to open for parallel I/O.
+      epoch: A unique string identifying the creation of this Dataset.
+        This should not normally need to be set by the user
+      structure_factory: A factory for creating MSv2Structure objects.
+        This should not normally need to be set by the user
+
+    Returns:
+      A :class:`~xarray.Dataset` referring to the unique
+      partition specified by :code:`partition_columns` and :code:`partition_key`.
+    """
+    filename_or_obj = _normalize_path(filename_or_obj)
+    store = MSv2Store.open(
+      filename_or_obj,
+      drop_variables=drop_variables,
+      partition_columns=partition_columns,
+      partition_key=partition_key,
+      auto_corrs=auto_corrs,
+      ninstances=ninstances,
+      epoch=epoch,
+      structure_factory=structure_factory,
+    )
+    store_entrypoint = StoreBackendEntrypoint()
+    return store_entrypoint.open_dataset(store)
+
+  @format_docstring(DEFAULT_PARTITION_COLUMNS=DEFAULT_PARTITION_COLUMNS)
+  def open_datatree(
+    self,
+    filename_or_obj: str | os.PathLike[Any] | BufferedIOBase | AbstractDataStore,
+    *,
+    drop_variables: str | Iterable[str] | None = None,
+    partition_columns: List[str] | None = None,
+    auto_corrs: bool = True,
+    ninstances: int = 8,
+    epoch: str | None = None,
+    **kwargs,
+  ) -> DataTree:
+    """Create a :class:`~xarray.core.datatree.DataTree` presenting an MSv4 view
+    over multiple partitions of a MSv2 CASA Measurement Set.
+
+    Args:
+      filename_or_obj: The path to the MSv2 CASA Measurement Set file.
+      drop_variables: Variables to drop from the dataset.
+      partition_columns: The columns to use for partitioning the Measurement set.
+        Defaults to :code:`{DEFAULT_PARTITION_COLUMNS}`.
+      auto_corrs: Include/Exclude auto-correlations.
+      ninstances: The number of Measurement Set instances to open for parallel I/O.
+      epoch: A unique string identifying the creation of this Dataset.
+        This should not normally need to be set by the user
+
+    Returns:
+      An xarray :class:`~xarray.core.datatree.DataTree`
+    """
+    if isinstance(filename_or_obj, os.PathLike):
+      ms = str(filename_or_obj)
+    elif isinstance(filename_or_obj, str):
+      ms = filename_or_obj
+    else:
+      raise ValueError("Measurement Set paths must be strings")
+
+    epoch, _, partition_columns, structure_factory = initialise_default_args(
+      ms, ninstances, auto_corrs, epoch, None, partition_columns, None
+    )
+
+    structure = structure_factory()
+    datasets = {}
+    chunks = kwargs.pop("chunks", None)
+    pchunks = promote_chunks(structure, chunks)
+
+    for partition_key in structure:
+      ds = xarray.open_dataset(
+        ms,
+        drop_variables=drop_variables,
+        partition_columns=partition_columns,
+        partition_key=partition_key,
+        auto_corrs=auto_corrs,
+        ninstances=ninstances,
+        epoch=epoch,
+        structure_factory=structure_factory,
+        chunks=None if pchunks is None else pchunks[partition_key],
+        **kwargs,
+      )
+
+      key = ",".join(f"{k}={v}" for k, v in sorted(partition_key))
+      datasets[key] = ds
+
+    return DataTree.from_dict(datasets)