pytme 0.2.2__cp311-cp311-macosx_14_0_arm64.whl → 0.2.4__cp311-cp311-macosx_14_0_arm64.whl

tests/test_orientations.py

@@ -0,0 +1,173 @@
+from tempfile import mkstemp
+
+import pytest
+import numpy as np
+
+from tme import Orientations
+
+
+class TestDensity:
+    def setup_method(self):
+        self.translations = np.random.rand(100, 3).astype(np.float32)
+        self.rotations = np.random.rand(100, 3).astype(np.float32)
+        self.scores = np.random.rand(100).astype(np.float32)
+        self.details = np.random.rand(100).astype(np.float32)
+
+        self.orientations = Orientations(
+            translations=self.translations,
+            rotations=self.rotations,
+            scores=self.scores,
+            details=self.details,
+        )
+
+    def teardown_method(self):
+        self.translations = None
+        self.rotations = None
+        self.scores = None
+        self.details = None
+        self.orientations = None
+
+    def test_initialization(self):
+        orientations = Orientations(
+            translations=self.translations,
+            rotations=self.rotations,
+            scores=self.scores,
+            details=self.details,
+        )
+
+        assert np.array_equal(self.translations, orientations.translations)
+        assert np.array_equal(self.rotations, orientations.rotations)
+        assert np.array_equal(self.scores, orientations.scores)
+        assert np.array_equal(self.details, orientations.details)
+
+    def test_initialization_type(self):
+        orientations = Orientations(
+            translations=self.translations.astype(int),
+            rotations=self.rotations.astype(int),
+            scores=self.scores.astype(int),
+            details=self.details.astype(int),
+        )
+        assert np.issubdtype(orientations.translations.dtype, np.floating)
+        assert np.issubdtype(orientations.rotations.dtype, np.floating)
+        assert np.issubdtype(orientations.scores.dtype, np.floating)
+        assert np.issubdtype(orientations.details.dtype, np.floating)
+
+    def test_initialization_error(self):
+        with pytest.raises(ValueError):
+            _ = Orientations(
+                translations=self.translations,
+                rotations=np.random.rand(self.translations.shape[0] + 1),
+                scores=self.scores,
+                details=self.details,
+            )
+
+        with pytest.raises(ValueError):
+            _ = Orientations(
+                translations=np.random.rand(self.translations.shape[0]),
+                rotations=np.random.rand(self.translations.shape[0] + 1),
+                scores=self.scores,
+                details=self.details,
+            )
+            _ = Orientations(
+                translations=self.translations,
+                rotations=np.random.rand(self.translations.shape[0]),
+                scores=self.scores,
+                details=self.details,
+            )
+
+        assert True
+
+    @pytest.mark.parametrize("file_format", ("text", "relion", "tbl"))
+    def test_to_file(self, file_format: str):
+        _, output_file = mkstemp(suffix=f".{file_format}")
+        self.orientations.to_file(output_file)
+        assert True
+
+    @pytest.mark.parametrize("file_format", ("text", "star", "tbl"))
+    def test_from_file(self, file_format: str):
+        _, output_file = mkstemp(suffix=f".{file_format}")
+        self.orientations.to_file(output_file)
+        orientations_new = Orientations.from_file(output_file)
+
+        assert np.array_equal(
+            self.orientations.translations, orientations_new.translations
+        )
+
+    @pytest.mark.parametrize("input_format", ("text", "star", "tbl"))
+    @pytest.mark.parametrize("output_format", ("text", "star", "tbl"))
+    def test_file_format_io(self, input_format: str, output_format: str):
+        _, output_file = mkstemp(suffix=f".{input_format}")
+        _, output_file2 = mkstemp(suffix=f".{output_format}")
+
+        self.orientations.to_file(output_file)
+        orientations_new = Orientations.from_file(output_file)
+        orientations_new.to_file(output_file2)
+
+        assert True
+
+    @pytest.mark.parametrize("drop_oob", (True, False))
+    @pytest.mark.parametrize("shape", (10, 40, 80))
+    @pytest.mark.parametrize("odd", (True, False))
+    def test_extraction(self, shape: int, drop_oob: bool, odd: bool):
+        if odd:
+            shape = shape + (1 - shape % 2)
+
+        data = np.random.rand(50, 50, 50)
+        translations = np.array([[25, 25, 25], [15, 25, 35], [35, 25, 15], [0, 15, 49]])
+        orientations = Orientations(
+            translations=translations,
+            rotations=np.random.rand(*translations.shape),
+            scores=np.random.rand(translations.shape[0]),
+            details=np.random.rand(translations.shape[0]),
+        )
+        extraction_shape = np.repeat(np.array(shape), data.ndim)
+        orientations, cand_slices, obs_slices = orientations.get_extraction_slices(
+            target_shape=data.shape,
+            extraction_shape=extraction_shape,
+            drop_out_of_box=drop_oob,
+            return_orientations=True,
+        )
+        assert orientations.translations.shape[0] == len(cand_slices)
+        assert len(cand_slices) == len(obs_slices)
+
+        cand_slices2, obs_slices2 = orientations.get_extraction_slices(
+            target_shape=data.shape,
+            extraction_shape=extraction_shape,
+            drop_out_of_box=drop_oob,
+            return_orientations=False,
+        )
+        assert cand_slices == cand_slices2
+        assert obs_slices == obs_slices2
+
+        # Check whether extraction slices are pasted in center
+        out = np.zeros(extraction_shape, dtype=data.dtype)
+        center = np.divide(extraction_shape, 2) + np.mod(extraction_shape, 2)
+        center = center.astype(int)
+        for index, (cand_slice, obs_slice) in enumerate(zip(cand_slices, obs_slices)):
+            out[cand_slice] = data[obs_slice]
+            assert np.allclose(
+                out[tuple(center)],
+                data[tuple(orientations.translations[index].astype(int))],
+            )
+
+    @pytest.mark.parametrize(
+        "order", (("x", "y", "z"), ("z", "y", "x"), ("y", "x", "z"))
+    )
+    def test_txt_sort(self, order: str):
+        _, output_file = mkstemp(suffix=".tsv")
+        translations = ((50, 30, 20), (10, 5, 30))
+
+        with open(output_file, mode="w", encoding="utf-8") as ofile:
+            ofile.write("\t".join([str(x) for x in order]) + "\n")
+            for translation in translations:
+                ofile.write("\t".join([str(x) for x in translation]) + "\n")
+
+        translations = np.array(translations).astype(np.float32)
+        orientations = Orientations.from_file(output_file)
+
+        out_order = zip(order, range(len(order)))
+        out_order = tuple(
+            x[1] for x in sorted(out_order, key=lambda x: x[0], reverse=True)
+        )
+
+        assert np.array_equal(translations[..., out_order], orientations.translations)

