xarray-ms 0.1.9__tar.gz → 0.2.0__tar.gz

xarray_ms-0.2.0/LICENSE

@@ -0,0 +1,28 @@
+BSD 3-Clause License
+
+Copyright (c) 2024, Rhodes University Centre for Radio Astronomy Techniques & Technologies (RATT)
+
+Redistribution and use in source and binary forms, with or without
+modification, are permitted provided that the following conditions are met:
+
+1. Redistributions of source code must retain the above copyright notice, this
+   list of conditions and the following disclaimer.
+
+2. Redistributions in binary form must reproduce the above copyright notice,
+   this list of conditions and the following disclaimer in the documentation
+   and/or other materials provided with the distribution.
+
+3. Neither the name of the copyright holder nor the names of its
+   contributors may be used to endorse or promote products derived from
+   this software without specific prior written permission.
+
+THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
+FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
+DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
+SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
+CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
+OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

xarray_ms-0.2.0/PKG-INFO

@@ -0,0 +1,149 @@
+Metadata-Version: 2.1
+Name: xarray-ms
+Version: 0.2.0
+Summary: xarray MSv4 views over MSv2 Measurement Sets
+Author: Simon Perkins
+Author-email: simon.perkins@gmail.com
+Requires-Python: >=3.10,<4.0
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.10
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Provides-Extra: testing
+Requires-Dist: arcae (>=0.2.4,<0.3.0)
+Requires-Dist: cacheout (>=0.16.0,<0.17.0)
+Requires-Dist: dask[testing] (>=2024.5.0,<2025.0.0) ; extra == "testing"
+Requires-Dist: distributed[testing] (>=2024.5.0,<2025.0.0) ; extra == "testing"
+Requires-Dist: pytest[testing] (>=8.0.0,<9.0.0) ; extra == "testing"
+Requires-Dist: typing-extensions (>=4.12.2,<5.0.0) ; python_version < "3.11"
+Requires-Dist: xarray (>=2024.3.0,<2025.0.0)
+Requires-Dist: zarr[testing] (>=2.18.3,<3.0.0) ; extra == "testing"
+Description-Content-Type: text/x-rst
+
+=========
+xarray-ms
+=========
+
+.. image:: https://img.shields.io/pypi/v/xarray-ms.svg
+        :target: https://pypi.python.org/pypi/xarray-ms
+
+.. image:: https://github.com/ratt-ru/xarray-ms/actions/workflows/ci.yml/badge.svg
+        :target: https://github.com/ratt-ru/xarray-ms/actions/workflows/ci.yml
+
+.. image:: https://readthedocs.org/projects/xarray-ms/badge/?version=latest
+        :target: https://xarray-ms.readthedocs.io/en/latest/?badge=latest
+        :alt: Documentation Status
+
+====
+
+xarray-ms presents a Measurement Set v4 view (MSv4) over
+`CASA Measurement Sets <https://casa.nrao.edu/Memos/229.html>`_ (MSv2).
+It provides access to MSv2 data via the xarray API, allowing MSv4 compliant applications
+to be developed on well-understood MSv2 data.
+
+.. code-block:: python
+
+  >>> import xarray_ms
+  >>> import xarray
+  >>> ds = xarray.open_dataset("/data/L795830_SB001_uv.MS/",
+                               chunks={"time": 2000, "baseline": 1000})
+  >>> ds
+    <xarray.Dataset> Size: 70GB
+    Dimensions:                     (time: 28760, baseline: 2775, frequency: 16,
+                                    polarization: 4, uvw_label: 3)
+    Coordinates:
+        antenna1_name               (baseline) object 22kB dask.array<chunksize=(1000,), meta=np.ndarray>
+        antenna2_name               (baseline) object 22kB dask.array<chunksize=(1000,), meta=np.ndarray>
+        baseline_id                 (baseline) int64 22kB dask.array<chunksize=(1000,), meta=np.ndarray>
+      * frequency                   (frequency) float64 128B 1.202e+08 ... 1.204e+08
+      * polarization                (polarization) <U2 32B 'XX' 'XY' 'YX' 'YY'
+      * time                        (time) float64 230kB 1.601e+09 ... 1.601e+09
+    Dimensions without coordinates: baseline, uvw_label
+    Data variables:
+        EFFECTIVE_INTEGRATION_TIME  (time, baseline) float64 638MB dask.array<chunksize=(2000, 1000), meta=np.ndarray>
+        FLAG                        (time, baseline, frequency, polarization) uint8 5GB dask.array<chunksize=(2000, 1000, 16, 4), meta=np.ndarray>
+        TIME_CENTROID               (time, baseline) float64 638MB dask.array<chunksize=(2000, 1000), meta=np.ndarray>
+        UVW                         (time, baseline, uvw_label) float64 2GB dask.array<chunksize=(2000, 1000, 3), meta=np.ndarray>
+        VISIBILITY                  (time, baseline, frequency, polarization) complex64 41GB dask.array<chunksize=(2000, 1000, 16, 4), meta=np.ndarray>
+        WEIGHT                      (time, baseline, frequency, polarization) float32 20GB dask.array<chunksize=(2000, 1000, 16, 4), meta=np.ndarray>
+    Attributes:
+        antenna_xds:          <xarray.Dataset> Size: 4kB\nDimensions:           (...
+        version:              0.0.1
+        creation_date:        2024-09-10T14:29:22.587984+00:00
+        data_description_id:  0
+
+Measurement Set v4
+------------------
+
+NRAO_/SKAO_ are developing a new xarray-based `Measurement Set v4 specification <msv4-spec_>`_.
+While there are many changes some of the major highlights are:
+
+* xarray_ is used to define the specification.
+* MSv4 data consists of Datasets of ndarrays on a regular time-channel grid.
+  MSv2 data is tabular and, while in many instances the time-channel grid is regular,
+  this was not guaranteed, especially after MSv2 datasets had been transformed by various tasks.
+
+
+xarray_ Datasets are self-describing and they are therefore easier to reason about and work with.
+Additionally, the regularity of data will make writing MSv4-based software less complex.
+
+xradio
+------
+
+`casangi/xradio <xradio_>`_ provides a reference implementation that converts
+CASA v2 Measurement Sets to Zarr v4 Measurement Sets using the python-casacore_
+package.
+
+Why xarray-ms?
+--------------
+
+* By developing against an MSv4 xarray view over MSv2 data,
+  developers can develop applications on well-understood data,
+  and then seamlessly transition to newer formats.
+  Data can also be exported to newer formats (principally zarr_) via xarray's
+  native I/O routines.
+  However, the xarray view of either format looks the same to the software developer.
+
+* xarray-ms builds on xarray's
+  `backend API <https://docs.xarray.dev/en/stable/internals/how-to-add-new-backend.html>`_:
+  Implementing a formal CASA MSv2 backend has a number of benefits:
+
+  * xarray's internal I/O routines such as ``open_dataset`` and ``open_datatree``
+    can dispatch to the backend to load data.
+  * Similarly xarray's `lazy loading mechanism <xarray_lazy_>`_ dispatches
+    through the backend.
+  * Automatic access to any `chunked array types <xarray_chunked_arrays_>`_
+    supported by xarray including, but not limited to dask_.
+  * Arbitrary chunking along any xarray dimension.
+
+* xarray-ms uses arcae_, a high-performance backend to CASA Tables implementing
+  a subset of python-casacore_'s interface.
+* Some limited support for irregular MSv2 data via padding.
+
+Work in Progress
+----------------
+
+The Measurement Set v4 specification is currently under active development.
+xarray-ms is currently under active development and does not yet
+have feature parity with xradio_.
+
+Most measures information and many secondary sub-tables are currently missing.
+However, the most important parts of the ``MAIN`` tables,
+as well as the ``ANTENNA``, ``POLARIZATON`` and ``SPECTRAL_WINDOW``
+sub-tables are implemented and should be sufficient
+for basic algorithm development.
+
+.. _SKAO: https://www.skao.int/
+.. _NRAO: https://public.nrao.edu/
+.. _msv4-spec: https://docs.google.com/spreadsheets/d/14a6qMap9M5r_vjpLnaBKxsR9TF4azN5LVdOxLacOX-s/
+.. _xradio: https://github.com/casangi/xradio
+.. _dask-ms: https://github.com/ratt-ru/dask-ms
+.. _arcae: https://github.com/ratt-ru/arcae
+.. _dask: https://www.dask.org/
+.. _python-casacore: https://github.com/casacore/python-casacore/
+.. _xarray: https://xarray.dev/
+.. _xarray_backend: https://docs.xarray.dev/en/stable/internals/how-to-add-new-backend.html
+.. _xarray_lazy: https://docs.xarray.dev/en/latest/internals/internal-design.html#lazy-indexing-classes
+.. _xarray_chunked_arrays: https://docs.xarray.dev/en/latest/internals/chunked-arrays.html
+.. _zarr: https://zarr.dev/
+

xarray_ms-0.2.0/README.rst

@@ -0,0 +1,126 @@
+=========
+xarray-ms
+=========
+
+.. image:: https://img.shields.io/pypi/v/xarray-ms.svg
+        :target: https://pypi.python.org/pypi/xarray-ms
+
+.. image:: https://github.com/ratt-ru/xarray-ms/actions/workflows/ci.yml/badge.svg
+        :target: https://github.com/ratt-ru/xarray-ms/actions/workflows/ci.yml
+
+.. image:: https://readthedocs.org/projects/xarray-ms/badge/?version=latest
+        :target: https://xarray-ms.readthedocs.io/en/latest/?badge=latest
+        :alt: Documentation Status
+
+====
+
+xarray-ms presents a Measurement Set v4 view (MSv4) over
+`CASA Measurement Sets <https://casa.nrao.edu/Memos/229.html>`_ (MSv2).
+It provides access to MSv2 data via the xarray API, allowing MSv4 compliant applications
+to be developed on well-understood MSv2 data.
+
+.. code-block:: python
+
+  >>> import xarray_ms
+  >>> import xarray
+  >>> ds = xarray.open_dataset("/data/L795830_SB001_uv.MS/",
+                               chunks={"time": 2000, "baseline": 1000})
+  >>> ds
+    <xarray.Dataset> Size: 70GB
+    Dimensions:                     (time: 28760, baseline: 2775, frequency: 16,
+                                    polarization: 4, uvw_label: 3)
+    Coordinates:
+        antenna1_name               (baseline) object 22kB dask.array<chunksize=(1000,), meta=np.ndarray>
+        antenna2_name               (baseline) object 22kB dask.array<chunksize=(1000,), meta=np.ndarray>
+        baseline_id                 (baseline) int64 22kB dask.array<chunksize=(1000,), meta=np.ndarray>
+      * frequency                   (frequency) float64 128B 1.202e+08 ... 1.204e+08
+      * polarization                (polarization) <U2 32B 'XX' 'XY' 'YX' 'YY'
+      * time                        (time) float64 230kB 1.601e+09 ... 1.601e+09
+    Dimensions without coordinates: baseline, uvw_label
+    Data variables:
+        EFFECTIVE_INTEGRATION_TIME  (time, baseline) float64 638MB dask.array<chunksize=(2000, 1000), meta=np.ndarray>
+        FLAG                        (time, baseline, frequency, polarization) uint8 5GB dask.array<chunksize=(2000, 1000, 16, 4), meta=np.ndarray>
+        TIME_CENTROID               (time, baseline) float64 638MB dask.array<chunksize=(2000, 1000), meta=np.ndarray>
+        UVW                         (time, baseline, uvw_label) float64 2GB dask.array<chunksize=(2000, 1000, 3), meta=np.ndarray>
+        VISIBILITY                  (time, baseline, frequency, polarization) complex64 41GB dask.array<chunksize=(2000, 1000, 16, 4), meta=np.ndarray>
+        WEIGHT                      (time, baseline, frequency, polarization) float32 20GB dask.array<chunksize=(2000, 1000, 16, 4), meta=np.ndarray>
+    Attributes:
+        antenna_xds:          <xarray.Dataset> Size: 4kB\nDimensions:           (...
+        version:              0.0.1
+        creation_date:        2024-09-10T14:29:22.587984+00:00
+        data_description_id:  0
+
+Measurement Set v4
+------------------
+
+NRAO_/SKAO_ are developing a new xarray-based `Measurement Set v4 specification <msv4-spec_>`_.
+While there are many changes some of the major highlights are:
+
+* xarray_ is used to define the specification.
+* MSv4 data consists of Datasets of ndarrays on a regular time-channel grid.
+  MSv2 data is tabular and, while in many instances the time-channel grid is regular,
+  this was not guaranteed, especially after MSv2 datasets had been transformed by various tasks.
+
+
+xarray_ Datasets are self-describing and they are therefore easier to reason about and work with.
+Additionally, the regularity of data will make writing MSv4-based software less complex.
+
+xradio
+------
+
+`casangi/xradio <xradio_>`_ provides a reference implementation that converts
+CASA v2 Measurement Sets to Zarr v4 Measurement Sets using the python-casacore_
+package.
+
+Why xarray-ms?
+--------------
+
+* By developing against an MSv4 xarray view over MSv2 data,
+  developers can develop applications on well-understood data,
+  and then seamlessly transition to newer formats.
+  Data can also be exported to newer formats (principally zarr_) via xarray's
+  native I/O routines.
+  However, the xarray view of either format looks the same to the software developer.
+
+* xarray-ms builds on xarray's
+  `backend API <https://docs.xarray.dev/en/stable/internals/how-to-add-new-backend.html>`_:
+  Implementing a formal CASA MSv2 backend has a number of benefits:
+
+  * xarray's internal I/O routines such as ``open_dataset`` and ``open_datatree``
+    can dispatch to the backend to load data.
+  * Similarly xarray's `lazy loading mechanism <xarray_lazy_>`_ dispatches
+    through the backend.
+  * Automatic access to any `chunked array types <xarray_chunked_arrays_>`_
+    supported by xarray including, but not limited to dask_.
+  * Arbitrary chunking along any xarray dimension.
+
+* xarray-ms uses arcae_, a high-performance backend to CASA Tables implementing
+  a subset of python-casacore_'s interface.
+* Some limited support for irregular MSv2 data via padding.
+
+Work in Progress
+----------------
+
+The Measurement Set v4 specification is currently under active development.
+xarray-ms is currently under active development and does not yet
+have feature parity with xradio_.
+
+Most measures information and many secondary sub-tables are currently missing.
+However, the most important parts of the ``MAIN`` tables,
+as well as the ``ANTENNA``, ``POLARIZATON`` and ``SPECTRAL_WINDOW``
+sub-tables are implemented and should be sufficient
+for basic algorithm development.
+
+.. _SKAO: https://www.skao.int/
+.. _NRAO: https://public.nrao.edu/
+.. _msv4-spec: https://docs.google.com/spreadsheets/d/14a6qMap9M5r_vjpLnaBKxsR9TF4azN5LVdOxLacOX-s/
+.. _xradio: https://github.com/casangi/xradio
+.. _dask-ms: https://github.com/ratt-ru/dask-ms
+.. _arcae: https://github.com/ratt-ru/arcae
+.. _dask: https://www.dask.org/
+.. _python-casacore: https://github.com/casacore/python-casacore/
+.. _xarray: https://xarray.dev/
+.. _xarray_backend: https://docs.xarray.dev/en/stable/internals/how-to-add-new-backend.html
+.. _xarray_lazy: https://docs.xarray.dev/en/latest/internals/internal-design.html#lazy-indexing-classes
+.. _xarray_chunked_arrays: https://docs.xarray.dev/en/latest/internals/chunked-arrays.html
+.. _zarr: https://zarr.dev/

xarray_ms-0.2.0/pyproject.toml

@@ -0,0 +1,96 @@
+[tool.poetry]
+name = "xarray-ms"
+version = "0.2.0"
+description = "xarray MSv4 views over MSv2 Measurement Sets"
+authors = ["Simon Perkins <simon.perkins@gmail.com>"]
+readme = "README.rst"
+
+[tool.poetry.dependencies]
+python = "^3.10"
+pytest = {version = "^8.0.0", optional = true, extras = ["testing"]}
+xarray = "^2024.3.0"
+dask = {version = "^2024.5.0", optional = true, extras = ["testing"]}
+distributed = {version = "^2024.5.0", optional = true, extras = ["testing"]}
+cacheout = "^0.16.0"
+arcae = "^0.2.4"
+typing-extensions = { version = "^4.12.2", python = "<3.11" }
+zarr = {version = "^2.18.3", optional = true, extras = ["testing"]}
+
+[tool.poetry.extras]
+testing = ["dask", "distributed", "pytest", "zarr"]
+
+[tool.poetry.plugins."xarray.backends"]
+"xarray-ms:msv2" = "xarray_ms.backend.msv2.entrypoint:MSv2PartitionEntryPoint"
+
+[tool.poetry.group.dev]
+optional = true
+
+[tool.poetry.group.dev.dependencies]
+pre-commit = "^3.8.0"
+tbump = "^6.11.0"
+
+[tool.poetry.group.doc]
+optional = true
+
+[tool.poetry.group.doc.dependencies]
+sphinx = "^8.0.2"
+pygments = "^2.18.0"
+sphinx-copybutton = "^0.5.2"
+pydata-sphinx-theme = "^0.15.4"
+
+[tool.ruff]
+line-length = 88
+indent-width = 2
+target-version = "py311"
+
+[tool.ruff.lint]
+select = ["F", "E", "W", "I001"]
+extend-select = ["I"]
+
+[build-system]
+requires = ["poetry-core"]
+build-backend = "poetry.core.masonry.api"
+[tool.tbump]
+# Uncomment this if your project is hosted on GitHub:
+# github_url = "https://github.com/<user or organization>/<project>/"
+
+[tool.tbump.version]
+current = "0.2.0"
+
+# Example of a semver regexp.
+# Make sure this matches current_version before
+# using tbump
+regex = '''
+  (?P<major>\d+)
+  \.
+  (?P<minor>\d+)
+  \.
+  (?P<patch>\d+)
+  '''
+
+[tool.tbump.git]
+message_template = "Bump to {new_version}"
+tag_template = "{new_version}"
+
+# For each file to patch, add a [[tool.tbump.file]] config
+# section containing the path of the file, relative to the
+# tbump.toml location.
+[[tool.tbump.file]]
+src = "pyproject.toml"
+
+[[tool.tbump.file]]
+src = "doc/source/conf.py"
+
+# You can specify a list of commands to
+# run after the files have been patched
+# and before the git commit is made
+
+#  [[tool.tbump.before_commit]]
+#  name = "check changelog"
+#  cmd = "grep -q {new_version} Changelog.rst"
+
+# Or run some commands after the git tag and the branch
+# have been pushed:
+#  [[tool.tbump.after_push]]
+#  name = "publish"
+#  cmd = "./publish.sh"

xarray_ms-0.2.0/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-0.2.0/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-0.2.0/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