From db7e21c4156ffc26f77eb35c55b797e9f0176bef Mon Sep 17 00:00:00 2001 From: John Helly Date: Thu, 20 Nov 2025 11:40:35 +0000 Subject: [PATCH 1/8] Added slicing code --- swiftsimio/slice_utils.py | 303 ++++++++++++++++++++++++++++++++++++++ 1 file changed, 303 insertions(+) create mode 100644 swiftsimio/slice_utils.py diff --git a/swiftsimio/slice_utils.py b/swiftsimio/slice_utils.py new file mode 100644 index 00000000..8c396415 --- /dev/null +++ b/swiftsimio/slice_utils.py @@ -0,0 +1,303 @@ +#!/bin/env python + +import h5py +import numpy as np + + +def validate_slices(starts, counts): + """ + Sanity check the supplied array of slices + + :param starts: 1D array with starting offset of each slice + :type starts: np.ndarray + :param counts: 1D array with length of each slice + :type counts: np.ndarray + """ + if starts.shape != counts.shape: + raise RuntimeError("start and count arrays must be the same shape") + if len(starts.shape) != 1 or len(counts.shape) != 1: + raise RuntimeError("start and count arrays must be 1D") + if len(starts) > 1: + if np.any(starts[1:] < starts[:-1]): + raise RuntimeError("slices must be in ascending order of start index") + ends = starts + counts + if np.any(starts[1:] < ends[:-1]): + raise RuntimeError("slices must not overlap") + if np.any(counts < 0): + raise RuntimeError("slices must have non-negative counts") + if np.any(starts < 0): + # We don't support negative indexes + raise RuntimeError("slices must have non-negative starts") + + +def merge_slices(starts, counts): + """ + Given a set of slices where slice i starts at index starts[i] and contains + counts[i] elements, merge any adjacent slices and return new starts and + counts arrays. + + :param starts: 1D array with starting offset of each slice + :type starts: np.ndarray + :param counts: 1D array with length of each slice + :type counts: np.ndarray + + :return: new (starts, counts) tuple with the merged slices + :rtype: (numpy.ndarray, numpy.ndarray) + """ + + starts = np.asarray(starts, dtype=int) + counts = np.asarray(counts, dtype=int) + + # First, eliminate any zero length slices + keep = counts > 0 + starts = starts[keep] + ends = starts + counts[keep] + + # Determine number of slices + nr_slices = len(starts) + if len(ends) != nr_slices: + raise ValueError("starts and counts arrays must be the same size!") + + # Determine starts to keep: every starting offset which is NOT + # equal to the end of the previous slice. Always keep the first. + keep_start = np.ones(nr_slices, dtype=bool) + keep_start[1:] = (starts[1:] != ends[:-1]) + + # Determine ends to keep: every end offset which is NOT equal + # to the start of the next slice. Always keep the last one. + keep_end = np.ones(nr_slices, dtype=bool) + keep_end[:-1] = (ends[:-1] != starts[1:]) + + # Discard unwanted elements + assert len(starts) == len(ends) + starts = starts[keep_start] + counts = ends[keep_end] - starts + + return starts, counts + + +def read_slices(dataset, starts, counts, result=None): + """ + Read the specified slices from a HDF5 dataset. Uses h5py low level calls + to read the slices with a single H5Dread(). Datasets can only be sliced + along the first dimension: we always read all elements in the remaining + dimensions. + + Slices must be in ascending order of starting index and must not overlap. + Python/numpy style negative indexes from the end of the dataset are not + supported. + + :param dataset: HDF5 dataset to read from + :type dataset: h5py.Dataset + :param starts: 1D array with starting offset of each slice + :type starts: np.ndarray + :param counts: 1D array with length of each slice + :type counts: np.ndarray + :param result: array to hold the result + :type result: np.ndarray, or None + + :return: a numpy array with the data + :rtype: numpy.ndarray + """ + + # Sanity check the slices + starts = np.asarray(starts, dtype=int) + counts = np.asarray(counts, dtype=int) + validate_slices(starts, counts) + + # Merged any adjacent slices + starts, counts = merge_slices(starts, counts) + + # Get dataset handle + dataset_id = dataset.id + + # Get file dataspace handle + file_space_id = dataset_id.get_space() + file_shape = file_space_id.get_simple_extent_dims() + + # Select the slices to read + nr_in_first_dim = 0 + file_space_id.select_none() + for start, count in zip(starts, counts): + if count > 0: + # Select this slice + slice_start = tuple([start,]+[0 for fs in file_shape[1:]]) + slice_count = tuple([count,]+[fs for fs in file_shape[1:]]) + file_space_id.select_hyperslab(slice_start, slice_count, op=h5py.h5s.SELECT_OR) + nr_in_first_dim += count + + # Allocate the output array, if necessary + result_shape = [nr_in_first_dim,]+list(file_shape[1:]) + result_shape = tuple([int(rs) for rs in result_shape]) + if result is None: + result = np.ndarray(result_shape, dtype=dataset.dtype) + + # Output array must be C contiguous + if not result.flags['C_CONTIGUOUS']: + raise RuntimeError("Can only read into C contiguous arrays!") + + # Output array must have the expected number of elements + nr_selected = file_space_id.get_select_npoints() + if nr_selected != result.size: + raise RuntimeError("Output buffer is not the right size for the selected slices!") + + # The output array must have the expected shape (could be wrong if it was passed in) + if result.shape != result_shape: + raise RuntimeError("Output buffer has the wrong shape!") + + # If we selected any elements, read the data + if nr_in_first_dim > 0: + mem_space_id = h5py.h5s.create_simple(result_shape) + dataset_id.read(mem_space_id, file_space_id, result) + + return result + + +class IndexedDatasetReader: + + def __init__(self, index, sorted_and_unique=False): + """ + Class for reading specified indexes from HDF5 datasets. Here we assume + that the requested indexes are likely to include runs of consecutive + values and so can be efficiently handled using hyperslab reads. The + array of indexes is converted into (start, count) pairs and datasets + are read using read_slices(). + + The supplied indexes are in the first dimension. We read all data in any + subsequent dimensions. Indexes must be unique and in ascending order if + sorted_and_unique is True. + + An instance of this class can be used to read the same elements from + multiple datasets. + + :param index: 1D array with indexes to read in the first dimension + :type index: np.ndarray + :param sorted_and_unique: set to True if index values are sorted and unique + :type sorted_and_unique: bool + """ + # Get sorted, unique indexes if necessary + index = np.asarray(index, dtype=int) + if sorted_and_unique: + self.unique_index = index + self.inverse_index = None + else: + self.unique_index, self.inverse_index = np.unique(index, return_inverse=True) + + # Every index is a range of length one. Merge any adjacent ranges. + self.starts, self.counts = merge_slices(self.unique_index, np.ones(len(self.unique_index), dtype=int)) + + def read(self, dataset): + """ + Read the specified indexes from a HDF5 dataset. + + :param dataset: HDF5 dataset to read from + :type dataset: h5py.Dataset + + :return: a numpy array with the data read from the dataset + :rtype: numpy.ndarray + """ + # Read in the specified ranges + result = read_slices(dataset, self.starts, self.counts) + + # And put the result into the order in which the indexes were requested + if self.inverse_index is not None: + result = result[self.inverse_index,...] + return result + + +class SlicedDatasetReader: + + def __init__(self, starts, counts): + """ + Class for reading specified slices from HDF5 datasets. Datasets are + read using read_slices(). The supplied slices are in the first + dimension. We read all data in any subsequent dimensions. + + An instance of this class can be used to read the same elements from + multiple datasets. + + :param starts: 1D array with starting offset of each slice + :type starts: np.ndarray + :param counts: 1D array with length of each slice + :type counts: np.ndarray + """ + # Merge and store any adjacent ranges + self.starts, self.counts = merge_slices(starts, counts) + + def read(self, dataset): + """ + Read the specified indexes from a HDF5 dataset. + + :param dataset: HDF5 dataset to read from + :type dataset: h5py.Dataset + + :return: a numpy array with the data read from the dataset + :rtype: numpy.ndarray + """ + # Read in the specified ranges + return read_slices(dataset, self.starts, self.counts) + + +def match(arr1, arr2, arr2_sorted=False, arr2_index=None): + """ + For each element in arr1 return the index of the element with the + same value in arr2, or -1 if there is no element with the same value. + Setting arr2_sorted=True will save some time if arr2 is already sorted + into ascending order. + + A precomputed sorting index for arr2 can be supplied using the + arr2_index parameter. This can save time if the routine is called + repeatedly with the same arr2 but arr2 is not already sorted. + + It is assumed that each element in arr1 only occurs once in arr2. + """ + + # Check for the case where we're searching an empty arr2 - can't be any matches + if len(arr2) == 0: + return -np.ones(len(arr1), dtype=int) + + # Workaround for a numpy bug (<=1.4): ensure arrays are native endian + # because searchsorted ignores endian flag + if not(arr1.dtype.isnative): + arr1_n = np.asarray(arr1, dtype=arr1.dtype.newbyteorder("=")) + else: + arr1_n = arr1 + if not(arr2.dtype.isnative): + arr2_n = np.asarray(arr2, dtype=arr2.dtype.newbyteorder("=")) + else: + arr2_n = arr2 + + # Sort arr2 into ascending order if necessary + tmp1 = arr1_n + if arr2_sorted: + tmp2 = arr2_n + idx = slice(0,len(arr2_n)) + else: + if arr2_index is None: + idx = np.argsort(arr2_n) + tmp2 = arr2_n[idx] + else: + # Use supplied sorting index + idx = arr2_index + tmp2 = arr2_n[arr2_index] + + # Find where elements of arr1 are in arr2 + ptr = np.searchsorted(tmp2, tmp1) + + # Make sure all elements in ptr are valid indexes into tmp2 + # (any out of range entries won't match so they'll get set to -1 + # in the next bit) + ptr[ptr>=len(tmp2)] = 0 + ptr[ptr<0] = 0 + + # Return -1 where no match is found + ind = tmp2[ptr] != tmp1 + ptr[ind] = -1 + + # Put ptr back into original order + ind = np.arange(len(arr2_n))[idx] + ptr = np.where(ptr>= 0, ind[ptr], -1) + + return ptr + + From 62009dff1b40275aa65361fa77c882a12938c13f Mon Sep 17 00:00:00 2001 From: John Helly Date: Thu, 20 Nov 2025 12:14:56 +0000 Subject: [PATCH 2/8] Call low level slicing code in accelerated.py --- swiftsimio/accelerated.py | 72 ++++++++++++++++++++++++++++++++------- 1 file changed, 59 insertions(+), 13 deletions(-) diff --git a/swiftsimio/accelerated.py b/swiftsimio/accelerated.py index 244d8175..e731b977 100644 --- a/swiftsimio/accelerated.py +++ b/swiftsimio/accelerated.py @@ -10,6 +10,7 @@ from .optional_packages import jit, prange, NUM_THREADS from ._ordered_slices import OrderedSlices +from .slice_utils import read_slices __all__ = [ "jit", @@ -170,6 +171,63 @@ def read_ranges_from_file_unchunked( return output +def read_ranges_from_file_low_level( + handle: Dataset, + ranges: np.ndarray, + output_shape: tuple, + output_type: type = np.float64, + columns: slice = np.s_[:], +) -> np.array: + """ + Read only a selection of index ranges from a dataset that is not chunked. + + Takes a hdf5 dataset, and the set of ranges from + ranges_from_array, and reads only those ranges from the file. + + This version uses the h5py low level API. + + Parameters + ---------- + handle : Dataset + HDF5 dataset to slice data from. + + ranges : np.ndarray + Array of ranges (see :func:`ranges_from_array`). + + output_shape : tuple + Resultant shape of output. + + output_type : type, optional + ``numpy`` type of output elements. If not supplied, we assume ``np.float64``. + + columns : slice, optional + Selector for columns if using a multi-dimensional array. If the array is only + a single dimension this is not used. + + Returns + ------- + np.ndarray + Result from reading only the relevant values from ``handle``. + """ + starts = ranges[:,0] + counts = ranges[:,1] - ranges[:,0] + output = read_slices(handle, starts, counts).astype(output_type) + if len(output.shape) > 1: + output = output[:, columns, ...] + + if not output.dtype.isnative: + # The data type we have read in is the opposite endian-ness to the + # machine we're on. Convert it here, to save pain down the line. + output = output.byteswap(inplace=True).newbyteorder() + + if not output.dtype.isnative: + raise RuntimeError( + "Unable to find a native type that is a match to read data." + ) + + return output + + def index_dataset(handle: Dataset, mask_array: np.array) -> np.array: """ Index the dataset using the mask array. @@ -522,20 +580,8 @@ def read_ranges_from_file( read_ranges_from_file_unchunked Reads data ranges for unchunked hdf5 file. """ - # It was found that the range size for which read_ranges_from_file_chunked was - # faster than unchunked was approximately 5e5. For ranges larger than this the - # overheads associated with read_ranges_from_file_chunked caused slightly worse - # performance than read_ranges_from_file_unchunked - cross_over_range_size = 5e5 - - average_range_size = np.diff(ranges).mean() - read_ranges = ( - read_ranges_from_file_chunked - if handle.chunks is not None and average_range_size < cross_over_range_size - else read_ranges_from_file_unchunked - ) return ( - read_ranges_from_hdfstream if hasattr(handle, "request_slices") else read_ranges + read_ranges_from_hdfstream if hasattr(handle, "request_slices") else read_ranges_from_file_low_level )(handle, ranges, output_shape, output_type, columns) From 5ca0c18393e52996581f49c3e63723e03476ac7b Mon Sep 17 00:00:00 2001 From: John Helly Date: Thu, 20 Nov 2025 12:47:09 +0000 Subject: [PATCH 3/8] Move slicing code into accelerated.py --- swiftsimio/accelerated.py | 56 +++++-- swiftsimio/slice_utils.py | 303 -------------------------------------- 2 files changed, 45 insertions(+), 314 deletions(-) delete mode 100644 swiftsimio/slice_utils.py diff --git a/swiftsimio/accelerated.py b/swiftsimio/accelerated.py index e731b977..8173e55a 100644 --- a/swiftsimio/accelerated.py +++ b/swiftsimio/accelerated.py @@ -10,7 +10,6 @@ from .optional_packages import jit, prange, NUM_THREADS from ._ordered_slices import OrderedSlices -from .slice_utils import read_slices __all__ = [ "jit", @@ -209,23 +208,58 @@ def read_ranges_from_file_low_level( np.ndarray Result from reading only the relevant values from ``handle``. """ - starts = ranges[:,0] - counts = ranges[:,1] - ranges[:,0] - output = read_slices(handle, starts, counts).astype(output_type) - if len(output.shape) > 1: - output = output[:, columns, ...] - - if not output.dtype.isnative: + # Get dataset handle + dataset_id = dataset.id + + # Get file dataspace handle + file_space_id = dataset_id.get_space() + file_shape = file_space_id.get_simple_extent_dims() + + # Select the slices to read + nr_in_first_dim = 0 + file_space_id.select_none() + for start, count in ranges: + if count > 0: + # Select this slice + slice_start = tuple([start,]+[0 for fs in file_shape[1:]]) + slice_count = tuple([count,]+[fs for fs in file_shape[1:]]) + file_space_id.select_hyperslab(slice_start, slice_count, op=h5py.h5s.SELECT_OR) + nr_in_first_dim += count + + # Allocate the output array + result_shape = [nr_in_first_dim,]+list(file_shape[1:]) + result_shape = tuple([int(rs) for rs in result_shape]) + result = np.ndarray(result_shape, dtype=output_type) + + # Output array must have the expected number of elements + nr_selected = file_space_id.get_select_npoints() + if nr_selected != result.size: + raise RuntimeError("Output buffer is not the right size for the selected slices!") + + # If we selected any elements, read the data + if nr_in_first_dim > 0: + mem_space_id = h5py.h5s.create_simple(result_shape) + dataset_id.read(mem_space_id, file_space_id, result) + mem_space_id.close() + + # Tidy up + file_space_id.close() + + # Select columns if necessary + if len(result.shape) > 1: + result = result[:, columns, ...] + + if not result.dtype.isnative: # The data type we have read in is the opposite endian-ness to the # machine we're on. Convert it here, to save pain down the line. - output = output.byteswap(inplace=True).newbyteorder() + result = result.byteswap(inplace=True).newbyteorder() - if not output.dtype.isnative: + if not result.dtype.isnative: raise RuntimeError( "Unable to find a native type that is a match to read data." ) - return output + return result def index_dataset(handle: Dataset, mask_array: np.array) -> np.array: diff --git a/swiftsimio/slice_utils.py b/swiftsimio/slice_utils.py deleted file mode 100644 index 8c396415..00000000 --- a/swiftsimio/slice_utils.py +++ /dev/null @@ -1,303 +0,0 @@ -#!/bin/env python - -import h5py -import numpy as np - - -def validate_slices(starts, counts): - """ - Sanity check the supplied array of slices - - :param starts: 1D array with starting offset of each slice - :type starts: np.ndarray - :param counts: 1D array with length of each slice - :type counts: np.ndarray - """ - if starts.shape != counts.shape: - raise RuntimeError("start and count arrays must be the same shape") - if len(starts.shape) != 1 or len(counts.shape) != 1: - raise RuntimeError("start and count arrays must be 1D") - if len(starts) > 1: - if np.any(starts[1:] < starts[:-1]): - raise RuntimeError("slices must be in ascending order of start index") - ends = starts + counts - if np.any(starts[1:] < ends[:-1]): - raise RuntimeError("slices must not overlap") - if np.any(counts < 0): - raise RuntimeError("slices must have non-negative counts") - if np.any(starts < 0): - # We don't support negative indexes - raise RuntimeError("slices must have non-negative starts") - - -def merge_slices(starts, counts): - """ - Given a set of slices where slice i starts at index starts[i] and contains - counts[i] elements, merge any adjacent slices and return new starts and - counts arrays. - - :param starts: 1D array with starting offset of each slice - :type starts: np.ndarray - :param counts: 1D array with length of each slice - :type counts: np.ndarray - - :return: new (starts, counts) tuple with the merged slices - :rtype: (numpy.ndarray, numpy.ndarray) - """ - - starts = np.asarray(starts, dtype=int) - counts = np.asarray(counts, dtype=int) - - # First, eliminate any zero length slices - keep = counts > 0 - starts = starts[keep] - ends = starts + counts[keep] - - # Determine number of slices - nr_slices = len(starts) - if len(ends) != nr_slices: - raise ValueError("starts and counts arrays must be the same size!") - - # Determine starts to keep: every starting offset which is NOT - # equal to the end of the previous slice. Always keep the first. - keep_start = np.ones(nr_slices, dtype=bool) - keep_start[1:] = (starts[1:] != ends[:-1]) - - # Determine ends to keep: every end offset which is NOT equal - # to the start of the next slice. Always keep the last one. - keep_end = np.ones(nr_slices, dtype=bool) - keep_end[:-1] = (ends[:-1] != starts[1:]) - - # Discard unwanted elements - assert len(starts) == len(ends) - starts = starts[keep_start] - counts = ends[keep_end] - starts - - return starts, counts - - -def read_slices(dataset, starts, counts, result=None): - """ - Read the specified slices from a HDF5 dataset. Uses h5py low level calls - to read the slices with a single H5Dread(). Datasets can only be sliced - along the first dimension: we always read all elements in the remaining - dimensions. - - Slices must be in ascending order of starting index and must not overlap. - Python/numpy style negative indexes from the end of the dataset are not - supported. - - :param dataset: HDF5 dataset to read from - :type dataset: h5py.Dataset - :param starts: 1D array with starting offset of each slice - :type starts: np.ndarray - :param counts: 1D array with length of each slice - :type counts: np.ndarray - :param result: array to hold the result - :type result: np.ndarray, or None - - :return: a numpy array with the data - :rtype: numpy.ndarray - """ - - # Sanity check the slices - starts = np.asarray(starts, dtype=int) - counts = np.asarray(counts, dtype=int) - validate_slices(starts, counts) - - # Merged any adjacent slices - starts, counts = merge_slices(starts, counts) - - # Get dataset handle - dataset_id = dataset.id - - # Get file dataspace handle - file_space_id = dataset_id.get_space() - file_shape = file_space_id.get_simple_extent_dims() - - # Select the slices to read - nr_in_first_dim = 0 - file_space_id.select_none() - for start, count in zip(starts, counts): - if count > 0: - # Select this slice - slice_start = tuple([start,]+[0 for fs in file_shape[1:]]) - slice_count = tuple([count,]+[fs for fs in file_shape[1:]]) - file_space_id.select_hyperslab(slice_start, slice_count, op=h5py.h5s.SELECT_OR) - nr_in_first_dim += count - - # Allocate the output array, if necessary - result_shape = [nr_in_first_dim,]+list(file_shape[1:]) - result_shape = tuple([int(rs) for rs in result_shape]) - if result is None: - result = np.ndarray(result_shape, dtype=dataset.dtype) - - # Output array must be C contiguous - if not result.flags['C_CONTIGUOUS']: - raise RuntimeError("Can only read into C contiguous arrays!") - - # Output array must have the expected number of elements - nr_selected = file_space_id.get_select_npoints() - if nr_selected != result.size: - raise RuntimeError("Output buffer is not the right size for the selected slices!") - - # The output array must have the expected shape (could be wrong if it was passed in) - if result.shape != result_shape: - raise RuntimeError("Output buffer has the wrong shape!") - - # If we selected any elements, read the data - if nr_in_first_dim > 0: - mem_space_id = h5py.h5s.create_simple(result_shape) - dataset_id.read(mem_space_id, file_space_id, result) - - return result - - -class IndexedDatasetReader: - - def __init__(self, index, sorted_and_unique=False): - """ - Class for reading specified indexes from HDF5 datasets. Here we assume - that the requested indexes are likely to include runs of consecutive - values and so can be efficiently handled using hyperslab reads. The - array of indexes is converted into (start, count) pairs and datasets - are read using read_slices(). - - The supplied indexes are in the first dimension. We read all data in any - subsequent dimensions. Indexes must be unique and in ascending order if - sorted_and_unique is True. - - An instance of this class can be used to read the same elements from - multiple datasets. - - :param index: 1D array with indexes to read in the first dimension - :type index: np.ndarray - :param sorted_and_unique: set to True if index values are sorted and unique - :type sorted_and_unique: bool - """ - # Get sorted, unique indexes if necessary - index = np.asarray(index, dtype=int) - if sorted_and_unique: - self.unique_index = index - self.inverse_index = None - else: - self.unique_index, self.inverse_index = np.unique(index, return_inverse=True) - - # Every index is a range of length one. Merge any adjacent ranges. - self.starts, self.counts = merge_slices(self.unique_index, np.ones(len(self.unique_index), dtype=int)) - - def read(self, dataset): - """ - Read the specified indexes from a HDF5 dataset. - - :param dataset: HDF5 dataset to read from - :type dataset: h5py.Dataset - - :return: a numpy array with the data read from the dataset - :rtype: numpy.ndarray - """ - # Read in the specified ranges - result = read_slices(dataset, self.starts, self.counts) - - # And put the result into the order in which the indexes were requested - if self.inverse_index is not None: - result = result[self.inverse_index,...] - return result - - -class SlicedDatasetReader: - - def __init__(self, starts, counts): - """ - Class for reading specified slices from HDF5 datasets. Datasets are - read using read_slices(). The supplied slices are in the first - dimension. We read all data in any subsequent dimensions. - - An instance of this class can be used to read the same elements from - multiple datasets. - - :param starts: 1D array with starting offset of each slice - :type starts: np.ndarray - :param counts: 1D array with length of each slice - :type counts: np.ndarray - """ - # Merge and store any adjacent ranges - self.starts, self.counts = merge_slices(starts, counts) - - def read(self, dataset): - """ - Read the specified indexes from a HDF5 dataset. - - :param dataset: HDF5 dataset to read from - :type dataset: h5py.Dataset - - :return: a numpy array with the data read from the dataset - :rtype: numpy.ndarray - """ - # Read in the specified ranges - return read_slices(dataset, self.starts, self.counts) - - -def match(arr1, arr2, arr2_sorted=False, arr2_index=None): - """ - For each element in arr1 return the index of the element with the - same value in arr2, or -1 if there is no element with the same value. - Setting arr2_sorted=True will save some time if arr2 is already sorted - into ascending order. - - A precomputed sorting index for arr2 can be supplied using the - arr2_index parameter. This can save time if the routine is called - repeatedly with the same arr2 but arr2 is not already sorted. - - It is assumed that each element in arr1 only occurs once in arr2. - """ - - # Check for the case where we're searching an empty arr2 - can't be any matches - if len(arr2) == 0: - return -np.ones(len(arr1), dtype=int) - - # Workaround for a numpy bug (<=1.4): ensure arrays are native endian - # because searchsorted ignores endian flag - if not(arr1.dtype.isnative): - arr1_n = np.asarray(arr1, dtype=arr1.dtype.newbyteorder("=")) - else: - arr1_n = arr1 - if not(arr2.dtype.isnative): - arr2_n = np.asarray(arr2, dtype=arr2.dtype.newbyteorder("=")) - else: - arr2_n = arr2 - - # Sort arr2 into ascending order if necessary - tmp1 = arr1_n - if arr2_sorted: - tmp2 = arr2_n - idx = slice(0,len(arr2_n)) - else: - if arr2_index is None: - idx = np.argsort(arr2_n) - tmp2 = arr2_n[idx] - else: - # Use supplied sorting index - idx = arr2_index - tmp2 = arr2_n[arr2_index] - - # Find where elements of arr1 are in arr2 - ptr = np.searchsorted(tmp2, tmp1) - - # Make sure all elements in ptr are valid indexes into tmp2 - # (any out of range entries won't match so they'll get set to -1 - # in the next bit) - ptr[ptr>=len(tmp2)] = 0 - ptr[ptr<0] = 0 - - # Return -1 where no match is found - ind = tmp2[ptr] != tmp1 - ptr[ind] = -1 - - # Put ptr back into original order - ind = np.arange(len(arr2_n))[idx] - ptr = np.where(ptr>= 0, ind[ptr], -1) - - return ptr - - From eeb7e61673e2277c3880e15bed84bb2e062e079b Mon Sep 17 00:00:00 2001 From: John Helly Date: Thu, 20 Nov 2025 13:46:18 +0000 Subject: [PATCH 4/8] Fixes to low level slicing code --- swiftsimio/accelerated.py | 40 +++++++++++++++++++++++++++------------ 1 file changed, 28 insertions(+), 12 deletions(-) diff --git a/swiftsimio/accelerated.py b/swiftsimio/accelerated.py index 8173e55a..6bfe2f83 100644 --- a/swiftsimio/accelerated.py +++ b/swiftsimio/accelerated.py @@ -5,7 +5,7 @@ """ import numpy as np - +import h5py from h5py._hl.dataset import Dataset from .optional_packages import jit, prange, NUM_THREADS @@ -209,26 +209,45 @@ def read_ranges_from_file_low_level( Result from reading only the relevant values from ``handle``. """ # Get dataset handle - dataset_id = dataset.id + dataset_id = handle.id # Get file dataspace handle file_space_id = dataset_id.get_space() file_shape = file_space_id.get_simple_extent_dims() + # Determine range of elements to read in the second dimension (if any) + if len(handle.shape) == 1: + column_start = () + column_count = () + elif len(handle.shape) == 2: + if isinstance(columns, slice): + start, stop, step = columns.indices(handle.shape[1]) + if step != 1: + raise RuntimeError("Can only handle column slices with step=1") + column_start = (start,) + column_count = (stop-start,) + elif isinstance(columns, int): + column_start = (columns,) + column_count = (1,) + else: + raise RuntimeError("columns parameter must be slice or integer") + else: + raise RuntimeError("Can only handle 1 or 2 dimensional datasets") + # Select the slices to read nr_in_first_dim = 0 file_space_id.select_none() - for start, count in ranges: + for start, stop in ranges: + count = stop - start if count > 0: # Select this slice - slice_start = tuple([start,]+[0 for fs in file_shape[1:]]) - slice_count = tuple([count,]+[fs for fs in file_shape[1:]]) + slice_start = (start,)+column_start + slice_count = (count,)+column_count file_space_id.select_hyperslab(slice_start, slice_count, op=h5py.h5s.SELECT_OR) nr_in_first_dim += count # Allocate the output array - result_shape = [nr_in_first_dim,]+list(file_shape[1:]) - result_shape = tuple([int(rs) for rs in result_shape]) + result_shape = (nr_in_first_dim,)+column_count result = np.ndarray(result_shape, dtype=output_type) # Output array must have the expected number of elements @@ -241,13 +260,10 @@ def read_ranges_from_file_low_level( mem_space_id = h5py.h5s.create_simple(result_shape) dataset_id.read(mem_space_id, file_space_id, result) mem_space_id.close() - - # Tidy up file_space_id.close() - # Select columns if necessary - if len(result.shape) > 1: - result = result[:, columns, ...] + # Reshape: if columns was an integer we need to remove a dimension + result = result.reshape(output_shape) if not result.dtype.isnative: # The data type we have read in is the opposite endian-ness to the From 5756450206e2626830ed88d3a21bc17d79bbeb35 Mon Sep 17 00:00:00 2001 From: John Helly Date: Thu, 20 Nov 2025 14:41:43 +0000 Subject: [PATCH 5/8] Read un-sorted ranges using the low level h5py API --- swiftsimio/accelerated.py | 28 ++++++++++++++++++++++++++++ 1 file changed, 28 insertions(+) diff --git a/swiftsimio/accelerated.py b/swiftsimio/accelerated.py index 6bfe2f83..caf57f8e 100644 --- a/swiftsimio/accelerated.py +++ b/swiftsimio/accelerated.py @@ -208,6 +208,14 @@ def read_ranges_from_file_low_level( np.ndarray Result from reading only the relevant values from ``handle``. """ + + # This will only work if slices do not overlap + order = np.argsort(ranges[:,0]) + sorted_starts = ranges[order,0] + sorted_stops = ranges[order,1] + if np.any(sorted_stops[:-1] > sorted_starts[1:]): + raise RuntimeError("slices must not overlap") + # Get dataset handle dataset_id = handle.id @@ -265,6 +273,26 @@ def read_ranges_from_file_low_level( # Reshape: if columns was an integer we need to remove a dimension result = result.reshape(output_shape) + # If the slices were not sorted by start index, we'll need to reorder the data + if np.any(ranges[1:,0] <= ranges[:-1,0]): + # Compute the offset into the result array for each slice. + # HDF5 reads the slices in order of starting index. + ranges_read = np.empty_like(ranges) + offset = 0 + for i in np.argsort(ranges[:,0]): + n = ranges[i,1] - ranges[i,0] + ranges_read[i,0] = offset + ranges_read[i,1] = offset + n + offset += n + # Copy the slices to a new array in the input slice order + result_sorted = np.empty_like(result) + offset = 0 + for start, stop in ranges_read: + n = stop - start + result_sorted[offset:offset+n,...] = result[start:stop,...] + offset += n + result = result_sorted + if not result.dtype.isnative: # The data type we have read in is the opposite endian-ness to the # machine we're on. Convert it here, to save pain down the line. From 1c0b99cbc464e8f476d47542cb06a0de4fd9627f Mon Sep 17 00:00:00 2001 From: John Helly Date: Fri, 21 Nov 2025 10:00:44 +0000 Subject: [PATCH 6/8] Correct docstring in read_ranges_from_file_low_level() --- swiftsimio/accelerated.py | 11 ++++++----- 1 file changed, 6 insertions(+), 5 deletions(-) diff --git a/swiftsimio/accelerated.py b/swiftsimio/accelerated.py index caf57f8e..cf2e4902 100644 --- a/swiftsimio/accelerated.py +++ b/swiftsimio/accelerated.py @@ -178,12 +178,12 @@ def read_ranges_from_file_low_level( columns: slice = np.s_[:], ) -> np.array: """ - Read only a selection of index ranges from a dataset that is not chunked. + Read only a selection of index ranges from a dataset. - Takes a hdf5 dataset, and the set of ranges from - ranges_from_array, and reads only those ranges from the file. - - This version uses the h5py low level API. + Takes a hdf5 dataset and the set of ranges from ranges_from_array, + selects all ranges with select_hyperslab(), then reads all of the + data with a single read() call. We do not distinguish between + chunked and un-chunked cases here. Parameters ---------- @@ -207,6 +207,7 @@ def read_ranges_from_file_low_level( ------- np.ndarray Result from reading only the relevant values from ``handle``. + """ # This will only work if slices do not overlap From 7bbec968ac31e40fc0be4250cf8da50e4c38d3b1 Mon Sep 17 00:00:00 2001 From: John Helly Date: Wed, 8 Apr 2026 17:03:06 +0100 Subject: [PATCH 7/8] Run ruff formatter on accelerated.py --- swiftsimio/accelerated.py | 39 ++++++++++++++++++++++----------------- 1 file changed, 22 insertions(+), 17 deletions(-) diff --git a/swiftsimio/accelerated.py b/swiftsimio/accelerated.py index cf2e4902..262784fd 100644 --- a/swiftsimio/accelerated.py +++ b/swiftsimio/accelerated.py @@ -207,13 +207,12 @@ def read_ranges_from_file_low_level( ------- np.ndarray Result from reading only the relevant values from ``handle``. - """ # This will only work if slices do not overlap - order = np.argsort(ranges[:,0]) - sorted_starts = ranges[order,0] - sorted_stops = ranges[order,1] + order = np.argsort(ranges[:, 0]) + sorted_starts = ranges[order, 0] + sorted_stops = ranges[order, 1] if np.any(sorted_stops[:-1] > sorted_starts[1:]): raise RuntimeError("slices must not overlap") @@ -234,7 +233,7 @@ def read_ranges_from_file_low_level( if step != 1: raise RuntimeError("Can only handle column slices with step=1") column_start = (start,) - column_count = (stop-start,) + column_count = (stop - start,) elif isinstance(columns, int): column_start = (columns,) column_count = (1,) @@ -250,19 +249,23 @@ def read_ranges_from_file_low_level( count = stop - start if count > 0: # Select this slice - slice_start = (start,)+column_start - slice_count = (count,)+column_count - file_space_id.select_hyperslab(slice_start, slice_count, op=h5py.h5s.SELECT_OR) + slice_start = (start,) + column_start + slice_count = (count,) + column_count + file_space_id.select_hyperslab( + slice_start, slice_count, op=h5py.h5s.SELECT_OR + ) nr_in_first_dim += count # Allocate the output array - result_shape = (nr_in_first_dim,)+column_count + result_shape = (nr_in_first_dim,) + column_count result = np.ndarray(result_shape, dtype=output_type) # Output array must have the expected number of elements nr_selected = file_space_id.get_select_npoints() if nr_selected != result.size: - raise RuntimeError("Output buffer is not the right size for the selected slices!") + raise RuntimeError( + "Output buffer is not the right size for the selected slices!" + ) # If we selected any elements, read the data if nr_in_first_dim > 0: @@ -275,22 +278,22 @@ def read_ranges_from_file_low_level( result = result.reshape(output_shape) # If the slices were not sorted by start index, we'll need to reorder the data - if np.any(ranges[1:,0] <= ranges[:-1,0]): + if np.any(ranges[1:, 0] <= ranges[:-1, 0]): # Compute the offset into the result array for each slice. # HDF5 reads the slices in order of starting index. ranges_read = np.empty_like(ranges) offset = 0 - for i in np.argsort(ranges[:,0]): - n = ranges[i,1] - ranges[i,0] - ranges_read[i,0] = offset - ranges_read[i,1] = offset + n + for i in np.argsort(ranges[:, 0]): + n = ranges[i, 1] - ranges[i, 0] + ranges_read[i, 0] = offset + ranges_read[i, 1] = offset + n offset += n # Copy the slices to a new array in the input slice order result_sorted = np.empty_like(result) offset = 0 for start, stop in ranges_read: n = stop - start - result_sorted[offset:offset+n,...] = result[start:stop,...] + result_sorted[offset : offset + n, ...] = result[start:stop, ...] offset += n result = result_sorted @@ -660,7 +663,9 @@ def read_ranges_from_file( Reads data ranges for unchunked hdf5 file. """ return ( - read_ranges_from_hdfstream if hasattr(handle, "request_slices") else read_ranges_from_file_low_level + read_ranges_from_hdfstream + if hasattr(handle, "request_slices") + else read_ranges_from_file_low_level )(handle, ranges, output_shape, output_type, columns) From f44aa02a391e03dfad342f290ccf43871c5d0b1d Mon Sep 17 00:00:00 2001 From: John Helly Date: Wed, 8 Apr 2026 17:06:38 +0100 Subject: [PATCH 8/8] Remove unused variable --- swiftsimio/accelerated.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/swiftsimio/accelerated.py b/swiftsimio/accelerated.py index 262784fd..317e50db 100644 --- a/swiftsimio/accelerated.py +++ b/swiftsimio/accelerated.py @@ -208,7 +208,6 @@ def read_ranges_from_file_low_level( np.ndarray Result from reading only the relevant values from ``handle``. """ - # This will only work if slices do not overlap order = np.argsort(ranges[:, 0]) sorted_starts = ranges[order, 0] @@ -221,7 +220,6 @@ def read_ranges_from_file_low_level( # Get file dataspace handle file_space_id = dataset_id.get_space() - file_shape = file_space_id.get_simple_extent_dims() # Determine range of elements to read in the second dimension (if any) if len(handle.shape) == 1: