xarray-ms 0.4.0a2__tar.gz → 0.4.0a4__tar.gz

{xarray_ms-0.4.0a2 → xarray_ms-0.4.0a4}/.github/workflows/ci.yml

@@ -29,29 +29,8 @@ jobs:
     strategy:
       fail-fast: false
       matrix:
-        os: ["ubuntu-22.04", "macos-14", "macos-15"]
-        python-version: ["3.11", "3.12", "3.13"]
-        is_main_or_release:
-          - ${{ contains(github.ref, 'main') || startsWith(github.ref, 'refs/tags')}}
-        exclude:
-          # Don't test on macos-{13,14} in PR's
-          - is_main_or_release: false
-            os: "macos-14"
-          # Don't test on 3.11 - 3.12 in PR's
-          - is_main_or_release: false
-            python-version: "3.11"
-          - is_main_or_release: false
-            python-version: "3.12"
-          # Don't test on macos on 3.11 - 3.12 in PRs
-          - os: "macos-14"
-            python-version: "3.11"
-          - os: "macos-14"
-            python-version: "3.12"
-          - os: "macos-15"
-            python-version: "3.11"
-          - os: "macos-15"
-            python-version: "3.12"
-
+        os: ["ubuntu-22.04", "macos-13", "macos-14", "macos-15"]
+        python-version: ["3.11", "3.12", "3.13", "3.14"]
 
     steps:
       - name: Set up Python ${{ matrix.python-version }}

{xarray_ms-0.4.0a2 → xarray_ms-0.4.0a4}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: xarray-ms
-Version: 0.4.0a2
+Version: 0.4.0a4
 Summary: xarray MSv4 views over MSv2 Measurement Sets
 Author-email: Simon Perkins <simon.perkins@gmail.com>
 License-File: LICENSE

{xarray_ms-0.4.0a2 → xarray_ms-0.4.0a4}/doc/source/changelog.rst

@@ -3,6 +3,14 @@
 Changelog
 =========
 
+X.Y.Z (DD-MM-YYYY)
+------------------
+* Handle duplicate antenna name (:pr:`151`)
+* Deprecate python 3.10 support (:pr:`150`)
+* Add python 3.14 to CI testing matrix (:pr:`150`)
+* Test more comprehensively in pull requests (:pr:`150`)
+* Depend on arcae ``>= 0.5.1, < 0.6.0`` (:pr:`149`)
+
 0.5.1 (06-03-2026)
 ------------------
 * Correct ``BroadcastMSv2Array`` initialisation of ``MSv2Array`` base class

{xarray_ms-0.4.0a2 → xarray_ms-0.4.0a4}/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.4.0-alpha.2"
+release = "0.4.0-alpha.4"
 
 # -- General configuration ---------------------------------------------------
 # https://www.sphinx-doc.org/en/master/usage/configuration.html#general-configuration

{xarray_ms-0.4.0a2 → xarray_ms-0.4.0a4}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "xarray-ms"
-version = "0.4.0-alpha.2"
+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"
@@ -55,7 +55,7 @@ extend-select = ["I"]
 # github_url = "https://github.com/<user or organization>/<project>/"
 
 [tool.tbump.version]
