xradio 0.0.41__py3-none-any.whl → 0.0.42__py3-none-any.whl

xradio/{correlated_data/_utils/_ms → measurement_set/_utils/_msv2}/conversion.py

@@ -11,18 +11,18 @@ import xarray as xr
 
 import toolviper.utils.logger as logger
 from casacore import tables
-from xradio.correlated_data._utils._ms.msv4_sub_xdss import (
+from xradio.measurement_set._utils._msv2.msv4_sub_xdss import (
     create_pointing_xds,
     create_system_calibration_xds,
     create_weather_xds,
 )
 from .msv4_info_dicts import create_info_dicts
-from xradio.correlated_data._utils._ms.create_antenna_xds import (
+from xradio.measurement_set._utils._msv2.create_antenna_xds import (
     create_antenna_xds,
     create_gain_curve_xds,
     create_phase_calibration_xds,
 )
-from xradio.correlated_data._utils._ms.create_field_and_source_xds import (
+from xradio.measurement_set._utils._msv2.create_field_and_source_xds import (
     create_field_and_source_xds,
 )
 from xradio._utils.schema import column_description_casacore_to_msv4_measure
@@ -422,9 +422,9 @@ def create_coordinates(
         "time": utime,
         "baseline_antenna1_id": ("baseline_id", baseline_ant1_id),
         "baseline_antenna2_id": ("baseline_id", baseline_ant2_id),
-        "uvw_label": ["u", "v", "w"],
         "baseline_id": np.arange(len(baseline_ant1_id)),
         "scan_number": ("time", scan_id),
+        "uvw_label": ["u", "v", "w"],
     }
 
     ddi_xds = load_generic_table(in_file, "DATA_DESCRIPTION").sel(row=ddi)
@@ -482,7 +482,6 @@ def create_coordinates(
     }
     xds.frequency.attrs["spectral_window_id"] = spectral_window_id
 
-    # xds.frequency.attrs["effective_channel_width"] = "EFFECTIVE_CHANNEL_WIDTH"
     # Add if doppler table is present
     # xds.frequency.attrs["doppler_velocity"] =
     # xds.frequency.attrs["doppler_type"] =
@@ -605,8 +604,8 @@ def create_data_variables(
                 logger.debug(
                     "Time to read column " + str(col) + " : " + str(time.time() - start)
                 )
-            except Exception as e:
-                logger.debug("Could not load column", col)
+            except Exception as exc:
+                logger.debug(f"Could not load column {col}, exception: {exc}")
 
                 if ("WEIGHT_SPECTRUM" == col) and (
                     "WEIGHT" in col_names
@@ -624,13 +623,27 @@ def create_data_variables(
 
 
 def add_missing_data_var_attrs(xds):
-    """Adds in attributes expected metadata that cannot be found
-    in the input MSv2. For now specifically for missing
-    single-dish/SPECTRUM metadata"""
+    """
+    Adds in the xds attributes expected metadata that cannot be found in the input MSv2.
+    For now:
+    - missing single-dish/SPECTRUM metadata
+    - missing interferometry/VISIBILITY_MODEL metadata
+    """
     data_var_names = ["SPECTRUM", "SPECTRUM_CORRECTED"]
     for var_name in data_var_names:
         if var_name in xds.data_vars:
-            xds.data_vars[var_name].attrs["units"] = ["Jy"]
+            xds.data_vars[var_name].attrs["units"] = [""]
+
+    vis_var_names = ["VISIBILITY_MODEL"]
+    for var_name in vis_var_names:
+        if var_name in xds.data_vars and "units" not in xds.data_vars[var_name].attrs:
+            # Assume MODEL uses the same units
+            if "VISIBILITY" in xds.data_vars:
+                xds.data_vars[var_name].attrs["units"] = xds.data_vars[
+                    "VISIBILITY"
+                ].attrs["units"]
+            else:
+                xds.data_vars[var_name].attrs["units"] = [""]
 
     return xds
 
@@ -688,6 +701,28 @@ def create_taql_query(partition_info):
     return taql_where
 
 
+def fix_uvw_frame(
+    xds: xr.Dataset, field_and_source_xds: xr.Dataset, is_single_dish: bool
+) -> xr.Dataset:
+    """
+    Fix UVW frame
+
+    From CASA fixvis docs: clean and the im tool ignore the reference frame claimed by the UVW column (it is often
+    mislabelled as ITRF when it is really FK5 (J2000)) and instead assume the (u, v, w)s are in the same frame as the phase
+    tracking center. calcuvw does not yet force the UVW column and field centers to use the same reference frame!
+    Blank = use the phase tracking frame of vis.
+    """
+    if xds.UVW.attrs["frame"] == "ITRF":
+        if is_single_dish:
+            center_var = "FIELD_REFERENCE_CENTER"
+        else:
+            center_var = "FIELD_PHASE_CENTER"
+
+        xds.UVW.attrs["frame"] = field_and_source_xds[center_var].attrs["frame"]
+
+    return xds
+
+
 def convert_and_write_partition(
     in_file: str,
     out_file: str,
@@ -798,9 +833,9 @@ def convert_and_write_partition(
             start = time.time()
             xds = xr.Dataset(
                 attrs={
-                    "creation_date": datetime.datetime.now().isoformat(),
+                    "creation_date": datetime.datetime.utcnow().isoformat(),
                     "xradio_version": importlib.metadata.version("xradio"),
-                    "schema_version": "4.0.-9999",
+                    "schema_version": "4.0.-9994",
                     "type": "visibility",
                 }
             )
@@ -847,7 +882,6 @@ def convert_and_write_partition(
 
             # Add data_groups
             xds, is_single_dish = add_data_groups(xds)
-
             xds = add_missing_data_var_attrs(xds)
 
             if (
@@ -895,8 +929,8 @@ def convert_and_write_partition(
                 antenna_id,
                 feed_id,
                 telescope_name,
+                xds.polarization,
             )
-
             logger.debug("Time antenna xds  " + str(time.time() - start))
 
             start = time.time()
@@ -943,6 +977,8 @@ def convert_and_write_partition(
             logger.debug("Time weather " + str(time.time() - start))
 
             # Create pointing_xds
+            pointing_xds = xr.Dataset()
+            print("with_pointing", with_pointing)
             if with_pointing:
                 start = time.time()
                 if pointing_interpolate:
@@ -1003,16 +1039,7 @@ def convert_and_write_partition(
             )
             logger.debug("Time field_and_source_xds " + str(time.time() - start))
 
-            # Fix UVW frame
-            # From CASA fixvis docs: clean and the im tool ignore the reference frame claimed by the UVW column (it is often mislabelled as ITRF when it is really FK5 (J2000)) and instead assume the (u, v, w)s are in the same frame as the phase tracking center. calcuvw does not yet force the UVW column and field centers to use the same reference frame! Blank = use the phase tracking frame of vis.
-            if is_single_dish:
-                xds.UVW.attrs["frame"] = field_and_source_xds[
-                    "FIELD_REFERENCE_CENTER"
-                ].attrs["frame"]
-            else:
-                xds.UVW.attrs["frame"] = field_and_source_xds[
-                    "FIELD_PHASE_CENTER"
-                ].attrs["frame"]
+            xds = fix_uvw_frame(xds, field_and_source_xds, is_single_dish)
 
             partition_info_misc_fields = {
                 "scan_id": scan_id,
@@ -1040,41 +1067,53 @@ def convert_and_write_partition(
             else:
                 mode = "w-"
 
+            if is_single_dish:
+                xds.attrs["type"] = "spectrum"
+                xds = xds.drop_vars(["UVW"])
+                del xds["uvw_label"]
+            else:
+                if any("WVR" in s for s in intents):
+                    xds.attrs["type"] = "wvr"
+                else:
+                    xds.attrs["type"] = "visibility"
+
             start = time.time()
             if storage_backend == "zarr":
-                xds.to_zarr(store=os.path.join(file_name, "MAIN"), mode=mode)
-                ant_xds.to_zarr(store=os.path.join(file_name, "ANTENNA"), mode=mode)
+                xds.to_zarr(store=os.path.join(file_name, "correlated_xds"), mode=mode)
+                ant_xds.to_zarr(store=os.path.join(file_name, "antenna_xds"), mode=mode)
                 for group_name in xds.attrs["data_groups"]:
                     field_and_source_xds.to_zarr(
                         store=os.path.join(
-                            file_name, f"FIELD_AND_SOURCE_{group_name.upper()}"
+                            file_name, f"field_and_source_xds_{group_name}"
                         ),
                         mode=mode,
                     )
 
-                if with_pointing and len(pointing_xds.data_vars) > 1:
+                if with_pointing and len(pointing_xds.data_vars) > 0:
                     pointing_xds.to_zarr(
-                        store=os.path.join(file_name, "POINTING"), mode=mode
+                        store=os.path.join(file_name, "pointing_xds"), mode=mode
                     )
 
                 if system_calibration_xds:
                     system_calibration_xds.to_zarr(
-                        store=os.path.join(file_name, "SYSCAL"), mode=mode
+                        store=os.path.join(file_name, "system_calibration_xds"),
+                        mode=mode,
                     )
 
                 if gain_curve_xds:
                     gain_curve_xds.to_zarr(
-                        store=os.path.join(file_name, "GAIN_CURVE"), mode=mode
+                        store=os.path.join(file_name, "gain_curve_xds"), mode=mode
                     )
 
                 if phase_calibration_xds:
                     phase_calibration_xds.to_zarr(
-                        store=os.path.join(file_name, "PHASE_CAL"), mode=mode
+                        store=os.path.join(file_name, "phase_calibration_xds"),
+                        mode=mode,
                     )
 
                 if weather_xds:
                     weather_xds.to_zarr(
-                        store=os.path.join(file_name, "WEATHER"), mode=mode
+                        store=os.path.join(file_name, "weather_xds"), mode=mode
                     )
 
             elif storage_backend == "netcdf":
@@ -1106,7 +1145,11 @@ def antenna_ids_to_names(
             }
         )
     else:
-        xds["baseline_id"] = ant_xds["antenna_name"].sel(antenna_id=xds["baseline_id"])
+        # baseline_antenna1_id will be removed soon below, but it is useful here to know the actual antenna_ids,
+        # as opposed to the baseline_ids which can mismatch when data is missing for some antennas
+        xds["baseline_id"] = ant_xds["antenna_name"].sel(
+            antenna_id=xds["baseline_antenna1_id"]
+        )
         unwanted_coords_from_ant_xds = [
             "antenna_id",
             "antenna_name",

xradio/{correlated_data/_utils/_ms → measurement_set/_utils/_msv2}/create_antenna_xds.py

@@ -6,8 +6,8 @@ import numpy as np
 import xarray as xr
 import os
 
-from xradio.correlated_data._utils._ms.subtables import subt_rename_ids
-from xradio.correlated_data._utils._ms._tables.read import (
+from xradio.measurement_set._utils._msv2.subtables import subt_rename_ids
+from xradio.measurement_set._utils._msv2._tables.read import (
     load_generic_table,
     convert_casacore_time,
     convert_casacore_time_to_mjd,
@@ -15,7 +15,7 @@ from xradio.correlated_data._utils._ms._tables.read import (
     table_exists,
 )
 from xradio._utils.schema import convert_generic_xds_to_xradio_schema
-from xradio.correlated_data._utils._ms.msv4_sub_xdss import interpolate_to_time
+from xradio.measurement_set._utils._msv2.msv4_sub_xdss import interpolate_to_time
 
 from xradio._utils.list_and_array import (
     check_if_consistent,
@@ -31,6 +31,7 @@ def create_antenna_xds(
     antenna_id: list,
     feed_id: list,
     telescope_name: str,
+    partition_polarization: xr.DataArray,
 ) -> xr.Dataset:
     """
     Create an Xarray Dataset containing antenna information.
@@ -47,6 +48,8 @@ def create_antenna_xds(
         List of feed IDs.
     telescope_name : str
         Name of the telescope.
+    partition_polarization: xr.DataArray
+        Polarization labels of this partition, needed if that info is not present in FEED
 
     Returns
     ----------
@@ -59,6 +62,18 @@ def create_antenna_xds(
     ant_xds = extract_feed_info(
         ant_xds, in_file, antenna_id, feed_id, spectral_window_id
     )
+    # Needed for special SPWs such as ALMA WVR or CHANNEL_AVERAGE data (have no feed info)
+    if "polarization_type" not in ant_xds:
+        pols_chars = list(partition_polarization.values[0])
+        pols_labels = [f"pol_{idx}" for idx in np.arange(0, len(pols_chars))]
+        ant_xds = ant_xds.assign_coords(receptor_label=pols_labels)
+        pol_type_values = [pols_chars] * len(ant_xds.antenna_name)
+        ant_xds = ant_xds.assign_coords(
+            polarization_type=(
+                ["antenna_name", "receptor_label"],
+                pol_type_values,
+            )
+        )
 
     ant_xds.attrs["overall_telescope_name"] = telescope_name
     return ant_xds
@@ -87,7 +102,6 @@ def extract_antenna_info(
     """
     to_new_data_variables = {
         "POSITION": ["ANTENNA_POSITION", ["antenna_name", "cartesian_pos_label"]],
-        "OFFSET": ["ANTENNA_FEED_OFFSET", ["antenna_name", "cartesian_pos_label"]],
         "DISH_DIAMETER": ["ANTENNA_DISH_DIAMETER", ["antenna_name"]],
     }
 
@@ -95,7 +109,7 @@ def extract_antenna_info(
         "NAME": ["antenna_name", ["antenna_name"]],
         "STATION": ["station", ["antenna_name"]],
         "MOUNT": ["mount", ["antenna_name"]],
-        "PHASED_ARRAY_ID": ["phased_array_id", ["antenna_name"]],
+        # "PHASED_ARRAY_ID": ["phased_array_id", ["antenna_name"]],
         "antenna_id": ["antenna_id", ["antenna_name"]],
     }
 
@@ -124,9 +138,8 @@ def extract_antenna_info(
 
     ant_xds["ANTENNA_DISH_DIAMETER"].attrs.update({"units": ["m"], "type": "quantity"})
 
-    ant_xds["ANTENNA_FEED_OFFSET"].attrs["type"] = "earth_location_offset"
-    ant_xds["ANTENNA_FEED_OFFSET"].attrs["coordinate_system"] = "geocentric"
     ant_xds["ANTENNA_POSITION"].attrs["coordinate_system"] = "geocentric"
+    ant_xds["ANTENNA_POSITION"].attrs["origin_object_name"] = "earth"
 
     if telescope_name in ["ALMA", "VLA", "NOEMA", "EVLA"]:
         # antenna_name = ant_xds["antenna_name"].values + "_" + ant_xds["station"].values
@@ -202,12 +215,17 @@ def extract_feed_info(
         taql_where=f" where (ANTENNA_ID IN [{','.join(map(str, ant_xds.antenna_id.values))}]) AND (FEED_ID IN [{','.join(map(str, feed_id))}])",
     )  # Some Lofar and MeerKAT data have the spw column set to -1 so we can't use '(SPECTRAL_WINDOW_ID = {spectral_window_id})'
 
+    if not generic_feed_xds:
+        # Some MSv2 have a FEED table that does not cover all antenna_id (and feed_id)
+        return ant_xds
+
     feed_spw = np.unique(generic_feed_xds.SPECTRAL_WINDOW_ID)
     if len(feed_spw) == 1 and feed_spw[0] == -1:
         generic_feed_xds = generic_feed_xds.isel(SPECTRAL_WINDOW_ID=0, drop=True)
     else:
         if spectral_window_id not in feed_spw:
-            return ant_xds  # For some spw the feed table is empty (this is the case with ALMA spw WVR#NOMINAL).
+            # For some spw the feed table is empty (this is the case with ALMA spw WVR#NOMINAL).
+            return ant_xds
         else:
             generic_feed_xds = generic_feed_xds.sel(
                 SPECTRAL_WINDOW_ID=spectral_window_id, drop=True
@@ -228,14 +246,14 @@ def extract_feed_info(
     ), "The number of receptors must be constant in feed table."
 
     to_new_data_variables = {
-        "BEAM_OFFSET": [
-            "BEAM_OFFSET",
-            ["antenna_name", "receptor_label", "sky_dir_label"],
+        "RECEPTOR_ANGLE": [
+            "ANTENNA_RECEPTOR_ANGLE",
+            ["antenna_name", "receptor_label"],
         ],
-        "RECEPTOR_ANGLE": ["RECEPTOR_ANGLE", ["antenna_name", "receptor_label"]],
-        # "pol_response": ["POLARIZATION_RESPONSE", ["antenna_name", "receptor_label", "receptor_name_"]] #repeated dim creates problems.
-        "FOCUS_LENGTH": ["FOCUS_LENGTH", ["antenna_name"]],  # optional
-        # "position": ["ANTENNA_FEED_OFFSET",["antenna_name", "cartesian_pos_label"]] #Will be added to the existing position in ant_xds
+        "FOCUS_LENGTH": [
+            "ANTENNA_FOCUS_LENGTH",
+            ["antenna_name"],
+        ],  # optional
     }
 
     to_new_coords = {
@@ -249,29 +267,31 @@ def extract_feed_info(
         to_new_coords=to_new_coords,
     )
 
-    # print('ant_xds["ANTENNA_FEED_OFFSET"]',ant_xds["ANTENNA_FEED_OFFSET"].data)
-    # print('generic_feed_xds["POSITION"].data',generic_feed_xds["POSITION"].data)
-    feed_offset_attrs = ant_xds["ANTENNA_FEED_OFFSET"].attrs
-    ant_xds["ANTENNA_FEED_OFFSET"] = (
-        ant_xds["ANTENNA_FEED_OFFSET"] + generic_feed_xds["POSITION"].data
-    )
-    # recover attrs after arithmetic operation
-    ant_xds["ANTENNA_FEED_OFFSET"].attrs.update(feed_offset_attrs)
-
-    coords = {}
     # coords["receptor_label"] = "pol_" + np.arange(ant_xds.sizes["receptor_label"]).astype(str) #Works on laptop but fails in github test runner.
-    coords["receptor_label"] = np.array(
-        list(
-            map(
-                lambda x, y: x + "_" + y,
-                ["pol"] * ant_xds.sizes["receptor_label"],
-                np.arange(ant_xds.sizes["receptor_label"]).astype(str),
-            )
+    coords = {
+        "receptor_label": np.array(
+            list(
+                map(
+                    lambda x, y: x + "_" + y,
+                    ["pol"] * ant_xds.sizes["receptor_label"],
+                    np.arange(ant_xds.sizes["receptor_label"]).astype(str),
+                )
+            ),
+            dtype=str,
         )
-    )
+    }
 
-    coords["sky_dir_label"] = ["ra", "dec"]
     ant_xds = ant_xds.assign_coords(coords)
+
+    # Correct to expected types. Some ALMA-SD (at least) leave receptor_label, polarization_type columns
+    # in the MS empty, causing a type mismatch
+    if (
+        "polarization_type" in ant_xds.coords
+        and ant_xds.coords["polarization_type"].dtype != str
+    ):
+        ant_xds.coords["polarization_type"] = ant_xds.coords[
+            "polarization_type"
+        ].astype(str)
     return ant_xds
 
 
@@ -370,6 +390,11 @@ def create_gain_curve_xds(
         }
     )
 
+    # correct expected types (for example "GAIN_CURVE" can be float32)
+    for data_var in gain_curve_xds:
+        if gain_curve_xds.data_vars[data_var].dtype != np.float64:
+            gain_curve_xds[data_var] = gain_curve_xds[data_var].astype(np.float64)
+
     return gain_curve_xds
 
 
@@ -445,7 +470,7 @@ def create_phase_calibration_xds(
         "TIME": ["time_phase_cal", ["time_phase_cal"]],
     }
 
-    phase_cal_xds = xr.Dataset(attrs={"type": "phase_cal"})
+    phase_cal_xds = xr.Dataset(attrs={"type": "phase_calibration"})
     phase_cal_xds = convert_generic_xds_to_xradio_schema(
         generic_phase_cal_xds, phase_cal_xds, to_new_data_variables, to_new_coords
     )
@@ -494,8 +519,8 @@ def create_phase_calibration_xds(
     time_coord_attrs = {
         "type": "time",
         "units": ["s"],
-        "scale": "UTC",
-        "format": "UNIX",
+        "scale": "utc",
+        "format": "unix",
     }
 
     # If we interpolate rename the time_phase_cal axis to time.

xradio/{correlated_data/_utils/_ms → measurement_set/_utils/_msv2}/create_field_and_source_xds.py

@@ -6,20 +6,27 @@ import numpy as np
 import xarray as xr
 
 import toolviper.utils.logger as logger
-from xradio.correlated_data._utils._ms.msv4_sub_xdss import interpolate_to_time
-from xradio.correlated_data._utils._ms.subtables import subt_rename_ids
-from xradio.correlated_data._utils._ms._tables.read import (
+from xradio.measurement_set._utils._msv2.msv4_sub_xdss import interpolate_to_time
+from xradio.measurement_set._utils._msv2.subtables import subt_rename_ids
+from xradio.measurement_set._utils._msv2._tables.read import (
     convert_casacore_time_to_mjd,
     make_taql_where_between_min_max,
     load_generic_table,
 )
-from xradio._utils.common import cast_to_str, convert_to_si_units, add_position_offsets
+from xradio._utils.list_and_array import cast_to_str
+from xradio._utils.coord_math import (
+    convert_to_si_units,
+    add_position_offsets,
+    wrap_to_pi,
+)
+
 from xradio._utils.list_and_array import (
     check_if_consistent,
     unique_1d,
     to_np_array,
 )
 from xradio._utils.schema import (
+    casacore_to_msv4_measure_type,
     column_description_casacore_to_msv4_measure,
     convert_generic_xds_to_xradio_schema,
 )
@@ -169,9 +176,21 @@ def extract_ephemeris_info(
     ), "Only geocentric observer ephemeris are supported."
 
     if "posrefsys" in ephemeris_meta:
-        sky_coord_frame = ephemeris_meta["posrefsys"].replace("ICRF/", "")
+        # Note the phase center can be given as "J2000" or "J2000.0"
+        ref_frame = (
+            ephemeris_meta["posrefsys"]
+            .replace("ICRF/", "", 1)
+            .replace("J2000.0", "J2000", 1)
+        )
+        if ref_frame in casacore_to_msv4_measure_type["direction"].get("Ref_map", {}):
+            ref_frame = casacore_to_msv4_measure_type["direction"]["Ref_map"][ref_frame]
+        else:
+            logger.debug(
+                f"Unrecognized casacore direction reference frame found in posrefsys: {ref_frame}"
+            )
+        sky_coord_frame = ref_frame.lower()
     else:
-        sky_coord_frame = "ICRS"  # We will have to just assume this.
+        sky_coord_frame = "icrs"  # We will have to just assume this.
 
     # Find out witch keyword is used for units (UNIT/QuantumUnits)
     if "UNIT" in ephemeris_column_description["RA"]["keywords"]:
@@ -195,7 +214,7 @@ def extract_ephemeris_info(
             "type": "location",
             "units": ["deg", "deg", "m"],
             "data": observer_position,
-            "ellipsoid": "WGS84",
+            "frame": "WGS84",
             "origin_object_name": "Earth",
             "coordinate_system": ephemeris_meta["obsloc"].lower(),
         }
@@ -260,7 +279,7 @@ def extract_ephemeris_info(
                 }
             )
 
-    # Add optional data: SUB_OBSERVER_POSITION and SUB_SOLAR_POSITION
+    # Add optional data: SUB_OBSERVER_DIRECTION and SUB_SOLAR_POSITION
     if "DiskLong" in ephemeris_column_description:
         key_lon = "DiskLong"
         key_lat = "DiskLat"
@@ -283,7 +302,7 @@ def extract_ephemeris_info(
         temp_xds["SUB_OBSERVER_DIRECTION"].attrs.update(
             {
                 "type": "location",
-                "ellipsoid": "NA",
+                "frame": "Undefined",
                 "origin_object_name": ephemeris_meta["NAME"],
                 "coordinate_system": "planetodetic",
                 "units": [
@@ -312,7 +331,7 @@ def extract_ephemeris_info(
         temp_xds["SUB_SOLAR_POSITION"].attrs.update(
             {
                 "type": "location",
-                "ellipsoid": "NA",
+                "frame": "Undefined",
                 "origin_object_name": "Sun",
                 "coordinate_system": "planetodetic",
                 "units": [
@@ -339,8 +358,8 @@ def extract_ephemeris_info(
     time_coord_attrs = {
         "type": "time",
         "units": ["s"],
-        "scale": "UTC",
-        "format": "UNIX",
+        "scale": "utc",
+        "format": "unix",
     }
     temp_xds["time_ephemeris"].attrs.update(time_coord_attrs)
 
@@ -374,21 +393,28 @@ def extract_ephemeris_info(
             interp_time is not None
         ), 'ephemeris_interpolate must be True if there is ephemeris data and multiple fields (this will occur if "FIELD_ID" is not in partition_scheme).'
 
+        field_phase_center = wrap_to_pi(
+            xds[center_dv].values + xds["SOURCE_LOCATION"][:, 0:2].values
+        )
+        field_phase_center = np.column_stack(
+            (field_phase_center, np.zeros(xds[center_dv].values.shape[0]))
+        )
+        field_phase_center[:, -1] = (
+            field_phase_center[:, -1] + xds["SOURCE_LOCATION"][:, -1].values
+        )
+
         xds[center_dv] = xr.DataArray(
-            add_position_offsets(
-                np.column_stack(
-                    (xds[center_dv].values, np.zeros(xds[center_dv].values.shape[0]))
-                ),
-                xds["SOURCE_LOCATION"].values,
-            ),
+            field_phase_center,
             dims=[xds["SOURCE_LOCATION"].dims[0], "sky_pos_label"],
         )
     else:
+        field_phase_center = (
+            np.append(xds[center_dv].values, 0) + xds["SOURCE_LOCATION"].values
+        )
+        field_phase_center[:, 0:2] = wrap_to_pi(field_phase_center[:, 0:2])
+
         xds[center_dv] = xr.DataArray(
-            add_position_offsets(
-                np.append(xds[center_dv].values, 0),
-                xds["SOURCE_LOCATION"].values,
-            ),
+            field_phase_center,
             dims=[xds["SOURCE_LOCATION"].dims[0], "sky_pos_label"],
         )
 
@@ -460,6 +486,65 @@ def make_line_dims_and_coords(
     return line_dims, line_coords
 
 
+def pad_missing_sources(
+    source_xds: xr.Dataset, unique_source_ids: np.array
+) -> xr.Dataset:
+    """
+    In some MSs there can be source IDs referenced from the field subtable which do not exist in
+    the source table: https://github.com/casangi/xradio/issues/266
+
+    This addresses the issue by padding/filling those IDs with "Unknown"/nan values. Produces a
+    source_xds that, in addition to the information loaded for the non-missing source IDs, has
+    padding for the IDs that are missing from the input MSv2 source table.
+    This function does not need to do anything when unique_source_ids is a single value
+    (partitioning by "FIELD_ID" or othwerwise single field/source)
+
+    Parameters:
+    ----------
+    xds: xr.Dataset
+        source dataset to fix/pad missing sources
+    unique_source_ids: np.array
+        IDs of the sources included in this partition
+
+    Returns:
+    -------
+    filled_source_xds : xr.Dataset
+        source dataset with padding in the originally missing sources
+    """
+
+    # Only fill gaps in multi-source xdss. If single source_id, no need to
+    if len(unique_source_ids) <= 1:
+        return source_xds
+
+    missing_source_ids = [
+        source_id
+        for source_id in unique_source_ids
+        if source_id not in source_xds.coords["SOURCE_ID"]
+    ]
+
+    # would like to use the new-ish xr.pad, but it creates issues with indices/coords and is
+    # also not free of overheads, as it for example changes all numeric types to float64
+    missing_source_xds = xr.full_like(source_xds.isel(SOURCE_ID=0), fill_value=np.nan)
+    pad_str = "Unknown"
+    pad_str_type = "<U9"
+    for var in missing_source_xds.data_vars:
+        if np.issubdtype(missing_source_xds.data_vars[var].dtype, np.str_):
+            # Avoid truncation to length of previously loaded strings
+            missing_source_xds[var] = missing_source_xds[var].astype(
+                np.dtype(pad_str_type)
+            )
+            missing_source_xds[var] = pad_str
+
+    concat_dim = "SOURCE_ID"
+    xdss_to_concat = [source_xds]
+    for missing_id in missing_source_ids:
+        missing_source_xds[concat_dim] = missing_id
+        xdss_to_concat.append(missing_source_xds)
+    filled_source_xds = xr.concat(xdss_to_concat, concat_dim).sortby(concat_dim)
+
+    return filled_source_xds
+
+
 def extract_source_info(
     xds: xr.Dataset,
     path: str,
@@ -542,6 +627,8 @@ def extract_source_info(
         "column_descriptions"
     ]
 
+    source_xds = pad_missing_sources(source_xds, unique_source_id)
+
     # Get source name (the time axis is optional and will probably be required if the partition scheme does not include 'FIELD_ID' or 'SOURCE_ID'.).
     # Note again that this optional time axis has nothing to do with the original time axis in the source table that we drop.
     if len(source_id) == 1: