cubevis 0.5.2__py3-none-any.whl

cubevis/data/measurement_set/processing_set/_ps_coords.py

@@ -0,0 +1,78 @@
+'''
+Modify/add xarray Dataset coordinates for plotting.
+'''
+
+import numpy as np
+from pandas import to_datetime
+
+def set_coordinates(ms_xdt):
+    ''' Convert coordinate units and add baseline coordinate for plotting.
+        Returns xarray.Dataset
+    '''
+    _set_coordinate_unit(ms_xdt)
+    _set_frequency_unit(ms_xdt)
+    return _add_baseline_coordinate(ms_xdt)
+
+def set_datetime_coordinate(ms_xds):
+    ''' Convert float time to datetime for plotting. '''
+    time_attrs = ms_xds.time.attrs
+    try:
+        ms_xds.coords['time'] = to_datetime(ms_xds.time, unit=time_attrs['units'], origin=time_attrs['format'])
+    except TypeError:
+        ms_xds.coords['time'] = to_datetime(ms_xds.time, unit=time_attrs['units'][0], origin=time_attrs['format'])
+    ms_xds.time.attrs = time_attrs
+
+def _set_coordinate_unit(ms_xdt):
+    ''' Set coordinate units attribute as string not list for plotting. '''
+    for coord in ms_xdt.coords:
+        # Plots need unit to be string not list
+        if 'units' in ms_xdt.coords[coord].attrs:
+            units = ms_xdt.coords[coord].units
+            if isinstance(units, list) and len(units) == 1:
+                ms_xdt.coords[coord].attrs['units'] = units[0]
+
+def set_index_coordinates(ms_xds, coordinates):
+    ''' Return ms_xds with new coordinate for string values (name) then replace coordinate with numerical index. '''
+    for coordinate in coordinates:
+        if coordinate == "polarization":
+            ms_xds = ms_xds.assign_coords({"polarization_name": (ms_xds.polarization.dims, ms_xds.polarization.values)})
+            ms_xds["polarization"] = np.array(range(ms_xds.polarization.size))
+        elif coordinate == "baseline":
+            ms_xds = ms_xds.assign_coords({"baseline_name": (ms_xds.baseline.dims, ms_xds.baseline.values)})
+            ms_xds["baseline"] = np.array(range(ms_xds.baseline.size))
+        elif coordinate == "antenna_name":
+            ms_xds = ms_xds.assign_coords({"antenna": (ms_xds.antenna_name.dims, ms_xds.antenna_name.values)})
+            ms_xds["antenna_name"] = np.array(range(ms_xds.antenna_name.size))
+    return ms_xds
+
+def _set_frequency_unit(ms_xdt):
+    ''' Convert frequency to GHz. Note attrs (channel_width, reference_frequency) still have Hz units in dict '''
+    if ms_xdt.frequency.attrs['units'] == "Hz":
+        frequency_xda = ms_xdt.frequency / 1e9
+        frequency_attrs = ms_xdt.frequency.attrs
+        frequency_attrs['units'] = "GHz"
+        frequency_xda = frequency_xda.assign_attrs(frequency_attrs)
+        ms_xdt.coords['frequency'] = frequency_xda
+
+def _add_baseline_coordinate(ms_xdt):
+    '''
+        Replace "baseline_id" (int) with "baseline" (string) coordinate "ant1 & ant2".
+        Baseline ids are not consistent across ms_xdts. 
+    '''
+    # Cannot assign coords to DataTree.
+    baseline_ms_xdt = ms_xdt.to_dataset() # mutable Dataset
+
+    if 'baseline_id' not in baseline_ms_xdt.coords:
+        return baseline_ms_xdt
+
+    ant1_names = ms_xdt.baseline_antenna1_name.values
+    ant2_names = ms_xdt.baseline_antenna2_name.values
+    if ant1_names.size == 1:
+        baseline_names = f"{ant1_names.item()} & {ant2_names.item()}"
+        baseline_ms_xdt = baseline_ms_xdt.assign_coords({"baseline": np.array(baseline_names)})
+    else:
+        baseline_names = [f"{ant1_names[idx]} & {ant2_names[idx]}" for idx in range(len(ant1_names))]
+        baseline_ms_xdt = baseline_ms_xdt.assign_coords({"baseline": ("baseline_id", np.array(baseline_names))})
+        baseline_ms_xdt = baseline_ms_xdt.swap_dims({"baseline_id": "baseline"})
+    baseline_ms_xdt = baseline_ms_xdt.drop("baseline_id")
+    return baseline_ms_xdt

