From 5d4d82b7bcc870bd5273058160cf17632111e4f2 Mon Sep 17 00:00:00 2001 From: vitous Date: Thu, 23 Jul 2026 16:45:02 +0200 Subject: [PATCH] random access fixes --- brukerapi/data.py | 6 ++--- brukerapi/schemas.py | 48 +++++++++++++++++++++----------------- test/conftest.py | 21 ++++++++++++++++- test/test_random_access.py | 41 +++++++------------------------- 4 files changed, 57 insertions(+), 59 deletions(-) diff --git a/brukerapi/data.py b/brukerapi/data.py index b10fb4b..c06b736 100644 --- a/brukerapi/data.py +++ b/brukerapi/data.py @@ -1,7 +1,7 @@ class DataRandomAccess: def __init__(self, dataset): self._dataset = dataset - self._scheme = dataset._scheme + self._schema = dataset._schema - def __getitem__(self, slice): - return self._scheme.ra(slice) + def __getitem__(self, key): + return self._schema.ra(key) diff --git a/brukerapi/schemas.py b/brukerapi/schemas.py index 4975a81..871cca9 100644 --- a/brukerapi/schemas.py +++ b/brukerapi/schemas.py @@ -260,8 +260,8 @@ def ra(self, slice_): """ random access """ - array_ra = np.zeros(layouts_ra["storage"], dtype=self.numpy_dtype) - fp = np.memmap(self._dataset.path, dtype=self.numpy_dtype, mode="r", shape=layouts["storage"], order="F") + array_ra = np.zeros(layouts_ra["storage"], dtype=self._dataset.numpy_dtype) + fp = np.memmap(self._dataset.path, dtype=self._dataset.numpy_dtype, mode="r", shape=layouts["storage"], order="F") for index_ra in np.ndindex(layouts_ra["k_space"][1:]): # index of line in the original k_space @@ -291,7 +291,7 @@ def ra(self, slice_): layouts_ra["k_space"] = (layouts_ra["k_space"][0] // 2,) + layouts_ra["k_space"][1:] layouts_ra["encoding_space"] = (layouts_ra["encoding_space"][0] // 2,) + layouts_ra["encoding_space"][1:] - array_ra = self.reshape_fw(array_ra, layouts_ra) + array_ra = self.deserialize(array_ra, layouts_ra) singletons = tuple(i for i, v in enumerate(slice_) if isinstance(v, int)) @@ -299,12 +299,12 @@ def ra(self, slice_): def get_ra_layouts(self, slice_): layouts = deepcopy(self.layouts) - layouts["k_space"] = (layouts["k_space"][0] * 2,) + layouts["k_space"][1:] - layouts["encoding_space"] = (layouts["encoding_space"][0] * 2,) + layouts["encoding_space"][1:] + layouts["k_space"] = (layouts["k_space"][0] * 2,) + tuple(layouts["k_space"][1:]) + layouts["encoding_space"] = (layouts["encoding_space"][0] * 2,) + tuple(layouts["encoding_space"][1:]) layouts["inverse_permute"] = tuple(self.permutation_inverse(layouts["permute"])) layouts["encoding_permute"] = tuple(layouts["encoding_space"][i] for i in layouts["permute"]) - layouts["channel_index"] = self.dim_type.index("channel") - layouts["channels"] = layouts["k_space"][layouts["channel_index"]] + layouts["channel_index"] = self._dataset.dim_type.index("channel") if "channel" in self._dataset.dim_type else None + layouts["channels"] = self._dataset.channels if layouts["channel_index"] is None else layouts["k_space"][layouts["channel_index"]] layouts["acquisition_position_ch"] = (layouts["acquisition_position"][0] // layouts["channels"], layouts["acquisition_position"][1] // layouts["channels"]) layouts["storage_clear"] = (layouts["acquisition_position"][1], layouts["storage"][1]) layouts["storage_clear_ch"] = (layouts["storage_clear"][0] // layouts["channels"], layouts["channels"], layouts["storage"][1]) @@ -313,8 +313,8 @@ def get_ra_layouts(self, slice_): layouts_ra = deepcopy(layouts) layouts_ra["k_space"], layouts_ra["k_space_offset"] = self._get_ra_k_space_info(layouts, slice_) - layouts_ra["channels"] = layouts_ra["k_space"][layouts_ra["channel_index"]] - layouts_ra["acquisition_position"] = (0, self.get_acquisition_length(channels=layouts_ra["channels"])) # delete offset + layouts_ra["channels"] = layouts["channels"] if layouts_ra["channel_index"] is None else layouts_ra["k_space"][layouts_ra["channel_index"]] + layouts_ra["acquisition_position"] = (0, self._get_acquisition_length(layouts, layouts_ra["channels"])) # delete offset # delete offset layouts_ra["encoding_space"], layouts_ra["storage"] = self._get_e_ra(layouts, layouts_ra) @@ -389,10 +389,14 @@ def _get_e_ra(self, layout_full, layout_ra): encoding_space_ra = max_enc_index - min_enc_index + 1 encoding_space_ra = (layout_full["encoding_space"][0],) + tuple(encoding_space_ra) - storage_ra = (self.get_acquisition_length(channels=layout_ra["channels"]), len(storage_ra)) + storage_ra = (self._get_acquisition_length(layout_full, layout_ra["channels"]), len(storage_ra)) return encoding_space_ra, storage_ra + @staticmethod + def _get_acquisition_length(layouts, channels): + return layouts["acquisition_position"][1] // layouts["channels"] * channels + def index_k_to_encode(self, layout, index): index = np.ravel_multi_index(index, layout["k_space"], order="F") # linear to encoding permuted @@ -484,9 +488,9 @@ def layouts(self): def get_rel_fg_index(self, fg_type): try: - return self.fg_list.index(fg_type) - except MissingProperty: - raise KeyError(f"Framegroup {fg_type} not found in fg_list") from MissingProperty + return self._dataset.dim_type[self._dataset.encoded_dim :].index(fg_type) + except ValueError: + raise KeyError(f"Framegroup {fg_type} not found in dim_type") from ValueError def scale(self): self._dataset.data = np.reshape(self._dataset.data, self._dataset.shape_storage, order="F") @@ -575,14 +579,14 @@ def ra(self, slice_): layouts, layouts_ra = self._get_ra_layouts(slice_) - array_ra = np.zeros(layouts_ra["shape_storage"], dtype=self.numpy_dtype) + array_ra = np.zeros(layouts_ra["shape_storage"], dtype=self._dataset.numpy_dtype) - fp = np.memmap(self._dataset.path, dtype=self.numpy_dtype, mode="r", shape=layouts["shape_storage"], order="F") + fp = np.memmap(self._dataset.path, dtype=self._dataset.numpy_dtype, mode="r", shape=layouts["shape_storage"], order="F") for slice_ra, slice_full in self._generate_ra_indices(layouts_ra, layouts): array_ra[slice_ra] = np.array(fp[slice_full]) - array_ra = self.reshape_fw(array_ra, layouts_ra) + array_ra = self.deserialize(array_ra, layouts_ra) singletons = tuple(i for i, v in enumerate(slice_) if isinstance(v, int)) @@ -593,7 +597,7 @@ def _get_ra_layouts(self, slice_full): layouts_ra = deepcopy(layouts) layouts_ra["mask"] = np.zeros(layouts["shape_fg"], dtype=bool, order="F") - layouts_ra["mask"][slice_full[self.encoded_dim :]] = True + layouts_ra["mask"][slice_full[self._dataset.encoded_dim :]] = True layouts_ra["shape_fg"], layouts_ra["offset_fg"] = self._get_ra_shape(layouts_ra["mask"]) layouts_ra["shape_frames"] = (np.prod(layouts_ra["shape_fg"], dtype=int),) layouts_ra["shape_storage"] = layouts_ra["shape_block"] + layouts_ra["shape_frames"] @@ -615,10 +619,10 @@ def _get_ra_shape(self, mask): return tuple(ra_shape), np.array(ra_offset) def _generate_ra_indices(self, layouts_ra, layouts): - for index_ra in np.ndindex(layouts_ra["shape_final"][self.encoded_dim :]): + for index_ra in np.ndindex(layouts_ra["shape_final"][self._dataset.encoded_dim :]): index = tuple(np.array(index_ra) + layouts_ra["offset_fg"]) - index = tuple(0 for i in range(self.encoded_dim)) + index - index_ra_f = tuple(0 for i in range(self.encoded_dim)) + index_ra + index = tuple(0 for i in range(self._dataset.encoded_dim)) + index + index_ra_f = tuple(0 for i in range(self._dataset.encoded_dim)) + index_ra index_ra_f = np.ravel_multi_index(index_ra_f, layouts_ra["shape_final"], order="F") index = np.ravel_multi_index(index, layouts["shape_final"], order="F") @@ -626,6 +630,6 @@ def _generate_ra_indices(self, layouts_ra, layouts): index_ra_f = np.unravel_index(index_ra_f, layouts_ra["shape_storage"], order="F") index = np.unravel_index(index, layouts["shape_storage"], order="F") - slice_ra = tuple(slice(None) for i in range(self.encoded_dim)) + index_ra_f[self.encoded_dim :] - slice_full = tuple(slice(None) for i in range(self.encoded_dim)) + index[self.encoded_dim :] + slice_ra = tuple(slice(None) for i in range(self._dataset.encoded_dim)) + index_ra_f[self._dataset.encoded_dim :] + slice_full = tuple(slice(None) for i in range(self._dataset.encoded_dim)) + index[self._dataset.encoded_dim :] yield slice_ra, slice_full diff --git a/test/conftest.py b/test/conftest.py index e9963e8..0380a8f 100644 --- a/test/conftest.py +++ b/test/conftest.py @@ -199,6 +199,25 @@ def _find_2dseq_datasets(dataset_name: str): return datasets +def _find_random_access_datasets(dataset_name: str): + dataset_root = TEST_DATA_ROOT / dataset_name + if not dataset_root.exists(): + return [] + + folder_obj = Folder(dataset_root, dataset_state={"parameter_files": [], "property_files": [], "load": 2}) + datasets = folder_obj.get_dataset_list_rec() + datasets_2dseq = [dataset for dataset in datasets if dataset.type == "2dseq"] + + representative_fids = {} + for dataset in (dataset for dataset in datasets if dataset.type == "fid"): + layout = "channel" in dataset.dim_type + current = representative_fids.get(layout) + if current is None or dataset.path.stat().st_size < current.path.stat().st_size: + representative_fids[layout] = dataset + + return [*datasets_2dseq, *representative_fids.values()] + + def _ensure_test_data(dataset_name: str): dataset_dir = TEST_DATA_ROOT / dataset_name if dataset_dir.exists() and any(dataset_dir.iterdir()): @@ -273,7 +292,7 @@ def pytest_generate_tests(metafunc): dataset_root = TEST_DATA_ROOT / dataset_name folder_obj = Folder(dataset_root, dataset_state={"parameter_files": [], "property_files": [], "load": 2}) - for dataset in _find_2dseq_datasets(dataset_name): + for dataset in _find_random_access_datasets(dataset_name): ra_ids.append(f"{dataset_name}/{dataset.id}") ra_items.append((dataset.path, ref_state.get(dataset.id, {}))) diff --git a/test/test_random_access.py b/test/test_random_access.py index e6502e3..c456f7b 100644 --- a/test/test_random_access.py +++ b/test/test_random_access.py @@ -4,39 +4,14 @@ def test_ra(test_ra_data): - d1 = Dataset(test_ra_data[0]) - core_index = tuple(slice(None) for i in range(d1.encoded_dim)) - d2 = Dataset(test_ra_data[0], random_access=True) + loaded = Dataset(test_ra_data[0]) + mmap = Dataset(test_ra_data[0], mmap=True) - # if "slices" in test_ra_data[0].keys(): - # for s in test_ra_data[0]['slices']: - # slice_ = json_to_slice(s) - # assert np.array_equal(d1.data[slice_], d2.data[slice_]) - # else: - # test by single slice - index - for index in np.ndindex(d1.shape[d1.encoded_dim :]): - assert np.array_equal(d1.data[core_index + index], d2.data[core_index + index]) + core = tuple(slice(None) for _ in range(loaded.encoded_dim)) + frame_shape = loaded.shape[loaded.encoded_dim :] - # test all possible slices - for slice_ in generate_slices(d1.shape[d1.encoded_dim :]): - assert np.array_equal(d1.data[core_index + slice_], d2.data[core_index + slice_]) + index = tuple(0 for _ in frame_shape) + assert np.array_equal(loaded.data[core + index], mmap.data[core + index]) - -def generate_slices(shape): - slices = [] - for i1 in np.ndindex(shape): - for i2 in np.ndindex(shape): - if np.all(np.array(i1) <= np.array(i2)): - slice_ = tuple(slice(i1_, i2_ + 1) for i1_, i2_ in zip(i1, i2)) - slices.append(slice_) - return slices - - -def json_to_slice(s): - slice_ = [] - for item in s: - if isinstance(item, str): - slice_.append(eval(item)) - elif isinstance(item, int): - slice_.append(item) - return tuple(slice_) + subarray = tuple(slice(0, min(size, 2)) for size in frame_shape) + assert np.array_equal(loaded.data[core + subarray], mmap.data[core + subarray])