Limited write support

doc/source/conf.py

@@ -11,7 +11,7 @@
 project = "xarray-ms"
 copyright = "2024 - 2025 NRF (SARAO) and Rhodes University (RATT) Centre"
 author = "Simon Perkins"
-release = "0.5.1"
+release = "0.4.0-alpha.4"
 
 # -- General configuration ---------------------------------------------------
 # https://www.sphinx-doc.org/en/master/usage/configuration.html#general-configuration

pyproject.toml

@@ -1,14 +1,14 @@
 [project]
 name = "xarray-ms"
-version = "0.5.1"
+version = "0.4.0-alpha.4"
 description = "xarray MSv4 views over MSv2 Measurement Sets"
 authors = [{name = "Simon Perkins", email = "simon.perkins@gmail.com"}]
 readme = "README.rst"
 requires-python = ">=3.11"
 dependencies = [
     "xarray>=2025.0",
     "cacheout>=0.16.0",
-    "arcae>=0.5.1, < 0.6.0",
+    "arcae>=0.4.0a4, < 0.5.0",
     "typing-extensions>=4.12.2",
 ]
 
@@ -55,18 +55,18 @@ extend-select = ["I"]
 # github_url = "https://github.com/<user or organization>/<project>/"
 
 [tool.tbump.version]
-current = "0.5.1"
+current = "0.4.0-alpha.4"
 