tests/test_packaging.py

@@ -0,0 +1,95 @@
+import pytest
+import numpy as np
+
+from tme.package import GaussianKernel, KernelFitting
+from tme.analyzer import PeakCallerFast
+
+class TestKernelFitting:
+    def setup_method(self):
+        self.number_of_peaks = 100
+        self.min_distance = 5
+        self.data = np.random.rand(100, 100, 100)
+        self.rotation_matrix = np.eye(3)
+
+        self.kernel_box = np.random.choice(np.arange(5, 15), size=self.data.ndim)
+        self.template = np.random.rand(*self.kernel_box)
+
+    @pytest.mark.parametrize(
+        "kernel",
+        [
+            (GaussianKernel),
+        ],
+    )
+    def test_initialization(self, kernel):
+        kernel_params = kernel.estimate_parameters(self.template)
+        _ = KernelFitting(
+            number_of_peaks=self.number_of_peaks,
+            min_distance=self.min_distance,
+            kernel_box=self.kernel_box,
+            kernel_class=kernel,
+            kernel_params=kernel_params,
+            peak_caller=PeakCallerFast,
+        )
+
+    @pytest.mark.parametrize("kernel",[GaussianKernel])
+    def test__call__(self, kernel):
+        kernel_params = kernel.estimate_parameters(self.template)
+        score_analyzer = KernelFitting(
+            number_of_peaks=self.number_of_peaks,
+            min_distance=self.min_distance,
+            kernel_box=self.kernel_box,
+            kernel_class=kernel,
+            kernel_params=kernel_params,
+            peak_caller=PeakCallerFast,
+        )
+        score_analyzer(self.data, rotation_matrix=self.rotation_matrix)
+
+
+class TestKernel:
+    def setup_method(self):
+        self.number_of_peaks = 100
+        self.min_distance = 5
+        self.data = np.random.rand(100, 100, 100)
+        self.rotation_matrix = np.eye(3)
+
+    def test_initialization(self):
+        _ = GaussianKernel()
+
+    @pytest.mark.parametrize(
+        "kernel",
+        [
+            GaussianKernel,
+        ],
+    )
+    def test_fit(self, kernel):
+        params, success, final_error = kernel.fit(self.data)
+        assert isinstance(params, tuple)
+        assert len(params) == 3
+
+    @pytest.mark.parametrize(
+        "kernel",
+        [
+            GaussianKernel,
+        ],
+    )
+    def test_fit_height(self, kernel):
+        params, success, final_error = kernel.fit(self.data)
+        height, mean, cov = params
+        params, succes, final_error = kernel.fit_height(
+            data=self.data, mean=mean, cov=cov
+        )
+        assert np.allclose(params[1], mean)
+        assert np.allclose(params[2], cov)
+
+    @pytest.mark.parametrize(
+        "kernel",
+        [
+            GaussianKernel,
+        ],
+    )
+    def test_score(self, kernel):
+        params, *_ = kernel.fit(self.data)
+        params2, *_ = kernel.fit(self.data)
+
+        score = kernel.score(params, params2)
+        assert score == 1