cubevis/data/measurement_set/processing_set/_ps_data.py

@@ -0,0 +1,213 @@
+'''
+MeasurementSet data backend using xradio Processing Set.
+'''
+
+import numpy as np
+import pandas as pd
+
+try:
+    from cubevis.data.measurement_set.processing_set._ps_io import get_processing_set
+    _HAVE_XRADIO = True
+    from cubevis.data.measurement_set.processing_set._ps_select import select_ps
+    from cubevis.data.measurement_set.processing_set._ps_stats import calculate_ps_stats
+    from cubevis.data.measurement_set.processing_set._ps_raster_data import raster_data
+    from cubevis.data.measurement_set.processing_set._xds_data import get_correlated_data
+except ImportError as e:
+    _HAVE_XRADIO = False
+
+
+class PsData:
+    '''
+    Class implementing data backend using xradio Processing Set for accessing and selecting MeasurementSet data.
+    '''
+
+    def __init__(self, ms, logger):
+        if not _HAVE_XRADIO:
+            raise RuntimeError("xradio package not available for reading MeasurementSet")
+
+        if not ms:
+            raise RuntimeError("MS path not available for reading MeasurementSet")
+
+        # Open processing set from zarr
+        # Converts msv2 if ms path is not zarr
+        self._ps_xdt, self._zarr_path = get_processing_set(ms, logger)
+
+        self._logger = logger
+        self._selection = {}
+        self._selected_ps_xdt = None # cumulative selection
+
+    def get_path(self):
+        ''' Return path to zarr file (input or converted from msv2) '''
+        return self._zarr_path
+
+    def summary(self, data_group='base', columns=None):
+        ''' Print full or selected summary of Processing Set metadata, optionally by ms '''
+        ps_summary = self._ps_xdt.xr_ps.summary(data_group=data_group)
+        pd.set_option("display.max_rows", len(self._ps_xdt))
+        pd.set_option("display.max_columns", len(ps_summary.columns))
+        pd.set_option("display.max_colwidth", None)
+
+        if columns is None:
+            print(ps_summary)
+        elif columns == "by_ms":
+            for row in ps_summary.itertuples(index=False):
+                print(f"MSv4 name: {row[0]}")
+                print(f"intent: {row[1]}")
+                shape = row[2]
+                print(f"shape: {shape[0]} times, {shape[1]} baselines, {shape[2]} channels, {shape[3]} polarizations")
+                print(f"polarization: {row[3]}")
+                scans = [str(scan) for scan in row[4]]
+                print(f"scan_name: {scans}")
+                print(f"spw_name: {row[5]}")
+                fields = [str(field) for field in row[6]]
+                print(f"field_name: {fields}")
+                sources = [str(source) for source in row[7]]
+                print(f"source_name: {sources}")
+                lines = [str(line) for line in row[8]]
+                print(f"line_name: {lines}")
+                field_coords = row[9]
+                print(f"field_coords: ({field_coords[0]}) {field_coords[1]} {field_coords[2]}")
+                print(f"frequency range: {row[10]:e} - {row[11]:e}")
+                print("-----")
+        else:
+            if isinstance(columns, str):
+                columns = [columns]
+            col_df = ps_summary[columns]
+            print(col_df)
+
+    def get_summary(self):
+        ''' Return summary of original ps '''
+        return self._ps_xdt.xr_ps.summary()
+
+    def get_data_groups(self):
+        ''' Returns set of data group names in Processing Set data. '''
+        data_groups = []
+        for ms_xdt_name in self._ps_xdt:
+            data_groups.extend(list(self._ps_xdt[ms_xdt_name].data_groups))
+        return set(data_groups)
+
+    def get_antennas(self, plot_positions=False, label_antennas=False):
+        ''' Returns list of antenna names in ProcessingSet antenna_xds.
+                plot_positions (bool): show plot of antenna positions.
+                label_antennas (bool): label positions with antenna names.
+        '''
+        if plot_positions:
+            self._ps_xdt.xr_ps.plot_antenna_positions(label_antennas)
+        return self._ps_xdt.xr_ps.get_combined_antenna_xds().antenna_name.values.tolist()
+
+    def plot_phase_centers(self, label_all_fields=False, data_group='base'):
+        ''' Plot the phase center locations of all fields in the Processing Set (original or selected) and label central field.
+                label_all_fields (bool); label all fields on the plot
+                data_group (str); data group to use for processing.
+        '''
+        self._ps_xdt.xr_ps.plot_phase_centers(label_all_fields, data_group)
+
+    def get_ps_len(self):
+        ''' Returns number of ms_xdt in selected ps_xdt (if selected) '''
+        return len(self._get_ps_xdt())
+
+    def get_max_dims(self):
+        ''' Returns maximum length of data dimensions in selected ps_xdt (if selected) '''
+        ps_xdt = self._get_ps_xdt()
+        return ps_xdt.xr_ps.get_max_dims()
+
+    def get_data_dimensions(self):
+        ''' Return the maximum dimensions in selected ps_xdt (if selected) '''
+        dims = list(self.get_max_dims().keys())
+        if 'uvw_label' in dims:
+            dims.remove('uvw_label') # not a VISIBILITY/SPECTRUM data dim
+        return dims
+
+    def get_dimension_values(self, dimension):
+        ''' Return sorted list of unique values for input dimension in ProcessingSet. '''
+        ps_xdt = self._get_ps_xdt()
+        dim_values = []
+        for ms_xdt in ps_xdt.values():
+            if dimension == 'baseline':
+                ant1_names = ms_xdt.baseline_antenna1_name.values
+                ant2_names = ms_xdt.baseline_antenna2_name.values
+                for baseline_id in ms_xdt.baseline_id:
+                    dim_values.append(f"{ant1_names[baseline_id]} & {ant2_names[baseline_id]}")
+            else:
+                try:
+                    dim_values.extend([value.item() for value in ms_xdt[dimension].values])
+                except TypeError:
+                    dim_values.append(ms_xdt[dimension].values.item())
+        return sorted(set(dim_values))
+
+    def get_dimension_attrs(self, dim):
+        ''' Return attributes dict for input dimension in ProcessingSet. '''
+        ps_xdt = self._get_ps_xdt()
+        return ps_xdt.get(0)[dim].attrs
+
+    def get_first_spw(self):
+        ''' Return first spw name by id '''
+        spw_id_names = {}
+        ps_xdt = self._get_ps_xdt()
+        for ms_xdt in ps_xdt.values():
+            freq_xds = ms_xdt.frequency
+            spw_id_names[freq_xds.spectral_window_id] = freq_xds.spectral_window_name
+
+        first_spw_id = min(spw_id_names)
+        first_spw_name = spw_id_names[first_spw_id]
+
+        summary = self.get_summary()
+        spw_df = summary[summary['spw_name'] == first_spw_name]
+        start_freq = spw_df.at[spw_df.index[0], 'start_frequency']
+        end_freq = spw_df.at[spw_df.index[0], 'end_frequency']
+        self._logger.info(f"Selecting first spw {first_spw_name} (id {first_spw_id}) with frequency range {start_freq:e} - {end_freq:e}")
+        return first_spw_name
+
+    def select_data(self, selection):
+        ''' Apply selection dict to ProcessingSet to create selected ps_xdt.
+            If previous selection done, apply to selected ps_xdt.
+            Add selection to previous selections. '''
+        ps_xdt = self._get_ps_xdt()
+        self._selected_ps_xdt = select_ps(ps_xdt, selection, self._logger)
+        if self._selection:
+            self._selection |= selection
+        else:
+            self._selection = selection
+
+    def clear_selection(self):
+        ''' Clear previous selections and use original ps_xdt '''
+        self._selection = None
+        self._selected_ps_xdt = None
+
+    def get_vis_stats(self, selection, vis_axis):
+        ''' Returns statistics (min, max, mean, std) for data selected by selection.
+                selection (dict): fields and values to select
+        '''
+        stats_ps_xdt = select_ps(self._ps_xdt, selection, self._logger)
+        data_group = selection['data_group'] if 'data_group' in selection else 'base'
+        return calculate_ps_stats(stats_ps_xdt, self._zarr_path, vis_axis, data_group, self._logger)
+
+    def get_correlated_data(self, data_group):
+        ''' Returns name of 'correlated_data' in Processing Set data_group '''
+        ps_xdt = self._get_ps_xdt()
+        for ms_xdt in ps_xdt.values():
+            if data_group in ms_xdt.attrs['data_groups']:
+                return get_correlated_data(ms_xdt.ds, data_group)
+        raise RuntimeError(f"No correlated data for data group {data_group}")
+
+    def get_raster_data(self, plot_inputs):
+        ''' Returns xarray Dataset after applying plot inputs and raster plane selection '''
+        return raster_data(self._get_ps_xdt(),
+            plot_inputs,
+            self._logger
+        )
+
+    def _get_ps_xdt(self):
+        ''' Returns selected ps_xdt if selection has been done, else original ps_xdt '''
+        return self._selected_ps_xdt if self._selected_ps_xdt else self._ps_xdt
+
+    def _get_unique_values(self, df_col):
+        ''' Return unique values in pandas Dataframe column, for summary '''
+        values = df_col.to_numpy()
+        try:
+            # numeric arrays
+            return np.unique(np.concatenate(values))
+        except ValueError:
+            # string arrays
+            all_values = [row[0] for row in values]
+            return np.unique(np.concatenate(all_values))

