diff --git a/src/graphviper/graph_tools/coordinate_utils.py b/src/graphviper/graph_tools/coordinate_utils.py index 2a4f39e..72bd407 100644 --- a/src/graphviper/graph_tools/coordinate_utils.py +++ b/src/graphviper/graph_tools/coordinate_utils.py @@ -258,6 +258,204 @@ def make_parallel_coord_by_gap(coord: Union[Dict, xr.DataArray], gap: float) -> return parallel_coord +def _array_split(data: Union[list, np.ndarray], n_chunks: int): + """Takes an input array and splits it into n_chunk arrays which are stored in a dictionary with numbered keys. + + Parameters + ---------- + data : Union[list,np.ndarray] + Array to be split. + n_chunks : int + Number of chunks to split array into. + + Returns + ------- + data_chunks + Dictionary with array broken into chunks. + """ + + chunk_size = int(np.ceil(len(data) / n_chunks)) + from itertools import islice + + data_iter = iter(data) + + data_chunks_list = [] + for _ in range(n_chunks): + chunk = list(islice(data_iter, chunk_size)) + if not chunk: + break + data_chunks_list.append(np.array(chunk)) + + data_chunks = dict(zip(np.arange(n_chunks), data_chunks_list)) + + return data_chunks + + +def _array_split_edges(data_chunks_dict: Dict): + """ + Creates a list of the start and end values of arrays in data_chunks_dict. + + Parameters + ---------- + data_chunks_dict : Dict + Dictionary of integer keys and array values. + + Returns + ------- + data_chunks_edges + List of start and end values from data_chunks_dict. + """ + data_chunks_edges = [] + for key in sorted(data_chunks_dict.keys()): + data_chunks_edges.append(data_chunks_dict[key][0]) + data_chunks_edges.append(data_chunks_dict[key][-1]) + return data_chunks_edges + + +def _array_split_slices(data_chunks_dict: Dict) -> Dict: + """ + Creates a dictionary of slices giving the start and end index of each chunk + relative to the full (unsplit) array. + + Parameters + ---------- + data_chunks_dict : Dict + Dictionary of integer keys and array values, as returned by _array_split. + + Returns + ------- + Dict + Dictionary with the same integer keys as data_chunks_dict, where each + value is a slice(start, stop) indexing that chunk in the original array. + """ + slices = {} + start = 0 + for key in sorted(data_chunks_dict.keys()): + size = len(data_chunks_dict[key]) + slices[key] = slice(start, start + size) + start += size + return slices + + +def _make_iter_chunks_indices(parallel_coords: Dict): + """Creates an iterator of all the combinations of the chunks in the parallel coordinates. + + Parameters + ---------- + parallel_coords : Dict + See make_parallel_coord for explanation. + Returns + ------- + iter_chunks_indxs, parallel_dims + The chunks of each + """ + + parallel_dims = [] + list_chunk_indxs = [] + n_chunks = 1 + for dim, pc in parallel_coords.items(): + chunk_indxs = list(pc["data_chunks"].keys()) + n_chunks = n_chunks * len(chunk_indxs) + list_chunk_indxs.append(chunk_indxs) + parallel_dims.append(dim) + + iter_chunks_indxs = itertools.product(*list_chunk_indxs) + return iter_chunks_indxs, parallel_dims + + +def _partition_ps_by_non_dimensions(ps, ps_partition_keys): + "This requires at least one member of ps_partition_keys!" + ps_split_map = {} + for name, xds in ps.items(): + partition_info = xds.xr_ms.get_partition_info() + for key in ps_partition_keys: + val_for_xds = partition_info[key] + # value could be a list + list_for_xds = None + if isinstance(val_for_xds, list): + list_for_xds = val_for_xds + # OK I think I can punt: the key should probably be an integer but that doesn't feel very Pythonic + # But I *can* reasonably demand it is hashable + if list_for_xds is not None: + for v in list_for_xds: + if not isinstance(v, Hashable): + raise ValueError( + f"Can't split by {key}; value {v} is not suitable for splitting" + ) + ps_split_map.setdefault(key, {}).setdefault(v, []).append(name) + else: + if not isinstance(val_for_xds, Hashable): + raise ValueError( + f"Can't split by {key}; value {val_for_xds} is not suitable for splitting" + ) + ps_split_map.setdefault(key, {}).setdefault(val_for_xds, []).append( + name + ) + + d = {} + # We loop over the cartersian product of the keys + for multi_index in itertools.product( + *[ps_split_map[key] for key in ps_partition_keys] + ): + # And for each key we look up the corresponding set of xds names + sets = [ + set(ps_split_map[key][i]) for i, key in zip(multi_index, ps_partition_keys) + ] + d[multi_index] = set.intersection(*sets) + + return d + + +def _nearest_interp_indices( + coord_values: np.ndarray, + query_values: np.ndarray, +) -> np.ndarray: + """Fast replacement for ``scipy.interpolate.interp1d(..., kind='nearest', + bounds_error=False, fill_value=-1)`` that maps coordinate values to their + nearest integer indices. + + Uses ``np.searchsorted`` (pure C) instead of constructing a scipy + interpolator object, giving ~10-100x less overhead per call when + ``coord_values`` is already sorted (i.e. ``assume_sorted=True``). + + Parameters + ---------- + coord_values : np.ndarray + Sorted 1-D array of coordinate values (the "x" knots). + query_values : np.ndarray + Values to look up. + + Returns + ------- + np.ndarray of int + Index of the nearest element in ``coord_values`` for each query, or + ``-1`` when the query is strictly outside the range of + ``coord_values``. + """ + n = len(coord_values) + query = np.asarray(query_values) + + # Insertion points in the sorted coord array + right = np.searchsorted(coord_values, query, side="left") + left = right - 1 + + # Clamp for safe distance comparison (out-of-range cases handled below) + rc = np.clip(right, 0, n - 1) + lc = np.clip(left, 0, n - 1) + + right_dist = np.abs(coord_values[rc] - query) + left_dist = np.abs(coord_values[lc] - query) + + # Prefer the left neighbour on an exact tie (consistent with floor behaviour) + interp_index = np.where(left_dist <= right_dist, lc, rc).astype(int) + + # Apply fill_value=-1 for strictly out-of-bounds queries + out_of_bounds = (query < coord_values[0]) | (query > coord_values[-1]) + interp_index[out_of_bounds] = -1 + + return interp_index + + def interpolate_data_coords_onto_parallel_coords( parallel_coords: dict, input_data: Union[Dict, xr.DataTree], @@ -273,9 +471,7 @@ def interpolate_data_coords_onto_parallel_coords( "next", } = "nearest", assume_sorted: bool = True, - ps_partition: Optional[ - list[str] - ] = None, # Current options are {'field_name', 'spectral_window_name'} + ps_partition: Optional[list[str]] = None, ) -> 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. @@ -372,9 +568,7 @@ def interpolate_data_coords_onto_parallel_coords( - ``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. """ - # Nested Dict keys: xds_name, dim, chunk_index. - - if ps_partition == None: + if ps_partition is None: ps_partition = [] if ("spectral_window_name" in ps_partition) and ("frequency" in parallel_coords): raise ValueError("Cannot split by both spw and frequency") @@ -382,91 +576,92 @@ def interpolate_data_coords_onto_parallel_coords( if len(ps_partition) > 0: partition_map = _partition_ps_by_non_dimensions(input_data, ps_partition) else: - # By default we iterate over everything partition_map = {0: [xds_name for xds_name in input_data]} xds_data_selection = {} - # Loop over every dataset and interpolate onto parallel_coords; - # The result of this loop is a filled-in xds_data_selection that maps each - # xds to a dictionary of mapping its dimensions to the range for each chunk - for partition in partition_map: - for xds_name in input_data: - for dim, pc in parallel_coords.items(): - - # Interpolate all the chunk edges. This is done for performance reasons. - if "data_chunks_edges" not in pc: - pc["data_chunks_edges"] = _array_split_edges(pc["data_chunks"]) + # ----------------------------------------------------------------------- + # Phase 1 — build xds_data_selection + # + # Loop order swapped to dim → xds_name so that per-dim constants + # (edges array, sorted chunk keys) are computed only once. + # The outer partition loop has been removed: partition is never referenced + # here, so repeating the work n_partitions times was wasteful. + # ----------------------------------------------------------------------- + for dim, pc in parallel_coords.items(): - if input_data[xds_name][dim].dtype.kind in ("U", "S", "O"): - # For string arrays, we perform exact matching only. - # We map each string to its integer index. - string_to_idx = { - val: i for i, val in enumerate(input_data[xds_name][dim].values) - } + # Pre-compute once per dim (not once per (partition, xds_name, dim)) + if "data_chunks_edges" not in pc: + pc["data_chunks_edges"] = _array_split_edges(pc["data_chunks"]) + edges = np.asarray(pc["data_chunks_edges"]) + sorted_chunk_keys = sorted(pc["data_chunks"].keys()) - # Create a vectorized lookup function - # Values not found in the input data return -1 - def string_interpolator(query_values, string_to_idx): - return np.array( - [string_to_idx.get(val, -1) for val in query_values] - ) + for xds_name in input_data: - interp_index = string_interpolator( - pc["data_chunks_edges"], string_to_idx - ).astype(int) - else: - interpolator = interp1d( - input_data[xds_name][dim].values, - np.arange(len(input_data[xds_name][dim].values)), - kind=interpolation_method, - bounds_error=False, - fill_value=-1, - # fill_value="extrapolate", - assume_sorted=assume_sorted, - ) - interp_index = interpolator(pc["data_chunks_edges"]).astype(int) - - chunk_indx_start_stop = {} - i = 0 - # Split the interp_index for each chunk and fix any boundary issues. - for chunk_index in sorted(pc["data_chunks"].keys()): - if interp_index[i] == -1 and interp_index[i + 1] == -1: - chunk_indx_start_stop[chunk_index] = slice(None) - if ( - pc["data_chunks_edges"][i] < input_data[xds_name][dim][0] - ) and ( - pc["data_chunks_edges"][i + 1] - > input_data[xds_name][dim][-1] - ): - interp_index[i] = 0 - interp_index[i + 1] = -2 - chunk_indx_start_stop[chunk_index] = slice( - interp_index[i], interp_index[i + 1] + 1 - ) - else: - if interp_index[i] == -1: - interp_index[i] = 0 - if interp_index[i + 1] == -1: - interp_index[i + 1] = -2 + # Cache coordinate array once — avoids repeated xarray/zarr access + # and eliminates per-chunk xarray element indexing later. + coord_da = input_data[xds_name][dim] + coord_values = coord_da.values + + if coord_values.dtype.kind in ("U", "S", "O"): + # String coordinates: exact matching only + string_to_idx = {val: i for i, val in enumerate(coord_values)} + interp_index = np.array( + [string_to_idx.get(val, -1) for val in edges] + ).astype(int) + elif interpolation_method == "nearest": + # Fast path: no interpolator object construction + interp_index = _nearest_interp_indices(coord_values, edges) + else: + # Fallback for non-nearest methods + interpolator = interp1d( + coord_values, + np.arange(len(coord_values)), + kind=interpolation_method, + bounds_error=False, + fill_value=-1, + assume_sorted=assume_sorted, + ) + interp_index = interpolator(edges).astype(int) + + # Cached boundary values — avoids per-chunk xarray element indexing + coord_min = coord_values[0] + coord_max = coord_values[-1] + + chunk_indx_start_stop = {} + i = 0 + for chunk_index in sorted_chunk_keys: + if interp_index[i] == -1 and interp_index[i + 1] == -1: + chunk_indx_start_stop[chunk_index] = slice(None) + if edges[i] < coord_min and edges[i + 1] > coord_max: + interp_index[i] = 0 + interp_index[i + 1] = -2 chunk_indx_start_stop[chunk_index] = slice( interp_index[i], interp_index[i + 1] + 1 ) - i = i + 2 - - xds_data_selection.setdefault(xds_name, {})[dim] = chunk_indx_start_stop + else: + if interp_index[i] == -1: + interp_index[i] = 0 + if interp_index[i + 1] == -1: + interp_index[i + 1] = -2 + chunk_indx_start_stop[chunk_index] = slice( + interp_index[i], interp_index[i + 1] + 1 + ) + i += 2 - # To create the node_task_data_mapping we have to get the slices for each dimension on node task level. - # Using itertools we can get all the combinations of the chunk indices that belong to a given node: - node_task_data_mapping = ( - {} - ) # Nested Dict keys: task_id, [data_selection,chunk_indices,parallel_dims], xds_name, dim. + xds_data_selection.setdefault(xds_name, {})[dim] = chunk_indx_start_stop - # Loop over every task node (each task node has a unique task_id): + # ----------------------------------------------------------------------- + # Phase 2 — build node_task_data_mapping + # + # Largely unchanged from v1; parallel_dims is extracted once since it is + # the same for every partition. logger.debug uses %s formatting to skip + # string construction when debug logging is inactive. + # ----------------------------------------------------------------------- + node_task_data_mapping = {} task_id = 0 for partition in partition_map.keys(): - # We redo this for every partition, because task number will have changed iter_chunks_indices, parallel_dims = _make_iter_chunks_indices(parallel_coords) for chunk_indices in iter_chunks_indices: logger.debug(f"chunk_index: {task_id}, {chunk_indices}") @@ -476,7 +671,6 @@ def string_interpolator(query_values, string_to_idx): node_task_data_mapping[task_id]["data_selection"] = {} task_coords = {} - # For task_id get the task_coords from parallel_coords: for i_dim, dim in enumerate(parallel_dims): chunk_coords = {} chunk_coords["data"] = parallel_coords[dim]["data_chunks"][ @@ -492,180 +686,291 @@ def string_interpolator(query_values, string_to_idx): chunk_coords["slice"] = slice(None) task_coords[dim] = chunk_coords - # breakpoint() node_task_data_mapping[task_id]["task_coords"] = task_coords partition_xds_names = partition_map[partition] - # For task_id get the selection slices for each dataset in the input data from xds_data_selection: for xds_name in partition_xds_names: node_task_data_mapping[task_id]["data_selection"][xds_name] = {} empty_chunk = False for i, chunk_index in enumerate(chunk_indices): if chunk_index in xds_data_selection[xds_name][parallel_dims[i]]: + sel = xds_data_selection[xds_name][parallel_dims[i]][ + chunk_index + ] node_task_data_mapping[task_id]["data_selection"][xds_name][ parallel_dims[i] - ] = xds_data_selection[xds_name][parallel_dims[i]][chunk_index] - - if xds_data_selection[xds_name][parallel_dims[i]][ - chunk_index - ] == slice(None): + ] = sel + if sel == slice(None): empty_chunk = True else: empty_chunk = True - # - if ( - empty_chunk - ): # The xds with xds_name has no data for the parallel chunk (no slice on one of the dims). + if empty_chunk: del node_task_data_mapping[task_id]["data_selection"][xds_name] task_id += 1 - # breakpoint() return node_task_data_mapping -def _array_split(data: Union[list, np.ndarray], n_chunks: int): - """Takes an input array and splits it into n_chunk arrays which are stored in a dictionary with numbered keys. - - Parameters - ---------- - data : Union[list,np.ndarray] - Array to be split. - n_chunks : int - Number of chunks to split array into. - - Returns - ------- - data_chunks - Dictionary with array broken into chunks. - """ - - chunk_size = int(np.ceil(len(data) / n_chunks)) - from itertools import islice - - data_iter = iter(data) - - data_chunks_list = [] - for _ in range(n_chunks): - chunk = list(islice(data_iter, chunk_size)) - if not chunk: - break - data_chunks_list.append(np.array(chunk)) - - data_chunks = dict(zip(np.arange(n_chunks), data_chunks_list)) - - return data_chunks - - -def _array_split_edges(data_chunks_dict: Dict): - """ - Creates a list of the start and end values of arrays in data_chunks_dict. - - Parameters - ---------- - data_chunks_dict : Dict - Dictionary of integer keys and array values. - - Returns - ------- - data_chunks_edges - List of start and end values from data_chunks_dict. - """ - data_chunks_edges = [] - for key in sorted(data_chunks_dict.keys()): - data_chunks_edges.append(data_chunks_dict[key][0]) - data_chunks_edges.append(data_chunks_dict[key][-1]) - return data_chunks_edges - - -def _array_split_slices(data_chunks_dict: Dict) -> Dict: - """ - Creates a dictionary of slices giving the start and end index of each chunk - relative to the full (unsplit) array. - - Parameters - ---------- - data_chunks_dict : Dict - Dictionary of integer keys and array values, as returned by _array_split. - - Returns - ------- - Dict - Dictionary with the same integer keys as data_chunks_dict, where each - value is a slice(start, stop) indexing that chunk in the original array. - """ - slices = {} - start = 0 - for key in sorted(data_chunks_dict.keys()): - size = len(data_chunks_dict[key]) - slices[key] = slice(start, start + size) - start += size - return slices - - -def _make_iter_chunks_indices(parallel_coords: Dict): - """Creates an iterator of all the combinations of the chunks in the parallel coordinates. - - Parameters - ---------- - parallel_coords : Dict - See make_parallel_coord for explanation. - Returns - ------- - iter_chunks_indxs, parallel_dims - The chunks of each - """ - - parallel_dims = [] - list_chunk_indxs = [] - n_chunks = 1 - for dim, pc in parallel_coords.items(): - chunk_indxs = list(pc["data_chunks"].keys()) - n_chunks = n_chunks * len(chunk_indxs) - list_chunk_indxs.append(chunk_indxs) - parallel_dims.append(dim) - - iter_chunks_indxs = itertools.product(*list_chunk_indxs) - return iter_chunks_indxs, parallel_dims - - -def _partition_ps_by_non_dimensions(ps, ps_partition_keys): - "This requires at least one member of ps_partition_keys!" - ps_split_map = {} - for name, xds in ps.items(): - partition_info = xds.xr_ms.get_partition_info() - for key in ps_partition_keys: - val_for_xds = partition_info[key] - # value could be a list - list_for_xds = None - if isinstance(val_for_xds, list): - list_for_xds = val_for_xds - # OK I think I can punt: the key should probably be an integer but that doesn't feel very Pythonic - # But I *can* reasonably demand it is hashable - if list_for_xds is not None: - for v in list_for_xds: - if not isinstance(v, Hashable): - raise ValueError( - f"Can't split by {key}; value {v} is not suitable for splitting" - ) - ps_split_map.setdefault(key, {}).setdefault(v, []).append(name) - else: - if not isinstance(val_for_xds, Hashable): - raise ValueError( - f"Can't split by {key}; value {val_for_xds} is not suitable for splitting" - ) - ps_split_map.setdefault(key, {}).setdefault(val_for_xds, []).append( - name - ) - - d = {} - # We loop over the cartersian product of the keys - for multi_index in itertools.product( - *[ps_split_map[key] for key in ps_partition_keys] - ): - # And for each key we look up the corresponding set of xds names - sets = [ - set(ps_split_map[key][i]) for i, key in zip(multi_index, ps_partition_keys) - ] - d[multi_index] = set.intersection(*sets) - - return d +# def interpolate_data_coords_onto_parallel_coords( +# parallel_coords: dict, +# input_data: Union[Dict, xr.DataTree], +# interpolation_method: { +# "linear", +# "nearest", +# "nearest-up", +# "zero", +# "slinear", +# "quadratic", +# "cubic", +# "previous", +# "next", +# } = "nearest", +# assume_sorted: bool = True, +# ps_partition: Optional[ +# list[str] +# ] = None, # Current options are {'field_name', 'spectral_window_name'} +# ) -> 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. + +# 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 +# ------- +# Dict : +# Node task data mapping dictionary. See :ref:`notes ` for structure of dictionary. + +# 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. + +# .. _parallel coords: +# 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_chunk_edges': list/np.ndarray of Number, +# '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 :ref:`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. +# - ``dims``: The dimension coordinate name. +# - `attrs``: The `XRADIO measures attributes `_ of the data. + +# .. _node task data mapping: +# 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, +# } +# ⋮ +# 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 :ref:`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. +# """ +# # Nested Dict keys: xds_name, dim, chunk_index. + +# if ps_partition == None: +# ps_partition = [] +# if ("spectral_window_name" in ps_partition) and ("frequency" in parallel_coords): +# raise ValueError("Cannot split by both spw and frequency") + +# if len(ps_partition) > 0: +# partition_map = _partition_ps_by_non_dimensions(input_data, ps_partition) +# else: +# # By default we iterate over everything +# partition_map = {0: [xds_name for xds_name in input_data]} + +# xds_data_selection = {} + +# # Loop over every dataset and interpolate onto parallel_coords; +# # The result of this loop is a filled-in xds_data_selection that maps each +# # xds to a dictionary of mapping its dimensions to the range for each chunk +# for partition in partition_map: +# for xds_name in input_data: +# for dim, pc in parallel_coords.items(): + +# # Interpolate all the chunk edges. This is done for performance reasons. +# if "data_chunks_edges" not in pc: +# pc["data_chunks_edges"] = _array_split_edges(pc["data_chunks"]) + +# if input_data[xds_name][dim].dtype.kind in ("U", "S", "O"): +# # For string arrays, we perform exact matching only. +# # We map each string to its integer index. +# string_to_idx = { +# val: i for i, val in enumerate(input_data[xds_name][dim].values) +# } + +# # Create a vectorized lookup function +# # Values not found in the input data return -1 +# def string_interpolator(query_values, string_to_idx): +# return np.array( +# [string_to_idx.get(val, -1) for val in query_values] +# ) + +# interp_index = string_interpolator( +# pc["data_chunks_edges"], string_to_idx +# ).astype(int) +# else: +# interpolator = interp1d( +# input_data[xds_name][dim].values, +# np.arange(len(input_data[xds_name][dim].values)), +# kind=interpolation_method, +# bounds_error=False, +# fill_value=-1, +# # fill_value="extrapolate", +# assume_sorted=assume_sorted, +# ) +# interp_index = interpolator(pc["data_chunks_edges"]).astype(int) + +# chunk_indx_start_stop = {} +# i = 0 +# # Split the interp_index for each chunk and fix any boundary issues. +# for chunk_index in sorted(pc["data_chunks"].keys()): +# if interp_index[i] == -1 and interp_index[i + 1] == -1: +# chunk_indx_start_stop[chunk_index] = slice(None) +# if ( +# pc["data_chunks_edges"][i] < input_data[xds_name][dim][0] +# ) and ( +# pc["data_chunks_edges"][i + 1] +# > input_data[xds_name][dim][-1] +# ): +# interp_index[i] = 0 +# interp_index[i + 1] = -2 +# chunk_indx_start_stop[chunk_index] = slice( +# interp_index[i], interp_index[i + 1] + 1 +# ) +# else: +# if interp_index[i] == -1: +# interp_index[i] = 0 +# if interp_index[i + 1] == -1: +# interp_index[i + 1] = -2 +# chunk_indx_start_stop[chunk_index] = slice( +# interp_index[i], interp_index[i + 1] + 1 +# ) +# i = i + 2 + +# xds_data_selection.setdefault(xds_name, {})[dim] = chunk_indx_start_stop + +# # To create the node_task_data_mapping we have to get the slices for each dimension on node task level. +# # Using itertools we can get all the combinations of the chunk indices that belong to a given node: +# node_task_data_mapping = ( +# {} +# ) # Nested Dict keys: task_id, [data_selection,chunk_indices,parallel_dims], xds_name, dim. + +# # Loop over every task node (each task node has a unique task_id): + +# task_id = 0 +# for partition in partition_map.keys(): +# # We redo this for every partition, because task number will have changed +# iter_chunks_indices, parallel_dims = _make_iter_chunks_indices(parallel_coords) +# for chunk_indices in iter_chunks_indices: +# logger.debug(f"chunk_index: {task_id}, {chunk_indices}") +# node_task_data_mapping[task_id] = {} +# node_task_data_mapping[task_id]["chunk_indices"] = chunk_indices +# node_task_data_mapping[task_id]["parallel_dims"] = parallel_dims +# node_task_data_mapping[task_id]["data_selection"] = {} + +# task_coords = {} +# # For task_id get the task_coords from parallel_coords: +# for i_dim, dim in enumerate(parallel_dims): +# chunk_coords = {} +# chunk_coords["data"] = parallel_coords[dim]["data_chunks"][ +# chunk_indices[i_dim] +# ] +# chunk_coords["dims"] = parallel_coords[dim]["dims"] +# chunk_coords["attrs"] = parallel_coords[dim]["attrs"] +# if "data_chunk_slices" in parallel_coords[dim]: +# chunk_coords["slice"] = parallel_coords[dim]["data_chunk_slices"][ +# chunk_indices[i_dim] +# ] +# else: +# chunk_coords["slice"] = slice(None) +# task_coords[dim] = chunk_coords + +# # breakpoint() +# node_task_data_mapping[task_id]["task_coords"] = task_coords +# partition_xds_names = partition_map[partition] +# # For task_id get the selection slices for each dataset in the input data from xds_data_selection: +# for xds_name in partition_xds_names: +# node_task_data_mapping[task_id]["data_selection"][xds_name] = {} +# empty_chunk = False +# for i, chunk_index in enumerate(chunk_indices): +# if chunk_index in xds_data_selection[xds_name][parallel_dims[i]]: +# node_task_data_mapping[task_id]["data_selection"][xds_name][ +# parallel_dims[i] +# ] = xds_data_selection[xds_name][parallel_dims[i]][chunk_index] + +# if xds_data_selection[xds_name][parallel_dims[i]][ +# chunk_index +# ] == slice(None): +# empty_chunk = True +# else: +# empty_chunk = True + +# # +# if ( +# empty_chunk +# ): # The xds with xds_name has no data for the parallel chunk (no slice on one of the dims). +# del node_task_data_mapping[task_id]["data_selection"][xds_name] +# task_id += 1 +# # breakpoint() + +# return node_task_data_mapping