change indexing

src/mrpro/data/AcqInfo.py

@@ -9,7 +9,6 @@
 from einops import rearrange
 from typing_extensions import Self
 
-import mrpro.utils.typing as type_utils
 from mrpro.data.MoveDataMixin import MoveDataMixin
 from mrpro.data.Rotation import Rotation
 from mrpro.data.SpatialDimension import SpatialDimension
@@ -260,50 +259,48 @@ def spatialdimension_2d(data: np.ndarray) -> SpatialDimension[torch.Tensor]:
         )
         return acq_info
 
-    def add_to_ismrmrd_acquisition(
-        self, acquisition: ismrmrd.Acquisition, other: type_utils.TorchIndexerType, k2: int, k1: int
-    ) -> ismrmrd.Acquisition:
+    def add_to_ismrmrd_acquisition(self, acquisition: ismrmrd.Acquisition, idx: Sequence[int]) -> ismrmrd.Acquisition:
         """ISMRMRD acquisition information for single acquisition."""
-        acquisition.idx.kspace_encode_step_1 = self.idx.k1[*other, k2, k1]
-        acquisition.idx.kspace_encode_step_2 = self.idx.k2[*other, k2, k1]
-        acquisition.idx.average = self.idx.average[*other, k2, k1]
-        acquisition.idx.slice = self.idx.slice[*other, k2, k1]
-        acquisition.idx.contrast = self.idx.contrast[*other, k2, k1]
-        acquisition.idx.phase = self.idx.phase[*other, k2, k1]
-        acquisition.idx.repetition = self.idx.repetition[*other, k2, k1]
-        acquisition.idx.set = self.idx.set[*other, k2, k1]
-        acquisition.idx.segment = self.idx.segment[*other, k2, k1]
+        acquisition.idx.kspace_encode_step_1 = self.idx.k1[*idx]
+        acquisition.idx.kspace_encode_step_2 = self.idx.k2[*idx]
+        acquisition.idx.average = self.idx.average[*idx]
+        acquisition.idx.slice = self.idx.slice[*idx]
+        acquisition.idx.contrast = self.idx.contrast[*idx]
+        acquisition.idx.phase = self.idx.phase[*idx]
+        acquisition.idx.repetition = self.idx.repetition[*idx]
+        acquisition.idx.set = self.idx.set[*idx]
+        acquisition.idx.segment = self.idx.segment[*idx]
         acquisition.idx.user = (
-            self.idx.user0[*other, k2, k1],
-            self.idx.user1[*other, k2, k1],
-            self.idx.user2[*other, k2, k1],
-            self.idx.user3[*other, k2, k1],
-            self.idx.user4[*other, k2, k1],
-            self.idx.user5[*other, k2, k1],
-            self.idx.user6[*other, k2, k1],
-            self.idx.user7[*other, k2, k1],
+            self.idx.user0[*idx],
+            self.idx.user1[*idx],
+            self.idx.user2[*idx],
+            self.idx.user3[*idx],
+            self.idx.user4[*idx],
+            self.idx.user5[*idx],
+            self.idx.user6[*idx],
+            self.idx.user7[*idx],
         )
 
         # active_channesl, number_of_samples and trajectory_dimensions are read-only and cannot be set
-        acquisition.acquisition_time_stamp = self.acquisition_time_stamp[*other, k2, k1, 0]
-        acquisition.available_channels = self.available_channels[*other, k2, k1, 0]
-        acquisition.center_sample = self.center_sample[*other, k2, k1, 0]
-        acquisition.channel_mask = tuple(self.channel_mask[*other, k2, k1, :])
-        acquisition.discard_post = self.discard_post[*other, k2, k1, 0]
-        acquisition.discard_pre = self.discard_pre[*other, k2, k1, 0]
-        acquisition.encoding_space_ref = self.encoding_space_ref[*other, k2, k1, 0]
-        acquisition.measurement_uid = self.measurement_uid[*other, k2, k1, 0]
+        acquisition.acquisition_time_stamp = self.acquisition_time_stamp[*idx][0]
+        acquisition.available_channels = self.available_channels[*idx][0]
+        acquisition.center_sample = self.center_sample[*idx][0]
+        acquisition.channel_mask = tuple(self.channel_mask[*idx])
+        acquisition.discard_post = self.discard_post[*idx][0]
+        acquisition.discard_pre = self.discard_pre[*idx][0]
+        acquisition.encoding_space_ref = self.encoding_space_ref[*idx][0]
+        acquisition.measurement_uid = self.measurement_uid[*idx][0]
         acquisition.patient_table_position = (
-            self.patient_table_position[*other, k2, k1, 0].apply_(m_to_mm).zyx[::-1]
+            self.patient_table_position[*idx][0].apply_(m_to_mm).zyx[::-1]
         )  # zyx -> xyz
-        acquisition.phase_dir = self.orientation.as_directions()[-2][*other, k2, k1, 0].zyx[::-1]  # zyx -> xyz
-        acquisition.physiology_time_stamp = tuple(self.physiology_time_stamp[*other, k2, k1, :])
-        acquisition.position = self.position[*other, k2, k1, 0].apply_(m_to_mm).zyx[::-1]  # zyx -> xyz
-        acquisition.read_dir = self.orientation.as_directions()[-1][*other, k2, k1, 0].zyx[::-1]  # zyx -> xyz
-        acquisition.sample_time_us = self.sample_time_us[*other, k2, k1, 0]
-        acquisition.scan_counter = self.scan_counter[*other, k2, k1, 0]
-        acquisition.slice_dir = self.orientation.as_directions()[-3][*other, k2, k1, 0].zyx[::-1]  # zyx -> xyz
-        acquisition.user_float = tuple(self.user_float[*other, k2, k1, :])
-        acquisition.user_int = tuple(self.user_int[*other, k2, k1, :])
-        acquisition.version = self.version[*other, k2, k1, 0]
+        acquisition.phase_dir = self.orientation.as_directions()[-2][*idx][0].zyx[::-1]  # zyx -> xyz
+        acquisition.physiology_time_stamp = tuple(self.physiology_time_stamp[*idx])
+        acquisition.position = self.position[*idx][0].apply_(m_to_mm).zyx[::-1]  # zyx -> xyz
+        acquisition.read_dir = self.orientation.as_directions()[-1][*idx][0].zyx[::-1]  # zyx -> xyz
+        acquisition.sample_time_us = self.sample_time_us[*idx][0]
+        acquisition.scan_counter = self.scan_counter[*idx][0]
+        acquisition.slice_dir = self.orientation.as_directions()[-3][*idx][0].zyx[::-1]  # zyx -> xyz
+        acquisition.user_float = tuple(self.user_float[*idx])
+        acquisition.user_int = tuple(self.user_int[*idx])
+        acquisition.version = self.version[*idx][0]
         return acquisition

src/mrpro/data/_kdata/KData.py

@@ -293,7 +293,7 @@ def to_file(self, filename: str | Path) -> None:
             for k2 in range(self.data.shape[-3]):
                 for k1 in range(self.data.shape[-2]):
                     acq.clear_all_flags()
-                    acq = self.header.acq_info.add_to_ismrmrd_acquisition(acq, other, k2, k1)
+                    acq = self.header.acq_info.add_to_ismrmrd_acquisition(acq, (*other, k2, k1))
 
                     # Rearrange, switch from zyx to xyz and set trajectory.
                     acq.traj[:] = rearrange(trajectory[:, *other, k2, k1, :].numpy(), 'dim k0->k0 dim')[:, ::-1]