cubevis/data/measurement_set/processing_set/_ps_io.py

@@ -0,0 +1,55 @@
+"""
+Functions for reading MSv2 or zarr MSv4 into xradio ProcessingSet
+"""
+
+import os.path
+
+from xradio.measurement_set.open_processing_set import open_processing_set
+
+try:
+    # requires python-casacore
+    from xradio.measurement_set.convert_msv2_to_processing_set import convert_msv2_to_processing_set
+    __HAVE_CASACORE = True
+except ImportError:
+    __HAVE_CASACORE = False
+
+def get_processing_set(ms_path, logger):
+    '''
+    Read msv2 or zarr file into processing set
+
+    Args:
+        ms_path (str): path to MSv2 or MSv4 zarr file
+    Returns:
+        xradio ProcessingSet
+    '''
+
+    if not os.path.exists(ms_path):
+        raise RuntimeError(f"Visibility file {ms_path} does not exist")
+
+    # Remove trailing / or entire path is basename with no extension
+    if ms_path[-1] == '/':
+        ms_path = ms_path[:-1]
+
+    basename, ext = os.path.splitext(ms_path)
+    if ext == ".zarr":
+        zarr_path = ms_path
+    else:
+        if not __HAVE_CASACORE:
+            raise RuntimeError("Cannot convert MSv2 to xradio zarr file: python-casacore not installed.")
+        zarr_path = basename + ".ps.zarr"
+        if not os.path.exists(zarr_path):
+            logger.info(f"Converting input MS {ms_path} to zarr {zarr_path}")
+            convert_msv2_to_processing_set(
+                in_file=ms_path,
+                out_file=zarr_path
+                )
+
+    if not os.path.exists(zarr_path):
+        raise RuntimeError("Zarr file does not exist")
+
+    ps = open_processing_set(zarr_path)
+    if not ps or len(ps) == 0:
+        raise RuntimeError("Failed to read measurement set into processing set.")
+    logger.info(f"Processing set contains {len(ps)} msv4 datasets.")
+
+    return ps, zarr_path

