From e9aea0eb463eb49418820eaecf612a61dfb5cb1d Mon Sep 17 00:00:00 2001 From: McHaillet Date: Wed, 15 Jul 2026 14:40:20 +0200 Subject: [PATCH 1/8] Rework whitening filter estimation - Support 1D shapes in radial_reduced_grid/create_gaussian_*_pass so the bandpass filter chain can be applied directly to a 1D radial profile. - Replace radial_average/power_spectrum_profile with estimate_whitening_filter, which estimates the noise power spectrum on overlapping, windowed, mean- subtracted patches, masked to only the Fourier voxels actually sampled by the tilt series (via create_wedge with CTF/dose stripped from the metadata), rejecting the highest-variance patches and returning an already-transformed whitening filter (1/sqrt, DC zeroed, high-frequency tapered, normalized). - Wire TMJob to call estimate_whitening_filter directly instead of manually transforming a raw power spectrum profile. - Fix an off-by-one in profile_to_weighting's map_coordinates interpolation that rolled the near-Nyquist weights off toward 0, and zero the DC component of the returned weighting. Co-Authored-By: Claude Sonnet 5 --- src/pytom_tm/tmjob.py | 24 ++-- src/pytom_tm/weights.py | 279 +++++++++++++++++++++++++++++----------- tests/test_tmjob.py | 10 +- tests/test_weights.py | 124 ++++++++++-------- 4 files changed, 290 insertions(+), 147 deletions(-) diff --git a/src/pytom_tm/tmjob.py b/src/pytom_tm/tmjob.py index f6b2dc0f..05512c96 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, ) @@ -514,19 +514,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(64, min(self.search_size)) + _, weights = estimate_whitening_filter( + 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], + ], + self.ts_metadata, + 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 126db52c..35e4fba8 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,8 @@ def wavelength_ev2m(voltage: float) -> float: def radial_reduced_grid( - shape: tuple[int, int, int] | tuple[int, int], shape_is_reduced: bool = False + shape: tuple[int, int, int] | tuple[int, int] | tuple[int], + shape_is_reduced: bool = False, ) -> npt.NDArray[float]: """Calculates a Fourier space radial reduced grid for the given input shape, with the 0 frequency in the center of the output image. Values range from 0 in the @@ -98,8 +100,8 @@ def radial_reduced_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 shape_is_reduced: bool, default False whether the shape is already in a reduced fourier format, False by default @@ -108,10 +110,12 @@ def radial_reduced_grid( radial_reduced_grid: npt.NDArray[float] fourier space frequency grid, 0 in center, 1 at nyquist """ - if len(shape) not in [2, 3]: - raise ValueError("radial_reduced_grid() only works for 2D or 3D shapes") + if len(shape) not in [1, 2, 3]: + raise ValueError("radial_reduced_grid() only works for 1D, 2D or 3D shapes") reduced_dim = shape[-1] if shape_is_reduced else shape[-1] // 2 + 1 - if len(shape) == 3: + if len(shape) == 1: + return np.arange(0, reduced_dim, 1.0) / (reduced_dim - 1) + elif len(shape) == 3: x = ( np.abs( np.arange( @@ -144,7 +148,7 @@ def radial_reduced_grid( 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 +157,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 @@ -171,11 +175,12 @@ def create_gaussian_low_pass( # then convert cutoff (hwhm) to sigma for gaussian function sigma_fourier = hwhm_to_sigma(2 * spacing / resolution) - return np.fft.ifftshift(np.exp(-(q**2) / (2 * sigma_fourier**2)), axes=(0, 1)) + low_pass = np.exp(-(q**2) / (2 * sigma_fourier**2)) + return low_pass if len(shape) == 1 else np.fft.ifftshift(low_pass, axes=(0, 1)) 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 +189,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 @@ -202,11 +207,12 @@ def create_gaussian_high_pass( # then convert cutoff (hwhm) to sigma for gaussian function sigma_fourier = hwhm_to_sigma(2 * spacing / resolution) - return np.fft.ifftshift(1 - np.exp(-(q**2) / (2 * sigma_fourier**2)), axes=(0, 1)) + high_pass = 1 - np.exp(-(q**2) / (2 * sigma_fourier**2)) + return high_pass if len(shape) == 1 else np.fft.ifftshift(high_pass, axes=(0, 1)) 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 +224,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 @@ -251,10 +257,11 @@ def create_gaussian_band_pass( sigma_high_pass = hwhm_to_sigma(2 * spacing / high_pass) sigma_low_pass = hwhm_to_sigma(2 * spacing / low_pass) - return np.fft.ifftshift( - (1 - np.exp(-(q**2) / (2 * sigma_high_pass**2))) - * np.exp(-(q**2) / (2 * sigma_low_pass**2)), - axes=(0, 1), + band_pass = (1 - np.exp(-(q**2) / (2 * sigma_high_pass**2))) * np.exp( + -(q**2) / (2 * sigma_low_pass**2) + ) + return ( + band_pass if len(shape) == 1 else np.fft.ifftshift(band_pass, axes=(0, 1)) ) @@ -705,70 +712,190 @@ def ctf_1d(q): return np.fft.ifftshift(ctf, axes=(0, 1) if len(shape) == 3 else 0) -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, + exclude: npt.NDArray[bool] | None = None, + 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. + """ + Masked radial noise power spectrum of a tomogram, estimated on cubic patches. + + Averages the power spectrum ONLY over Fourier voxels the tilt series actually + sampled (via the tilt 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; the highest-variance patches (gold, carbon, + ice, volume edges) are rejected. Parameters ---------- - weights: npt.NDArray[float] - 3D array to be radially averaged: in fourier reduced form and with the origin - in the corner. + tomogram: npt.NDArray[float] + 3D real-space array. + 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 mask is + built for. Sensible default ~64 for typical binned ET (~10 A/vox): reaches + the low frequencies a template box needs, is a clean FFT size, and still + fits several times in a thin z so you keep many patches. Drop to 48 if z + is tight; rise to 96 only for thick volumes with large targets. A masked + profile can be interpolated to a different correlation-box size afterwards, + so estimation size and matching size are decoupled -- pick this purely for + estimation quality. + overlap: float, default 0.5 + fractional overlap between patches (0.5 = 50%). + reject_frac: float, default 0.10 + drop this fraction of the highest-variance patches. 0 disables. + exclude: Optional[npt.NDArray[bool]], default None + optional bool array, same shape as tomogram, True where voxels must not be + used (gold / carbon / contamination). Do NOT mask out biological + structure -- under H0 the crowded cell IS the noise, and excluding it + biases P_n low and leaves the background sigma above 1. + statistic: str, default "median" + "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 ------- - (q, mean): tuple[npt.NDArray[float], npt.NDArray[float]] - A tuple of two 1d numpy arrays. Their length equals half of largest - input dimension. + (q, w): tuple[npt.NDArray[float], npt.NDArray[float]] + q is the frequency of each shell, cycles/Angstrom (0..Nyquist), with + shape (patch_size // 2 + 1,). w is the whitening filter (already + transformed from the power spectrum profile, DC-zeroed, high-frequency + tapered and normalized to a maximum of 1), with the same shape as q. """ - 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_reduced_grid(weights.shape, shape_is_reduced=True) - * (sampling_points - 1) - + 0.5 - ).astype(int) - mean = ndimage.mean( - np.fft.fftshift(weights, axes=(0, 1) if len(weights.shape) == 3 else 0), - labels=q_grid, - index=q, - ) + 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 + 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(round(length * (1.0 - overlap)))) + psds, variances = [], [] + for sl in _patch_slices(tomogram.shape, length, step): + if exclude is not None and exclude[sl].mean() > 0.01: + continue + 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())) + + if len(psds) < 4: + raise RuntimeError(f"only {len(psds)} usable patches; reduce patch_size") + + # reject high-variance patches (gold / carbon / ice / edges) + 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() < 4: + keep = np.ones(len(psds), bool) + + psd = np.mean([p for p, k in zip(psds, keep) if k], 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_reduced_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), np.zeros_like(prof)) + 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 + + zs, ys, xs = (starts(n) for n in shape) + for z in zs: + for y in ys: + for x in xs: + yield ( + slice(z, z + length), + slice(y, y + length), + slice(x, x + length), + ) - return q, mean +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]]: + kz = np.fft.fftfreq(length, d=voxel_size)[:, None, None] + ky = np.fft.fftfreq(length, d=voxel_size)[None, :, None] + kx = np.fft.rfftfreq(length, d=voxel_size)[None, None, :] + r = np.sqrt(kz**2 + ky**2 + kx**2) -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). + k_nyq = 0.5 / voxel_size + nb = length // 2 + 1 + shell = np.floor(r / k_nyq * (nb - 1) + 0.5).astype(int) + q = np.arange(nb) / (nb - 1) * k_nyq - Parameters - ---------- - image: npt.NDArray[float] - 2D/3D array in real space for which the power spectrum profile needs to be - calculated + labels = shell.copy() + labels[(shell >= nb) | (shell < 0)] = -1 + if mask is not None: + labels[~mask] = -1 - Returns - ------- - profile: npt.NDArray[float] - A 1D numpy array - """ - if len(image.shape) not in [2, 3]: - raise ValueError( - "Power spectrum profile calculation only works for 2d/3d arrays." - ) + 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 - _, power_profile = radial_average(np.abs(np.fft.rfftn(image)) ** 2) + if statistic == "median": + prof /= (1.0 - 1.0 / (9.0 * max(dof, 1))) ** 3 # median -> mean, Gamma(dof) - return power_profile + good = np.isfinite(prof) & (prof > 0) + if good.sum() >= 2: + prof = np.interp(q, q[good], prof[good]) + return q, prof def profile_to_weighting( @@ -788,7 +915,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.") @@ -798,9 +925,15 @@ def profile_to_weighting( q_grid = radial_reduced_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 - return np.fft.ifftshift(weights, axes=(0, 1) if len(shape) == 3 else 0) + weights = np.fft.ifftshift(weights, axes=(0, 1) if len(shape) == 3 else 0) + weights[(0,) * len(shape)] = 0 + + return weights diff --git a/tests/test_tmjob.py b/tests/test_tmjob.py index 6655769f..0a64d006 100644 --- a/tests/test_tmjob.py +++ b/tests/test_tmjob.py @@ -542,20 +542,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 8f7af5fd..245407ea 100644 --- a/tests/test_weights.py +++ b/tests/test_weights.py @@ -6,8 +6,7 @@ create_ctf, create_gaussian_band_pass, radial_reduced_grid, - radial_average, - power_spectrum_profile, + estimate_whitening_filter, profile_to_weighting, ) from pytom_tm.dataclass import CtfData, TiltSeriesMetaData @@ -25,6 +24,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 +39,7 @@ def test_radial_reduced_grid(self): with self.assertRaises( ValueError, msg="Radial reduced grid should raise ValueError if the shape is " - "not 2- or 3-dimensional.", - ): - radial_reduced_grid((5,)) - with self.assertRaises( - ValueError, - msg="Radial reduced grid should raise ValueError if the shape is " - "not 2- or 3-dimensional.", + "not 1-, 2- or 3-dimensional.", ): radial_reduced_grid((5,) * 4) @@ -59,6 +53,11 @@ def test_radial_reduced_grid(self): self.reduced_even_shape_2d, msg="2D radial reduced grid does not have the correct shape", ) + self.assertEqual( + radial_reduced_grid(self.volume_shape_even[:1]).shape, + self.reduced_even_shape_1d, + msg="1D radial reduced grid does not have the correct shape", + ) def test_band_pass(self): with self.assertRaises( @@ -128,6 +127,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) @@ -333,61 +345,47 @@ def test_ctf(self): "crossing", ) - 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 too few " + "patches are usable.", ): - radial_average(np.zeros(x)) - q, m = radial_average(np.zeros((x, y))) - self.assertEqual( - m.shape[0], - x // 2 + 1, - msg="Radial average shape should equal largest sampling dimension.", + estimate_whitening_filter( + tomogram, + 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 ) - 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.", + 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((20, x, y))) self.assertEqual( - m.shape[0], - x // 2 + 1, - msg="Radial average shape should equal largest sampling dimension.", + w.shape, + (patch_size // 2 + 1,), + msg="Whitening filter does not have the expected shape", ) - 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.", + w[0], + 0.0, + msg="Whitening filter should have the DC component set to 0", ) - - 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)) - 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,) * 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.", + self.assertAlmostEqual( + w.max(), + 1.0, + msg="Whitening filter should be normalized to a maximum of 1", ) def test_profile_to_weighting(self): @@ -410,7 +408,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, @@ -421,3 +419,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.", + ) From b489202975b09300f97f49a13718cdceb6d120eb Mon Sep 17 00:00:00 2001 From: McHaillet Date: Wed, 15 Jul 2026 15:04:34 +0200 Subject: [PATCH 2/8] Scale whitening filter patch size with template box Use max(template_shape[0], 64) instead of a fixed 64 so larger template boxes get a correspondingly larger estimation patch, still clamped to the smallest search region dimension. Co-Authored-By: Claude Sonnet 5 --- src/pytom_tm/tmjob.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/pytom_tm/tmjob.py b/src/pytom_tm/tmjob.py index 05512c96..0704de63 100644 --- a/src/pytom_tm/tmjob.py +++ b/src/pytom_tm/tmjob.py @@ -514,7 +514,7 @@ def __init__( ) if self.whiten_spectrum and not job_loaded_for_extraction: logging.info("Estimating whitening filter...") - patch_size = min(64, min(self.search_size)) + patch_size = min(max(self.template_shape[0], 64), min(self.search_size)) _, weights = estimate_whitening_filter( read_mrc(self.tomogram)[ self.search_origin[0] : self.search_origin[0] + self.search_size[0], From 1e9fd3153205aa5b7c9745a2e9f0d631a0baf451 Mon Sep 17 00:00:00 2001 From: McHaillet Date: Wed, 15 Jul 2026 15:41:33 +0200 Subject: [PATCH 3/8] Reuse radial_reduced_grid in _masked_radial's shell binning Replace the hand-rolled fftfreq/rfftfreq radial-distance computation with radial_reduced_grid, matching the normalized-frequency convention already used by the bandpass filters, CTF, and wedge in this module. This also fixes a small inconsistency for odd patch sizes, where the fftfreq-based Nyquist normalization diverged from radial_reduced_grid's convention. Co-Authored-By: Claude Sonnet 5 --- src/pytom_tm/weights.py | 9 ++++----- 1 file changed, 4 insertions(+), 5 deletions(-) diff --git a/src/pytom_tm/weights.py b/src/pytom_tm/weights.py index 35e4fba8..622de477 100644 --- a/src/pytom_tm/weights.py +++ b/src/pytom_tm/weights.py @@ -868,14 +868,13 @@ def _masked_radial( statistic: str, dof: int, ) -> tuple[npt.NDArray[float], npt.NDArray[float]]: - kz = np.fft.fftfreq(length, d=voxel_size)[:, None, None] - ky = np.fft.fftfreq(length, d=voxel_size)[None, :, None] - kx = np.fft.rfftfreq(length, d=voxel_size)[None, None, :] - r = np.sqrt(kz**2 + ky**2 + kx**2) + r_frac = np.fft.ifftshift( + radial_reduced_grid((length, length, length)), axes=(0, 1) + ) k_nyq = 0.5 / voxel_size nb = length // 2 + 1 - shell = np.floor(r / k_nyq * (nb - 1) + 0.5).astype(int) + shell = np.floor(r_frac * (nb - 1) + 0.5).astype(int) q = np.arange(nb) / (nb - 1) * k_nyq labels = shell.copy() From fd4acc647917c7d0bf22d6f62ecedb09ae51fb01 Mon Sep 17 00:00:00 2001 From: McHaillet Date: Wed, 15 Jul 2026 15:46:57 +0200 Subject: [PATCH 4/8] Restyle estimate_whitening_filter docstring to match module conventions Match the numpydoc formatting used elsewhere in weights.py: single-line opening summary, single spaces after periods, Optional[type] for None defaults, and no trailing periods on single-sentence parameter entries. No behavior change. Co-Authored-By: Claude Sonnet 5 --- src/pytom_tm/weights.py | 59 +++++++++++++++++++---------------------- 1 file changed, 28 insertions(+), 31 deletions(-) diff --git a/src/pytom_tm/weights.py b/src/pytom_tm/weights.py index 622de477..d26e4739 100644 --- a/src/pytom_tm/weights.py +++ b/src/pytom_tm/weights.py @@ -722,55 +722,52 @@ def estimate_whitening_filter( statistic: str = "median", voxel_size: float = 1.0, ) -> tuple[npt.NDArray[float], npt.NDArray[float]]: - """ - Masked radial noise power spectrum of a tomogram, estimated on cubic patches. + """Estimate a whitening filter from the radially averaged noise power spectrum + of a tomogram, sampled on overlapping cubic patches. - Averages the power spectrum ONLY over Fourier voxels the tilt series actually - sampled (via the tilt 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; the highest-variance patches (gold, carbon, - ice, volume edges) are rejected. + 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 ---------- tomogram: npt.NDArray[float] - 3D real-space array. + 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. + 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 mask is - built for. Sensible default ~64 for typical binned ET (~10 A/vox): reaches - the low frequencies a template box needs, is a clean FFT size, and still - fits several times in a thin z so you keep many patches. Drop to 48 if z - is tight; rise to 96 only for thick volumes with large targets. A masked - profile can be interpolated to a different correlation-box size afterwards, - so estimation size and matching size are decoupled -- pick this purely for - estimation quality. + 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 (0.5 = 50%). + fractional overlap between patches reject_frac: float, default 0.10 - drop this fraction of the highest-variance patches. 0 disables. + fraction of the highest-variance patches to reject, 0 disables rejection exclude: Optional[npt.NDArray[bool]], default None - optional bool array, same shape as tomogram, True where voxels must not be - used (gold / carbon / contamination). Do NOT mask out biological - structure -- under H0 the crowded cell IS the noise, and excluding it - biases P_n low and leaves the background sigma above 1. + boolean array, same shape as tomogram, True where voxels must not be used + (e.g. gold, carbon, contamination). Should not be used to mask out + biological structure, as that would bias the estimate low statistic: str, default "median" - "median" (robust; bias-corrected to the mean) or "mean". + 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. + voxel size in Angstrom, sets the physical units of the returned frequency + axis Returns ------- (q, w): tuple[npt.NDArray[float], npt.NDArray[float]] - q is the frequency of each shell, cycles/Angstrom (0..Nyquist), with - shape (patch_size // 2 + 1,). w is the whitening filter (already - transformed from the power spectrum profile, DC-zeroed, high-frequency - tapered and normalized to a maximum of 1), with the same shape as q. + 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 """ tomogram = np.asarray(tomogram, dtype=np.float32) length = int(patch_size) From 1b5313596aa2f9dcdc1d1e2eddf39a11292906d1 Mon Sep 17 00:00:00 2001 From: McHaillet Date: Wed, 15 Jul 2026 15:49:22 +0200 Subject: [PATCH 5/8] add comment about why the wedge is thresholded --- src/pytom_tm/weights.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/pytom_tm/weights.py b/src/pytom_tm/weights.py index d26e4739..bc1f69a5 100644 --- a/src/pytom_tm/weights.py +++ b/src/pytom_tm/weights.py @@ -777,7 +777,7 @@ def estimate_whitening_filter( # 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 + 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) From c775181dec642d32610c4fd048a1740ee83673e4 Mon Sep 17 00:00:00 2001 From: Marten Chaillet <58044494+McHaillet@users.noreply.github.com> Date: Thu, 16 Jul 2026 11:24:13 +0200 Subject: [PATCH 6/8] Update src/pytom_tm/weights.py Co-authored-by: Sander Roet --- src/pytom_tm/weights.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/pytom_tm/weights.py b/src/pytom_tm/weights.py index bc1f69a5..4948c8d0 100644 --- a/src/pytom_tm/weights.py +++ b/src/pytom_tm/weights.py @@ -110,7 +110,7 @@ def radial_reduced_grid( radial_reduced_grid: npt.NDArray[float] fourier space frequency grid, 0 in center, 1 at nyquist """ - if len(shape) not in [1, 2, 3]: + if len(shape) not in {1, 2, 3}: raise ValueError("radial_reduced_grid() only works for 1D, 2D or 3D shapes") reduced_dim = shape[-1] if shape_is_reduced else shape[-1] // 2 + 1 if len(shape) == 1: From 07ba9289505bfb5c3cc9ca01a80984f5f81e180b Mon Sep 17 00:00:00 2001 From: McHaillet Date: Thu, 16 Jul 2026 12:32:33 +0200 Subject: [PATCH 7/8] Address PR review comments on whitening filter estimation - Reject non-cubic templates with a clear error (tmjob.py) - Use keyword arguments for estimate_whitening_filter call - Use int steps instead of float in radial_reduced_grid's np.arange calls - Drop the unused exclude/focus-mask parameter from estimate_whitening_filter - Fix step rounding to avoid Python's banker's rounding - Simplify patch rejection: drop the silent auto-turnoff fallback, raise RuntimeError only when zero patches survive - Correct _patch_slices axis naming to match pytom's internal xyz convention - Raise instead of silently leaving NaNs when too few radial shells are valid - Add/extend tests for the above, including previously uncovered branches (_masked_radial's new error path and the statistic="mean" code path) Co-Authored-By: Claude Sonnet 5 --- src/pytom_tm/tmjob.py | 12 +++++++--- src/pytom_tm/weights.py | 52 +++++++++++++++++++---------------------- tests/test_tmjob.py | 29 +++++++++++++++++++++++ tests/test_weights.py | 39 +++++++++++++++++++++++++++++-- 4 files changed, 99 insertions(+), 33 deletions(-) diff --git a/src/pytom_tm/tmjob.py b/src/pytom_tm/tmjob.py index 0704de63..ef5415e1 100644 --- a/src/pytom_tm/tmjob.py +++ b/src/pytom_tm/tmjob.py @@ -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. " @@ -516,13 +522,13 @@ def __init__( logging.info("Estimating whitening filter...") patch_size = min(max(self.template_shape[0], 64), min(self.search_size)) _, weights = estimate_whitening_filter( - read_mrc(self.tomogram)[ + 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], ], - self.ts_metadata, - patch_size, + ts_metadata=self.ts_metadata, + patch_size=patch_size, voxel_size=self.voxel_size, ) np.save(self.whitening_filter, weights) diff --git a/src/pytom_tm/weights.py b/src/pytom_tm/weights.py index 4948c8d0..22b7262f 100644 --- a/src/pytom_tm/weights.py +++ b/src/pytom_tm/weights.py @@ -114,12 +114,12 @@ def radial_reduced_grid( raise ValueError("radial_reduced_grid() only works for 1D, 2D or 3D shapes") reduced_dim = shape[-1] if shape_is_reduced else shape[-1] // 2 + 1 if len(shape) == 1: - return np.arange(0, reduced_dim, 1.0) / (reduced_dim - 1) + return np.arange(0, reduced_dim, 1) / (reduced_dim - 1) elif len(shape) == 3: x = ( np.abs( np.arange( - -shape[0] // 2 + shape[0] % 2, shape[0] // 2 + shape[0] % 2, 1.0 + -shape[0] // 2 + shape[0] % 2, shape[0] // 2 + shape[0] % 2, 1 ) ) / (shape[0] // 2) @@ -127,23 +127,23 @@ def radial_reduced_grid( y = ( np.abs( np.arange( - -shape[1] // 2 + shape[1] % 2, shape[1] // 2 + shape[1] % 2, 1.0 + -shape[1] // 2 + shape[1] % 2, shape[1] // 2 + shape[1] % 2, 1 ) ) / (shape[1] // 2) )[:, np.newaxis] - z = np.arange(0, reduced_dim, 1.0) / (reduced_dim - 1) + z = np.arange(0, reduced_dim, 1) / (reduced_dim - 1) return np.sqrt(x**2 + y**2 + z**2) elif len(shape) == 2: x = ( np.abs( np.arange( - -shape[0] // 2 + shape[0] % 2, shape[0] // 2 + shape[0] % 2, 1.0 + -shape[0] // 2 + shape[0] % 2, shape[0] // 2 + shape[0] % 2, 1 ) ) / (shape[0] // 2) )[:, np.newaxis] - y = np.arange(0, reduced_dim, 1.0) / (reduced_dim - 1) + y = np.arange(0, reduced_dim, 1) / (reduced_dim - 1) return np.sqrt(x**2 + y**2) @@ -718,7 +718,6 @@ def estimate_whitening_filter( patch_size: int, overlap: float = 0.5, reject_frac: float = 0.10, - exclude: npt.NDArray[bool] | None = None, statistic: str = "median", voxel_size: float = 1.0, ) -> tuple[npt.NDArray[float], npt.NDArray[float]]: @@ -750,10 +749,6 @@ def estimate_whitening_filter( fractional overlap between patches reject_frac: float, default 0.10 fraction of the highest-variance patches to reject, 0 disables rejection - exclude: Optional[npt.NDArray[bool]], default None - boolean array, same shape as tomogram, True where voxels must not be used - (e.g. gold, carbon, contamination). Should not be used to mask out - biological structure, as that would bias the estimate low statistic: str, default "median" radial averaging statistic, either "median" (robust, bias-corrected to the mean) or "mean" @@ -785,11 +780,9 @@ def estimate_whitening_filter( win_power = float((win**2).sum()) # collect windowed periodograms over overlapping patches - step = max(1, int(round(length * (1.0 - overlap)))) + step = max(1, int(length * (1.0 - overlap) + 0.5)) # round half up psds, variances = [], [] for sl in _patch_slices(tomogram.shape, length, step): - if exclude is not None and exclude[sl].mean() > 0.01: - continue v = tomogram[sl] if not np.all(np.isfinite(v)): continue @@ -798,18 +791,18 @@ def estimate_whitening_filter( psds.append((f.real**2 + f.imag**2) / win_power) variances.append(float(v.var())) - if len(psds) < 4: - raise RuntimeError(f"only {len(psds)} usable patches; reduce patch_size") - - # reject high-variance patches (gold / carbon / ice / edges) + # 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() < 4: - keep = np.ones(len(psds), bool) - psd = np.mean([p for p, k in zip(psds, keep) if k], axis=0) + 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 @@ -846,14 +839,14 @@ def starts(n): s.append(n - length) return s - zs, ys, xs = (starts(n) for n in shape) - for z in zs: + xs, ys, zs = (starts(n) for n in shape) + for x in xs: for y in ys: - for x in xs: + for z in zs: yield ( - slice(z, z + length), - slice(y, y + length), slice(x, x + length), + slice(y, y + length), + slice(z, z + length), ) @@ -889,8 +882,11 @@ def _masked_radial( 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: - prof = np.interp(q, q[good], prof[good]) + 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 diff --git a/tests/test_tmjob.py b/tests/test_tmjob.py index 0a64d006..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, diff --git a/tests/test_weights.py b/tests/test_weights.py index 245407ea..0fea1941 100644 --- a/tests/test_weights.py +++ b/tests/test_weights.py @@ -3,6 +3,7 @@ from pytom_tm.weights import ( create_wedge, _create_symmetric_wedge, + _masked_radial, create_ctf, create_gaussian_band_pass, radial_reduced_grid, @@ -352,11 +353,11 @@ def test_estimate_whitening_filter(self): with self.assertRaises( RuntimeError, - msg="Estimate whitening filter should raise RuntimeError if too few " + msg="Estimate whitening filter should raise RuntimeError if no " "patches are usable.", ): estimate_whitening_filter( - tomogram, + np.full((64, 64, 64), np.nan, dtype=np.float32), self.ts_metadata, patch_size=64, voxel_size=self.voxel_size, @@ -388,6 +389,40 @@ def test_estimate_whitening_filter(self): 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", + ) + self.assertEqual( + 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_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( + RuntimeError, + msg="_masked_radial should raise RuntimeError if fewer than 2 radial " + "shells have valid data.", + ): + _masked_radial(psd, mask, length, self.voxel_size, "median", 10) + def test_profile_to_weighting(self): with self.assertRaises( ValueError, From a226dedd5fb71f297af55c10128ab4766159a741 Mon Sep 17 00:00:00 2001 From: Marten Chaillet <58044494+McHaillet@users.noreply.github.com> Date: Thu, 16 Jul 2026 15:33:16 +0200 Subject: [PATCH 8/8] Update src/pytom_tm/weights.py Co-authored-by: Sander Roet --- src/pytom_tm/weights.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/pytom_tm/weights.py b/src/pytom_tm/weights.py index 7fe27ecc..a727d636 100644 --- a/src/pytom_tm/weights.py +++ b/src/pytom_tm/weights.py @@ -803,7 +803,7 @@ def cosine_cutoff(): 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), np.zeros_like(prof)) + 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()