tests/test_parser.py

@@ -0,0 +1,33 @@
+import pytest
+from importlib_resources import files
+
+from tme.parser import Parser, PDBParser
+
+
+class TestParser:
+    def setup_method(self):
+        self.pdb_file = str(files("tests.data").joinpath("Structures/5khe.pdb"))
+
+    def teardown_method(self):
+        self.pdb_file = None
+
+    def test_initialize_parser_error(self):
+        with pytest.raises(TypeError):
+            _ = Parser(self.pdb_file)
+
+    def test_parser_get(self):
+        parser = PDBParser(self.pdb_file)
+        assert parser.get("NOTPRESENT", None) is None
+        assert parser.get("record_type", None) is not None
+
+    def test_parser_keys(self):
+        parser = PDBParser(self.pdb_file)
+        assert parser.keys() == parser._data.keys()
+
+    def test_parser_values(self):
+        parser = PDBParser(self.pdb_file)
+        assert str(parser.values()) == str(parser._data.values())
+
+    def test_parser_items(self):
+        parser = PDBParser(self.pdb_file)
+        assert parser.items() == parser._data.items()

tests/test_structure.py

@@ -0,0 +1,243 @@
+from os import remove
+from tempfile import mkstemp
+from importlib_resources import files
+
+import pytest
+import numpy as np
+
+from tme import Structure
+from tme.matching_utils import euler_to_rotationmatrix, minimum_enclosing_box
+
+
+STRUCTURE_ATTRIBUTES = [
+    "record_type",
+    "atom_serial_number",
+    "atom_name",
+    "atom_coordinate",
+    "alternate_location_indicator",
+    "residue_name",
+    "chain_identifier",
+    "residue_sequence_number",
+    "code_for_residue_insertion",
+    "occupancy",
+    "temperature_factor",
+    "segment_identifier",
+    "element_symbol",
+    "charge",
+    "metadata",
+]
+
+
+class TestStructure:
+    def setup_method(self):
+        self.structure = Structure.from_file(
+            str(files("tests.data").joinpath("Structures/5khe.cif"))
+        )
+        _, self.path = mkstemp()
+
+    def teardown_method(self):
+        del self.structure
+        remove(self.path)
+
+    def compare_structures(self, structure1, structure2, exclude_attributes=[]):
+        for attribute in STRUCTURE_ATTRIBUTES:
+            if attribute in exclude_attributes:
+                continue
+            value = getattr(structure1, attribute)
+            value_comparison = getattr(structure2, attribute)
+            if isinstance(value, np.ndarray):
+                assert np.all(value_comparison == value)
+            else:
+                assert value == value_comparison
+
+    def test_initialization(self):
+        structure = Structure(
+            record_type=self.structure.record_type,
+            atom_serial_number=self.structure.atom_serial_number,
+            atom_name=self.structure.atom_name,
+            atom_coordinate=self.structure.atom_coordinate,
+            alternate_location_indicator=self.structure.alternate_location_indicator,
+            residue_name=self.structure.residue_name,
+            chain_identifier=self.structure.chain_identifier,
+            residue_sequence_number=self.structure.residue_sequence_number,
+            code_for_residue_insertion=self.structure.code_for_residue_insertion,
+            occupancy=self.structure.occupancy,
+            temperature_factor=self.structure.temperature_factor,
+            segment_identifier=self.structure.segment_identifier,
+            element_symbol=self.structure.element_symbol,
+            charge=self.structure.charge,
+            metadata=self.structure.metadata,
+        )
+
+        for attribute in STRUCTURE_ATTRIBUTES:
+            value = getattr(self.structure, attribute)
+            value_comparison = getattr(structure, attribute)
+            if isinstance(value, np.ndarray):
+                assert np.all(value_comparison == value)
+            else:
+                assert value == value_comparison
+
+    @pytest.mark.parametrize(
+        "modified_attribute",
+        [
+            ("record_type"),
+            ("atom_serial_number"),
+            ("atom_name"),
+            ("atom_coordinate"),
+            ("alternate_location_indicator"),
+            ("residue_name"),
+            ("chain_identifier"),
+            ("residue_sequence_number"),
+            ("code_for_residue_insertion"),
+            ("occupancy"),
+            ("temperature_factor"),
+            ("segment_identifier"),
+            ("element_symbol"),
+        ],
+    )
+    def test_initialization_errors(self, modified_attribute):
+        kwargs = {
+            attribute: getattr(self.structure, attribute)
+            for attribute in STRUCTURE_ATTRIBUTES
+            if attribute != modified_attribute
+        }
+        kwargs[modified_attribute] = getattr(self.structure, modified_attribute)[:1]
+
+        with pytest.raises(ValueError):
+            Structure(**kwargs)
+
+    def test__getitem__(self):
+        ret_single_index = self.structure[1]
+        ret = self.structure[[1]]
+        self.compare_structures(ret_single_index, ret)
+
+        ret = self.structure[self.structure.record_type == "ATOM"]
+        assert np.all(ret.record_type == "ATOM")
+
+        ret = self.structure[self.structure.element_symbol == "C"]
+        assert np.all(ret.element_symbol == "C")
+
+    def test__repr__(self):
+        unique_chains = "-".join(
+            [
+                ",".join([str(x) for x in entity])
+                for entity in self.structure.metadata["unique_chains"]
+            ]
+        )
+
+        min_atom = np.min(self.structure.atom_serial_number)
+        max_atom = np.max(self.structure.atom_serial_number)
+        n_atom = self.structure.atom_serial_number.size
+
+        min_residue = np.min(self.structure.residue_sequence_number)
+        max_residue = np.max(self.structure.residue_sequence_number)
+        n_residue = np.unique(self.structure.residue_sequence_number).size
+
+        repr_str = (
+            f"Structure object at {id(self.structure)}\n"
+            f"Unique Chains: {unique_chains}, "
+            f"Atom Range: {min_atom}-{max_atom} [N = {n_atom}], "
+            f"Residue Range: {min_residue}-{max_residue} [N = {n_residue}]"
+        )
+        assert repr_str == self.structure.__repr__()
+
+    @pytest.mark.parametrize(
+        "path",
+        [
+            str(files("tests.data").joinpath("Structures/5khe.cif")),
+            str(files("tests.data").joinpath("Structures/5khe.pdb")),
+        ],
+    )
+    def test_fromfile(self, path):
+        _ = Structure.from_file(path)
+
+    def test_fromfile_error(self):
+        with pytest.raises(NotImplementedError):
+            _ = Structure.from_file("madeup.extension")
+
+    @pytest.mark.parametrize("file_format", [("cif"), ("pdb")])
+    def test_to_file(self, file_format):
+        _, path = mkstemp()
+        path = f"{path}.{file_format}"
+        self.structure.to_file(path)
+        read = self.structure.from_file(path)
+        comparison = self.structure.copy()
+
+        self.compare_structures(comparison, read, exclude_attributes=["metadata"])
+
+    def test_to_file_error(self):
+        _, path = mkstemp()
+        path = f"{path}.RAISERROR"
+        with pytest.raises(NotImplementedError):
+            self.structure.to_file(path)
+
+    def test_subset_by_chain(self):
+        chain = "A"
+        ret = self.structure.subset_by_chain(chain=chain)
+        assert np.all(ret.chain_identifier == chain)
+
+    def test_subset_by_chain_range(self):
+        chain, start, stop = "A", 0, 20
+        ret = self.structure.subset_by_range(chain=chain, start=start, stop=stop)
+        assert np.all(ret.chain_identifier == chain)
+        assert np.all(
+            np.logical_and(
+                ret.residue_sequence_number >= start,
+                ret.residue_sequence_number <= stop,
+            )
+        )
+
+    def test_center_of_mass(self):
+        center_of_mass = self.structure.center_of_mass()
+        assert center_of_mass.shape[0] == self.structure.atom_coordinate.shape[1]
+        assert np.allclose(center_of_mass, [-0.89391639, 29.94908928, -2.64736741])
+
+    def test_centered(self):
+        ret, translation = self.structure.centered()
+        box = minimum_enclosing_box(coordinates=self.structure.atom_coordinate.T)
+        assert np.allclose(ret.center_of_mass(), np.divide(box, 2), atol=1)
+
+    def test__get_atom_weights_error(self):
+        with pytest.raises(NotImplementedError):
+            self.structure._get_atom_weights(
+                self.structure.atom_name, weight_type="RAISEERROR"
+            )
+
+    def test_compare_structures(self):
+        rmsd = Structure.compare_structures(self.structure, self.structure)
+        assert rmsd == 0
+
+        rmsd = Structure.compare_structures(
+            self.structure, self.structure, weighted=True
+        )
+        assert rmsd == 0
+
+        translation = (3, 0, 0)
+        structure_transform = self.structure.rigid_transform(
+            translation=translation,
+            rotation_matrix=np.eye(self.structure.atom_coordinate.shape[1]),
+        )
+        rmsd = Structure.compare_structures(self.structure, structure_transform)
+        assert np.allclose(rmsd, np.linalg.norm(translation))
+
+    def test_comopare_structures_error(self):
+        ret = self.structure[[1, 2, 3, 4, 5]]
+        with pytest.raises(ValueError):
+            Structure.compare_structures(self.structure, ret)
+
+    def test_align_structures(self):
+        rotation_matrix = euler_to_rotationmatrix((20, -10, 45))
+        translation = (10, 0, -15)
+
+        structure_transform = self.structure.rigid_transform(
+            rotation_matrix=rotation_matrix, translation=translation
+        )
+        aligned, final_rmsd = Structure.align_structures(
+            self.structure, structure_transform
+        )
+        assert final_rmsd <= 0.1
+
+        aligned, final_rmsd = Structure.align_structures(
+            self.structure, structure_transform, sampling_rate=1
+        )
+        assert final_rmsd <= 1

