Merge pull request #352 from casangi/327-parition_info_optional_field…

…_antenna_name Add optional partition_info/antenna_name
casangi · Jan 24, 2025 · 6fbfc96 · 6fbfc96
2 parents 38e4635 + ea30e29
commit 6fbfc96
Show file tree

Hide file tree

Showing 4 changed files with 18 additions and 1 deletion.
diff --git a/src/xradio/measurement_set/_utils/_msv2/conversion.py b/src/xradio/measurement_set/_utils/_msv2/conversion.py
@@ -1029,7 +1029,7 @@ def get_observation_info(in_file, observation_id, intents):
                         datetime.timezone.utc
                     ).isoformat(),
                     "xradio_version": importlib.metadata.version("xradio"),
-                    "schema_version": "4.0.-9988",
+                    "schema_version": "4.0.-9987",
                     "type": "visibility",
                 }
             )
@@ -1242,6 +1242,10 @@ def get_observation_info(in_file, observation_id, intents):
                 "intents": intents,
                 "taql_where": taql_where,
             }
+            if with_antenna_partitioning:
+                partition_info_misc_fields["antenna_name"] = xds.coords[
+                    "antenna_name"
+                ].data[0]
             info_dicts = create_info_dicts(
                 in_file, xds, field_and_source_xds, partition_info_misc_fields, tb_tool
             )

diff --git a/src/xradio/measurement_set/_utils/_msv2/msv4_info_dicts.py b/src/xradio/measurement_set/_utils/_msv2/msv4_info_dicts.py
@@ -62,6 +62,10 @@ def create_info_dicts(
         "taql": partition_info_misc_fields["taql_where"],
         "line_name": line_name,
     }
+    if "antenna_name" in partition_info_misc_fields:
+        info_dicts["partition_info"]["antenna_name"] = partition_info_misc_fields[
+            "antenna_name"
+        ]
 
     observation_id = check_if_consistent(
         tb_tool.getcol("OBSERVATION_ID"), "OBSERVATION_ID"

diff --git a/src/xradio/measurement_set/schema.py b/src/xradio/measurement_set/schema.py
@@ -1221,6 +1221,8 @@ class PartitionInfoDict:
     """ The taql query used if converted from MSv2. """
     line_name: list[str]
     """ Spectral line names """
+    antenna_name: Optional[str]
+    """ Name of antenna when partitioning also by antenna (single-dish). """
 
 
 @dict_schema

diff --git a/tests/stakeholder/test_measure_set_stakeholder.py b/tests/stakeholder/test_measure_set_stakeholder.py
@@ -135,6 +135,13 @@ def base_test(
         combined_antenna = ps.get_combined_antenna_xds()
         assert type(combined_antenna) == xr.Dataset
         ps.get_combined_field_and_source_xds()
+        assert all(
+            [
+                "antenna_name" in ps[xds_name].attrs["partition_info"]
+                for xds_name in ps.keys()
+                if "ANTENNA1" in partition_scheme
+            ]
+        )
 
         sum = 0.0
         sum_lazy = 0.0