graphviper.graph_tools.coordinate_utils

Functions

`make_time_coord`(→ Dict)	Convenience function that creates a time coordinate measures dictionary that can be used to create parallel_coords using `make_parallel_coord()` function.
`make_frequency_coord`(→ Dict)
`make_parallel_coord`(→ Dict)
`make_parallel_coord_by_gap`(→ Dict)
`interpolate_data_coords_onto_parallel_coords`(→ Dict)	Interpolate data_coords onto parallel_coords to create the `node_task_data_mapping`. For the case of string coordinates (for example antenna_name), only exact matching is performed.
`get_disk_chunk_sizes`(→ Dict[str, int])	Determine the native on-disk chunk size for each parallel coordinate dimension.

Module Contents

make_time_coord(time_start: str = '2019-10-03T19:00:00.000', time_delta: numbers.Number = 3600, n_samples: int = 10, time_scale: {'tai', 'tcb', 'tcg', 'tdb', 'tt', 'ut1', 'utc', 'local'} = 'utc') → Dict[source]

Convenience function that creates a time coordinate measures dictionary that can be used to create parallel_coords using make_parallel_coord() function.

Parameters:

time_start (str, optional) – Start time string in format YYYY-MM-DDTHH:mm:ss.SSS, by default “2019-10-03T19:00:00.000”
time_delta (numbers.Number, optional) – The increment between time samples in seconds, by default 3600.
n_samples (int, optional) – Number of time steps, by default 10.
time_scale ({"tai", "tcb", "tcg", "tdb", "tt", "ut1", "utc", "local"}, optional) – Time scale, by default “utc”.

Returns:

Time coordinate measures dictionary.

Return type:

Dict

Notes

Time coordinate measures dictionary format:

{ "dims": "time",
    "data": time_array,
    "attrs": {
    "units": "s",
    "type": "time",
    "format": "unix",
    "time_scale": time_scale,
    },
}

The time_array values are in Unix seconds.

make_frequency_coord(freq_start: numbers.Number = 3 * 10**9, freq_delta: numbers.Number = 0.4 * 10**9, n_channels: int = 50, velocity_frame: {'gcrs', 'icrs', 'hcrs', 'lsrk', 'lsrd', 'lsr'} = 'lsrk') → Dict[source]

Convenience function that creates a frequency coordinate measures dictionary that can be used to create parallel_coords using make_parallel_coord() function.

Parameters:

freq_start (numbers.Number, optional) – Start frequency in Hz, by default 3 * 10**9.
freq_delta (numbers.Number, optional) – The increment between frequency samples, by default 0.4 * 10**9.
n_channels (int, optional) – Number of frequency steps, by default 50.
velocity_frame ({'gcrs','icrs','hcrs','lsrk','lsrd','lsr'}, optional) – Velocity frame , by default lsrk.

Returns:

Frequency coordinate measures dictionary.

Return type:

Dict

Notes

Frequency coordinate measures dictionary format:

{ "dims": "frequency",
    "data": frequency_array,
    "attrs": {
    "units": "Hz",
    "type": "spectral_coord",
    "velocity_frame": velocity_frame,
    },
}

make_parallel_coord(coord: Dict | xarray.DataArray, n_chunks: None | int = None, gap: None | float = None) → Dict[source]

Creates a single parallel coordinate from a measures dictionary or a xarray.DataArray with measures attributes.

This function only returns a single parallel_coord to create parallel_coords a dictionary must be created where the keys are the dimension coordinate names and the values are the respective parallel_coord.

Parameters:

coord (Union[Dict, xr.DataArray]) –
The input measures dictionary or xarray.DataArray with measures attributes.
n_chunks (Union[None, int]) – If specified, how many chunks to divide coord into.
gap (Union[None, float]) – If specified, gaps in coordinate values greater than the given value will be used to split chunks in the coordinate.

Returns:

Parallel coordinate dictionary. See notes for structure of dictionary.

Return type:

Dict

Notes

Nomenclature used: - dim: The dimension name. - n_dim_chunks: Number of chunks into which the dimension coordinate dim has been divided.

The structure of a parallel coordinate:

parallel_coord = {
        'data': 1D list/np.ndarray of Number,
        'data_chunks': {
            0 : 1D list/np.ndarray of Number,
            ⋮
            n_dim_chunks-1 : ...,
        }
        'data_chunks_edges': 1D list/np.ndarray of Number,
        'data_chunk_slices': {
            0 : slice,
            ⋮
            n_dim_chunks-1 : slice,
        }
        'dims': (dim,),
        'attrs': measure attribute,
    }

The keys with the following meanings:

data: An array containing all the coordinate values associated with that dimension. These values do not necessarily have to match the values in the coordinates of the input data (dictionary of xarray.Datasets or ProcessingSet), as those are interpolated onto these values. The minimum and maximum values can be respectively larger or smaller than the values in the coordinates of individual xarray.Datasets; this will simply exclude that data from being processed. It’s important to note that the parallel_coords and the input data coordinates must have the same measures attributes (reference frame, units, etc.).
data_chunks: A dictionary where the values are chunks of the data and the keys are integers. This chunking determines the parallelism of the graph. The values in the chunks can overlap.
data_chunks_edges: An array with the start and end values of each chunk.
data_chunk_slices: A dictionary with the same integer chunk keys as data_chunks, where each value is a slice(start, stop) indexing that chunk into the full (unsplit) data array.
dims: The dimension coordinate name.
attrs`: The XRADIO measures attributes of the data.

Parallel coordinates can be combined into a single dictionary called the parallel_coords where the keys are the coordinate dimension names (dim, see interpolate_data_coords_onto_parallel_coords() notes).

make_parallel_coord_by_gap(coord: Dict | xarray.DataArray, gap: float) → Dict[source]

Creates a single parallel coordinate from from a measures dictionary or a xarray.DataArray with measures attributes.

This function only returns a single parallel_coord; to create parallel_coords a dictionary must be created where the keys are the dimension coordinate names and the values are the respective parallel_coord.

Parameters:

coord (Union[Dict, xr.DataArray]) –
The input measures dictionary or xarray.DataArray with measures attributes.
gap (float) – The coordinate will be split into chunks where successive values (which are assumed to be monotonically increasing) are more than gap apart.

interpolate_data_coords_onto_parallel_coords(parallel_coords: dict, input_data: Dict | xarray.DataTree, interpolation_method: {'linear', 'nearest', 'nearest-up', 'zero', 'slinear', 'quadratic', 'cubic', 'previous', 'next'} = 'nearest', assume_sorted: bool = True, ps_partition: list[str] | None = None) → Dict[source]

Interpolate data_coords onto parallel_coords to create the node_task_data_mapping. For the case of string coordinates (for example antenna_name), only exact matching is performed.

Parameters:

parallel_coords (Dict) – The parallel coordinates determine the parallelism of the map graph. The keys in the parallel coordinates can by any combination of the dimension coordinates in the input data. See notes in docstring for structure.
input_data (Union[Dict, ProcessingSet]) –
Can either be a ProcessingSet or a Dictionary of xarray.Datasets. Only coordinates are needed so no actual data is loaded into memory.
interpolation_method ({"linear", "nearest", "nearest-up", "zero", "slinear", "quadratic", "cubic", "previous", "next",}, optional) – The kind of interpolation method to use as described in Scipy documentation , by default nearest.
assume_sorted (bool, optional) – Are the data in parallel_coords and input_data monotonically increasing in value, by default True.
ps_partition (An optional list of strings ('spectral_window_name' and/or 'field_name' are currently supported); if non-empty, the function will use the meta-data of each Dataset to partition the parallel sets by these pseudo-dimensions as well as the actual Dataset dimensions specified.)

Returns:

Node task data mapping dictionary. See notes for structure of dictionary.

Return type:

Dict

Notes

Nomenclature used:

input data: A dictionary of xarray.Datasets or a ProcessingSet.
n_datasets: The number of xarray.Datasets in the input data.
dim_i: The ith dimension name.
n_dims: The number of dimensions over which parallelism will occur.
n_dim_i_chunks: Number of chunks into which the dimension coordinate dim_i has been divided.
n_nodes: Number of nodes in the mapping stage of a Map Reduce graph.
_{}: If curly brackets are preceded by an underscore, it indicates a subscript and not a dictionary value.

The structure of the parallel coordinates:

parallel_coords = {
    dim_0: {
        'data': list/np.ndarray of Number,
        'data_chunks': {
            0 : list/np.ndarray of Number,
            ⋮
            n_dim_0_chunks-1 : ...,
        }
        'data_chunks_edges': list/np.ndarray of Number,
        'data_chunk_slices': {
            0 : slice,
            ⋮
            n_dim_0_chunks-1 : slice,
        }
        'dims': (dim_0,),
        'attrs': measure attribute,
    }
    ⋮
    dim_(n_dims-1): ...
}

The dim_i dictionaries have keys with the following meanings:

data: An array containing all the coordinate values associated with that dimension. These values do not necessarily have to match the values in the coordinates of the input data (dictionary of xarray.Datasets or ProcessingSet), as those are interpolated onto these values. The minimum and maximum values can be respectively larger or smaller than the values in the coordinates of individual xarray.Datasets; this will simply exclude that data from being processed. It’s important to note that the parallel_coords and the input data coordinates must have the same measured attributes (reference frame, units, etc.).
data_chunks: A dictionary where the data is broken into chunks with integer keys. This chunking determines the parallelism of the graph. The values in the chunks can overlap.
data_chunks_edges: An array with the start and end values of each chunk.
data_chunk_slices: A dictionary with the same integer chunk keys as data_chunks, where each value is a slice(start, stop) indexing that chunk into the full (unsplit) data array. Only present when the parallel coordinate was built with n_chunks (see make_parallel_coord()).
dims: The dimension coordinate name.
attrs`: The XRADIO measures attributes of the data.

