pytme 0.1.9__cp311-cp311-macosx_14_0_arm64.whl → 0.2.0b0__cp311-cp311-macosx_14_0_arm64.whl

tme/matching_data.py

@@ -47,6 +47,9 @@ class MatchingData:
         self.target_filter = {}
 
         self._invert_target = False
+        self._rotations = None
+
+        self._set_batch_dimension()
 
     @staticmethod
     def _shape_to_slice(shape: Tuple[int]):
@@ -149,13 +152,13 @@ class MatchingData:
         def _warn_on_mismatch(
             expectation: float, computation: float, name: str
         ) -> float:
-            expectation, computation = float(expectation), float(computation)
             if expectation is None:
                 expectation = computation
+            expectation, computation = float(expectation), float(computation)
 
             if abs(computation) > abs(expectation):
                 warnings.warn(
-                    f"Computed {name} value is more extreme than value specified in file"
+                    f"Computed {name} value is more extreme than value in file header"
                     f" (|{computation}| > |{expectation}|). This may lead to issues"
                     " with padding and contrast inversion."
                 )
@@ -176,11 +179,9 @@ class MatchingData:
             arr_max = _warn_on_mismatch(arr_max, arr.max(), "max")
 
             # Avoid in-place operation in case ret is not floating point
-            ret = np.divide(
-                np.subtract(-ret, arr_min),
-                np.subtract(arr_max, arr_min)
+            ret = (
+                -np.divide(np.subtract(ret, arr_min), np.subtract(arr_max, arr_min)) + 1
             )
-
         return ret
 
     def subset_by_slice(
@@ -223,14 +224,16 @@ class MatchingData:
         if target_pad is None:
             target_pad = np.zeros(len(self._target.shape), dtype=int)
         if template_pad is None:
-            template_pad = np.zeros(len(self._target.shape), dtype=int)
-
-        indices = compute_full_convolution_index(
-            outer_shape=self._target.shape,
-            inner_shape=self._template.shape,
-            outer_split=target_slice,
-            inner_split=template_slice,
-        )
+            template_pad = np.zeros(len(self._template.shape), dtype=int)
+
+        indices = None
+        if len(self._target.shape) == len(self._template.shape):
+            indices = compute_full_convolution_index(
+                outer_shape=self._target.shape,
+                inner_shape=self._template.shape,
+                outer_split=target_slice,
+                inner_split=template_slice,
+            )
 
         target_subset = self.subset_array(
             arr=self._target,
@@ -246,13 +249,19 @@ class MatchingData:
         )
         ret = self.__class__(target=target_subset, template=template_subset)
 
-        ret._translation_offset = np.add(
-            [x.start for x in target_slice],
-            [x.start for x in template_slice],
-        )
-        ret.template_filter = self.template_filter
+        target_offset = np.zeros(len(self._output_target_shape), dtype=int)
+        target_offset[(target_offset.size - len(target_slice)) :] = [
+            x.start for x in target_slice
+        ]
+        template_offset = np.zeros(len(self._output_target_shape), dtype=int)
+        template_offset[(template_offset.size - len(template_slice)) :] = [
+            x.start for x in template_slice
+        ]
+        ret._translation_offset = np.add(target_offset, template_offset)
 
-        ret.rotations, ret.indices = self.rotations, indices
+        ret.template_filter = self.template_filter
+        ret.target_filter = self.target_filter
+        ret._rotations, ret.indices = self.rotations, indices
         ret._target_pad, ret._template_pad = target_pad, template_pad
         ret._invert_target = self._invert_target
 
@@ -282,10 +291,14 @@ class MatchingData:
         """
         Transfer the class instance's numpy arrays to the current backend.
         """
+        backend_arr = type(backend.zeros((1), dtype=backend._default_dtype))
         for attr_name, attr_value in vars(self).items():
             if isinstance(attr_value, np.ndarray):
                 converted_array = backend.to_backend_array(attr_value.copy())
                 setattr(self, attr_name, converted_array)
+            elif isinstance(attr_value, backend_arr):
+                converted_array = backend.to_backend_array(attr_value)
+                setattr(self, attr_name, converted_array)
 
         self._default_dtype = backend._default_dtype
         self._complex_dtype = backend._complex_dtype
@@ -427,13 +440,13 @@ class MatchingData:
             An array indicating the padding for each dimension of the target.
         """
         target_padding = backend.zeros(
-            len(self.target.shape), dtype=backend._default_dtype_int
+            len(self._output_target_shape), dtype=backend._default_dtype_int
         )
 
         if pad_target:
             backend.subtract(
-                self._template.shape,
-                backend.mod(self._template.shape, 2),
+                self._output_template_shape,
+                backend.mod(self._output_template_shape, 2),
                 out=target_padding,
             )
             if hasattr(self, "_is_target_batch"):
@@ -475,20 +488,30 @@ class MatchingData:
         fourier_shift = backend.zeros(len(fourier_pad))
 
         if not pad_fourier:
-            fourier_pad = backend.full(shape=len(fourier_pad), fill_value=1, dtype=int)
+            fourier_pad = backend.full(
+                shape=(len(fourier_pad),),
+                fill_value=1,
+                dtype=backend._default_dtype_int,
+            )
 
         fourier_pad = backend.to_backend_array(fourier_pad)
         if hasattr(self, "_batch_mask"):
             batch_mask = backend.to_backend_array(self._batch_mask)
             fourier_pad[batch_mask] = 1
 
-        ret = backend.compute_convolution_shapes(target_shape, fourier_pad)
+        pad_shape = backend.maximum(target_shape, template_shape)
+        ret = backend.compute_convolution_shapes(pad_shape, fourier_pad)
         convolution_shape, fast_shape, fast_ft_shape = ret
         if not pad_fourier:
             fourier_shift = 1 - backend.astype(backend.divide(template_shape, 2), int)
             fourier_shift -= backend.mod(template_shape, 2)
             shape_diff = backend.subtract(fast_shape, convolution_shape)
             shape_diff = backend.astype(backend.divide(shape_diff, 2), int)
+
+            if hasattr(self, "_batch_mask"):
+                batch_mask = backend.to_backend_array(self._batch_mask)
+                shape_diff[batch_mask] = 0
+
             backend.add(fourier_shift, shape_diff, out=fourier_shift)
 
         return fast_shape, fast_ft_shape, fourier_shift
@@ -523,15 +546,21 @@ class MatchingData:
     def target(self):
         """Returns the target NDArray."""
         if isinstance(self._target, Density):
-            return self._target.data
-        return self._target
+            target = self._target.data
+        else:
+            target = self._target
+        out_shape = backend.to_numpy_array(self._output_target_shape)
+        return target.reshape(tuple(int(x) for x in out_shape))
 
     @property
     def template(self):
         """Returns the reversed template NDArray."""
+        template = self._template
         if isinstance(self._template, Density):
-            return backend.reverse(self._template.data)
-        return backend.reverse(self._template)
+            template = self._template.data
+        template = backend.reverse(template)
+        out_shape = backend.to_numpy_array(self._output_template_shape)
+        return template.reshape(tuple(int(x) for x in out_shape))
 
     @template.setter
     def template(self, template: NDArray):
@@ -558,9 +587,15 @@ class MatchingData:
     @property
     def target_mask(self):
         """Returns the target mask NDArray."""
+        target_mask = self._target_mask
         if isinstance(self._target_mask, Density):
-            return self._target_mask.data
-        return self._target_mask
+            target_mask = self._target_mask.data
+
+        if target_mask is not None:
+            out_shape = backend.to_numpy_array(self._output_target_shape)
+            target_mask = target_mask.reshape(tuple(int(x) for x in out_shape))
+
+        return target_mask
 
     @target_mask.setter
     def target_mask(self, mask: NDArray):
@@ -587,9 +622,15 @@ class MatchingData:
         template : NDArray
             Array to set as the template.
         """
+        mask = self._template_mask
         if isinstance(self._template_mask, Density):
-            return backend.reverse(self._template_mask.data)
-        return backend.reverse(self._template_mask)
+            mask = self._template_mask.data
+
+        if mask is not None:
+            mask = backend.reverse(mask)
+            out_shape = backend.to_numpy_array(self._output_template_shape)
+            mask = mask.reshape(tuple(int(x) for x in out_shape))
+        return mask
 
     @template_mask.setter
     def template_mask(self, mask: NDArray):