cubevis/data/measurement_set/processing_set/_ps_raster_data.py

@@ -0,0 +1,154 @@
+'''
+Functions to create a raster xarray Dataset from xradio ProcessingSet after applying plot inputs
+'''
+
+import numpy as np
+
+from xradio.measurement_set._utils._utils.stokes_types import stokes_types
+
+from cubevis.data.measurement_set.processing_set._ps_concat import concat_ps_xdt
+from cubevis.data.measurement_set.processing_set._ps_coords import set_datetime_coordinate
+from cubevis.data.measurement_set.processing_set._ps_select import select_ps
+from cubevis.data.measurement_set.processing_set._xds_data import get_axis_data
+
+def raster_data(ps_xdt, plot_inputs, logger):
+    '''
+    Create raster xds: y_axis vs x_axis for vis axis.
+        ps_xdt (xarray DataTree): input datasets.
+        plot_inputs (dict): user inputs for plot
+        logger (graphviper logger): logger
+    Returns: selected xarray Dataset of visibility component and updated selection
+    '''
+    raster_xdt, dim_selection = _select_raster_ps_xdt(ps_xdt, plot_inputs, logger)
+    plot_inputs['dim_selection'] = dim_selection
+
+    # Create xds from concat ms_xds in ps
+    raster_xds = concat_ps_xdt(raster_xdt, logger)
+    correlated_data = plot_inputs['correlated_data']
+    if raster_xds[correlated_data].count() == 0:
+        raise RuntimeError("Plot failed: raster plane selection yielded data with all nan values.")
+
+    # Set complex component of vis data
+    raster_xds[correlated_data] = get_axis_data(raster_xds,
+        plot_inputs['vis_axis'],
+        plot_inputs['selection']['data_group_name']
+    )
+
+    # Convert float time to datetime
+    set_datetime_coordinate(raster_xds)
+
+    # Apply aggregator
+    raster_xds = aggregate_data(raster_xds, plot_inputs, logger)
+
+    logger.debug(f"Plotting visibility data with shape: {raster_xds[correlated_data].shape}")
+    return raster_xds
+
+def _select_raster_ps_xdt(ps_xdt, plot_inputs, logger):
+    ''' Select default dimensions if needed for raster data '''
+    # Determine which dims must be selected, add to selection, and do selection
+    input_selection = plot_inputs['selection']
+    dims_to_select = _get_raster_selection_dims(plot_inputs)
+    dim_selection = {} # return value
+
+    if dims_to_select:
+        for dim in dims_to_select:
+            # Select first value (by index) and add to dim selection, or apply iter_axis value
+            # (user selection would have been applied previously)
+            if dim not in input_selection:
+                dim_selection[dim] = _get_first_dim_value(ps_xdt, dim, plot_inputs, logger)
+            elif dim == plot_inputs['iter_axis']:
+                dim_selection[dim] = input_selection[dim]
+        if dim_selection:
+            logger.info(f"Applying raster plane selection (using first index or iter value): {dim_selection}")
+            return select_ps(ps_xdt, dim_selection, logger), dim_selection
+    return ps_xdt, dim_selection
+
+def _get_raster_selection_dims(plot_inputs):
+    ''' Return which dimensions should be selected for raster plot.
+        List of dimensions which are not x, y, or agg axis. '''
+    data_dims = plot_inputs['data_dims'].copy()
+    if plot_inputs['x_axis'] in data_dims:
+        data_dims.remove(plot_inputs['x_axis'])
+    if plot_inputs['y_axis'] in data_dims:
+        data_dims.remove(plot_inputs['y_axis'])
+    if plot_inputs['aggregator'] and plot_inputs['agg_axis']:
+        for axis in plot_inputs['agg_axis']:
+            data_dims.remove(axis)
+    return data_dims
+
+def _get_first_dim_value(ps_xdt, dim, plot_inputs, logger):
+    ''' Return value of first dimension by index for polarization or by value for others. '''
+    # If iter_axis, get first dim value after iter value is selected to avoid empty selected ps
+    iter_axis = plot_inputs['iter_axis'] if 'iter_axis' in plot_inputs else None
+    iter_ps = ps_xdt
+    if iter_axis:
+        iter_selection = {iter_axis: plot_inputs['selection'][iter_axis]}
+        iter_ps = select_ps(ps_xdt, iter_selection, logger)
+
+    values = []
+    if dim == "polarization":
+        # Get sorted list of polarization ids used
+        stokes = list(stokes_types.values())
+        for ms_xdt in iter_ps.values():
+            for val in ms_xdt.polarization.values:
+                values.append(stokes.index(val))
+        sorted_values = sorted(list(set(values)))
+        # Get pol name for pol id
+        return stokes[sorted_values[0]]
+
+    if dim == 'baseline':
+        for ms_xdt in iter_ps.values():
+            ant1_names = ms_xdt.baseline_antenna1_name.values
+            ant2_names = ms_xdt.baseline_antenna2_name.values
+            for baseline_id in ms_xdt.baseline_id.values:
+                values.append(ant1_names[baseline_id] + " & " + ant2_names[baseline_id])
+            return sorted(list(set(values)))[0]
+
+    # Get sorted values list
+    for ms_xdt in iter_ps.values():
+        values.extend(ms_xdt[dim].values.tolist())
+    return sorted(list(set(values)))[0]
+
+def _add_index_dimensions(xds):
+    ''' Add index coordinates for plotting if coord has dimension '''
+    # Baseline/antenna id
+    if "baseline" in xds.coords and xds.baseline.dims:
+        xds = xds.assign_coords({"baseline_id": (xds.baseline.dims, np.array(range(xds.baseline.size)))})
+        xds = xds.swap_dims({"baseline": "baseline_id"})
+    elif "antenna_name" in xds.coords and xds.antenna_name.dims:
+        xds = xds.assign_coords({"antenna_id": (xds.antenna_name.dims, np.array(range(xds.antenna_name.size)))})
+
+    # Polarization id
+    if "polarization" in xds.coords and xds.polarization.dims:
+        xds = xds.assign_coords({"polarization_id": (xds.polarization.dims, np.array(range(xds.polarization.size)))})
+
+    return xds
+
+def aggregate_data(xds, plot_inputs, logger):
+    ''' Apply aggregator to agg axis list. '''
+    if not plot_inputs['aggregator']:
+        return xds
+
+    aggregator = plot_inputs['aggregator']
+    agg_axis = plot_inputs['agg_axis']
+    agg_xds = xds
+
+    # Check if agg axes have been selected (selection or iteration) and are no longer a dimension
+    apply_agg_axis = [axis for axis in agg_axis if axis in xds.dims]
+    logger.debug(f"Applying {aggregator} to {apply_agg_axis}.")
+
+    if aggregator == 'max':
+        agg_xds = xds.max(dim=apply_agg_axis, keep_attrs=True)
+    elif aggregator == 'mean':
+        agg_xds = xds.mean(dim=apply_agg_axis, keep_attrs=True)
+    elif aggregator == 'median':
+        agg_xds = xds.median(dim=apply_agg_axis, keep_attrs=True)
+    elif aggregator == 'min':
+        agg_xds = xds.min(dim=apply_agg_axis, keep_attrs=True)
+    elif aggregator == 'std':
+        agg_xds = xds.std(dim=apply_agg_axis, keep_attrs=True)
+    elif aggregator == 'sum':
+        agg_xds = xds.sum(dim=apply_agg_axis, keep_attrs=True)
+    elif aggregator == 'var':
+        agg_xds = xds.var(dim=apply_agg_axis, keep_attrs=True)
+    return agg_xds