-# Example of a semver regexp.
-# Make sure this matches current_version before
-# using tbump
+# https://semver.org/#spec-item-9
 regex = '''
   (?P<major>\d+)
   \.
   (?P<minor>\d+)
   \.
   (?P<patch>\d+)
-  '''
+  -?
+  (?P<ptype>(alpha|beta|rc)\.(?P<prerelease>\d+))?
+'''
 
 [tool.tbump.git]
 message_template = "Bump to {new_version}"
@@ -77,9 +77,11 @@ tag_template = "{new_version}"
 # tbump.toml location.
 [[tool.tbump.file]]
 src = "pyproject.toml"
+search = "^(current|version)\\s*=\\s*\"{current_version}\""
 
 [[tool.tbump.file]]
 src = "doc/source/conf.py"
+search = "^release\\s*=\\s*\"{current_version}\""
 
 # You can specify a list of commands to
 # run after the files have been patched

tests/conftest.py

@@ -1,3 +1,5 @@
+import gc
+
 import numpy as np
 import pytest
 from arcae.lib.arrow_tables import Table, ms_descriptor
@@ -38,14 +40,18 @@ def pytest_collection_modifyitems(config, items):
 @pytest.fixture(autouse=True)
 def clear_caches():
   yield
-  Multiton._INSTANCE_CACHE.clear()
+
+  # Structure Factories have references to Multitons
   MSv2StructureFactory._STRUCTURE_CACHE.clear()
+  Multiton._INSTANCE_CACHE.clear()
+  gc.collect()
 
 
 @pytest.fixture(scope="session", params=[DEFAULT_SIM_PARAMS])
 def simmed_ms(request, tmp_path_factory):
-  ms = tmp_path_factory.mktemp("simulated") / request.param.pop("name", "test.ms")
-  simulator = MSStructureSimulator(**{**DEFAULT_SIM_PARAMS, **request.param})
+  params = request.param.copy()
+  ms = tmp_path_factory.mktemp("simulated") / params.pop("name", "test.ms")
+  simulator = MSStructureSimulator(**{**DEFAULT_SIM_PARAMS, **params})
   simulator.simulate_ms(str(ms))
   return str(ms)
 

tests/test_write.py

@@ -0,0 +1,154 @@
+from contextlib import ExitStack
+
+import arcae
+import numpy as np
+import pytest
+import xarray
+from xarray import DataTree
+
+import xarray_ms
+from xarray_ms.errors import MismatchedWriteRegion
+from xarray_ms.msv4_types import CORRELATED_DATASET_TYPES
+
+
+def test_write_support():
+  assert xarray_ms.multithreaded_writes()
+
+
+@pytest.mark.parametrize("simmed_ms", [{"name": "test_store.ms"}], indirect=True)
+def test_store(simmed_ms):
+  read = written = False
+
+  with xarray.open_datatree(simmed_ms, auto_corrs=True) as xdt:
+    # Overwrite UVW coordinates with zeroes
+    # Add a CORRECTED column
+    for node in xdt.subtree:
+      if node.attrs.get("type") in CORRELATED_DATASET_TYPES:
+        assert not np.all(node.UVW == 0)
+        node.UVW[:] = 0
+        assert len(node.encoding) > 0
+        ds = node.ds.assign(CORRECTED=xarray.full_like(node.VISIBILITY, 2 + 3j))
+        xdt[node.path] = DataTree(ds)
+        assert len(node.encoding) > 0
+
+    xdt.sync_msv2()
+    xdt.to_msv2()
+    written = written or True
+
+  with xarray.open_datatree(simmed_ms, auto_corrs=True) as xdt:
+    for node in xdt.subtree:
+      if node.attrs.get("type") in CORRELATED_DATASET_TYPES:
+        np.testing.assert_array_equal(node.UVW, 0)
+        np.testing.assert_array_equal(node.CORRECTED, 2 + 3j)
+        read = read or True
+
+  assert read
+  assert written
+
+  # We can check that CORRECTED has been written correctly
+  with arcae.table(simmed_ms) as T:
+    np.testing.assert_array_equal(T.getcol("UVW"), 0)
+    np.testing.assert_array_equal(T.getcol("CORRECTED"), 2 + 3j)
+
+
+@pytest.mark.parametrize("simmed_ms", [{"name": "test_store_region.ms"}], indirect=True)
+def test_store_region(simmed_ms):
+  region = {"time": slice(0, 2), "frequency": slice(2, 4)}
+
+  with xarray.open_datatree(simmed_ms, auto_corrs=True) as xdt:
+    # Add a CORRECTED column
+    for node in xdt.subtree:
+      if node.attrs.get("type") in CORRELATED_DATASET_TYPES:
+        ds = node.ds.assign(CORRECTED=xarray.zeros_like(node.VISIBILITY))
+        xdt[node.path] = DataTree(ds)
+        assert len(node.encoding) > 0
+
+    # Create the new MS columns
+    xdt.sync_msv2()
+
+    for node in xdt.subtree:
+      if node.attrs.get("type") in CORRELATED_DATASET_TYPES:
+        sizes = node.sizes
+        ds = ds.isel(**region)
+        ds = ds.assign(CORRECTED=xarray.full_like(ds.CORRECTED, 1 + 2j))
+        # Now write it out
+        ds.to_msv2(compute=False, region=region)
+
+  # We can check that CORRECTED has been written correctly
+  with arcae.table(simmed_ms) as T:
+    corrected = T.getcol("CORRECTED")
+    nt, nbl, nf, npol = (
+      sizes[d] for d in ("time", "baseline_id", "frequency", "polarization")
+    )
+    corrected = corrected.reshape((nt, nbl, nf, npol))
+    ts, fs = (region[d] for d in ("time", "frequency"))
+    mask = np.full(corrected.shape, False, np.bool_)
+    mask[ts, :, fs, :] = True
+    np.testing.assert_array_equal(corrected[mask], 1 + 2j)
+    np.testing.assert_array_equal(corrected[~mask], 0 + 0j)
+
+
+@pytest.mark.parametrize("chunks", [{"time": 2, "frequency": 2}])
+@pytest.mark.parametrize("simmed_ms", [{"name": "distributed-write.ms"}], indirect=True)
+@pytest.mark.parametrize("nworkers", [4])
+@pytest.mark.parametrize("processes", [True, False])
+def test_distributed_write(simmed_ms, processes, nworkers, chunks):
+  da = pytest.importorskip("dask.array")
+  distributed = pytest.importorskip("dask.distributed")
+
+  with ExitStack() as stack:
+    cluster = stack.enter_context(
+      distributed.LocalCluster(processes=processes, n_workers=nworkers)
+    )
+    stack.enter_context(distributed.Client(cluster))
+    dt = stack.enter_context(
+      xarray.open_datatree(simmed_ms, chunks=chunks, auto_corrs=True)
+    )
+    for node in dt.subtree:
+      if node.attrs.get("type") in CORRELATED_DATASET_TYPES:
+        vis = node.VISIBILITY
+        sizes = node.sizes
+        corrected = da.arange(np.prod(vis.shape), dtype=np.int32)
+        corrected = corrected.reshape(vis.shape).rechunk(vis.data.chunks)
+        ds = node.ds.assign(CORRECTED=(vis.dims, corrected))
+        dt[node.path] = DataTree(ds)
+        assert len(node.encoding) > 0
+
+    # Create the new MS columns
+    dt.sync_msv2()
+    dt.to_msv2(compute=False)
+
+    for node in dt.subtree:
+      if node.attrs.get("type") in CORRELATED_DATASET_TYPES:
+        node.ds.to_msv2(compute=True)
+
+  with arcae.table(simmed_ms) as T:
+    corrected = T.getcol("CORRECTED")
+    shape = tuple(
+      sizes[d] for d in ("time", "baseline_id", "frequency", "polarization")
+    )
+    expected = np.arange(np.prod(vis.shape), dtype=np.int32)
+    expected = expected.reshape((-1,) + shape[2:])
+    np.testing.assert_array_equal(corrected, expected)
+
+
+@pytest.mark.parametrize("simmed_ms", [{"name": "indexed-write.ms"}], indirect=True)
+def test_indexed_write(simmed_ms):
+  """Check that we throw if we select a variable out with an integer index
+  and then try write that sub-selection out"""
+  dt = xarray.open_datatree(simmed_ms)
+  assert len(dt.children) == 1
+
+  for node in dt.subtree:
+    if node.attrs.get("type") in CORRELATED_DATASET_TYPES:
+      ds = node.ds.assign(CORRECTED=xarray.full_like(node.VISIBILITY, 1 + 2j))
+      dt[node.path] = DataTree(ds)
+
+  dt.sync_msv2()
+  dt.to_msv2(compute=False)
+
+  for node in dt.subtree:
+    if node.attrs.get("type") in CORRELATED_DATASET_TYPES:
+      ds = node.ds.isel(time=slice(0, 2), baseline_id=slice(0, 2), frequency=1)
+      with pytest.raises(MismatchedWriteRegion):
+        ds.to_msv2(compute=True)

xarray_ms/__init__.py

@@ -0,0 +1,36 @@
+import warnings
+
+HAS_WRITE_SUPPORT = False
+
+
+def _install_write_support():
+  """Patch write support methods onto xarray.Dataset and xarray.DataTree"""
+  global HAS_WRITE_SUPPORT
+
+  if HAS_WRITE_SUPPORT:
+    return
+
+  try:
+    from xarray import Dataset, DataTree
+
+    from xarray_ms.backend.msv2.writes import (
+      dataset_to_msv2,
+      datatree_to_msv2,
+      sync_msv2,
+    )
+  except ImportError as e:
+    warnings.warn(f"Engaging write support failed due to {e}", UserWarning)
+    HAS_WRITE_SUPPORT = False
+  else:
+    Dataset.to_msv2 = dataset_to_msv2
+    DataTree.to_msv2 = datatree_to_msv2
+    DataTree.sync_msv2 = sync_msv2
+    HAS_WRITE_SUPPORT = True
+
+
+def multithreaded_writes() -> bool:
+  """Return True if multithreaded write support is enabled"""
+  return HAS_WRITE_SUPPORT
+
+
+_install_write_support()