From 2664b04daec51dc10867bfbb1c2a41c11c21c87c Mon Sep 17 00:00:00 2001
From: McHaillet <martenchaillet@gmail.com>
Date: Wed, 8 May 2024 11:56:26 +0200
Subject: [PATCH 01/21] add phase randomziation functionfor templates

---
 src/pytom_tm/template.py | 47 +++++++++++++++++++++++++++++++++++++++-
 1 file changed, 46 insertions(+), 1 deletion(-)

diff --git a/src/pytom_tm/template.py b/src/pytom_tm/template.py
index 940d6f4e..23ef3977 100644
--- a/src/pytom_tm/template.py
+++ b/src/pytom_tm/template.py
@@ -5,7 +5,12 @@
 from scipy.ndimage import center_of_mass, zoom
 from scipy.fft import rfftn, irfftn
 from typing import Optional
-from pytom_tm.weights import create_gaussian_low_pass
+from pytom_tm.weights import (
+    create_ctf,
+    create_gaussian_low_pass,
+    radial_average,
+    radial_reduced_grid,
+)
 
 
 def generate_template_from_map(
@@ -103,3 +108,43 @@ def generate_template_from_map(
         irfftn(rfftn(input_map) * lpf, s=input_map.shape),
         input_spacing / output_spacing
     )
+
+
+def phase_randomize_template(
+        template: npt.NDArray[float],
+        seed: int = 321,
+):
+    """Create a version of the template that has its phases randomly
+    permuted in Fourier space.
+
+    Parameters
+    ----------
+    template: npt.NDArray[float]
+        input structure
+    seed: int, default 321
+        seed for random number generator for phase permutation
+
+    Returns
+    -------
+    result: npt.NDArray[float]
+        phase randomized version of the template
+    """
+    ft = np.fft.rfftn(template)
+    amplitude = np.abs(ft)
+
+    # permute the phases in flattened version of the array
+    phase = np.angle(ft).flatten()
+    grid = np.fft.ifftshift(
+        radial_reduced_grid(template.shape), axes=(0, 1)
+    ).flatten()
+    relevant_freqs = grid <= 1  # permute only up to Nyquist
+    noise = np.zeros_like(phase)
+    rng = np.random.default_rng(seed)
+    noise[relevant_freqs] = rng.permutation(phase[relevant_freqs])
+
+    # construct the new template
+    noise = np.reshape(noise, amplitude.shape)
+    result = np.fft.irfftn(
+        amplitude * np.exp(1j * noise), s=template.shape
+    )
+    return result

From 72733738072ac669728965eeca44431752d1cc2d Mon Sep 17 00:00:00 2001
From: McHaillet <martenchaillet@gmail.com>
Date: Wed, 8 May 2024 12:06:50 +0200
Subject: [PATCH 02/21] merge tests

---
 tests/test_template.py | 15 ++++++++++++++-
 1 file changed, 14 insertions(+), 1 deletion(-)

diff --git a/tests/test_template.py b/tests/test_template.py
index 775f8929..82a7d858 100644
--- a/tests/test_template.py
+++ b/tests/test_template.py
@@ -1,7 +1,9 @@
 import unittest
 import numpy as np
-from pytom_tm.template import generate_template_from_map
 from scipy.ndimage import center_of_mass
+from pytom_tm.template import (
+    phase_randomize_template, generate_template_from_map
+)
 
 
 class TestTemplate(unittest.TestCase):
@@ -73,3 +75,14 @@ def test_lowpass_resolution(self):
         with self.assertNoLogs(level='WARNING'):
             _ = generate_template_from_map(self.template, 1., 1.,
                                            filter_to_resolution=2.5)
+
+    def test_phase_randomize_template(self):
+        randomized = phase_randomize_template(
+            self.template,  # use default seed
+        )
+        self.assertEqual(self.template.shape, randomized.shape)
+        self.assertGreater(
+            (self.template != randomized).sum(), 0,
+            msg='After phase randomization the template should '
+                'no longer be equal to the input.'
+        )

From bbb2db4b8d8a06a243507825a15b0ad314558f4a Mon Sep 17 00:00:00 2001
From: McHaillet <martenchaillet@gmail.com>
Date: Wed, 8 May 2024 12:12:43 +0200
Subject: [PATCH 03/21] focus just on this PR

---
 tests/test_template.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tests/test_template.py b/tests/test_template.py
index 82a7d858..a90499fb 100644
--- a/tests/test_template.py
+++ b/tests/test_template.py
@@ -2,7 +2,7 @@
 import numpy as np
 from scipy.ndimage import center_of_mass
 from pytom_tm.template import (
-    phase_randomize_template, generate_template_from_map
+    phase_randomize_template
 )
 
 

From b232879054c873b7bafa0a9badec2d4619552b33 Mon Sep 17 00:00:00 2001
From: McHaillet <martenchaillet@gmail.com>
Date: Wed, 8 May 2024 12:13:42 +0200
Subject: [PATCH 04/21] remove unused imports

---
 src/pytom_tm/matching.py | 3 +--
 1 file changed, 1 insertion(+), 2 deletions(-)

diff --git a/src/pytom_tm/matching.py b/src/pytom_tm/matching.py
index a41baf07..8c1e0fe9 100644
--- a/src/pytom_tm/matching.py
+++ b/src/pytom_tm/matching.py
@@ -4,10 +4,9 @@
 import voltools as vt
 import gc
 from typing import Optional
-from cupyx.scipy.fft import rfftn, irfftn, fftshift
+from cupyx.scipy.fft import rfftn, irfftn
 from tqdm import tqdm
 from pytom_tm.correlation import mean_under_mask, std_under_mask
-from packaging import version
 
 
 class TemplateMatchingPlan:

From 491fd61c3a978c98ad1dc47b88032d857927f63a Mon Sep 17 00:00:00 2001
From: McHaillet <martenchaillet@gmail.com>
Date: Wed, 8 May 2024 12:38:12 +0200
Subject: [PATCH 05/21] add some stuff for the template matching plan setup for
 noise correction

---
 src/pytom_tm/matching.py | 39 ++++++++++++++++++++++++++++++++++++---
 1 file changed, 36 insertions(+), 3 deletions(-)

diff --git a/src/pytom_tm/matching.py b/src/pytom_tm/matching.py
index 8c1e0fe9..ba67075d 100644
--- a/src/pytom_tm/matching.py
+++ b/src/pytom_tm/matching.py
@@ -7,6 +7,7 @@
 from cupyx.scipy.fft import rfftn, irfftn
 from tqdm import tqdm
 from pytom_tm.correlation import mean_under_mask, std_under_mask
+from pytom_tm.template import phase_randomize_template
 
 
 class TemplateMatchingPlan:
@@ -16,7 +17,8 @@ def __init__(
             template: npt.NDArray[float],
             mask: npt.NDArray[float],
             device_id: int,
-            wedge: Optional[npt.NDArray[float]] = None
+            wedge: Optional[npt.NDArray[float]] = None,
+            phase_randomized_template: Optional[npt.NDArray[float]] = None,
     ):
         """Initialize a template matching plan. All the necessary cupy arrays will be allocated on the GPU.
 
@@ -33,6 +35,8 @@ def __init__(
         wedge: Optional[npt.NDArray[float]], default None
             3D numpy array that contains the Fourier space weighting for the template, it should be in Fourier
             reduced form, with dimensions (sx, sx, sx // 2 + 1)
+        phase_randomized_template: Optional[npt.NDArray[float]], default None
+            initialize the plan with a phase randomized version of the template for noise correction
         """
         # Search volume + and fft transform plan for the volume
         volume_shape = volume.shape
@@ -64,6 +68,16 @@ def __init__(
         self.scores = cp.ones(volume_shape, dtype=cp.float32)*-1000
         self.angles = cp.ones(volume_shape, dtype=cp.float32)*-1000
 
+        self.random_phase_template_texture = None
+        self.noise_scores = None
+        if phase_randomized_template is not None:
+            self.random_phase_template_texture = vt.StaticVolume(
+                cp.asarray(phase_randomized_template, dtype=cp.float32, order='C'),
+                interpolation='filt_bspline',
+                device=f'gpu:{device_id}',
+            )
+            self.noise_scores = cp.ones(volume_shape, dtype=cp.float32)*-1000
+
         # wait for stream to complete the work
         cp.cuda.stream.get_current_stream().synchronize()
 
@@ -91,7 +105,9 @@ def __init__(
             angle_ids: list[int],
             mask_is_spherical: bool = True,
             wedge: Optional[npt.NDArray[float]] = None,
-            stats_roi: Optional[tuple[slice, slice, slice]] = None
+            stats_roi: Optional[tuple[slice, slice, slice]] = None,
+            noise_correction: bool = False,
+            rng_seed: int = 321,
     ):
         """Initialize a template matching run.
 
@@ -126,6 +142,11 @@ def __init__(
         stats_roi: Optional[tuple[slice, slice, slice]], default None
             region of interest to calculate statistics on the search volume, default will just take the full search
             volume
+        noise_correction: bool, default False
+            initialize template matching with a phase randomized version of the template that is used to subtract
+            background noise from the score map; expense is more gpu memory and computation time
+        rng_seed: int, default 321
+            seed for rng in phase randomization
         """
         cp.cuda.Device(device_id).use()
 
@@ -145,8 +166,20 @@ def __init__(
             )
         else:
             self.stats_roi = stats_roi
+        self.noise_correction = noise_correction
+        shuffled_template = (
+            phase_randomize_template(template, rng_seed)
+            if noise_correction else None
+        )
 
-        self.plan = TemplateMatchingPlan(volume, template, mask, device_id, wedge=wedge)
+        self.plan = TemplateMatchingPlan(
+            volume,
+            template,
+            mask,
+            device_id,
+            wedge=wedge,
+            phase_randomized_template=shuffled_template,
+        )
 
     def run(self) -> tuple[npt.NDArray[float], npt.NDArray[float], dict]:
         """Run the template matching job.

From 74be2d5d410de76c2f3bc54cd88a4315b0c7ebc1 Mon Sep 17 00:00:00 2001
From: McHaillet <martenchaillet@gmail.com>
Date: Wed, 8 May 2024 14:18:25 +0200
Subject: [PATCH 06/21] add the noise correction to the template matching loop

---
 src/pytom_tm/matching.py | 61 ++++++++++++++++++++++++++++++++++++++--
 1 file changed, 58 insertions(+), 3 deletions(-)

diff --git a/src/pytom_tm/matching.py b/src/pytom_tm/matching.py
index ba67075d..c792e41f 100644
--- a/src/pytom_tm/matching.py
+++ b/src/pytom_tm/matching.py
@@ -167,6 +167,8 @@ def __init__(
         else:
             self.stats_roi = stats_roi
         self.noise_correction = noise_correction
+
+        # create a 'random noise' version of the template
         shuffled_template = (
             phase_randomize_template(template, rng_seed)
             if noise_correction else None
@@ -294,6 +296,50 @@ def run(self) -> tuple[npt.NDArray[float], npt.NDArray[float], dict]:
                 square_sum_kernel(self.plan.ccc_map * roi_mask) / roi_size
             )
 
+            if self.noise_correction:
+                # Rotate noise template texture into template
+                self.plan.random_phase_template_texture.transform(
+                    rotation=(rotation[0], rotation[1], rotation[2]),
+                    rotation_order='rzxz',
+                    output=self.plan.template,
+                    rotation_units='rad'
+                )
+
+                if self.plan.wedge is not None:
+                    # Add wedge to the template after rotating
+                    self.plan.template = irfftn(
+                        rfftn(self.plan.template) * self.plan.wedge,
+                        s=self.plan.template.shape
+                    )
+
+                # Normalize and mask template
+                mean = mean_under_mask(self.plan.template, self.plan.mask, mask_weight=self.plan.mask_weight)
+                std = std_under_mask(self.plan.template, self.plan.mask, mean, mask_weight=self.plan.mask_weight)
+                self.plan.template = ((self.plan.template - mean) / std) * self.plan.mask
+
+                # Paste in center
+                self.plan.template_padded[pad_index] = self.plan.template
+
+                # Fast local correlation function between volume and template, norm is the standard deviation at each
+                # point in the volume in the masked area
+                self.plan.ccc_map = (
+                        irfftn(self.plan.volume_rft_conj * rfftn(self.plan.template_padded),
+                               s=self.plan.template_padded.shape)
+                        / self.plan.std_volume
+                )
+                # update noise scores results
+                update_noise_template_results_kernel(
+                    self.plan.noise_scores,
+                    self.plan.ccc_map,
+                    self.plan.noise_scores,
+                )
+
+        # do the noise correction on the scores map: substract the noise scores first,
+        # and then add the noise mean to ensure stats are consistent
+        self.plan.scores = (
+            (self.plan.scores - self.plan.noise_scores) + self.plan.noise_scores.mean()
+        )
+
         # Get correct orientation back!
         # Use same method as William Wan's STOPGAP
         # (https://doi.org/10.1107/S205979832400295X): the search volume is Fourier
@@ -353,13 +399,22 @@ def std_under_mask_convolution(
 
 """Update scores and angles if a new maximum is found."""
 update_results_kernel = cp.ElementwiseKernel(
-    'float32 scores, float32 ccc_map, float32 angle_id',
-    'float32 out1, float32 out2',
-    'if (scores < ccc_map) {out1 = ccc_map; out2 = angle_id;}',
+    'float32 scores, float32 ccc_map, float32 angles',
+    'float32 scores_out, float32 angles_out',
+    'if (scores < ccc_map) {scores_out = ccc_map; angles_out = angles;}',
     'update_results'
 )
 
 
+"""Update scores for noise template"""
+update_noise_template_results_kernel = cp.ElementwiseKernel(
+    'float32 scores, float32 ccc_map',
+    'float32 scores_out',
+    'if (scores < ccc_map) {scores_out = ccc_map}',
+    'update_noise_template_results'
+)
+
+
 """Calculate the sum of squares in a volume. Mean is assumed to be 0 which makes this operation a lot faster."""
 square_sum_kernel = cp.ReductionKernel(
     'T x',  # input params

From 9f7d05e27ae8c865703d3e74f3d199a65991e9a8 Mon Sep 17 00:00:00 2001
From: McHaillet <martenchaillet@gmail.com>
Date: Wed, 8 May 2024 14:20:51 +0200
Subject: [PATCH 07/21] add test using the phase randomized template for noise
 correction

---
 tests/test_template_matching.py | 33 +++++++++++++++++++++++++++++++++
 1 file changed, 33 insertions(+)

diff --git a/tests/test_template_matching.py b/tests/test_template_matching.py
index 118cbd3d..8d9f1f8b 100644
--- a/tests/test_template_matching.py
+++ b/tests/test_template_matching.py
@@ -92,3 +92,36 @@ def test_search_non_spherical_mask(self):
         self.assertAlmostEqual(stats['std'], 0.005187, places=4,
                 msg='Standard deviation of the search should be almost equal')
 
+    def test_search_noise_correction(self):
+        angle_id = 100
+        rotation = self.angles[angle_id]
+        loc = (77, 26, 40)
+        self.volume[loc[0] - self.t_size // 2: loc[0] + self.t_size // 2,
+        loc[1] - self.t_size // 2: loc[1] + self.t_size // 2,
+        loc[2] - self.t_size // 2: loc[2] + self.t_size // 2] = vt.transform(
+            self.template,
+            rotation=rotation,
+            rotation_units='rad',
+            rotation_order='rzxz',
+            device='cpu'
+        )
+
+        tm = TemplateMatchingGPU(
+            0,
+            0,
+            self.volume,
+            self.template,
+            self.mask,
+            self.angles,
+            list(range(len(self.angles))),
+            noise_correction=True,
+        )
+        score_volume, angle_volume, stats = tm.run()
+
+        ind = np.unravel_index(score_volume.argmax(), self.volume.shape)
+        self.assertTrue(score_volume.max() > 0.99, msg='lcc max value lower than expected')
+        self.assertEqual(angle_id, angle_volume[ind])
+        self.assertSequenceEqual(loc, ind)
+        self.assertEqual(stats['search_space'], 256000000, msg='Search space should exactly equal this value')
+        self.assertAlmostEqual(stats['std'], 0.005187, places=5,
+                               msg='Standard deviation of the search should be almost equal')

From 06507bc136d816d8ec493d53f52ba7370a96205a Mon Sep 17 00:00:00 2001
From: McHaillet <martenchaillet@gmail.com>
Date: Wed, 8 May 2024 14:26:57 +0200
Subject: [PATCH 08/21] substraction of noise scores should be conditional of
 course

---
 src/pytom_tm/matching.py | 9 +++++----
 1 file changed, 5 insertions(+), 4 deletions(-)

diff --git a/src/pytom_tm/matching.py b/src/pytom_tm/matching.py
index c792e41f..db165811 100644
--- a/src/pytom_tm/matching.py
+++ b/src/pytom_tm/matching.py
@@ -336,9 +336,10 @@ def run(self) -> tuple[npt.NDArray[float], npt.NDArray[float], dict]:
 
         # do the noise correction on the scores map: substract the noise scores first,
         # and then add the noise mean to ensure stats are consistent
-        self.plan.scores = (
-            (self.plan.scores - self.plan.noise_scores) + self.plan.noise_scores.mean()
-        )
+        if self.noise_correction:
+            self.plan.scores = (
+                (self.plan.scores - self.plan.noise_scores) + self.plan.noise_scores.mean()
+            )
 
         # Get correct orientation back!
         # Use same method as William Wan's STOPGAP
@@ -410,7 +411,7 @@ def std_under_mask_convolution(
 update_noise_template_results_kernel = cp.ElementwiseKernel(
     'float32 scores, float32 ccc_map',
     'float32 scores_out',
-    'if (scores < ccc_map) {scores_out = ccc_map}',
+    'if (scores < ccc_map) {scores_out = ccc_map;}',
     'update_noise_template_results'
 )
 

From 39f0df0c434c23b029cbefb8eb7461729c48b46c Mon Sep 17 00:00:00 2001
From: McHaillet <martenchaillet@gmail.com>
Date: Wed, 8 May 2024 15:14:24 +0200
Subject: [PATCH 09/21] add the noise correction options to the entry points
 and job

---
 src/pytom_tm/entry_points.py | 21 +++++++++++++++++++++
 src/pytom_tm/tmjob.py        | 17 +++++++++++++++--
 2 files changed, 36 insertions(+), 2 deletions(-)

diff --git a/src/pytom_tm/entry_points.py b/src/pytom_tm/entry_points.py
index a46350ea..31898fd1 100644
--- a/src/pytom_tm/entry_points.py
+++ b/src/pytom_tm/entry_points.py
@@ -756,6 +756,25 @@ def match_template(argv=None):
         "apply it to the tomogram patch and template. Effectively puts more weight on "
         "high resolution features and sharpens the correlation peaks.",
     )
+    additional_group = parser.add_argument_group('Additional options')
+    additional_group.add_argument(
+        "--random-phase-correction",
+        action="store_true",
+        default=False,
+        required=False,
+        help="Run template matching simultaneously with a phase randomized version of "
+             "the template, and subtract this 'noise' map from the final score map. "
+             "For this method please see STOPGAP as a reference: "
+             "https://doi.org/10.1107/S205979832400295X ."
+    )
+    additional_group.add_argument(
+        "--rng-seed",
+        action=LargerThanZero,
+        default=321,
+        required=False,
+        help="Specify a seed for the random number generator used for phase "
+             "randomization for consistent results!"
+    )
     device_group = parser.add_argument_group('Device control')
     device_group.add_argument(
         "-g",
@@ -836,6 +855,8 @@ def match_template(argv=None):
         whiten_spectrum=args.spectral_whitening,
         rotational_symmetry=args.z_axis_rotational_symmetry,
         particle_diameter=args.particle_diameter,
+        random_phase_correction=args.random_phase_correction,
+        rng_seed=args.rng_seed,
     )
 
     score_volume, angle_volume = run_job_parallel(
diff --git a/src/pytom_tm/tmjob.py b/src/pytom_tm/tmjob.py
index 159e7011..ee693e1d 100644
--- a/src/pytom_tm/tmjob.py
+++ b/src/pytom_tm/tmjob.py
@@ -66,6 +66,8 @@ def load_json_to_tmjob(file_name: pathlib.Path, load_for_extraction: bool = True
         pytom_tm_version_number=data.get('pytom_tm_version_number', '0.3.0'),
         job_loaded_for_extraction=load_for_extraction,
         particle_diameter=data.get('particle_diameter', None),
+        random_phase_correction=data.get('random_phase_correction', False),
+        rng_seed=data.get('rng_seed', 321),
     )
     # if the file originates from an old version set the phase shift for compatibility
     if (
@@ -212,6 +214,8 @@ def __init__(
             pytom_tm_version_number: str = PYTOM_TM_VERSION,
             job_loaded_for_extraction: bool = False,
             particle_diameter: Optional[float] = None,
+            random_phase_correction: bool = False,
+            rng_seed: int = 321,
     ):
         """
         Parameters
@@ -263,9 +267,12 @@ def __init__(
             a string with the version number of pytom_tm for backward compatibility
         job_loaded_for_extraction: bool, default False
             flag to set for finished template matching jobs that are loaded back for extraction, it prevents
-            recalculation of the whitening filter which is unnecessary at this stage
         particle_diameter: Optional[float], default None
             particle diameter (in Angstrom) to calculate angular search
+        random_phase_correction: bool, default False,
+            run matching with a phase randomized version of the template to correct scores for noise
+        rng_seed: int, default 321
+            set a seed for the rng for phase randomization
         """
         self.mask = mask
         self.mask_is_spherical = mask_is_spherical
@@ -384,6 +391,10 @@ def __init__(
             weights /= weights.max()  # scale to 1
             np.save(self.whitening_filter, weights)
 
+        # phase randomization options
+        self.random_phase_correction = random_phase_correction
+        self.rng_seed = rng_seed
+
         # Job details
         self.job_key = job_key
         self.leader = None  # the job that spawned this job
@@ -746,7 +757,9 @@ def start_job(
             angle_ids=angle_ids,
             mask_is_spherical=self.mask_is_spherical,
             wedge=template_wedge,
-            stats_roi=search_volume_roi
+            stats_roi=search_volume_roi,
+            noise_correction=self.random_phase_correction,
+            rng_seed=self.rng_seed,
         )
         results = tm.run()
         score_volume = results[0][:self.search_size[0], :self.search_size[1], :self.search_size[2]]

From 222225c216b6ae0e4e027fd1cae12c3644d715a6 Mon Sep 17 00:00:00 2001
From: McHaillet <martenchaillet@gmail.com>
Date: Wed, 8 May 2024 15:17:13 +0200
Subject: [PATCH 10/21] update docstring

---
 src/pytom_tm/matching.py | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/src/pytom_tm/matching.py b/src/pytom_tm/matching.py
index db165811..9f401700 100644
--- a/src/pytom_tm/matching.py
+++ b/src/pytom_tm/matching.py
@@ -116,7 +116,8 @@ def __init__(
         - pyTME: https://github.com/KosinskiLab/pyTME
 
         The precalculation of conjugated FTs of the tomo was (AFAIK) introduced
-        by STOPGAP!
+        by STOPGAP! Also, they introduced simultaneous matching with a phase randomized
+        version of the template.
 
         Parameters
         ----------

From ef7315ab48f3e6c5be8e6e453fa7c0ba7be99f80 Mon Sep 17 00:00:00 2001
From: McHaillet <martenchaillet@gmail.com>
Date: Wed, 8 May 2024 15:20:23 +0200
Subject: [PATCH 11/21] add shorthand input for phase randomization

---
 src/pytom_tm/entry_points.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/src/pytom_tm/entry_points.py b/src/pytom_tm/entry_points.py
index 31898fd1..28d6e459 100644
--- a/src/pytom_tm/entry_points.py
+++ b/src/pytom_tm/entry_points.py
@@ -758,6 +758,7 @@ def match_template(argv=None):
     )
     additional_group = parser.add_argument_group('Additional options')
     additional_group.add_argument(
+        "-r",
         "--random-phase-correction",
         action="store_true",
         default=False,

From 8578dad42c6588025ea6026d026f46a0926793c6 Mon Sep 17 00:00:00 2001
From: McHaillet <martenchaillet@gmail.com>
Date: Wed, 8 May 2024 15:21:46 +0200
Subject: [PATCH 12/21] remove unused import

---
 src/pytom_tm/tmjob.py | 1 -
 1 file changed, 1 deletion(-)

diff --git a/src/pytom_tm/tmjob.py b/src/pytom_tm/tmjob.py
index ee693e1d..d2192eac 100644
--- a/src/pytom_tm/tmjob.py
+++ b/src/pytom_tm/tmjob.py
@@ -1,5 +1,4 @@
 from __future__ import annotations
-from importlib import metadata
 from packaging import version
 import pathlib
 import os

From 5d76bdf515fc074c783c4c4ae397aabf3dee2f96 Mon Sep 17 00:00:00 2001
From: McHaillet <martenchaillet@gmail.com>
Date: Thu, 6 Jun 2024 14:51:28 +0200
Subject: [PATCH 13/21] add correct import back

---
 tests/test_template.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tests/test_template.py b/tests/test_template.py
index a90499fb..2b57aa83 100644
--- a/tests/test_template.py
+++ b/tests/test_template.py
@@ -2,7 +2,7 @@
 import numpy as np
 from scipy.ndimage import center_of_mass
 from pytom_tm.template import (
-    phase_randomize_template
+    generate_template_from_map, phase_randomize_template
 )
 
 

From d243a9726e36eed1ddb7fb285c691119efbb67ec Mon Sep 17 00:00:00 2001
From: McHaillet <martenchaillet@gmail.com>
Date: Thu, 6 Jun 2024 14:53:08 +0200
Subject: [PATCH 14/21] add randomization test

---
 tests/test_template.py | 7 +++++++
 1 file changed, 7 insertions(+)

diff --git a/tests/test_template.py b/tests/test_template.py
index 2b57aa83..f8031529 100644
--- a/tests/test_template.py
+++ b/tests/test_template.py
@@ -86,3 +86,10 @@ def test_phase_randomize_template(self):
             msg='After phase randomization the template should '
                 'no longer be equal to the input.'
         )
+
+        randomized_seeded = phase_randomize_template(
+            self.template, 11  # use default seed
+        )
+        diff = np.abs(randomized_seeded - randomized).sum()
+        self.assertEqual(diff, 0, msg='Different seed should return different '
+                                      'randomization')

From b1a765d553da26eab7390655569cf3d01be8e16a Mon Sep 17 00:00:00 2001
From: McHaillet <martenchaillet@gmail.com>
Date: Thu, 6 Jun 2024 14:54:57 +0200
Subject: [PATCH 15/21] add entry point test

---
 tests/test_entry_points.py | 5 +++++
 1 file changed, 5 insertions(+)

diff --git a/tests/test_entry_points.py b/tests/test_entry_points.py
index 5c7724e7..e88aaf5d 100644
--- a/tests/test_entry_points.py
+++ b/tests/test_entry_points.py
@@ -147,6 +147,11 @@ def start(arg_dict):  # simplify run
         arguments['--low-pass'] = '50'
         start(arguments)
 
+        # phase randomization test
+        arguments = defaults.copy()
+        arguments['-r'] = ''
+        start(arguments)
+
         # test debug files
         arguments = defaults.copy()
         arguments['--log'] = 'debug'

From d6c7829db65e6ff4fa860b861f3c6a36321b3a7e Mon Sep 17 00:00:00 2001
From: McHaillet <martenchaillet@gmail.com>
Date: Thu, 6 Jun 2024 14:56:48 +0200
Subject: [PATCH 16/21] fix test assertion

---
 tests/test_template.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/tests/test_template.py b/tests/test_template.py
index f8031529..e283c99e 100644
--- a/tests/test_template.py
+++ b/tests/test_template.py
@@ -91,5 +91,5 @@ def test_phase_randomize_template(self):
             self.template, 11  # use default seed
         )
         diff = np.abs(randomized_seeded - randomized).sum()
-        self.assertEqual(diff, 0, msg='Different seed should return different '
-                                      'randomization')
+        self.assertNotEqual(diff, 0,
+                            msg='Different seed should return different randomization')

From 3179a4687783ffdd4b0d53ddf26b35001e79ad38 Mon Sep 17 00:00:00 2001
From: McHaillet <martenchaillet@gmail.com>
Date: Thu, 6 Jun 2024 15:01:44 +0200
Subject: [PATCH 17/21] fix search tests

---
 tests/test_template_matching.py | 6 ++++--
 1 file changed, 4 insertions(+), 2 deletions(-)

diff --git a/tests/test_template_matching.py b/tests/test_template_matching.py
index 8d9f1f8b..a4b2da1b 100644
--- a/tests/test_template_matching.py
+++ b/tests/test_template_matching.py
@@ -122,6 +122,8 @@ def test_search_noise_correction(self):
         self.assertTrue(score_volume.max() > 0.99, msg='lcc max value lower than expected')
         self.assertEqual(angle_id, angle_volume[ind])
         self.assertSequenceEqual(loc, ind)
-        self.assertEqual(stats['search_space'], 256000000, msg='Search space should exactly equal this value')
-        self.assertAlmostEqual(stats['std'], 0.005187, places=5,
+        expected_search_space = len(self.angles) * self.volume.size
+        self.assertEqual(stats['search_space'], expected_search_space,
+                         msg='Search space should exactly equal this value')
+        self.assertAlmostEqual(stats['std'], 0.005187, places=4,
                                msg='Standard deviation of the search should be almost equal')

From cb020663441691bdade960272f1a1316caa04297 Mon Sep 17 00:00:00 2001
From: Marten Chaillet <58044494+McHaillet@users.noreply.github.com>
Date: Fri, 7 Jun 2024 13:52:24 +0200
Subject: [PATCH 18/21] scipy rfftn and doi link

Co-authored-by: Sander Roet <sanderroet@hotmail.com>
---
 src/pytom_tm/matching.py | 2 +-
 src/pytom_tm/template.py | 2 +-
 2 files changed, 2 insertions(+), 2 deletions(-)

diff --git a/src/pytom_tm/matching.py b/src/pytom_tm/matching.py
index 9f401700..00a018fb 100644
--- a/src/pytom_tm/matching.py
+++ b/src/pytom_tm/matching.py
@@ -117,7 +117,7 @@ def __init__(
 
         The precalculation of conjugated FTs of the tomo was (AFAIK) introduced
         by STOPGAP! Also, they introduced simultaneous matching with a phase randomized
-        version of the template.
+        version of the template. https://doi.org/10.1107/S205979832400295X
 
         Parameters
         ----------
diff --git a/src/pytom_tm/template.py b/src/pytom_tm/template.py
index 23ef3977..7806de28 100644
--- a/src/pytom_tm/template.py
+++ b/src/pytom_tm/template.py
@@ -144,7 +144,7 @@ def phase_randomize_template(
 
     # construct the new template
     noise = np.reshape(noise, amplitude.shape)
-    result = np.fft.irfftn(
+    result = irfftn(
         amplitude * np.exp(1j * noise), s=template.shape
     )
     return result

From 189ccf961408b8351c89fb5ac16af05e6b568ce9 Mon Sep 17 00:00:00 2001
From: McHaillet <martenchaillet@gmail.com>
Date: Fri, 7 Jun 2024 13:55:52 +0200
Subject: [PATCH 19/21] use scipy rfftn at start as well

---
 src/pytom_tm/template.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/pytom_tm/template.py b/src/pytom_tm/template.py
index 7806de28..a7d102a4 100644
--- a/src/pytom_tm/template.py
+++ b/src/pytom_tm/template.py
@@ -129,7 +129,7 @@ def phase_randomize_template(
     result: npt.NDArray[float]
         phase randomized version of the template
     """
-    ft = np.fft.rfftn(template)
+    ft = rfftn(template)
     amplitude = np.abs(ft)
 
     # permute the phases in flattened version of the array

From 66ffcc06e937c052bc23f4ad8fad639394e7ea6d Mon Sep 17 00:00:00 2001
From: McHaillet <martenchaillet@gmail.com>
Date: Fri, 7 Jun 2024 14:14:58 +0200
Subject: [PATCH 20/21] split correlation to function

---
 src/pytom_tm/matching.py | 80 ++++++++++++++++++----------------------
 1 file changed, 36 insertions(+), 44 deletions(-)

diff --git a/src/pytom_tm/matching.py b/src/pytom_tm/matching.py
index 00a018fb..2311f78e 100644
--- a/src/pytom_tm/matching.py
+++ b/src/pytom_tm/matching.py
@@ -261,28 +261,7 @@ def run(self) -> tuple[npt.NDArray[float], npt.NDArray[float], dict]:
                 rotation_units='rad'
             )
 
-            if self.plan.wedge is not None:
-                # Add wedge to the template after rotating
-                self.plan.template = irfftn(
-                    rfftn(self.plan.template) * self.plan.wedge,
-                    s=self.plan.template.shape
-                )
-
-            # Normalize and mask template
-            mean = mean_under_mask(self.plan.template, self.plan.mask, mask_weight=self.plan.mask_weight)
-            std = std_under_mask(self.plan.template, self.plan.mask, mean, mask_weight=self.plan.mask_weight)
-            self.plan.template = ((self.plan.template - mean) / std) * self.plan.mask
-
-            # Paste in center
-            self.plan.template_padded[pad_index] = self.plan.template
-
-            # Fast local correlation function between volume and template, norm is the standard deviation at each
-            # point in the volume in the masked area
-            self.plan.ccc_map = (
-                irfftn(self.plan.volume_rft_conj * rfftn(self.plan.template_padded),
-                       s=self.plan.template_padded.shape)
-                / self.plan.std_volume
-            )
+            self.correlate(pad_index)
 
             # Update the scores and angle_lists
             update_results_kernel(
@@ -306,28 +285,8 @@ def run(self) -> tuple[npt.NDArray[float], npt.NDArray[float], dict]:
                     rotation_units='rad'
                 )
 
-                if self.plan.wedge is not None:
-                    # Add wedge to the template after rotating
-                    self.plan.template = irfftn(
-                        rfftn(self.plan.template) * self.plan.wedge,
-                        s=self.plan.template.shape
-                    )
-
-                # Normalize and mask template
-                mean = mean_under_mask(self.plan.template, self.plan.mask, mask_weight=self.plan.mask_weight)
-                std = std_under_mask(self.plan.template, self.plan.mask, mean, mask_weight=self.plan.mask_weight)
-                self.plan.template = ((self.plan.template - mean) / std) * self.plan.mask
-
-                # Paste in center
-                self.plan.template_padded[pad_index] = self.plan.template
-
-                # Fast local correlation function between volume and template, norm is the standard deviation at each
-                # point in the volume in the masked area
-                self.plan.ccc_map = (
-                        irfftn(self.plan.volume_rft_conj * rfftn(self.plan.template_padded),
-                               s=self.plan.template_padded.shape)
-                        / self.plan.std_volume
-                )
+                self.correlate(pad_index)
+
                 # update noise scores results
                 update_noise_template_results_kernel(
                     self.plan.noise_scores,
@@ -363,6 +322,39 @@ def run(self) -> tuple[npt.NDArray[float], npt.NDArray[float], dict]:
 
         return results
 
+    def correlate(self, padding_index: tuple[slice, slice, slice]):
+        """Correlate template and tomogram.
+
+        Parameters
+        ----------
+        padding_index: tuple[slice, slice, slice]
+            Location to pad template after weighting and normalization
+        """
+        if self.plan.wedge is not None:
+            # Add wedge to the template after rotating
+            self.plan.template = irfftn(
+                rfftn(self.plan.template) * self.plan.wedge,
+                s=self.plan.template.shape
+            )
+
+        # Normalize and mask template
+        mean = mean_under_mask(self.plan.template, self.plan.mask,
+                               mask_weight=self.plan.mask_weight)
+        std = std_under_mask(self.plan.template, self.plan.mask, mean,
+                             mask_weight=self.plan.mask_weight)
+        self.plan.template = ((self.plan.template - mean) / std) * self.plan.mask
+
+        # Paste in center
+        self.plan.template_padded[padding_index] = self.plan.template
+
+        # Fast local correlation function between volume and template, norm is the standard deviation at each
+        # point in the volume in the masked area
+        self.plan.ccc_map = (
+                irfftn(self.plan.volume_rft_conj * rfftn(self.plan.template_padded),
+                       s=self.plan.template_padded.shape)
+                / self.plan.std_volume
+        )
+
 
 def std_under_mask_convolution(
         volume_rft_conj: cpt.NDArray[float],

From 979cacdbbc92d1178c3e6775bcc424a00df904ff Mon Sep 17 00:00:00 2001
From: McHaillet <martenchaillet@gmail.com>
Date: Fri, 7 Jun 2024 14:24:42 +0200
Subject: [PATCH 21/21] add urandom seed

---
 src/pytom_tm/entry_points.py | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/src/pytom_tm/entry_points.py b/src/pytom_tm/entry_points.py
index 28d6e459..130a9924 100644
--- a/src/pytom_tm/entry_points.py
+++ b/src/pytom_tm/entry_points.py
@@ -20,6 +20,7 @@
     BetweenZeroAndOne,
 )
 from pytom_tm.tmjob import load_json_to_tmjob
+from os import urandom
 
 
 def _parse_argv(argv=None):
@@ -771,7 +772,7 @@ def match_template(argv=None):
     additional_group.add_argument(
         "--rng-seed",
         action=LargerThanZero,
-        default=321,
+        default=int.from_bytes(urandom(8)),
         required=False,
         help="Specify a seed for the random number generator used for phase "
              "randomization for consistent results!"