The node_task_data_mapping is a dictionary where each key is the node id of the nodes in the mapping stage of the graph and has the following structure:

node_task_data_mapping = {
    0 : {
        'chunk_indices': tuple of int,
        'parallel_dims': (dim_0, ..., dim_{n_dims-1}),
        'data_selection': {
                dataset_name_0: {
                        dim_0: slice,
                        ⋮
                        dim_(n_dims-1): slice
                }
                ⋮
                dataset_name_{n_dataset-1}: ...
        }
        'task_coords':
            dim_0:{
                'data': list/np.ndarray of Number,
                'dims': str,
                'attrs': measure attribute,
                'slice': slice,
            }
            ⋮
            dim_(n_dims-1): ...
        }
    ⋮
    n_nodes-1 : ...
}

Each node id dictionary has the keys with the following meaning:

chunk_indices: The indices assigned to the data chunks in the parallel_coords. There must be an index for each parallel_dims.
parallel_dims: The dimension coordinates over which parallelism will occur.
data_selection: A dictionary where the keys are the names of the datasets in the ProcessingSet, and the values are dictionaries with the coordinates and accompanying slices. If a coordinate is not included, all values will be selected.
task_coords: The chunk of the parallel_coord that is assigned to this node. Each per-dimension entry also carries a 'slice' key giving that chunk’s slice into the full coordinate (slice(None) when the parallel coordinate has no data_chunk_slices).

get_disk_chunk_sizes(input_data: Dict | xarray.DataTree, parallel_coords: Dict) → Dict[str, int][source]

Determine the native on-disk chunk size for each parallel coordinate dimension.

When a Zarr store is opened lazily with xarray (chunks={} or equivalent, as xradio.measurement_set.open_processing_set() does), the resulting dask arrays of the data variables carry the native zarr chunk sizes. This function reads those chunk sizes so that the caller can pass them as disk_chunk_sizes to graphviper.graph_tools.map.map(), enabling the data loading layer to coalesce I/O at disk-chunk granularity.

Data variables — not coordinates — are inspected because xarray always loads dimension coordinates eagerly into memory as numpy-backed indexes, so coordinate arrays carry no on-disk chunk information.

Only the dimensions listed in parallel_coords are inspected; non-parallel dimensions are ignored.

Parameters:

input_data (Union[Dict, xr.DataTree]) – The opened processing set — either an xr.DataTree or a plain dict of xr.Dataset. The datasets must have been opened lazily (chunks={} or equivalent) so that their data variables are backed by dask arrays with chunk information. If no dask-backed data variable carries a given dimension, the full axis length is used as the chunk size, which is equivalent to treating the entire axis as a single disk chunk.
parallel_coords (Dict) – The parallel coordinates dict as produced by make_parallel_coord(). Only the keys (dimension names) are used.

Returns:

{dim: native_chunk_size} for every dimension in parallel_coords that is found in at least one dataset. The returned chunk size is the minimum first-chunk size observed across all data variables and datasets, which is the conservative choice when arrays in the same processing set have different on-disk chunking for the same dimension.

Return type:

Dict[str, int]

Notes

Zarr stores the chunk size uniformly except for the last chunk, which may be smaller if the array length is not a multiple of the chunk size. This function always reads the first chunk size (index 0) to get the nominal chunk size and ignores the trailing remainder chunk.

Examples

>>> ps_xdt = open_processing_set(ps_store)
>>> parallel_coords = {"frequency": make_parallel_coord(coord=img_xds.frequency, n_chunks=n_chunks)}
>>> disk_chunk_sizes = get_disk_chunk_sizes(ps_xdt, parallel_coords)
>>> # e.g. {"frequency": 200}
>>> viper_graph = map(..., disk_chunk_sizes=disk_chunk_sizes)