cubevis/data/measurement_set/processing_set/_ps_select.py

@@ -0,0 +1,91 @@
+''' Apply selection dict to ProcessingSet and MeasurementSetXds '''
+
+def select_ps(ps_xdt, selection, logger):
+    '''
+        Apply selection dict to Processing Set.
+        Select Processing Set first (ps summary columns), then each MeasurementSetXds (data_group etc.).
+        Returns dict of selected name, ms_xdt.
+        Throws exception for empty Processing Set (null selection).
+    '''
+    if not selection:
+        return ps_xdt
+
+    # Separate PS selection and MS selection
+    ps_summary = ps_xdt.xr_ps.summary()
+    ps_selection_keys = list(ps_summary.columns.array)
+    ps_selection_keys.append('data_group')
+
+    first_ms = ps_summary['name'][0]
+    ms_selection_keys = list(ps_xdt[first_ms].coords.keys())
+
+    # Sort selections into categories
+    ps_selection = {}
+    ms_selection = {}
+    antenna_selection = {}
+    for key, val in selection.items():
+        if key in ps_selection_keys and selection[key]:
+            ps_selection[key] = val
+        if key in ms_selection_keys and selection[key]:
+            ms_selection[key] = val
+        if 'antenna1' in key or 'antenna2' in key:
+            antenna_selection[key] = val
+        elif key == 'baseline':
+            ant1, ant2 = val.split('&')
+            antenna_selection['antenna1'] = ant1.strip()
+            antenna_selection['antenna2'] = ant2.strip()
+
+    # Do PSXdt selection
+    logger.debug(f"Applying selection to ProcessingSet: {ps_selection}")
+
+    if ps_selection:
+        selected_ps_xdt = ps_xdt.xr_ps.query(**ps_selection)
+        if len(selected_ps_xdt) == 0:
+            raise RuntimeError("Selection failed: ps selection yielded empty processing set.")
+    else:
+        selected_ps_xdt = ps_xdt.copy()
+
+    # Do MSXdt selection
+    return _select_ms_xdt(selected_ps_xdt, ms_selection, antenna_selection, logger)
+
+def _select_ms_xdt(ps_xdt, ms_selection, antenna_selection, logger):
+    ''' Apply selection to each MeasurementSetXds and return ProcessingSet.
+        Remove ms_xds which do not contain selection.
+    '''
+    # Done if no selection to apply
+    if not ms_selection and not antenna_selection:
+        return ps_xdt
+
+    logger.debug(f"Applying selection to measurement set: {ms_selection}, {antenna_selection}")
+
+    # Drop ms_xdt where selection fails
+    names_to_drop = []
+
+    for name, ms_xdt in ps_xdt.items():
+        try:
+            if antenna_selection:
+                for antenna, val in antenna_selection.items():
+                    if antenna == 'antenna1':
+                        ms_xdt = ms_xdt.sel(baseline_id=ms_xdt.baseline_antenna1_name==val)
+                    else:
+                        ms_xdt = ms_xdt.sel(baseline_id=ms_xdt.baseline_antenna2_name==val)
+
+                if ms_xdt.baseline_id.size == 1:
+                    # Select baseline_id to remove dimension
+                    ms_xdt = ms_xdt.sel(baseline_id=ms_xdt.baseline_id.item())
+                elif ms_xdt.baseline_id.size == 0:
+                    names_to_drop.append(name)
+                    continue
+
+            ms_xdt = ms_xdt.xr_ms.sel(**ms_selection)
+            ps_xdt[name] = ms_xdt
+        except KeyError:
+            # selection not in this ms_xdt, do not include in returned ps_xdt
+            names_to_drop.append(name)
+
+    if names_to_drop:
+        ps_xdt = ps_xdt.drop_nodes(names_to_drop)
+
+    if len(ps_xdt) == 0:
+        raise RuntimeError("Selection failed: ms selection yielded empty processing set.")
+
+    return ps_xdt