helena 0.1.0

Core types and component interfaces for helena, a latent data-to-waveform generation platform.
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134

use proptest::prelude::*;

use super::*;

proptest! {
    #[test]
    fn variance_is_preserved(gamma in 0.0f32..=1.0) {
        let level = NoiseLevel::from_signal_variance(gamma).unwrap();
        let sum = level.signal().powi(2) + level.noise().powi(2);
        prop_assert!((sum - 1.0).abs() < 1e-4, "signal^2 + noise^2 = {sum}");
    }

    #[test]
    fn velocity_round_trip_prop(
        gamma in 0.0f32..=1.0,
        xs in prop::collection::vec(-50.0f32..50.0, 1..24),
        ns in prop::collection::vec(-50.0f32..50.0, 1..24),
    ) {
        let dim = xs.len().min(ns.len());
        let x0 = tensor([dim, 1], xs[..dim].to_vec());
        let eps = tensor([dim, 1], ns[..dim].to_vec());
        let level = NoiseLevel::from_signal_variance(gamma).unwrap();

        let x_t = level.diffuse(&x0, &eps).unwrap();
        let v = level.velocity_target(&x0, &eps).unwrap();
        let x0_hat = level.signal_from_velocity(&x_t, &v).unwrap();
        let eps_hat = level.noise_from_velocity(&x_t, &v).unwrap();

        for (a, b) in x0.data().iter().zip(x0_hat.data()) {
            let tol = 1e-3 * (1.0 + a.abs());
            prop_assert!((a - b).abs() <= tol, "x0 {a} vs {b}");
        }
        for (a, b) in eps.data().iter().zip(eps_hat.data()) {
            let tol = 1e-3 * (1.0 + a.abs());
            prop_assert!((a - b).abs() <= tol, "eps {a} vs {b}");
        }
    }

    #[test]
    fn schedule_monotone_prop(offset in 0.0f32..0.1, a in 0.0f32..=1.0, b in 0.0f32..=1.0) {
        let (lo, hi) = if a <= b { (a, b) } else { (b, a) };
        let schedule = CosineSchedule::new(offset).unwrap();
        let g_lo = schedule.signal_variance(lo).unwrap();
        let g_hi = schedule.signal_variance(hi).unwrap();
        prop_assert!((0.0..=1.0).contains(&g_lo) && (0.0..=1.0).contains(&g_hi));
        prop_assert!(g_lo + 1e-6 >= g_hi, "gamma({lo})={g_lo} < gamma({hi})={g_hi}");
    }

    /// The v-corrected weight times `SNR + 1` recovers the Min-SNR clip on the
    /// clean-target error: `w_v(gamma) * (SNR + 1) = min(SNR, clip)`.
    #[test]
    fn min_snr_weight_yields_clipped_clean_target_weight(
        gamma in 1e-3f32..=1.0 - 1e-3,
        clip in 0.1f32..50.0,
    ) {
        let level = NoiseLevel::from_signal_variance(gamma).unwrap();
        let clip = MinSnrGamma::new(clip).unwrap();
        let snr = gamma / (1.0 - gamma);
        let effective = level.velocity_min_snr_weight(clip) * (snr + 1.0);
        let expected = snr.min(clip.get());
        let tol = 1e-3 * (1.0 + expected);
        prop_assert!(
            (effective - expected).abs() <= tol,
            "effective {effective} vs min(SNR, clip) {expected}"
        );
    }

    /// The weight is bounded in `[0, clip / (1 + clip)] ⊂ [0, 1)` for every level.
    #[test]
    fn min_snr_weight_is_bounded(gamma in 0.0f32..=1.0, clip in 0.1f32..50.0) {
        let level = NoiseLevel::from_signal_variance(gamma).unwrap();
        let clip = MinSnrGamma::new(clip).unwrap();
        let w = level.velocity_min_snr_weight(clip);
        let peak = clip.get() / (1.0 + clip.get());
        prop_assert!(w >= 0.0, "weight {w} negative");
        prop_assert!(w <= peak + 1e-6, "weight {w} exceeds peak {peak}");
        prop_assert!(w <= gamma + 1e-6 && w <= clip.get() * (1.0 - gamma) + 1e-6);
    }

    /// A logit-normal draw is always a valid schedule time: it lands in `[0, 1]`
    /// for every finite input and so never trips the schedule's range check.
    #[test]
    fn logit_normal_sample_is_a_valid_time(
        location in -10.0f32..10.0,
        scale in 1e-2f32..10.0,
        n in -10.0f32..10.0,
    ) {
        let density = LogitNormalDensity::new(location, scale).unwrap();
        let t = density.sample(n);
        prop_assert!((0.0..=1.0).contains(&t), "t={t} out of [0, 1]");
        prop_assert!(CosineSchedule::default().level(t).is_ok());
    }

    /// The draw is order-preserving in the standard-normal input: a monotone
    /// reparameterization of the schedule time, never a folding of it.
    #[test]
    fn logit_normal_sample_is_monotone(
        location in -5.0f32..5.0,
        scale in 1e-2f32..5.0,
        a in -8.0f32..8.0,
        b in -8.0f32..8.0,
    ) {
        let (lo, hi) = if a <= b { (a, b) } else { (b, a) };
        let density = LogitNormalDensity::new(location, scale).unwrap();
        prop_assert!(density.sample(lo) <= density.sample(hi) + 1e-6);
    }

    /// A scheduled guidance never leaves the segment between the neutral
    /// pass-through `1.0` and the run's base scale, so scheduling only ever
    /// attenuates toward neutral. This is what keeps the sampler's unconditional
    /// pass gated on the base scale alone, unchanged by the schedule.
    #[test]
    fn guidance_schedule_stays_on_the_neutral_to_base_segment(
        gamma in 0.0f32..=1.0,
        base in 0.0f32..8.0,
        lo in 0.0f32..0.99,
        span in 5e-3f32..1.0,
    ) {
        let hi = (lo + span).min(1.0);
        let level = NoiseLevel::from_signal_variance(gamma).unwrap();
        let band = GuidanceInterval::new(lo, hi).unwrap();
        let (min, max) = if base <= 1.0 { (base, 1.0) } else { (1.0, base) };
        for schedule in [
            GuidanceSchedule::Constant,
            GuidanceSchedule::Monotone,
            GuidanceSchedule::Interval(band),
        ] {
            let w = schedule.guidance_at(level, base);
            prop_assert!(w >= min - 1e-6 && w <= max + 1e-6, "{w} outside [{min}, {max}]");
        }
        // Constant is the base scale exactly, the byte-identical replay guarantee.
        prop_assert_eq!(GuidanceSchedule::Constant.guidance_at(level, base), base);
    }
}