cubevis 0.5.14__py3-none-any.whl → 0.5.16__py3-none-any.whl

cubevis/data/measurement_set/processing_set/_ps_data.py

@@ -1,213 +0,0 @@
-'''
-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

@@ -1,55 +0,0 @@
-"""
-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

@@ -1,154 +0,0 @@
-'''
-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

@@ -1,91 +0,0 @@
-''' 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