-current = "0.4.0-alpha.2"
+current = "0.4.0-alpha.4"
 
 # https://semver.org/#spec-item-9
 regex = '''

{xarray_ms-0.4.0a2 → xarray_ms-0.4.0a4}/tests/test_antenna.py

@@ -4,6 +4,8 @@ import pytest
 import xarray
 from arcae.lib.arrow_tables import Table, ms_descriptor
 
+from xarray_ms.errors import DuplicateAntennaNameWarning
+
 NANTENNA = 6
 
 
@@ -138,3 +140,60 @@ def test_antenna_feed_join(simmed_ms, auto_corrs):
   npt.assert_array_equal(
     p1["antenna_xds"].antenna_name, np.array(["ANTENNA-0", "ANTENNA-2"], dtype=object)
   )
+
+
+@pytest.mark.parametrize(
+  "simmed_ms",
+  [{"name": "dup.ms", "nantenna": 3, "auto_corrs": True}],
+  indirect=True,
+)
+def test_duplicate_antenna_names(simmed_ms):
+  """When ANTENNA::NAME contains duplicates, names are made unique and a
+  warning is emitted.  The names in antenna_xds and the baseline coordinate
+  must be consistent with each other."""
+
+  # Overwrite ANTENNA::NAME so that antenna 0 and antenna 1 share a name.
+  with Table.from_filename(f"{simmed_ms}::ANTENNA") as T:
+    T.putcol("NAME", np.asarray(["SAME", "SAME", "OTHER"]))
+
+  with pytest.warns(DuplicateAntennaNameWarning, match="SAME"):
+    dt = xarray.open_datatree(simmed_ms, auto_corrs=True).load()
+
+  partition = dt[list(dt.children)[0]]
+  antenna_xds = partition["antenna_xds"]
+
+  # All antenna_name coordinate values must be unique
+  ant_names = antenna_xds.antenna_name.values
+  assert len(ant_names) == len(set(ant_names)), "antenna_name coordinate has duplicates"
+
+  # Duplicated base names get -N suffixes; unique names are unchanged
+  assert "SAME-1" in ant_names
+  assert "SAME-2" in ant_names
+  assert "OTHER" in ant_names
+
+  # correlated dataset baseline_antenna{i}_name maps
+  # appropriately to the antenna dataset antenna names.
+  concat_ant_names = np.concatenate(
+    [partition.baseline_antenna1_name, partition.baseline_antenna2_name]
+  )
+  unique_ant_names, inv = np.unique(concat_ant_names, return_inverse=True)
+  _, lookup = np.where(unique_ant_names[:, None] == ant_names[None, :])
+  ant1 = lookup[inv[: len(inv) // 2]]
+  ant2 = lookup[inv[len(inv) // 2 :]]
+
+  np.testing.assert_array_equal(partition.baseline_antenna1_name, ant_names[ant1])
+  np.testing.assert_array_equal(partition.baseline_antenna2_name, ant_names[ant2])
+
+
+@pytest.mark.parametrize(
+  "simmed_ms",
+  [{"name": "nodup.ms", "nantenna": 3, "auto_corrs": True}],
+  indirect=True,
+)
+def test_unique_antenna_names_no_warning(simmed_ms):
+  """When all antenna names are already unique no warning should be emitted."""
+  import warnings
+
+  with warnings.catch_warnings():
+    warnings.simplefilter("error", DuplicateAntennaNameWarning)
+    xarray.open_datatree(simmed_ms, auto_corrs=True).load()

xarray_ms-0.4.0a4/tests/test_table_utils.py

@@ -0,0 +1,89 @@
+import warnings
+
+import numpy as np
+import numpy.testing as npt
+import pytest
+
+from xarray_ms.backend.msv2.table_utils import unique_antenna_names
+from xarray_ms.errors import DuplicateAntennaNameWarning
+
+
+def test_all_unique_no_change():
+  names = np.array(["A", "B", "C"])
+  result = unique_antenna_names(names)
+  npt.assert_array_equal(result, names)
+
+
+def test_all_unique_no_warning():
+  names = np.array(["A", "B", "C"])
+  with warnings.catch_warnings():
+    warnings.simplefilter("error", DuplicateAntennaNameWarning)
+    unique_antenna_names(names)  # should not raise
+
+
+def test_two_duplicates():
+  names = np.array(["A", "A", "B"])
+  with pytest.warns(DuplicateAntennaNameWarning):
+    result = unique_antenna_names(names)
+  npt.assert_array_equal(result, ["A-1", "A-2", "B"])
+
+
+def test_triple_duplicate():
+  names = np.array(["A", "A", "A"])
+  with pytest.warns(DuplicateAntennaNameWarning):
+    result = unique_antenna_names(names)
+  npt.assert_array_equal(result, ["A-1", "A-2", "A-3"])
+
+
+def test_non_contiguous_duplicates():
+  names = np.array(["A", "B", "A", "C", "A"])
+  with pytest.warns(DuplicateAntennaNameWarning):
+    result = unique_antenna_names(names)
+  npt.assert_array_equal(result, ["A-1", "B", "A-2", "C", "A-3"])
+
+
+def test_unique_name_unchanged_when_mixed():
+  names = np.array(["X", "Y", "X"])
+  with pytest.warns(DuplicateAntennaNameWarning):
+    result = unique_antenna_names(names)
+  npt.assert_array_equal(result, ["X-1", "Y", "X-2"])
+
+
+def test_underscore_in_base_name():
+  """Dash separator avoids collision with names already containing underscores."""
+  names = np.array(["ANTENNA_1", "ANTENNA_1"])
+  with pytest.warns(DuplicateAntennaNameWarning):
+    result = unique_antenna_names(names)
+  npt.assert_array_equal(result, ["ANTENNA_1-1", "ANTENNA_1-2"])
+
+
+def test_empty_array():
+  names = np.array([], dtype=str)
+  result = unique_antenna_names(names)
+  npt.assert_array_equal(result, [])
+
+
+def test_single_element():
+  names = np.array(["A"])
+  result = unique_antenna_names(names)
+  npt.assert_array_equal(result, ["A"])
+
+
+def test_warning_raised_on_duplicates():
+  names = np.array(["A", "A", "B"])
+  with pytest.warns(DuplicateAntennaNameWarning, match="Duplicate antenna names"):
+    unique_antenna_names(names)
+
+
+def test_warning_lists_duplicate_names():
+  names = np.array(["FOO", "BAR", "FOO"])
+  with pytest.warns(DuplicateAntennaNameWarning, match="FOO"):
+    unique_antenna_names(names)
+
+
+def test_result_is_always_unique():
+  """Output names must all be distinct regardless of input."""
+  names = np.array(["A", "A", "B", "B", "C"])
+  with pytest.warns(DuplicateAntennaNameWarning):
+    result = unique_antenna_names(names)
+  assert len(set(result.tolist())) == len(result)

{xarray_ms-0.4.0a2 → xarray_ms-0.4.0a4}/tests/test_write.py

@@ -31,8 +31,8 @@ def test_store(simmed_ms):
         xdt[node.path] = DataTree(ds)
         assert len(node.encoding) > 0
 
-    xdt.sync_msv2(["CORRECTED"])
-    xdt.to_msv2(["UVW", "CORRECTED"])
+    xdt.sync_msv2()
+    xdt.to_msv2()
     written = written or True
 
   with xarray.open_datatree(simmed_ms, auto_corrs=True) as xdt:
@@ -64,7 +64,7 @@ def test_store_region(simmed_ms):
         assert len(node.encoding) > 0
 
     # Create the new MS columns
-    xdt.sync_msv2("CORRECTED")
+    xdt.sync_msv2()
 
     for node in xdt.subtree:
       if node.attrs.get("type") in CORRELATED_DATASET_TYPES:
@@ -72,7 +72,7 @@ def test_store_region(simmed_ms):
         ds = ds.isel(**region)
         ds = ds.assign(CORRECTED=xarray.full_like(ds.CORRECTED, 1 + 2j))
         # Now write it out
-        ds.to_msv2(["CORRECTED"], compute=False, region=region)
+        ds.to_msv2(compute=False, region=region)
 
   # We can check that CORRECTED has been written correctly
   with arcae.table(simmed_ms) as T:
@@ -115,12 +115,12 @@ def test_distributed_write(simmed_ms, processes, nworkers, chunks):
         assert len(node.encoding) > 0
 
     # Create the new MS columns
-    dt.sync_msv2("CORRECTED")
-    dt.to_msv2(["CORRECTED"], compute=False)
+    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(["CORRECTED"], compute=True)
+        node.ds.to_msv2(compute=True)
 
   with arcae.table(simmed_ms) as T:
     corrected = T.getcol("CORRECTED")
@@ -144,11 +144,11 @@ def test_indexed_write(simmed_ms):
       ds = node.ds.assign(CORRECTED=xarray.full_like(node.VISIBILITY, 1 + 2j))
       dt[node.path] = DataTree(ds)
 
-  dt.sync_msv2("CORRECTED")
-  dt.to_msv2(["CORRECTED"], compute=False)
+  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(["CORRECTED"], compute=True)
+        ds.to_msv2(compute=True)

{xarray_ms-0.4.0a2 → xarray_ms-0.4.0a4}/xarray_ms/backend/msv2/factories/antenna.py

@@ -4,6 +4,7 @@ from xarray import Dataset, Variable
 from xarray_ms.backend.msv2.factories.core import DatasetFactory
 from xarray_ms.backend.msv2.imputation import maybe_impute_observation_table
 from xarray_ms.backend.msv2.measures_encoders import MSv2CoderFactory
+from xarray_ms.backend.msv2.table_utils import unique_antenna_names
 from xarray_ms.errors import InvalidMeasurementSet
 
 RELOCATABLE_ARRAY = {"ALMA", "VLA", "NOEMA", "EVLA"}
@@ -53,7 +54,10 @@ class AntennaFactory(DatasetFactory):
       )
 
     ant_coder_factory = MSv2CoderFactory.from_arrow_table(filtered_ants)
-    antenna_names = filtered_ants["NAME"].to_numpy().astype(str)
+    # Deduplicate against the full ANTENNA table so suffix assignments are
+    # consistent with those produced in the correlated dataset factory.
+    all_ant_names = unique_antenna_names(ants["NAME"].to_numpy().astype(str))
+    antenna_names = all_ant_names[feed_ant_id[mask]]
     telescope_names = np.asarray([telescope_name] * len(antenna_names), dtype=str)
     position = pac.list_flatten(filtered_ants["POSITION"]).to_numpy().reshape(-1, 3)
     diameter = filtered_ants["DISH_DIAMETER"].to_numpy()

{xarray_ms-0.4.0a2 → xarray_ms-0.4.0a4}/xarray_ms/backend/msv2/factories/correlated.py

@@ -26,6 +26,7 @@ from xarray_ms.backend.msv2.measures_encoders import (
   VisibilityCoder,
 )
 from xarray_ms.backend.msv2.structure import MSv2StructureFactory, PartitionKeyT
+from xarray_ms.backend.msv2.table_utils import unique_antenna_names
 from xarray_ms.casa_types import ColumnDesc, Polarisations
 from xarray_ms.errors import (
   ColumnShapeImputationWarning,
@@ -266,7 +267,7 @@ class CorrelatedFactory(DatasetFactory):
     assert (partition.nbl,) == ant1.shape
 
     antenna = self._subtable_factories["ANTENNA"].instance
-    ant_names = antenna["NAME"].to_numpy().astype(str)
+    ant_names = unique_antenna_names(antenna["NAME"].to_numpy().astype(str))
     ant1_names = ant_names[ant1]
     ant2_names = ant_names[ant2]
 

xarray_ms-0.4.0a4/xarray_ms/backend/msv2/table_utils.py

@@ -0,0 +1,53 @@
+import json
+import warnings
+from typing import Any, Dict
+
+import numpy as np
+import pyarrow as pa
+
+from xarray_ms.errors import DuplicateAntennaNameWarning
+
+
+def unique_antenna_names(names: np.ndarray) -> np.ndarray:
+  """Return antenna names with duplicates made unique by appending -N suffixes.
+
+  All occurrences of a duplicated name are renamed: the first gets ``-1``,
+  the second ``-2``, and so on.  Names that are already unique are left
+  unchanged.  A :class:`~xarray_ms.errors.DuplicateAntennaNameWarning` is
+  emitted when any renaming is performed.
+  """
+  unique, counts = np.unique(names, return_counts=True)
+  if not (duplicates := set(unique[counts > 1].tolist())):
+    return names
+
+  warnings.warn(
+    f"Duplicate antenna names detected in the ANTENNA table: "
+    f"{sorted(duplicates)}. "
+    f"A numeric suffix will be appended to make names unique.",
+    DuplicateAntennaNameWarning,
+    stacklevel=2,
+  )
+  # Build as a Python list to avoid numpy fixed-width string truncation,
+  # then convert back to an array with a wide enough dtype.
+  result = names.tolist()
+  counters: Dict[str, int] = {}
+  for i, name in enumerate(result):
+    if name in duplicates:
+      counters[name] = counters.get(name, 0) + 1
+      result[i] = f"{name}-{counters[name]}"
+  return np.array(result, dtype=str)
+
+
+def extract_table_desc(table: pa.Table) -> Dict[str, Any]:
+  """Extract the CASA table descriptor stored in an Arrow table constructed by arcae"""
+  try:
+    arcae_metadata = table.schema.metadata[b"__arcae_metadata__"]
+  except KeyError:
+    raise KeyError("__arcae_metadata__ was not present in the table metadata")
+
+  try:
+    return json.loads(arcae_metadata)["__casa_descriptor__"]
+  except KeyError:
+    raise KeyError(
+      f"arcae metadata {arcae_metadata} does not contain a __casa_descriptor__"
+    )

{xarray_ms-0.4.0a2 → xarray_ms-0.4.0a4}/xarray_ms/backend/msv2/writes.py

@@ -1,7 +1,8 @@
 import warnings
 from collections import defaultdict
+from dataclasses import dataclass
 from importlib.metadata import version as package_version
-from typing import Any, Dict, Iterable, List, Literal, Mapping, Set, Tuple
+from typing import Any, Dict, List, Literal, Mapping, Set, Tuple
 
 import numpy as np
 import numpy.typing as npt
@@ -25,7 +26,40 @@ else:
 
 ShapeSetType = Set[Tuple[int, ...]]
 DTypeSetType = Set[npt.DTypeLike]
-ShapeAndDTypeType = Dict[Tuple[str, str], Tuple[ShapeSetType, DTypeSetType]]
+
+
+@dataclass
+class DataVariableInfo:
+  counts: int
+  shapes: ShapeSetType
+  dtypes: DTypeSetType
+
+
+DataVariableInfoMap = Dict[Tuple[str, str], DataVariableInfo]
+
+
+MSV4_WRITE_MAP = {
+  "UVW": "UVW",
+  "VISIBILITY": "DATA",
+  "WEIGHT": "WEIGHT_SPECTRUM",
+  "FLAG": "FLAG",
+}
+
+
+# MSv4 variables that will not be written back to
+# the MSv2 as they are effectively metadata
+IGNORED_VARIABLES = [
+  # Standard MSv4 variables
+  "EFFECTIVE_INTEGRATION_TIME",
+  "TIME_CENTROID",
+  "TIME_CENTROID_EXTRA_PRECISION",
+  "EFFECTIVE_CHANNEL_WIDTH",
+  "FREQUENCY_CENTROID",
+  # Added by Correlated Factory if grids are irregular
+  "TIME",
+  "INTEGRATION_TIME",
+  "CHANNEL_WIDTH",
+]
 
 
 def validate_column_desc(
@@ -100,7 +134,7 @@ def fit_tile_shape(shape: Tuple[int, ...], dtype: npt.DTypeLike) -> Dict[str, np
 
 def generate_column_descriptor(
   table_desc: Dict[str, Any],
-  shapes_and_dtypes: ShapeAndDTypeType,
+  data_var_map: DataVariableInfoMap,
 ) -> Tuple[Dict[str, Any], Dict[str, Any]]:
   """
   Synthesises a column descriptor from:
@@ -113,9 +147,8 @@ def generate_column_descriptor(
   Args:
     table_desc: Table descriptor containing existing
       column definitions.
-    shapes_and_dtypes: A :code:`{var_name: (set(shapes), set(dtypes))}`
-      mapping of shapes and data types associated with each xarray
-      Variable.
+    data_var_map: A mapping of shapes and data types
+      associated with each xarray Variable.
 
   Returns:
     A (descriptor, dminfo) tuple containing any columns
@@ -125,23 +158,22 @@ def generate_column_descriptor(
   actual_desc = {}
   dm_groups = []
 
-  for (var_name, column), (shapes, dtypes) in shapes_and_dtypes.items():
+  for (var_name, msv2_column), data_var_info in data_var_map.items():
     # If there are existing descriptors, either for
     # columns present on the table, or in the canonical definition
     # validate that the variable shape matches the column
-    if column_desc := table_desc.get(column):
-      validate_column_desc(var_name, column, shapes, column_desc)
-    elif column_desc := canonical_table_desc.get(column):
-      validate_column_desc(var_name, column, shapes, column_desc)
+    if column_desc := table_desc.get(msv2_column):
+      validate_column_desc(var_name, msv2_column, data_var_info.shapes, column_desc)
+    elif column_desc := canonical_table_desc.get(msv2_column):
+      validate_column_desc(var_name, msv2_column, data_var_info.shapes, column_desc)
     else:
-      # Construct a column descriptor and possibly an
-      # associated data manager
+      # Construct a column descriptor and possibly an associated data manager
       # Unify variable numpy types
-      dtype = np.result_type(*dtypes)
+      dtype = np.result_type(*data_var_info.dtypes)
 
       if dtype is object:
         raise NotImplementedError(
-          f"Types of variable {var_name} ({list(dtypes)}) "
+          f"Types of variable {var_name} ({list(data_var_info.dtypes)}) "
           f"resolves to an object. "
           f"Writing of objects is not supported"
         )
@@ -155,21 +187,21 @@ def generate_column_descriptor(
 
       column_desc = {"valueType": casa_type, "option": 0}
 
-      if len(shapes) == 1:
+      if len(data_var_info.shapes) == 1:
         # If the shape is fixed, Tile the column
         # column descriptor shapes are fortran ordered
-        fixed_shape = tuple(reversed(next(iter(shapes))))
+        fixed_shape = tuple(reversed(next(iter(data_var_info.shapes))))
         row_only = len(fixed_shape) == 0
         if not row_only:
           column_desc["option"] |= 4
           column_desc["shape"] = list(fixed_shape)
           column_desc["ndim"] = len(fixed_shape)
-        column_desc["dataManagerGroup"] = dm_group = f"{column}_GROUP"
+        column_desc["dataManagerGroup"] = dm_group = f"{msv2_column}_GROUP"
         column_desc["dataManagerType"] = dm_type = "TiledColumnStMan"
 
         dm_groups.append(
           {
-            "COLUMNS": [column],
+            "COLUMNS": [msv2_column],
             "NAME": dm_group,
             "TYPE": dm_type,
             "SPEC": fit_tile_shape(fixed_shape, dtype),
@@ -182,7 +214,7 @@ def generate_column_descriptor(
         column_desc["dataManagerGroup"] = "StandardStMan"
         column_desc["dataManagerType"] = "StandardStMan"
 
-      actual_desc[column] = column_desc
+      actual_desc[msv2_column] = column_desc
 
   dminfo = {f"*{i + 1}": g for i, g in enumerate(dm_groups)}
   return actual_desc, dminfo
@@ -216,68 +248,39 @@ def msv2_store_from_dataset(ds: Dataset, region="auto") -> MSv2Store:
   )
 
 
-WriteVariablesT = str | Tuple[str, str] | List[str | Tuple[str, str]] | Dict[str, str]
+WriteMapT = Tuple[str, str] | List[Tuple[str, str]] | Dict[str, str] | None
 
 
-def promote_write_variables(variables: WriteVariablesT) -> Dict[str, str]:
+def promote_write_map(variables: WriteMapT) -> Dict[str, str]:
+  """Promotes the supplied write variable mapping into a Dict[str, str]"""
   type_except = TypeError(
     f"{variables} should be one of:\n"
-    f"1. a str representing a variable\n"
-    f"2. a Tuple[str, str] variable -> column mapping\n"
-    f"3. A List containing either 1. or 2.\n"
-    f"4. A Dict[str, str] variable -> column mapping"
+    f"1. a Tuple[str, str] variable -> column mapping\n"
+    f"2. An Iterable[Tuple[str, str]].\n"
+    f"3. A Dict[str, str] variable -> column mapping"
   )
 
-  def check_var(var, on_type_error="raise"):
-    if isinstance(var, str):
-      return (var, var)
-    elif isinstance(var, tuple):
-      if not len(var) == 2 and all(isinstance(v, str) for v in var):
-        raise type_except
-
-      return (var[0], var[1])
-
-    if on_type_error == "raise":
-      raise type_except
-
-    return None
-
-  if result := check_var(variables, on_type_error="none"):
-    return result
-
-  if isinstance(variables, list):
-    return {src: dest for src, dest in map(check_var, variables)}
-
-  if isinstance(variables, dict):
-    for k, v in variables.items():
-      if not isinstance(k, str) or not isinstance(v, str):
-        raise type_except
-
-  raise type_except
-
-
-IGNORED_VARIABLES = [
-  # Standard MSv4 variables
-  # spectrum_xds
-  "SPECTRUM",
-  # visibility_xds
-  "VISIBILITY",
-  "FLAG",
-  "WEIGHT",
-  "UVW",
-  "EFFECTIVE_INTEGRATION_TIME",
-  "TIME_CENTROID",
-  "TIME_CENTROID_EXTRA_PRECISION",
-  "EFFECTIVE_CHANNEL_WIDTH",
-  "FREQUENCY_CENTROID",
-  # Added by Correlated Factory if grids are irregular
-  "TIME",
-  "INTEGRATION_TIME",
-  "CHANNEL_WIDTH",
-]
+  def check_tuple(tup: Tuple[str, str]) -> Tuple[str, str]:
+    if (
+      isinstance(tup, tuple) and len(tup) == 2 and all(isinstance(t, str) for t in tup)
+    ):
+      return tup
+
+    raise type_except
+
+  if variables is None:
+    return {}
+  elif isinstance(variables, dict):
+    return variables
+  elif isinstance(variables, tuple):
+    return dict([check_tuple(variables)])
+  elif isinstance(variables, (list, set)):
+    return dict([check_tuple(v) for v in variables])
+  else:
+    raise type_except
 
 
-def sync_msv2(dt: DataTree, variables: WriteVariablesT):
+def sync_msv2(dt: DataTree, write_map: WriteMapT = None):
   assert isinstance(dt, DataTree)
   vis_datasets = [
     n for n in dt.subtree if n.attrs.get("type") in CORRELATED_DATASET_TYPES
@@ -291,61 +294,56 @@ def sync_msv2(dt: DataTree, variables: WriteVariablesT):
   table_factory = msv2_store_from_dataset(next(iter(vis_datasets)).ds).table_factory
   table_desc = table_factory.instance.tabledesc()
 
-  # write_var_map = promote_write_variables(variables)
-
-  # if variables is None:
-  #   columns = table_factory.instance.columns()
-
-  list_var_names: List[Any] = (
-    [variables] if isinstance(variables, str) else list(variables)
+  write_map = {**MSV4_WRITE_MAP, **promote_write_map(write_map)}
+  var_info_map: DataVariableInfoMap = defaultdict(
+    lambda: DataVariableInfo(0, set(), set())
   )
-  set_var_names = set(list_var_names)
-
-  if len(set_var_names) == 0:
-    warnings.warn("Empty 'variables'")
-    return
-
-  shapes_and_dtypes: ShapeAndDTypeType = defaultdict(lambda: (set(), set()))
 
-  for node in vis_datasets:
+  for nodes_visited, node in enumerate(vis_datasets, 1):
     assert isinstance(node, DataTree)
-    if not set_var_names.issubset(node.data_vars.keys()):
-      raise ValueError(
-        f"{set_var_names} are not present in all visibility DataTree nodes"
-      )
-
-    for n in set_var_names:
-      var = node.data_vars[n]
+    for name, var in node.data_vars.items():
+      if name in IGNORED_VARIABLES:
+        continue
 
+      # Don't try to create anything that doesn't translate
+      # to a MAIN MSv2 row dimension
       if var.dims[: len(MAIN_PREFIX_DIMS)] != MAIN_PREFIX_DIMS:
-        raise ValueError(
-          f"{n} dimensions {var.dims} do not start with {MAIN_PREFIX_DIMS}"
+        warnings.warn(
+          f"Ignoring {name} in {node.path}  "
+          f"dimensions {var.dims} do not begin with {MAIN_PREFIX_DIMS}",
+          UserWarning,
         )
+        continue
+
+      entry = var_info_map[(name, write_map.get(name, name))]
+      entry.counts += 1
+      entry.shapes.add(var.shape[len(MAIN_PREFIX_DIMS) :])
+      entry.dtypes.add(var.dtype)
+
+  # Ensure that the identified variables are universally
+  # defined across all DataTree.
+  for key, entry in list(var_info_map.items()):
+    if entry.counts != nodes_visited:
+      warnings.warn(
+        f"Ignoring {key[0]} which does not appear in all "
+        f"visibility datasets {list(dt.children)}",
+        UserWarning,
+      )
+      var_info_map.pop(key)
 
-      shapes, dtypes = shapes_and_dtypes[(n, n)]
-      shapes.add(var.shape[len(MAIN_PREFIX_DIMS) :])
-      dtypes.add(var.dtype)
-
-  table_factory = msv2_store_from_dataset(next(iter(vis_datasets)).ds)._table_factory
-  table_desc = table_factory.instance.tabledesc()
-  column_descs, dminfo = generate_column_descriptor(table_desc, shapes_and_dtypes)
+  # Generate column descriptors and add the columns
+  column_descs, dminfo = generate_column_descriptor(table_desc, var_info_map)
   table_factory.instance.addcols(column_descs, dminfo)
   assert set(column_descs.keys()).issubset(table_factory.instance.columns())
 
 
 def datatree_to_msv2(
   dt: DataTree,
-  variables: str | Iterable[str],
+  write_map: WriteMapT = None,
   compute: Literal[True] = True,
   write_inherited_coords: bool = False,
-):
+) -> None:
   assert isinstance(dt, DataTree)
-  list_var_names = [variables] if isinstance(variables, str) else list(variables)
-  set_var_names = set(list_var_names)
-
-  if len(set_var_names) == 0:
-    warnings.warn("Empty 'variables'")
-    return
 
   if (
     len(
@@ -363,32 +361,32 @@ def datatree_to_msv2(
   for node in vis_datasets:
     at_root = node is dt.root
     ds = node.to_dataset(inherit=write_inherited_coords or at_root)
-    dataset_to_msv2(ds, list_var_names, compute)
+    dataset_to_msv2(ds, write_map=write_map, compute=compute)
 
 
 def dataset_to_msv2(
   ds: Dataset,
-  variables: str | Iterable[str],
+  write_map: WriteMapT = None,
   compute: Literal[True] = True,
   region: Mapping[str, slice | Literal["auto"]] | Literal["auto"] = "auto",
-):
+) -> None:
   assert isinstance(ds, Dataset)
-  list_vars = [variables] if isinstance(variables, str) else list(variables)
-
-  if len(list_vars) == 0:
-    return
 
   # Strip out
-  # 1. variables that will not be written
+  # 1. ancilliary variables
   # 2. coordinates
   # 3. attributes
-  ignored_vars = set(ds.data_vars) - set(list_vars)
-  ds = ds.drop_vars(ignored_vars | set(ds.coords)).drop_attrs()
+  write_map = {**MSV4_WRITE_MAP, **promote_write_map(write_map)}
+  ignored = set(IGNORED_VARIABLES) & set(ds.data_vars)
+  write_ds = ds.drop_vars(ignored | set(ds.coords)).drop_attrs()
+  write_ds = write_ds.rename_vars(
+    {k: v for k, v in write_map.items() if k in write_ds.data_vars}
+  )
 
-  msv2_store = msv2_store_from_dataset(ds, region)
-  msv2_store.set_write_region(ds)
+  msv2_store = msv2_store_from_dataset(write_ds, region)
+  msv2_store.set_write_region(write_ds)
   writer = ArrayWriter()
-  dump_to_store(ds, msv2_store, writer)
+  dump_to_store(write_ds, msv2_store, writer)
   writes = writer.sync(compute=compute)
 
   if compute:

{xarray_ms-0.4.0a2 → xarray_ms-0.4.0a4}/xarray_ms/errors.py

@@ -30,6 +30,11 @@ class ImputedMetadataWarning(MissingMetadataWarning):
   if the original metadata is missing"""
 
 
+class DuplicateAntennaNameWarning(UserWarning):
+  """Warning raised when duplicate antenna names are found in the ANTENNA table
+  and are made unique by appending a numeric suffix"""
+
+
 class FrameConversionWarning(UserWarning):
   """Warning raised if there's no defined conversion from a CASA
   to astropy reference frame"""

xarray_ms-0.4.0a2/xarray_ms/backend/msv2/table_utils.py

@@ -1,19 +0,0 @@
-import json
-from typing import Any, Dict
-
-import pyarrow as pa
-
-
-def extract_table_desc(table: pa.Table) -> Dict[str, Any]:
-  """Extract the CASA table descriptor stored in an Arrow table constructed by arcae"""
-  try:
-    arcae_metadata = table.schema.metadata[b"__arcae_metadata__"]
-  except KeyError:
-    raise KeyError("__arcae_metadata__ was not present in the table metadata")
-
-  try:
-    return json.loads(arcae_metadata)["__casa_descriptor__"]
-  except KeyError:
-    raise KeyError(
-      f"arcae metadata {arcae_metadata} does not contain a __casa_descriptor__"
-    )