tme/__init__.py

@@ -1,4 +1,3 @@
-from . import extensions
 from .__version__ import __version__
 from .density import Density
 from .preprocessor import Preprocessor

tme/__version__.py

@@ -1 +1 @@
-__version__ = "0.2.2"
+__version__ = "0.2.4"

tme/analyzer.py

@@ -459,14 +459,14 @@ class PeakCaller(ABC):
 
         final_order = top_scores[
             filter_points_indices(
-                coordinates=peak_positions[top_scores],
+                coordinates=peak_positions[top_scores, :],
                 min_distance=self.min_distance,
                 batch_dims=self.batch_dims,
             )
         ]
 
-        self.peak_list[0] = peak_positions[final_order,]
-        self.peak_list[1] = rotations[final_order,]
+        self.peak_list[0] = peak_positions[final_order, :]
+        self.peak_list[1] = rotations[final_order, :]
         self.peak_list[2] = peak_scores[final_order]
         self.peak_list[3] = peak_details[final_order]
 
@@ -475,6 +475,7 @@ class PeakCaller(ABC):
         fast_shape: Tuple[int],
         targetshape: Tuple[int],
         templateshape: Tuple[int],
+        convolution_shape: Tuple[int] = None,
         fourier_shift: Tuple[int] = None,
         convolution_mode: str = None,
         shared_memory_handler=None,
