diff --git a/src/pytom_tm/tmjob.py b/src/pytom_tm/tmjob.py index f6b2dc0f..ef5415e1 100644 --- a/src/pytom_tm/tmjob.py +++ b/src/pytom_tm/tmjob.py @@ -13,7 +13,7 @@ from pytom_tm.matching import TemplateMatchingGPU from pytom_tm.weights import ( create_wedge, - power_spectrum_profile, + estimate_whitening_filter, profile_to_weighting, create_gaussian_band_pass, ) @@ -383,6 +383,12 @@ def __init__( self.template_shape = meta_data_template["shape"] self.mask_shape = meta_data_mask["shape"] + if len(set(self.template_shape)) != 1: + raise ValueError( + "Template is not cubic, all dimensions should be equal. " + f"Found template shape: {self.template_shape}." + ) + if self.template_shape != self.mask_shape: raise ValueError( "Template and mask have a different shape in pixels. " @@ -514,19 +520,17 @@ def __init__( ) if self.whiten_spectrum and not job_loaded_for_extraction: logging.info("Estimating whitening filter...") - weights = 1 / np.sqrt( - power_spectrum_profile( - read_mrc(self.tomogram)[ - self.search_origin[0] : self.search_origin[0] - + self.search_size[0], - self.search_origin[1] : self.search_origin[1] - + self.search_size[1], - self.search_origin[2] : self.search_origin[2] - + self.search_size[2], - ] - ) + patch_size = min(max(self.template_shape[0], 64), min(self.search_size)) + _, weights = estimate_whitening_filter( + tomogram=read_mrc(self.tomogram)[ + self.search_origin[0] : self.search_origin[0] + self.search_size[0], + self.search_origin[1] : self.search_origin[1] + self.search_size[1], + self.search_origin[2] : self.search_origin[2] + self.search_size[2], + ], + ts_metadata=self.ts_metadata, + patch_size=patch_size, + voxel_size=self.voxel_size, ) - weights /= weights.max() # scale to 1 np.save(self.whitening_filter, weights) # phase randomization options diff --git a/src/pytom_tm/weights.py b/src/pytom_tm/weights.py index c45577fe..a727d636 100644 --- a/src/pytom_tm/weights.py +++ b/src/pytom_tm/weights.py @@ -2,6 +2,7 @@ import numpy.typing as npt import scipy.ndimage as ndimage import voltools as vt +from typing import Generator from pytom_tm.io import UnequalSpacingError from pytom_tm.dataclass import CtfData, TiltSeriesMetaData from itertools import pairwise @@ -84,7 +85,7 @@ def wavelength_ev2m(voltage: float) -> float: def radial_grid( - shape: tuple[int, int, int] | tuple[int, int], + shape: tuple[int, int, int] | tuple[int, int] | tuple[int], reduced: bool = True, fftshifted: bool = False, shape_is_reduced: bool = False, @@ -100,8 +101,8 @@ def radial_grid( Parameters ---------- - shape: Union[tuple[int, int, int], tuple[int, int]] - 2D/3D input shape, usually the .shape attribute of a numpy array + shape: Union[tuple[int, int, int], tuple[int, int], tuple[int]] + 1D/2D/3D input shape, usually the .shape attribute of a numpy array reduced: bool, default True whether the last dimension should be reduced to shape[-1] // 2 + 1, as in the output of numpy.fft.rfftn @@ -117,8 +118,8 @@ def radial_grid( radial_grid: npt.NDArray[float] fourier space frequency grid """ - if len(shape) not in [2, 3]: - raise ValueError("radial_grid() only works for 2D or 3D shapes") + if len(shape) not in {1, 2, 3}: + raise ValueError("radial_grid() only works for 1D, 2D or 3D shapes") def full_axis(n: int) -> npt.NDArray[float]: # magnitude of frequency for a full (non-reduced) axis: 0 in the center, @@ -137,14 +138,16 @@ def last_axis(n: int) -> npt.NDArray[float]: y = full_axis(shape[1])[:, np.newaxis] z = last_axis(shape[2]) return np.sqrt(x**2 + y**2 + z**2) - else: + elif len(shape) == 2: x = full_axis(shape[0])[:, np.newaxis] y = last_axis(shape[1]) return np.sqrt(x**2 + y**2) + else: + return last_axis(shape[0]) def create_gaussian_low_pass( - shape: tuple[int, int, int] | tuple[int, int], + shape: tuple[int, int, int] | tuple[int, int] | tuple[int], spacing: float, resolution: float, ) -> npt.NDArray[float]: @@ -153,8 +156,8 @@ def create_gaussian_low_pass( Parameters ---------- - shape: Union[tuple[int, int, int], tuple[int, int]] - shape tuple with x,y or x,y,z dimension + shape: Union[tuple[int, int, int], tuple[int, int], tuple[int]] + shape tuple with x,y,z or x,y or x dimension spacing: float voxel size in real space resolution: float @@ -175,7 +178,7 @@ def create_gaussian_low_pass( def create_gaussian_high_pass( - shape: tuple[int, int, int] | tuple[int, int], + shape: tuple[int, int, int] | tuple[int, int] | tuple[int], spacing: float, resolution: float, ) -> npt.NDArray[float]: @@ -184,8 +187,8 @@ def create_gaussian_high_pass( Parameters ---------- - shape: Union[tuple[int, int, int], tuple[int, int]] - shape tuple with x,y or x,y,z dimension + shape: Union[tuple[int, int, int], tuple[int, int], tuple[int]] + shape tuple with x,y,z or x,y or x dimension spacing: float voxel size in real space resolution: float @@ -206,7 +209,7 @@ def create_gaussian_high_pass( def create_gaussian_band_pass( - shape: tuple[int, int, int] | tuple[int, int], + shape: tuple[int, int, int] | tuple[int, int] | tuple[int], spacing: float, low_pass: float | None = None, high_pass: float | None = None, @@ -218,8 +221,8 @@ def create_gaussian_band_pass( Parameters ---------- - shape: Union[tuple[int, int, int], tuple[int, int]] - shape tuple with x,y or x,y,z dimension + shape: Union[tuple[int, int, int], tuple[int, int], tuple[int]] + shape tuple with x,y,z or x,y or x dimension spacing: float voxel size in real space low_pass: Optional[float], default None @@ -697,64 +700,180 @@ def create_ctf( return ctf -def radial_average( - weights: npt.NDArray[float], +def estimate_whitening_filter( + tomogram: npt.NDArray[float], + ts_metadata: TiltSeriesMetaData, + patch_size: int, + overlap: float = 0.5, + reject_frac: float = 0.10, + statistic: str = "median", + voxel_size: float = 1.0, ) -> tuple[npt.NDArray[float], npt.NDArray[float]]: - """This calculates the radial average of a reduced fourier space function. + """Estimate a whitening filter from the radially averaged noise power spectrum + of a tomogram, sampled on overlapping cubic patches. - Parameters - ---------- - weights: npt.NDArray[float] - 3D array to be radially averaged: in fourier reduced form and with the origin - in the corner. - - Returns - ------- - (q, mean): tuple[npt.NDArray[float], npt.NDArray[float]] - A tuple of two 1d numpy arrays. Their length equals half of largest - input dimension. - """ - if len(weights.shape) not in [2, 3]: - raise ValueError("Radial average calculation only works for 2d/3d arrays") - - # get the number of sampling points from the largest fourier dimension, - # unless the reduced dimensions is already the largest one - sampling_points = max(max(weights.shape[:-1]) // 2 + 1, weights.shape[-1]) - - q = np.arange(sampling_points) - q_grid = np.floor( - # convert to radial indices in the fourier power spectrum, - # 0.5 is added to obtain the correct ring - radial_grid(weights.shape, shape_is_reduced=True) * (sampling_points - 1) + 0.5 - ).astype(int) - mean = ndimage.mean(weights, labels=q_grid, index=q) - - return q, mean - - -def power_spectrum_profile(image: npt.NDArray[float]) -> npt.NDArray[float]: - """Calculate the power spectrum for a real space array and then find the profile - (radial average). + The power spectrum is averaged only over Fourier voxels the tilt series + actually sampled (via a tilt coverage mask), so the estimate is not biased by + the empty missing-wedge / inter-tilt-gap voxels. Estimation is done on many + overlapping, windowed, mean-subtracted patches, and the highest-variance + patches (e.g. gold, carbon, ice, volume edges) are rejected before averaging. Parameters ---------- - image: npt.NDArray[float] - 2D/3D array in real space for which the power spectrum profile needs to be - calculated + tomogram: npt.NDArray[float] + 3D real-space array to estimate the whitening filter from + ts_metadata: TiltSeriesMetaData + tilt series metadata of the tomogram, used to build the coverage mask that + restricts the radial average to Fourier voxels the tilt series actually + sampled. CTF and dose weighting are not applied to this mask, only the + tilt geometry matters here, regardless of what ts_metadata specifies + patch_size: int + edge length of the cubic estimation box, also the box the tilt coverage + mask is built for. A larger patch reaches lower frequencies but leaves + fewer patches to average over in a thin tomogram. The estimated profile + can be interpolated to a different box size afterwards, so estimation + size and matching size are decoupled + overlap: float, default 0.5 + fractional overlap between patches + reject_frac: float, default 0.10 + fraction of the highest-variance patches to reject, 0 disables rejection + statistic: str, default "median" + radial averaging statistic, either "median" (robust, bias-corrected to + the mean) or "mean" + voxel_size: float, default 1.0 + voxel size in Angstrom, sets the physical units of the returned frequency + axis Returns ------- - profile: npt.NDArray[float] - A 1D numpy array + (q, w): tuple[npt.NDArray[float], npt.NDArray[float]] + q is the frequency of each shell in cycles/Angstrom (0 to Nyquist), with + shape (patch_size // 2 + 1,). w is the whitening filter derived from the + power spectrum profile (DC set to 0, high frequencies tapered, and + normalized to a maximum of 1), with the same shape as q """ - if len(image.shape) not in [2, 3]: - raise ValueError( - "Power spectrum profile calculation only works for 2d/3d arrays." - ) + tomogram = np.asarray(tomogram, dtype=np.float32) + length = int(patch_size) + + # tilt coverage mask, built once for the cubic patch (geometry identical per + # patch). CTF and dose data are stripped from the metadata so only the tilt + # geometry determines which Fourier voxels count as sampled. + mask_metadata = ts_metadata.replace(ctf_data=None, dose_accumulation=None) + wedge = create_wedge((length, length, length), mask_metadata, voxel_size) + mask = wedge > 0.05 # binarize the mask and set DC to zero + mask[0, 0, 0] = 0 + + # window (mandatory: controls leakage across the 3+ decade dynamic range) + win = _hann3d(length) + win_power = float((win**2).sum()) + + # collect windowed periodograms over overlapping patches + step = max(1, int(length * (1.0 - overlap) + 0.5)) # round half up + psds, variances = [], [] + for sl in _patch_slices(tomogram.shape, length, step): + v = tomogram[sl] + if not np.all(np.isfinite(v)): + continue + v = v - v.mean() + f = np.fft.rfftn(v * win) + psds.append((f.real**2 + f.imag**2) / win_power) + variances.append(float(v.var())) + + # reject high-variance patches (gold / carbon / ice / edges). Only attempted + # with enough patches for the quantile estimate to be meaningful. + psds = np.asarray(psds) + variances = np.asarray(variances) + keep = np.ones(len(psds), bool) + if reject_frac > 0 and len(psds) >= 10: + keep = variances <= np.quantile(variances, 1.0 - reject_frac) + + if keep.sum() == 0: + raise RuntimeError("no usable patches for whitening filter estimation") + + psd = psds[keep].mean(axis=0) + dof = int(keep.sum()) + + # masked radial average + q, prof = _masked_radial(psd, mask, length, voxel_size, statistic, dof) + + def cosine_cutoff(): + r_frac = radial_grid((length,)) + lo, hi = 0.9, 1.0 + ramp = np.clip((r_frac - lo) / (hi - lo), 0, 1) + return 0.5 * (1 + np.cos(np.pi * ramp)) # 1 below lo, 0 at hi + + # transform into a whitening filter + w = np.where(prof > 0, 1 / np.sqrt(prof), 0.0) + w[0] = 0.0 # zero DC because its not needed + w = w * cosine_cutoff() # tamper high frequency estimate + w /= w.max() + + return q, w + + +def _hann3d(length: int) -> npt.NDArray[float]: + w = np.hanning(length + 2)[1:-1].astype(np.float32) # drop the exact zeros + return (w[:, None, None] * w[None, :, None] * w[None, None, :]).astype(np.float32) + + +def _patch_slices( + shape: tuple[int, int, int], length: int, step: int +) -> Generator[tuple[slice, slice, slice], None, None]: + def starts(n): + s = list(range(0, n - length + 1, step)) + if not s: + raise ValueError(f"patch_size {length} larger than tomogram extent {n}") + if s[-1] != n - length: + s.append(n - length) + return s + + xs, ys, zs = (starts(n) for n in shape) + for x in xs: + for y in ys: + for z in zs: + yield ( + slice(x, x + length), + slice(y, y + length), + slice(z, z + length), + ) - _, power_profile = radial_average(np.abs(np.fft.rfftn(image)) ** 2) - return power_profile +def _masked_radial( + psd: npt.NDArray[float], + mask: npt.NDArray[bool], + length: int, + voxel_size: float, + statistic: str, + dof: int, +) -> tuple[npt.NDArray[float], npt.NDArray[float]]: + r_frac = radial_grid((length, length, length)) + + k_nyq = 0.5 / voxel_size + nb = length // 2 + 1 + shell = np.floor(r_frac * (nb - 1) + 0.5).astype(int) + q = np.arange(nb) / (nb - 1) * k_nyq + + labels = shell.copy() + labels[(shell >= nb) | (shell < 0)] = -1 + if mask is not None: + labels[~mask] = -1 + + idx = np.arange(nb) + counts = ndimage.sum(np.ones_like(psd), labels=labels, index=idx) + fn = ndimage.median if statistic == "median" else ndimage.mean + prof = np.asarray(fn(psd, labels=labels, index=idx), dtype=float) + prof[counts == 0] = np.nan # empty shells -> NaN, not 0 + + if statistic == "median": + prof /= (1.0 - 1.0 / (9.0 * max(dof, 1))) ** 3 # median -> mean, Gamma(dof) + + good = np.isfinite(prof) & (prof > 0) + if good.sum() < 2: + raise RuntimeError( + "too few valid radial shells to interpolate a whitening filter profile" + ) + prof = np.interp(q, q[good], prof[good]) + return q, prof def profile_to_weighting( @@ -774,7 +893,7 @@ def profile_to_weighting( Returns ------- weighting: npt.NDArray[float] - Reduced Fourier space weighting for shape + Reduced Fourier space weighting for shape, with the DC component set to 0 """ if len(profile.shape) != 1: raise ValueError("Profile passed to profile_to_weighting is not 1-dimensional.") @@ -784,9 +903,13 @@ def profile_to_weighting( q_grid = radial_grid(shape) weights = ndimage.map_coordinates( - profile, q_grid.flatten()[np.newaxis, :] * profile.shape[0], order=1 + profile, + q_grid.flatten()[np.newaxis, :] * (profile.shape[0] - 1), + order=1, + mode="nearest", ).reshape(q_grid.shape) weights[q_grid > 1] = 0 + weights[(0,) * len(shape)] = 0 return weights diff --git a/tests/test_tmjob.py b/tests/test_tmjob.py index 6655769f..f95c4732 100644 --- a/tests/test_tmjob.py +++ b/tests/test_tmjob.py @@ -43,6 +43,8 @@ TEST_EXTRACTION_MASK_INT8 = TEST_DATA_DIR.joinpath("extraction_mask_int8.mrc") TEST_TEMPLATE = TEST_DATA_DIR.joinpath("template.mrc") TEST_TEMPLATE_UNEQUAL_SPACING = TEST_DATA_DIR.joinpath("template_unequal_spacing.mrc") +TEST_TEMPLATE_NOT_CUBIC = TEST_DATA_DIR.joinpath("template_not_cubic.mrc") +TEST_MASK_NOT_CUBIC = TEST_DATA_DIR.joinpath("mask_not_cubic.mrc") TEST_TEMPLATE_WRONG_VOXEL_SIZE = TEST_DATA_DIR.joinpath("template_voxel_error_test.mrc") TEST_MASK = TEST_DATA_DIR.joinpath("mask.mrc") TEST_MASK_WRONG_SIZE = TEST_DATA_DIR.joinpath("mask_wrong_size.mrc") @@ -112,6 +114,17 @@ def setUpClass(cls) -> None: write_mrc(TEST_MASK_WRONG_SIZE, mask_wrong_size, 1.0) write_mrc(TEST_TEMPLATE, template, 1.0) write_mrc(TEST_TEMPLATE_WRONG_VOXEL_SIZE, template, 1.5) + not_cubic_shape = (TEMPLATE_SIZE, TEMPLATE_SIZE, TEMPLATE_SIZE - 2) + write_mrc( + TEST_TEMPLATE_NOT_CUBIC, + np.zeros(not_cubic_shape, dtype=np.float32), + 1.0, + ) + write_mrc( + TEST_MASK_NOT_CUBIC, + np.zeros(not_cubic_shape, dtype=np.float32), + 1.0, + ) mrcfile.write( TEST_TEMPLATE_UNEQUAL_SPACING, template, voxel_size=(1.5, 1.0, 2.0) ) @@ -356,6 +369,22 @@ def test_tm_job_errors(self): voxel_size=1.0, tomogram_mask=TEST_WRONG_SIZE_TOMO_MASK, ) + # Test non-cubic template + with self.assertRaisesRegex( + ValueError, + "not cubic", + ): + TMJob( + "0", + 10, + TEST_TOMOGRAM, + TEST_TEMPLATE_NOT_CUBIC, + TEST_MASK_NOT_CUBIC, + TEST_DATA_DIR, + ts_metadata=TS_METADATA, + angle_increment=ANGULAR_SEARCH, + voxel_size=1.0, + ) # Test template mask mismatch with self.assertRaisesRegex( ValueError, @@ -542,20 +571,20 @@ def test_tm_job_weighting_options(self): angle_increment=90.00, voxel_size=1.0, whiten_spectrum=True, - search_y=[10, 90], + search_z=[0, 30], ) new_whitening_filter = np.load(TEST_WHITENING_FILTER) self.assertNotEqual( whitening_filter.shape, new_whitening_filter.shape, - msg="After reducing the search region along the largest dimension the " - "whitening filter should have less sampling points", + msg="After reducing the search region below the default estimation " + "patch size the whitening filter should have less sampling points", ) self.assertEqual( new_whitening_filter.shape, - (max(job.search_size) // 2 + 1,), + (min(64, min(job.search_size)) // 2 + 1,), msg="The whitening filter does not have the expected size, it should be " - "equal (x // 2) + 1, where x is the largest dimension of the search box.", + "equal (x // 2) + 1, where x is min(64, smallest search box dimension).", ) # TMJob with none of these weighting options is tested in all other runs diff --git a/tests/test_weights.py b/tests/test_weights.py index 711979c3..5865ef9c 100644 --- a/tests/test_weights.py +++ b/tests/test_weights.py @@ -3,11 +3,11 @@ from pytom_tm.weights import ( create_wedge, _create_symmetric_wedge, + _masked_radial, create_ctf, create_gaussian_band_pass, radial_grid, - radial_average, - power_spectrum_profile, + estimate_whitening_filter, profile_to_weighting, ) from pytom_tm.dataclass import CtfData, TiltSeriesMetaData @@ -25,6 +25,7 @@ def setUp(self): self.reduced_even_shape_3d = (10, 10, 6) self.reduced_even_shape_2d = (10, 6) + self.reduced_even_shape_1d = (6,) self.reduced_uneven_shape_3d = (11, 11, 6) self.reduced_uneven_shape_2d = (11, 6) self.reduced_irregular_shape_3d = (7, 12, 6 // 2 + 1) @@ -39,13 +40,7 @@ def test_radial_grid(self): with self.assertRaises( ValueError, msg="Radial grid should raise ValueError if the shape is " - "not 2- or 3-dimensional.", - ): - radial_grid((5,)) - with self.assertRaises( - ValueError, - msg="Radial grid should raise ValueError if the shape is " - "not 2- or 3-dimensional.", + "not 1-, 2- or 3-dimensional.", ): radial_grid((5,) * 4) @@ -59,6 +54,11 @@ def test_radial_grid(self): self.reduced_even_shape_2d, msg="2D radial grid does not have the correct shape", ) + self.assertEqual( + radial_grid(self.volume_shape_even[:1]).shape, + self.reduced_even_shape_1d, + msg="1D radial grid does not have the correct shape", + ) def test_band_pass(self): with self.assertRaises( @@ -128,6 +128,19 @@ def test_band_pass(self): msg="Low-pass and high-pass filter should be different", ) + # 1D band-pass, e.g. for filtering a radially averaged profile + band_pass_1d = create_gaussian_band_pass( + self.volume_shape_even[:1], + self.voxel_size, + self.low_pass, + self.high_pass, + ) + self.assertEqual( + band_pass_1d.shape, + self.reduced_even_shape_1d, + msg="1D bandpass filter does not have expected output shape", + ) + def test_create_symmetric_wedge(self): with self.assertRaisesRegex(ValueError, "bigger than 90 degrees"): _create_symmetric_wedge(self.volume_shape_even, 4, 1.0) @@ -322,62 +335,82 @@ def test_ctf(self): msg="CTF does not have expected output shape", ) - def test_radial_average(self): - x, y = 100, 50 + def test_estimate_whitening_filter(self): + patch_size = 32 + rng = np.random.default_rng(0) + tomogram = rng.normal(size=(64, 64, 64)).astype(np.float32) + with self.assertRaises( - ValueError, - msg="Radial average should raise error if something other than 2d/3d " - "array is provided.", + RuntimeError, + msg="Estimate whitening filter should raise RuntimeError if no " + "patches are usable.", ): - radial_average(np.zeros(x)) - q, m = radial_average(np.zeros((x, y))) + estimate_whitening_filter( + np.full((64, 64, 64), np.nan, dtype=np.float32), + self.ts_metadata, + patch_size=64, + voxel_size=self.voxel_size, + ) + + q, w = estimate_whitening_filter( + tomogram, self.ts_metadata, patch_size, voxel_size=self.voxel_size + ) + self.assertEqual( - m.shape[0], - x // 2 + 1, - msg="Radial average shape should equal largest sampling dimension.", + q.shape, + (patch_size // 2 + 1,), + msg="Frequency axis of the whitening filter does not have the expected " + "shape", ) - q, m = radial_average(np.zeros((30, y))) self.assertEqual( - m.shape[0], - y, - msg="Radial average shape should equal largest sampling dimension, " - "considering Fourier reduced form.", + w.shape, + (patch_size // 2 + 1,), + msg="Whitening filter does not have the expected shape", ) - q, m = radial_average(np.zeros((20, x, y))) self.assertEqual( - m.shape[0], - x // 2 + 1, - msg="Radial average shape should equal largest sampling dimension.", + w[0], + 0.0, + msg="Whitening filter should have the DC component set to 0", + ) + self.assertAlmostEqual( + w.max(), + 1.0, + msg="Whitening filter should be normalized to a maximum of 1", + ) + + # statistic="mean" should also run without error + q_mean, w_mean = estimate_whitening_filter( + tomogram, + self.ts_metadata, + patch_size, + voxel_size=self.voxel_size, + statistic="mean", ) - q, m = radial_average(np.zeros((20, 30, y))) self.assertEqual( - m.shape[0], - y, - msg="Radial average shape should equal largest sampling dimension, " - "considering Fourier reduced form.", + q_mean.shape, + (patch_size // 2 + 1,), + msg="Frequency axis of the whitening filter does not have the expected " + "shape when using the mean statistic", + ) + self.assertEqual( + w_mean.shape, + (patch_size // 2 + 1,), + msg="Whitening filter does not have the expected shape when using the " + "mean statistic", ) - def test_power_spectrum_profile(self): - with self.assertRaises( - ValueError, - msg="Power spectrum profile should raise ValueError if input image is " - "not 2- or 3-dimensional.", - ): - power_spectrum_profile(np.zeros(5)) + def test_masked_radial(self): + length = 8 + nb = length // 2 + 1 + psd = np.ones((length, length, nb)) + mask = np.zeros((length, length, nb), dtype=bool) + with self.assertRaises( - ValueError, - msg="Power spectrum profile should raise ValueError if input image is " - "not 2- or 3-dimensional.", + RuntimeError, + msg="_masked_radial should raise RuntimeError if fewer than 2 radial " + "shells have valid data.", ): - power_spectrum_profile(np.zeros((5,) * 4)) - profile = power_spectrum_profile(np.zeros(self.volume_shape_irregular)) - self.assertEqual( - profile.shape, - (max(self.volume_shape_irregular) // 2 + 1,), - msg="Power spectrum profile output shape should be a 1-dimensional array " - "with length equal to max(input_shape) // 2 + 1, corresponding to largest " - "sampling component in Fourier space.", - ) + _masked_radial(psd, mask, length, self.voxel_size, "median", 10) def test_profile_to_weighting(self): with self.assertRaises( @@ -399,7 +432,7 @@ def test_profile_to_weighting(self): ): profile_to_weighting(np.zeros(5), (5,) * 4) - profile = power_spectrum_profile(np.zeros(self.volume_shape_irregular)) + profile = np.ones(7) self.assertEqual( profile_to_weighting(profile, self.volume_shape_irregular).shape, self.reduced_irregular_shape_3d, @@ -410,3 +443,17 @@ def test_profile_to_weighting(self): self.reduced_irregular_shape_2d, msg="Profile to weighting should return 2D Fourier reduced array.", ) + + ramp_profile = np.arange(1, 7, dtype=float) # monotonic, non-zero everywhere + weights_3d = profile_to_weighting(ramp_profile, self.volume_shape_even) + self.assertEqual( + weights_3d[0, 0, 0], + 0, + msg="Profile to weighting should set the DC component to 0.", + ) + self.assertAlmostEqual( + weights_3d.max(), + ramp_profile.max(), + msg="Nyquist frequency should map to the last profile value instead of " + "rolling off towards 0.", + )