Skip to main content

av_denoise_core/nlmeans/
params.rs

1use super::{MotionCompensationMode, PrefilterMode, prefilter};
2
3/// The reference value patch distances are normalised against, matching
4/// FFmpeg's nlmeans at 255 squared.
5///
6/// Distances here are measured in `[0, 1]` units, so this constant folds
7/// in the scale-up back to 8-bit terms.
8pub(super) const NLM_NORM: f32 = 255.0 * 255.0;
9
10/// A scaling factor inherited from FFmpeg's nlmeans, kept so our
11/// `strength` parameter means the same thing theirs does.
12pub(super) const NLM_LEGACY: f32 = 3.0;
13
14/// The measured HQ default `strength` for luma at each temporal radius,
15/// indexed by `temporal_radius.min(8)`.
16///
17/// `hq_default_strength` reads this for `ChannelMode::Luma` and
18/// `ChannelMode::Yuv`.
19const HQ_DEFAULT_STRENGTH_LUMA: [f32; 9] = [0.45, 0.45, 0.42, 0.42, 0.35, 0.35, 0.35, 0.30, 0.30];
20
21/// The measured HQ default `strength` for chroma at each temporal
22/// radius, indexed by `temporal_radius.min(8)`.
23///
24/// `hq_default_strength` reads this for `ChannelMode::Chroma`.
25const HQ_DEFAULT_STRENGTH_CHROMA: [f32; 9] = [1.00, 0.85, 0.70, 0.70, 0.70, 0.70, 0.70, 0.70, 0.70];
26
27/// The calibrated default `strength` multiplier for `nlmeans-hq`'s
28/// auto-strength mode.
29///
30/// The answer depends on which plane is being denoised and how far the
31/// temporal window reaches.
32///
33/// # Where the numbers come from
34///
35/// Every entry is a measured value rather than a fitted curve. Both
36/// tables come from quality-harness sweeps that score a grid of
37/// strengths at three noise levels per radius.
38///
39/// The luma sweep covers each radius from 0 to 8 directly.
40///
41/// The chroma sweep pins luma at the value already chosen for that
42/// radius, so the chroma numbers stay clean, and covers radii 0, 1, 2,
43/// 4, and 8 with a bracketed peak at each.
44///
45/// At each measured radius the chosen value is the one whose worst XPSNR
46/// gain across the tested noise levels is highest, so it holds up at
47/// whichever noise level is hardest to serve.
48///
49/// # Shape of the tables
50///
51/// Luma never rises with radius, because a wider temporal window already
52/// gathers more samples to average over.
53///
54/// Chroma falls the same way out to radius 2 and then holds flat at
55/// 0.70. Radii 3, 5, 6, and 7 sit on that measured plateau rather than
56/// being swept directly.
57///
58/// `ChannelMode::Yuv` reads the luma table, on the assumption that a
59/// fused pass is dominated by luma. That mode was not part of the sweep,
60/// so this is an assumption rather than a measurement.
61///
62/// # Clamping
63///
64/// `temporal_radius` is clamped to the last table index, which matches
65/// [`MAX_TEMPORAL_RADIUS`]. That is only a safety net, because
66/// [`NlmParams::validate`] already rejects anything larger.
67pub fn hq_default_strength(channels: ChannelMode, temporal_radius: u32) -> f32 {
68    let idx = temporal_radius.min(MAX_TEMPORAL_RADIUS) as usize;
69    match channels {
70        ChannelMode::Luma | ChannelMode::Yuv => HQ_DEFAULT_STRENGTH_LUMA[idx],
71        ChannelMode::Chroma => HQ_DEFAULT_STRENGTH_CHROMA[idx],
72    }
73}
74
75/// The smallest frame side length the denoiser supports.
76///
77/// The Immerkær noise estimate only reads interior pixels, because its
78/// 3x3 mask cannot reach the one-pixel border. A frame under 3 pixels
79/// across has no interior at all, which leaves the estimate undefined.
80pub const MIN_FRAME_DIM: u32 = 3;
81
82/// Rejects frame dimensions the kernels cannot handle.
83///
84/// Both denoiser constructors call this before allocating any buffer.
85pub fn validate_dimensions(width: u32, height: u32) -> Result<(), anyhow::Error> {
86    if width < MIN_FRAME_DIM || height < MIN_FRAME_DIM {
87        anyhow::bail!(
88            "frame dimensions {width}x{height} are below the supported minimum of \
89             {MIN_FRAME_DIM}x{MIN_FRAME_DIM}, because the noise estimate needs at \
90             least one interior pixel"
91        );
92    }
93    Ok(())
94}
95
96/// The patch radius above which the dispatcher switches to the separable
97/// path, so per-pixel cost stays linear in `patch_radius`.
98pub(super) const SEPARABLE_THRESHOLD: u32 = 8;
99
100/// The hard ceiling on `patch_radius`.
101///
102/// The fused kernels load a `(block + 2 * patch_radius)^2` tile into
103/// shared memory, and anything larger runs out of it on RDNA-class GPUs.
104pub const MAX_PATCH_RADIUS: u32 = 16;
105
106/// The hard ceiling on `search_radius`.
107///
108/// Shared memory is not the limit here. The windowed kernel's tile is
109/// `(block + 2 * patch_radius + 2 * search_radius)^2 * stored_ch * 4`
110/// bytes, which stays comfortably inside hardware limits at every
111/// supported size.
112///
113/// The real cost is the kernel's fully unrolled
114/// `(2 * search_radius + 1)^2` window loop. Both its compiled size and
115/// how long it takes to generate grow with the radius. See the
116/// stack-size note in `.cargo/config.toml`.
117///
118/// The per-offset dispatch path, used when `patch_radius` forces the
119/// separable fallback, is limited by the same constant, which keeps its
120/// `(2 * search_radius + 1)^2` launches per temporal offset reasonable.
121pub const MAX_SEARCH_RADIUS: u32 = 8;
122
123/// The hard ceiling on `temporal_radius`.
124///
125/// The ring buffer holds `2 * radius + 1` frames, so device memory grows
126/// with it. At radius 16, 1080p YUV would need roughly 540 MB for the
127/// input alone.
128pub const MAX_TEMPORAL_RADIUS: u32 = 8;
129
130/// The hard ceiling on the radius the bilateral prefilter derives from
131/// `sigma_s`, where `radius = ceil(2 * sigma_s).max(1)`. See
132/// `prefilter::bilateral_radius`.
133///
134/// `nlm_bilateral` loads a `(32 + 2r) x (8 + 2r)` tile of
135/// `Vector<f32, N>` into shared memory. `N` reaches 4 for YUV storage at
136/// 4 bytes per `f32`, so the tile costs `16 * (32 + 2r) * (8 + 2r)`
137/// bytes.
138///
139/// RDNA-class hardware gives 64 KiB of shared memory to work with. At
140/// `r = 22` the tile is 63,232 bytes, or 61.75 KiB, which fits. At
141/// `r = 23` it is 67,392 bytes, or 65.8 KiB, which does not.
142///
143/// So 22 is the largest radius that fits, and `bilateral_radius` reaches
144/// it at `sigma_s = 11.0`.
145pub const MAX_BILATERAL_RADIUS: u32 = 22;
146
147#[derive(Debug, Clone, Copy, PartialEq, Eq)]
148/// Which channels of a frame the denoiser works on.
149pub enum ChannelMode {
150    /// The single brightness channel, with distances scaled by 3.0.
151    Luma,
152    /// The two colour channels, U and V, with distances scaled by 1.5.
153    Chroma,
154    /// All three channels together, with distances left unscaled.
155    Yuv,
156}
157
158impl ChannelMode {
159    /// How many channels take part in the distance and the output.
160    pub fn count(self) -> u32 {
161        match self {
162            ChannelMode::Luma => 1,
163            ChannelMode::Chroma => 2,
164            ChannelMode::Yuv => 3,
165        }
166    }
167
168    /// How many channels each pixel occupies in GPU storage.
169    ///
170    /// This is padded up to the next supported vector width so kernels
171    /// can read whole `Line<f32>` values at once. Backends only support
172    /// power-of-two widths, so YUV pads from 3 up to 4.
173    pub fn storage_count(self) -> u32 {
174        match self {
175            ChannelMode::Luma => 1,
176            ChannelMode::Chroma => 2,
177            ChannelMode::Yuv => 4,
178        }
179    }
180}
181
182/// Parameters for the quality-focused `nlmeans-hq` variant.
183///
184/// The measured noise level drives both the effective strength and the
185/// distance floor, so the weighting adapts to how noisy the source
186/// really is.
187#[derive(Debug, Clone, Copy, PartialEq)]
188pub struct HqParams {
189    /// Reads `strength` as a multiplier on the noise level, making the
190    /// effective FFmpeg-style strength `strength * sigma_eff * 255`.
191    ///
192    /// Defaults to true.
193    pub auto_strength: bool,
194    /// Subtracts the expected noise floor from patch distances before
195    /// weighting, so a match is not penalised for the noise it carries.
196    ///
197    /// Defaults to true.
198    pub noise_floor: bool,
199    /// A fixed noise standard deviation in `[0, 1]` units, replacing the
200    /// automatic per-frame estimate.
201    ///
202    /// `None`, the default, measures the noise in each pushed frame and
203    /// smooths it over time. `Some` applies one fixed value to every
204    /// frame instead.
205    ///
206    /// The CLI takes this in 8-bit units through `--hq-sigma` and
207    /// divides by 255.
208    pub sigma_override: Option<f32>,
209    /// Weights each temporal neighbour by how well it block-matches the
210    /// centre frame, so occlusion or a change of content collapses its
211    /// contribution rather than blurring it in.
212    ///
213    /// Only has an effect when `temporal_radius` is above 0. Defaults to
214    /// true.
215    pub temporal_confidence: bool,
216    /// A multiplier on the per-pixel mismatch threshold, which sets how
217    /// much extra SAD a block tolerates before its confidence starts to
218    /// fall.
219    ///
220    /// Higher values tolerate larger mismatches. Defaults to 1.0.
221    pub thsad_scale: f32,
222    /// A multiplier applied to each channel's measured sigma before it
223    /// folds into the running estimate.
224    ///
225    /// `1.0`, the default, keeps the measurement as it is. This does
226    /// nothing when `sigma_override` is set, because the estimator never
227    /// runs in that case.
228    ///
229    /// The CLI takes this as `--hq-sigma-scale`.
230    pub sigma_scale: f32,
231    /// Estimates noise fresh from each submit's own window, instead of
232    /// smoothing it with an exponential average carried forward from
233    /// every earlier frame the stream has folded.
234    ///
235    /// `false`, the default, keeps the temporal EMA every calibrated
236    /// preset assumes. `true` makes the automatic estimate depend only
237    /// on the frames currently in the window, so a
238    /// [`crate::frame::PlanarDenoiser::reseed`] targeting frame `n` and
239    /// a continuous stream that reaches frame `n` compute the same
240    /// sigma, regardless of how each got there. This does nothing when
241    /// `sigma_override` pins a fixed value, because the estimator never
242    /// runs in that case.
243    pub windowed_noise_estimation: bool,
244}
245
246impl Default for HqParams {
247    fn default() -> Self {
248        Self {
249            auto_strength: true,
250            noise_floor: true,
251            sigma_override: None,
252            temporal_confidence: true,
253            thsad_scale: 1.0,
254            sigma_scale: 1.0,
255            windowed_noise_estimation: false,
256        }
257    }
258}
259
260impl HqParams {
261    /// The HQ defaults with a fixed noise level in `[0, 1]` units, which
262    /// skips automatic estimation.
263    pub fn with_sigma(sigma: f32) -> Self {
264        Self {
265            sigma_override: Some(sigma),
266            ..Self::default()
267        }
268    }
269}
270
271/// The low-level parameters a denoiser is built from.
272#[derive(Debug, Clone)]
273pub struct NlmParams {
274    /// How many frames on each side of the current one to look at.
275    ///
276    /// 0 means each frame is cleaned on its own. Anything higher uses a
277    /// window of `2 * radius + 1` frames.
278    pub temporal_radius: u32,
279    /// Half the width of the search window, which covers
280    /// `(2 * radius + 1)^2` pixels. Defaults to 2.
281    pub search_radius: u32,
282    /// Half the width of a compared patch, which covers
283    /// `(2 * radius + 1)^2` pixels. Defaults to 4.
284    pub patch_radius: u32,
285    /// How hard to filter. Higher values smooth more. Defaults to 1.2.
286    pub strength: f32,
287    /// How much weight the centre pixel gets in the average.
288    ///
289    /// Defaults to 1.0. Set it to 0 for pure NLM, where the centre pixel
290    /// only counts through the patches that match it.
291    pub self_weight: f32,
292    /// Which channels to process.
293    pub channels: ChannelMode,
294    /// Which image the patch distances are measured against.
295    ///
296    /// Defaults to `None`. When set, the weights come from a prefiltered
297    /// or externally supplied image while the pixels being averaged
298    /// still come from the original input.
299    pub prefilter: PrefilterMode,
300    /// Whether temporal denoising follows motion between frames.
301    ///
302    /// Defaults to `None`. Set to `Mvtools`, each submit warps the
303    /// temporal neighbours into line with the centre frame before the
304    /// NLM weighting runs.
305    ///
306    /// Only has an effect when `temporal_radius` is above 0.
307    pub motion_compensation: MotionCompensationMode,
308    /// The quality-mode parameters. `None` runs the fast path unchanged.
309    pub hq: Option<HqParams>,
310}
311
312impl Default for NlmParams {
313    fn default() -> Self {
314        Self {
315            temporal_radius: 0,
316            search_radius: 2,
317            patch_radius: 4,
318            strength: 1.2,
319            self_weight: 1.0,
320            channels: ChannelMode::Yuv,
321            prefilter: PrefilterMode::None,
322            motion_compensation: MotionCompensationMode::None,
323            hq: None,
324        }
325    }
326}
327
328impl NlmParams {
329    /// The FFmpeg-style strength the weighting actually uses.
330    ///
331    /// With HQ auto-strength the user's value multiplies the noise
332    /// level, so one setting follows sources of different noisiness.
333    ///
334    /// `sigma_eff` is the scale-weighted RMS of the per-channel noise
335    /// estimates.
336    pub(super) fn effective_strength_with(&self, sigma_eff: Option<f32>) -> f32 {
337        match (self.hq, sigma_eff) {
338            (Some(hq), Some(sigma)) if hq.auto_strength => self.strength * sigma * 255.0,
339            _ => self.strength,
340        }
341    }
342
343    /// `h2_inv_norm` for a noise estimate given here, ignoring whatever
344    /// `self.hq.sigma_override` holds.
345    ///
346    /// The denoiser calls this each submit to refresh the value from a
347    /// freshly measured sigma.
348    pub fn h2_inv_norm_with(&self, sigma_eff: Option<f32>) -> f32 {
349        let s_size = (2 * self.patch_radius + 1) * (2 * self.patch_radius + 1);
350        let s = self.effective_strength_with(sigma_eff);
351        NLM_NORM / (NLM_LEGACY * s * s * s_size as f32)
352    }
353
354    /// `h2_inv_norm` using `self.hq.sigma_override` as the noise
355    /// estimate, or no estimate at all on the fast path.
356    ///
357    /// HQ denoisers that estimate noise automatically call
358    /// [`Self::h2_inv_norm_with`] each submit instead.
359    pub fn h2_inv_norm(&self) -> f32 {
360        self.h2_inv_norm_with(self.hq.and_then(|hq| hq.sigma_override))
361    }
362
363    /// The patch distance two noisy copies of the same content are
364    /// expected to show, for a given set of per-channel sigmas.
365    ///
366    /// Each active channel contributes `2 * channel_scale * sigma^2`,
367    /// summed over all `(2 * patch_radius + 1)^2` taps.
368    ///
369    /// Returns 0 when the HQ noise floor is off, or when there is no
370    /// estimate to apply.
371    pub(super) fn noise_offset_with(&self, sigmas: Option<&[f32]>) -> f32 {
372        match (self.hq, sigmas) {
373            (Some(hq), Some(sigmas)) if hq.noise_floor => {
374                let s_size = (2 * self.patch_radius + 1) * (2 * self.patch_radius + 1);
375                let scale = channel_scale(self.channels);
376                let count = self.channels.count() as usize;
377                let sum_sq: f32 = sigmas.iter().take(count).map(|&s| s * s).sum();
378                2.0 * scale * sum_sq * s_size as f32
379            },
380            _ => 0.0,
381        }
382    }
383
384    /// `noise_offset` with `self.hq.sigma_override` applied to every
385    /// active channel, or no estimate at all on the fast path.
386    ///
387    /// HQ denoisers that estimate noise automatically call
388    /// [`Self::noise_offset_with`] each submit instead.
389    pub(super) fn noise_offset(&self) -> f32 {
390        match self.hq.and_then(|hq| hq.sigma_override) {
391            Some(sigma) => {
392                let sigmas = [sigma; 3];
393                self.noise_offset_with(Some(&sigmas[..self.channels.count() as usize]))
394            },
395            None => 0.0,
396        }
397    }
398
399    pub(super) fn total_frames(&self) -> u32 {
400        1 + 2 * self.temporal_radius
401    }
402
403    /// Rejects parameter combinations that would fail to launch, by
404    /// running the kernels past their shared-memory or register limits,
405    /// or that would produce meaningless output.
406    ///
407    /// `NlmDenoiser::new` calls this for you. Callers building params by
408    /// hand can call it directly to see errors before construction.
409    pub fn validate(&self) -> Result<(), anyhow::Error> {
410        if self.patch_radius > MAX_PATCH_RADIUS {
411            anyhow::bail!(
412                "patch_radius={} exceeds the supported maximum of {}, because larger \
413                 patches exhaust on-chip shared memory in the fused and windowed kernels",
414                self.patch_radius,
415                MAX_PATCH_RADIUS,
416            );
417        }
418
419        if self.search_radius > MAX_SEARCH_RADIUS {
420            anyhow::bail!(
421                "search_radius={} exceeds the supported maximum of {}. The windowed \
422                 kernel's search window loop is fully unrolled, so both its compiled \
423                 size and its build time grow with search_radius",
424                self.search_radius,
425                MAX_SEARCH_RADIUS,
426            );
427        }
428
429        if self.temporal_radius > MAX_TEMPORAL_RADIUS {
430            anyhow::bail!(
431                "temporal_radius={} exceeds the supported maximum of {}, because the \
432                 ring buffer grows in step with the window size",
433                self.temporal_radius,
434                MAX_TEMPORAL_RADIUS,
435            );
436        }
437
438        if !(self.strength.is_finite() && self.strength > 0.0) {
439            anyhow::bail!(
440                "strength must be finite and greater than 0, got {}. A strength of 0 \
441                 produces an infinite Welsch normalisation factor",
442                self.strength,
443            );
444        }
445
446        if !self.self_weight.is_finite() || self.self_weight < 0.0 {
447            anyhow::bail!(
448                "self_weight must be finite and 0 or greater, got {}",
449                self.self_weight,
450            );
451        }
452
453        if let Some(hq) = self.hq
454            && let Some(sigma) = hq.sigma_override
455            && (!sigma.is_finite() || sigma <= 0.0 || sigma > 1.0)
456        {
457            anyhow::bail!(
458                "hq sigma_override must be finite and in (0, 1] in normalised units, got {}",
459                sigma,
460            );
461        }
462
463        if let Some(hq) = self.hq
464            && !(hq.thsad_scale.is_finite() && hq.thsad_scale > 0.0)
465        {
466            anyhow::bail!(
467                "hq thsad_scale must be finite and greater than 0, got {}. A thsad_scale \
468                 of 0 collapses every block's confidence to zero no matter how well it \
469                 matches",
470                hq.thsad_scale,
471            );
472        }
473
474        if let Some(hq) = self.hq
475            && !(hq.sigma_scale.is_finite() && (0.1..=10.0).contains(&hq.sigma_scale))
476        {
477            anyhow::bail!(
478                "hq sigma_scale must be finite and in [0.1, 10.0], got {}",
479                hq.sigma_scale,
480            );
481        }
482
483        if let PrefilterMode::Bilateral { sigma_s, sigma_r } = self.prefilter {
484            if !sigma_s.is_finite() || sigma_s <= 0.0 {
485                anyhow::bail!(
486                    "bilateral prefilter sigma_s must be finite and greater than 0, got \
487                     {}. A sigma_s of 0 produces an infinite spatial-weight \
488                     normalisation factor",
489                    sigma_s,
490                );
491            }
492            if !sigma_r.is_finite() || sigma_r <= 0.0 {
493                anyhow::bail!(
494                    "bilateral prefilter sigma_r must be finite and greater than 0, got \
495                     {}. A sigma_r of 0 produces an infinite range-weight normalisation \
496                     factor, which turns the centre tap into NaN",
497                    sigma_r,
498                );
499            }
500            // sigma_s decides the shared-memory tile radius through
501            // `prefilter::bilateral_radius`. Checking that derived
502            // radius, rather than working out an equivalent sigma_s
503            // threshold here, keeps this in step if the formula ever
504            // changes.
505            //
506            // A very large sigma_s can also overflow the radius inside
507            // the tile-size expression. This check catches that too,
508            // because an overflowed radius always lands far past the
509            // maximum.
510            let bilateral_radius = prefilter::bilateral_radius(sigma_s);
511            if bilateral_radius > MAX_BILATERAL_RADIUS {
512                anyhow::bail!(
513                    "bilateral prefilter sigma_s={} implies a shared-memory tile radius \
514                     of {}, from radius = ceil(2 * sigma_s) with a minimum of 1. That \
515                     is past the supported maximum of {}, and larger radii exhaust \
516                     on-chip shared memory in the bilateral kernel",
517                    sigma_s,
518                    bilateral_radius,
519                    MAX_BILATERAL_RADIUS,
520                );
521            }
522            // A sigma can be finite and positive yet small enough that
523            // `sigma * sigma` underflows to 0.0 in f32, which happens
524            // below roughly 3.8e-20. That makes the reciprocal
525            // normalisation factor the kernel uses infinite.
526            //
527            // Checking the same derived factor `run_bilateral` computes
528            // for the launch catches this wherever the underflow
529            // threshold actually falls, without picking a sigma cutoff
530            // by hand or repeating the expression here.
531            if !prefilter::inv_two_sigma_sq(sigma_s).is_finite() {
532                anyhow::bail!(
533                    "bilateral prefilter sigma_s is too small, got {}. Squaring it \
534                     underflows to 0 in f32, which makes the spatial-weight \
535                     normalisation factor infinite",
536                    sigma_s,
537                );
538            }
539            if !prefilter::inv_two_sigma_sq(sigma_r).is_finite() {
540                anyhow::bail!(
541                    "bilateral prefilter sigma_r is too small, got {}. Squaring it \
542                     underflows to 0 in f32, which makes the range-weight normalisation \
543                     factor infinite and the centre tap NaN",
544                    sigma_r,
545                );
546            }
547        }
548
549        if let PrefilterMode::NlmSpatial { strength_scale } = self.prefilter {
550            if !strength_scale.is_finite() || strength_scale <= 0.0 {
551                anyhow::bail!(
552                    "nlm pilot strength_scale must be finite and greater than 0, got {}",
553                    strength_scale,
554                );
555            }
556            if self.patch_radius > SEPARABLE_THRESHOLD {
557                anyhow::bail!(
558                    "the nlm pilot uses the windowed spatial kernel, which supports \
559                     patch_radius up to {} (got {})",
560                    SEPARABLE_THRESHOLD,
561                    self.patch_radius,
562                );
563            }
564        }
565
566        self.motion_compensation.validate()?;
567
568        Ok(())
569    }
570}
571
572/// The per-channel distance scale for a channel mode, which is 3 for
573/// luma, 1.5 for chroma, and 1 for full YUV.
574///
575/// This matches the `channel_scale` the weighting kernels use on the
576/// GPU, and it is the same for every channel within a given mode.
577pub(super) fn channel_scale(channels: ChannelMode) -> f32 {
578    match channels {
579        ChannelMode::Luma => 3.0,
580        ChannelMode::Chroma => 1.5,
581        ChannelMode::Yuv => 1.0,
582    }
583}
584
585/// The scale-weighted RMS of the per-channel noise estimates, over the
586/// channels a mode actually uses.
587///
588/// Because `channel_scale` is the same for every channel in a mode, the
589/// weighting cancels out and this is really just a plain RMS.
590///
591/// Anything in `sigmas` past the mode's channel count is ignored.
592pub(super) fn sigma_eff(sigmas: &[f32], channels: ChannelMode) -> f32 {
593    let count = channels.count() as usize;
594    let sum_sq: f32 = sigmas.iter().take(count).map(|&s| s * s).sum();
595    (sum_sq / count as f32).sqrt()
596}
597
598#[cfg(test)]
599mod tests {
600    use super::*;
601
602    #[test]
603    fn noise_offset_scales_with_sigma_and_patch_size() {
604        let sigma = 4.0 / 255.0;
605        let params = NlmParams {
606            patch_radius: 4,
607            hq: Some(HqParams::with_sigma(sigma)),
608            ..NlmParams::default()
609        };
610
611        let expected = 6.0 * sigma * sigma * 81.0;
612        assert!(
613            (params.noise_offset() - expected).abs() < 1e-6,
614            "expected {expected}, got {}",
615            params.noise_offset()
616        );
617    }
618
619    #[test]
620    fn noise_offset_zero_without_noise_floor() {
621        let params = NlmParams {
622            hq: Some(HqParams {
623                auto_strength: true,
624                noise_floor: false,
625                sigma_override: Some(4.0 / 255.0),
626                temporal_confidence: true,
627                thsad_scale: 1.0,
628                sigma_scale: 1.0,
629                windowed_noise_estimation: false,
630            }),
631            ..NlmParams::default()
632        };
633
634        assert_eq!(params.noise_offset(), 0.0);
635    }
636
637    #[test]
638    fn noise_offset_zero_without_hq() {
639        let params = NlmParams::default();
640        assert_eq!(params.noise_offset(), 0.0);
641    }
642
643    #[test]
644    fn h2_inv_norm_with_auto_strength_matches_hand_computed() {
645        let sigma = 8.0 / 255.0;
646        let params = NlmParams {
647            strength: 1.0,
648            hq: Some(HqParams::with_sigma(sigma)),
649            ..NlmParams::default()
650        };
651
652        let s_size = (2 * params.patch_radius + 1) * (2 * params.patch_radius + 1);
653        let effective_strength = 1.0 * sigma * 255.0;
654        let expected = NLM_NORM / (NLM_LEGACY * effective_strength * effective_strength * s_size as f32);
655
656        assert!(
657            (params.h2_inv_norm() - expected).abs() < 1e-6,
658            "expected {expected}, got {}",
659            params.h2_inv_norm()
660        );
661    }
662
663    #[test]
664    fn validate_rejects_zero_hq_sigma() {
665        let params = NlmParams {
666            hq: Some(HqParams::with_sigma(0.0)),
667            ..NlmParams::default()
668        };
669        assert!(params.validate().is_err());
670    }
671
672    #[test]
673    fn validate_rejects_hq_sigma_above_one() {
674        let params = NlmParams {
675            hq: Some(HqParams::with_sigma(1.5)),
676            ..NlmParams::default()
677        };
678        assert!(params.validate().is_err());
679    }
680
681    #[test]
682    fn validate_rejects_nan_hq_sigma() {
683        let params = NlmParams {
684            hq: Some(HqParams::with_sigma(f32::NAN)),
685            ..NlmParams::default()
686        };
687        assert!(params.validate().is_err());
688    }
689
690    #[test]
691    fn validate_rejects_zero_thsad_scale() {
692        let params = NlmParams {
693            hq: Some(HqParams {
694                thsad_scale: 0.0,
695                ..HqParams::default()
696            }),
697            ..NlmParams::default()
698        };
699        assert!(params.validate().is_err());
700    }
701
702    #[test]
703    fn validate_rejects_negative_thsad_scale() {
704        let params = NlmParams {
705            hq: Some(HqParams {
706                thsad_scale: -1.0,
707                ..HqParams::default()
708            }),
709            ..NlmParams::default()
710        };
711        assert!(params.validate().is_err());
712    }
713
714    #[test]
715    fn validate_rejects_nan_thsad_scale() {
716        let params = NlmParams {
717            hq: Some(HqParams {
718                thsad_scale: f32::NAN,
719                ..HqParams::default()
720            }),
721            ..NlmParams::default()
722        };
723        assert!(params.validate().is_err());
724    }
725
726    #[test]
727    fn validate_accepts_default_thsad_scale() {
728        let params = NlmParams {
729            hq: Some(HqParams::default()),
730            ..NlmParams::default()
731        };
732        assert!(params.validate().is_ok());
733    }
734
735    #[test]
736    fn hq_params_default_sigma_scale_is_one() {
737        assert_eq!(HqParams::default().sigma_scale, 1.0);
738    }
739
740    #[test]
741    fn validate_rejects_sigma_scale_below_the_minimum() {
742        let params = NlmParams {
743            hq: Some(HqParams {
744                sigma_scale: 0.05,
745                ..HqParams::default()
746            }),
747            ..NlmParams::default()
748        };
749        let err = params.validate().expect_err("0.05 is below the 0.1 minimum");
750        assert!(
751            err.to_string().contains("hq sigma_scale"),
752            "error should name the field, got {err}"
753        );
754    }
755
756    #[test]
757    fn validate_rejects_sigma_scale_above_the_maximum() {
758        let params = NlmParams {
759            hq: Some(HqParams {
760                sigma_scale: 10.5,
761                ..HqParams::default()
762            }),
763            ..NlmParams::default()
764        };
765        assert!(params.validate().is_err());
766    }
767
768    #[test]
769    fn validate_rejects_nan_sigma_scale() {
770        let params = NlmParams {
771            hq: Some(HqParams {
772                sigma_scale: f32::NAN,
773                ..HqParams::default()
774            }),
775            ..NlmParams::default()
776        };
777        assert!(params.validate().is_err());
778    }
779
780    #[test]
781    fn validate_accepts_sigma_scale_at_the_bounds() {
782        let low = NlmParams {
783            hq: Some(HqParams {
784                sigma_scale: 0.1,
785                ..HqParams::default()
786            }),
787            ..NlmParams::default()
788        };
789        assert!(low.validate().is_ok());
790
791        let high = NlmParams {
792            hq: Some(HqParams {
793                sigma_scale: 10.0,
794                ..HqParams::default()
795            }),
796            ..NlmParams::default()
797        };
798        assert!(high.validate().is_ok());
799    }
800
801    #[test]
802    fn noise_offset_with_handles_distinct_per_channel_sigmas() {
803        let sigma_u = 4.0 / 255.0;
804        let sigma_v = 10.0 / 255.0;
805        let params = NlmParams {
806            patch_radius: 4,
807            channels: ChannelMode::Chroma,
808            hq: Some(HqParams {
809                auto_strength: true,
810                noise_floor: true,
811                sigma_override: None,
812                temporal_confidence: true,
813                thsad_scale: 1.0,
814                sigma_scale: 1.0,
815                windowed_noise_estimation: false,
816            }),
817            ..NlmParams::default()
818        };
819
820        let s_size = (2 * params.patch_radius + 1) * (2 * params.patch_radius + 1);
821        // The chroma scale of 1.5 applies per channel, and each channel
822        // keeps its own sigma rather than sharing one.
823        let expected = 2.0 * 1.5 * (sigma_u * sigma_u + sigma_v * sigma_v) * s_size as f32;
824
825        let got = params.noise_offset_with(Some(&[sigma_u, sigma_v]));
826        assert!((got - expected).abs() < 1e-9, "expected {expected}, got {got}");
827    }
828
829    #[test]
830    fn sigma_eff_is_rms_over_active_channels() {
831        let sigmas = [3.0 / 255.0, 4.0 / 255.0];
832        let got = sigma_eff(&sigmas, ChannelMode::Chroma);
833        let expected = ((sigmas[0] * sigmas[0] + sigmas[1] * sigmas[1]) / 2.0).sqrt();
834        assert!((got - expected).abs() < 1e-9, "expected {expected}, got {got}");
835    }
836
837    #[test]
838    fn validate_rejects_non_positive_pilot_strength_scale() {
839        let zero = NlmParams {
840            prefilter: PrefilterMode::NlmSpatial { strength_scale: 0.0 },
841            ..NlmParams::default()
842        };
843        assert!(zero.validate().is_err());
844
845        let nan = NlmParams {
846            prefilter: PrefilterMode::NlmSpatial {
847                strength_scale: f32::NAN,
848            },
849            ..NlmParams::default()
850        };
851        assert!(nan.validate().is_err());
852    }
853
854    #[test]
855    fn validate_rejects_pilot_with_patch_radius_above_separable_threshold() {
856        let params = NlmParams {
857            prefilter: PrefilterMode::NlmSpatial { strength_scale: 1.0 },
858            patch_radius: SEPARABLE_THRESHOLD + 1,
859            ..NlmParams::default()
860        };
861        assert!(params.validate().is_err());
862    }
863
864    #[test]
865    fn validate_accepts_pilot_within_limits() {
866        let params = NlmParams {
867            prefilter: PrefilterMode::NlmSpatial { strength_scale: 1.0 },
868            patch_radius: SEPARABLE_THRESHOLD,
869            ..NlmParams::default()
870        };
871        assert!(params.validate().is_ok());
872    }
873
874    #[test]
875    fn validate_rejects_non_positive_bilateral_sigma_r() {
876        // A sigma_r of 0 makes inv_two_sigma_r_sq infinite. The centre
877        // tap's range_sq is 0, and 0 times infinity is NaN, which
878        // poisons every pixel of the reference image.
879        let params = NlmParams {
880            prefilter: PrefilterMode::Bilateral {
881                sigma_s: 3.0,
882                sigma_r: 0.0,
883            },
884            ..NlmParams::default()
885        };
886        assert!(params.validate().is_err());
887
888        let negative = NlmParams {
889            prefilter: PrefilterMode::Bilateral {
890                sigma_s: 3.0,
891                sigma_r: -0.02,
892            },
893            ..NlmParams::default()
894        };
895        assert!(negative.validate().is_err());
896
897        let nan = NlmParams {
898            prefilter: PrefilterMode::Bilateral {
899                sigma_s: 3.0,
900                sigma_r: f32::NAN,
901            },
902            ..NlmParams::default()
903        };
904        assert!(nan.validate().is_err());
905
906        let inf = NlmParams {
907            prefilter: PrefilterMode::Bilateral {
908                sigma_s: 3.0,
909                sigma_r: f32::INFINITY,
910            },
911            ..NlmParams::default()
912        };
913        assert!(inf.validate().is_err());
914    }
915
916    #[test]
917    fn validate_rejects_non_positive_bilateral_sigma_s() {
918        // A sigma_s of 0 makes inv_two_sigma_s_sq infinite. The centre
919        // tap's spatial_dist_sq is 0, so the spatial term is poisoned by
920        // the same 0 times infinity NaN.
921        let params = NlmParams {
922            prefilter: PrefilterMode::Bilateral {
923                sigma_s: 0.0,
924                sigma_r: 0.02,
925            },
926            ..NlmParams::default()
927        };
928        assert!(params.validate().is_err());
929
930        let negative = NlmParams {
931            prefilter: PrefilterMode::Bilateral {
932                sigma_s: -3.0,
933                sigma_r: 0.02,
934            },
935            ..NlmParams::default()
936        };
937        assert!(negative.validate().is_err());
938
939        let nan = NlmParams {
940            prefilter: PrefilterMode::Bilateral {
941                sigma_s: f32::NAN,
942                sigma_r: 0.02,
943            },
944            ..NlmParams::default()
945        };
946        assert!(nan.validate().is_err());
947
948        let inf = NlmParams {
949            prefilter: PrefilterMode::Bilateral {
950                sigma_s: f32::INFINITY,
951                sigma_r: 0.02,
952            },
953            ..NlmParams::default()
954        };
955        assert!(inf.validate().is_err());
956    }
957
958    #[test]
959    fn validate_accepts_positive_finite_bilateral_sigmas() {
960        let params = NlmParams {
961            prefilter: PrefilterMode::Bilateral {
962                sigma_s: 3.0,
963                sigma_r: 0.02,
964            },
965            ..NlmParams::default()
966        };
967        assert!(params.validate().is_ok());
968    }
969
970    #[test]
971    fn validate_accepts_a_small_positive_bilateral_sigma_at_the_boundary() {
972        // Pins the guard to `<= 0.0` rather than `< 0.0`. A value that
973        // is small but strictly positive, and far enough from the f32
974        // underflow cliff that squaring it stays a normal float, has to
975        // be accepted for either field on its own.
976        //
977        // 1e-6 squares to 1e-12, nowhere near the smallest normal f32 of
978        // about 1.18e-38, so `inv_two_sigma_sq` stays finite here.
979        let safe_small = 1e-6_f32;
980        assert!(
981            (safe_small * safe_small).is_normal(),
982            "the test value itself must not underflow"
983        );
984
985        let small_sigma_s = NlmParams {
986            prefilter: PrefilterMode::Bilateral {
987                sigma_s: safe_small,
988                sigma_r: 0.02,
989            },
990            ..NlmParams::default()
991        };
992        assert!(small_sigma_s.validate().is_ok());
993
994        let small_sigma_r = NlmParams {
995            prefilter: PrefilterMode::Bilateral {
996                sigma_s: 3.0,
997                sigma_r: safe_small,
998            },
999            ..NlmParams::default()
1000        };
1001        assert!(small_sigma_r.validate().is_ok());
1002    }
1003
1004    #[test]
1005    fn validate_rejects_a_subnormal_bilateral_sigma_that_underflows_on_squaring() {
1006        // `f32::MIN_POSITIVE`, the smallest normal positive f32 at about
1007        // 1.1754944e-38, is finite and above 0, so a guard that only
1008        // checked the raw value let it through.
1009        //
1010        // Squaring it underflows to exactly 0.0 in f32, because its true
1011        // square of about 1.38e-76 is far below the smallest subnormal
1012        // of about 1.4e-45. `inv_two_sigma_sq` then divides by zero and
1013        // returns infinity.
1014        //
1015        // That is the same NaN poisoning this validation exists to
1016        // prevent, reached through a sigma other than exactly 0.0.
1017        let sq = f32::MIN_POSITIVE * f32::MIN_POSITIVE;
1018        assert_eq!(sq, 0.0, "this test assumes MIN_POSITIVE underflows on squaring");
1019        let inv = prefilter::inv_two_sigma_sq(f32::MIN_POSITIVE);
1020        assert!(
1021            !inv.is_finite(),
1022            "this test assumes the derived factor is infinite here"
1023        );
1024
1025        let sigma_s = NlmParams {
1026            prefilter: PrefilterMode::Bilateral {
1027                sigma_s: f32::MIN_POSITIVE,
1028                sigma_r: 0.02,
1029            },
1030            ..NlmParams::default()
1031        };
1032        assert!(
1033            sigma_s.validate().is_err(),
1034            "a subnormal sigma_s that underflows to an infinite normalisation factor must be rejected"
1035        );
1036
1037        let sigma_r = NlmParams {
1038            prefilter: PrefilterMode::Bilateral {
1039                sigma_s: 3.0,
1040                sigma_r: f32::MIN_POSITIVE,
1041            },
1042            ..NlmParams::default()
1043        };
1044        assert!(
1045            sigma_r.validate().is_err(),
1046            "a subnormal sigma_r that underflows to an infinite normalisation factor must be rejected"
1047        );
1048    }
1049
1050    /// A `sigma_s` of 16.0 gives a bilateral radius of 32, which is the
1051    /// worked example in `prefilter.rs` and well past the maximum of 22.
1052    #[test]
1053    fn validate_rejects_bilateral_sigma_s_above_the_smem_ceiling() {
1054        let params = NlmParams {
1055            prefilter: PrefilterMode::Bilateral {
1056                sigma_s: 16.0,
1057                sigma_r: 0.02,
1058            },
1059            ..NlmParams::default()
1060        };
1061        let err = params.validate().expect_err("radius 32 exceeds the 22 ceiling");
1062        assert!(
1063            err.to_string().contains("sigma_s"),
1064            "error should name the field, got {err}"
1065        );
1066    }
1067
1068    /// A `sigma_s` of 1e9 overflows the tile-size arithmetic if it ever
1069    /// reaches the kernel launch.
1070    ///
1071    /// Validation has to reject it long before that, through the same
1072    /// radius check any other oversized `sigma_s` hits.
1073    #[test]
1074    fn validate_rejects_extreme_bilateral_sigma_s() {
1075        let params = NlmParams {
1076            prefilter: PrefilterMode::Bilateral {
1077                sigma_s: 1e9,
1078                sigma_r: 0.02,
1079            },
1080            ..NlmParams::default()
1081        };
1082        assert!(params.validate().is_err());
1083    }
1084
1085    /// The boundary pair for [`MAX_BILATERAL_RADIUS`], written as
1086    /// literal `sigma_s` values rather than derived from the constant.
1087    ///
1088    /// A `sigma_s` of 11.0 gives a radius of 22, right at the ceiling,
1089    /// so it is accepted. A `sigma_s` of 11.01 gives 23, one past it, so
1090    /// it is rejected.
1091    #[test]
1092    fn validate_accepts_bilateral_sigma_s_at_the_smem_ceiling() {
1093        let params = NlmParams {
1094            prefilter: PrefilterMode::Bilateral {
1095                sigma_s: 11.0,
1096                sigma_r: 0.02,
1097            },
1098            ..NlmParams::default()
1099        };
1100        assert!(params.validate().is_ok());
1101    }
1102
1103    #[test]
1104    fn validate_rejects_bilateral_sigma_s_just_above_the_smem_ceiling() {
1105        let params = NlmParams {
1106            prefilter: PrefilterMode::Bilateral {
1107                sigma_s: 11.01,
1108                sigma_r: 0.02,
1109            },
1110            ..NlmParams::default()
1111        };
1112        assert!(params.validate().is_err());
1113    }
1114
1115    #[test]
1116    fn sigma_eff_ignores_channels_past_the_mode_count() {
1117        // Luma only reads the first element, even when handed extra
1118        // chroma samples.
1119        let sigmas = [6.0 / 255.0, 100.0 / 255.0, 200.0 / 255.0];
1120        let got = sigma_eff(&sigmas, ChannelMode::Luma);
1121        assert!(
1122            (got - sigmas[0]).abs() < 1e-9,
1123            "expected {}, got {got}",
1124            sigmas[0]
1125        );
1126    }
1127
1128    #[test]
1129    fn hq_default_strength_matches_the_measured_luma_table() {
1130        const EXPECTED: [f32; 9] = [0.45, 0.45, 0.42, 0.42, 0.35, 0.35, 0.35, 0.30, 0.30];
1131        for (radius, &expected) in EXPECTED.iter().enumerate() {
1132            let got = hq_default_strength(ChannelMode::Luma, radius as u32);
1133            assert!(
1134                (got - expected).abs() < f32::EPSILON,
1135                "at radius {radius} expected {expected}, got {got}"
1136            );
1137        }
1138    }
1139
1140    #[test]
1141    fn hq_default_strength_matches_the_measured_chroma_table() {
1142        const EXPECTED: [f32; 9] = [1.00, 0.85, 0.70, 0.70, 0.70, 0.70, 0.70, 0.70, 0.70];
1143        for (radius, &expected) in EXPECTED.iter().enumerate() {
1144            let got = hq_default_strength(ChannelMode::Chroma, radius as u32);
1145            assert!(
1146                (got - expected).abs() < f32::EPSILON,
1147                "at radius {radius} expected {expected}, got {got}"
1148            );
1149        }
1150    }
1151
1152    #[test]
1153    fn hq_default_strength_yuv_reads_the_luma_table() {
1154        for radius in 0..=8u32 {
1155            let yuv = hq_default_strength(ChannelMode::Yuv, radius);
1156            let luma = hq_default_strength(ChannelMode::Luma, radius);
1157            assert!(
1158                (yuv - luma).abs() < f32::EPSILON,
1159                "at radius {radius} yuv is {yuv} but luma is {luma}"
1160            );
1161        }
1162    }
1163
1164    #[test]
1165    fn validate_dimensions_rejects_frames_below_the_minimum() {
1166        assert!(validate_dimensions(2, 64).is_err());
1167        assert!(validate_dimensions(64, 2).is_err());
1168        assert!(validate_dimensions(0, 0).is_err());
1169    }
1170
1171    #[test]
1172    fn validate_dimensions_accepts_the_minimum() {
1173        assert!(validate_dimensions(MIN_FRAME_DIM, MIN_FRAME_DIM).is_ok());
1174        assert!(validate_dimensions(1920, 1080).is_ok());
1175    }
1176
1177    #[test]
1178    fn hq_default_strength_clamps_radius_above_the_table() {
1179        let at_max = hq_default_strength(ChannelMode::Luma, MAX_TEMPORAL_RADIUS);
1180        let above_max = hq_default_strength(ChannelMode::Luma, MAX_TEMPORAL_RADIUS + 5);
1181        assert!(
1182            (at_max - above_max).abs() < f32::EPSILON,
1183            "expected clamping to hold the last table entry, got {at_max} vs {above_max}"
1184        );
1185    }
1186}