@@ -488,6 +489,7 @@ class PeakCaller(ABC):
             return self
 
         # Wrap peaks around score space
+        convolution_shape = be.to_backend_array(convolution_shape)
         fast_shape = be.to_backend_array(fast_shape)
         if fourier_shift is not None:
             fourier_shift = be.to_backend_array(fourier_shift)
@@ -501,10 +503,9 @@ class PeakCaller(ABC):
             )
 
         # Remove padding to fast Fourier (and potential full convolution) shape
+        output_shape = convolution_shape
         targetshape = be.to_backend_array(targetshape)
         templateshape = be.to_backend_array(templateshape)
-        fast_shape = be.minimum(be.add(targetshape, templateshape) - 1, fast_shape)
-        output_shape = fast_shape
         if convolution_mode == "same":
             output_shape = targetshape
         elif convolution_mode == "valid":
@@ -515,7 +516,7 @@ class PeakCaller(ABC):
 
         output_shape = be.to_backend_array(output_shape)
         starts = be.astype(
-            be.divide(be.subtract(fast_shape, output_shape), 2),
+            be.divide(be.subtract(convolution_shape, output_shape), 2),
             be._int_dtype,
         )
         stops = be.add(starts, output_shape)
@@ -1019,6 +1020,7 @@ class MaxScoreOverRotations:
         self,
         targetshape: Tuple[int],
         templateshape: Tuple[int],
