From b9a563025445d4ba5c8e92306e252f8d07fc4c94 Mon Sep 17 00:00:00 2001 From: McHaillet Date: Tue, 16 Apr 2024 16:47:49 +0200 Subject: [PATCH 01/14] speeeeed --- src/pytom_tm/matching.py | 117 +++++++++++++++++---------------------- 1 file changed, 50 insertions(+), 67 deletions(-) diff --git a/src/pytom_tm/matching.py b/src/pytom_tm/matching.py index ea2c83a6..879bf4a6 100644 --- a/src/pytom_tm/matching.py +++ b/src/pytom_tm/matching.py @@ -36,32 +36,33 @@ def __init__( reduced form, with dimensions (sx, sx, sx // 2 + 1) """ # Search volume + and fft transform plan for the volume - self.volume = cp.asarray(volume, dtype=cp.float32, order='C') - self.volume_rft = rfftn(self.volume) + volume_shape = volume.shape + self.volume_rft_conj = rfftn(cp.asarray(volume, dtype=cp.float32, order='C')).conj() + self.volume_sq_rft_conj = rfftn(cp.asarray(volume, dtype=cp.float32, order='C') ** 2).conj() # Explicit fft plan is no longer necessary as cupy generates a plan behind the scene which leads to # comparable timings # Array for storing local standard deviations - self.std_volume = cp.zeros_like(volume, dtype=cp.float32) + self.std_volume = cp.zeros(volume_shape, dtype=cp.float32) # Data for the mask self.mask = cp.asarray(mask, dtype=cp.float32, order='C') self.mask_texture = vt.StaticVolume(self.mask, interpolation='filt_bspline', device=f'gpu:{device_id}') - self.mask_padded = cp.zeros_like(self.volume).astype(cp.float32) + self.mask_padded = cp.zeros(volume_shape, dtype=cp.float32) self.mask_weight = self.mask.sum() # weight of the mask # Init template data self.template = cp.asarray(template, dtype=cp.float32, order='C') self.template_texture = vt.StaticVolume(self.template, interpolation='filt_bspline', device=f'gpu:{device_id}') - self.template_padded = cp.zeros_like(self.volume) + self.template_padded = cp.zeros(volume_shape, dtype=cp.float32) # fourier binary wedge weight for the template self.wedge = cp.asarray(wedge, order='C', dtype=cp.float32) if wedge is not None else None # Initialize result volumes - self.ccc_map = cp.zeros_like(self.volume) - self.scores = cp.ones_like(self.volume)*-1000 - self.angles = cp.ones_like(self.volume)*-1000 + self.ccc_map = cp.zeros(volume_shape, dtype=cp.float32) + self.scores = cp.ones(volume_shape, dtype=cp.float32)*-1000 + self.angles = cp.ones(volume_shape, dtype=cp.float32)*-1000 # wait for stream to complete the work cp.cuda.stream.get_current_stream().synchronize() @@ -69,9 +70,11 @@ def __init__( def clean(self) -> None: """Remove all stored cupy arrays from the GPU's memory pool.""" gpu_memory_pool = cp.get_default_memory_pool() - del self.volume, self.volume_rft, self.mask, self.mask_texture, self.mask_padded, self.template, ( - self.template_texture), self.template_padded, self.wedge, self.ccc_map, self.scores, self.angles, ( - self.std_volume) + del ( + self.volume_rft_conj, self.volume_sq_rft_conj, self.mask, self.mask_texture, self.mask_padded, + self.template, self.template_texture, self.template_padded, self.wedge, self.ccc_map, self.scores, + self.angles, self.std_volume + ) gc.collect() gpu_memory_pool.free_all_blocks() @@ -162,19 +165,24 @@ def run(self) -> tuple[npt.NDArray[float], npt.NDArray[float], dict]: sxv, syv, szv = self.plan.template_padded.shape cxv, cyv, czv = sxv // 2, syv // 2, szv // 2 + # create slice for padding + pad_index = ( + slice(cxv - cxt, cxv + cxt + mx), + slice(cyv - cyt, cyv + cyt + my), + slice(czv - czt, czv + czt + mz), + ) + # calculate roi size - roi_size = self.plan.volume[self.stats_roi].size + roi_size = self.plan.scores[self.stats_roi].size if self.mask_is_spherical: # Then we only need to calculate std volume once - self.plan.mask_padded[cxv - cxt:cxv + cxt + mx, - cyv - cyt:cyv + cyt + my, - czv - czt:czv + czt + mz] = self.plan.mask + self.plan.mask_padded[pad_index] = self.plan.mask self.plan.std_volume = std_under_mask_convolution( - self.plan.volume, + self.plan.volume_rft_conj, + self.plan.volume_sq_rft_conj, self.plan.mask_padded, self.plan.mask_weight, - volume_rft=self.plan.volume_rft - ) + ) * self.plan.mask_weight # Track iterations with a tqdm progress bar for i in tqdm(range(len(self.angle_ids))): @@ -189,16 +197,14 @@ def run(self) -> tuple[npt.NDArray[float], npt.NDArray[float], dict]: output=self.plan.mask, rotation_units='rad' ) - self.plan.mask_padded[cxv - cxt:cxv + cxt + mx, - cyv - cyt:cyv + cyt + my, - czv - czt:czv + czt + mz] = self.plan.mask + self.plan.mask_padded[pad_index] = self.plan.mask # Std volume needs to be recalculated for every rotation of the mask, expensive step self.plan.std_volume = std_under_mask_convolution( - self.plan.volume, + self.plan.volume_rft_conj, + self.plan.volume_sq_rft_conj, self.plan.mask_padded, self.plan.mask_weight, - volume_rft=self.plan.volume_rft, - ) + ) * self.plan.mask_weight # Rotate template self.plan.template_texture.transform( @@ -213,7 +219,7 @@ def run(self) -> tuple[npt.NDArray[float], npt.NDArray[float], dict]: self.plan.template = irfftn( rfftn(self.plan.template) * self.plan.wedge, s=self.plan.template.shape - ).real + ) # Normalize and mask template mean = mean_under_mask(self.plan.template, self.plan.mask, mask_weight=self.plan.mask_weight) @@ -221,16 +227,14 @@ def run(self) -> tuple[npt.NDArray[float], npt.NDArray[float], dict]: self.plan.template = ((self.plan.template - mean) / std) * self.plan.mask # Paste in center - self.plan.template_padded[cxv - cxt:cxv + cxt + mx, - cyv - cyt:cyv + cyt + my, - czv - czt:czv + czt + mz] = self.plan.template + 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 = fftshift( - irfftn(self.plan.volume_rft * rfftn(self.plan.template_padded).conj(), - s=self.plan.template_padded.shape).real - / (self.plan.mask_weight * self.plan.std_volume) + 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 the scores and angle_lists @@ -248,6 +252,11 @@ def run(self) -> tuple[npt.NDArray[float], npt.NDArray[float], dict]: ) / roi_size ) + # get correct orientation back + center = cp.floor(cp.array(self.plan.scores.shape) / 2).astype(int) + 1 + self.plan.scores = cp.roll(cp.flip(self.plan.scores), center, axis=(0, 1, 2)) + self.plan.angles = cp.roll(cp.flip(self.plan.angles), center, axis=(0, 1, 2)) + self.stats['search_space'] = int(roi_size * len(self.angle_ids)) self.stats['variance'] = float(self.stats['variance'] / len(self.angle_ids)) self.stats['std'] = float(cp.sqrt(self.stats['variance'])) @@ -262,66 +271,40 @@ def run(self) -> tuple[npt.NDArray[float], npt.NDArray[float], dict]: def std_under_mask_convolution( - volume: cpt.NDArray[float], + volume_rft_conj: cpt.NDArray[float], + volume_sq_rft_conj: cpt.NDArray[float], padded_mask: cpt.NDArray[float], mask_weight: float, - volume_rft: Optional[cpt.NDArray[complex]] = None ) -> cpt.NDArray[float]: """Calculate the local standard deviation under the mask for each position in the volume. Calculation is done in Fourier space as this is a convolution between volume and mask. Parameters ---------- - volume: cpt.NDArray[float] - cupy array to calculate local std in + volume_rft_conj: cpt.NDArray[float] + complex conjugate of the rft of the search volume + volume_sq_rft_conj: cpt.NDArray[float] + complex conjugate of the rft of the squared search volume padded_mask: cpt.NDArray[float] template mask that has been padded to dimensions of volume mask_weight: float weight of the mask, usually calculated as mask.sum() - volume_rft: Optional[cpt.NDArray[float]], default None - optionally provide a precalculated reduced Fourier transform of volume to save computation Returns ------- std_v: cpt.NDArray[float] array with local standard deviations in volume """ - volume_rft = rfftn(volume) if volume_rft is None else volume_rft + padded_mask_rft = rfftn(padded_mask) std_v = ( - mean_under_mask_convolution(rfftn(volume ** 2), padded_mask, mask_weight) - - mean_under_mask_convolution(volume_rft, padded_mask, mask_weight) ** 2 + irfftn(volume_sq_rft_conj * padded_mask_rft, s=padded_mask.shape) / mask_weight - + (irfftn(volume_rft_conj * padded_mask_rft, s=padded_mask.shape) / mask_weight) ** 2 ) std_v[std_v <= cp.float32(1e-18)] = 1 # prevent potential sqrt of negative value and division by zero std_v = cp.sqrt(std_v) return std_v -def mean_under_mask_convolution( - volume_rft: cpt.NDArray[complex], - mask: cpt.NDArray[float], - mask_weight: float -) -> cpt.NDArray[float]: - """Calculate local mean in volume under the masked region. - - Parameters - ---------- - volume_rft: cpt.NDArray[complex] - array containing the rfftn of the volume - mask: cpt.NDArray[float] - mask to calculate the mean under - mask_weight: float - weight of the mask, usually calculated as mask.sum() - - Returns - ------- - mean: cpt.NDArray[float] - array with local means under the mask - """ - return irfftn( - volume_rft * rfftn(mask).conj(), s=mask.shape - ).real / mask_weight - - """Update scores and angles if a new maximum is found.""" update_results_kernel = cp.ElementwiseKernel( 'float32 scores, float32 ccc_map, float32 angle_id', From 4a25e9980ad3581f854060c4962b95668fb2380c Mon Sep 17 00:00:00 2001 From: McHaillet Date: Tue, 16 Apr 2024 17:34:31 +0200 Subject: [PATCH 02/14] shift roi_mask to correct for rolling at the end --- src/pytom_tm/matching.py | 22 +++++++++++++++------- 1 file changed, 15 insertions(+), 7 deletions(-) diff --git a/src/pytom_tm/matching.py b/src/pytom_tm/matching.py index 879bf4a6..e95414d5 100644 --- a/src/pytom_tm/matching.py +++ b/src/pytom_tm/matching.py @@ -166,6 +166,7 @@ def run(self) -> tuple[npt.NDArray[float], npt.NDArray[float], dict]: cxv, cyv, czv = sxv // 2, syv // 2, szv // 2 # create slice for padding + shift = cp.floor(cp.array(self.plan.scores.shape) / 2).astype(int) + 1 pad_index = ( slice(cxv - cxt, cxv + cxt + mx), slice(cyv - cyt, cyv + cyt + my), @@ -173,7 +174,10 @@ def run(self) -> tuple[npt.NDArray[float], npt.NDArray[float], dict]: ) # calculate roi size - roi_size = self.plan.scores[self.stats_roi].size + roi_mask = cp.zeros(self.plan.scores.shape, dtype=bool) + roi_mask[self.stats_roi] = True + roi_mask = cp.flip(cp.roll(roi_mask, -shift, (0, 1, 2))) + roi_size = self.plan.scores[roi_mask].size if self.mask_is_spherical: # Then we only need to calculate std volume once self.plan.mask_padded[pad_index] = self.plan.mask @@ -247,15 +251,14 @@ def run(self) -> tuple[npt.NDArray[float], npt.NDArray[float], dict]: ) self.stats['variance'] += ( - square_sum_kernel( - self.plan.ccc_map[self.stats_roi] + masked_square_sum( + self.plan.ccc_map, roi_mask ) / roi_size ) # get correct orientation back - center = cp.floor(cp.array(self.plan.scores.shape) / 2).astype(int) + 1 - self.plan.scores = cp.roll(cp.flip(self.plan.scores), center, axis=(0, 1, 2)) - self.plan.angles = cp.roll(cp.flip(self.plan.angles), center, axis=(0, 1, 2)) + self.plan.scores = cp.roll(cp.flip(self.plan.scores), shift, axis=(0, 1, 2)) + self.plan.angles = cp.roll(cp.flip(self.plan.angles), shift, axis=(0, 1, 2)) self.stats['search_space'] = int(roi_size * len(self.angle_ids)) self.stats['variance'] = float(self.stats['variance'] / len(self.angle_ids)) @@ -314,6 +317,11 @@ def std_under_mask_convolution( ) +@cp.fuse() +def masked_square_sum(x, i): + return cp.sum(x * x * i) + + # Temporary workaround for ReductionKernel issue in cupy 13.0.0 (see: https://github.com/cupy/cupy/issues/8184) if version.parse(cp.__version__) == version.parse('13.0.0'): def square_sum_kernel(x): @@ -322,7 +330,7 @@ def square_sum_kernel(x): """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 - 'T y', # output params + 'T z', # output params 'x * x', # pre-processing expression 'a + b', # reduction operation 'y = a', # post-reduction output processing From 97cd60a9a39c20802e7449f473e9d12c776f4f82 Mon Sep 17 00:00:00 2001 From: McHaillet Date: Tue, 16 Apr 2024 17:39:20 +0200 Subject: [PATCH 03/14] switch to new kernel --- src/pytom_tm/matching.py | 31 +++++-------------------------- 1 file changed, 5 insertions(+), 26 deletions(-) diff --git a/src/pytom_tm/matching.py b/src/pytom_tm/matching.py index e95414d5..c59812b3 100644 --- a/src/pytom_tm/matching.py +++ b/src/pytom_tm/matching.py @@ -166,14 +166,14 @@ def run(self) -> tuple[npt.NDArray[float], npt.NDArray[float], dict]: cxv, cyv, czv = sxv // 2, syv // 2, szv // 2 # create slice for padding - shift = cp.floor(cp.array(self.plan.scores.shape) / 2).astype(int) + 1 pad_index = ( slice(cxv - cxt, cxv + cxt + mx), slice(cyv - cyt, cyv + cyt + my), slice(czv - czt, czv + czt + mz), ) - # calculate roi size + # calculate roi mask + shift = cp.floor(cp.array(self.plan.scores.shape) / 2).astype(int) + 1 roi_mask = cp.zeros(self.plan.scores.shape, dtype=bool) roi_mask[self.stats_roi] = True roi_mask = cp.flip(cp.roll(roi_mask, -shift, (0, 1, 2))) @@ -250,11 +250,7 @@ def run(self) -> tuple[npt.NDArray[float], npt.NDArray[float], dict]: self.plan.angles ) - self.stats['variance'] += ( - masked_square_sum( - self.plan.ccc_map, roi_mask - ) / roi_size - ) + self.stats['variance'] += masked_square_sum(self.plan.ccc_map, roi_mask, roi_size) # get correct orientation back self.plan.scores = cp.roll(cp.flip(self.plan.scores), shift, axis=(0, 1, 2)) @@ -318,22 +314,5 @@ def std_under_mask_convolution( @cp.fuse() -def masked_square_sum(x, i): - return cp.sum(x * x * i) - - -# Temporary workaround for ReductionKernel issue in cupy 13.0.0 (see: https://github.com/cupy/cupy/issues/8184) -if version.parse(cp.__version__) == version.parse('13.0.0'): - def square_sum_kernel(x): - return (x ** 2).sum() -else: - """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 - 'T z', # output params - 'x * x', # pre-processing expression - 'a + b', # reduction operation - 'y = a', # post-reduction output processing - '0', # identity value - 'variance' # kernel name - ) +def masked_square_sum(x, i, n): + return cp.sum(x * x * i) / n From 99cc4e53d53d54b757fed247314ddea7464994fe Mon Sep 17 00:00:00 2001 From: McHaillet Date: Tue, 16 Apr 2024 17:47:43 +0200 Subject: [PATCH 04/14] add some docs and refs --- src/pytom_tm/matching.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/src/pytom_tm/matching.py b/src/pytom_tm/matching.py index c59812b3..b944bd10 100644 --- a/src/pytom_tm/matching.py +++ b/src/pytom_tm/matching.py @@ -252,7 +252,11 @@ def run(self) -> tuple[npt.NDArray[float], npt.NDArray[float], dict]: self.stats['variance'] += masked_square_sum(self.plan.ccc_map, roi_mask, roi_size) - # get correct orientation back + # Get correct orientation back! + # Use same method as William Wan's STOPGAP (https://doi.org/10.1107/S205979832400295X): + # the search volume is Fourier transformed and conjugated before the iterations + # this means the eventual score map needs to be flipped back. The map is also rolled due to the ftshift + # effect of a Fourier space correlation function. self.plan.scores = cp.roll(cp.flip(self.plan.scores), shift, axis=(0, 1, 2)) self.plan.angles = cp.roll(cp.flip(self.plan.angles), shift, axis=(0, 1, 2)) From ea53a530adeca1fe445a609931826a3dcb39ff30 Mon Sep 17 00:00:00 2001 From: McHaillet Date: Tue, 30 Apr 2024 12:05:56 +0200 Subject: [PATCH 05/14] square_sum_kernel --- src/pytom_tm/matching.py | 20 ++++++++++++++++---- 1 file changed, 16 insertions(+), 4 deletions(-) diff --git a/src/pytom_tm/matching.py b/src/pytom_tm/matching.py index b944bd10..2d9e9fbc 100644 --- a/src/pytom_tm/matching.py +++ b/src/pytom_tm/matching.py @@ -250,7 +250,7 @@ def run(self) -> tuple[npt.NDArray[float], npt.NDArray[float], dict]: self.plan.angles ) - self.stats['variance'] += masked_square_sum(self.plan.ccc_map, roi_mask, roi_size) + self.stats['variance'] += (square_sum_kernel(self.plan.ccc_map * roi_mask) / roi_size) # Get correct orientation back! # Use same method as William Wan's STOPGAP (https://doi.org/10.1107/S205979832400295X): @@ -317,6 +317,18 @@ def std_under_mask_convolution( ) -@cp.fuse() -def masked_square_sum(x, i, n): - return cp.sum(x * x * i) / n +# Temporary workaround for ReductionKernel issue in cupy 13.0.0 (see: https://github.com/cupy/cupy/issues/8184) +if version.parse(cp.__version__) == version.parse('13.0.0'): + def square_sum_kernel(x): + return (x ** 2).sum() +else: + """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 + 'T y', # output params + 'x * x', # pre-processing expression + 'a + b', # reduction operation + 'y = a', # post-reduction output processing + '0', # identity value + 'variance' # kernel name + ) \ No newline at end of file From f3f80ec46effc718184d86a89df1b1c61693ca25 Mon Sep 17 00:00:00 2001 From: McHaillet Date: Tue, 30 Apr 2024 12:37:07 +0200 Subject: [PATCH 06/14] line fix --- src/pytom_tm/matching.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/src/pytom_tm/matching.py b/src/pytom_tm/matching.py index 2d9e9fbc..1c1e6f36 100644 --- a/src/pytom_tm/matching.py +++ b/src/pytom_tm/matching.py @@ -250,7 +250,9 @@ def run(self) -> tuple[npt.NDArray[float], npt.NDArray[float], dict]: self.plan.angles ) - self.stats['variance'] += (square_sum_kernel(self.plan.ccc_map * roi_mask) / roi_size) + self.stats['variance'] += ( + square_sum_kernel(self.plan.ccc_map * roi_mask) / roi_size + ) # Get correct orientation back! # Use same method as William Wan's STOPGAP (https://doi.org/10.1107/S205979832400295X): @@ -331,4 +333,4 @@ def square_sum_kernel(x): 'y = a', # post-reduction output processing '0', # identity value 'variance' # kernel name - ) \ No newline at end of file + ) From 1fbbfb6ea8ecbe969ac558fa189e61118890e742 Mon Sep 17 00:00:00 2001 From: McHaillet Date: Fri, 3 May 2024 20:38:02 +0200 Subject: [PATCH 07/14] add refs to other codebases --- src/pytom_tm/matching.py | 14 +++++++++++--- 1 file changed, 11 insertions(+), 3 deletions(-) diff --git a/src/pytom_tm/matching.py b/src/pytom_tm/matching.py index 1c1e6f36..29c9dba3 100644 --- a/src/pytom_tm/matching.py +++ b/src/pytom_tm/matching.py @@ -95,6 +95,13 @@ def __init__( ): """Initialize a template matching run. + For other great implementations see: + - STOPGAP: https://github.com/wan-lab-vanderbilt/STOPGAP + - pyTME: https://github.com/KosinskiLab/pyTME + + The precalculation of conjugated FTs of the tomo was (AFAIK) introduced + by STOPGAP! + Parameters ---------- job_id: str @@ -255,9 +262,10 @@ def run(self) -> tuple[npt.NDArray[float], npt.NDArray[float], dict]: ) # Get correct orientation back! - # Use same method as William Wan's STOPGAP (https://doi.org/10.1107/S205979832400295X): - # the search volume is Fourier transformed and conjugated before the iterations - # this means the eventual score map needs to be flipped back. The map is also rolled due to the ftshift + # Use same method as William Wan's STOPGAP + # (https://doi.org/10.1107/S205979832400295X): the search volume is Fourier + # transformed and conjugated before the iterations this means the eventual + # score map needs to be flipped back. The map is also rolled due to the ftshift # effect of a Fourier space correlation function. self.plan.scores = cp.roll(cp.flip(self.plan.scores), shift, axis=(0, 1, 2)) self.plan.angles = cp.roll(cp.flip(self.plan.angles), shift, axis=(0, 1, 2)) From 6cd38246e429fad661c65f2964d502ed949607e0 Mon Sep 17 00:00:00 2001 From: Marten Chaillet <58044494+McHaillet@users.noreply.github.com> Date: Tue, 7 May 2024 14:49:07 +0200 Subject: [PATCH 08/14] Update src/pytom_tm/matching.py Co-authored-by: Sander Roet --- src/pytom_tm/matching.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/src/pytom_tm/matching.py b/src/pytom_tm/matching.py index 29c9dba3..810deaa2 100644 --- a/src/pytom_tm/matching.py +++ b/src/pytom_tm/matching.py @@ -37,8 +37,9 @@ def __init__( """ # Search volume + and fft transform plan for the volume volume_shape = volume.shape - self.volume_rft_conj = rfftn(cp.asarray(volume, dtype=cp.float32, order='C')).conj() - self.volume_sq_rft_conj = rfftn(cp.asarray(volume, dtype=cp.float32, order='C') ** 2).conj() + cp_vol = cp.asarray(volume, dtype=cp.float32, order='C') + self.volume_rft_conj = rfftn(cp_vol)).conj() + self.volume_sq_rft_conj = rfftn(cp_vol ** 2).conj() # Explicit fft plan is no longer necessary as cupy generates a plan behind the scene which leads to # comparable timings From 5138ab85b0e61a2bf52bc7cb015aefb12f1c7b72 Mon Sep 17 00:00:00 2001 From: Marten Chaillet <58044494+McHaillet@users.noreply.github.com> Date: Tue, 7 May 2024 14:55:00 +0200 Subject: [PATCH 09/14] remove bracket --- src/pytom_tm/matching.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/pytom_tm/matching.py b/src/pytom_tm/matching.py index 810deaa2..6d6372e8 100644 --- a/src/pytom_tm/matching.py +++ b/src/pytom_tm/matching.py @@ -38,7 +38,7 @@ def __init__( # Search volume + and fft transform plan for the volume volume_shape = volume.shape cp_vol = cp.asarray(volume, dtype=cp.float32, order='C') - self.volume_rft_conj = rfftn(cp_vol)).conj() + self.volume_rft_conj = rfftn(cp_vol).conj() self.volume_sq_rft_conj = rfftn(cp_vol ** 2).conj() # Explicit fft plan is no longer necessary as cupy generates a plan behind the scene which leads to # comparable timings From 9aed81e30e86de9b6ddb325cd2ac4d65fc50b636 Mon Sep 17 00:00:00 2001 From: McHaillet Date: Tue, 7 May 2024 18:38:21 +0200 Subject: [PATCH 10/14] update version --- pyproject.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pyproject.toml b/pyproject.toml index cc3bb372..e069be96 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta" [project] name = "pytom-match-pick" -version = "0.6.1" +version = "0.7.0" description = "PyTOM's GPU template matching module as an independent package" readme = "README.md" license = {file = "LICENSE"} From 8dfda3df90c4a4179d88ce98d13a6d2a0981c083 Mon Sep 17 00:00:00 2001 From: McHaillet Date: Tue, 7 May 2024 18:40:39 +0200 Subject: [PATCH 11/14] add matching test for non spherical mask --- tests/test_template_matching.py | 47 ++++++++++++++++++++++++++++++--- 1 file changed, 44 insertions(+), 3 deletions(-) diff --git a/tests/test_template_matching.py b/tests/test_template_matching.py index 9573619e..ed695404 100644 --- a/tests/test_template_matching.py +++ b/tests/test_template_matching.py @@ -17,7 +17,7 @@ def setUp(self): self.gpu_id = 'gpu:0' self.angles = load_angle_list(files('pytom_tm.angle_lists').joinpath('angles_38.53_256.txt')) - def test_search(self): + def test_search_spherical_mask(self): angle_id = 100 rotation = self.angles[angle_id] loc = (77, 26, 40) @@ -31,8 +31,49 @@ def test_search(self): device='cpu' ) - tm = TemplateMatchingGPU(0, 0, self.volume, self.template, self.mask, self.angles, list(range(len( - self.angles)))) + tm = TemplateMatchingGPU( + 0, + 0, + self.volume, + self.template, + self.mask, + self.angles, + list(range(len(self.angles))), + ) + 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.005175, places=5, + msg='Standard deviation of the search should be almost equal') + + def test_search_non_spherical_mask(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))), + mask_is_spherical=True, + ) score_volume, angle_volume, stats = tm.run() ind = np.unravel_index(score_volume.argmax(), self.volume.shape) From e453b1e690234327447acbd986f46d662f48f3cb Mon Sep 17 00:00:00 2001 From: Marten Chaillet <58044494+McHaillet@users.noreply.github.com> Date: Tue, 7 May 2024 19:02:50 +0200 Subject: [PATCH 12/14] mask should be non spherical :[ --- tests/test_template_matching.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/test_template_matching.py b/tests/test_template_matching.py index ed695404..7b176ce4 100644 --- a/tests/test_template_matching.py +++ b/tests/test_template_matching.py @@ -72,7 +72,7 @@ def test_search_non_spherical_mask(self): self.mask, self.angles, list(range(len(self.angles))), - mask_is_spherical=True, + mask_is_spherical=False, ) score_volume, angle_volume, stats = tm.run() From a2b28c36985ad88c717e3ee6b21bf76b70b9ce0d Mon Sep 17 00:00:00 2001 From: Marten Chaillet <58044494+McHaillet@users.noreply.github.com> Date: Tue, 7 May 2024 19:09:39 +0200 Subject: [PATCH 13/14] Update test_template_matching.py --- tests/test_template_matching.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/test_template_matching.py b/tests/test_template_matching.py index 7b176ce4..d6960949 100644 --- a/tests/test_template_matching.py +++ b/tests/test_template_matching.py @@ -81,5 +81,5 @@ def test_search_non_spherical_mask(self): 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.005175, places=5, + self.assertAlmostEqual(stats['std'], 0.005175, places=4, msg='Standard deviation of the search should be almost equal') From 724392b1b9315d4b57448d80e6724b7ac3e0a918 Mon Sep 17 00:00:00 2001 From: Marten Chaillet <58044494+McHaillet@users.noreply.github.com> Date: Tue, 7 May 2024 19:14:57 +0200 Subject: [PATCH 14/14] set coverage to 80 --- .github/workflows/coverage.yml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.github/workflows/coverage.yml b/.github/workflows/coverage.yml index 2e693b4e..a00fe57b 100644 --- a/.github/workflows/coverage.yml +++ b/.github/workflows/coverage.yml @@ -44,7 +44,7 @@ jobs: path: coverage.xml repo_token: ${{ secrets.GITHUB_TOKEN }} pull_request_number: ${{ steps.get-pr.outputs.PR }} - minimum_coverage: 79 + minimum_coverage: 80 show_missing: True fail_below_threshold: True link_missing_lines: True