+        convolution_shape: Tuple[int],
         fourier_shift: Tuple[int] = None,
         convolution_mode: str = None,
         shared_memory_handler=None,
@@ -1039,6 +1041,7 @@ class MaxScoreOverRotations:
             "s1": targetshape,
             "s2": templateshape,
             "convolution_mode": convolution_mode,
+            "convolution_shape": convolution_shape,
         }
         if convolution_mode is not None:
             scores = apply_convolution_mode(scores, **convargs)

tme/backends/__init__.py

@@ -147,7 +147,7 @@ class BackendManager:
         _dependencies = {
             "numpyfftw": "numpy",
             "cupy": "cupy",
-            "pytorch": "pytorch",
+            "pytorch": "torch",
             "mlx": "mlx",
             "jax": "jax",
         }

tme/backends/_jax_utils.py

@@ -19,9 +19,9 @@ def _correlate(template: BackendArray, ft_target: BackendArray) -> BackendArray:
     """
     Computes :py:meth:`tme.matching_exhaustive.cc_setup`.
     """
-    template_ft = jnp.fft.rfftn(template)
+    template_ft = jnp.fft.rfftn(template, s=template.shape)
     template_ft = template_ft.at[:].multiply(ft_target)
-    correlation = jnp.fft.irfftn(template_ft)
+    correlation = jnp.fft.irfftn(template_ft, s=template.shape)
     return correlation
 
 
@@ -77,14 +77,15 @@ def _reciprocal_target_std(
     --------
     :py:meth:`tme.matching_exhaustive.flc_scoring`.
     """
-    ft_template_mask = jnp.fft.rfftn(template_mask)
+    ft_shape = template_mask.shape
+    ft_template_mask = jnp.fft.rfftn(template_mask, s=ft_shape)
 
     # E(X^2)- E(X)^2
-    exp_sq = jnp.fft.irfftn(ft_target2 * ft_template_mask)
+    exp_sq = jnp.fft.irfftn(ft_target2 * ft_template_mask, s=ft_shape)
     exp_sq = exp_sq.at[:].divide(n_observations)
 
     ft_template_mask = ft_template_mask.at[:].multiply(ft_target)
-    sq_exp = jnp.fft.irfftn(ft_template_mask)
+    sq_exp = jnp.fft.irfftn(ft_template_mask, s=ft_shape)
     sq_exp = sq_exp.at[:].divide(n_observations)
     sq_exp = sq_exp.at[:].power(2)
 
@@ -99,7 +100,7 @@ def _reciprocal_target_std(
 
 
 def _apply_fourier_filter(arr: BackendArray, arr_filter: BackendArray) -> BackendArray:
-    arr_ft = jnp.fft.rfftn(arr)
+    arr_ft = jnp.fft.rfftn(arr, s=arr.shape)
     arr_ft = arr_ft.at[:].multiply(arr_filter)
     return arr.at[:].set(jnp.fft.irfftn(arr_ft, s=arr.shape))
 
@@ -107,6 +108,7 @@ def _apply_fourier_filter(arr: BackendArray, arr_filter: BackendArray) -> Backen
 def _identity(arr: BackendArray, arr_filter: BackendArray) -> BackendArray:
     return arr
 
+
 @partial(
     pmap,
     in_axes=(0,) + (None,) * 6,
@@ -127,8 +129,8 @@ def scan(
     if hasattr(target_filter, "shape"):
         target = _apply_fourier_filter(target, target_filter)
 
-    ft_target = jnp.fft.rfftn(target)
-    ft_target2 = jnp.fft.rfftn(jnp.square(target))
+    ft_target = jnp.fft.rfftn(target, s=fast_shape)
+    ft_target2 = jnp.fft.rfftn(jnp.square(target), s=fast_shape)
     inv_denominator, target, scoring_func = None, None, _flc_scoring
     if not rotate_mask:
         n_observations = jnp.sum(template_mask)